diff --git a/src/dask_awkward/lib/structure.py b/src/dask_awkward/lib/structure.py index af72d527..6b70c8b9 100644 --- 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 index 9b761339..9f86b747 100644 --- 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):