Created
June 9, 2016 21:39
-
-
Save wesm/b9c1bcf1dfd7456ce052da6aca9a2e29 to your computer and use it in GitHub Desktop.
Implementing a pandas idiom using ibis + window functions
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{ | |
"cells": [ | |
{ | |
"cell_type": "code", | |
"execution_count": 10, | |
"metadata": { | |
"collapsed": true | |
}, | |
"outputs": [], | |
"source": [ | |
"import ibis\n", | |
"import pandas as pd\n", | |
"\n", | |
"hdfs = ibis.hdfs_connect(port=5070)\n", | |
"con = ibis.impala.connect(hdfs_client=hdfs)" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 6, | |
"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>key</th>\n", | |
" <th>value1</th>\n", | |
" <th>value2</th>\n", | |
" </tr>\n", | |
" </thead>\n", | |
" <tbody>\n", | |
" <tr>\n", | |
" <th>0</th>\n", | |
" <td>a</td>\n", | |
" <td>0</td>\n", | |
" <td>6</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>1</th>\n", | |
" <td>a</td>\n", | |
" <td>1</td>\n", | |
" <td>7</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>2</th>\n", | |
" <td>b</td>\n", | |
" <td>2</td>\n", | |
" <td>8</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>3</th>\n", | |
" <td>b</td>\n", | |
" <td>3</td>\n", | |
" <td>9</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>4</th>\n", | |
" <td>c</td>\n", | |
" <td>4</td>\n", | |
" <td>10</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>5</th>\n", | |
" <td>c</td>\n", | |
" <td>5</td>\n", | |
" <td>11</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>0</th>\n", | |
" <td>a</td>\n", | |
" <td>0</td>\n", | |
" <td>6</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>1</th>\n", | |
" <td>a</td>\n", | |
" <td>1</td>\n", | |
" <td>7</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>2</th>\n", | |
" <td>b</td>\n", | |
" <td>2</td>\n", | |
" <td>8</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>3</th>\n", | |
" <td>b</td>\n", | |
" <td>3</td>\n", | |
" <td>9</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>4</th>\n", | |
" <td>c</td>\n", | |
" <td>4</td>\n", | |
" <td>10</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>5</th>\n", | |
" <td>c</td>\n", | |
" <td>5</td>\n", | |
" <td>11</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>0</th>\n", | |
" <td>a</td>\n", | |
" <td>0</td>\n", | |
" <td>6</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>1</th>\n", | |
" <td>a</td>\n", | |
" <td>1</td>\n", | |
" <td>7</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>2</th>\n", | |
" <td>b</td>\n", | |
" <td>2</td>\n", | |
" <td>8</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>3</th>\n", | |
" <td>b</td>\n", | |
" <td>3</td>\n", | |
" <td>9</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>4</th>\n", | |
" <td>c</td>\n", | |
" <td>4</td>\n", | |
" <td>10</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>5</th>\n", | |
" <td>c</td>\n", | |
" <td>5</td>\n", | |
" <td>11</td>\n", | |
" </tr>\n", | |
" </tbody>\n", | |
"</table>\n", | |
"</div>" | |
], | |
"text/plain": [ | |
" key value1 value2\n", | |
"0 a 0 6\n", | |
"1 a 1 7\n", | |
"2 b 2 8\n", | |
"3 b 3 9\n", | |
"4 c 4 10\n", | |
"5 c 5 11\n", | |
"0 a 0 6\n", | |
"1 a 1 7\n", | |
"2 b 2 8\n", | |
"3 b 3 9\n", | |
"4 c 4 10\n", | |
"5 c 5 11\n", | |
"0 a 0 6\n", | |
"1 a 1 7\n", | |
"2 b 2 8\n", | |
"3 b 3 9\n", | |
"4 c 4 10\n", | |
"5 c 5 11" | |
] | |
}, | |
"execution_count": 6, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"df = pd.DataFrame({\n", | |
" 'key': ['a', 'a', 'b', 'b', 'c', 'c'],\n", | |
" 'value1': range(6),\n", | |
" 'value2': range(6, 12)\n", | |
"})\n", | |
"\n", | |
"df = pd.concat([df] * 3)\n", | |
"df" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 7, | |
"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>key</th>\n", | |
" <th>value1</th>\n", | |
" <th>value2</th>\n", | |
" </tr>\n", | |
" </thead>\n", | |
" <tbody>\n", | |
" <tr>\n", | |
" <th>0</th>\n", | |
" <td>a</td>\n", | |
" <td>0</td>\n", | |
" <td>6</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>2</th>\n", | |
" <td>b</td>\n", | |
" <td>2</td>\n", | |
" <td>8</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>4</th>\n", | |
" <td>c</td>\n", | |
" <td>4</td>\n", | |
" <td>10</td>\n", | |
" </tr>\n", | |
" </tbody>\n", | |
"</table>\n", | |
"</div>" | |
], | |
"text/plain": [ | |
" key value1 value2\n", | |
"0 a 0 6\n", | |
"2 b 2 8\n", | |
"4 c 4 10" | |
] | |
}, | |
"execution_count": 7, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"df.drop_duplicates(['key'], keep='first')" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 8, | |
"metadata": { | |
"collapsed": true | |
}, | |
"outputs": [], | |
"source": [ | |
"con.create_database('my_test')" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 11, | |
"metadata": { | |
"collapsed": false | |
}, | |
"outputs": [], | |
"source": [ | |
"db = con.database('my_test')\n", | |
"db.create_table('df_test', df)" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 12, | |
"metadata": { | |
"collapsed": false | |
}, | |
"outputs": [ | |
{ | |
"data": { | |
"text/plain": [ | |
"DatabaseTable[table]\n", | |
" name: my_test.`df_test`\n", | |
" schema:\n", | |
" key : string\n", | |
" value1 : int64\n", | |
" value2 : int64" | |
] | |
}, | |
"execution_count": 12, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"t = db.table('df_test')\n", | |
"\n", | |
"t" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 16, | |
"metadata": { | |
"collapsed": false | |
}, | |
"outputs": [ | |
{ | |
"name": "stdout", | |
"output_type": "stream", | |
"text": [ | |
"SELECT *,\n", | |
" row_number() OVER (PARTITION BY `key` ORDER BY `key`) - 1 AS `row_index`\n", | |
"FROM my_test.`df_test`\n", | |
"LIMIT 10000\n" | |
] | |
}, | |
{ | |
"data": { | |
"text/html": [ | |
"<div>\n", | |
"<table border=\"1\" class=\"dataframe\">\n", | |
" <thead>\n", | |
" <tr style=\"text-align: right;\">\n", | |
" <th></th>\n", | |
" <th>key</th>\n", | |
" <th>value1</th>\n", | |
" <th>value2</th>\n", | |
" <th>row_index</th>\n", | |
" </tr>\n", | |
" </thead>\n", | |
" <tbody>\n", | |
" <tr>\n", | |
" <th>0</th>\n", | |
" <td>a</td>\n", | |
" <td>0</td>\n", | |
" <td>6</td>\n", | |
" <td>0</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>1</th>\n", | |
" <td>a</td>\n", | |
" <td>1</td>\n", | |
" <td>7</td>\n", | |
" <td>1</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>2</th>\n", | |
" <td>a</td>\n", | |
" <td>1</td>\n", | |
" <td>7</td>\n", | |
" <td>2</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>3</th>\n", | |
" <td>a</td>\n", | |
" <td>0</td>\n", | |
" <td>6</td>\n", | |
" <td>3</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>4</th>\n", | |
" <td>a</td>\n", | |
" <td>0</td>\n", | |
" <td>6</td>\n", | |
" <td>4</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>5</th>\n", | |
" <td>a</td>\n", | |
" <td>1</td>\n", | |
" <td>7</td>\n", | |
" <td>5</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>6</th>\n", | |
" <td>b</td>\n", | |
" <td>3</td>\n", | |
" <td>9</td>\n", | |
" <td>0</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>7</th>\n", | |
" <td>b</td>\n", | |
" <td>2</td>\n", | |
" <td>8</td>\n", | |
" <td>1</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>8</th>\n", | |
" <td>b</td>\n", | |
" <td>3</td>\n", | |
" <td>9</td>\n", | |
" <td>2</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>9</th>\n", | |
" <td>b</td>\n", | |
" <td>2</td>\n", | |
" <td>8</td>\n", | |
" <td>3</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>10</th>\n", | |
" <td>b</td>\n", | |
" <td>3</td>\n", | |
" <td>9</td>\n", | |
" <td>4</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>11</th>\n", | |
" <td>b</td>\n", | |
" <td>2</td>\n", | |
" <td>8</td>\n", | |
" <td>5</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>12</th>\n", | |
" <td>c</td>\n", | |
" <td>4</td>\n", | |
" <td>10</td>\n", | |
" <td>0</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>13</th>\n", | |
" <td>c</td>\n", | |
" <td>5</td>\n", | |
" <td>11</td>\n", | |
" <td>1</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>14</th>\n", | |
" <td>c</td>\n", | |
" <td>5</td>\n", | |
" <td>11</td>\n", | |
" <td>2</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>15</th>\n", | |
" <td>c</td>\n", | |
" <td>4</td>\n", | |
" <td>10</td>\n", | |
" <td>3</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>16</th>\n", | |
" <td>c</td>\n", | |
" <td>4</td>\n", | |
" <td>10</td>\n", | |
" <td>4</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>17</th>\n", | |
" <td>c</td>\n", | |
" <td>5</td>\n", | |
" <td>11</td>\n", | |
" <td>5</td>\n", | |
" </tr>\n", | |
" </tbody>\n", | |
"</table>\n", | |
"</div>" | |
], | |
"text/plain": [ | |
" key value1 value2 row_index\n", | |
"0 a 0 6 0\n", | |
"1 a 1 7 1\n", | |
"2 a 1 7 2\n", | |
"3 a 0 6 3\n", | |
"4 a 0 6 4\n", | |
"5 a 1 7 5\n", | |
"6 b 3 9 0\n", | |
"7 b 2 8 1\n", | |
"8 b 3 9 2\n", | |
"9 b 2 8 3\n", | |
"10 b 3 9 4\n", | |
"11 b 2 8 5\n", | |
"12 c 4 10 0\n", | |
"13 c 5 11 1\n", | |
"14 c 5 11 2\n", | |
"15 c 4 10 3\n", | |
"16 c 4 10 4\n", | |
"17 c 5 11 5" | |
] | |
}, | |
"execution_count": 16, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"ibis.options.verbose = True\n", | |
"\n", | |
"expr = (t.group_by('key')\n", | |
" .order_by('key')\n", | |
" .mutate(row_index=ibis.row_number()))\n", | |
"\n", | |
"expr.execute()" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 29, | |
"metadata": { | |
"collapsed": false | |
}, | |
"outputs": [ | |
{ | |
"name": "stdout", | |
"output_type": "stream", | |
"text": [ | |
"SELECT t0.*\n", | |
"FROM (\n", | |
" SELECT *,\n", | |
" first_value(`row_index`) OVER (PARTITION BY `key` ORDER BY `row_index`) AS `first_row_index`\n", | |
" FROM (\n", | |
" SELECT *,\n", | |
" row_number() OVER (PARTITION BY `key` ORDER BY `key`) - 1 AS `row_index`\n", | |
" FROM my_test.`df_test`\n", | |
" ) t2\n", | |
") t0\n", | |
"WHERE t0.`first_row_index` = t0.`row_index`\n", | |
"LIMIT 10000\n" | |
] | |
}, | |
{ | |
"data": { | |
"text/html": [ | |
"<div>\n", | |
"<table border=\"1\" class=\"dataframe\">\n", | |
" <thead>\n", | |
" <tr style=\"text-align: right;\">\n", | |
" <th></th>\n", | |
" <th>key</th>\n", | |
" <th>value1</th>\n", | |
" <th>value2</th>\n", | |
" <th>row_index</th>\n", | |
" <th>first_row_index</th>\n", | |
" </tr>\n", | |
" </thead>\n", | |
" <tbody>\n", | |
" <tr>\n", | |
" <th>0</th>\n", | |
" <td>a</td>\n", | |
" <td>0</td>\n", | |
" <td>6</td>\n", | |
" <td>0</td>\n", | |
" <td>0</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>1</th>\n", | |
" <td>b</td>\n", | |
" <td>3</td>\n", | |
" <td>9</td>\n", | |
" <td>0</td>\n", | |
" <td>0</td>\n", | |
" </tr>\n", | |
" <tr>\n", | |
" <th>2</th>\n", | |
" <td>c</td>\n", | |
" <td>4</td>\n", | |
" <td>10</td>\n", | |
" <td>0</td>\n", | |
" <td>0</td>\n", | |
" </tr>\n", | |
" </tbody>\n", | |
"</table>\n", | |
"</div>" | |
], | |
"text/plain": [ | |
" key value1 value2 row_index first_row_index\n", | |
"0 a 0 6 0 0\n", | |
"1 b 3 9 0 0\n", | |
"2 c 4 10 0 0" | |
] | |
}, | |
"execution_count": 29, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"expr2 = (expr.group_by('key')\n", | |
" .mutate(first_row_index=lambda x: x.row_index.first()))\n", | |
"\n", | |
"result = expr2[expr2.first_row_index == expr2.row_index]\n", | |
"result.execute()" | |
] | |
} | |
], | |
"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 | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment