chore: import upstream snapshot with attribution

python/ray/data/tests/unit/test_block.py

@@ -0,0 +1,164 @@
+from datetime import datetime
+
+import numpy as np
+import pyarrow as pa
+import pytest
+
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+from ray.data.block import BlockAccessor, BlockColumnAccessor
+
+
+def test_find_partitions_single_column_ascending():
+    table = pa.table({"value": [1, 2, 2, 3, 5]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    boundaries = [(0,), (2,), (4,), (6,)]
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+    assert len(partitions) == 5
+    assert partitions[0].to_pydict() == {"value": []}  # <0
+    assert partitions[1].to_pydict() == {"value": [1]}  # [0,2)
+    assert partitions[2].to_pydict() == {"value": [2, 2, 3]}  # [2,4)
+    assert partitions[3].to_pydict() == {"value": [5]}  # [4,6)
+    assert partitions[4].to_pydict() == {"value": []}  # >=6
+
+
+def test_find_partitions_single_column_descending():
+    table = pa.table({"value": [5, 3, 2, 2, 1]})
+    sort_key = SortKey(key=["value"], descending=[True])
+    boundaries = [(6,), (3,), (2,), (0,)]
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+    assert len(partitions) == 5
+    assert partitions[0].to_pydict() == {"value": []}  # >=6
+    assert partitions[1].to_pydict() == {"value": [5, 3]}  # [3, 6)
+    assert partitions[2].to_pydict() == {"value": [2, 2]}  # [2, 3)
+    assert partitions[3].to_pydict() == {"value": [1]}  # [0, 2)
+    assert partitions[4].to_pydict() == {"value": []}  # <0
+
+
+def test_find_partitions_multi_column_ascending_first():
+    table = pa.table({"col1": [1, 1, 1, 1, 1, 2, 2], "col2": [4, 3, 2.5, 2, 1, 2, 1]})
+    sort_key = SortKey(key=["col1", "col2"], descending=[False, True])
+    boundaries = [(1, 3), (1, 2), (2, 2), (2, 0)]
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+    assert len(partitions) == 5
+    assert partitions[0].to_pydict() == {"col1": [1], "col2": [4]}
+    assert partitions[1].to_pydict() == {"col1": [1, 1], "col2": [3, 2.5]}
+    assert partitions[2].to_pydict() == {"col1": [1, 1], "col2": [2, 1]}
+    assert partitions[3].to_pydict() == {"col1": [2, 2], "col2": [2, 1]}
+    assert partitions[4].to_pydict() == {"col1": [], "col2": []}
+
+
+def test_find_partitions_multi_column_descending_first():
+    table = pa.table({"col1": [2, 2, 1, 1, 1, 1, 1], "col2": [1, 2, 1, 2, 3, 4, 5]})
+    sort_key = SortKey(key=["col1", "col2"], descending=[True, False])
+    boundaries = [(2, 0), (2, 2), (1, 2), (1, 6)]
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+    assert len(partitions) == 5
+    assert partitions[0].to_pydict() == {"col1": [], "col2": []}
+    assert partitions[1].to_pydict() == {"col1": [2, 2], "col2": [1, 2]}
+    assert partitions[2].to_pydict() == {"col1": [1, 1], "col2": [1, 2]}
+    assert partitions[3].to_pydict() == {"col1": [1, 1, 1], "col2": [3, 4, 5]}
+    assert partitions[4].to_pydict() == {"col1": [], "col2": []}
+
+
+@pytest.mark.parametrize("null", [None, np.nan])
+def test_find_partitions_table_with_nulls(null):
+    table = pa.table({"value": [1, 2, 3, null, null]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    boundaries = [(2,), (4,)]
+
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+
+    assert len(partitions) == 3
+    assert partitions[0].to_pydict() == {"value": [1]}  # <2
+    assert partitions[1].to_pydict() == {"value": [2, 3]}  # [2, 4)
+
+    if null is None:
+        assert partitions[2].to_pydict() == {"value": [null, null]}  # >=4
+    else:
+        # NOTE: NaNs couldn't be compared directly
+        result = partitions[2].to_pydict()
+        assert len(result["value"]) == 2 and all([np.isnan(v) for v in result["value"]])
+
+
+@pytest.mark.parametrize("dtype", ["str", "datetime"])
+def test_find_partitions_object_table_with_nulls(dtype):
+    if dtype == "str":
+        col = pa.array(["a", "b", "c", None, None])
+    elif dtype == "datetime":
+        col = pa.array(
+            [
+                datetime.fromordinal(1),
+                datetime.fromordinal(2),
+                datetime.fromordinal(3),
+                None,
+                None,
+            ]
+        )
+    else:
+        raise ValueError(f"Unexpected dtype={dtype}")
+
+    table = pa.table({"value": col})
+
+    sort_key = SortKey(key=["value"], descending=[False])
+
+    ndarray = BlockColumnAccessor.for_column(col).to_numpy(zero_copy_only=False)
+
+    # Compose boundaries
+    boundaries = [(ndarray[1],), (None,)]
+
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+
+    assert len(partitions) == 3
+    assert (
+        partitions[0]["value"].to_pylist() == col[0:1].to_pylist()
+    )  # < second element
+    assert partitions[1]["value"].to_pylist() == col[1:3].to_pylist()  # < None
+    assert partitions[2]["value"].to_pylist() == col[3:].to_pylist()  # remaining
+
+
+@pytest.mark.parametrize("null", [None, np.nan])
+def test_find_partitions_null_boundary(null):
+    table = pa.table({"value": [1, 2, 3]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    boundaries = [(2,), (null,)]
+
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+
+    assert len(partitions) == 3
+    assert partitions[0].to_pydict() == {"value": [1]}  # <2
+    assert partitions[1].to_pydict() == {"value": [2, 3]}  # < null (nulls go last)
+    assert partitions[2].to_pydict() == {"value": []}
+
+
+def test_find_partitions_duplicates():
+    table = pa.table({"value": [2, 2, 2, 2, 2]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    boundaries = [(1,), (2,), (3,)]
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+    assert len(partitions) == 4
+    assert partitions[0].to_pydict() == {"value": []}  # <1
+    assert partitions[1].to_pydict() == {"value": []}  # [1,2)
+    assert partitions[2].to_pydict() == {"value": [2, 2, 2, 2, 2]}  # [2,3)
+    assert partitions[3].to_pydict() == {"value": []}  # >=3
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))