imports

src/dask_awkward/lib/__init__.py

@@ -1,6 +1,11 @@
-from dask_awkward.lib.core import Array, Record, Scalar
+from dask_awkward.lib.core import Array, PartitionCompatibility, Record, Scalar
 from dask_awkward.lib.core import _type as type
-from dask_awkward.lib.core import map_partitions, typetracer_from_form
+from dask_awkward.lib.core import (
+    compatible_partitions,
+    map_partitions,
+    partition_compatibility,
+    typetracer_from_form,
+)
 from dask_awkward.lib.describe import fields
 from dask_awkward.lib.inspect import necessary_columns, sample
 from dask_awkward.lib.io.io import (

tests/test_core.py

@@ -17,11 +17,9 @@
 
 import dask_awkward as dak
 from dask_awkward.lib.core import (
-    PartitionCompatibility,
     Record,
     Scalar,
     calculate_known_divisions,
-    compatible_partitions,
     compute_typetracer,
     is_typetracer,
     meta_or_identity,
@@ -30,7 +28,6 @@
     new_record_object,
     new_scalar_object,
     normalize_single_outer_inner_index,
-    partition_compatibility,
     to_meta,
     typetracer_array,
 )
@@ -453,24 +450,24 @@ def test_compatible_partitions(ndjson_points_file: str) -> None:
     daa1 = dak.from_json([ndjson_points_file] * 5)
     daa2 = dak.from_awkward(daa1.compute(), npartitions=4)
     with pytest.warns(DeprecationWarning):
-        assert compatible_partitions(daa1, daa1)
+        assert dak.compatible_partitions(daa1, daa1)
     with pytest.warns(DeprecationWarning):
-        assert compatible_partitions(daa1, daa1, daa1)
+        assert dak.compatible_partitions(daa1, daa1, daa1)
     with pytest.warns(DeprecationWarning):
-        assert not compatible_partitions(daa1, daa2)
+        assert not dak.compatible_partitions(daa1, daa2)
     daa1.eager_compute_divisions()
     with pytest.warns(DeprecationWarning):
-        assert compatible_partitions(daa1, daa1)
+        assert dak.compatible_partitions(daa1, daa1)
     x = ak.Array([[1, 2, 3], [1, 2, 3], [3, 4, 5]])
     y = ak.Array([[1, 2, 3], [3, 4, 5]])
     x = dak.from_awkward(x, npartitions=2)
     y = dak.from_awkward(y, npartitions=2)
     with pytest.warns(DeprecationWarning):
-        assert not compatible_partitions(x, y)
+        assert not dak.compatible_partitions(x, y)
     with pytest.warns(DeprecationWarning):
-        assert not compatible_partitions(x, x, y)
+        assert not dak.compatible_partitions(x, x, y)
     with pytest.warns(DeprecationWarning):
-        assert compatible_partitions(y, y)
+        assert dak.compatible_partitions(y, y)
 
 
 def test_compatible_partitions_after_slice() -> None:
@@ -484,14 +481,14 @@ def test_compatible_partitions_after_slice() -> None:
 
     # sanity
     with pytest.warns(DeprecationWarning):
-        assert compatible_partitions(lazy, lazy + 2)
+        assert dak.compatible_partitions(lazy, lazy + 2)
     with pytest.warns(DeprecationWarning):
-        assert compatible_partitions(lazy, dak.num(lazy, axis=1) > 2)
+        assert dak.compatible_partitions(lazy, dak.num(lazy, axis=1) > 2)
 
     with pytest.warns(DeprecationWarning):
-        assert not compatible_partitions(lazy[:-2], lazy)
+        assert not dak.compatible_partitions(lazy[:-2], lazy)
     with pytest.warns(DeprecationWarning):
-        assert not compatible_partitions(lazy[:-2], dak.num(lazy, axis=1) != 3)
+        assert not dak.compatible_partitions(lazy[:-2], dak.num(lazy, axis=1) != 3)
 
     with pytest.raises(IncompatiblePartitions, match="incompatibly partitioned"):
         (lazy[:-2] + lazy).compute()
@@ -531,11 +528,14 @@ def test_partition_compatiblity() -> None:
     b = dak.from_awkward(a, npartitions=2)
     c = b[dak.sum(b, axis=1) == 0]
     d = b[dak.sum(b, axis=1) == 6]
-    assert partition_compatibility(c, d) == PartitionCompatibility.MAYBE
-    assert partition_compatibility(b, c, d) == PartitionCompatibility.MAYBE
-    assert partition_compatibility(b, dak.num(b, axis=1)) == PartitionCompatibility.YES
+    assert dak.partition_compatibility(c, d) == dak.PartitionCompatibility.MAYBE
+    assert dak.partition_compatibility(b, c, d) == dak.PartitionCompatibility.MAYBE
+    assert (
+        dak.partition_compatibility(b, dak.num(b, axis=1))
+        == dak.PartitionCompatibility.YES
+    )
     c.eager_compute_divisions()
-    assert partition_compatibility(b, c) == PartitionCompatibility.NO
+    assert dak.partition_compatibility(b, c) == dak.PartitionCompatibility.NO
 
 
 @pytest.mark.parametrize("meta", [5, False, [1, 2, 3]])