Pandas_tosql_upsert.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sqlalchemy import create_engine\n",
    "import threading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>A</th>\n",
       "      <th>B</th>\n",
       "      <th>C</th>\n",
       "      <th>D</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>8668</th>\n",
       "      <td>47</td>\n",
       "      <td>50</td>\n",
       "      <td>46</td>\n",
       "      <td>33</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       A   B   C   D\n",
       "8668  47  50  46  33"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Get sample data\n",
    "df = pd.DataFrame(np.random.randint(0,100,size=(100000, 4)), columns=list('ABCD'))\n",
    "df.drop_duplicates(['A', 'B'], keep='last', inplace=True)\n",
    "\n",
    "df.head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Create our test database for upserts (this is postgreSQL)\n",
    "DB_TYPE = 'postgresql'\n",
    "DB_DRIVER = 'psycopg2'\n",
    "DB_USER = 'admin'\n",
    "DB_PASS = 'password'\n",
    "DB_HOST = 'localhost'\n",
    "DB_PORT = '5432'\n",
    "DB_NAME = 'pandas_upsert'\n",
    "POOL_SIZE = 50\n",
    "### Config update complete ###\n",
    "SQLALCHEMY_DATABASE_URI = '%s+%s://%s:%s@%s:%s/%s' %(DB_TYPE, DB_DRIVER, DB_USER,\n",
    "                                                     DB_PASS, DB_HOST, DB_PORT, DB_NAME)\n",
    "#Add more threads to the pool for execution\n",
    "engine = create_engine(SQLALCHEMY_DATABASE_URI, pool_size=POOL_SIZE, max_overflow=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def to_sql_newrows(df, pool_size, *args, **kargs):\n",
    "    \"\"\"\n",
    "    Extend the Python pandas to_sql() method to thread database insertion\n",
    "\n",
    "    Required: \n",
    "        df : pandas dataframe to insert new rows into a database table\n",
    "        POOL_SIZE : your sqlalchemy max connection pool size.  Set < your db connection limit.\n",
    "                    Example where this matters: your cloud DB has a connection limit.\n",
    "    *args:\n",
    "        Pandas to_sql() arguments.  \n",
    "\n",
    "        Required arguments are:\n",
    "            tablename : Database table name to write results to\n",
    "            engine : SqlAlchemy engine\n",
    "\n",
    "        Optional arguments are:\n",
    "            'if_exists' : 'append' or 'replace'.  If table already exists, use append.\n",
    "            'index' : True or False.  True if you want to write index values to the db.\n",
    "\n",
    "\n",
    "    Credits for intial threading code:\n",
    "        http://techyoubaji.blogspot.com/2015/10/speed-up-pandas-tosql-with.html\n",
    "    \"\"\"\n",
    "\n",
    "    CHUNKSIZE = 1000\n",
    "    INITIAL_CHUNK = 100\n",
    "    if len(df) > CHUNKSIZE:\n",
    "        #write the initial chunk to the database if df is bigger than chunksize\n",
    "        df.iloc[:INITIAL_CHUNK, :].to_sql(*args, **kargs)\n",
    "    else:\n",
    "        #if df is smaller than chunksize, just write it to the db now\n",
    "        df.to_sql(*args, **kargs)\n",
    "\n",
    "    workers, i = [], 0\n",
    "\n",
    "    for i in range((df.shape[0] - INITIAL_CHUNK)/CHUNKSIZE):\n",
    "        t = threading.Thread(target=lambda: df.iloc[INITIAL_CHUNK+i*CHUNKSIZE:INITIAL_CHUNK+(i+1)*CHUNKSIZE].to_sql(*args, **kargs))\n",
    "        t.start()\n",
    "        workers.append(t)\n",
    "        \n",
    "    df.iloc[INITIAL_CHUNK+(i+1)*CHUNKSIZE:, :].to_sql(*args, **kargs)\n",
    "    [t.join() for t in workers]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Create a table with a unique constraint on A and B columns.\n",
    "\n",
    "engine.execute(\"\"\"DROP TABLE IF EXISTS \"test_upsert\" \"\"\")\n",
    "\n",
    "engine.execute(\"\"\"CREATE TABLE \"test_upsert\" (\n",
    "                  \"A\" INTEGER,\n",
    "                  \"B\" INTEGER,\n",
    "                  \"C\" INTEGER,\n",
    "                  \"D\" INTEGER,\n",
    "                  CONSTRAINT pk_A_B PRIMARY KEY (\"A\",\"B\")) \n",
    "                  \"\"\")\n",
    "\"\"\"\n",
    "#Add unique constraint to table\n",
    "try:\n",
    "    args = ' ALTER TABLE test_upsert ADD CONSTRAINT uk_a_b UNIQUE (\"A\", \"B\") '\n",
    "    results = engine.execute(args)\n",
    "    print 'success'\n",
    "except:\n",
    "    print 'unique constraint already exists'\n",
    "\"\"\"\n",
    "\n",
    "#Insert data using pandas.to_sql\n",
    "df.to_sql('test_upsert', engine, if_exists='append', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>A</th>\n",
       "      <th>B</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>47</td>\n",
       "      <td>50</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    A   B\n",
       "0  47  50"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Check that the table exists and there is data in it\n",
    "df_in_db = pd.read_sql_query('SELECT \"A\", \"B\" FROM test_upsert', engine)\n",
    "df_in_db.head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>A</th>\n",
       "      <th>B</th>\n",
       "      <th>C</th>\n",
       "      <th>D</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>905</td>\n",
       "      <td>368</td>\n",
       "      <td>145</td>\n",
       "      <td>477</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     A    B    C    D\n",
       "0  905  368  145  477"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Now lets bring in some new data to insert, along with the same old data\n",
    "df_new = df = pd.DataFrame(np.random.randint(0,1000,size=(100000, 4)), columns=list('ABCD'))\n",
    "df2 = pd.concat([df, df_new])\n",
    "df2.head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new df is 200000 rows, and data in db is 10000 rows\n"
     ]
    }
   ],
   "source": [
    "#First let's get the length of both dataframe\n",
    "len_df2 = df2.shape[0]\n",
    "len_df_in_db = df_in_db.shape[0]\n",
    "print ('new df is %s rows, and data in db is %s rows') %(len_df2, len_df_in_db)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "left joined df is 94267 rows\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>A</th>\n",
       "      <th>B</th>\n",
       "      <th>C</th>\n",
       "      <th>D</th>\n",
       "      <th>_merge</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>905</td>\n",
       "      <td>368</td>\n",
       "      <td>145</td>\n",
       "      <td>477</td>\n",
       "      <td>left_only</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     A    B    C    D     _merge\n",
       "0  905  368  145  477  left_only"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Now let's find out what rows are duplicates using a combination of left outer join and select where the join type is 'left only'\n",
    "#The new _merge column added via the new indicator column in pandas will help us greatly here\n",
    "\n",
    "#Lets use a self join to make sure our sample data doesnt have duplicates\n",
    "df2.drop_duplicates(['A', 'B'], keep='last', inplace=True)\n",
    "df_all = pd.merge(df2, df_in_db, how='left', on=['A', 'B'], \n",
    "                     copy=False, indicator=True, suffixes=['', '_in_db'])\n",
    "df_all.reset_index(inplace=True, drop=True)\n",
    "df_all = df_all[df_all['_merge']=='left_only']\n",
    "\n",
    "print 'left joined df is %s rows' %(df_all.shape[0])\n",
    "df_all.head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['_merge']\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>A</th>\n",
       "      <th>B</th>\n",
       "      <th>C</th>\n",
       "      <th>D</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>905</td>\n",
       "      <td>368</td>\n",
       "      <td>145</td>\n",
       "      <td>477</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     A    B    C    D\n",
       "0  905  368  145  477"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Now let's drop any columns that are in \"both\" or the \"right only (in the datbase)\n",
    "cols_to_drop = list([col for col in df_all.columns \\\n",
    "                     if '_in_db' in col \\\n",
    "                     or 'ID' in col\n",
    "                     or 'index' in col\n",
    "                     or '_merge' in col])\n",
    "print cols_to_drop\n",
    "df_unique = df_all.drop(cols_to_drop, axis=1)\n",
    "df_unique.head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "checking df size\n",
      "connection pool not large enough - setting to max pool size\n",
      "threadcount set to 50, initial chunk set to 44267\n",
      "threading...\n",
      "complete - joining results!\n",
      "threading job ended!\n"
     ]
    }
   ],
   "source": [
    "#Trying multithreaded insert new rows into database...\n",
    "tosql(df_unique, POOL_SIZE, upsert=True, 'test_upsert', engine, if_exists='append', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Now let's put the new records in df_unique back into the database!\n",
    "df_unique.to_sql('test_upsert', engine, if_exists='append', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>count</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>104314</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    count\n",
       "0  104314"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_in_db_2 = pd.read_sql_query('SELECT count(\"A\") FROM test_upsert', engine)\n",
    "df_in_db_2.head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>A</th>\n",
       "      <th>row</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Empty DataFrame\n",
       "Columns: [A, row]\n",
       "Index: []"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Awesome, that seems to work.  We only inserted new rows!  Let's check to make sure it's unique\n",
    "df_dupscount = pd.read_sql_query(\"\"\"\n",
    "                                select * from (\n",
    "                                      SELECT \"A\",\n",
    "                                      ROW_NUMBER() OVER(PARTITION BY \"A\", \"B\" ORDER BY \"A\" asc) AS Row\n",
    "                                      FROM test_upsert\n",
    "                                    ) dups\n",
    "                                    where \n",
    "                                    dups.Row > 1\n",
    "                                \"\"\", engine)\n",
    "df_dupscount.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}