chore: import upstream snapshot with attribution

python/ray/data/tests/test_sort.py

@@ -0,0 +1,348 @@
+import random
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey, SortTaskSpec
+from ray.data.block import BlockAccessor
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.util import extract_values
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.mark.parametrize(
+    "descending,boundaries",
+    [
+        (True, list(range(100, 1000, 200))),
+        (False, list(range(100, 1000, 200))),
+        (True, [1, 998]),
+        (False, [1, 998]),
+        # Test float.
+        (True, [501.5]),
+        (False, [501.5]),
+    ],
+)
+def test_sort_with_specified_boundaries(ray_start_regular, descending, boundaries):
+    num_items = 1000
+    ds = ray.data.range(num_items)
+    ds = ds.sort("id", descending, boundaries).materialize()
+
+    items = range(num_items)
+    boundaries = [0] + sorted([round(b) for b in boundaries]) + [num_items]
+    expected_blocks = [
+        items[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)
+    ]
+    if descending:
+        expected_blocks = [list(reversed(block)) for block in reversed(expected_blocks)]
+
+    blocks = list(ds.iter_batches(batch_size=None))
+    assert len(blocks) == len(expected_blocks)
+    for block, expected_block in zip(blocks, expected_blocks):
+        assert np.all(block["id"] == expected_block)
+
+
+def test_sort_multiple_keys_produces_equally_sized_blocks(ray_start_regular):
+    # Test for https://github.com/ray-project/ray/issues/45303.
+    ds = ray.data.from_items(
+        [{"a": i, "b": j} for i in range(2) for j in range(5)], override_num_blocks=5
+    )
+
+    ds_sorted = ds.sort(["a", "b"], descending=[False, True])
+
+    num_rows_per_block = [
+        bundle.num_rows() for bundle in ds_sorted.iter_internal_ref_bundles()
+    ]
+    # Number of output blocks should be equal to the number of input blocks.
+    assert len(num_rows_per_block) == 5, len(num_rows_per_block)
+    # Ideally we should have 10 rows / 5 blocks = 2 rows per block, but to make this
+    # test less fragile we allow for a small deviation.
+    assert all(
+        1 <= num_rows <= 3 for num_rows in num_rows_per_block
+    ), num_rows_per_block
+
+
+def test_sort_simple(ray_start_regular, configure_shuffle_method):
+    num_items = 100
+    parallelism = 4
+    xs = list(range(num_items))
+    random.shuffle(xs)
+    ds = ray.data.from_items(xs, override_num_blocks=parallelism)
+    assert extract_values("item", ds.sort("item").take(num_items)) == list(
+        range(num_items)
+    )
+    # Make sure we have rows in each block.
+    assert len([n for n in ds.sort("item")._block_num_rows() if n > 0]) == parallelism
+
+    assert extract_values(
+        "item", ds.sort("item", descending=True).take(num_items)
+    ) == list(reversed(range(num_items)))
+
+    # Test empty dataset.
+    ds = ray.data.from_items([])
+    s1 = ds.sort("item")
+    assert s1.count() == 0
+    assert s1.take() == ds.take()
+    ds = ray.data.range(10).filter(lambda r: r["id"] > 10).sort("id")
+    assert ds.count() == 0
+
+
+def test_sort_partition_same_key_to_same_block(
+    ray_start_regular, configure_shuffle_method
+):
+    num_items = 100
+    xs = [1] * num_items
+    ds = ray.data.from_items(xs)
+    sorted_ds = ds.repartition(num_items).sort("item")
+
+    # We still have 100 blocks
+    assert len(sorted_ds._block_num_rows()) == num_items
+    # Only one of them is non-empty
+    count = sum(1 for x in sorted_ds._block_num_rows() if x > 0)
+    assert count == 1
+    # That non-empty block contains all rows
+    total = sum(x for x in sorted_ds._block_num_rows() if x > 0)
+    assert total == num_items
+
+
+@pytest.mark.parametrize("num_items,parallelism", [(100, 1), (1000, 4)])
+def test_sort_arrow(
+    ray_start_regular,
+    num_items,
+    parallelism,
+    configure_shuffle_method,
+    use_polars_sort,
+):
+    ctx = ray.data.context.DataContext.get_current()
+
+    try:
+        original_use_polars = ctx.use_polars_sort
+        ctx.use_polars_sort = use_polars_sort
+
+        a = list(reversed(range(num_items)))
+        b = [f"{x:03}" for x in range(num_items)]
+        shard = int(np.ceil(num_items / parallelism))
+        offset = 0
+        dfs = []
+        while offset < num_items:
+            dfs.append(
+                pd.DataFrame(
+                    {"a": a[offset : offset + shard], "b": b[offset : offset + shard]}
+                )
+            )
+            offset += shard
+        if offset < num_items:
+            dfs.append(pd.DataFrame({"a": a[offset:], "b": b[offset:]}))
+        ds = ray.data.from_blocks(dfs).map_batches(
+            lambda t: t, batch_format="pyarrow", batch_size=None
+        )
+
+        def assert_sorted(sorted_ds, expected_rows):
+            assert [tuple(row.values()) for row in sorted_ds.iter_rows()] == list(
+                expected_rows
+            )
+
+        assert_sorted(ds.sort(key="a"), zip(reversed(a), reversed(b)))
+        # Make sure we have rows in each block.
+        assert (
+            len([n for n in ds.sort(key="a")._block_num_rows() if n > 0]) == parallelism
+        )
+        assert_sorted(ds.sort(key="b"), zip(a, b))
+        assert_sorted(ds.sort(key="a", descending=True), zip(a, b))
+    finally:
+        ctx.use_polars_sort = original_use_polars
+
+
+def test_sort(ray_start_regular, use_polars_sort):
+    import random
+
+    import pyarrow as pa
+
+    N = 100
+    r = random.Random(0xDEED)
+
+    ints = [r.randint(0, 10) for _ in range(N)]
+    floats = [r.normalvariate(0, 5) for _ in range(N)]
+    t = pa.Table.from_pydict({"ints": ints, "floats": floats})
+
+    sorted_block = BlockAccessor.for_block(t).sort(SortKey(["ints", "floats"]))
+
+    sorted_tuples = list(zip(*sorted(zip(ints, floats))))
+
+    assert sorted_block == pa.Table.from_pydict(
+        {"ints": sorted_tuples[0], "floats": sorted_tuples[1]}
+    )
+
+
+def test_sort_arrow_with_empty_blocks(
+    ray_start_regular, configure_shuffle_method, use_polars_sort
+):
+    ctx = ray.data.context.DataContext.get_current()
+
+    try:
+        original_use_polars = ctx.use_polars_sort
+        ctx.use_polars_sort = use_polars_sort
+
+        assert (
+            BlockAccessor.for_block(pa.Table.from_pydict({}))
+            .sample(10, SortKey("A"))
+            .num_rows
+            == 0
+        )
+
+        partitions = BlockAccessor.for_block(
+            pa.Table.from_pydict({})
+        ).sort_and_partition([1, 5, 10], SortKey("A"))
+        assert len(partitions) == 4
+        for partition in partitions:
+            assert partition.num_rows == 0
+
+        assert (
+            BlockAccessor.for_block(pa.Table.from_pydict({}))
+            .merge_sorted_blocks([pa.Table.from_pydict({})], SortKey("A"))[1]
+            .metadata.num_rows
+            == 0
+        )
+
+        ds = ray.data.from_items(
+            [{"A": (x % 3), "B": x} for x in range(3)], override_num_blocks=3
+        )
+        ds = ds.filter(lambda r: r["A"] == 0)
+        assert list(ds.sort("A").iter_rows()) == [{"A": 0, "B": 0}]
+
+        # Test empty dataset.
+        ds = ray.data.range(10).filter(lambda r: r["id"] > 10)
+        assert (
+            len(
+                SortTaskSpec.sample_boundaries(
+                    ds._execute().block_refs, SortKey("id"), 3
+                )
+            )
+            == 2
+        )
+        assert ds.sort("id").count() == 0
+    finally:
+        ctx.use_polars_sort = original_use_polars
+
+
+@pytest.mark.parametrize("descending", [False, True])
+@pytest.mark.parametrize("batch_format", ["pyarrow", "pandas"])
+def test_sort_with_multiple_keys(ray_start_regular, descending, batch_format):
+    num_items = 1000
+    num_blocks = 100
+    df = pd.DataFrame(
+        {
+            "a": [random.choice("ABCD") for _ in range(num_items)],
+            "b": [x % 3 for x in range(num_items)],
+            "c": [bool(random.getrandbits(1)) for _ in range(num_items)],
+        }
+    )
+    ds = ray.data.from_pandas(df).map_batches(
+        lambda t: t,
+        batch_format=batch_format,
+        batch_size=None,
+    )
+    df.sort_values(
+        ["a", "b", "c"],
+        inplace=True,
+        ascending=[not descending, descending, not descending],
+    )
+    sorted_ds = ds.repartition(num_blocks).sort(
+        ["a", "b", "c"], descending=[descending, not descending, descending]
+    )
+
+    # Number of blocks is preserved
+    assert len(sorted_ds._block_num_rows()) == num_blocks
+    # Rows are sorted over the dimensions
+    assert [tuple(row.values()) for row in sorted_ds.iter_rows()] == list(
+        zip(df["a"], df["b"], df["c"])
+    )
+
+
+@pytest.mark.parametrize("num_items,parallelism", [(100, 1), (1000, 4)])
+def test_sort_pandas(
+    ray_start_regular, num_items, parallelism, configure_shuffle_method
+):
+    a = list(reversed(range(num_items)))
+    b = [f"{x:03}" for x in range(num_items)]
+    shard = int(np.ceil(num_items / parallelism))
+    offset = 0
+    dfs = []
+    while offset < num_items:
+        dfs.append(
+            pd.DataFrame(
+                {"a": a[offset : offset + shard], "b": b[offset : offset + shard]}
+            )
+        )
+        offset += shard
+    if offset < num_items:
+        dfs.append(pd.DataFrame({"a": a[offset:], "b": b[offset:]}))
+    ds = ray.data.from_blocks(dfs)
+
+    def assert_sorted(sorted_ds, expected_rows):
+        assert [tuple(row.values()) for row in sorted_ds.iter_rows()] == list(
+            expected_rows
+        )
+
+    assert_sorted(ds.sort(key="a"), zip(reversed(a), reversed(b)))
+    # Make sure we have rows in each block.
+    assert len([n for n in ds.sort(key="a")._block_num_rows() if n > 0]) == parallelism
+    assert_sorted(ds.sort(key="b"), zip(a, b))
+    assert_sorted(ds.sort(key="a", descending=True), zip(a, b))
+
+
+def test_sort_pandas_with_empty_blocks(ray_start_regular, configure_shuffle_method):
+    assert (
+        BlockAccessor.for_block(pa.Table.from_pydict({}))
+        .sample(10, SortKey("A"))
+        .num_rows
+        == 0
+    )
+
+    partitions = BlockAccessor.for_block(pa.Table.from_pydict({})).sort_and_partition(
+        [1, 5, 10], SortKey("A")
+    )
+    assert len(partitions) == 4
+    for partition in partitions:
+        assert partition.num_rows == 0
+
+    assert (
+        BlockAccessor.for_block(pa.Table.from_pydict({}))
+        .merge_sorted_blocks([pa.Table.from_pydict({})], SortKey("A"))[1]
+        .metadata.num_rows
+        == 0
+    )
+
+    ds = ray.data.from_items(
+        [{"A": (x % 3), "B": x} for x in range(3)], override_num_blocks=3
+    )
+    ds = ds.filter(lambda r: r["A"] == 0)
+    assert list(ds.sort("A").iter_rows()) == [{"A": 0, "B": 0}]
+
+    # Test empty dataset.
+    ds = ray.data.range(10).filter(lambda r: r["id"] > 10)
+    assert (
+        len(SortTaskSpec.sample_boundaries(ds._execute().block_refs, SortKey("id"), 3))
+        == 2
+    )
+    assert ds.sort("id").count() == 0
+
+
+def test_sort_with_one_block(shutdown_only, configure_shuffle_method):
+    ray.init(num_cpus=8)
+    ctx = ray.data.DataContext.get_current()
+    ctx.execution_options.verbose_progress = True
+    ctx.use_push_based_shuffle = True
+
+    # Use a dataset that will produce only one block to sort.
+    ray.data.range(1024).map_batches(
+        lambda _: pa.table([pa.array([1])], ["token_counts"])
+    ).sum("token_counts")
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-sv", __file__]))