multiple_buckets_repartitioning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook to re-partition the historical data of SOS reports into the final destination: DH-SECURE-SOSREPORTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pyspark==2.4.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "import os\n",
    "import pyspark\n",
    "import pandas\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql import functions as f\n",
    "from pyspark.sql import types as t\n",
    "import socket\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from source_buckets import source_buckets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['PYSPARK_PYTHON'] = '/opt/app-root/bin/python3'\n",
    "# os.environ['PYSPARK_DRIVER_PYTHON'] = '/opt/app-root/bin/python3'\n",
    "# spark.jars.ivy={os.environ['HOME']}\n",
    "SPARK_CLUSTER = 'spark://172.44.13.154:7077'\n",
    "S3_ENDPOINT = 'https://s3.upshift.redhat.com/'\n",
    "SPARK_APP_NAME = f'repartition-multiple-buckets-{datetime.datetime.now().strftime(\"%Y-%m-%d %H:%M\")}'\n",
    "HOSTNAME = socket.gethostbyname(socket.gethostname())\n",
    "print('Spark Cluster: {}'.format(SPARK_CLUSTER))\n",
    "print('S3 endpoint: {}'.format(S3_ENDPOINT))\n",
    "print('Spark App Name: {}'.format(SPARK_APP_NAME))\n",
    "print('Hostname: {}'.format(HOSTNAME))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_spark_config(spark_cluster, executor_memory='16g', executor_cores='4', max_cores='16'):\n",
    "    print('Spark cluster is: {}'.format(spark_cluster))\n",
    "    sc_conf = (\n",
    "        pyspark.SparkConf().setMaster(spark_cluster) \\\n",
    "        .set('spark.driver.host', HOSTNAME) \\\n",
    "        .set('spark.driver.port', 42000) \\\n",
    "        .set('spark.driver.bindAddress', '0.0.0.0') \\\n",
    "        .set('spark.driver.blockManager.port', 42100) \\\n",
    "        .set('spark.executor.cores', '3') \\\n",
    "        .set('spark.executor.memory', '4500M') \\\n",
    "        .set('spark.driver.memory', '4G') \\\n",
    "        .set('spark.sql.parquet.enableVectorizedReader', True) \\ # Turn this to False if low on Memory or faving OOM issues\n",
    "        .set('spark.kubernetes.memoryOverheadFactor', '0.20')\n",
    "    )\n",
    "    return sc_conf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_spark():\n",
    "    spark_config = create_spark_config(SPARK_CLUSTER)\n",
    "    print('spark_config is: {}'.format(spark_config))\n",
    "    print(\"Creating Spark Session at cluster: {}\".format(SPARK_CLUSTER))\n",
    "    spark = SparkSession.builder.appName(SPARK_APP_NAME).enableHiveSupport().config(conf=spark_config).getOrCreate()\n",
    "    spark.sparkContext.setLogLevel('ERROR')\n",
    "    hadoopConf = spark.sparkContext._jsc.hadoopConfiguration()\n",
    "    hadoopConf.set('fs.s3a.endpoint', S3_ENDPOINT)\n",
    "    hadoopConf.set('fs.s3a.path.style.access', 'true')\n",
    "    hadoopConf.set('fs.s3a.access.key', os.environ.get('AWS_ACCESS_KEY_ID'))\n",
    "    hadoopConf.set('fs.s3a.secret.key', os.environ.get('AWS_SECRET_ACCESS_KEY'))\n",
    "    hadoopConf.set('fs.s3a.impl', 'org.apache.hadoop.fs.s3a.S3AFileSystem')\n",
    "    print(\"hadoop is configured!\")\n",
    "    return spark"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Change the path of the source tables in the `if` block, for the destination buckets in the `else` blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def form_path_string(bucket_name, repartitioned_data=False):\n",
    "    table_name = bucket_name[3:-4].lower().replace('-', '_')\n",
    "    print('table_name: {}'.format(table_name))\n",
    "    if not repartitioned_data:\n",
    "        path_string = 's3a://{}/extraction/sos/parquet/{}/'.format(bucket_name, table_name)\n",
    "    else:\n",
    "        path_string = 's3a://{}/extraction/sos/parquet/{}/'.format('DH-SECURE-SOSREPORTS', table_name)\n",
    "    return path_string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_dataframe_from_bucket(bucket_name, repartitioned_data=False):\n",
    "    src_path = form_path_string(bucket_name, repartitioned_data)\n",
    "    print('src_path: {}'.format(src_path))\n",
    "    df = spark.read.parquet(f'{src_path}')\n",
    "    count = 0 #df.count()\n",
    "    distinct_count = 0 #df.distinct().count()\n",
    "    num_partitions = df.rdd.getNumPartitions()\n",
    "    print('count: {}'.format(count))\n",
    "    print('unique_count: {}'.format(distinct_count))\n",
    "    print('partitions: {}'.format(num_partitions))\n",
    "    return (df, count, distinct_count, num_partitions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function below takes the source dataframe, re-partitions it in a way that ensures one object per partition and saves it at the destination. Change the mode to 'overwrite' or 'append' depending on the use case.\n",
    "For incoming new data, use append."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_repartitioned_dataframe(bucket_name, df):\n",
    "    dest_path = form_path_string(bucket_name, repartitioned_data=True)\n",
    "    print('Trying to save repartitioned data at: {}'.format(dest_path))\n",
    "    df.repartition(1, \"created_year\", \"created_month\", \"created_day\").write.partitionBy(\n",
    "        \"created_year\", \"created_month\", \"created_day\").mode('append').parquet(dest_path)\n",
    "    print('Data repartitioning complete with at the following location: ')\n",
    "    print(dest_path)\n",
    "    _, count, distinct_count, num_partitions = read_dataframe_from_bucket(bucket_name, repartitioned_data=True)\n",
    "    return count, distinct_count, num_partitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    spark.stop()\n",
    "    spark = setup_spark()\n",
    "except:\n",
    "    spark = setup_spark()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_failed_bucket_to_file(bucket_name, error):\n",
    "    print(bucket_name)\n",
    "    print(error)\n",
    "    file_name = bucket_name + '.txt'\n",
    "    with open(file_name, 'a') as file_:\n",
    "        file_.write('\\n')\n",
    "        file_.write('----------{}----------\\n'.format(bucket_name))\n",
    "#         file_.write(error)\n",
    "#         file_.writeline('\\n----------------------------------------------')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for count, bucket_name in enumerate(source_buckets):\n",
    "    try:\n",
    "        print('trying to partition {}th bucket: {}'.format(count+1, bucket_name))\n",
    "        initial_df, initial_count, initial_distinct_count, initial_num_partitions = read_dataframe_from_bucket(bucket_name)\n",
    "        final_count, final_distinct_count, final_num_partitions = save_repartitioned_dataframe(bucket_name, initial_df)\n",
    "        print('partition count difference: {}'.format(final_num_partitions-initial_num_partitions))\n",
    "        if final_distinct_count != initial_distinct_count:\n",
    "            data_lost = initial_distinct_count - final_distinct_count\n",
    "            print('Data lost: {}'.format(data_lost))\n",
    "            with open('{}_data_list.txt'.format(bucket_name), 'w') as file_:\n",
    "                file_.write('\\n')\n",
    "                file_.write('{}'.format(data_lost))\n",
    "    except Exception as e:\n",
    "        print('Bucket failed: {}'.format(bucket_name))\n",
    "        save_failed_bucket_to_file(bucket_name, e)\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}