From 9ebaa3786a38930942d08eec670a66003572a905 Mon Sep 17 00:00:00 2001 From: Angus Hollands Date: Wed, 8 Nov 2023 18:23:49 +0000 Subject: [PATCH] feat: add attrs and highlevel --- src/dask_awkward/lib/core.py | 9 ++ src/dask_awkward/lib/io/io.py | 14 +- src/dask_awkward/lib/io/json.py | 1 + src/dask_awkward/lib/io/parquet.py | 37 ++++- src/dask_awkward/lib/operations.py | 5 +- src/dask_awkward/lib/reducers.py | 226 ++++++++++++++++++++++------- src/dask_awkward/lib/str.py | 50 +++++++ src/dask_awkward/lib/structure.py | 15 +- 8 files changed, 285 insertions(+), 72 deletions(-) diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py index 86fd8f021..843095e3f 100644 --- a/src/dask_awkward/lib/core.py +++ b/src/dask_awkward/lib/core.py @@ -1489,6 +1489,7 @@ def new_array_object( *, meta: ak.Array | None = None, behavior: dict | None = None, + attrs: dict | None = None, npartitions: int | None = None, divisions: tuple[int, ...] | tuple[None, ...] | None = None, ) -> Array: @@ -1507,6 +1508,10 @@ def new_array_object( typetracer for the new Array. If the configuration option ``awkward.compute-unknown-meta`` is set to ``False``, undefined `meta` will be assigned an empty typetracer. + behavior : dict, optional + Custom #ak.behavior for the output array. + attrs : dict, optional + Custom attributes for the output array. npartitions : int, optional Total number of partitions; if used `divisions` will be a tuple of length `npartitions` + 1 with all elements``None``. @@ -1550,6 +1555,8 @@ def new_array_object( if behavior is not None: actual_meta.behavior = behavior + if attrs is not None: + actual_meta.attrs = attrs out = Array(dsk, name, actual_meta, divs) if actual_meta.__doc__ != actual_meta.__class__.__doc__: @@ -1879,6 +1886,8 @@ def non_trivial_reduction( keepdims: bool, mask_identity: bool, reducer: Callable, + behavior: dict | None = None, + attrs: dict | None = None, combiner: Callable | None = None, token: str | None = None, dtype: Any | None = None, diff --git a/src/dask_awkward/lib/io/io.py b/src/dask_awkward/lib/io/io.py index 6153021b6..af23b6627 100644 --- a/src/dask_awkward/lib/io/io.py +++ b/src/dask_awkward/lib/io/io.py @@ -104,14 +104,17 @@ def from_awkward( class _FromListsFn: - def __init__(self, behavior: dict | None = None): + def __init__(self, behavior: dict | None = None, attrs: dict | None = None): self.behavior = behavior + self.attrs = attrs def __call__(self, x: list) -> ak.Array: - return ak.Array(x, behavior=self.behavior) + return ak.Array(x, behavior=self.behavior, attrs=self.attrs) -def from_lists(source: list, behavior: dict | None = None) -> Array: +def from_lists( + source: list, behavior: dict | None = None, attrs: dict | None = None +) -> Array: """Create an Array collection from a list of lists. Parameters @@ -140,11 +143,10 @@ def from_lists(source: list, behavior: dict | None = None) -> Array: lists = list(source) divs = (0, *np.cumsum(list(map(len, lists)))) return from_map( - _FromListsFn(behavior=behavior), + _FromListsFn(behavior=behavior, attrs=attrs), lists, - meta=typetracer_array(ak.Array(lists[0])), + meta=typetracer_array(ak.Array(lists[0], attrs=attrs, behavior=behavior)), divisions=divs, - behavior=behavior, label="from-lists", ) diff --git a/src/dask_awkward/lib/io/json.py b/src/dask_awkward/lib/io/json.py index 43ba19e37..5c06cd1a4 100644 --- a/src/dask_awkward/lib/io/json.py +++ b/src/dask_awkward/lib/io/json.py @@ -433,6 +433,7 @@ def from_json( resize: float = 8, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, blocksize: int | str | None = None, delimiter: bytes | None = None, compression: str | None = "infer", diff --git a/src/dask_awkward/lib/io/parquet.py b/src/dask_awkward/lib/io/parquet.py index 8a3c04175..265a20a45 100644 --- a/src/dask_awkward/lib/io/parquet.py +++ b/src/dask_awkward/lib/io/parquet.py @@ -42,6 +42,7 @@ def __init__( unnamed_root: bool = False, original_form: Form | None = None, behavior: dict | None = None, + attrs: dict | None = None, **kwargs: Any, ) -> None: self.fs = fs @@ -53,6 +54,7 @@ def __init__( self.columns = [f".{c}" for c in self.columns] self.original_form = original_form self.behavior = behavior + self.attrs = attrs self.kwargs = kwargs @abc.abstractmethod @@ -111,17 +113,22 @@ def __init__( ) def __call__(self, source: Any) -> Any: - array = ak_from_parquet._load( + layout = ak_from_parquet._load( [source], parquet_columns=self.columns, subrg=[None], subform=self.form, - highlevel=True, + highlevel=False, + attrs=None, + behavior=None, fs=self.fs, - behavior=self.behavior, **self.kwargs, ) - return ak.Array(unproject_layout(self.original_form, array.layout)) + return ak.Array( + unproject_layout(self.original_form, layout), + attrs=self.attrs, + behavior=self.behavior, + ) def project_columns(self, columns): return _FromParquetFileWiseFn( @@ -130,6 +137,7 @@ def project_columns(self, columns): listsep=self.listsep, unnamed_root=self.unnamed_root, original_form=self.form, + attrs=self.attrs, behavior=self.behavior, **self.kwargs, ) @@ -145,6 +153,7 @@ def __init__( unnamed_root: bool = False, original_form: Form | None = None, behavior: dict | None = None, + attrs: dict | None = None, **kwargs: Any, ) -> None: super().__init__( @@ -154,6 +163,7 @@ def __init__( unnamed_root=unnamed_root, original_form=original_form, behavior=behavior, + attrs=attrs, **kwargs, ) @@ -161,17 +171,22 @@ def __call__(self, pair: Any) -> ak.Array: subrg, source = pair if isinstance(subrg, int): subrg = [[subrg]] - array = ak_from_parquet._load( + layout = ak_from_parquet._load( [source], parquet_columns=self.columns, subrg=subrg, subform=self.form, - highlevel=True, + highlevel=False, + attrs=None, + behavior=None, fs=self.fs, - behavior=self.behavior, **self.kwargs, ) - return ak.Array(unproject_layout(self.original_form, array.layout)) + return ak.Array( + unproject_layout(self.original_form, layout), + behavior=self.behavior, + attrs=self.attrs, + ) def project_columns(self, columns): return _FromParquetFragmentWiseFn( @@ -180,6 +195,7 @@ def project_columns(self, columns): unnamed_root=self.unnamed_root, original_form=self.form, behavior=self.behavior, + attrs=self.attrs, **self.kwargs, ) @@ -194,6 +210,7 @@ def from_parquet( generate_bitmasks: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ignore_metadata: bool = True, scan_files: bool = False, split_row_groups: bool | None = False, @@ -263,6 +280,8 @@ def from_parquet( ignore_metadata, scan_files, split_row_groups, + behavior, + attrs, ) ( @@ -307,6 +326,7 @@ def from_parquet( footer_sample_size=footer_sample_size, generate_bitmasks=generate_bitmasks, behavior=behavior, + attrs=attrs, ), actual_paths, label=label, @@ -345,6 +365,7 @@ def from_parquet( footer_sample_size=footer_sample_size, generate_bitmasks=generate_bitmasks, behavior=behavior, + attrs=attrs, ), pairs, label=label, diff --git a/src/dask_awkward/lib/operations.py b/src/dask_awkward/lib/operations.py index 32e944793..929ead251 100644 --- a/src/dask_awkward/lib/operations.py +++ b/src/dask_awkward/lib/operations.py @@ -33,6 +33,7 @@ def concatenate( mergebool: bool = True, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: label = "concatenate" token = tokenize(arrays, axis, mergebool, highlevel, behavior) @@ -49,7 +50,7 @@ def concatenate( g[(name, i)] = k i += 1 - meta = ak.concatenate(metas) + meta = ak.concatenate(metas, behavior=behavior, attrs=attrs) assert isinstance(meta, ak.Array) prev_names = [iarr.name for iarr in arrays] @@ -65,7 +66,7 @@ def concatenate( if partition_compatibility(*arrays) == PartitionCompatibility.NO: raise IncompatiblePartitions("concatenate", *arrays) - fn = _ConcatenateFnAxisGT0(axis=axis) + fn = _ConcatenateFnAxisGT0(axis=axis, behavior=behavior, attrs=attrs) return map_partitions(fn, *arrays) else: diff --git a/src/dask_awkward/lib/reducers.py b/src/dask_awkward/lib/reducers.py index 716327cc1..ccf75d944 100644 --- a/src/dask_awkward/lib/reducers.py +++ b/src/dask_awkward/lib/reducers.py @@ -39,6 +39,8 @@ def all( axis: int | None = None, keepdims: bool = False, mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if axis is None or axis == 0 or axis == -1 * array.ndim: return non_trivial_reduction( @@ -49,6 +51,8 @@ def all( is_positional=False, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) else: return map_partitions( @@ -58,6 +62,8 @@ def all( axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) @@ -67,6 +73,8 @@ def any( axis: int | None = None, keepdims: bool = False, mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if axis is None or axis == 0 or axis == -1 * array.ndim: return non_trivial_reduction( @@ -77,6 +85,8 @@ def any( is_positional=False, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) else: return map_partitions( @@ -86,6 +96,8 @@ def any( axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) @@ -95,6 +107,8 @@ def argmax( axis: int | None = None, keepdims: bool = False, mask_identity: bool = True, + behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if axis is None or axis == 0 or axis == -1 * array.ndim: return non_trivial_reduction( @@ -105,6 +119,8 @@ def argmax( is_positional=True, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) else: return map_partitions( @@ -114,6 +130,8 @@ def argmax( axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) @@ -123,6 +141,8 @@ def argmin( axis: int | None = None, keepdims: bool = False, mask_identity: bool = True, + behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if axis is None or axis == 0 or axis == -1 * array.ndim: return non_trivial_reduction( @@ -133,6 +153,8 @@ def argmin( is_positional=True, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) else: return map_partitions( @@ -142,18 +164,22 @@ def argmin( axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) @borrow_docstring(ak.corr) def corr( - x, - y, - weight=None, - axis=None, - keepdims=False, - mask_identity=False, -): + x: Array, + y: Array, + weight: Array | int | float | complex | None = None, + axis: int | None = None, + keepdims: bool = False, + mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, +) -> Any: raise DaskAwkwardNotImplemented("TODO") @@ -163,6 +189,8 @@ def count( axis: int | None = None, keepdims: bool = False, mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if axis is None or axis == 0 or axis == -1 * array.ndim: return non_trivial_reduction( @@ -174,6 +202,8 @@ def count( is_positional=False, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) else: return map_partitions( @@ -183,6 +213,8 @@ def count( axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) @@ -192,6 +224,8 @@ def count_nonzero( axis: int | None = None, keepdims: bool = False, mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if axis is None or axis == 0 or axis == -1 * array.ndim: return non_trivial_reduction( @@ -203,6 +237,8 @@ def count_nonzero( is_positional=False, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) else: return map_partitions( @@ -212,33 +248,47 @@ def count_nonzero( axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) @borrow_docstring(ak.covar) def covar( - x, - y, - weight=None, - axis=None, - keepdims=False, - mask_identity=False, -): + x: Array, + y: Array, + weight: Array | int | float | complex | None = None, + axis: int | None = None, + keepdims: bool = False, + mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, +) -> Any: raise DaskAwkwardNotImplemented("TODO") @borrow_docstring(ak.linear_fit) def linear_fit( - x, - y, - weight=None, - axis=None, - keepdims=False, - mask_identity=False, -): + x: Array, + y: Array, + weight: Array | int | float | complex | None = None, + axis: int | None = None, + keepdims: bool = False, + mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, +) -> Any: raise DaskAwkwardNotImplemented("TODO") +class _MaxFn: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def __call__(self, array, **kwargs): + return ak.max(array, **self.kwargs, **kwargs) + + @borrow_docstring(ak.max) def max( array: Array, @@ -246,36 +296,49 @@ def max( keepdims: bool = False, initial: float | None = None, mask_identity: bool = True, + behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if axis is None or axis == 0 or axis == -1 * array.ndim: return non_trivial_reduction( axis=axis, label="max", array=array, - reducer=ak.max, + reducer=_MaxFn(initial=initial), is_positional=False, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) else: return map_partitions( - ak.max, + _MaxFn(initial=initial), array, output_divisions=1, axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) @borrow_docstring(ak.mean) def mean( - array, - weight=None, - axis=None, - keepdims=False, - mask_identity=False, -): + array: Array, + weight: Array | int | float | complex | None = None, + axis: int | None = None, + keepdims: bool = False, + mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, +) -> Array: + if weight is not None: + raise DaskAwkwardNotImplemented( + f"weight={weight} is not supported for this array yet." + ) + if axis == 0 or axis == -1 * array.ndim: raise DaskAwkwardNotImplemented( f"axis={axis} is not supported for this array yet." @@ -288,10 +351,20 @@ def mean( axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) raise DaskAwkwardNotImplemented("TODO") +class _MinFn: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def __call__(self, array, **kwargs): + return ak.min(array, **self.kwargs, **kwargs) + + @borrow_docstring(ak.min) def min( array: Array, @@ -299,42 +372,57 @@ def min( keepdims: bool = False, initial: float | None = None, mask_identity: bool = True, + behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if axis is None or axis == 0 or axis == -1 * array.ndim: return non_trivial_reduction( axis=axis, label="min", array=array, - reducer=ak.min, + reducer=_MinFn(initial=initial), is_positional=False, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) else: return map_partitions( - ak.min, + _MinFn(initial=initial), array, output_divisions=1, axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) @borrow_docstring(ak.moment) def moment( - x, - n, - weight=None, - axis=None, - keepdims=False, - mask_identity=False, -): + x: Array, + n: int, + weight: Array | int | float | complex | None = None, + axis: int | None = None, + keepdims: bool = False, + mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, +) -> Any: raise DaskAwkwardNotImplemented("TODO") @borrow_docstring(ak.prod) -def prod(array, axis=None, keepdims=False, mask_identity=False): +def prod( + array: Array, + axis: int | None = None, + keepdims: bool = False, + mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, +) -> Any: if axis is None or axis == 0 or axis == -1 * array.ndim: return non_trivial_reduction( axis=axis, @@ -344,6 +432,8 @@ def prod(array, axis=None, keepdims=False, mask_identity=False): is_positional=False, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) else: return map_partitions( @@ -353,16 +443,28 @@ def prod(array, axis=None, keepdims=False, mask_identity=False): axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) @borrow_docstring(ak.ptp) -def ptp(arr, axis=None, keepdims=False, mask_identity=True): +def ptp( + arr: Array, + axis: int | None = None, + keepdims: bool = False, + mask_identity: bool = True, +) -> Any: raise DaskAwkwardNotImplemented("TODO") @borrow_docstring(ak.softmax) -def softmax(x, axis=None, keepdims=False, mask_identity=False): +def softmax( + x: Array, + axis: int | None = None, + keepdims: bool = False, + mask_identity: bool = False, +) -> Any: raise DaskAwkwardNotImplemented("TODO") @@ -376,13 +478,15 @@ def __call__(self, array): @borrow_docstring(ak.std) def std( - x, - weight=None, - ddof=0, - axis=None, - keepdims=False, - mask_identity=False, -): + x: Array, + weight: Array | int | float | complex | None = None, + ddof: int = 0, + axis: int | None = None, + keepdims: bool = False, + mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, +) -> Any: if weight is not None: raise DaskAwkwardNotImplemented("weight argument is not supported.") if axis is None or axis == 0 or axis == -1 * x.ndim: @@ -396,6 +500,8 @@ def std( axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ), x, output_divisions=1, @@ -409,6 +515,8 @@ def sum( axis: int | None = None, keepdims: bool = False, mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if axis is None or axis == 0 or axis == -1 * array.ndim: return non_trivial_reduction( @@ -419,6 +527,8 @@ def sum( is_positional=False, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) else: return map_partitions( @@ -428,6 +538,8 @@ def sum( axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ) @@ -441,13 +553,15 @@ def __call__(self, array): @borrow_docstring(ak.var) def var( - x, - weight=None, - ddof=0, - axis=None, - keepdims=False, - mask_identity=False, -): + x: Array, + weight: Array | int | float | complex | None = None, + ddof: int = 0, + axis: int | None = None, + keepdims: bool = False, + mask_identity: bool = False, + behavior: dict | None = None, + attrs: dict | None = None, +) -> Any: if weight is not None: raise DaskAwkwardNotImplemented("weight argument is not supported.") if axis is None or axis == 0 or axis == -1 * x.ndim: @@ -461,6 +575,8 @@ def var( axis=axis, keepdims=keepdims, mask_identity=mask_identity, + behavior=behavior, + attrs=attrs, ), x, output_divisions=1, diff --git a/src/dask_awkward/lib/str.py b/src/dask_awkward/lib/str.py index 85324ae92..8fbfb203c 100644 --- a/src/dask_awkward/lib/str.py +++ b/src/dask_awkward/lib/str.py @@ -29,6 +29,7 @@ def capitalize( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.capitalize, @@ -46,6 +47,7 @@ def center( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.center, @@ -65,6 +67,7 @@ def count_substring( ignore_case: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.count_substring, @@ -84,6 +87,7 @@ def count_substring_regex( ignore_case: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.count_substring_regex, @@ -103,6 +107,7 @@ def ends_with( ignore_case: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.ends_with, @@ -121,6 +126,7 @@ def extract_regex( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.extract_regex, @@ -139,6 +145,7 @@ def find_substring( ignore_case: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.find_substring, @@ -158,6 +165,7 @@ def find_substring_regex( ignore_case: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.find_substring_regex, @@ -177,6 +185,7 @@ def index_in( skip_nones: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.index_in, @@ -194,6 +203,7 @@ def is_alnum( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_alnum, @@ -209,6 +219,7 @@ def is_alpha( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_alpha, @@ -224,6 +235,7 @@ def is_ascii( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_ascii, @@ -239,6 +251,7 @@ def is_decimal( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_decimal, @@ -254,6 +267,7 @@ def is_digit( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_digit, @@ -271,6 +285,7 @@ def is_in( skip_nones: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_in, @@ -288,6 +303,7 @@ def is_lower( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_lower, @@ -303,6 +319,7 @@ def is_numeric( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_numeric, @@ -318,6 +335,7 @@ def is_printable( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_printable, @@ -333,6 +351,7 @@ def is_space( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_space, @@ -348,6 +367,7 @@ def is_title( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_title, @@ -363,6 +383,7 @@ def is_upper( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.is_upper, @@ -379,6 +400,7 @@ def join( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.join, @@ -394,6 +416,7 @@ def join_element_wise( *arrays: Array, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.join_element_wise, @@ -409,6 +432,7 @@ def length( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.length, @@ -424,6 +448,7 @@ def lower( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.lower, @@ -441,6 +466,7 @@ def lpad( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.lpad, @@ -459,6 +485,7 @@ def ltrim( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.ltrim, @@ -475,6 +502,7 @@ def ltrim_whitespace( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.ltrim_whitespace, @@ -492,6 +520,7 @@ def match_like( ignore_case: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.match_like, @@ -511,6 +540,7 @@ def match_substring( ignore_case: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.match_substring, @@ -530,6 +560,7 @@ def match_substring_regex( ignore_case: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.match_substring_regex, @@ -548,6 +579,7 @@ def repeat( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.repeat, @@ -567,6 +599,7 @@ def replace_slice( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.replace_slice, @@ -588,6 +621,7 @@ def replace_substring( max_replacements: int | None = None, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.replace_substring, @@ -609,6 +643,7 @@ def replace_substring_regex( max_replacements: int | None = None, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.replace_substring_regex, @@ -627,6 +662,7 @@ def reverse( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.reverse, @@ -644,6 +680,7 @@ def rpad( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.rpad, @@ -662,6 +699,7 @@ def rtrim( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.rtrim, @@ -678,6 +716,7 @@ def rtrim_whitespace( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.rtrim_whitespace, @@ -696,6 +735,7 @@ def slice( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.slice, @@ -717,6 +757,7 @@ def split_pattern( reverse: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.split_pattern, @@ -738,6 +779,7 @@ def split_pattern_regex( reverse: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.split_pattern_regex, @@ -758,6 +800,7 @@ def split_whitespace( reverse: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.split_whitespace, @@ -776,6 +819,7 @@ def starts_with( ignore_case: bool = False, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.starts_with, @@ -793,6 +837,7 @@ def swapcase( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.swapcase, @@ -808,6 +853,7 @@ def title( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.title, @@ -823,6 +869,7 @@ def to_categorical( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.to_categorical, @@ -839,6 +886,7 @@ def trim( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.trim, @@ -855,6 +903,7 @@ def trim_whitespace( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.trim_whitespace, @@ -870,6 +919,7 @@ def upper( *, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( akstr.upper, diff --git a/src/dask_awkward/lib/structure.py b/src/dask_awkward/lib/structure.py index e781985a4..83a7aa362 100644 --- a/src/dask_awkward/lib/structure.py +++ b/src/dask_awkward/lib/structure.py @@ -291,6 +291,7 @@ def combinations( with_name: str | None = None, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: if not highlevel: raise ValueError("Only highlevel=True is supported") @@ -341,6 +342,7 @@ def fill_none( axis: int = -1, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: if not highlevel: raise ValueError("Only highlevel=True is supported") @@ -363,6 +365,7 @@ def drop_none( axis: int | None = None, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: if not highlevel: raise ValueError("Only highlevel=True is supported") @@ -385,6 +388,7 @@ def firsts( axis: int = 1, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if axis == 1: return map_partitions( @@ -416,6 +420,7 @@ def flatten( axis: int | None = 1, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: if not highlevel: raise ValueError("Only highlevel=True is supported") @@ -574,6 +579,7 @@ def num( axis: int = 1, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Any: if not highlevel: raise ValueError("Only highlevel=True is supported") @@ -596,7 +602,7 @@ def num( return new_scalar_object( hlg, name, - meta=TypeTracerArray._new(dtype=np.int64, shape=()), + meta=TypeTracerArray._new(dtype=np.dtype(np.int64), shape=()), ) else: return map_partitions( @@ -614,6 +620,7 @@ def ones_like( array: Array, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, dtype: DTypeLike | None = None, ) -> Array: if not highlevel: @@ -865,6 +872,7 @@ def where( mergebool: bool = True, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: if not highlevel: raise ValueError("Only highlevel=True is supported") @@ -952,6 +960,7 @@ def with_name( name: str, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: if not highlevel: raise ValueError("Only highlevel=True is supported") @@ -986,6 +995,7 @@ def with_parameter( value: Any, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( _WithParameterFn(parameter=parameter, value=value, behavior=behavior), @@ -1008,6 +1018,7 @@ def without_parameters( array: Array, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, ) -> Array: return map_partitions( _WithoutParameterFn(behavior=behavior), @@ -1022,6 +1033,7 @@ def zeros_like( array: Array, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, dtype: DTypeLike | None = None, ) -> Array: if not highlevel: @@ -1064,6 +1076,7 @@ def zip( with_name: str | None = None, highlevel: bool = True, behavior: dict | None = None, + attrs: dict | None = None, right_broadcast: bool = False, optiontype_outside_record: bool = False, ) -> Array: