(ignore_index, npartitions, max_branch, df)
| 76 | @pytest.mark.parametrize("npartitions", [8, 12]) |
| 77 | @pytest.mark.parametrize("max_branch", [32, 6]) |
| 78 | def test_task_shuffle(ignore_index, npartitions, max_branch, df): |
| 79 | df2 = df.shuffle( |
| 80 | "x", |
| 81 | shuffle_method="tasks", |
| 82 | npartitions=npartitions, |
| 83 | ignore_index=ignore_index, |
| 84 | max_branch=max_branch, |
| 85 | ) |
| 86 | |
| 87 | # Check that the output partition count is correct |
| 88 | assert df2.npartitions == (npartitions or df.npartitions) |
| 89 | |
| 90 | # Check the computed (re-ordered) result |
| 91 | assert_eq(df, df2, check_index=not ignore_index, check_divisions=False) |
| 92 | |
| 93 | # Check that df was really partitioned by "x". |
| 94 | # If any values of "x" can be found in multiple |
| 95 | # partitions, this will fail |
| 96 | df3 = df2["x"].map_partitions(lambda x: x.drop_duplicates()) |
| 97 | assert sorted(df3.compute().values) == list(range(20)) |
| 98 | |
| 99 | # Check `partitions` after shuffle |
| 100 | a = df2.partitions[1] |
| 101 | b = df.shuffle( |
| 102 | "y", |
| 103 | shuffle_method="tasks", |
| 104 | npartitions=npartitions, |
| 105 | ignore_index=ignore_index, |
| 106 | ).partitions[1] |
| 107 | assert set(a["x"].compute().values.tolist()).issubset( |
| 108 | b["y"].compute().values.tolist() |
| 109 | ) |
| 110 | |
| 111 | # Check for culling |
| 112 | assert len(a.optimize().dask) < len(df2.optimize().dask) |
| 113 | |
| 114 | |
| 115 | @pytest.mark.parametrize("npartitions", [3, 12]) |
nothing calls this directly
no test coverage detected