chore: import upstream snapshot with attribution

python/ray/data/tests/datasource/test_pandas.py

@@ -0,0 +1,223 @@
+from typing import Iterator
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
+from ray.data._internal.tensor_extensions.arrow import (
+    ArrowTensorArray,
+    get_arrow_extension_fixed_shape_tensor_types,
+)
+from ray.data.block import Block
+from ray.data.extensions import TensorDtype
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+from ray.types import ObjectRef
+
+
+def _get_first_block(bundles: Iterator[RefBundle]) -> ObjectRef[Block]:
+    return next(bundles).block_refs[0]
+
+
+@pytest.mark.parametrize("enable_pandas_block", [False, True])
+def test_from_pandas(ray_start_regular_shared, enable_pandas_block):
+    ctx = ray.data.context.DataContext.get_current()
+    old_enable_pandas_block = ctx.enable_pandas_block
+    ctx.enable_pandas_block = enable_pandas_block
+    try:
+        df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+        df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
+        ds = ray.data.from_pandas([df1, df2])
+        block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
+        assert (
+            isinstance(block, pd.DataFrame)
+            if enable_pandas_block
+            else isinstance(block, pa.Table)
+        )
+        values = [(r["one"], r["two"]) for r in ds.take(6)]
+        rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
+        assert values == rows
+        # Check that metadata fetch is included in stats.
+        assert "FromPandas" in ds.stats()
+
+        # test from single pandas dataframe
+        ds = ray.data.from_pandas(df1)
+        block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
+        assert (
+            isinstance(block, pd.DataFrame)
+            if enable_pandas_block
+            else isinstance(block, pa.Table)
+        )
+        values = [(r["one"], r["two"]) for r in ds.take(3)]
+        rows = [(r.one, r.two) for _, r in df1.iterrows()]
+        assert values == rows
+        # Check that metadata fetch is included in stats.
+        assert "FromPandas" in ds.stats()
+    finally:
+        ctx.enable_pandas_block = old_enable_pandas_block
+
+
+@pytest.mark.parametrize("num_inputs", [1, 2])
+def test_from_pandas_override_num_blocks(num_inputs, ray_start_regular_shared):
+    df = pd.DataFrame({"number": [0]})
+
+    ds = ray.data.from_pandas([df] * num_inputs, override_num_blocks=2)
+
+    assert ds.materialize().num_blocks() == 2
+
+
+@pytest.mark.parametrize("enable_pandas_block", [False, True])
+def test_from_pandas_refs(ray_start_regular_shared, enable_pandas_block):
+    ctx = ray.data.context.DataContext.get_current()
+    old_enable_pandas_block = ctx.enable_pandas_block
+    ctx.enable_pandas_block = enable_pandas_block
+    try:
+        df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+        df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
+        ds = ray.data.from_pandas_refs([ray.put(df1), ray.put(df2)])
+        block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
+        assert (
+            isinstance(block, pd.DataFrame)
+            if enable_pandas_block
+            else isinstance(block, pa.Table)
+        )
+        values = [(r["one"], r["two"]) for r in ds.take(6)]
+        rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
+        assert values == rows
+        # Check that metadata fetch is included in stats.
+        assert "FromPandas" in ds.stats()
+
+        # test from single pandas dataframe ref
+        ds = ray.data.from_pandas_refs(ray.put(df1))
+        block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
+        assert (
+            isinstance(block, pd.DataFrame)
+            if enable_pandas_block
+            else isinstance(block, pa.Table)
+        )
+        values = [(r["one"], r["two"]) for r in ds.take(3)]
+        rows = [(r.one, r.two) for _, r in df1.iterrows()]
+        assert values == rows
+        # Check that metadata fetch is included in stats.
+        assert "FromPandas" in ds.stats()
+    finally:
+        ctx.enable_pandas_block = old_enable_pandas_block
+
+
+def test_to_pandas(ray_start_regular_shared):
+    n = 5
+    df = pd.DataFrame({"id": list(range(n))})
+    ds = ray.data.range(n)
+    dfds = ds.to_pandas()
+    pd.testing.assert_frame_equal(df.astype(dfds.dtypes.to_dict()), dfds)
+
+    # Test limit.
+    with pytest.raises(ValueError):
+        dfds = ds.to_pandas(limit=3)
+
+    # Test limit greater than number of rows.
+    dfds = ds.to_pandas(limit=6)
+    pd.testing.assert_frame_equal(df.astype(dfds.dtypes.to_dict()), dfds)
+
+
+def test_to_pandas_different_block_types(ray_start_regular_shared):
+    # Test for https://github.com/ray-project/ray/issues/48575.
+    df = pd.DataFrame({"a": [0]})
+    ds1 = ray.data.from_pandas(df)
+
+    table = pa.Table.from_pandas(df)
+    ds2 = ray.data.from_arrow(table)
+
+    actual_df = ds1.union(ds2).to_pandas()
+
+    expected_df = pd.DataFrame({"a": [0, 0]}).astype(actual_df.dtypes.to_dict())
+    pd.testing.assert_frame_equal(actual_df, expected_df)
+
+
+def test_to_pandas_refs(ray_start_regular_shared):
+    n = 5
+    df = pd.DataFrame({"id": list(range(n))})
+    ds = ray.data.range(n)
+    dfds = pd.concat(ray.get(ds.to_pandas_refs()), ignore_index=True)
+    pd.testing.assert_frame_equal(df.astype(dfds.dtypes.to_dict()), dfds)
+
+
+def test_pandas_roundtrip(ray_start_regular_shared, tmp_path):
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
+    ds = ray.data.from_pandas([df1, df2], override_num_blocks=2)
+    dfds = ds.to_pandas()
+    expected = pd.concat([df1, df2], ignore_index=True)
+    pd.testing.assert_frame_equal(expected.astype(dfds.dtypes.to_dict()), dfds)
+
+
+def test_to_pandas_tensor_column_cast_pandas(ray_start_regular_shared):
+    # Check that tensor column casting occurs when converting a Dataset to a Pandas
+    # DataFrame.
+    data = np.arange(12).reshape((3, 2, 2))
+    ctx = ray.data.context.DataContext.get_current()
+    original = ctx.enable_tensor_extension_casting
+    try:
+        ctx.enable_tensor_extension_casting = True
+        in_df = pd.DataFrame({"a": [data]})
+        ds = ray.data.from_pandas(in_df)
+        dtypes = ds.schema().base_schema.types
+        assert len(dtypes) == 1
+        # Tensor column should be automatically cast to Tensor extension.
+        assert isinstance(dtypes[0], TensorDtype)
+        # Original df should not be changed.
+        assert not isinstance(in_df.dtypes[0], TensorDtype)
+        out_df = ds.to_pandas()
+        # Column should be cast back to object dtype when returning back to user.
+        assert out_df["a"].dtype.type is np.object_
+        expected_df = pd.DataFrame({"a": [data]})
+        pd.testing.assert_frame_equal(out_df, expected_df)
+    finally:
+        ctx.enable_tensor_extension_casting = original
+
+
+def test_to_pandas_tensor_column_cast_arrow(ray_start_regular_shared):
+    # Check that tensor column casting occurs when converting a Dataset to a Pandas
+    # DataFrame.
+    data = np.arange(12).reshape((3, 2, 2))
+    ctx = ray.data.context.DataContext.get_current()
+    original = ctx.enable_tensor_extension_casting
+    try:
+        ctx.enable_tensor_extension_casting = True
+        in_table = pa.table({"a": ArrowTensorArray.from_numpy(data)})
+        ds = ray.data.from_arrow(in_table)
+        dtype = ds.schema().base_schema.field(0).type
+        assert isinstance(dtype, get_arrow_extension_fixed_shape_tensor_types())
+        out_df = ds.to_pandas()
+        assert out_df["a"].dtype.type is np.object_
+        expected_df = pd.DataFrame({"a": list(data)})
+        pd.testing.assert_frame_equal(out_df, expected_df)
+    finally:
+        ctx.enable_tensor_extension_casting = original
+
+
+def test_read_pandas_data_array_column(ray_start_regular_shared):
+    df = pd.DataFrame(
+        {
+            "one": [1, 2, 3],
+            "array": [
+                np.array([1, 1, 1]),
+                np.array([2, 2, 2]),
+                np.array([3, 3, 3]),
+            ],
+        }
+    )
+    ds = ray.data.from_pandas(df)
+    row = ds.take(1)[0]
+    assert row["one"] == 1
+    assert all(row["array"] == [1, 1, 1])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))