Merge pull request #518 from martindurant/repartition_again

fix: allow for Nones in repartition n_to_one

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):

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):