fix: allow for Nones in repartition n_to_one

dask-contrib · Jun 19, 2024 · 188483e · 188483e
1 parent 38c065b
commit 188483e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/src/dask_awkward/lib/structure.py b/src/dask_awkward/lib/structure.py
@@ -1421,7 +1421,7 @@ def simple_repartition_layer(
         if arr.npartitions % n_to_one:
             new_divisions = new_divisions + (arr.divisions[-1],)
             layer[(key, i0 + 1)] = (_subcat,) + tuple(
-                (arr.name, part) for part in range(new_divisions[-2], new_divisions[-1])
+                (arr.name, part0) for part0 in range(len(layer), arr.npartitions)
             )
     elif one_to_n:
         for i in range(arr.npartitions):

diff --git a/tests/test_structure.py b/tests/test_structure.py
@@ -555,7 +555,7 @@ def test_repartition_one_to_n(daa):
 
 
 def test_repartition_n_to_one():
-    daa = dak.from_lists([[[1, 2, 3], [], [4, 5]]] * 52)
+    daa = dak.from_lists([[[1, 2, 3], [], [4, 5]] * 2] * 52)
     daa2 = daa.repartition(n_to_one=52)
     assert daa2.npartitions == 1
     assert daa.compute().to_list() == daa2.compute().to_list()
@@ -568,6 +568,9 @@ def test_repartition_n_to_one():
     daa2 = daa.repartition(n_to_one=10)
     assert daa2.npartitions == 6
     assert daa.compute().to_list() == daa2.compute().to_list()
+    daa._divisions = (None,) * len(daa.divisions)
+    assert daa2.npartitions == 6
+    assert daa.compute().to_list() == daa2.compute().to_list()
 
 
 def test_repartition_no_change(daa):