Experimental support dask_awkward collections, formatting and organization (#12)

Author: douglasdavis

.pre-commit-config.yaml

@@ -22,3 +22,7 @@ repos:
   - id: pyupgrade
     args:
     - --py37-plus
+- repo: https://github.com/MarcoGorelli/absolufy-imports
+  rev: v0.3.1
+  hooks:
+  - id: absolufy-imports

src/dask_histogram/__init__.py

@@ -3,9 +3,9 @@
 import boost_histogram.axis as _axis
 import boost_histogram.storage as _storage
 
-from ._version import version as __version__
-from .core import AggHistogram, PartitionedHistogram, factory
-from .routines import histogram, histogram2d, histogramdd
+from dask_histogram._version import version as __version__
+from dask_histogram.core import AggHistogram, PartitionedHistogram, factory
+from dask_histogram.routines import histogram, histogram2d, histogramdd
 
 version_info = tuple(__version__.split("."))
 

src/dask_histogram/bins.py

@@ -10,7 +10,7 @@
 if TYPE_CHECKING:
     from typing import Sequence
 
-    from .typing import BinArg, BinType, RangeArg, RangeType
+    from dask_histogram.typing import BinArg, BinType, RangeArg, RangeType
 
 
 class BinsStyle(Enum):

src/dask_histogram/boost.py

@@ -14,11 +14,17 @@
 from dask.delayed import Delayed, delayed
 from dask.utils import is_arraylike, is_dataframe_like
 
-from .bins import normalize_bins_range
-from .core import AggHistogram, factory
+from dask_histogram.bins import normalize_bins_range
+from dask_histogram.core import AggHistogram, factory
 
 if TYPE_CHECKING:
-    from .typing import BinArg, BinType, DaskCollection, RangeArg, RangeType
+    from dask_histogram.typing import (
+        BinArg,
+        BinType,
+        DaskCollection,
+        RangeArg,
+        RangeType,
+    )
 
 import dask_histogram
 

src/dask_histogram/core.py

@@ -8,18 +8,19 @@
 import boost_histogram as bh
 import dask.array as da
 import numpy as np
-from dask.bag.core import empty_safe_aggregate, partition_all
+from dask.bag.core import empty_safe_aggregate
 from dask.base import DaskMethodsMixin, is_dask_collection, tokenize
 from dask.dataframe.core import partitionwise_graph as partitionwise
 from dask.delayed import Delayed
 from dask.highlevelgraph import HighLevelGraph
 from dask.threaded import get as tget
 from dask.utils import is_dataframe_like, key_split
+from tlz import partition_all
 
 if TYPE_CHECKING:
     from numpy.typing import NDArray
 
-    from .typing import DaskCollection
+    from dask_histogram.typing import DaskCollection
 
 __all__ = (
     "AggHistogram",
@@ -29,7 +30,7 @@
 )
 
 
-def clone(histref: bh.Histogram = None) -> bh.Histogram:
+def clone(histref: bh.Histogram | None = None) -> bh.Histogram:
     """Create a Histogram object based on another.
 
     The axes and storage of the `histref` will be used to create a new
@@ -54,7 +55,7 @@ def clone(histref: bh.Histogram = None) -> bh.Histogram:
 def _blocked_sa(
     data: Any,
     *,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     """Blocked calculation; single argument; unweighted; no sample."""
     if data.ndim == 1:
@@ -69,7 +70,7 @@ def _blocked_sa_s(
     data: Any,
     sample: Any,
     *,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     """Blocked calculation; single argument; unweighted; with sample."""
     if data.ndim == 1:
@@ -84,7 +85,7 @@ def _blocked_sa_w(
     data: Any,
     weights: Any,
     *,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     """Blocked calculation; single argument; weighted; no sample."""
     if data.ndim == 1:
@@ -100,7 +101,7 @@ def _blocked_sa_w_s(
     weights: Any,
     sample: Any,
     *,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     """Blocked calculation; single argument; weighted; with sample."""
     if data.ndim == 1:
@@ -113,15 +114,15 @@ def _blocked_sa_w_s(
 
 def _blocked_ma(
     *data: Any,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     """Blocked calculation; multiargument; unweighted; no sample."""
     return clone(histref).fill(*data)
 
 
 def _blocked_ma_s(
     *data: Any,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     """Blocked calculation; multiargument; unweighted; with sample."""
     sample = data[-1]
@@ -131,7 +132,7 @@ def _blocked_ma_s(
 
 def _blocked_ma_w(
     *data: Any,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     """Blocked calculation; multiargument; weighted; no sample."""
     weights = data[-1]
@@ -141,7 +142,7 @@ def _blocked_ma_w(
 
 def _blocked_ma_w_s(
     *data: Any,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     """Blocked calculation; multiargument; weighted; with sample."""
     weights = data[-2]
@@ -153,7 +154,7 @@ def _blocked_ma_w_s(
 def _blocked_df(
     data: Any,
     *,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     return clone(histref).fill(*(data[c] for c in data.columns), weight=None)
 
@@ -162,7 +163,7 @@ def _blocked_df_s(
     data: Any,
     sample: Any,
     *,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     return clone(histref).fill(*(data[c] for c in data.columns), sample=sample)
 
@@ -171,7 +172,7 @@ def _blocked_df_w(
     data: Any,
     weights: Any,
     *,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     """Blocked calculation; single argument; weighted; no sample."""
     return clone(histref).fill(*(data[c] for c in data.columns), weight=weights)
@@ -182,14 +183,18 @@ def _blocked_df_w_s(
     weights: Any,
     sample: Any,
     *,
-    histref: bh.Histogram = None,
+    histref: bh.Histogram | None = None,
 ) -> bh.Histogram:
     """Blocked calculation; single argument; weighted; with sample."""
     return clone(histref).fill(
         *(data[c] for c in data.columns), weight=weights, sample=sample
     )
 
 
+def _blocked_dak(data: Any, *, histref: bh.Histogram | None = None) -> bh.Histogram:
+    return clone(histref).fill(data)
+
+
 class AggHistogram(DaskMethodsMixin):
     """Aggregated Histogram collection.
 
@@ -346,13 +351,13 @@ def to_delayed(self) -> Delayed:
         return Delayed(self.name, dsk, layer=self._layer)
 
     def values(self, flow: bool = False) -> NDArray[Any]:
-        return self.to_boost().values()
+        return self.to_boost().values(flow=flow)
 
     def variances(self, flow: bool = False) -> NDArray[Any] | None:
-        return self.to_boost().variances()
+        return self.to_boost().variances(flow=flow)
 
     def counts(self, flow: bool = False) -> NDArray[Any]:
-        return self.to_boost().counts()
+        return self.to_boost().counts(flow=flow)
 
     def __array__(self) -> NDArray[Any]:
         return self.compute().__array__()
@@ -483,12 +488,15 @@ def histref(self) -> bh.Histogram:
         """boost_histogram.Histogram: reference histogram."""
         return self._histref
 
-    def to_agg(self, split_every: int = None) -> AggHistogram:
+    def to_agg(self, split_every: int | None = None) -> AggHistogram:
         """Translate into a reduced aggregated histogram."""
         return _reduction(self, split_every=split_every)
 
 
-def _reduction(ph: PartitionedHistogram, split_every: int = None) -> AggHistogram:
+def _reduction(
+    ph: PartitionedHistogram,
+    split_every: int | None = None,
+) -> AggHistogram:
     if split_every is None:
         split_every = 4
     if split_every is False:
@@ -568,7 +576,19 @@ def _partitioned_histogram(
     name = f"hist-on-block-{tokenize(data, histref, weights, sample)}"
     data_is_df = is_dataframe_like(data[0])
     _weight_sample_check(*data, weights=weights)
-    if len(data) == 1 and not data_is_df:
+    if len(data) == 1 and hasattr(data[0], "_typetracer"):
+        from dask_awkward.core import partitionwise_layer as pwlayer
+
+        x = data[0]
+        if weights is not None and sample is not None:
+            raise NotImplementedError()
+        elif weights is not None and sample is None:
+            raise NotImplementedError()
+        elif weights is None and sample is not None:
+            raise NotImplementedError()
+        else:
+            g = pwlayer(_blocked_dak, name, x, histref=histref)
+    elif len(data) == 1 and not data_is_df:
         x = data[0]
         if weights is not None and sample is not None:
             g = partitionwise(
@@ -621,7 +641,6 @@ def _reduced_histogram(
         histref=histref,
         weights=weights,
         sample=sample,
-        split_every=split_every,
     )
     return ph.to_agg(split_every=split_every)
 
@@ -683,7 +702,11 @@ def to_dask_array(
 
 
 class BinaryOpAgg:
-    def __init__(self, func: Callable[[Any, Any], Any], name: str = None) -> None:
+    def __init__(
+        self,
+        func: Callable[[Any, Any], Any],
+        name: str | None = None,
+    ) -> None:
         self.func = func
         self.__name__ = func.__name__ if name is None else name
 

src/dask_histogram/routines.py

@@ -10,11 +10,17 @@
 from dask.base import is_dask_collection
 from dask.utils import is_arraylike, is_dataframe_like
 
-from .bins import normalize_bins_range
-from .core import AggHistogram, factory
+from dask_histogram.bins import normalize_bins_range
+from dask_histogram.core import AggHistogram, factory
 
 if TYPE_CHECKING:
-    from .typing import BinArg, BinType, DaskCollection, RangeArg, RangeType
+    from dask_histogram.typing import (
+        BinArg,
+        BinType,
+        DaskCollection,
+        RangeArg,
+        RangeType,
+    )
 else:
     DaskCollection = object