{ "cells": [ { "cell_type": "markdown", "id": "59204560", "metadata": {}, "source": [ "# Quick Start" ] }, { "cell_type": "markdown", "id": "f52a6720", "metadata": {}, "source": [ "This first example uses `pandas`, because it remains the most popular dataframe library for python. However, exactly the same methods would be called for any of the other supported backends.\n", "\n", "Import the packages that we'll use: " ] }, { "cell_type": "code", "execution_count": 1, "id": "c371c93d", "metadata": {}, "outputs": [], "source": [ "import awkward as ak\n", "import akimbo.pandas\n", "import pandas as pd\n", "import numpy as np" ] }, { "cell_type": "markdown", "id": "62af7184", "metadata": {}, "source": [ "## Vectorizing ragged data" ] }, { "cell_type": "markdown", "id": "199cc9d4", "metadata": {}, "source": [ "Consider a series made of python lists. This may happen a lot in ``pandas``. It is also possible in other dataframe libraries, but less likely. This data is only a couple of MB big, and the simplest amount of ragged nesting imaginable." ] }, { "cell_type": "code", "execution_count": 2, "id": "a4d0d2dc", "metadata": {}, "outputs": [], "source": [ "s = pd.Series([[1, 2, 3], [0], [4, 5]] * 100000)" ] }, { "cell_type": "code", "execution_count": 3, "id": "aca61be9", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 [1, 2, 3]\n", "1 [0]\n", "2 [4, 5]\n", "3 [1, 2, 3]\n", "4 [0]\n", " ... \n", "299995 [0]\n", "299996 [4, 5]\n", "299997 [1, 2, 3]\n", "299998 [0]\n", "299999 [4, 5]\n", "Length: 300000, dtype: object" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "s" ] }, { "cell_type": "markdown", "id": "f8bc0a8f-4716-4c28-a14b-a57455e84c8a", "metadata": {}, "source": [ "First let's do a super simple operation: get the maximum of each list. There are a number of different ways to do this, we'll comment and time several." ] }, { "cell_type": "markdown", "id": "9a065083", "metadata": {}, "source": [ "We can put the series in a DataFrame with another built-in pandas type, e.g. a column of integers:" ] }, { "cell_type": "code", "execution_count": 4, "id": "59b6fc51-4863-42a9-927c-cdc7ef97e25b", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "numpy function\n", "911 ms ± 27 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n", "\n", "python function\n", "49.6 ms ± 402 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "\n", "comprehension/iteration\n", "34.5 ms ± 1.26 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "\n", "ak with conversion\n", "34.8 ms ± 285 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "\n", "ak after conversion\n", "3.49 ms ± 148 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" ] } ], "source": [ "print(\"\\nnumpy function\")\n", "%timeit s.map(np.max);\n", "print(\"\\npython function\")\n", "%timeit s.map(max);\n", "print(\"\\ncomprehension/iteration\")\n", "%timeit [max(_) for _ in s];\n", "print(\"\\nak with conversion\")\n", "%timeit s.ak.max(axis=1);\n", "print(\"\\nak after conversion\")\n", "s2 = s.ak.to_output()\n", "%timeit s2.ak.max(axis=1)" ] }, { "cell_type": "markdown", "id": "8f60fd45-8f5a-4900-9f45-2b66b457310c", "metadata": {}, "source": [ "Some interesting results!\n", "- numpy is terrible at this, where most of the cost is converting the lists to arrays. numpy is not esigned for tiny arrays\n", "- using builting python functions and iteraction is OK when the data size isn't too big; this doesn't scale to millions of elements or lists-of-lists\n", "- sometimes you can shave off runtime when you ignore the index; both ak versions maintain the index\n", "- ak is just as fast even accounting for converting the data; but if the data is already in optimized form (which also uses less memory), ak is **much** faster than any other method. There is no equivalent numpy representation of the data." ] }, { "cell_type": "markdown", "id": "179f068f-1a1e-4331-860a-3fca1b07f51b", "metadata": {}, "source": [ "**NOTE**:\n", " Pandas supports arrow storage of data such as this, and some IO functions can create it\n", "without intermediate python objects with the argument dtype_backend=\"pyarrow\". For dask,\n", "arrow is already the default, but object are still common, and for polars and cuDF,\n", "arrow is the only storage available, so you are guaranteed fast operations.\n" ] }, { "cell_type": "markdown", "id": "a2fc3169-bf85-431c-84fe-18bc78ca656e", "metadata": {}, "source": [ "### Nested records\n", "\n", "Let's look at a tiny example of nested record-oriented data. " ] }, { "cell_type": "markdown", "id": "603472af-68c7-49cf-afd1-cbbf6aadb00e", "metadata": {}, "source": [ "This small fake sports dataset contains some players names, their team, and how many goals they've scored in some variable number of games that they've appeared in.\n", "\n", "\n", "The raw data:" ] }, { "cell_type": "code", "execution_count": 5, "id": "65c9764d-ce9a-4330-8dd6-49d185c1232c", "metadata": {}, "outputs": [], "source": [ "text = \"\"\"- name: Bob\\n team: tigers\\n goals: [0, 0, 0, 1, 2, 0, 1]\\n\\n- name: Alice\\n team: bears\\n goals: [3, 2, 1, 0, 1]\\n\\n- name: Jack\\n team: bears\\n goals: [0, 0, 0, 0, 0, 0, 0, 0, 1]\\n\\n- name: Jill\\n team: bears\\n goals: [3, 0, 2]\\n\\n- name: Ted\\n team: tigers\\n goals: [0, 0, 0, 0, 0]\\n\\n- name: Ellen\\n team: tigers\\n goals: [1, 0, 0, 0, 2, 0, 1]\\n\\n- name: Dan\\n team: bears\\n goals: [0, 0, 3, 1, 0, 2, 0, 0]\\n\\n- name: Brad\\n team: bears\\n goals: [0, 0, 4, 0, 0, 1]\\n\\n- name: Nancy\\n team: tigers\\n goals: [0, 0, 1, 1, 1, 1, 0]\\n\\n- name: Lance\\n team: bears\\n goals: [1, 1, 1, 1, 1]\\n\\n- name: Sara\\n team: tigers\\n goals: [0, 1, 0, 2, 0, 3]\\n\\n- name: Ryan\\n team: tigers\\n goals: [1, 2, 3, 0, 0, 0, 0]\\n\"\"\"" ] }, { "cell_type": "markdown", "id": "0a016204-e2e3-43e4-805f-f923967696a9", "metadata": {}, "source": [ "This is in YAML format, so that we can include it in a single line. Notice that YAML allows us to see the nesting and variable-length of the data clearly. \n", "The data in YAML format:" ] }, { "cell_type": "code", "execution_count": 6, "id": "22b42724-4ca0-44cb-b2cf-a095ea00dac8", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "- name: Bob\n", " team: tigers\n", " goals: [0, 0, 0, 1, 2, 0, 1]\n", "\n", "- name: Alice\n", " team: bears\n", " goals: [3, 2, 1, 0, 1]\n", "\n", "- name: Jack\n", " team: bears\n", " goals: [0, 0, 0, 0, 0, 0, 0, 0, 1]\n", "\n", "- name: Jill\n", " team: bears\n", " goals: [3, 0, 2]\n", "\n", "- name: Ted\n", " team: tigers\n", " goals: [0, 0, 0, 0, 0]\n", "\n", "- name: Ellen\n", " team: tigers\n", " goals: [1, 0, 0, 0, 2, 0, 1]\n", "\n", "- name: Dan\n", " team: bears\n", " goals: [0, 0, 3, 1, 0, 2, 0, 0]\n", "\n", "- name: Brad\n", " team: bears\n", " goals: [0, 0, 4, 0, 0, 1]\n", "\n", "- name: Nancy\n", " team: tigers\n", " goals: [0, 0, 1, 1, 1, 1, 0]\n", "\n", "- name: Lance\n", " team: bears\n", " goals: [1, 1, 1, 1, 1]\n", "\n", "- name: Sara\n", " team: tigers\n", " goals: [0, 1, 0, 2, 0, 3]\n", "\n", "- name: Ryan\n", " team: tigers\n", " goals: [1, 2, 3, 0, 0, 0, 0]\n", "\n" ] } ], "source": [ "print(text)" ] }, { "cell_type": "markdown", "id": "56f7c0b5-693c-418b-9cee-cb85904e290b", "metadata": {}, "source": [ "Awkward Array happily deals with this kind of data:" ] }, { "cell_type": "code", "execution_count": 7, "id": "101d69c2-6d37-4994-9afc-6c197041a2c3", "metadata": {}, "outputs": [], "source": [ "import yaml\n", "\n", "dicts = yaml.safe_load(text)\n", "data = pd.Series(dicts).ak.to_output()" ] }, { "cell_type": "code", "execution_count": 8, "id": "3578b97d-aa57-4c5a-8a83-6f8107cdf7fa", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 {'goals': array([0, 0, 0, 1, 2, 0, 1]), 'name'...\n", "1 {'goals': array([3, 2, 1, 0, 1]), 'name': 'Ali...\n", "2 {'goals': array([0, 0, 0, 0, 0, 0, 0, 0, 1]), ...\n", "3 {'goals': array([3, 0, 2]), 'name': 'Jill', 't...\n", "4 {'goals': array([0, 0, 0, 0, 0]), 'name': 'Ted...\n", "5 {'goals': array([1, 0, 0, 0, 2, 0, 1]), 'name'...\n", "6 {'goals': array([0, 0, 3, 1, 0, 2, 0, 0]), 'na...\n", "7 {'goals': array([0, 0, 4, 0, 0, 1]), 'name': '...\n", "8 {'goals': array([0, 0, 1, 1, 1, 1, 0]), 'name'...\n", "9 {'goals': array([1, 1, 1, 1, 1]), 'name': 'Lan...\n", "10 {'goals': array([0, 1, 0, 2, 0, 3]), 'name': '...\n", "11 {'goals': array([1, 2, 3, 0, 0, 0, 0]), 'name'...\n", "dtype: struct, name: string, team: string>[pyarrow]" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "data" ] }, { "cell_type": "markdown", "id": "dfabb8e8-0764-4bb8-8b29-1ed60868e1a9", "metadata": {}, "source": [ "but we use `akimbo` to transform it into an arrow-backed Series. This will allow us to use dataframe functionality such as groupby, below." ] }, { "cell_type": "markdown", "id": "4bdaf488-1ab9-4c8c-8e0d-4ecc47aaa00a", "metadata": {}, "source": [ "The dataset in Awkward Array form as three fields: \"name\", \"team\" and \"goals\"" ] }, { "cell_type": "markdown", "id": "d60887e8-582b-474f-a79f-173bc62c4bd1", "metadata": {}, "source": [ "Of these, two are \"normal\" fields - they can be made into dataframe columns containing no nesting. To unwrap the top record-like structure of the data, we can use ``unpack``." ] }, { "cell_type": "code", "execution_count": 9, "id": "5139d4fd-c967-4a07-9734-8a235490c6f0", "metadata": {}, "outputs": [], "source": [ "df = data.ak.unpack()" ] }, { "cell_type": "markdown", "id": "494081c8-21e4-4228-8010-b1bdf7563c47", "metadata": {}, "source": [ "We can use pure Pandas to investigate the dataset, but since Pandas doesn't have a builtin ability to handle the nested structure of our `goals` column, we're limited to some coarse information.\n", "\n", "For example, we can group by the team and see the average number of goals _total_ goals scored. Here we use the ``.ak`` accessor _on each group_, to be able to do arithmetic on the variable-length data, but while maintaining the pandas index." ] }, { "cell_type": "code", "execution_count": 10, "id": "350b5b98-e91c-475d-b5f5-5fc7bda73194", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "team name \n", "bears Jill 1.666667\n", " Alice 1.400000\n", " Lance 1.000000\n", "tigers Sara 1.000000\n", " Ryan 0.857143\n", "bears Brad 0.833333\n", " Dan 0.750000\n", "tigers Bob 0.571429\n", " Ellen 0.571429\n", " Nancy 0.571429\n", "bears Jack 0.111111\n", "tigers Ted 0.000000\n", "dtype: double[pyarrow]" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "df.set_index(\"name\") \\\n", " .groupby(\"team\", group_keys=True) \\\n", " .apply(lambda x: x.goals.ak.mean(axis=1)) \\\n", " .sort_values(ascending=False)" ] }, { "cell_type": "markdown", "id": "c480ffeb-971b-4af8-b01e-dd6f6f868c5e", "metadata": {}, "source": [ "Determine how many games each player has appeared in is simpler, using a direct method:" ] }, { "cell_type": "code", "execution_count": 11, "id": "a6ac5ba1-b205-4620-835e-ea3acd2a20d0", "metadata": {}, "outputs": [], "source": [ "df[\"n_games\"] = df.goals.ak.num(axis=1)" ] }, { "cell_type": "code", "execution_count": 12, "id": "8abda5da-4214-44a1-bcab-7c9926f0ddae", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
goalsnameteamn_games
0[0 0 0 1 2 0 1]Bobtigers7
1[3 2 1 0 1]Alicebears5
2[0 0 0 0 0 0 0 0 1]Jackbears9
3[3 0 2]Jillbears3
4[0 0 0 0 0]Tedtigers5
5[1 0 0 0 2 0 1]Ellentigers7
6[0 0 3 1 0 2 0 0]Danbears8
7[0 0 4 0 0 1]Bradbears6
8[0 0 1 1 1 1 0]Nancytigers7
9[1 1 1 1 1]Lancebears5
10[0 1 0 2 0 3]Saratigers6
11[1 2 3 0 0 0 0]Ryantigers7
\n", "
" ], "text/plain": [ " goals name team n_games\n", "0 [0 0 0 1 2 0 1] Bob tigers 7\n", "1 [3 2 1 0 1] Alice bears 5\n", "2 [0 0 0 0 0 0 0 0 1] Jack bears 9\n", "3 [3 0 2] Jill bears 3\n", "4 [0 0 0 0 0] Ted tigers 5\n", "5 [1 0 0 0 2 0 1] Ellen tigers 7\n", "6 [0 0 3 1 0 2 0 0] Dan bears 8\n", "7 [0 0 4 0 0 1] Brad bears 6\n", "8 [0 0 1 1 1 1 0] Nancy tigers 7\n", "9 [1 1 1 1 1] Lance bears 5\n", "10 [0 1 0 2 0 3] Sara tigers 6\n", "11 [1 2 3 0 0 0 0] Ryan tigers 7" ] }, "execution_count": 12, "metadata": {}, "output_type": "execute_result" } ], "source": [ "df" ] }, { "cell_type": "markdown", "id": "241d6cd9-5a4d-4c17-ac10-fa6d26d3d0f0", "metadata": {}, "source": [ "We can also convert the entire dataframe (any dataframe, in fact) back to a `Series`, which is convenient if we want to drop down to the Awkward library for further operations." ] }, { "cell_type": "code", "execution_count": 13, "id": "a8e4dbd9-8a77-4b4a-a05c-18621aa81ead", "metadata": {}, "outputs": [], "source": [ "s = df.ak.pack()" ] }, { "cell_type": "code", "execution_count": 14, "id": "4fe023c4-b36a-4e8c-bc56-486241349497", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 {'goals': array([0, 0, 0, 1, 2, 0, 1]), 'name'...\n", "1 {'goals': array([3, 2, 1, 0, 1]), 'name': 'Ali...\n", "2 {'goals': array([0, 0, 0, 0, 0, 0, 0, 0, 1]), ...\n", "3 {'goals': array([3, 0, 2]), 'name': 'Jill', 't...\n", "4 {'goals': array([0, 0, 0, 0, 0]), 'name': 'Ted...\n", "5 {'goals': array([1, 0, 0, 0, 2, 0, 1]), 'name'...\n", "6 {'goals': array([0, 0, 3, 1, 0, 2, 0, 0]), 'na...\n", "7 {'goals': array([0, 0, 4, 0, 0, 1]), 'name': '...\n", "8 {'goals': array([0, 0, 1, 1, 1, 1, 0]), 'name'...\n", "9 {'goals': array([1, 1, 1, 1, 1]), 'name': 'Lan...\n", "10 {'goals': array([0, 1, 0, 2, 0, 3]), 'name': '...\n", "11 {'goals': array([1, 2, 3, 0, 0, 0, 0]), 'name'...\n", "dtype: struct not null, name: string not null, team: string not null, n_games: int32 not null>[pyarrow]" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "s # look at that complex dtype!" ] }, { "cell_type": "markdown", "id": "e3f6d87f-d427-4ae2-b9c6-6e38f0f975a2", "metadata": {}, "source": [ "And go back to pure awkward (now with our new `n_games` column) using the accessor:" ] }, { "cell_type": "code", "execution_count": 15, "id": "9fa44a70-63ee-4c4d-80eb-3e979c0337d7", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[{goals: [0, 0, 0, 1, 2, 0, 1], name: 'Bob', team: 'tigers', n_games: 7},\n",
       " {goals: [3, 2, 1, 0, 1], name: 'Alice', team: 'bears', n_games: 5},\n",
       " {goals: [0, 0, 0, 0, 0, 0, 0, 0, 1], name: 'Jack', team: 'bears', ...},\n",
       " {goals: [3, 0, 2], name: 'Jill', team: 'bears', n_games: 3},\n",
       " {goals: [0, 0, 0, 0, 0], name: 'Ted', team: 'tigers', n_games: 5},\n",
       " {goals: [1, 0, 0, 0, 2, 0, 1], name: 'Ellen', team: 'tigers', ...},\n",
       " {goals: [0, 0, 3, 1, 0, 2, 0, 0], name: 'Dan', team: 'bears', ...},\n",
       " {goals: [0, 0, 4, 0, 0, 1], name: 'Brad', team: 'bears', n_games: 6},\n",
       " {goals: [0, 0, 1, 1, 1, 1, 0], name: 'Nancy', team: 'tigers', ...},\n",
       " {goals: [1, 1, 1, 1, 1], name: 'Lance', team: 'bears', n_games: 5},\n",
       " {goals: [0, 1, 0, 2, 0, 3], name: 'Sara', team: 'tigers', n_games: 6},\n",
       " {goals: [1, 2, 3, 0, 0, 0, 0], name: 'Ryan', team: 'tigers', ...}]\n",
       "-------------------------------------------------------------------------\n",
       "type: 12 * {\n",
       "    goals: var * ?int64,\n",
       "    name: string,\n",
       "    team: string,\n",
       "    n_games: int32\n",
       "}
" ], "text/plain": [ "" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "s.ak.array" ] }, { "cell_type": "code", "execution_count": 16, "id": "cd165bf6-ae01-45b3-b290-3b091fc4b9e3", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "['goals', 'name', 'team', 'n_games']" ] }, "execution_count": 16, "metadata": {}, "output_type": "execute_result" } ], "source": [ "s.ak.array.fields" ] }, { "cell_type": "code", "execution_count": 17, "id": "bca4d5c4-5802-43f8-bc24-e59612465890", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 7\n", "1 5\n", "2 9\n", "3 3\n", "4 5\n", "5 7\n", "6 8\n", "7 6\n", "8 7\n", "9 5\n", "10 6\n", "11 7\n", "dtype: int32[pyarrow]" ] }, "execution_count": 17, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# as series\n", "s.ak[\"n_games\"]" ] }, { "cell_type": "code", "execution_count": 18, "id": "7efeb90b-6395-4b70-9134-1cd4842876ae", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[7,\n",
       " 5,\n",
       " 9,\n",
       " 3,\n",
       " 5,\n",
       " 7,\n",
       " 8,\n",
       " 6,\n",
       " 7,\n",
       " 5,\n",
       " 6,\n",
       " 7]\n",
       "----------------\n",
       "type: 12 * int32
" ], "text/plain": [ "" ] }, "execution_count": 18, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# as awkward\n", "s.ak.array[\"n_games\"]" ] }, { "cell_type": "markdown", "id": "bf84c1d3-e950-4359-a080-9ef0dbe09fbc", "metadata": {}, "source": [ "### Strings and datetimes\n", "\n", "Arrow provides \"compute\" functions for working with strings and datetimes, closely modelled on Pandas. However, ``akimbo`` gives you a way to use these functions without breaking up the original structure of your data, so that you can then do any aggregations you might need." ] }, { "cell_type": "code", "execution_count": 19, "id": "36058654-cf88-4d37-bb88-d6fa4ca840c7", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 ['Hello' 'Oi' 'hi' 'hola']\n", "1 ['yo' 'watcha']\n", "2 ['Hello' 'Oi' 'hi' 'hola']\n", "3 ['yo' 'watcha']\n", "4 ['Hello' 'Oi' 'hi' 'hola']\n", " ... \n", "19995 ['yo' 'watcha']\n", "19996 ['Hello' 'Oi' 'hi' 'hola']\n", "19997 ['yo' 'watcha']\n", "19998 ['Hello' 'Oi' 'hi' 'hola']\n", "19999 ['yo' 'watcha']\n", "Length: 20000, dtype: list[pyarrow]" ] }, "execution_count": 19, "metadata": {}, "output_type": "execute_result" } ], "source": [ "data = [[\"Hello\", \"Oi\", \"hi\", \"hola\"], [\"yo\", \"watcha\"]] * 10000\n", "s = pd.Series(data).ak.to_output()\n", "s" ] }, { "cell_type": "code", "execution_count": 20, "id": "f4893f92-b1f1-49e8-aa41-650b748d39be", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 True\n", "1 False\n", "2 True\n", "3 False\n", "4 True\n", " ... \n", "19995 False\n", "19996 True\n", "19997 False\n", "19998 True\n", "19999 False\n", "Length: 20000, dtype: bool[pyarrow]" ] }, "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# in each list, does any of the greetings start with \"h\" or \"H\"?\n", "s.ak.str.lower().ak.str.startswith(\"h\").ak.any(axis=1)\n", "# equivalent with harder syntax but fewer arrow roundtrips\n", "# s.ak.apply(lambda x: ak.any(ak.str.starts_with(ak.str.lower(x), \"h\"), axis=1))" ] }, { "cell_type": "code", "execution_count": 21, "id": "0e757441-1954-412c-8000-bffb83343214", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 ['2024-07-09T10:00:00' '2024-07-09T10:30:00' '...\n", "1 ['2024-07-08T10:00:00' '2024-07-09T10:00:00']\n", "2 ['2024-07-09T10:00:00' '2024-07-09T10:30:00' '...\n", "3 ['2024-07-08T10:00:00' '2024-07-09T10:00:00']\n", "4 ['2024-07-09T10:00:00' '2024-07-09T10:30:00' '...\n", " ... \n", "19995 ['2024-07-08T10:00:00' '2024-07-09T10:00:00']\n", "19996 ['2024-07-09T10:00:00' '2024-07-09T10:30:00' '...\n", "19997 ['2024-07-08T10:00:00' '2024-07-09T10:00:00']\n", "19998 ['2024-07-09T10:00:00' '2024-07-09T10:30:00' '...\n", "19999 ['2024-07-08T10:00:00' '2024-07-09T10:00:00']\n", "Length: 20000, dtype: list[pyarrow]" ] }, "execution_count": 21, "metadata": {}, "output_type": "execute_result" } ], "source": [ "times = [[\"2024-07-09T10:00:00\", \"2024-07-09T10:30:00\", \"2024-07-09T11:00:00\", \"2024-07-09T12:00:00\"],\n", " [\"2024-07-08T10:00:00\", \"2024-07-09T10:00:00\"]] * 10000\n", "s = pd.Series(times).ak.dt.cast(\"timestamp[s]\")\n", "s" ] }, { "cell_type": "code", "execution_count": 22, "id": "bd0d72d3-3153-4337-9680-6b4798cfd81f", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 [30 30 60]\n", "1 [1440]\n", "2 [30 30 60]\n", "3 [1440]\n", "4 [30 30 60]\n", " ... \n", "19995 [1440]\n", "19996 [30 30 60]\n", "19997 [1440]\n", "19998 [30 30 60]\n", "19999 [1440]\n", "Length: 20000, dtype: list[pyarrow]" ] }, "execution_count": 22, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# set of differences\n", "s.ak[:, :-1].ak.dt.minutes_between(s.ak[:, 1:])" ] }, { "cell_type": "markdown", "id": "97aa0ac7-f1d5-4344-ae2c-72f617ededea", "metadata": {}, "source": [ "### Behaviours\n", "\n", "Let's take an example from upsrteam documentation: vectors are made of two fields, `(x, y)`. We know that adding and the the size of a vector are easily expressed. Let's encode this in a class and apply it to data in a dataframe." ] }, { "cell_type": "code", "execution_count": 23, "id": "4620a04a-e719-4b34-b633-0325b9121e25", "metadata": {}, "outputs": [], "source": [ "from akimbo import mixin_class, mixin_class_method, behavior\n", "import akimbo.pandas\n", "import numpy as np\n", "import pandas as pd\n", "\n", "\n", "@mixin_class(behavior)\n", "class Point:\n", "\n", " @mixin_class_method(np.abs)\n", " def point_abs(self):\n", " return np.sqrt(self.x ** 2 + self.y ** 2)\n", "\n", " @mixin_class_method(np.add, {\"Point\"})\n", " def point_add(self, other):\n", " return ak.zip(\n", " {\"x\": self.x + other.x, \"y\": self.y + other.y}, with_name=\"Point\",\n", " )" ] }, { "cell_type": "code", "execution_count": 24, "id": "c793edea-ff27-4a80-9023-fe8ec045718d", "metadata": {}, "outputs": [], "source": [ "data = [{\"x\": 1, \"y\": 2}] * 100000\n", "s = pd.Series(data).ak.to_output() # store as arrow" ] }, { "cell_type": "code", "execution_count": 25, "id": "1891f7c8-eb75-4294-902a-08a12a4dc914", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "True" ] }, "execution_count": 25, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# check that the unary method is there; so tab-complete will work\n", "\"point_abs\" in dir(s.ak.with_behavior(\"Point\"))" ] }, { "cell_type": "code", "execution_count": 26, "id": "7f0f32c0-5bc8-4965-ae9d-42047923bb7e", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 2.236068\n", "1 2.236068\n", "2 2.236068\n", "3 2.236068\n", "4 2.236068\n", " ... \n", "99995 2.236068\n", "99996 2.236068\n", "99997 2.236068\n", "99998 2.236068\n", "99999 2.236068\n", "Length: 100000, dtype: double[pyarrow]" ] }, "execution_count": 26, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# call to get vector sizes\n", "s.ak.with_behavior(\"Point\").point_abs()" ] }, { "cell_type": "code", "execution_count": 27, "id": "23bd6b30-ed37-4581-aaa8-579e9856a4ee", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 2.236068\n", "1 2.236068\n", "2 2.236068\n", "3 2.236068\n", "4 2.236068\n", " ... \n", "99995 2.236068\n", "99996 2.236068\n", "99997 2.236068\n", "99998 2.236068\n", "99999 2.236068\n", "Length: 100000, dtype: double[pyarrow]" ] }, "execution_count": 27, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# or do the same with numpy ufunc\n", "np.abs(s.ak.with_behavior(\"Point\"))" ] }, { "cell_type": "code", "execution_count": 28, "id": "557925c5-dbb1-4b4f-a85a-0f79a2183b6e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "3.7 ms ± 218 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n", "62.1 ms ± 924 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" ] } ], "source": [ "import math\n", "%timeit np.abs(s.ak.with_behavior(\"Point\"))\n", "%timeit s.apply(lambda struct: math.sqrt(struct[\"x\"] ** 2 + struct[\"y\"] ** 2))" ] }, { "cell_type": "markdown", "id": "41cc90b2-20a6-486c-aebf-0584926758bc", "metadata": {}, "source": [ "Of course, we could have done the same by extracting out the arrays (e.g., `s.ak[\"x\"]`) and using numpy directly, which would have been as fast, but this way we have an object-like experience." ] }, { "cell_type": "markdown", "id": "c37ef35e-c46d-473a-bc62-9da382cecb6e", "metadata": {}, "source": [ "Similarly, we defined an overload to add Point arrays together (both operands must be Point type). \n", "A vector addition is performed. This also happens at obviously vectorized speed - but I am not even sure how you would perform the same thing using python dicts." ] }, { "cell_type": "code", "execution_count": 29, "id": "92d5dcfa-9eb0-4bf0-915e-d7e5c530afd5", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "2.59 ms ± 50.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" ] }, { "data": { "text/plain": [ "0 {'x': 2, 'y': 4}\n", "1 {'x': 2, 'y': 4}\n", "2 {'x': 2, 'y': 4}\n", "3 {'x': 2, 'y': 4}\n", "4 {'x': 2, 'y': 4}\n", " ... \n", "99995 {'x': 2, 'y': 4}\n", "99996 {'x': 2, 'y': 4}\n", "99997 {'x': 2, 'y': 4}\n", "99998 {'x': 2, 'y': 4}\n", "99999 {'x': 2, 'y': 4}\n", "Length: 100000, dtype: struct[pyarrow]" ] }, "execution_count": 29, "metadata": {}, "output_type": "execute_result" } ], "source": [ "%timeit s.ak.with_behavior(\"Point\") + s.ak.with_behavior(\"Point\")\n", "s.ak.with_behavior(\"Point\") + s.ak.with_behavior(\"Point\")" ] }, { "cell_type": "markdown", "id": "33947d87-ab7d-4ac9-996e-c8f83fc66c66", "metadata": {}, "source": [ "### Numba integration\n", "\n", "The numpy API is very nice and can do most things you will need. The object-oriented behaviours are very convenient.\n", "\n", "However, some algorithms are complex enough that you need to process with a custom function, and the data may be big and complex enough that python iteration over rows is simply not an option. A functional approach may also allow compute operations in a single-pass and without temporaries that cannot be done with the numpy API.\n", "\n", "Enter [`numba`](https://numba.pydata.org/) , a JIT-compiler for numerical python, turning iterative, loopy functions into their C-language equivalent. Let's take an example." ] }, { "cell_type": "code", "execution_count": 30, "id": "9eb1257e-4777-47ca-ae42-9b2db2719a23", "metadata": {}, "outputs": [], "source": [ "import numba\n", "\n", "\n", "def mean_of_second_biggest(arr):\n", " count = 0\n", " total = 0\n", " for row in arr:\n", " if len(row) < 2:\n", " continue\n", " max = row[0]\n", " if row[1] > max:\n", " second = max\n", " max = row[1]\n", " else:\n", " second = row[1]\n", " for x in row[2:]:\n", " if x > max:\n", " second = max\n", " max = x\n", " elif x > second:\n", " second = x\n", " count += 1\n", " total += second\n", " return total / count\n", " \n" ] }, { "cell_type": "code", "execution_count": 31, "id": "7dde1591-1436-44a1-9bbb-1e5f228f5173", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "2.0" ] }, "execution_count": 31, "metadata": {}, "output_type": "execute_result" } ], "source": [ "mean_of_second_biggest([[1], [3, 2, 1], [3, 4, 4], [0, 1, 0]]) # mean of 2, 4, and 0" ] }, { "cell_type": "code", "execution_count": 32, "id": "c358c23b-fa5f-409f-bc58-67333fbd9c23", "metadata": {}, "outputs": [], "source": [ "jsecond = numba.njit(mean_of_second_biggest)" ] }, { "cell_type": "code", "execution_count": 33, "id": "c893840a-2236-4678-8ac2-be5b2e7bb05f", "metadata": {}, "outputs": [], "source": [ "s = pd.Series([[1], [3, 2, 1], [3, 4, 4], [0, 1, 0]] * 100000).ak.to_output()" ] }, { "cell_type": "code", "execution_count": 34, "id": "a65feec8-a2cc-4341-b1e3-1fd93174a1bd", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 [1]\n", "1 [3 2 1]\n", "2 [3 4 4]\n", "3 [0 1 0]\n", "4 [1]\n", " ... \n", "399995 [0 1 0]\n", "399996 [1]\n", "399997 [3 2 1]\n", "399998 [3 4 4]\n", "399999 [0 1 0]\n", "Length: 400000, dtype: list[pyarrow]" ] }, "execution_count": 34, "metadata": {}, "output_type": "execute_result" } ], "source": [ "s" ] }, { "cell_type": "code", "execution_count": 35, "id": "227af214-3950-403b-ae6d-6521c59942cc", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "2.0" ] }, "execution_count": 35, "metadata": {}, "output_type": "execute_result" } ], "source": [ "s.ak.apply(jsecond)" ] }, { "cell_type": "code", "execution_count": 36, "id": "55eb2870-ab7b-4595-8e67-573cefe72d72", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "1.08 ms ± 37.7 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n", "501 ms ± 2.72 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] } ], "source": [ "# timings\n", "%timeit s.ak.apply(jsecond)\n", "%timeit mean_of_second_biggest(s)" ] }, { "cell_type": "markdown", "id": "43234e7a-794a-4bb3-b163-1df8d11afe09", "metadata": {}, "source": [ "Question: could you have done this with vectorized numpy calls?" ] }, { "cell_type": "code", "execution_count": null, "id": "b9c11c8c-9c9b-45c4-ba7a-86f0d741cfb6", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "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.10.9" } }, "nbformat": 4, "nbformat_minor": 5 }