databricks_notebooks/mount_s3_and_get_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": [
    {
     "data": {
      "text/html": [
       "<style scoped>\n",
       "  .ansiout {\n",
       "    display: block;\n",
       "    unicode-bidi: embed;\n",
       "    white-space: pre-wrap;\n",
       "    word-wrap: break-word;\n",
       "    word-break: break-all;\n",
       "    font-family: \"Source Code Pro\", \"Menlo\", monospace;;\n",
       "    font-size: 13px;\n",
       "    color: #555;\n",
       "    margin-left: 4px;\n",
       "    line-height: 19px;\n",
       "  }\n",
       "</style>\n",
       "<div class=\"ansiout\"></div>"
      ]
     },
     "metadata": {
      "application/vnd.databricks.v1+output": {
       "addedWidgets": {},
       "arguments": {},
       "data": "<div class=\"ansiout\"></div>",
       "datasetInfos": [],
       "metadata": {},
       "removedWidgets": [],
       "type": "html"
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# import pyspark functions\n",
    "from pyspark.sql.functions import *\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": [
    {
     "data": {
      "text/html": [
       "<style scoped>\n",
       "  .ansiout {\n",
       "    display: block;\n",
       "    unicode-bidi: embed;\n",
       "    white-space: pre-wrap;\n",
       "    word-wrap: break-word;\n",
       "    word-break: break-all;\n",
       "    font-family: \"Source Code Pro\", \"Menlo\", monospace;;\n",
       "    font-size: 13px;\n",
       "    color: #555;\n",
       "    margin-left: 4px;\n",
       "    line-height: 19px;\n",
       "  }\n",
       "</style>\n",
       "<div class=\"ansiout\">Out[2]: [FileInfo(path=&#39;dbfs:/FileStore/tables/authentication_credentials.csv&#39;, name=&#39;authentication_credentials.csv&#39;, size=202, modificationTime=1687110999000)]</div>"
      ]
     },
     "metadata": {
      "application/vnd.databricks.v1+output": {
       "addedWidgets": {},
       "arguments": {},
       "data": "<div class=\"ansiout\">Out[2]: [FileInfo(path=&#39;dbfs:/FileStore/tables/authentication_credentials.csv&#39;, name=&#39;authentication_credentials.csv&#39;, size=202, modificationTime=1687110999000)]</div>",
       "datasetInfos": [],
       "metadata": {},
       "removedWidgets": [],
       "type": "html"
      }
     },
     "output_type": "display_data"
    }
   ],
   "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": [
    {
     "data": {
      "text/html": [
       "<style scoped>\n",
       "  .ansiout {\n",
       "    display: block;\n",
       "    unicode-bidi: embed;\n",
       "    white-space: pre-wrap;\n",
       "    word-wrap: break-word;\n",
       "    word-break: break-all;\n",
       "    font-family: \"Source Code Pro\", \"Menlo\", monospace;;\n",
       "    font-size: 13px;\n",
       "    color: #555;\n",
       "    margin-left: 4px;\n",
       "    line-height: 19px;\n",
       "  }\n",
       "</style>\n",
       "<div class=\"ansiout\"></div>"
      ]
     },
     "metadata": {
      "application/vnd.databricks.v1+output": {
       "addedWidgets": {},
       "arguments": {},
       "data": "<div class=\"ansiout\"></div>",
       "datasetInfos": [],
       "metadata": {},
       "removedWidgets": [],
       "type": "html"
      }
     },
     "output_type": "display_data"
    }
   ],
   "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": [
    {
     "data": {
      "text/html": [
       "<style scoped>\n",
       "  .ansiout {\n",
       "    display: block;\n",
       "    unicode-bidi: embed;\n",
       "    white-space: pre-wrap;\n",
       "    word-wrap: break-word;\n",
       "    word-break: break-all;\n",
       "    font-family: \"Source Code Pro\", \"Menlo\", monospace;;\n",
       "    font-size: 13px;\n",
       "    color: #555;\n",
       "    margin-left: 4px;\n",
       "    line-height: 19px;\n",
       "  }\n",
       "</style>\n",
       "<div class=\"ansiout\"></div>"
      ]
     },
     "metadata": {
      "application/vnd.databricks.v1+output": {
       "addedWidgets": {},
       "arguments": {},
       "data": "<div class=\"ansiout\"></div>",
       "datasetInfos": [],
       "metadata": {},
       "removedWidgets": [],
       "type": "html"
      }
     },
     "output_type": "display_data"
    }
   ],
   "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": [
    "## Mount S3 bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "782d3e17-4834-4f8c-8a61-2e0efa859875",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style scoped>\n",
       "  .ansiout {\n",
       "    display: block;\n",
       "    unicode-bidi: embed;\n",
       "    white-space: pre-wrap;\n",
       "    word-wrap: break-word;\n",
       "    word-break: break-all;\n",
       "    font-family: \"Source Code Pro\", \"Menlo\", monospace;;\n",
       "    font-size: 13px;\n",
       "    color: #555;\n",
       "    margin-left: 4px;\n",
       "    line-height: 19px;\n",
       "  }\n",
       "</style>\n",
       "<div class=\"ansiout\">Out[5]: True</div>"
      ]
     },
     "metadata": {
      "application/vnd.databricks.v1+output": {
       "addedWidgets": {},
       "arguments": {},
       "data": "<div class=\"ansiout\">Out[5]: True</div>",
       "datasetInfos": [],
       "metadata": {},
       "removedWidgets": [],
       "type": "html"
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# AWS S3 bucket name\n",
    "AWS_S3_BUCKET = \"user-1215be80977f-bucket\"\n",
    "# Mount name for the bucket\n",
    "MOUNT_NAME = \"/mnt/user-1215be80977f-bucket\"\n",
    "# Source url\n",
    "SOURCE_URL = \"s3n://{0}:{1}@{2}\".format(ACCESS_KEY, ENCODED_SECRET_KEY, AWS_S3_BUCKET)\n",
    "# Mount the drive\n",
    "dbutils.fs.mount(SOURCE_URL, MOUNT_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "e905edd9-4080-461c-a9b9-370928aef20e",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style scoped>\n",
       "  .table-result-container {\n",
       "    max-height: 300px;\n",
       "    overflow: auto;\n",
       "  }\n",
       "  table, th, td {\n",
       "    border: 1px solid black;\n",
       "    border-collapse: collapse;\n",
       "  }\n",
       "  th, td {\n",
       "    padding: 5px;\n",
       "  }\n",
       "  th {\n",
       "    text-align: left;\n",
       "  }\n",
       "</style><div class='table-result-container'><table class='table-result'><thead style='background-color: white'><tr><th>path</th><th>name</th><th>size</th><th>modificationTime</th></tr></thead><tbody><tr><td>dbfs:/mnt/user-1215be80977f-bucket/topics/1215be80977f.geo/</td><td>1215be80977f.geo/</td><td>0</td><td>1692041344013</td></tr><tr><td>dbfs:/mnt/user-1215be80977f-bucket/topics/1215be80977f.pin/</td><td>1215be80977f.pin/</td><td>0</td><td>1692041344013</td></tr><tr><td>dbfs:/mnt/user-1215be80977f-bucket/topics/1215be80977f.user/</td><td>1215be80977f.user/</td><td>0</td><td>1692041344013</td></tr></tbody></table></div>"
      ]
     },
     "metadata": {
      "application/vnd.databricks.v1+output": {
       "addedWidgets": {},
       "aggData": [],
       "aggError": "",
       "aggOverflow": false,
       "aggSchema": [],
       "aggSeriesLimitReached": false,
       "aggType": "",
       "arguments": {},
       "columnCustomDisplayInfos": {},
       "data": [
        [
         "dbfs:/mnt/user-1215be80977f-bucket/topics/1215be80977f.geo/",
         "1215be80977f.geo/",
         0,
         1692041344013
        ],
        [
         "dbfs:/mnt/user-1215be80977f-bucket/topics/1215be80977f.pin/",
         "1215be80977f.pin/",
         0,
         1692041344013
        ],
        [
         "dbfs:/mnt/user-1215be80977f-bucket/topics/1215be80977f.user/",
         "1215be80977f.user/",
         0,
         1692041344013
        ]
       ],
       "datasetInfos": [],
       "dbfsResultPath": null,
       "isJsonSchema": true,
       "metadata": {},
       "overflow": false,
       "plotOptions": {
        "customPlotOptions": {},
        "displayType": "table",
        "pivotAggregation": null,
        "pivotColumns": null,
        "xColumns": null,
        "yColumns": null
       },
       "removedWidgets": [],
       "schema": [
        {
         "metadata": "{}",
         "name": "path",
         "type": "\"string\""
        },
        {
         "metadata": "{}",
         "name": "name",
         "type": "\"string\""
        },
        {
         "metadata": "{}",
         "name": "size",
         "type": "\"long\""
        },
        {
         "metadata": "{}",
         "name": "modificationTime",
         "type": "\"long\""
        }
       ],
       "type": "table"
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# list the topics stored on the mounted S3 bucket\n",
    "display(dbutils.fs.ls(\"/mnt/user-1215be80977f-bucket/topics\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read contents of S3 bucket into dataframes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "0c182ac9-cd6b-453e-ac7f-5a42b19dcf5a",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# list of topic suffixes\n",
    "topics = [\".pin\", \".geo\", \".user\"]\n",
    "\n",
    "def read_topics_into_dataframe(topic):\n",
    "    # create path to topic files\n",
    "    file_path = f\"/mnt/user-1215be80977f-bucket/topics/1215be80977f{topic}/partition=0/*.json\"\n",
    "    # specify file type\n",
    "    file_type = \"json\"\n",
    "    # Ask Spark to infer the schema\n",
    "    infer_schema = \"true\"\n",
    "    # load JSONs from mounted S3 bucket to Spark dataframe\n",
    "    df = spark.read.format(file_type) \\\n",
    "        .option(\"inferSchema\", infer_schema) \\\n",
    "        .load(file_path)\n",
    "    return df\n",
    "\n",
    "for item in topics:\n",
    "    # create statement strings for naming and displaying dataframes\n",
    "    make_df_statement = f\"df_{item[1:]} = read_topics_into_dataframe('{item}')\"\n",
    "    display_df_statement = f\"display(df_{item[1:]})\"\n",
    "    # execute statements\n",
    "    exec(make_df_statement)\n",
    "    exec(display_df_statement)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unmount S3 bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "98045be2-eea1-4d34-b60b-abb3a17c2a57",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style scoped>\n",
       "  .ansiout {\n",
       "    display: block;\n",
       "    unicode-bidi: embed;\n",
       "    white-space: pre-wrap;\n",
       "    word-wrap: break-word;\n",
       "    word-break: break-all;\n",
       "    font-family: \"Source Code Pro\", \"Menlo\", monospace;;\n",
       "    font-size: 13px;\n",
       "    color: #555;\n",
       "    margin-left: 4px;\n",
       "    line-height: 19px;\n",
       "  }\n",
       "</style>\n",
       "<div class=\"ansiout\">/mnt/user-1215be80977f-bucket has been unmounted.\n",
       "Out[8]: True</div>"
      ]
     },
     "metadata": {
      "application/vnd.databricks.v1+output": {
       "addedWidgets": {},
       "arguments": {},
       "data": "<div class=\"ansiout\">/mnt/user-1215be80977f-bucket has been unmounted.\nOut[8]: True</div>",
       "datasetInfos": [],
       "metadata": {},
       "removedWidgets": [],
       "type": "html"
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# unmount the bucket from the filestore\n",
    "dbutils.fs.unmount(\"/mnt/user-1215be80977f-bucket\")"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "Mount S3",
   "widgets": {}
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}