load_data/Load_Bank_transact_data.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import uuid\n",
        "from notebookutils import mssparkutils\n",
        "from pyspark.sql import DataFrame\n",
        "import pyspark.sql.functions as F\n",
        "from pyspark.sql.types import StringType,BooleanType\n",
        "# from graphframes import *\n",
        "\n",
        "f_uuid = F.udf(lambda: str(uuid.uuid4()), StringType())\n",
        "f_bool = F.udf(lambda: True, BooleanType())\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "# update the values from your infra\r\n",
        "cosmosEndpoint = \"https://cosmos-7w6s3xqb4piyi.documents.azure.com:443/\"\r\n",
        "cosmosMasterKey = \"cosmos-key\" \r\n",
        "cosmosDatabaseName = \"database01\"\r\n",
        "cosmosContainerName = \"graph01\""
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": true
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# update csv file path based on your infra\n",
        "df = spark.read.load('abfss://data@7w6s3xqb4piyi.dfs.core.windows.net/PS_20174392719_1491204439457_log.csv', format='csv',header=True)\n",
        "display(df.limit(10))"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "collapsed": false
      }
    },
    {
      "cell_type": "code",
      "source": [
        "raw_data = df.selectExpr(\"type\", \r\n",
        "                        \"cast(amount as int) amount\", \r\n",
        "                        \"nameOrig\", \r\n",
        "                        \"cast(oldbalanceOrg as int) oldbalanceOrg\", \r\n",
        "                        \"cast(newbalanceOrig as int) newbalanceOrig\",\r\n",
        "                        \"nameDest\",\r\n",
        "                        \"cast(oldbalanceDest as int) oldbalanceDest\", \r\n",
        "                        \"cast(newbalanceDest as int) newbalanceDest\"\r\n",
        ")                         \r\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "raw_data.show()"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cfg = {\n",
        "  \"spark.cosmos.accountEndpoint\" : cosmosEndpoint,\n",
        "  \"spark.cosmos.accountKey\" : cosmosMasterKey,\n",
        "  \"spark.cosmos.database\" : cosmosDatabaseName,\n",
        "  \"spark.cosmos.container\" : cosmosContainerName,\n",
        "}\n",
        "# Configure Catalog Api to be used\n",
        "spark.conf.set(\"spark.sql.catalog.cosmosCatalog\", \"com.azure.cosmos.spark.CosmosCatalog\")\n",
        "spark.conf.set(\"spark.sql.catalog.cosmosCatalog.spark.cosmos.accountEndpoint\", cosmosEndpoint)\n",
        "spark.conf.set(\"spark.sql.catalog.cosmosCatalog.spark.cosmos.accountKey\", cosmosMasterKey)\n",
        "spark.conf.set(\"spark.cosmos.throughputControl.enabled\",True)\n",
        "spark.conf.set(\"spark.cosmos.throughputControl.targetThroughput\",40000)\n",
        "\n",
        "def write_to_cosmos_graph(df: DataFrame, data_type: str, save: bool = False):\n",
        "    if (save):\n",
        "        df.write.format(\"delta\").mode(\"overwrite\").option(\"overwriteSchema\", \"true\").save(f\"synfs:/{job_id}/mydata/{data_type}/\")\n",
        "        \n",
        "    df.write\\\n",
        "   .format(\"cosmos.oltp\")\\\n",
        "   .options(**cfg)\\\n",
        "   .mode(\"APPEND\")\\\n",
        "   .save()"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# PySpark function to create veritces and edges dataframes in a format accepted by Cosmos SQL api from raw dataframe. \n",
        "# TODO: Add vertex properties \n",
        "def prepare_vertices_edge_df(\n",
        "    df: DataFrame,\n",
        "    source_col_name: str,\n",
        "    dest_col_name: str,\n",
        "    parition_key_col_name: str,\n",
        "    cosmos_parition_name: str,\n",
        "    edge_properties_col_name: list,\n",
        "    vertex_properties_col_name: list,\n",
        "    vertex_label: str = \"account\",\n",
        "    edge_label_col_name: str = \"type\",\n",
        "    sample: bool = False,\n",
        "):\n",
        "    if (sample):\n",
        "        df = df.limit(100)\n",
        "    nameOrig = df.select(source_col_name).withColumnRenamed(source_col_name, \"id\")\n",
        "    nameDest = df.select(dest_col_name).withColumnRenamed(dest_col_name, \"id\")\n",
        "    all_vertices = nameOrig.union(nameDest).distinct()\n",
        "    cosmos_vertices_df = (\n",
        "        all_vertices.withColumn(cosmos_parition_name, all_vertices[\"id\"])\n",
        "        .withColumn(\"label\", F.lit(vertex_label))\n",
        "        .select(\"label\", \"id\", cosmos_parition_name, *vertex_properties_col_name)\n",
        "        .distinct()\n",
        "    )\n",
        "    # Create dataframe with required columns\n",
        "    # _sink => target account => nameDest\n",
        "    # _sinkLabel => target label => \"account\"\n",
        "    # _vertexId => source account => nameOrig\n",
        "    # _vertexLabel => source label => \"account\"\n",
        "    # cosmos_parition_name => partition key defined in Cosmos => \"accountId\"\n",
        "    cosmos_edges_df = (\n",
        "        df.withColumn(\"id\", f_uuid())\n",
        "        .withColumn(cosmos_parition_name, df[parition_key_col_name])\n",
        "        .withColumn(\"label\", df[edge_label_col_name])\n",
        "        .withColumn(\"_sinkPartition\", df[dest_col_name])\n",
        "        .withColumn(\"_vertexId\", df[source_col_name])\n",
        "        .withColumn(\"_sink\", df[dest_col_name])\n",
        "        .withColumn(\"_sinkLabel\", F.lit(vertex_label))\n",
        "        .withColumn(\"_vertexLabel\", F.lit(vertex_label))\n",
        "        .withColumn(\"_isEdge\", f_bool())\n",
        "        .select(\n",
        "            \"id\",\n",
        "            \"label\",\n",
        "            \"_sink\",\n",
        "            \"_sinkLabel\",\n",
        "            \"_sinkPartition\",\n",
        "            \"_vertexId\",\n",
        "            \"_vertexLabel\",\n",
        "            \"_isEdge\",\n",
        "            cosmos_parition_name,\n",
        "            *edge_properties_col_name\n",
        "        )\n",
        "    )\n",
        "    return cosmos_vertices_df, cosmos_edges_df\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "v, e = prepare_vertices_edge_df(\n",
        "    df=raw_data,\n",
        "    source_col_name=\"nameOrig\",\n",
        "    dest_col_name=\"nameDest\",\n",
        "    parition_key_col_name=\"nameOrig\",\n",
        "    cosmos_parition_name=\"accountId\",\n",
        "    edge_properties_col_name=[\n",
        "        \"amount\",\n",
        "        \"oldbalanceOrg\",\n",
        "        \"oldbalanceDest\",\n",
        "        \"newbalanceDest\",\n",
        "    ],\n",
        "    vertex_properties_col_name=[],\n",
        "    sample=False\n",
        ")\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "write_to_cosmos_graph(v,\"vertices\",False)\n",
        "write_to_cosmos_graph(e,\"edges\",False)"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    }
  ],
  "metadata": {
    "kernelspec": {
      "name": "synapse_pyspark",
      "language": "Python",
      "display_name": "Synapse PySpark"
    },
    "language_info": {
      "name": "python"
    },
    "kernel_info": {
      "name": "synapse_pyspark"
    },
    "description": null,
    "save_output": true,
    "synapse_widget": {
      "version": "0.1",
      "state": {}
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}