1.data_load.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import pandas as pd\n",
    "import datetime\n",
    "import os\n",
    "from neo4j import GraphDatabase\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "uri = os.getenv(\"NEO4J_URI\")\n",
    "user = os.getenv(\"NEO4J_USERNAME\")\n",
    "password = os.getenv(\"NEO4J_PASSWORD\")\n",
    "\n",
    "driver = GraphDatabase.driver(uri, auth=(user, password))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # To Delete all nodes\n",
    "# with driver.session() as session:\n",
    "#     session.run(\"MATCH (n) DETACH DELETE n\")\n",
    "\n",
    "# # Delete all relationships\n",
    "# with driver.session() as session:\n",
    "#     session.run(\"MATCH ()-[r]->() DELETE r\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "int_datasets = ['weather','foot_traffic','web_traffic','social_media','sec_master']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_dataset_info = pd.read_excel('sample_data.xlsx',sheet_name='dataset_info')\n",
    "df_column_info = pd.read_excel('sample_data.xlsx',sheet_name='column_info')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating the nodes and relationship in neo4j (I am using the free Aura instance). \n",
    "\n",
    "We are adding an additional tag called 'embeddable' to the dataset nodes and the reference column nodes so that we can add embeddings to them in the next stage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "weather\n",
      "foot_traffic\n",
      "web_traffic\n",
      "social_media\n",
      "sec_master\n"
     ]
    }
   ],
   "source": [
    "# Create a session\n",
    "with driver.session() as session:\n",
    "    for dataset in int_datasets:\n",
    "        print(dataset)\n",
    "        df = pd.read_excel('sample_data.xlsx',sheet_name=dataset)\n",
    "\n",
    "        # Create a node for the dataset\n",
    "        dataset_description = df_dataset_info[df_dataset_info['dataset']==dataset]['description'].iloc[0]\n",
    "        session.run(f\"\"\"CREATE ({dataset}:dataset:embeddable {{name:'{dataset}'\n",
    "                    ,type:'dataset'\n",
    "                    ,source:'snowflake'\n",
    "                    ,table:'graph_db.public.{dataset}'\n",
    "                    ,description:'{dataset_description}'\n",
    "                    }})\"\"\")\n",
    "\n",
    "        # Create nodes for each column in the DataFrame\n",
    "        for column in df.columns:\n",
    "            row = df_column_info[(df_column_info['dataset']==dataset)&(df_column_info['column']==column)].iloc[0]\n",
    "            col_type = row['type']\n",
    "            col_description = row['description']\n",
    "            unique_values = ','.join(list(df[column].astype(str).unique()))\n",
    "            if col_type == 'reference':\n",
    "                session.run(f\"\"\"CREATE ({column}:column:embeddable {{name:'{column}'\n",
    "                            ,type:'column'\n",
    "                            ,source:'{dataset}'\n",
    "                            ,col_type:'{col_type}'\n",
    "                            ,description:'{col_description}'\n",
    "                            ,values:'{unique_values}'\n",
    "                            }})\"\"\")\n",
    "            else:\n",
    "                session.run(f\"\"\"CREATE ({column}:column {{name:'{column}'\n",
    "                            ,type:'column'\n",
    "                            ,source:'{dataset}'\n",
    "                            ,col_type:'{col_type}'\n",
    "                            ,description:'{col_description}'\n",
    "                            ,values:'{unique_values}'\n",
    "                            }})\"\"\")\n",
    "\n",
    "\n",
    "        # Create relationships between the dataset and each column\n",
    "        for column in df.columns:\n",
    "            session.run(f\"MATCH ({dataset}:dataset {{name: '{dataset}'}}), ({column}:column {{name: '{column}', source: '{dataset}'}}) CREATE ({dataset})-[:HAS_COLUMN]->({column})\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>count(n)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>46</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   count(n)\n",
       "0        46"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Sample Query\n",
    "query = \"\"\"\n",
    "MATCH (n) \n",
    "RETURN count(n)\n",
    "\"\"\"\n",
    "with driver.session() as session:\n",
    "  result = session.run(query)\n",
    "  df_query = result.to_df()\n",
    "\n",
    "df_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # this is in case if you want to manually create the relationships instead of using the LLM to infer it in the next step. Mostly for testing purpose.\n",
    "# with driver.session() as session:\n",
    "#     session.run(\"\"\"MATCH (a:column {name:'post_code',source:'foot_traffic'}), (b:column {name:'zip_code',source:'weather'})\n",
    "#                  CREATE (a)-[:RELATED_TO {confidence:'90',type:'same'}]->(b)\"\"\")\n",
    "#     session.run(\"\"\"MATCH (a:column {name:'symbol',source:'foot_traffic'}), (b:column {name:'ticker',source:'sec_master'})\n",
    "#                  CREATE (a)-[:RELATED_TO {confidence:'90',type:'same'}]->(b)\"\"\")\n",
    "#     session.run(\"\"\"MATCH (a:column {name:'website_owner',source:'web_traffic'}), (b:column {name:'entity_name',source:'sec_master'})\n",
    "#                  CREATE (a)-[:RELATED_TO {confidence:'90',type:'same'}]->(b)\"\"\")\n",
    "#     session.run(\"\"\"MATCH (a:column {name:'website_brand',source:'web_traffic'}), (b:column {name:'page_owner',source:'social_media'})\n",
    "#                  CREATE (a)-[:RELATED_TO {confidence:'90',type:'same'}]->(b)\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading Data into Snowflake\n",
    "Optional if you want to load build chains and agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to import ArrowResult. No Apache Arrow result set format can be used. ImportError: DLL load failed while importing arrow_iterator: The specified procedure could not be found.\n"
     ]
    }
   ],
   "source": [
    "import snowflake.connector\n",
    "from snowflake.snowpark import Session\n",
    "\n",
    "USER = os.getenv(\"SNOW_USER\")\n",
    "PASSWORD = os.getenv(\"SNOW_PASSWORD\")\n",
    "ACCOUNT = os.getenv(\"SNOW_ACCOUNT\")\n",
    "WAREHOUSE = os.getenv(\"SNOW_WAREHOUSE\")\n",
    "DATABASE = os.getenv(\"SNOW_DATABASE\")\n",
    "ROLE = os.getenv(\"SNOW_ROLE\")\n",
    "SCHEMA = os.getenv(\"SNOW_SCHEMA\")\n",
    "\n",
    "con = snowflake.connector.connect(\n",
    "    user = USER,\n",
    "    password = PASSWORD,\n",
    "    account = ACCOUNT,\n",
    "    warehouse = WAREHOUSE,\n",
    "    database = DATABASE,\n",
    "    role = ROLE,\n",
    "    schema = SCHEMA\n",
    ")\n",
    "connection_parameters = {\n",
    "\"account\": ACCOUNT,\n",
    "\"user\": USER,\n",
    "\"password\": PASSWORD,\n",
    "\"role\": ROLE,\n",
    "\"warehouse\": WAREHOUSE,\n",
    "\"database\": DATABASE,\n",
    "\"schema\":SCHEMA\n",
    "}  \n",
    "snow_session = Session.builder.configs(connection_parameters).create()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "sql = \"\"\"\n",
    "create database if not exists graph_db\n",
    "\"\"\"\n",
    "df = pd.read_sql(sql,con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "weather\n",
      "foot_traffic\n",
      "web_traffic\n",
      "social_media\n",
      "sec_master\n"
     ]
    }
   ],
   "source": [
    "db = 'graph_db'\n",
    "for d in int_datasets:\n",
    "    print(d)\n",
    "    df = pd.read_excel('sample_data.xlsx',sheet_name=d)\n",
    "    df.columns = [c.upper() for c in df.columns]\n",
    "    df_snowpark = snow_session.create_dataframe(df)\n",
    "    df_snowpark.write.mode(\"overwrite\").save_as_table(f\"{db}.public.{d}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda38",
   "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}