Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: return proxy of attrs, behavior; dak.from_awkward correctly sets attrs #448

Merged
merged 9 commits into from
Jan 29, 2024
10 changes: 6 additions & 4 deletions src/dask_awkward/lib/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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."
)
Expand Down
7 changes: 6 additions & 1 deletion src/dask_awkward/lib/io/columnar.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
...

Expand All @@ -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:
Expand Down
6 changes: 3 additions & 3 deletions src/dask_awkward/lib/io/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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),
Expand Down
8 changes: 8 additions & 0 deletions src/dask_awkward/lib/io/json.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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__(
Expand All @@ -105,6 +108,7 @@ def __init__(
schema=schema,
form=form,
behavior=behavior,
attrs=attrs,
**kwargs,
)

Expand All @@ -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__(
Expand All @@ -139,6 +144,7 @@ def __init__(
schema=schema,
form=form,
behavior=behavior,
attrs=attrs,
**kwargs,
)

Expand Down Expand Up @@ -169,13 +175,15 @@ 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__(
storage=storage,
compression=compression,
schema=schema,
behavior=behavior,
attrs=attrs,
form=form,
**kwargs,
)
Expand Down
22 changes: 22 additions & 0 deletions tests/test_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
Loading