feat: update to latest awkward==2.5.0

src/dask_awkward/lib/core.py

@@ -7,7 +7,7 @@
 import operator
 import sys
 import warnings
-from collections.abc import Callable, Hashable, Sequence
+from collections.abc import Callable, Hashable, Mapping, Sequence
 from enum import IntEnum
 from functools import cached_property, partial, wraps
 from numbers import Number
@@ -17,13 +17,14 @@
 import dask.config
 import numpy as np
 from awkward._do import remove_structure as ak_do_remove_structure
-from awkward._nplikes.typetracer import (
+from awkward.highlevel import NDArrayOperatorsMixin, _dir_pattern
+from awkward.typetracer import (
  MaybeNone,
  OneOf,
  TypeTracerArray,
+ create_unknown_scalar,
  is_unknown_scalar,
 )
-from awkward.highlevel import NDArrayOperatorsMixin, _dir_pattern
 from dask.base import (
  DaskMethodsMixin,
  dont_optimize,
@@ -140,7 +141,9 @@ def key(self) -> Key:
  return (self._name, 0)
 
  def _check_meta(self, m: Any) -> Any | None:
- if isinstance(m, (MaybeNone, OneOf)) or is_unknown_scalar(m):
+ if m is None:
+ return m
+ elif isinstance(m, (MaybeNone, OneOf)) or is_unknown_scalar(m):
  return m
  elif isinstance(m, ak.Array) and len(m) == 1:
  return m
@@ -348,12 +351,9 @@ def new_scalar_object(dsk: HighLevelGraph, name: str, *, meta: Any) -> Scalar:
  Resulting collection.
 
  """
- if meta is None:
- meta = ak.Array(TypeTracerArray._new(dtype=np.dtype(None), shape=()))
-
  if isinstance(meta, MaybeNone):
  meta = ak.Array(meta.content)
- else:
+ elif meta is not None:
  try:
  if ak.backend(meta) != "typetracer":
  raise TypeError(
@@ -411,9 +411,7 @@ def new_known_scalar(
  dtype = np.dtype(dtype)
  llg = AwkwardMaterializedLayer({(name, 0): s}, previous_layer_names=[])
  hlg = HighLevelGraph.from_collections(name, llg, dependencies=())
- return Scalar(
- hlg, name, meta=TypeTracerArray._new(dtype=dtype, shape=()), known_value=s
- )
+ return Scalar(hlg, name, meta=create_unknown_scalar(dtype), known_value=s)
 
 
 class Record(Scalar):
@@ -1486,7 +1484,8 @@ def new_array_object(
  name: str,
  *,
  meta: ak.Array | None = None,
- behavior: dict | None = None,
+ behavior: Mapping | None = None,
+ attrs: Mapping[str, Any] | None = None,
  npartitions: int | None = None,
  divisions: tuple[int, ...] | tuple[None, ...] | None = None,
 ) -> Array:
@@ -1505,6 +1504,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``.
@@ -1548,6 +1551,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__:
@@ -1877,6 +1882,8 @@ def non_trivial_reduction(
  keepdims: bool,
  mask_identity: bool,
  reducer: Callable,
+ behavior: Mapping | None = None,
+ attrs: Mapping[str, Any] | None = None,
  combiner: Callable | None = None,
  token: str | None = None,
  dtype: Any | None = None,

src/dask_awkward/lib/io/io.py

@@ -3,7 +3,7 @@
 import functools
 import logging
 import math
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Mapping
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, cast
 
@@ -65,7 +65,7 @@ def __call__(self, start: int, stop: int, **kwargs: Any) -> ak.Array:
 def from_awkward(
  source: ak.Array,
  npartitions: int,
- behavior: dict | None = None,
+ behavior: Mapping | None = None,
  label: str | None = None,
 ) -> Array:
  """Create an Array collection from a concrete :class:`awkward.Array` object.
@@ -116,21 +116,30 @@ def from_awkward(
 
 
 class _FromListsFn:
- def __init__(self, behavior: dict | None = None):
+ def __init__(self, behavior: Mapping | None, attrs: Mapping[str, Any] | 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) -> Array:
+def from_lists(
+ source: list,
+ behavior: Mapping | None = None,
+ attrs: Mapping[str, Any] | None = None,
+) -> Array:
  """Create an Array collection from a list of lists.
 
  Parameters
  ----------
  source : list[list[Any]]
  List of lists, each outer list will become a partition in the
  collection.
+ behavior : dict, optional
+ Custom ak.behavior for the output array.
+ attrs : mapping, optional
+ Custom attributes for the output array.
 
  Returns
  -------
@@ -152,9 +161,9 @@ def from_lists(source: list) -> Array:
  lists = list(source)
  divs = (0, *np.cumsum(list(map(len, lists))))
  return from_map(
- _FromListsFn(),
+ _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,
  label="from-lists",
  )
@@ -163,9 +172,10 @@ def from_lists(source: list) -> Array:
 def from_delayed(
  source: list[Delayed] | Delayed,
  meta: ak.Array | None = None,
- behavior: dict | None = None,
+ behavior: Mapping | None = None,
  divisions: tuple[int, ...] | tuple[None, ...] | None = None,
  prefix: str = "from-delayed",
+ attrs: Mapping[str, Any] | None = None,
 ) -> Array:
  """Create an Array collection from a set of :class:`~dask.delayed.Delayed` objects.
 
@@ -179,10 +189,14 @@ def from_delayed(
  Metadata (typetracer array) if known, if ``None`` the first
  partition (first element of the list of ``Delayed`` objects)
  will be computed to determine the metadata.
+ behavior : dict, optional
+ Custom ak.behavior for the output array.
  divisions : tuple[int | None, ...], optional
  Partition boundaries (if known).
  prefix : str
  Prefix for the keys in the task graph.
+ attrs : mapping, optional
+ Custom attributes for the output array.
 
  Returns
  -------
@@ -206,11 +220,7 @@ def from_delayed(
  raise ValueError("divisions must be a tuple of length len(source) + 1")
  hlg = HighLevelGraph.from_collections(name, dsk, dependencies=parts)
  return new_array_object(
- hlg,
- name=name,
- meta=meta,
- behavior=behavior,
- divisions=divs,
+ hlg, name=name, meta=meta, behavior=behavior, divisions=divs, attrs=attrs
  )
 
 
@@ -346,13 +356,21 @@ def to_dask_array(
  return new_da_object(graph, name, meta=None, chunks=chunks, dtype=dtype)
 
 
-def from_dask_array(array: DaskArray, behavior: dict | None = None) -> Array:
+def from_dask_array(
+ array: DaskArray,
+ behavior: Mapping | None = None,
+ attrs: Mapping[str, Any] | None = None,
+) -> Array:
  """Convert a Dask Array collection to a Dask Awkard Array collection.
 
  Parameters
  ----------
  array : dask.array.Array
  Array to convert.
+ behavior : dict, optional
+ Custom ak.behavior for the output array.
+ attrs : mapping, optional
+ Custom attributes for the output array.
 
  Returns
  -------
@@ -389,11 +407,18 @@ def from_dask_array(array: DaskArray, behavior: dict | None = None) -> Array:
  hlg = HighLevelGraph.from_collections(name, layer, dependencies=[array])
  if np.any(np.isnan(array.chunks)):
  return new_array_object(
- hlg, name, npartitions=array.npartitions, meta=meta, behavior=behavior
+ hlg,
+ name,
+ npartitions=array.npartitions,
+ meta=meta,
+ behavior=behavior,
+ attrs=attrs,
  )
  else:
  divs = (0, *np.cumsum(array.chunks))
- return new_array_object(hlg, name, divisions=divs, meta=meta, behavior=behavior)
+ return new_array_object(
+ hlg, name, divisions=divs, meta=meta, behavior=behavior, attrs=attrs
+ )
 
 
 def to_dataframe(
@@ -522,6 +547,8 @@ def from_map(
  number of partitions in the output collection (only one
  element of each iterable will be passed to `func` for each
  partition).
+ args : tuple
+ Tuple of positional arguments to append after mapped arguments.
  label : str, optional
  String to use as the function-name label in the output
  collection-key names.

src/dask_awkward/lib/io/json.py

@@ -3,7 +3,7 @@
 import abc
 import logging
 import math
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from typing import TYPE_CHECKING, Any, Literal, overload
 
 import awkward as ak
@@ -43,7 +43,7 @@ def __init__(
  form: Form,
  compression: str | None = None,
  schema: str | dict | list | None = None,
- behavior: dict | None = None,
+ behavior: Mapping | None = None,
  **kwargs: Any,
  ) -> None:
  self.compression = compression
@@ -91,7 +91,7 @@ def __init__(
  form: Form,
  compression: str | None = None,
  schema: str | dict | list | None = None,
- behavior: dict | None = None,
+ behavior: Mapping | None = None,
  **kwargs: Any,
  ) -> None:
  super().__init__(
@@ -125,7 +125,7 @@ def __init__(
  form: Form,
  compression: str | None = None,
  schema: str | dict | list | None = None,
- behavior: dict | None = None,
+ behavior: Mapping | None = None,
  **kwargs: Any,
  ) -> None:
  super().__init__(
@@ -163,7 +163,7 @@ def __init__(
  form: Form,
  compression: str | None = None,
  schema: str | dict | list | None = None,
- behavior: dict | None = None,
+ behavior: Mapping | None = None,
  **kwargs: Any,
  ) -> None:
  super().__init__(
@@ -432,7 +432,8 @@ def from_json(
  initial: int = 1024,
  resize: float = 8,
  highlevel: bool = True,
- behavior: dict | None = None,
+ behavior: Mapping | None = None,
+ attrs: Mapping[str, Any] | None = None,
  blocksize: int | str | None = None,
  delimiter: bytes | None = None,
  compression: str | None = "infer",