chore: import upstream snapshot with attribution

python/ray/data/tests/test_arrow_serialization.py

@@ -0,0 +1,514 @@
+import logging
+import os
+import types
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+from packaging.version import parse as parse_version
+from pytest_lazy_fixtures import lf as lazy_fixture
+
+import ray
+import ray.cloudpickle as pickle
+import ray.data
+import ray.train
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+from ray.data.extensions.object_extension import (
+    ArrowPythonObjectArray,
+)
+from ray.data.extensions.tensor_extension import (
+    ArrowTensorArray,
+    ArrowVariableShapedTensorArray,
+)
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.fixture
+def null_array():
+    return pa.array([])
+
+
+@pytest.fixture
+def int_array():
+    return pa.array(list(range(1000)))
+
+
+@pytest.fixture
+def int_array_with_nulls():
+    return pa.array((list(range(9)) + [None]) * 100)
+
+
+@pytest.fixture
+def float_array():
+    return pa.array([float(i) for i in range(1000)])
+
+
+@pytest.fixture
+def boolean_array():
+    return pa.array([True, False] * 500)
+
+
+@pytest.fixture
+def string_array():
+    return pa.array(["foo", "bar", "bz", None, "quux"] * 200)
+
+
+@pytest.fixture
+def large_string_array():
+    return pa.array(["foo", "bar", "bz", None, "quux"] * 200, type=pa.large_string())
+
+
+@pytest.fixture
+def binary_array():
+    return pa.array([b"foo", b"bar", b"bz", None, b"quux"] * 200)
+
+
+@pytest.fixture
+def fixed_size_binary_array():
+    return pa.array([b"foo", b"bar", b"baz", None, b"qux"] * 200, type=pa.binary(3))
+
+
+@pytest.fixture
+def large_binary_array():
+    return pa.array(
+        [b"foo", b"bar", b"bz", None, b"quux"] * 200, type=pa.large_binary()
+    )
+
+
+@pytest.fixture
+def list_array():
+    return pa.array(([None] + [list(range(9)) + [None]] * 9) * 100)
+
+
+@pytest.fixture
+def large_list_array():
+    # Large list array with nulls
+    return pa.array(
+        ([None] + [list(range(9)) + [None]] * 9) * 100,
+        type=pa.large_list(pa.int64()),
+    )
+
+
+@pytest.fixture
+def fixed_size_list_array():
+    # Fixed size list array
+    return pa.FixedSizeListArray.from_arrays(
+        pa.array((list(range(9)) + [None]) * 1000), 10
+    )
+
+
+@pytest.fixture
+def map_array():
+    return pa.array(
+        [list(zip("abcdefghij", range(10))) for _ in range(1000)],
+        type=pa.map_(pa.string(), pa.int64()),
+    )
+
+
+@pytest.fixture
+def struct_array():
+    # Struct array
+    return pa.array({"a": i} for i in range(1000))
+
+
+@pytest.fixture
+def sparse_union_array():
+    return pa.UnionArray.from_sparse(
+        pa.array([0, 1] * 500, type=pa.int8()),
+        [pa.array(list(range(1000))), pa.array([True, False] * 500)],
+    )
+
+
+@pytest.fixture
+def dense_union_array():
+    return pa.UnionArray.from_dense(
+        pa.array([0, 1] * 500, type=pa.int8()),
+        pa.array(
+            [i if i % 2 == 0 else (i % 3) % 2 for i in range(1000)], type=pa.int32()
+        ),
+        [pa.array(list(range(1000))), pa.array([True, False])],
+    )
+
+
+@pytest.fixture
+def dictionary_array():
+    return pa.DictionaryArray.from_arrays(
+        pa.array((list(range(9)) + [None]) * 100),
+        pa.array(["a", "b", "c", "d", "e", "f", "g", "h", "i"]),
+    )
+
+
+@pytest.fixture
+def tensor_array():
+    return ArrowTensorArray.from_numpy(np.arange(1000 * 4 * 4).reshape((1000, 4, 4)))
+
+
+@pytest.fixture
+def boolean_tensor_array():
+    return ArrowTensorArray.from_numpy(
+        np.array(
+            [True, False, False, True, False, False, True, True] * 2 * 1000
+        ).reshape((1000, 4, 4))
+    )
+
+
+@pytest.fixture
+def variable_shaped_tensor_array():
+    return ArrowVariableShapedTensorArray.from_numpy(
+        np.array(
+            [
+                np.arange(4).reshape((2, 2)),
+                np.arange(4, 13).reshape((3, 3)),
+            ]
+            * 500,
+            dtype=object,
+        ),
+    )
+
+
+@pytest.fixture
+def boolean_variable_shaped_tensor_array():
+    return ArrowVariableShapedTensorArray.from_numpy(
+        np.array(
+            [
+                np.array([[True, False], [False, True]]),
+                np.array(
+                    [
+                        [False, True, False],
+                        [True, True, False],
+                        [False, False, False],
+                    ],
+                ),
+            ]
+            * 500,
+            dtype=object,
+        )
+    )
+
+
+@pytest.fixture
+def list_of_struct_array():
+    return pa.array([{"a": i}, {"a": -i}] for i in range(1000))
+
+
+@pytest.fixture
+def list_of_empty_struct_array():
+    return pa.array([{}, {}] for i in range(1000))
+
+
+@pytest.fixture
+def complex_nested_array():
+    return pa.UnionArray.from_sparse(
+        pa.array([0, 1] * 500, type=pa.int8()),
+        [
+            pa.array(
+                [
+                    {
+                        "a": i % 2 == 0,
+                        "b": i,
+                        "c": "bar",
+                    }
+                    for i in range(1000)
+                ]
+            ),
+            pa.array(
+                [list(zip("abcdefghij", range(10))) for _ in range(1000)],
+                type=pa.map_(pa.string(), pa.int64()),
+            ),
+        ],
+    )
+
+
+@pytest.fixture
+def pickled_objects_array():
+    elements = ["test", 20, False, {"some": "value"}, None, np.zeros((10, 10))]
+    elements *= 1 + 1000 // len(elements)
+    elements = elements[:1000]
+
+    arr = np.array(elements, dtype=object)
+    return ArrowPythonObjectArray.from_objects(arr)
+
+
+pytest_custom_serialization_arrays = [
+    # Null array
+    (lazy_fixture("null_array"), 1.0),
+    # Int array
+    (lazy_fixture("int_array"), 0.1),
+    # Array with nulls
+    (lazy_fixture("int_array_with_nulls"), 0.1),
+    # Float array
+    (lazy_fixture("float_array"), 0.1),
+    # Boolean array
+    # Due to bit-packing, most of the pickle bytes are metadata.
+    (lazy_fixture("boolean_array"), 0.8),
+    # String array
+    (lazy_fixture("string_array"), 0.1),
+    # Large string array
+    (lazy_fixture("large_string_array"), 0.1),
+    # Binary array
+    (lazy_fixture("binary_array"), 0.1),
+    # Fixed size binary array
+    (lazy_fixture("fixed_size_binary_array"), 0.1),
+    # Large binary array
+    (lazy_fixture("large_binary_array"), 0.1),
+    # List array with nulls
+    (lazy_fixture("list_array"), 0.1),
+    # Large list array with nulls
+    (lazy_fixture("large_list_array"), 0.1),
+    # Fixed size list array
+    (lazy_fixture("fixed_size_list_array"), 0.1),
+    # Map array
+    (lazy_fixture("map_array"), 0.1),
+    # Struct array
+    (lazy_fixture("struct_array"), 0.1),
+    # Union array (sparse)
+    (lazy_fixture("sparse_union_array"), 0.1),
+    # Union array (dense)
+    (lazy_fixture("dense_union_array"), 0.1),
+    # Dictionary array
+    (lazy_fixture("dictionary_array"), 0.1),
+    # Tensor extension array
+    (lazy_fixture("tensor_array"), 0.1),
+    # Boolean tensor extension array
+    (lazy_fixture("boolean_tensor_array"), 0.25),
+    # Variable-shaped tensor extension array
+    (lazy_fixture("variable_shaped_tensor_array"), 0.1),
+    # Boolean variable-shaped tensor extension array
+    (lazy_fixture("boolean_variable_shaped_tensor_array"), 0.25),
+    # List of struct array
+    (lazy_fixture("list_of_struct_array"), 0.1),
+    # List of empty struct array
+    (lazy_fixture("list_of_empty_struct_array"), 0.1),
+    # Complex nested array
+    (lazy_fixture("complex_nested_array"), 0.1),
+    # Array of pickled objects
+    (lazy_fixture("pickled_objects_array"), 0.1),
+]
+
+
+@pytest.mark.parametrize("data,cap_mult", pytest_custom_serialization_arrays)
+def test_custom_arrow_data_serializer(ray_start_regular_shared, data, cap_mult):
+    if len(data) == 0:
+        data = pa.table({"a": []})
+    else:
+        data = pa.Table.from_arrays(
+            [data, data, pa.array(range(1000), type=pa.int32())],
+            schema=pa.schema(
+                [
+                    pa.field("arr1", data.type),
+                    pa.field("arr2", data.type),
+                    pa.field("arr3", pa.int32()),
+                ],
+                metadata={b"foo": b"bar"},
+            ),
+        )
+    ray._private.worker.global_worker.get_serialization_context()
+    data.validate()
+    pyarrow_version = get_pyarrow_version()
+    if pyarrow_version >= parse_version("7.0.0"):
+        # get_total_buffer_size API was added in Arrow 7.0.0.
+        buf_size = data.get_total_buffer_size()
+    # Create a zero-copy slice view of data.
+    view = data.slice(10, 10)
+    s_arr = pickle.dumps(data)
+    s_view = pickle.dumps(view)
+    post_slice = pickle.loads(s_view)
+    post_slice.validate()
+    # Check for round-trip equality.
+    assert view.equals(post_slice), post_slice
+    # Check that the slice view was truncated upon serialization.
+    assert len(s_view) <= cap_mult * len(s_arr)
+    for column, pre_column in zip(post_slice.columns, view.columns):
+        # Check that offset was reset on slice.
+        if column.num_chunks > 0:
+            assert column.chunk(0).offset == 0
+        # Check that null count was either properly cached or recomputed.
+        assert column.null_count == pre_column.null_count
+    if pyarrow_version >= parse_version("7.0.0"):
+        # Check that slice buffer only contains slice data.
+        slice_buf_size = post_slice.get_total_buffer_size()
+        if buf_size > 0:
+            assert buf_size / slice_buf_size - len(data) / len(post_slice) < 100
+
+
+def test_custom_arrow_data_serializer_fallback(
+    ray_start_regular_shared, propagate_logs, caplog
+):
+    # Reset serialization fallback set so warning is logged.
+    import ray._private.arrow_serialization as arrow_ser_module
+
+    arrow_ser_module._serialization_fallback_set = set()
+
+    data = pa.table(
+        {
+            "a": pa.UnionArray.from_dense(
+                pa.array([0, 1] * 500, type=pa.int8()),
+                pa.array(
+                    [i if i % 2 == 0 else (i % 3) % 2 for i in range(1000)],
+                    type=pa.int32(),
+                ),
+                [pa.array(list(range(1000))), pa.array([True, False])],
+            )
+        }
+    )
+    cap_mult = 0.1
+    ray._private.worker.global_worker.get_serialization_context()
+    data.validate()
+    pyarrow_version = get_pyarrow_version()
+    if pyarrow_version >= parse_version("7.0.0"):
+        # get_total_buffer_size API was added in Arrow 7.0.0.
+        buf_size = data.get_total_buffer_size()
+    # Create a zero-copy slice view of data.
+    view = data.slice(10, 10)
+    # Confirm that (1) fallback works, and (2) warning is logged.
+    with caplog.at_level(
+        logging.WARNING,
+        logger="ray.data._internal.arrow_serialization",
+    ):
+        s_arr = pickle.dumps(data)
+    assert "Failed to complete optimized serialization" in caplog.text
+    caplog.clear()
+    # Confirm that we only warn once per process.
+    with caplog.at_level(
+        logging.WARNING,
+        logger="ray.data._internal.arrow_serialization",
+    ):
+        s_view = pickle.dumps(view)
+    assert "Failed to complete optimized serialization" not in caplog.text
+    post_slice = pickle.loads(s_view)
+    post_slice.validate()
+    # Check for round-trip equality.
+    assert view.equals(post_slice), post_slice
+    # Check that the slice view was truncated upon serialization.
+    assert len(s_view) <= cap_mult * len(s_arr)
+    for column, pre_column in zip(post_slice.columns, view.columns):
+        # Check that offset was reset on slice.
+        if column.num_chunks > 0:
+            assert column.chunk(0).offset == 0
+        # Check that null count was either properly cached or recomputed.
+        assert column.null_count == pre_column.null_count
+    if pyarrow_version >= parse_version("7.0.0"):
+        # Check that slice buffer only contains slice data.
+        slice_buf_size = post_slice.get_total_buffer_size()
+        if buf_size > 0:
+            assert buf_size / slice_buf_size - len(data) / len(post_slice) < 100
+
+
+def test_arrow_scalar_conversion(ray_start_regular_shared):
+    ds = ray.data.from_items([1])
+
+    def fn(batch: list):
+        return {"id": np.array([1])}
+
+    ds = ds.map_batches(fn)
+    res = ds.take()
+    assert res == [{"id": 1}], res
+
+
+def test_arrow_object_and_array_support(ray_start_regular_shared):
+    obj = types.SimpleNamespace(some_attribute="test")
+
+    def f(batch):
+        batch_size = len(batch["id"])
+        return {
+            "array": np.zeros((batch_size, 32, 32, 3)),
+            "unsupported": [obj] * batch_size,
+        }
+
+    res = ray.data.range(5).map_batches(f, batch_size=None).take(1)
+    assert res[0]["array"].shape == (32, 32, 3)
+    assert np.all(res[0]["array"] == 0)
+    assert res[0]["unsupported"] == obj
+
+
+def test_custom_arrow_data_serializer_parquet_roundtrip(
+    ray_start_regular_shared, tmp_path
+):
+    ray._private.worker.global_worker.get_serialization_context()
+    t = pa.table({"a": list(range(10000000))})
+    pq.write_table(t, f"{tmp_path}/test.parquet")
+    t2 = pq.read_table(f"{tmp_path}/test.parquet")
+    s_t = pickle.dumps(t)
+    s_t2 = pickle.dumps(t2)
+    # Check that the post-Parquet slice view chunks don't cause a serialization blow-up.
+    assert len(s_t2) < 1.1 * len(s_t)
+    # Check for round-trip equality.
+    assert t2.equals(pickle.loads(s_t2))
+
+
+def test_arrow_schema_ipc_serialization(ray_start_regular_shared):
+    """Test that Arrow Schema uses IPC serialization for performance."""
+    from ray._private.arrow_serialization import (
+        _arrow_schema_reduce,
+        _restore_schema_from_ipc,
+    )
+
+    # Verify the reducer is registered
+    ray._private.worker.global_worker.get_serialization_context()
+    assert pa.Schema in pickle.CloudPickler.dispatch
+    assert pickle.CloudPickler.dispatch[pa.Schema] == _arrow_schema_reduce
+
+    # Create a complex schema with various types
+    schema = pa.schema(
+        [
+            pa.field("id", pa.int64()),
+            pa.field("name", pa.string()),
+            pa.field("timestamp", pa.timestamp("us", tz="UTC")),
+            pa.field("tags", pa.list_(pa.string())),
+            pa.field("metadata", pa.map_(pa.string(), pa.string())),
+            pa.field(
+                "nested",
+                pa.struct(
+                    [
+                        pa.field("x", pa.float64()),
+                        pa.field("y", pa.float64()),
+                    ]
+                ),
+            ),
+            pa.field("category", pa.dictionary(pa.int8(), pa.string())),
+            pa.field("decimal_val", pa.decimal128(18, 6)),
+        ],
+        metadata={b"foo": b"bar"},
+    )
+
+    # Test roundtrip serialization
+    serialized = pickle.dumps(schema)
+    deserialized = pickle.loads(serialized)
+    assert schema.equals(deserialized)
+    assert schema.metadata == deserialized.metadata
+
+    # Verify the reducer uses IPC format (check via direct call)
+    restore_func, (ipc_bytes,) = _arrow_schema_reduce(schema)
+    assert restore_func == _restore_schema_from_ipc
+    # IPC bytes should match what schema.serialize() produces
+    assert ipc_bytes == schema.serialize().to_pybytes()
+    # Verify restore works
+    restored = restore_func(ipc_bytes)
+    assert schema.equals(restored)
+
+
+def test_custom_arrow_data_serializer_disable(shutdown_only):
+    ray.shutdown()
+    ray.worker._post_init_hooks = []
+    context = ray.worker.global_worker.get_serialization_context()
+    context._unregister_cloudpickle_reducer(pa.Table)
+    # Disable custom Arrow array serialization.
+    os.environ["RAY_DISABLE_CUSTOM_ARROW_ARRAY_SERIALIZATION"] = "1"
+    ray.init()
+    # Create a zero-copy slice view of table.
+    t = pa.table({"a": list(range(10000000))})
+    view = t.slice(10, 10)
+    s_t = pickle.dumps(t)
+    s_view = pickle.dumps(view)
+    # Check that the slice view contains the full buffer of the underlying array.
+    d_view = pickle.loads(s_view)
+    assert d_view["a"].chunk(0).buffers()[1].size == t["a"].chunk(0).buffers()[1].size
+    # Check that the serialized slice view is large
+    assert len(s_view) > 0.8 * len(s_t)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))