Skip to content

Instantly share code, notes, and snippets.

@wesm
Created June 9, 2016 21:39
Show Gist options
  • Save wesm/b9c1bcf1dfd7456ce052da6aca9a2e29 to your computer and use it in GitHub Desktop.
Save wesm/b9c1bcf1dfd7456ce052da6aca9a2e29 to your computer and use it in GitHub Desktop.
Implementing a pandas idiom using ibis + window functions
Display the source blob
Display the rendered blob
Raw
{
"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