Stock_Injest.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "incorporate-description",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas\n",
    "import math\n",
    "import re\n",
    "\n",
    "from multiprocessing import Pool\n",
    "from functools import partial\n",
    "from itertools import chain\n",
    "from neo4j import GraphDatabase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "religious-horizontal",
   "metadata": {},
   "outputs": [],
   "source": [
    "SUBREDDIT = 'investing'\n",
    "GRAPH_PASSWORD = 'test1234'\n",
    "\n",
    "SUBMISSIONS_LOC = 'data/investing/investing_2020_1_1_2021_2_28.csv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "given-multimedia",
   "metadata": {},
   "outputs": [],
   "source": [
    "THREADS = 48\n",
    "\n",
    "POSTS_LOC = f'data/{SUBREDDIT}/posts.csv'\n",
    "USERS_LOC = f'data/{SUBREDDIT}/users.csv'\n",
    "SYMBOLS_LOC = f'data/{SUBREDDIT}/symbols.csv'\n",
    "AUTHORS_TO_POSTS_LOC = f'data/{SUBREDDIT}/authors_to_posts.csv'\n",
    "POSTS_TO_SYMBOLS_LOC = f'data/{SUBREDDIT}/posts_to_symbols.csv'\n",
    "AUTHORS_TO_AUTHORS_LOC = f'data/{SUBREDDIT}/author_to_authors.csv'\n",
    "AUTHORS_TO_SYMBOLS_LOC = f'data/{SUBREDDIT}/author_to_symbols.csv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "swiss-prefix",
   "metadata": {},
   "outputs": [],
   "source": [
    "driver = GraphDatabase.driver('bolt://localhost:7687', auth=('neo4j', 'test1234'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "polished-blind",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9738"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def load_symbols():\n",
    "    exchanges = ['amex', 'nasdaq', 'nyse']\n",
    "\n",
    "    # create a dict of symbols to exchanges\n",
    "    symbols = dict()\n",
    "    for exchange in exchanges:\n",
    "        exchange_symbols = pandas.read_csv(f'data/{exchange}_symbols.tsv', sep='\\t')\n",
    "        for index, row in exchange_symbols.iterrows():\n",
    "            symbols[row['Symbol']] = exchange\n",
    "            \n",
    "    return symbols\n",
    "\n",
    "symbols = load_symbols()\n",
    "len(symbols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "friendly-charlotte",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of submissions: 8411\n"
     ]
    },
    {
     "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>type</th>\n",
       "      <th>depth</th>\n",
       "      <th>parent_id</th>\n",
       "      <th>submission_id</th>\n",
       "      <th>content</th>\n",
       "      <th>author_id</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>up_votes</th>\n",
       "      <th>ratio</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>submission_title</td>\n",
       "      <td>0</td>\n",
       "      <td>&lt;NA&gt;</td>\n",
       "      <td>ltky2b</td>\n",
       "      <td>Daily Advice Thread - All basic help or advice...</td>\n",
       "      <td>AutoModerator</td>\n",
       "      <td>1614420015</td>\n",
       "      <td>22</td>\n",
       "      <td>0.97</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>submission_body</td>\n",
       "      <td>0</td>\n",
       "      <td>&lt;NA&gt;</td>\n",
       "      <td>ltky2b</td>\n",
       "      <td>If your question is \"I have $10,000, what do I...</td>\n",
       "      <td>AutoModerator</td>\n",
       "      <td>1614420015</td>\n",
       "      <td>22</td>\n",
       "      <td>0.97</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>comment</td>\n",
       "      <td>1</td>\n",
       "      <td>ltky2b</td>\n",
       "      <td>goytzuy</td>\n",
       "      <td>**Hi, welcome to /r/investing. Please note tha...</td>\n",
       "      <td>AutoModerator</td>\n",
       "      <td>1614420016</td>\n",
       "      <td>1</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>comment</td>\n",
       "      <td>1</td>\n",
       "      <td>ltky2b</td>\n",
       "      <td>gp4h7pb</td>\n",
       "      <td>whats a reasonable portion of a portfolio to a...</td>\n",
       "      <td>S7EFEN</td>\n",
       "      <td>1614488027</td>\n",
       "      <td>3</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>comment</td>\n",
       "      <td>2</td>\n",
       "      <td>gp4h7pb</td>\n",
       "      <td>gp4ouvn</td>\n",
       "      <td>I'm a millennial, half of my IRA is in safe ET...</td>\n",
       "      <td>kfuzion</td>\n",
       "      <td>1614491062</td>\n",
       "      <td>1</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>794333</th>\n",
       "      <td>comment</td>\n",
       "      <td>2</td>\n",
       "      <td>fcniw6m</td>\n",
       "      <td>fco14w8</td>\n",
       "      <td>I'd like to have a steady stream. Monthly if p...</td>\n",
       "      <td>killthebatman15</td>\n",
       "      <td>1577822203</td>\n",
       "      <td>1</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>794334</th>\n",
       "      <td>comment</td>\n",
       "      <td>1</td>\n",
       "      <td>ei44tp</td>\n",
       "      <td>fcnnqgw</td>\n",
       "      <td>I absolutely wouldn't recommend investing on F...</td>\n",
       "      <td>dvdmovie1</td>\n",
       "      <td>1577814970</td>\n",
       "      <td>1</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>794335</th>\n",
       "      <td>comment</td>\n",
       "      <td>2</td>\n",
       "      <td>fcnnqgw</td>\n",
       "      <td>fd0emr6</td>\n",
       "      <td>They advertise to us on Facebook... that is th...</td>\n",
       "      <td>peepee21</td>\n",
       "      <td>1578117725</td>\n",
       "      <td>1</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>794336</th>\n",
       "      <td>comment</td>\n",
       "      <td>1</td>\n",
       "      <td>ei44tp</td>\n",
       "      <td>fco9y2c</td>\n",
       "      <td>What I don't see mentioned here is that if you...</td>\n",
       "      <td>mercury187</td>\n",
       "      <td>1577827119</td>\n",
       "      <td>0</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>794337</th>\n",
       "      <td>comment</td>\n",
       "      <td>2</td>\n",
       "      <td>fco9y2c</td>\n",
       "      <td>fcrplh4</td>\n",
       "      <td>You only pay taxes when you take a profit. I n...</td>\n",
       "      <td>boyinahouse</td>\n",
       "      <td>1577904917</td>\n",
       "      <td>0</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>794338 rows × 9 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                    type  depth parent_id submission_id  \\\n",
       "0       submission_title      0      <NA>        ltky2b   \n",
       "1        submission_body      0      <NA>        ltky2b   \n",
       "2                comment      1    ltky2b       goytzuy   \n",
       "3                comment      1    ltky2b       gp4h7pb   \n",
       "4                comment      2   gp4h7pb       gp4ouvn   \n",
       "...                  ...    ...       ...           ...   \n",
       "794333           comment      2   fcniw6m       fco14w8   \n",
       "794334           comment      1    ei44tp       fcnnqgw   \n",
       "794335           comment      2   fcnnqgw       fd0emr6   \n",
       "794336           comment      1    ei44tp       fco9y2c   \n",
       "794337           comment      2   fco9y2c       fcrplh4   \n",
       "\n",
       "                                                  content        author_id  \\\n",
       "0       Daily Advice Thread - All basic help or advice...    AutoModerator   \n",
       "1       If your question is \"I have $10,000, what do I...    AutoModerator   \n",
       "2       **Hi, welcome to /r/investing. Please note tha...    AutoModerator   \n",
       "3       whats a reasonable portion of a portfolio to a...           S7EFEN   \n",
       "4       I'm a millennial, half of my IRA is in safe ET...          kfuzion   \n",
       "...                                                   ...              ...   \n",
       "794333  I'd like to have a steady stream. Monthly if p...  killthebatman15   \n",
       "794334  I absolutely wouldn't recommend investing on F...        dvdmovie1   \n",
       "794335  They advertise to us on Facebook... that is th...         peepee21   \n",
       "794336  What I don't see mentioned here is that if you...       mercury187   \n",
       "794337  You only pay taxes when you take a profit. I n...      boyinahouse   \n",
       "\n",
       "         timestamp  up_votes  ratio  \n",
       "0       1614420015        22   0.97  \n",
       "1       1614420015        22   0.97  \n",
       "2       1614420016         1   0.00  \n",
       "3       1614488027         3   0.00  \n",
       "4       1614491062         1   0.00  \n",
       "...            ...       ...    ...  \n",
       "794333  1577822203         1   0.00  \n",
       "794334  1577814970         1   0.00  \n",
       "794335  1578117725         1   0.00  \n",
       "794336  1577827119         0   0.00  \n",
       "794337  1577904917         0   0.00  \n",
       "\n",
       "[794338 rows x 9 columns]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "submissions = pandas.read_csv(\n",
    "    SUBMISSIONS_LOC,\n",
    "    names=['type', 'depth', 'parent_id', 'submission_id', 'content', 'author_id', 'timestamp', 'up_votes', 'ratio'],\n",
    "    dtype={\n",
    "        'type': 'string',\n",
    "        'depth': int,\n",
    "        'parent_id': 'string',\n",
    "        'submission_id': 'string',\n",
    "        'content': 'string',\n",
    "        'author_id': 'string', \n",
    "        'timestamp': int, \n",
    "        'up_votes': int,\n",
    "        'ratio': float\n",
    "    }\n",
    ")\n",
    "\n",
    "print(f'Number of submissions: {len(submissions[submissions.type == \"submission_title\"].index)}')\n",
    "submissions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "governing-implement",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "type\n",
       "comment             104058\n",
       "submission_body       4705\n",
       "submission_title      4705\n",
       "Name: author_id, dtype: int64"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "submissions.groupby(['type'])['author_id'].nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "dynamic-distributor",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "48"
      ]
     },
     "execution_count": 128,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def chunk_submissions(submissions, num_chunks):\n",
    "    # figure out the aprox size of each chunk, this is aprox since we want to chunk by full submissions\n",
    "    approx_size = math.floor(len(submissions.index) / num_chunks)\n",
    "\n",
    "    splits = [0]\n",
    "    for i in range(1, num_chunks):\n",
    "        start_row = i * approx_size\n",
    "        \n",
    "        # search for the start of the first submission up until the last chunk, if we reach the last chunk, skip this\n",
    "        for j in range(start_row, splits[-1], -1):\n",
    "            if submissions.iloc[j]['type'] == 'submission_title':\n",
    "                splits.append(j)\n",
    "                break\n",
    "                \n",
    "    splits.append(len(submissions.index))\n",
    "\n",
    "    return list(submissions[splits[i - 1]:splits[i]] for i in range(1, len(splits)))\n",
    "\n",
    "submission_chunks = chunk_submissions(submissions, THREADS)\n",
    "len(submission_chunks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "demonstrated-shift",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(['GME', 'AMC'], ['GME', 'COM', 'MOON', 'WWW', 'AMC'])"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# extract symbols from text \n",
    "def extract_symbols(content, symbols, strict=True):\n",
    "    if pandas.isna(content):\n",
    "        return []\n",
    "    \n",
    "    if not strict:\n",
    "        tokens = re.split(r'\\W+', content.upper())\n",
    "    else:\n",
    "        tokens = content.split(' ')\n",
    "    \n",
    "    found = set()\n",
    "    for token in tokens:\n",
    "        if len(token) > 1:\n",
    "            if token[0] == '$' or token[0] == ':':\n",
    "                token = token[1:].upper()\n",
    "            \n",
    "            if token in symbols:\n",
    "                found.add(token)\n",
    "        \n",
    "    return list(found)\n",
    "\n",
    "(\n",
    "    extract_symbols('$amc and GME to the www.moon.com.', symbols), \n",
    "    extract_symbols('$amc and GME to the www.moon.com.', symbols, False)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "entertaining-finish",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "square-colombia",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_chunk(symbols, chunk):\n",
    "    # id of the root post\n",
    "    post_row = None\n",
    "    \n",
    "    # seen symbols in ancestors\n",
    "    ancestor_symbols = []\n",
    "    # seen authors in ancestors\n",
    "    ancestor_authors = []\n",
    "    \n",
    "    # accumulate top level post to symbols in all its descendents\n",
    "    post_to_symbols = []\n",
    "    # accumulate authors to interaction with posts\n",
    "    author_to_posts = []\n",
    "    # accumulate author to author interactions per branch\n",
    "    author_to_authors = []\n",
    "    # accumulate author to symbol interactions per branch\n",
    "    author_to_symbols = []\n",
    "    \n",
    "    # top level details about each post\n",
    "    posts = []\n",
    "    # dataframes of the aggregates of the post level accumulators for the entire chunk\n",
    "    post_to_symbols_dfs = []\n",
    "    author_to_posts_dfs = []\n",
    "    author_to_authors_dfs = []\n",
    "    author_to_symbols_dfs = []\n",
    "    \n",
    "    def aggregate_post():\n",
    "        aggregates = {\n",
    "            'max_depth_score': pandas.NamedAgg(column='depth_score', aggfunc='max'),\n",
    "            'mean_depth_score': pandas.NamedAgg(column='depth_score', aggfunc='mean'),\n",
    "            'depth_score': pandas.NamedAgg(column='depth_score', aggfunc='sum'),\n",
    "            'up_votes': pandas.NamedAgg(column='up_votes', aggfunc='sum'),\n",
    "            'up_votes_with_depth': pandas.NamedAgg(column='up_votes_with_depth', aggfunc='sum'),\n",
    "            'count': pandas.NamedAgg(column='depth_score', aggfunc='count')\n",
    "        }\n",
    "        \n",
    "        post_to_symbols_df = None\n",
    "        try:\n",
    "            post_to_symbols_df = pandas.DataFrame(post_to_symbols, columns=['post_id', 'symbol', 'depth_score', 'up_votes', 'up_votes_with_depth']) \\\n",
    "                .groupby([\n",
    "                    pandas.Grouper('post_id'),\n",
    "                    pandas.Grouper('symbol')\n",
    "                ]).agg(**aggregates)\n",
    "            post_to_symbols_dfs.append(post_to_symbols_df)\n",
    "        except:\n",
    "            display(f'Error with post_to_symbols {post_row.submission_id}.', post_to_symbols)\n",
    "            \n",
    "        author_to_posts_df = pandas.DataFrame(author_to_posts, columns=['author_id', 'post_id', 'depth_score', 'up_votes', 'up_votes_with_depth']) \\\n",
    "            .groupby([\n",
    "                pandas.Grouper('author_id'),\n",
    "                pandas.Grouper('post_id')\n",
    "            ]).agg(**aggregates)\n",
    "        author_to_posts_dfs.append(author_to_posts_df)\n",
    "            \n",
    "        author_to_authors_df = pandas.DataFrame(author_to_authors, columns=['from_id', 'to_id', 'depth_score', 'up_votes', 'up_votes_with_depth']) \\\n",
    "            .groupby([\n",
    "                pandas.Grouper('from_id'),\n",
    "                pandas.Grouper('to_id')\n",
    "            ]).agg(**aggregates)\n",
    "        author_to_authors_dfs.append(author_to_authors_df)\n",
    "            \n",
    "        try:\n",
    "            author_to_symbols_df = pandas.DataFrame(author_to_symbols, columns=['author_id', 'symbol', 'depth_score', 'up_votes', 'up_votes_with_depth']) \\\n",
    "                .groupby([\n",
    "                    pandas.Grouper('author_id'),\n",
    "                    pandas.Grouper('symbol')\n",
    "                ]).agg(**aggregates)\n",
    "            author_to_symbols_dfs.append(author_to_symbols_df)\n",
    "        except:\n",
    "            display(f'Error with author_to_symbols {post_row.submission_id}.', author_to_symbols)\n",
    "            \n",
    "        posts.append([\n",
    "            # copy some stuff from the main post row\n",
    "            post_row.submission_id, post_row.content, post_row.author_id, post_row.timestamp, post_row.up_votes, post_row.ratio,\n",
    "            # some aggregate stats\n",
    "            len(author_to_posts),\n",
    "            len(author_to_posts_df.index),\n",
    "            len(post_to_symbols_df.index) if post_to_symbols_df is not None else 0,\n",
    "            author_to_posts_df['up_votes'].sum()\n",
    "        ])\n",
    "    \n",
    "    rows_processed = 0\n",
    "    for (i, row) in chunk.iterrows():\n",
    "        row_symbols = extract_symbols(row.content, symbols)\n",
    "        \n",
    "        if row.type == 'submission_title':\n",
    "            # if not the first, process current accumulators\n",
    "            if post_row is not None:\n",
    "                aggregate_post()\n",
    "            \n",
    "            # new submission, reset the tree accumulators\n",
    "            post_row = row\n",
    "            \n",
    "            ancestor_symbols = [row_symbols]\n",
    "            ancestor_authors = [row.author_id]\n",
    "            \n",
    "            post_to_symbols = []\n",
    "            author_to_posts = []\n",
    "            author_to_authors = []\n",
    "            author_to_symbols = []\n",
    "        elif row.type == 'submission_body':\n",
    "            # just an extension of the root, just extend with any symbols in the body\n",
    "            row_symbols = list(set(row_symbols + ancestor_symbols[0]))\n",
    "            ancestor_symbols[0] = row_symbols\n",
    "        elif row.depth <= previous_depth:\n",
    "            # sibling or cousin, remove some context from the accumulators\n",
    "            pops = previous_depth - row.depth + 1\n",
    "            ancestor_symbols = ancestor_symbols[:-pops]\n",
    "            ancestor_authors = ancestor_authors[:-pops]\n",
    "\n",
    "        # adjust upvotes to 0, stored as -1 if not retrieved in original dataset\n",
    "        up_votes = max(row.up_votes, 0)\n",
    "\n",
    "        # if not the first row of a post, start building the relationships\n",
    "        if row.type != 'submission_title':\n",
    "            # interactions are inversely propertional to the depth\n",
    "            depth_score = round(1 / (row.depth + 1), 2) \n",
    "\n",
    "            # relationships to symbols in the current row\n",
    "            for symbol in row_symbols:\n",
    "                # we want to weight symbols in the main post and those closer to it heigher so use the depth score\n",
    "                post_to_symbols.append([post_row.submission_id, symbol, depth_score, up_votes, round(depth_score * up_votes, 2)])\n",
    "                # author to symbols they posted about so depth score of 1\n",
    "                author_to_symbols.append([row.author_id, symbol, 1, up_votes, up_votes])\n",
    "\n",
    "            # author to current post\n",
    "            author_to_posts.append([row.author_id, post_row.submission_id, depth_score, up_votes, round(depth_score * up_votes, 2)])\n",
    "\n",
    "        # for comments we also build relationships to authors and symbols in ancestors\n",
    "        if row.type == 'comment':\n",
    "            # for each previously seen author, create a relationship with a weight thats the inverse of their distance\n",
    "            for cur_depth, cur_author_id in enumerate(reversed(ancestor_authors)):\n",
    "                # ignore yourself in the ancestors\n",
    "                if cur_author_id != row.author_id:\n",
    "                    cur_depth_score = round(1 / (cur_depth + 1), 2) \n",
    "                    author_to_authors.append([row.author_id, cur_author_id, cur_depth_score, up_votes, round(cur_depth_score * up_votes, 2)])\n",
    "\n",
    "            # similarly for previously seen symbols\n",
    "            for cur_depth, cur_symbols in enumerate(reversed(ancestor_symbols)):\n",
    "                cur_depth_score = round(1 / (cur_depth + 1), 2) \n",
    "                for symbol in cur_symbols:\n",
    "                    author_to_symbols.append([row.author_id, symbol, cur_depth_score, up_votes, round(cur_depth_score * up_votes, 2)])\n",
    "\n",
    "            # update ancestor accumulators\n",
    "            ancestor_symbols.append(row_symbols)\n",
    "            ancestor_authors.append(row.author_id)\n",
    "\n",
    "        # track depth so we know when we popped up from a leaf\n",
    "        previous_depth = row.depth\n",
    "            \n",
    "        rows_processed += 1\n",
    "        if rows_processed % 25000 == 0:\n",
    "            print(f'Processed {rows_processed} rows.')\n",
    "            \n",
    "    # aggregate the last post\n",
    "    aggregate_post()\n",
    "    \n",
    "    return [\n",
    "        pandas.DataFrame(posts, columns=['post_id', 'title', 'author_id', 'timestamp', 'up_votes', 'ratio', 'replies', 'authors', 'symbols', 'sum_up_votes']),\n",
    "        pandas.concat(post_to_symbols_dfs),\n",
    "        pandas.concat(author_to_posts_dfs),\n",
    "        pandas.concat(author_to_authors_dfs),\n",
    "        pandas.concat(author_to_symbols_dfs)\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "suitable-parameter",
   "metadata": {},
   "outputs": [],
   "source": [
    "with Pool(THREADS) as pool:\n",
    "    processed_posts = pool.map(partial(process_chunk, symbols), submission_chunks)\n",
    "    \n",
    "# concat the dataframes from each chunk together and write to file\n",
    "pandas.concat([s[0] for s in processed_posts]).to_csv(POSTS_LOC, index=False, header=True)\n",
    "pandas.concat([s[1] for s in processed_posts]).reset_index().to_csv(POSTS_TO_SYMBOLS_LOC, index=False, header=True)\n",
    "pandas.concat([s[2] for s in processed_posts]).reset_index().to_csv(AUTHORS_TO_POSTS_LOC, index=False, header=True)\n",
    "\n",
    "# these are not indexed by posts so actually need to be reaggregated\n",
    "aggregates = {\n",
    "    'max_depth_score': pandas.NamedAgg(column='depth_score', aggfunc='max'),\n",
    "    'depth_score': pandas.NamedAgg(column='depth_score', aggfunc='sum'),\n",
    "    'up_votes': pandas.NamedAgg(column='up_votes', aggfunc='sum'),\n",
    "    'up_votes_with_depth': pandas.NamedAgg(column='up_votes_with_depth', aggfunc='sum'),\n",
    "    'count': pandas.NamedAgg(column='depth_score', aggfunc='count')\n",
    "}\n",
    "\n",
    "authors_to_authors = pandas.concat([s[3] for s in processed_posts]).reset_index()\\\n",
    "    .groupby([\n",
    "        pandas.Grouper('from_id'),\n",
    "        pandas.Grouper('to_id')\n",
    "    ]).agg(**aggregates)\n",
    "authors_to_authors['mean_depth_score'] = authors_to_authors['depth_score'] / authors_to_authors['count']\n",
    "authors_to_authors.reset_index().to_csv(AUTHORS_TO_AUTHORS_LOC, index=False, header=True)\n",
    "\n",
    "authors_to_symbols = pandas.concat([s[4] for s in processed_posts]).reset_index()\\\n",
    "    .groupby([\n",
    "        pandas.Grouper('author_id'),\n",
    "        pandas.Grouper('symbol')\n",
    "    ]).agg(**aggregates)\n",
    "authors_to_symbols['mean_depth_score'] = authors_to_symbols['depth_score'] / authors_to_symbols['count']\n",
    "authors_to_symbols.reset_index().to_csv(AUTHORS_TO_SYMBOLS_LOC, index=False, header=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "revised-czech",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "owned-concentrate",
   "metadata": {},
   "outputs": [],
   "source": [
    "authors_to_posts = pandas.read_csv(AUTHORS_TO_POSTS_LOC)\n",
    "authors_to_authors = pandas.read_csv(AUTHORS_TO_AUTHORS_LOC)\n",
    "authors_to_symbols = pandas.read_csv(AUTHORS_TO_SYMBOLS_LOC)\n",
    "posts_to_symbols = pandas.read_csv(POSTS_TO_SYMBOLS_LOC)\n",
    "\n",
    "authors_to_posts_agg = authors_to_posts.groupby([\n",
    "    pandas.Grouper('author_id')\n",
    "]).agg(\n",
    "    posts = pandas.NamedAgg(column='post_id', aggfunc='count'),\n",
    "    posts_depth_score = pandas.NamedAgg(column='depth_score', aggfunc='sum'),\n",
    "    posts_up_votes = pandas.NamedAgg(column='up_votes', aggfunc='sum'),\n",
    "    posts_up_votes_with_depth = pandas.NamedAgg(column='up_votes_with_depth', aggfunc='sum')\n",
    ").sort_values(by=['posts'])\n",
    "\n",
    "authors_to_authors_agg = authors_to_authors.groupby([\n",
    "    pandas.Grouper('from_id')\n",
    "]).agg(\n",
    "    replied_to = pandas.NamedAgg(column='to_id', aggfunc='count'),\n",
    "    replied_depth_score = pandas.NamedAgg(column='depth_score', aggfunc='sum'),\n",
    "    replied_up_votes = pandas.NamedAgg(column='up_votes', aggfunc='sum'),\n",
    "    replied_up_votes_with_depth = pandas.NamedAgg(column='up_votes_with_depth', aggfunc='sum')\n",
    ").rename(columns={'from_id': 'author_id'})\n",
    "\n",
    "authors_to_symbols_agg = authors_to_symbols.groupby([\n",
    "    pandas.Grouper('author_id')\n",
    "]).agg(\n",
    "    symbols = pandas.NamedAgg(column='symbol', aggfunc='count'),\n",
    "    symbols_depth_score = pandas.NamedAgg(column='depth_score', aggfunc='sum'),\n",
    "    symbols_up_votes = pandas.NamedAgg(column='up_votes', aggfunc='sum'),\n",
    "    symbols_up_votes_with_depth = pandas.NamedAgg(column='up_votes_with_depth', aggfunc='sum')\n",
    ")\n",
    "\n",
    "# join all the user metrics into a single csv\n",
    "authors_to_posts_agg.join(authors_to_authors_agg).join(authors_to_symbols_agg).reset_index().to_csv(USERS_LOC, index=False, header=True)\n",
    "\n",
    "# aggregate some symbols metrics as well\n",
    "symbols_to_authors_agg = authors_to_symbols.groupby([\n",
    "    pandas.Grouper('symbol')\n",
    "]).agg(\n",
    "    authors = pandas.NamedAgg(column='author_id', aggfunc='count'),\n",
    "    authors_depth_score = pandas.NamedAgg(column='depth_score', aggfunc='sum'),\n",
    "    authors_up_votes = pandas.NamedAgg(column='up_votes', aggfunc='sum'),\n",
    "    authors_up_votes_with_depth = pandas.NamedAgg(column='up_votes_with_depth', aggfunc='sum')\n",
    ").sort_values(by=['authors'])\n",
    "\n",
    "symbols_to_posts_agg = posts_to_symbols.groupby([\n",
    "    pandas.Grouper('symbol')\n",
    "]).agg(\n",
    "    posts = pandas.NamedAgg(column='post_id', aggfunc='count'),\n",
    "    posts_depth_score = pandas.NamedAgg(column='depth_score', aggfunc='sum'),\n",
    "    posts_up_votes = pandas.NamedAgg(column='up_votes', aggfunc='sum'),\n",
    "    posts_up_votes_with_depth = pandas.NamedAgg(column='up_votes_with_depth', aggfunc='sum')\n",
    ")\n",
    "\n",
    "# join some symbol metrics into a since csv\n",
    "symbols_to_authors_agg.join(symbols_to_posts_agg).reset_index().to_csv(SYMBOLS_LOC, index=False, header=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wireless-blanket",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "medical-triumph",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_users(file, tx):\n",
    "    tx.run(f\"LOAD CSV WITH HEADERS FROM 'file:///{file}' \" +\n",
    "        \"\"\"\n",
    "            AS line\n",
    "            CREATE (:User {\n",
    "                id: line.author_id,\n",
    "                posts: toInteger(line.posts),\n",
    "                posts_up_votes: toInteger(line.posts_up_votes),\n",
    "                posts_up_votes_with_depth: toFloat(line.posts_up_votes_with_depth),\n",
    "                replied_to: toInteger(line.replied_to),\n",
    "                symbols: toInteger(line.symbols)\n",
    "            })\n",
    "        \"\"\",\n",
    "        file=file\n",
    "    )\n",
    "    \n",
    "def load_posts(file, tx):\n",
    "    tx.run(f\"LOAD CSV WITH HEADERS FROM 'file:///{file}' \" +\n",
    "        \"\"\"\n",
    "            AS line\n",
    "            CREATE (:Post {\n",
    "                id: line.post_id,\n",
    "                title: line.title,\n",
    "                author_id: line.author_id,\n",
    "                timestamp: datetime({epochSeconds:toInteger(line.timestamp)}),\n",
    "                up_votes: toInteger(line.up_votes),\n",
    "                ratio: toInteger(line.ratio),\n",
    "                replies: toInteger(line.replies),\n",
    "                authors: toInteger(line.authors),\n",
    "                symbols: toInteger(line.symbols), \n",
    "                sum_up_votes: toInteger(line.sum_up_votes)\n",
    "            })\n",
    "        \"\"\"\n",
    "    )\n",
    "    \n",
    "def load_symbols(file, tx):\n",
    "    tx.run(f\"LOAD CSV WITH HEADERS FROM 'file:///{file}' \" +\n",
    "        \"\"\" \n",
    "            AS line\n",
    "            CREATE (:Symbol {\n",
    "                id: line.symbol,\n",
    "                authors: toInteger(line.authors),\n",
    "                posts: toInteger(line.posts),\n",
    "                posts_up_votes_with_depth: toFloat(line.posts_up_votes_with_depth)\n",
    "            })\n",
    "        \"\"\"\n",
    "    )\n",
    "    \n",
    "def create_indices(tx):\n",
    "    # create indices for things were going to query by\n",
    "    for node_label in ['User', 'Submission', 'Symbol']:\n",
    "        for node_property in ['id', 'communityUpVotes']:\n",
    "            tx.run(f\"CREATE INDEX {node_label}_{node_property} IF NOT EXISTS FOR (n:{node_label}) ON (n.{node_property})\")\n",
    "    \n",
    "def load_posts_to_symbols(file, tx):\n",
    "    tx.run('CALL apoc.periodic.iterate(' +\n",
    "        f\"\"\"'LOAD CSV WITH HEADERS FROM \"file:///{file}\" AS line RETURN line', \"\"\" +  \n",
    "        \"\"\"\n",
    "            'MATCH (post:Post {id: line.post_id}), (symbol:Symbol {id: line.symbol})\n",
    "            CREATE (post)-[:Mentioned {\n",
    "                count: toInteger(line.count),\n",
    "                depth_score: toInteger(line.depth_score),\n",
    "                up_votes: toInteger(line.up_votes),\n",
    "                up_votes_with_depth: toInteger(line.up_votes_with_depth)\n",
    "            }]->(symbol)'\n",
    "        , {batchSize:100000, iterateList:true, parallel:false})\"\"\"\n",
    "    )\n",
    "    \n",
    "def load_authors_to_authors(file, tx):\n",
    "    tx.run('CALL apoc.periodic.iterate(' +\n",
    "        f\"\"\"'LOAD CSV WITH HEADERS FROM \"file:///{file}\" AS line RETURN line', \"\"\" +  \n",
    "        \"\"\"\n",
    "            'MATCH (from:User {id: line.from_id}), (to:User {id: line.to_id})\n",
    "            CREATE (from)-[:Replied {\n",
    "                count: toInteger(line.count),\n",
    "                depth_score: toInteger(line.depth_score),\n",
    "                up_votes: toInteger(line.up_votes),\n",
    "                up_votes_with_depth: toInteger(line.up_votes_with_depth)\n",
    "            }]->(to)'\n",
    "        , {batchSize:100000, iterateList:true, parallel:false})\"\"\"\n",
    "     )\n",
    "    \n",
    "def load_authors_to_symbols(file, tx):\n",
    "    tx.run('CALL apoc.periodic.iterate(' +\n",
    "        f\"\"\"'LOAD CSV WITH HEADERS FROM \"file:///{file}\" AS line RETURN line', \"\"\" +  \n",
    "        \"\"\"\n",
    "            'MATCH (user:User {id: line.author_id}), (symbol:Symbol {id: line.symbol})\n",
    "            CREATE (user)-[:Mentioned {\n",
    "                count: toInteger(line.count),\n",
    "                depth_score: toInteger(line.depth_score),\n",
    "                up_votes: toInteger(line.up_votes),\n",
    "                up_votes_with_depth: toInteger(line.up_votes_with_depth)\n",
    "            }]->(symbol)'\n",
    "        , {batchSize:100000, iterateList:true, parallel:false})\"\"\"\n",
    "     )\n",
    "    \n",
    "def load_authors_to_posts(file, tx):\n",
    "    tx.run('CALL apoc.periodic.iterate(' +\n",
    "        f\"\"\"'LOAD CSV WITH HEADERS FROM \"file:///{file}\" AS line RETURN line', \"\"\" +  \n",
    "        \"\"\"\n",
    "            'MATCH (user:User {id: line.author_id}), (post:Post {id: line.post_id})\n",
    "            CREATE (user)-[:Post {\n",
    "                count: toInteger(line.count),\n",
    "                depth_score: toInteger(line.depth_score),\n",
    "                up_votes: toInteger(line.up_votes),\n",
    "                up_votes_with_depth: toInteger(line.up_votes_with_depth)\n",
    "            }]->(post)'\n",
    "        , {batchSize:100000, iterateList:true, parallel:false})\"\"\"\n",
    "     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "related-office",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the vertices and indices\n",
    "with driver.session() as session:\n",
    "    session.write_transaction(partial(load_users, USERS_LOC))\n",
    "    session.write_transaction(partial(load_posts, POSTS_LOC))\n",
    "    session.write_transaction(partial(load_symbols, SYMBOLS_LOC))\n",
    "    \n",
    "    session.write_transaction(create_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ruled-trout",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the edges\n",
    "with driver.session() as session:\n",
    "    session.write_transaction(partial(load_authors_to_posts, AUTHORS_TO_POSTS_LOC))\n",
    "    session.write_transaction(partial(load_posts_to_symbols, POSTS_TO_SYMBOLS_LOC))\n",
    "    session.write_transaction(partial(load_authors_to_authors, AUTHORS_TO_AUTHORS_LOC))\n",
    "    session.write_transaction(partial(load_authors_to_symbols, AUTHORS_TO_SYMBOLS_LOC))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compact-cinema",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "coastal-guest",
   "metadata": {},
   "outputs": [],
   "source": [
    "driver.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "reddit",
   "language": "python",
   "name": "reddit"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}