Merge branch 'main' into repartition

.pre-commit-config.yaml

@@ -21,7 +21,7 @@ repos:
     - --target-version=py38
 
 - repo: https://github.com/charliermarsh/ruff-pre-commit
-  rev: v0.0.272
+  rev: v0.0.275
   hooks:
   - id: ruff
 

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
   "Topic :: Scientific/Engineering",
 ]
 dependencies = [
-  "awkward >=2.2.2",
+  "awkward >=2.2.4",
   "dask >=2023.04.0",
 ]
 dynamic = ["version"]

src/dask_awkward/lib/__init__.py

@@ -47,6 +47,7 @@
     cartesian,
     combinations,
     copy,
+    drop_none,
     fill_none,
     firsts,
     flatten,

src/dask_awkward/lib/core.py

@@ -1242,6 +1242,19 @@ def clear_divisions(self) -> None:
         """Clear the divisions of a Dask Awkward Collection."""
         self._divisions = (None,) * (self.npartitions + 1)
 
+    def __awkward_function__(self, func, array_likes, args, kwargs):
+        import dask_awkward
+
+        if any(isinstance(arg, ak.Array) for arg in array_likes):
+            raise TypeError("cannot mix awkward.Array and dask_awkward.Array")
+
+        fn_name = func.__qualname__
+        try:
+            fn = getattr(dask_awkward, fn_name)
+        except AttributeError:
+            return NotImplemented
+        return fn(*args, **kwargs)
+
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         if method != "__call__":
             raise RuntimeError("Array ufunc supports only method == '__call__'")

src/dask_awkward/lib/structure.py

@@ -29,6 +29,7 @@
     "cartesian",
     "combinations",
     "copy",
+    "drop_none",
     "fill_none",
     "firsts",
     "flatten",
@@ -334,6 +335,28 @@ def fill_none(
     return map_partitions(fn, array, label="fill-none", output_divisions=1)
 
 
+class _DropNoneFn:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+    def __call__(self, arr):
+        return ak.drop_none(arr, **self.kwargs)
+
+
+@borrow_docstring(ak.drop_none)
+def drop_none(
+    array: Array,
+    axis: int | None = None,
+    highlevel: bool = True,
+    behavior: dict | None = None,
+) -> Array:
+    if not highlevel:
+        raise ValueError("Only highlevel=True is supported")
+
+    fn = _DropNoneFn(axis=axis, highlevel=highlevel, behavior=behavior)
+    return map_partitions(fn, array, label="drop-none", output_divisions=1)
+
+
 class _FirstsFn:
     def __init__(self, **kwargs):
         self.kwargs = kwargs
@@ -431,6 +454,7 @@ def full_like(array, fill_value, highlevel=True, behavior=None, dtype=None):
         highlevel=highlevel,
         behavior=behavior,
         dtype=dtype,
+        output_divisions=1,
     )
 
 
@@ -454,6 +478,7 @@ def isclose(
         highlevel=highlevel,
         behavior=behavior,
         label="is-close",
+        output_divisions=1,
     )
 
 
@@ -541,6 +566,7 @@ def num(
             axis=axis,
             highlevel=highlevel,
             behavior=behavior,
+            output_divisions=1,
         )
     if axis == 0:
         return len(array)
@@ -559,10 +585,10 @@ def ones_like(
     return map_partitions(
         ak.ones_like,
         array,
-        output_divisions=1,
         label="ones-like",
         behavior=behavior,
         dtype=dtype,
+        output_divisions=1,
     )
 
 
@@ -967,10 +993,10 @@ def zeros_like(
     return map_partitions(
         ak.zeros_like,
         array,
-        output_divisions=1,
         label="zeros-like",
         behavior=behavior,
         dtype=dtype,
+        output_divisions=1,
     )
 
 

tests/test_core.py

@@ -463,6 +463,23 @@ def test_compatible_partitions(ndjson_points_file: str) -> None:
     assert compatible_partitions(y, y)
 
 
+def test_compatible_partitions_after_slice() -> None:
+    a = [[1, 2, 3], [4, 5]]
+    b = [[5, 6, 7, 8], [], [9]]
+    lazy = dak.from_lists([a, b])
+    ccrt = ak.Array(a + b)
+
+    # sanity
+    assert_eq(lazy, ccrt)
+
+    # sanity
+    assert compatible_partitions(lazy, lazy + 2)
+    assert compatible_partitions(lazy, dak.num(lazy, axis=1) > 2)
+
+    assert not compatible_partitions(lazy[:-2], lazy)
+    assert not compatible_partitions(lazy[:-2], dak.num(lazy, axis=1) != 3)
+
+
 @pytest.mark.parametrize("meta", [5, False, [1, 2, 3]])
 def test_bad_meta_type(ndjson_points_file: str, meta: Any) -> None:
     with pytest.raises(TypeError, match="meta must be an instance of an Awkward Array"):
@@ -522,7 +539,6 @@ def test_scalar_persist_and_rebuild(daa: Array) -> None:
 
 def test_output_divisions(daa: Array) -> None:
     assert dak.max(daa.points.y, axis=1).divisions == daa.divisions
-    assert dak.num(daa.points.y, axis=1).divisions == (None,) * (daa.npartitions + 1)
     assert daa["points"][["x", "y"]].divisions == daa.divisions
     assert daa["points"].divisions == daa.divisions
 

tests/test_structure.py

@@ -138,6 +138,16 @@ def test_fill_none(vf: int | float | str, axis: int | None) -> None:
     assert_eq(d, e, check_forms=(not isinstance(vf, str)))
 
 
+@pytest.mark.parametrize("axis", [None, 0, 1, -1])
+def test_drop_none(axis: int) -> None:
+    a = [[1, 2, None], [], [None], [5, 6, 7, None], [1, 2], None]
+    b = [[None, 2, 1], [None], [], None, [7, 6, None, 5], [None, None]]
+    c = dak.from_lists([a, b])
+    d = dak.drop_none(c)
+    e = ak.drop_none(ak.from_iter(a + b))
+    assert_eq(d, e)
+
+
 @pytest.mark.parametrize("axis", [0, 1, -1])
 def test_is_none(axis: int) -> None:
     a: list[Any] = [[1, 2, None], None, None, [], [None], [5, 6, 7, None], [1, 2], None]
@@ -468,7 +478,6 @@ def test_from_regular(caa):
     )
 
 
-@pytest.mark.xfail(reason="typetracer")
 def test_to_regular(caa):
     regular = ak.to_packed(caa[[0, 4, 5, 9, 10, 14]].points.x)
     dregular = dak.from_awkward(regular, 3)