diff --git a/docs/api.rst b/docs/api.rst index c961fbe..e33ff25 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -37,6 +37,7 @@ Backends akimbo.pandas.PandasAwkwardAccessor akimbo.dask.DaskAwkwardAccessor akimbo.polars.PolarsAwkwardAccessor + akimbo.cudf.CudfAwkwardAccessor .. autoclass:: akimbo.pandas.PandasAwkwardAccessor @@ -44,6 +45,8 @@ Backends .. autoclass:: akimbo.polars.PolarsAwkwardAccessor +.. autoclass:: akimbo.polars.CudfAwkwardAccessor + Extensions ~~~~~~~~~~ @@ -63,3 +66,12 @@ to see the operations available. + +The cuDF backend also has these implemented with GPU-specific variants, +``akimbo.cudf.CudfStringAccessor`` and ``akimbo.cudf.CudfDatetimeAccessor``. + +Adding Extensions (advanced) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The patterns used in the two builtin extensions above can be used to add +type-specific functionality to ``akimbo``. One diff --git a/src/akimbo/apply_tree.py b/src/akimbo/apply_tree.py index 1b946a7..6db39cd 100644 --- a/src/akimbo/apply_tree.py +++ b/src/akimbo/apply_tree.py @@ -1,16 +1,15 @@ +from __future__ import annotations + import functools import inspect -from typing import Sequence +from typing import Callable, Literal, Sequence import awkward as ak import pyarrow as pa -class NoDtype: - kind = "" - - def leaf(*layout, **_): + """True for the lowest elements of any akwward layout tree""" return layout[0].is_leaf @@ -42,8 +41,20 @@ def func(layout, **kwargs): return ak.transform(func, arr, *others) -def dec(func, match=leaf, outtype=None, inmode="arrow"): - """Make a nested/ragged version of an operation to apply throughout a tree""" +def dec(func: callable, match: Callable[[ak.contents.Content], bool] = leaf, + outtype: Callable[[ak.contents.Content], ak.contents.Content] | None = None, + inmode: Literal["arrow", "numpy", "ak"] = "arrow"): + + """Make a nested/ragged version of an operation to apply throughout a tree + + Parameters + ---------- + func: which we want to apply to (parts of) inputted data + match: function to determine if a part of the data structure matches the type we want to + operate on + outtype: postprocessing function after transform + inmode: how ``func`` expects its inputs: as awkward arrays (ak), numpy or arrow + """ @functools.wraps(func) def f(self, *args, where=None, match_kwargs=None, **kwargs):