Quick Start¶
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.
Import the packages that we’ll use:
[1]:
import awkward as ak
import akimbo.pandas
import pandas as pd
import numpy as np
Vectorizing ragged data¶
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.
[2]:
s = pd.Series([[1, 2, 3], [0], [4, 5]] * 100000)
[3]:
s
[3]:
0 [1, 2, 3]
1 [0]
2 [4, 5]
3 [1, 2, 3]
4 [0]
...
299995 [0]
299996 [4, 5]
299997 [1, 2, 3]
299998 [0]
299999 [4, 5]
Length: 300000, dtype: object
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.
We can put the series in a DataFrame with another built-in pandas type, e.g. a column of integers:
[4]:
print("\nnumpy function")
%timeit s.map(np.max);
print("\npython function")
%timeit s.map(max);
print("\ncomprehension/iteration")
%timeit [max(_) for _ in s];
print("\nak with conversion")
%timeit s.ak.max(axis=1);
print("\nak after conversion")
s2 = s.ak.to_output()
%timeit s2.ak.max(axis=1)
numpy function
911 ms ± 27 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
python function
49.6 ms ± 402 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
comprehension/iteration
34.5 ms ± 1.26 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
ak with conversion
34.8 ms ± 285 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
ak after conversion
3.49 ms ± 148 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Some interesting results!
numpy is terrible at this, where most of the cost is converting the lists to arrays. numpy is not esigned for tiny arrays
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
sometimes you can shave off runtime when you ignore the index; both ak versions maintain the index
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.
NOTE: Pandas supports arrow storage of data such as this, and some IO functions can create it without intermediate python objects with the argument dtype_backend=”pyarrow”. For dask, arrow is already the default, but object are still common, and for polars and cuDF, arrow is the only storage available, so you are guaranteed fast operations.
Nested records¶
Let’s look at a tiny example of nested record-oriented data.
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.
The raw data:
[5]:
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"""
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. The data in YAML format:
[6]:
print(text)
- name: Bob
team: tigers
goals: [0, 0, 0, 1, 2, 0, 1]
- name: Alice
team: bears
goals: [3, 2, 1, 0, 1]
- name: Jack
team: bears
goals: [0, 0, 0, 0, 0, 0, 0, 0, 1]
- name: Jill
team: bears
goals: [3, 0, 2]
- name: Ted
team: tigers
goals: [0, 0, 0, 0, 0]
- name: Ellen
team: tigers
goals: [1, 0, 0, 0, 2, 0, 1]
- name: Dan
team: bears
goals: [0, 0, 3, 1, 0, 2, 0, 0]
- name: Brad
team: bears
goals: [0, 0, 4, 0, 0, 1]
- name: Nancy
team: tigers
goals: [0, 0, 1, 1, 1, 1, 0]
- name: Lance
team: bears
goals: [1, 1, 1, 1, 1]
- name: Sara
team: tigers
goals: [0, 1, 0, 2, 0, 3]
- name: Ryan
team: tigers
goals: [1, 2, 3, 0, 0, 0, 0]
Awkward Array happily deals with this kind of data:
[7]:
import yaml
dicts = yaml.safe_load(text)
data = pd.Series(dicts).ak.to_output()
[8]:
data
[8]:
0 {'goals': array([0, 0, 0, 1, 2, 0, 1]), 'name'...
1 {'goals': array([3, 2, 1, 0, 1]), 'name': 'Ali...
2 {'goals': array([0, 0, 0, 0, 0, 0, 0, 0, 1]), ...
3 {'goals': array([3, 0, 2]), 'name': 'Jill', 't...
4 {'goals': array([0, 0, 0, 0, 0]), 'name': 'Ted...
5 {'goals': array([1, 0, 0, 0, 2, 0, 1]), 'name'...
6 {'goals': array([0, 0, 3, 1, 0, 2, 0, 0]), 'na...
7 {'goals': array([0, 0, 4, 0, 0, 1]), 'name': '...
8 {'goals': array([0, 0, 1, 1, 1, 1, 0]), 'name'...
9 {'goals': array([1, 1, 1, 1, 1]), 'name': 'Lan...
10 {'goals': array([0, 1, 0, 2, 0, 3]), 'name': '...
11 {'goals': array([1, 2, 3, 0, 0, 0, 0]), 'name'...
dtype: struct<goals: list<item: int64>, name: string, team: string>[pyarrow]
but we use akimbo to transform it into an arrow-backed Series. This will allow us to use dataframe functionality such as groupby, below.
The dataset in Awkward Array form as three fields: “name”, “team” and “goals”
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.
[9]:
df = data.ak.unpack()
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.
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.
[10]:
df.set_index("name") \
.groupby("team", group_keys=True) \
.apply(lambda x: x.goals.ak.mean(axis=1)) \
.sort_values(ascending=False)
[10]:
team name
bears Jill 1.666667
Alice 1.400000
Lance 1.000000
tigers Sara 1.000000
Ryan 0.857143
bears Brad 0.833333
Dan 0.750000
tigers Bob 0.571429
Ellen 0.571429
Nancy 0.571429
bears Jack 0.111111
tigers Ted 0.000000
dtype: double[pyarrow]
Determine how many games each player has appeared in is simpler, using a direct method:
[11]:
df["n_games"] = df.goals.ak.num(axis=1)
[12]:
df
[12]:
| goals | name | team | n_games | |
|---|---|---|---|---|
| 0 | [0 0 0 1 2 0 1] | Bob | tigers | 7 |
| 1 | [3 2 1 0 1] | Alice | bears | 5 |
| 2 | [0 0 0 0 0 0 0 0 1] | Jack | bears | 9 |
| 3 | [3 0 2] | Jill | bears | 3 |
| 4 | [0 0 0 0 0] | Ted | tigers | 5 |
| 5 | [1 0 0 0 2 0 1] | Ellen | tigers | 7 |
| 6 | [0 0 3 1 0 2 0 0] | Dan | bears | 8 |
| 7 | [0 0 4 0 0 1] | Brad | bears | 6 |
| 8 | [0 0 1 1 1 1 0] | Nancy | tigers | 7 |
| 9 | [1 1 1 1 1] | Lance | bears | 5 |
| 10 | [0 1 0 2 0 3] | Sara | tigers | 6 |
| 11 | [1 2 3 0 0 0 0] | Ryan | tigers | 7 |
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.
[13]:
s = df.ak.pack()
[14]:
s # look at that complex dtype!
[14]:
0 {'goals': array([0, 0, 0, 1, 2, 0, 1]), 'name'...
1 {'goals': array([3, 2, 1, 0, 1]), 'name': 'Ali...
2 {'goals': array([0, 0, 0, 0, 0, 0, 0, 0, 1]), ...
3 {'goals': array([3, 0, 2]), 'name': 'Jill', 't...
4 {'goals': array([0, 0, 0, 0, 0]), 'name': 'Ted...
5 {'goals': array([1, 0, 0, 0, 2, 0, 1]), 'name'...
6 {'goals': array([0, 0, 3, 1, 0, 2, 0, 0]), 'na...
7 {'goals': array([0, 0, 4, 0, 0, 1]), 'name': '...
8 {'goals': array([0, 0, 1, 1, 1, 1, 0]), 'name'...
9 {'goals': array([1, 1, 1, 1, 1]), 'name': 'Lan...
10 {'goals': array([0, 1, 0, 2, 0, 3]), 'name': '...
11 {'goals': array([1, 2, 3, 0, 0, 0, 0]), 'name'...
dtype: struct<goals: list<item: int64> not null, name: string not null, team: string not null, n_games: int32 not null>[pyarrow]
And go back to pure awkward (now with our new n_games column) using the accessor:
[15]:
s.ak.array
[15]:
[{goals: [0, 0, 0, 1, 2, 0, 1], name: 'Bob', team: 'tigers', n_games: 7},
{goals: [3, 2, 1, 0, 1], name: 'Alice', team: 'bears', n_games: 5},
{goals: [0, 0, 0, 0, 0, 0, 0, 0, 1], name: 'Jack', team: 'bears', ...},
{goals: [3, 0, 2], name: 'Jill', team: 'bears', n_games: 3},
{goals: [0, 0, 0, 0, 0], name: 'Ted', team: 'tigers', n_games: 5},
{goals: [1, 0, 0, 0, 2, 0, 1], name: 'Ellen', team: 'tigers', ...},
{goals: [0, 0, 3, 1, 0, 2, 0, 0], name: 'Dan', team: 'bears', ...},
{goals: [0, 0, 4, 0, 0, 1], name: 'Brad', team: 'bears', n_games: 6},
{goals: [0, 0, 1, 1, 1, 1, 0], name: 'Nancy', team: 'tigers', ...},
{goals: [1, 1, 1, 1, 1], name: 'Lance', team: 'bears', n_games: 5},
{goals: [0, 1, 0, 2, 0, 3], name: 'Sara', team: 'tigers', n_games: 6},
{goals: [1, 2, 3, 0, 0, 0, 0], name: 'Ryan', team: 'tigers', ...}]
-------------------------------------------------------------------------
type: 12 * {
goals: var * ?int64,
name: string,
team: string,
n_games: int32
}[16]:
s.ak.array.fields
[16]:
['goals', 'name', 'team', 'n_games']
[17]:
# as series
s.ak["n_games"]
[17]:
0 7
1 5
2 9
3 3
4 5
5 7
6 8
7 6
8 7
9 5
10 6
11 7
dtype: int32[pyarrow]
[18]:
# as awkward
s.ak.array["n_games"]
[18]:
[7, 5, 9, 3, 5, 7, 8, 6, 7, 5, 6, 7] ---------------- type: 12 * int32
Strings and datetimes¶
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.
[19]:
data = [["Hello", "Oi", "hi", "hola"], ["yo", "watcha"]] * 10000
s = pd.Series(data).ak.to_output()
s
[19]:
0 ['Hello' 'Oi' 'hi' 'hola']
1 ['yo' 'watcha']
2 ['Hello' 'Oi' 'hi' 'hola']
3 ['yo' 'watcha']
4 ['Hello' 'Oi' 'hi' 'hola']
...
19995 ['yo' 'watcha']
19996 ['Hello' 'Oi' 'hi' 'hola']
19997 ['yo' 'watcha']
19998 ['Hello' 'Oi' 'hi' 'hola']
19999 ['yo' 'watcha']
Length: 20000, dtype: list<item: string>[pyarrow]
[20]:
# in each list, does any of the greetings start with "h" or "H"?
s.ak.str.lower().ak.str.startswith("h").ak.any(axis=1)
# equivalent with harder syntax but fewer arrow roundtrips
# s.ak.apply(lambda x: ak.any(ak.str.starts_with(ak.str.lower(x), "h"), axis=1))
[20]:
0 True
1 False
2 True
3 False
4 True
...
19995 False
19996 True
19997 False
19998 True
19999 False
Length: 20000, dtype: bool[pyarrow]
[21]:
times = [["2024-07-09T10:00:00", "2024-07-09T10:30:00", "2024-07-09T11:00:00", "2024-07-09T12:00:00"],
["2024-07-08T10:00:00", "2024-07-09T10:00:00"]] * 10000
s = pd.Series(times).ak.dt.cast("timestamp[s]")
s
[21]:
0 ['2024-07-09T10:00:00' '2024-07-09T10:30:00' '...
1 ['2024-07-08T10:00:00' '2024-07-09T10:00:00']
2 ['2024-07-09T10:00:00' '2024-07-09T10:30:00' '...
3 ['2024-07-08T10:00:00' '2024-07-09T10:00:00']
4 ['2024-07-09T10:00:00' '2024-07-09T10:30:00' '...
...
19995 ['2024-07-08T10:00:00' '2024-07-09T10:00:00']
19996 ['2024-07-09T10:00:00' '2024-07-09T10:30:00' '...
19997 ['2024-07-08T10:00:00' '2024-07-09T10:00:00']
19998 ['2024-07-09T10:00:00' '2024-07-09T10:30:00' '...
19999 ['2024-07-08T10:00:00' '2024-07-09T10:00:00']
Length: 20000, dtype: list<item: timestamp[s]>[pyarrow]
[22]:
# set of differences
s.ak[:, :-1].ak.dt.minutes_between(s.ak[:, 1:])
[22]:
0 [30 30 60]
1 [1440]
2 [30 30 60]
3 [1440]
4 [30 30 60]
...
19995 [1440]
19996 [30 30 60]
19997 [1440]
19998 [30 30 60]
19999 [1440]
Length: 20000, dtype: list<item: int64>[pyarrow]
Behaviours¶
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.
[23]:
from akimbo import mixin_class, mixin_class_method, behavior
import akimbo.pandas
import numpy as np
import pandas as pd
@mixin_class(behavior)
class Point:
@mixin_class_method(np.abs)
def point_abs(self):
return np.sqrt(self.x ** 2 + self.y ** 2)
@mixin_class_method(np.add, {"Point"})
def point_add(self, other):
return ak.zip(
{"x": self.x + other.x, "y": self.y + other.y}, with_name="Point",
)
[24]:
data = [{"x": 1, "y": 2}] * 100000
s = pd.Series(data).ak.to_output() # store as arrow
[25]:
# check that the unary method is there; so tab-complete will work
"point_abs" in dir(s.ak.with_behavior("Point"))
[25]:
True
[26]:
# call to get vector sizes
s.ak.with_behavior("Point").point_abs()
[26]:
0 2.236068
1 2.236068
2 2.236068
3 2.236068
4 2.236068
...
99995 2.236068
99996 2.236068
99997 2.236068
99998 2.236068
99999 2.236068
Length: 100000, dtype: double[pyarrow]
[27]:
# or do the same with numpy ufunc
np.abs(s.ak.with_behavior("Point"))
[27]:
0 2.236068
1 2.236068
2 2.236068
3 2.236068
4 2.236068
...
99995 2.236068
99996 2.236068
99997 2.236068
99998 2.236068
99999 2.236068
Length: 100000, dtype: double[pyarrow]
[28]:
import math
%timeit np.abs(s.ak.with_behavior("Point"))
%timeit s.apply(lambda struct: math.sqrt(struct["x"] ** 2 + struct["y"] ** 2))
3.7 ms ± 218 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
62.1 ms ± 924 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
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.
Similarly, we defined an overload to add Point arrays together (both operands must be Point type). 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.
[29]:
%timeit s.ak.with_behavior("Point") + s.ak.with_behavior("Point")
s.ak.with_behavior("Point") + s.ak.with_behavior("Point")
2.59 ms ± 50.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
[29]:
0 {'x': 2, 'y': 4}
1 {'x': 2, 'y': 4}
2 {'x': 2, 'y': 4}
3 {'x': 2, 'y': 4}
4 {'x': 2, 'y': 4}
...
99995 {'x': 2, 'y': 4}
99996 {'x': 2, 'y': 4}
99997 {'x': 2, 'y': 4}
99998 {'x': 2, 'y': 4}
99999 {'x': 2, 'y': 4}
Length: 100000, dtype: struct<x: int64, y: int64>[pyarrow]
Numba integration¶
The numpy API is very nice and can do most things you will need. The object-oriented behaviours are very convenient.
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.
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.
[30]:
import numba
def mean_of_second_biggest(arr):
count = 0
total = 0
for row in arr:
if len(row) < 2:
continue
max = row[0]
if row[1] > max:
second = max
max = row[1]
else:
second = row[1]
for x in row[2:]:
if x > max:
second = max
max = x
elif x > second:
second = x
count += 1
total += second
return total / count
[31]:
mean_of_second_biggest([[1], [3, 2, 1], [3, 4, 4], [0, 1, 0]]) # mean of 2, 4, and 0
[31]:
2.0
[32]:
jsecond = numba.njit(mean_of_second_biggest)
[33]:
s = pd.Series([[1], [3, 2, 1], [3, 4, 4], [0, 1, 0]] * 100000).ak.to_output()
[34]:
s
[34]:
0 [1]
1 [3 2 1]
2 [3 4 4]
3 [0 1 0]
4 [1]
...
399995 [0 1 0]
399996 [1]
399997 [3 2 1]
399998 [3 4 4]
399999 [0 1 0]
Length: 400000, dtype: list<item: int64>[pyarrow]
[35]:
s.ak.apply(jsecond)
[35]:
2.0
[36]:
# timings
%timeit s.ak.apply(jsecond)
%timeit mean_of_second_biggest(s)
1.08 ms ± 37.7 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
501 ms ± 2.72 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Question: could you have done this with vectorized numpy calls?
[ ]: