databricks_notebooks/process_kinesis_streaming_data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a03162a2-2d0e-47d5-81f2-0b62f7d2cf12",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# import pyspark functions\n",
    "from pyspark.sql.functions import *\n",
    "from pyspark.sql.types import StructType, StructField, IntegerType, StringType, TimestampType, FloatType\n",
    "# import URL processing\n",
    "import urllib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AWS Credentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "86b63fb1-5aac-4f1a-bfb1-ba5a7ae6b4ee",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# list tables in filestore to get name of credentials file\n",
    "dbutils.fs.ls(\"/FileStore/tables\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "529a0bc3-adb1-4b3f-aadd-fd278bb45533",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Read the CSV file to spark dataframe, passing in options for the header row and separator\n",
    "aws_keys_df = spark.read.format(\"csv\")\\\n",
    ".option(\"header\", \"true\")\\\n",
    ".option(\"sep\", \",\")\\\n",
    ".load(\"/FileStore/tables/authentication_credentials.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "6ee1af56-daba-44c5-bd34-5265934a3316",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Get the AWS access key and secret key from the spark dataframe\n",
    "ACCESS_KEY = aws_keys_df.where(col('User name')=='databricks-user').select('Access key ID').collect()[0]['Access key ID']\n",
    "SECRET_KEY = aws_keys_df.where(col('User name')=='databricks-user').select('Secret access key').collect()[0]['Secret access key']\n",
    "# Encode the secret key\n",
    "ENCODED_SECRET_KEY = urllib.parse.quote(string=SECRET_KEY, safe=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "02664ec8-e540-4701-a287-c947297303af",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def get_stream(stream_name: str):\n",
    "    '''Uses spark.readStream to retrieve Kinesis stream and returns stream as dataframe'''\n",
    "    dataframe = spark \\\n",
    "    .readStream \\\n",
    "    .format('kinesis') \\\n",
    "    .option('streamName', stream_name) \\\n",
    "    .option('initialPosition','earliest') \\\n",
    "    .option('region','us-east-1') \\\n",
    "    .option('awsAccessKey', ACCESS_KEY) \\\n",
    "    .option('awsSecretKey', SECRET_KEY) \\\n",
    "    .load()\n",
    "    return dataframe\n",
    "\n",
    "def deserialize_stream(stream, schema):\n",
    "    '''Takes stream dataframe and schema, deserializes data from stream and returns data as dataframe'''\n",
    "    dataframe = stream \\\n",
    "    .selectExpr(\"CAST(data as STRING)\") \\\n",
    "    .withColumn(\"data\", from_json(col(\"data\"), schema)) \\\n",
    "    .select(col(\"data.*\"))\n",
    "    return dataframe\n",
    "\n",
    "def add_nulls_to_dataframe_column(dataframe, column, value_to_replace):\n",
    "    '''Converts matched values in column of dataframe to null based on expression'''\n",
    "    dataframe = dataframe.withColumn(column, when(col(column).like(value_to_replace), None).otherwise(col(column)))\n",
    "    return dataframe\n",
    "\n",
    "def write_stream_df_to_table(dataframe, name: str):\n",
    "    '''Takes dataframe and name string and writes dataframe to delta table using name in options and table name'''\n",
    "    dataframe.writeStream \\\n",
    "    .format(\"delta\") \\\n",
    "    .outputMode(\"append\") \\\n",
    "    .option(\"checkpointLocation\", f\"/tmp/kinesis/1215be80977f_{name}_table_checkpoints/\") \\\n",
    "    .table(f\"1215be80977f_{name}_table\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Schema definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "bc0fb2aa-6563-483c-9d34-8511509858bc",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# define schemas for each of the dataframes\n",
    "pin_schema = StructType([\n",
    "    StructField(\"index\", IntegerType()),\n",
    "    StructField(\"unique_id\", StringType()),\n",
    "    StructField(\"title\", StringType()),\n",
    "    StructField(\"description\", StringType()),\n",
    "    StructField(\"poster_name\", StringType()),\n",
    "    StructField(\"follower_count\", StringType()),\n",
    "    StructField(\"tag_list\", StringType()),\n",
    "    StructField(\"is_image_or_video\", StringType()),\n",
    "    StructField(\"image_src\", StringType()),\n",
    "    StructField(\"downloaded\", IntegerType()),\n",
    "    StructField(\"save_location\", StringType()),\n",
    "    StructField(\"category\", StringType())\n",
    "])\n",
    "geo_schema = StructType([\n",
    "    StructField(\"ind\", IntegerType()),\n",
    "    StructField(\"timestamp\", TimestampType()),\n",
    "    StructField(\"latitude\", FloatType()),\n",
    "    StructField(\"longitude\", FloatType()),\n",
    "    StructField(\"country\", StringType())\n",
    "])\n",
    "user_schema = StructType([\n",
    "    StructField(\"ind\", IntegerType()),\n",
    "    StructField(\"first_name\", StringType()),\n",
    "    StructField(\"last_name\", StringType()),\n",
    "    StructField(\"age\", StringType()),\n",
    "    StructField(\"date_joined\", TimestampType())\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get streams from Kinesis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d7390da9-7815-4a3c-92b0-24556592b227",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "pin_stream = get_stream('streaming-1215be80977f-pin')\n",
    "geo_stream = get_stream('streaming-1215be80977f-geo')\n",
    "user_stream = get_stream('streaming-1215be80977f-user')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deserialize streams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "dc6dddc5-c174-48a2-b945-e02631a5aed2",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "df_pin = deserialize_stream(pin_stream, pin_schema)\n",
    "df_geo = deserialize_stream(geo_stream, geo_schema)\n",
    "df_user = deserialize_stream(user_stream, user_schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "182e9c7e-90ee-4d47-b5a4-35412a71d00a",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# replace empty entries and entries with no relevant data in each column with Nones\n",
    "# column names and values to change to null\n",
    "columns_and_values_for_null = {\n",
    "    \"description\": \"No description available%\",\n",
    "    \"follower_count\": \"User Info Error\",\n",
    "    \"image_src\": \"Image src error.\",\n",
    "    \"poster_name\": \"User Info Error\",\n",
    "    \"tag_list\": \"N,o, ,T,a,g,s, ,A,v,a,i,l,a,b,l,e\",\n",
    "    \"title\": \"No Title Data Available\"\n",
    "}\n",
    "# loop through dictionary, calling function with dictionary values as arguments\n",
    "for key, value in columns_and_values_for_null.items():\n",
    "    df_pin = add_nulls_to_dataframe_column(df_pin, key, value)\n",
    "# Perform the necessary transformations on the follower_count to ensure every entry is a number\n",
    "df_pin = df_pin.withColumn(\"follower_count\", regexp_replace(\"follower_count\", \"k\", \"000\"))\n",
    "df_pin = df_pin.withColumn(\"follower_count\", regexp_replace(\"follower_count\", \"M\", \"000000\"))\n",
    "# cast follower_count column to integer type\n",
    "df_pin = df_pin.withColumn(\"follower_count\", col(\"follower_count\").cast('int'))\n",
    "# convert save_location column to include only the save location path\n",
    "df_pin = df_pin.withColumn(\"save_location\", regexp_replace(\"save_location\", \"Local save in \", \"\"))\n",
    "# rename the index column to ind\n",
    "df_pin = df_pin.withColumnRenamed(\"index\", \"ind\")\n",
    "# reorder columns\n",
    "new_pin_column_order = [\n",
    "    \"ind\",\n",
    "    \"unique_id\",\n",
    "    \"title\",\n",
    "    \"description\",\n",
    "    \"follower_count\",\n",
    "    \"poster_name\",\n",
    "    \"tag_list\",\n",
    "    \"is_image_or_video\",\n",
    "    \"image_src\",\n",
    "    \"save_location\",\n",
    "    \"category\"\n",
    "]\n",
    "df_pin = df_pin.select(new_pin_column_order)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "71d1f5d7-6d33-4e31-8b36-86fc38ba4bad",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# import types\n",
    "from pyspark.sql.types import ArrayType, DoubleType\n",
    "# define function for returning list containing two values\n",
    "def combine_lat_and_long(latitude, longitude):\n",
    "    return [latitude, longitude]\n",
    "# define new user-defined function\n",
    "new_func = udf(combine_lat_and_long, ArrayType(DoubleType()))\n",
    "# apply new udf to combine latitude and longitude columns\n",
    "df_geo = df_geo.withColumn(\"coordinates\", new_func(\"latitude\", \"longitude\"))\n",
    "# drop the latitude and longitude columns\n",
    "cols_to_drop = (\"latitude\", \"longitude\")\n",
    "df_geo = df_geo.drop(*cols_to_drop)\n",
    "# convert timestamp column from type string to type timestamp\n",
    "df_geo = df_geo.withColumn(\"timestamp\", to_timestamp(\"timestamp\"))\n",
    "# change column order\n",
    "new_geo_column_order = [\n",
    "    \"ind\",\n",
    "    \"country\",\n",
    "    \"coordinates\",\n",
    "    \"timestamp\",\n",
    "]\n",
    "df_geo = df_geo.select(new_geo_column_order)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b88961d8-32db-40b4-8e6c-cdf4a4b431e7",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# create new column for full name\n",
    "df_user = df_user.withColumn(\"user_name\", concat_ws(\" \", \"first_name\", \"last_name\"))\n",
    "# drop the first_name and last_name columns\n",
    "cols_to_drop = (\"first_name\", \"last_name\")\n",
    "df_user = df_user.drop(*cols_to_drop)\n",
    "# convert date_joined column from type string to type timestamp\n",
    "df_user = df_user.withColumn(\"date_joined\", to_timestamp(\"date_joined\"))\n",
    "# change column order\n",
    "new_user_column_order = [\n",
    "    \"ind\",\n",
    "    \"user_name\",\n",
    "    \"age\",\n",
    "    \"date_joined\",\n",
    "]\n",
    "df_user = df_user.select(new_user_column_order)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Display data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "4f75477c-7a37-4f59-bae8-c06422eb1d85",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "display(df_pin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "ebf591ef-c60d-4304-98ac-bcfbb0e36376",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "display(df_geo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "3f08efc2-b770-4ecb-8644-4df8dd2bcee6",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "display(df_user)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write data to Delta tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "283594d1-3123-4a10-b03e-a14a3f752362",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "write_stream_df_to_table(df_pin, \"pin\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "ac97113e-f437-4467-a52d-78705c6b0704",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "write_stream_df_to_table(df_geo, \"geo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "935e4d34-f5b7-4910-a278-d2dcb89c6eaf",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "write_stream_df_to_table(df_user, \"user\")"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "Processing Kinesis Streaming Data",
   "widgets": {}
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}