diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py index d540bb6c..9dae18a0 100644 --- a/src/dask_awkward/lib/core.py +++ b/src/dask_awkward/lib/core.py @@ -11,6 +11,7 @@ from functools import cached_property, partial, wraps from inspect import getattr_static from numbers import Number +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Literal, TypeVar, Union, overload import awkward as ak @@ -1090,17 +1091,18 @@ def layout(self) -> Content: raise ValueError("This collection's meta is None; unknown layout.") @property - def attrs(self) -> dict: + def attrs(self) -> Mapping[str, Any]: """awkward Array attrs dictionary.""" if self._meta is not None: - return self._meta.attrs + return MappingProxyType(self._meta.attrs) raise ValueError("This collection's meta is None; no attrs property available.") @property - def behavior(self) -> Mapping: + def behavior(self) -> Mapping | None: """awkward Array behavior dictionary.""" if self._meta is not None: - return self._meta.behavior + behavior = self._meta.behavior + return None if behavior is None else MappingProxyType(behavior) raise ValueError( "This collection's meta is None; no behavior property available." ) diff --git a/src/dask_awkward/lib/io/columnar.py b/src/dask_awkward/lib/io/columnar.py index 0fd95592..47fde251 100644 --- a/src/dask_awkward/lib/io/columnar.py +++ b/src/dask_awkward/lib/io/columnar.py @@ -41,6 +41,10 @@ def form(self) -> Form: def behavior(self) -> dict | None: ... + @property + def attrs(self) -> dict | None: + ... + def project_columns(self: T, columns: frozenset[str]) -> T: ... @@ -62,7 +66,8 @@ class ColumnProjectionMixin(ImplementsNecessaryColumns[FormStructure]): def mock(self: S) -> AwkwardArray: return cast( - AwkwardArray, typetracer_from_form(self.form, behavior=self.behavior) + AwkwardArray, + typetracer_from_form(self.form, behavior=self.behavior, attrs=self.attrs), ) def mock_empty(self: S, backend: BackendT = "cpu") -> AwkwardArray: diff --git a/src/dask_awkward/lib/io/io.py b/src/dask_awkward/lib/io/io.py index e32bbfab..b3108081 100644 --- a/src/dask_awkward/lib/io/io.py +++ b/src/dask_awkward/lib/io/io.py @@ -55,8 +55,8 @@ def __init__( ) -> None: self.arr = arr self.form = arr.layout.form - self.behavior = behavior - self.attrs = attrs + self.behavior = behavior if behavior else arr.behavior + self.attrs = attrs if attrs else arr.attrs @property def use_optimization(self): @@ -119,7 +119,7 @@ def from_awkward( return cast( Array, from_map( - FromAwkwardFn(source, behavior=behavior), + FromAwkwardFn(source, behavior=behavior, attrs=attrs), starts_stops, label=label or "from-awkward", token=tokenize(source, npartitions), diff --git a/src/dask_awkward/lib/io/json.py b/src/dask_awkward/lib/io/json.py index b7cf6a99..d6695a60 100644 --- a/src/dask_awkward/lib/io/json.py +++ b/src/dask_awkward/lib/io/json.py @@ -49,6 +49,7 @@ def __init__( compression: str | None = None, schema: str | dict | list | None = None, behavior: Mapping | None = None, + attrs: Mapping[str, Any] | None = None, **kwargs: Any, ) -> None: self.compression = compression @@ -57,6 +58,7 @@ def __init__( self.kwargs = kwargs self.form = form self.behavior = behavior + self.attrs = attrs @abc.abstractmethod def __call__(self, source: Any) -> ak.Array: @@ -97,6 +99,7 @@ def __init__( compression: str | None = None, schema: str | dict | list | None = None, behavior: Mapping | None = None, + attrs: Mapping[str, Any] | None = None, **kwargs: Any, ) -> None: super().__init__( @@ -105,6 +108,7 @@ def __init__( schema=schema, form=form, behavior=behavior, + attrs=attrs, **kwargs, ) @@ -131,6 +135,7 @@ def __init__( compression: str | None = None, schema: str | dict | list | None = None, behavior: Mapping | None = None, + attrs: Mapping[str, Any] | None = None, **kwargs: Any, ) -> None: super().__init__( @@ -139,6 +144,7 @@ def __init__( schema=schema, form=form, behavior=behavior, + attrs=attrs, **kwargs, ) @@ -169,6 +175,7 @@ def __init__( compression: str | None = None, schema: str | dict | list | None = None, behavior: Mapping | None = None, + attrs: Mapping[str, Any] | None = None, **kwargs: Any, ) -> None: super().__init__( @@ -176,6 +183,7 @@ def __init__( compression=compression, schema=schema, behavior=behavior, + attrs=attrs, form=form, **kwargs, ) diff --git a/tests/test_core.py b/tests/test_core.py index 5d9c9a4e..26ae1944 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -891,6 +891,28 @@ def test_shape_only_ops(fn: Callable, tmp_path_factory: pytest.TempPathFactory) result.compute() +def test_assign_behavior() -> None: + behavior = {"test": "hello"} + x = ak.Array([{"a": 1, "b": 2}, {"a": 3, "b": 4}], behavior=behavior, attrs={}) + dx = dak.from_awkward(x, 3) + with pytest.raises( + TypeError, match="'mappingproxy' object does not support item assignment" + ): + dx.behavior["should_fail"] = None # type: ignore + assert dx.behavior == behavior + + +def test_assign_attrs() -> None: + attrs = {"test": "hello"} + x = ak.Array([{"a": 1, "b": 2}, {"a": 3, "b": 4}], behavior={}, attrs=attrs) + dx = dak.from_awkward(x, 3) + with pytest.raises( + TypeError, match="'mappingproxy' object does not support item assignment" + ): + dx.attrs["should_fail"] = None # type: ignore + assert dx.attrs == attrs + + @delayed def a_delayed_array(): return ak.Array([2, 4])