chore: import upstream snapshot with attribution

python/ray/data/tests/test_pandas_block.py

@@ -0,0 +1,540 @@
+import pickle
+import random
+import sys
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+from packaging.version import parse as parse_version
+
+import ray
+import ray.data
+from ray.data._internal.pandas_block import (
+    PandasBlockAccessor,
+    PandasBlockBuilder,
+    PandasBlockColumnAccessor,
+)
+from ray.data._internal.util import is_null
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+from ray.data.context import DataContext
+
+# Set seed for the test for size as it related to sampling
+np.random.seed(42)
+
+
+def simple_series(null):
+    return pd.Series([1, 2, null, 6])
+
+
+@pytest.mark.parametrize("arr", [simple_series(None), simple_series(np.nan)])
+class TestPandasBlockColumnAccessor:
+    @pytest.mark.parametrize(
+        "ignore_nulls, expected",
+        [
+            (True, 3),
+            (False, 4),
+        ],
+    )
+    def test_count(self, arr, ignore_nulls, expected):
+        accessor = PandasBlockColumnAccessor(arr)
+        result = accessor.count(ignore_nulls=ignore_nulls, as_py=True)
+        assert result == expected
+
+    @pytest.mark.parametrize(
+        "ignore_nulls, expected",
+        [
+            (True, 9),
+            (False, np.nan),
+        ],
+    )
+    def test_sum(self, arr, ignore_nulls, expected):
+        accessor = PandasBlockColumnAccessor(arr)
+        result = accessor.sum(ignore_nulls=ignore_nulls, as_py=True)
+        assert result == expected or is_null(result) and is_null(expected)
+
+    @pytest.mark.parametrize(
+        "ignore_nulls, expected",
+        [
+            (True, 1),
+            (False, np.nan),
+        ],
+    )
+    def test_min(self, arr, ignore_nulls, expected):
+        accessor = PandasBlockColumnAccessor(arr)
+        result = accessor.min(ignore_nulls=ignore_nulls, as_py=True)
+        assert result == expected or is_null(result) and is_null(expected)
+
+    @pytest.mark.parametrize(
+        "ignore_nulls, expected",
+        [
+            (True, 6),
+            (False, np.nan),
+        ],
+    )
+    def test_max(self, arr, ignore_nulls, expected):
+        accessor = PandasBlockColumnAccessor(arr)
+        result = accessor.max(ignore_nulls=ignore_nulls, as_py=True)
+        assert result == expected or is_null(result) and is_null(expected)
+
+    @pytest.mark.parametrize(
+        "ignore_nulls, expected",
+        [
+            (True, 3.0),
+            (False, np.nan),
+        ],
+    )
+    def test_mean(self, arr, ignore_nulls, expected):
+        accessor = PandasBlockColumnAccessor(arr)
+        result = accessor.mean(ignore_nulls=ignore_nulls, as_py=True)
+        assert result == expected or is_null(result) and is_null(expected)
+
+    @pytest.mark.parametrize(
+        "provided_mean, expected",
+        [
+            (3.0, 14.0),
+            (None, 14.0),
+        ],
+    )
+    def test_sum_of_squared_diffs_from_mean(self, arr, provided_mean, expected):
+        accessor = PandasBlockColumnAccessor(arr)
+        result = accessor.sum_of_squared_diffs_from_mean(
+            ignore_nulls=True, mean=provided_mean, as_py=True
+        )
+        assert result == expected or is_null(result) and is_null(expected)
+
+    def test_to_pylist(self, arr):
+        accessor = PandasBlockColumnAccessor(arr)
+
+        result = accessor.to_pylist()
+        expected = arr.to_list()
+
+        assert all(
+            [a == b or is_null(a) and is_null(b) for a, b in zip(expected, result)]
+        )
+
+
+class TestPandasBlockColumnAccessorAllNullSeries:
+    @pytest.fixture
+    def all_null_series(self):
+        return pd.Series([None] * 3, dtype=np.float64)
+
+    def test_count_all_null(self, all_null_series):
+        accessor = PandasBlockColumnAccessor(all_null_series)
+        # When ignoring nulls, count should be 0; otherwise, count returns length.
+        assert accessor.count(ignore_nulls=True, as_py=True) == 0
+        assert accessor.count(ignore_nulls=False, as_py=True) == len(all_null_series)
+
+    @pytest.mark.parametrize("ignore_nulls", [True, False])
+    def test_sum_all_null(self, all_null_series, ignore_nulls):
+        accessor = PandasBlockColumnAccessor(all_null_series)
+        result = accessor.sum(ignore_nulls=ignore_nulls)
+        assert is_null(result)
+
+    @pytest.mark.parametrize("ignore_nulls", [True, False])
+    def test_min_all_null(self, all_null_series, ignore_nulls):
+        accessor = PandasBlockColumnAccessor(all_null_series)
+        result = accessor.min(ignore_nulls=ignore_nulls, as_py=True)
+        assert is_null(result)
+
+    @pytest.mark.parametrize("ignore_nulls", [True, False])
+    def test_max_all_null(self, all_null_series, ignore_nulls):
+        accessor = PandasBlockColumnAccessor(all_null_series)
+        result = accessor.max(ignore_nulls=ignore_nulls)
+        assert is_null(result)
+
+    @pytest.mark.parametrize("ignore_nulls", [True, False])
+    def test_mean_all_null(self, all_null_series, ignore_nulls):
+        accessor = PandasBlockColumnAccessor(all_null_series)
+        result = accessor.mean(ignore_nulls=ignore_nulls)
+        assert is_null(result)
+
+    @pytest.mark.parametrize("ignore_nulls", [True, False])
+    def test_sum_of_squared_diffs_all_null(self, all_null_series, ignore_nulls):
+        accessor = PandasBlockColumnAccessor(all_null_series)
+        result = accessor.sum_of_squared_diffs_from_mean(
+            ignore_nulls=ignore_nulls, mean=None
+        )
+        assert is_null(result)
+
+
+@pytest.mark.parametrize(
+    "input_block, fill_column_name, fill_value, expected_output_block",
+    [
+        (
+            pd.DataFrame({"a": [0, 1]}),
+            "b",
+            2,
+            pd.DataFrame({"a": [0, 1], "b": [2, 2]}),
+        ),
+    ],
+)
+def test_fill_column(input_block, fill_column_name, fill_value, expected_output_block):
+    block_accessor = PandasBlockAccessor.for_block(input_block)
+
+    actual_output_block = block_accessor.fill_column(fill_column_name, fill_value)
+
+    assert actual_output_block.equals(expected_output_block)
+
+
+def test_pandas_block_timestamp_ns(ray_start_regular_shared):
+    # Input data with nanosecond precision timestamps
+    data_rows = [
+        {"col1": 1, "col2": pd.Timestamp("2023-01-01T00:00:00.123456789")},
+        {"col1": 2, "col2": pd.Timestamp("2023-01-01T01:15:30.987654321")},
+        {"col1": 3, "col2": pd.Timestamp("2023-01-01T02:30:15.111111111")},
+        {"col1": 4, "col2": pd.Timestamp("2023-01-01T03:45:45.222222222")},
+        {"col1": 5, "col2": pd.Timestamp("2023-01-01T05:00:00.333333333")},
+    ]
+
+    # Initialize PandasBlockBuilder
+    pandas_builder = PandasBlockBuilder()
+    for row in data_rows:
+        pandas_builder.add(row)
+    pandas_block = pandas_builder.build()
+
+    assert pd.api.types.is_datetime64_ns_dtype(pandas_block["col2"])
+
+    for original_row, result_row in zip(
+        data_rows, pandas_block.to_dict(orient="records")
+    ):
+        assert (
+            original_row["col2"] == result_row["col2"]
+        ), "Timestamp mismatch in PandasBlockBuilder output"
+
+
+def test_dict_and_none_use_arrow_block(ray_start_regular_shared, restore_data_context):
+    # Dicts are represented as Arrow struct types, so the block should remain Arrow
+    # even if object-extension fallback is disabled.
+    DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
+
+    def fn(batch):
+        batch["data_dict"] = [{"data": 0} for _ in range(len(batch["id"]))]
+        return batch
+
+    ds = ray.data.range(10).map_batches(fn)
+    ds = ds.materialize()
+    block = ray.get(next(ds.iter_internal_ref_bundles()).block_refs[0])
+    assert isinstance(block, pa.Table)
+    assert block.schema.field("data_dict").type == pa.struct([("data", pa.int64())])
+
+    def fn2(batch):
+        batch["data_none"] = [None for _ in range(len(batch["id"]))]
+        return batch
+
+    ds2 = ray.data.range(10).map_batches(fn2)
+    ds2 = ds2.materialize()
+    block = ray.get(next(ds2.iter_internal_ref_bundles()).block_refs[0])
+    assert isinstance(block, pa.Table)
+
+
+class TestSizeBytes:
+    def test_small(ray_start_regular_shared):
+        animals = ["Flamingo", "Centipede"]
+        block = pd.DataFrame({"animals": animals})
+
+        block_accessor = PandasBlockAccessor.for_block(block)
+        bytes_size = block_accessor.size_bytes()
+
+        # check that memory usage is within 10% of the size_bytes
+        # For strings, Pandas seems to be fairly accurate, so let's use that.
+        memory_usage = block.memory_usage(index=True, deep=True).sum()
+        assert bytes_size == pytest.approx(memory_usage, rel=0.1), (
+            bytes_size,
+            memory_usage,
+        )
+
+    def test_large_str(ray_start_regular_shared):
+        animals = [
+            random.choice(["alligator", "crocodile", "centipede", "flamingo"])
+            for i in range(100_000)
+        ]
+        block = pd.DataFrame({"animals": animals})
+        block["animals"] = block["animals"].astype("string")
+
+        block_accessor = PandasBlockAccessor.for_block(block)
+        bytes_size = block_accessor.size_bytes()
+
+        memory_usage = block.memory_usage(index=True, deep=True).sum()
+        assert bytes_size == pytest.approx(memory_usage, rel=0.1), (
+            bytes_size,
+            memory_usage,
+        )
+
+    def test_large_str_object(ray_start_regular_shared):
+        """Note - this test breaks if you refactor/move the list of animals."""
+        num = 100_000
+        animals = [
+            random.choice(["alligator", "crocodile", "centipede", "flamingo"])
+            for i in range(num)
+        ]
+        block = pd.DataFrame({"animals": animals})
+
+        block_accessor = PandasBlockAccessor.for_block(block)
+        bytes_size = block_accessor.size_bytes()
+
+        # The actual usage should be the index usage + the string data usage.
+        memory_usage = block.memory_usage(index=True, deep=False).sum() + sum(
+            [sys.getsizeof(animal) for animal in animals]
+        )
+
+        assert bytes_size == pytest.approx(memory_usage, rel=0.1), (
+            bytes_size,
+            memory_usage,
+        )
+
+    def test_large_floats(ray_start_regular_shared):
+        animals = [random.random() for i in range(100_000)]
+        block = pd.DataFrame({"animals": animals})
+
+        block_accessor = PandasBlockAccessor.for_block(block)
+        bytes_size = block_accessor.size_bytes()
+
+        memory_usage = pickle.dumps(block).__sizeof__()
+        # check that memory usage is within 10% of the size_bytes
+        assert bytes_size == pytest.approx(memory_usage, rel=0.1), (
+            bytes_size,
+            memory_usage,
+        )
+
+    def test_bytes_object(ray_start_regular_shared):
+        def generate_data(batch):
+            for _ in range(8):
+                yield {"data": [[b"\x00" * 128 * 1024 * 128]]}
+
+        ds = (
+            ray.data.range(1, override_num_blocks=1)
+            .map_batches(generate_data, batch_size=1)
+            .map_batches(lambda batch: batch, batch_format="pandas")
+        )
+
+        true_value = 128 * 1024 * 128 * 8
+        for bundle in ds.iter_internal_ref_bundles():
+            size = bundle.size_bytes()
+            # assert that true_value is within 10% of bundle.size_bytes()
+            assert size == pytest.approx(true_value, rel=0.1), (
+                size,
+                true_value,
+            )
+
+    def test_nested_numpy(ray_start_regular_shared):
+        size = 1024
+        rows = 1_000
+        data = [
+            np.random.randint(size=size, low=0, high=100, dtype=np.int8)
+            for _ in range(rows)
+        ]
+        df = pd.DataFrame({"data": data})
+
+        block_accessor = PandasBlockAccessor.for_block(df)
+        block_size = block_accessor.size_bytes()
+        true_value = rows * size
+        assert block_size == pytest.approx(true_value, rel=0.1), (
+            block_size,
+            true_value,
+        )
+
+    def test_nested_objects(ray_start_regular_shared):
+        size = 10
+        rows = 10_000
+        lists = [[random.randint(0, 100) for _ in range(size)] for _ in range(rows)]
+        data = {"lists": lists}
+        block = pd.DataFrame(data)
+
+        block_accessor = PandasBlockAccessor.for_block(block)
+        bytes_size = block_accessor.size_bytes()
+
+        # List overhead + 10 integers per list
+        true_size = rows * (
+            sys.getsizeof([random.randint(0, 100) for _ in range(size)]) + size * 28
+        )
+
+        assert bytes_size == pytest.approx(true_size, rel=0.1), (
+            bytes_size,
+            true_size,
+        )
+
+    def test_mixed_types(ray_start_regular_shared):
+        rows = 10_000
+
+        data = {
+            "integers": [random.randint(0, 100) for _ in range(rows)],
+            "floats": [random.random() for _ in range(rows)],
+            "strings": [
+                random.choice(["apple", "banana", "cherry"]) for _ in range(rows)
+            ],
+            "object": [b"\x00" * 128 for _ in range(rows)],
+        }
+        block = pd.DataFrame(data)
+        block_accessor = PandasBlockAccessor.for_block(block)
+        bytes_size = block_accessor.size_bytes()
+
+        # Manually calculate the size
+        int_size = rows * 8
+        float_size = rows * 8
+        str_size = sum(sys.getsizeof(string) for string in data["strings"])
+        object_size = rows * sys.getsizeof(b"\x00" * 128)
+
+        true_size = int_size + float_size + str_size + object_size
+        assert bytes_size == pytest.approx(true_size, rel=0.1), (bytes_size, true_size)
+
+    def test_nested_lists_strings(ray_start_regular_shared):
+        rows = 5_000
+        nested_lists = ["a"] * 3 + ["bb"] * 4 + ["ccc"] * 3
+        data = {
+            "nested_lists": [nested_lists for _ in range(rows)],
+        }
+        block = pd.DataFrame(data)
+        block_accessor = PandasBlockAccessor.for_block(block)
+        bytes_size = block_accessor.size_bytes()
+
+        # Manually calculate the size
+        list_overhead = sys.getsizeof(block["nested_lists"].iloc[0]) + sum(
+            [sys.getsizeof(x) for x in nested_lists]
+        )
+        true_size = rows * list_overhead
+        assert bytes_size == pytest.approx(true_size, rel=0.1), (bytes_size, true_size)
+
+    @pytest.mark.parametrize("size", [10, 1024])
+    def test_multi_level_nesting(ray_start_regular_shared, size):
+        rows = 1_000
+        data = {
+            "complex": [
+                {"list": [np.random.rand(size)], "value": {"key": "val"}}
+                for _ in range(rows)
+            ],
+        }
+        block = pd.DataFrame(data)
+        block_accessor = PandasBlockAccessor.for_block(block)
+        bytes_size = block_accessor.size_bytes()
+
+        numpy_size = np.random.rand(size).nbytes
+
+        values = ["list", "value", "key", "val"]
+        str_size = sum([sys.getsizeof(v) for v in values])
+
+        list_ref_overhead = sys.getsizeof([np.random.rand(size)])
+
+        dict_overhead1 = sys.getsizeof({"key": "val"})
+
+        dict_overhead3 = sys.getsizeof(
+            {"list": [np.random.rand(size)], "value": {"key": "val"}}
+        )
+
+        true_size = (
+            numpy_size + str_size + list_ref_overhead + dict_overhead1 + dict_overhead3
+        ) * rows
+        assert bytes_size == pytest.approx(true_size, rel=0.15), (
+            bytes_size,
+            true_size,
+        )
+
+    def test_boolean(ray_start_regular_shared):
+        data = [random.choice([True, False, None]) for _ in range(100_000)]
+        block = pd.DataFrame({"flags": pd.Series(data, dtype="boolean")})
+        block_accessor = PandasBlockAccessor.for_block(block)
+        bytes_size = block_accessor.size_bytes()
+
+        # No object case
+        true_size = block.memory_usage(index=True, deep=True).sum()
+        assert bytes_size == pytest.approx(true_size, rel=0.1), (bytes_size, true_size)
+
+    @pytest.mark.skipif(
+        get_pyarrow_version() < parse_version("10.0.1"),
+        reason="ArrowDtype requires pyarrow>=10.0.1",
+    )
+    def test_arrow(ray_start_regular_shared):
+        data = [
+            random.choice(["alligator", "crocodile", "flamingo"]) for _ in range(50_000)
+        ]
+        arrow_dtype = pd.ArrowDtype(pa.string())
+        block = pd.DataFrame({"animals": pd.Series(data, dtype=arrow_dtype)})
+        block_accessor = PandasBlockAccessor.for_block(block)
+        bytes_size = block_accessor.size_bytes()
+
+        true_size = block.memory_usage(index=True, deep=False).sum() + sum(
+            sys.getsizeof(x) for x in data
+        )
+        assert bytes_size == pytest.approx(true_size, rel=0.1), (bytes_size, true_size)
+
+    def test_deterministic_across_blocks(self):
+        """size_bytes() must return the same value for two blocks holding
+        identical data. Non-determinism here can cause a streaming generator
+        task to produce different block counts across replay attempts (e.g.
+        lineage reconstruction), since each attempt rebuilds the block and
+        re-estimates its size. That surfaces as a silent hang or silent data
+        loss downstream.
+        """
+        # Use enough rows to trigger sampling (sample_size < total_size), and
+        # vary the string lengths so a different sample yields a different
+        # estimate (this is what makes the non-deterministic case observable).
+        data = [f"str_{i}" for i in range(10_000)]
+        block1 = pd.DataFrame({"col": pd.Series(data, dtype="string")})
+        block2 = pd.DataFrame({"col": pd.Series(data, dtype="string")})
+
+        first = PandasBlockAccessor.for_block(block1).size_bytes()
+        second = PandasBlockAccessor.for_block(block2).size_bytes()
+
+        assert (
+            first == second
+        ), f"size_bytes() is non-deterministic: first={first}, second={second}"
+
+
+def test_iter_rows_with_na(ray_start_regular_shared):
+    block = pd.DataFrame({"col": [pd.NA]})
+    block_accessor = PandasBlockAccessor.for_block(block)
+
+    rows = block_accessor.iter_rows(public_row_format=True)
+
+    # We should return None for NaN values.
+    assert list(rows) == [{"col": None}]
+
+
+def test_empty_dataframe_with_object_columns(ray_start_regular_shared):
+    """Test that size_bytes handles empty DataFrames with object/string columns.
+
+    The warning log:
+    "Error calculating size for column 'parent': cannot call `vectorize`
+    on size 0 inputs unless `otypes` is set"
+    should not be logged in the presence of empty columns.
+    """
+    from unittest.mock import patch
+
+    # Create an empty DataFrame but with defined columns and dtypes
+    block = pd.DataFrame(
+        {
+            "parent": pd.Series([], dtype=object),
+            "child": pd.Series([], dtype="string"),
+            "data": pd.Series([], dtype=object),
+        }
+    )
+
+    block_accessor = PandasBlockAccessor.for_block(block)
+
+    # Check that NO warning is logged after calling size_bytes
+    with patch("ray.data._internal.pandas_block.logger.warning") as mock_warning:
+        bytes_size = block_accessor.size_bytes()
+        mock_warning.assert_not_called()
+
+    assert bytes_size >= 0
+
+
+def test_tensor_column_with_all_nan_preserves_type(ray_start_regular_shared):
+    from ray.data.extensions import ArrowTensorType, ArrowTensorTypeV2, TensorArray
+
+    # Create a DataFrame with all-NaN tensor column
+    df = pd.DataFrame(
+        {"foo": TensorArray([np.array([np.nan, np.nan]), np.array([np.nan, np.nan])])}
+    )
+
+    block_accessor = PandasBlockAccessor.for_block(df)
+    arrow_table = block_accessor.to_arrow()
+
+    # The column should preserve tensor type
+    assert isinstance(
+        arrow_table.schema.field("foo").type, (ArrowTensorType, ArrowTensorTypeV2)
+    ), "TensorDtype column with all-NaN values should preserve tensor type"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))