Merge pull request #542 from pfackeldey/feat/align_array_behavior_of_…

…_bool_nonzero_long_int_float_complex_index__with_dask

[feat] Explicitly raise `TracerConversionError` when a tracer is used in a context that requires a concrete value

docs/api/utils.rst

@@ -0,0 +1,23 @@
+Errors
+------
+
+Utilities to implement array behaviors for dask-awkward arrays.
+
+
+.. currentmodule:: dask_awkward
+
+
+.. autosummary::
+   :toctree: generated/
+
+   utils.IncompatiblePartitions
+
+.. autosummary::
+   :toctree: generated/
+
+   utils.ConcretizationTypeError
+
+.. raw:: html
+
+   <script data-goatcounter="https://dask-awkward.goatcounter.com/count"
+           async src="//gc.zgo.at/count.js"></script>

docs/index.rst

@@ -50,6 +50,7 @@ Table of Contents
    api/reducers.rst
    api/structure.rst
    api/behavior.rst
+   api/utils.rst
 
 .. toctree::
    :maxdepth: 1

src/dask_awkward/lib/core.py

@@ -48,6 +48,7 @@
 from dask_awkward.layers import AwkwardBlockwiseLayer, AwkwardMaterializedLayer
 from dask_awkward.lib.optimize import all_optimizations
 from dask_awkward.utils import (
+    ConcretizationTypeError,
     DaskAwkwardNotImplemented,
     IncompatiblePartitions,
     field_access_to_front,
@@ -921,6 +922,23 @@ def __dask_postpersist__(self):
 
     __dask_scheduler__ = staticmethod(threaded_get)
 
+    def __bool__(self):
+        raise ConcretizationTypeError(f"The __bool__() method was called on {self!r}.")
+
+    def __int__(self):
+        raise ConcretizationTypeError(f"The __int__() method was called on {self!r}.")
+
+    def __float__(self):
+        raise ConcretizationTypeError(f"The __float__() method was called on {self!r}.")
+
+    def __complex__(self):
+        raise ConcretizationTypeError(
+            f"The __complex__() method was called on {self!r}."
+        )
+
+    def __index__(self):
+        raise ConcretizationTypeError(f"The __index__() method was called on {self!r}.")
+
     def __setitem__(self, where: Any, what: Any) -> None:
         if not (
             isinstance(where, str)

src/dask_awkward/utils.py

@@ -8,7 +8,6 @@
 if TYPE_CHECKING:
     from dask_awkward.lib.core import Array
 
-
 T = TypeVar("T")
 P = ParamSpec("P")
 
@@ -38,6 +37,51 @@ def divisions_msg(name: str, *args: Array) -> str:
         return msg
 
 
+class ConcretizationTypeError(TypeError):
+    """
+    This error occurs when a ``dask_awkward.Array`` is used in a context that requires a concrete
+    value.
+
+
+    There are several reasons why this error might occur:
+
+    Examples
+    --------
+
+    - When a ``dask_awkward.Array`` is used in a conditional statement:
+
+    >>> import dask_awkward as dak
+    >>> import awkward as ak
+    >>> dask_arr = dak.from_awkward(ak.Array([1]), npartitions=1)
+    >>> if dask_arr > 2:
+    >>>     dask_arr += 1
+    Traceback (most recent call last): ...
+    dask_awkward.utils.ConcretizationTypeError: A dask_awkward.Array is encountered in a computation where a concrete value is expected. If you intend to convert the dask_awkward.Array to a concrete value, use the `.compute()` method. The __bool__() method was called on dask.awkward<greater, npartitions=1>.
+
+    - When a ``dask_awkward.Array`` is cast to a Python type:
+
+    >>> import dask_awkward as dak
+    >>> import awkward as ak
+    >>> dask_arr = dak.from_awkward(ak.Array([1]), npartitions=1)
+    >>> int(dask_arr)
+    Traceback (most recent call last): ...
+    dask_awkward.utils.ConcretizationTypeError: A dask_awkward.Array is encountered in a computation where a concrete value is expected. If you intend to convert the dask_awkward.Array to a concrete value, use the `.compute()` method. The __int__() method was called on dask.awkward<from-awkward, npartitions=1>.
+
+    These errors can be resolved by explicitely converting the tracer to a concrete value:
+
+    >>> import dask_awkward as dak
+    >>> dask_arr = dak.from_awkward(ak.Array([1]), npartitions=1)
+    >>> bool(dask_arr.compute())
+    True
+    """
+
+    def __init__(self, msg: str):
+        self.message = "A dask_awkward.Array is encountered in a computation where a concrete value is expected. "
+        self.message += "If you intend to convert the dask_awkward.Array to a concrete value, use the `.compute()` method. "
+        self.message += msg
+        super().__init__(self.message)
+
+
 class LazyInputsDict(Mapping):
     """Dictionary with lazy key value pairs
 

tests/test_core.py

@@ -34,7 +34,7 @@
     typetracer_array,
 )
 from dask_awkward.lib.testutils import assert_eq
-from dask_awkward.utils import IncompatiblePartitions
+from dask_awkward.utils import ConcretizationTypeError, IncompatiblePartitions
 
 if TYPE_CHECKING:
     from dask_awkward.lib.core import Array
@@ -968,3 +968,16 @@ def test_map_partitions_bad_arguments():
             array2,
             meta=empty_typetracer(),
         )
+
+
+def test_array__bool_nonzero_long_int_float_complex_index():
+    import operator
+
+    dask_arr = dak.from_awkward(ak.Array([1]), npartitions=1)
+
+    for fun in bool, int, float, complex, operator.index:
+        with pytest.raises(
+            ConcretizationTypeError,
+            match=r"A dask_awkward.Array is encountered in a computation where a concrete value is expected. If you intend to convert the dask_awkward.Array to a concrete value, use the `.compute\(\)` method. The .+ method was called on .+.",
+        ):
+            fun(dask_arr)

tests/test_io_json.py

@@ -68,9 +68,9 @@ def test_json_sanity(json_data_dir: Path, concrete_data: ak.Array) -> None:
             ", use `\\.eager_compute_divisions\\(\\)` on the collection."
         ),
     ):
-        assert ds
+        assert len(ds)
     ds.eager_compute_divisions()
-    assert ds
+    assert len(ds)
 
     assert_eq(ds, concrete_data)