chore: import upstream snapshot with attribution

python/ray/data/tests/test_ecosystem_dask.py

@@ -0,0 +1,199 @@
+import sys
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.tensor_extensions.arrow import (
+    get_arrow_extension_fixed_shape_tensor_types,
+)
+from ray.data.extensions.tensor_extension import (
+    ArrowTensorArray,
+    TensorArray,
+    TensorDtype,
+)
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.test_util import _check_usage_record
+from ray.tests.conftest import *  # noqa
+
+
+def test_from_dask(ray_start_regular_shared):
+    import dask.dataframe as dd
+
+    df = pd.DataFrame({"one": list(range(100)), "two": list(range(100))})
+    ddf = dd.from_pandas(df, npartitions=10)
+    ds = ray.data.from_dask(ddf)
+    dfds = ds.to_pandas()
+    assert df.equals(dfds)
+
+
+def test_from_dask_e2e(ray_start_regular_shared):
+    import dask.dataframe as dd
+
+    df = pd.DataFrame({"one": list(range(100)), "two": list(range(100))})
+    ddf = dd.from_pandas(df, npartitions=10)
+    ds = ray.data.from_dask(ddf)
+    # `ds.take_all()` triggers execution with new backend, which is
+    # needed for checking operator usage below.
+    assert len(ds.take_all()) == len(df)
+    dfds = ds.to_pandas()
+    assert df.equals(dfds)
+
+    # Underlying implementation uses `FromPandas` operator
+    assert "FromPandas" in ds.stats()
+    assert ds._logical_plan.dag.name == "FromPandas"
+    _check_usage_record(["FromPandas"])
+
+
+def test_to_dask_simple(ray_start_regular_shared):
+    ds = ray.data.range(100)
+    assert ds.to_dask().sum().compute()[0] == 4950
+
+
+@pytest.mark.parametrize("ds_format", ["pandas", "arrow"])
+def test_to_dask(ray_start_regular_shared, ds_format):
+    # Since 2023.7.1, Dask DataFrame automatically converts text data using object data types to string[pyarrow]
+    # For the purpose of this test, we need to disable this behavior.
+    import dask
+
+    dask.config.set({"dataframe.convert-string": False})
+
+    from ray.util.dask import ray_dask_get
+
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
+    df = pd.concat([df1, df2])
+    ds = ray.data.from_blocks([df1, df2])
+    if ds_format == "arrow":
+        ds = ds.map_batches(lambda df: df, batch_format="pyarrow", batch_size=None)
+        # ArrowBlockAccessor.to_pandas() preserves Arrow dtypes via types_mapper,
+        # so the dask partitions (and inferred meta) come back Arrow-backed.
+        expected_dtypes = [pd.ArrowDtype(pa.int64()), pd.ArrowDtype(pa.string())]
+        df = df.astype(dict(zip(df.columns, expected_dtypes)))
+    else:
+        expected_dtypes = [np.int64, object]
+    ddf = ds.to_dask()
+    meta = ddf._meta
+    # Check metadata.
+    assert isinstance(meta, pd.DataFrame)
+    assert meta.empty
+    assert list(meta.columns) == ["one", "two"]
+    assert list(meta.dtypes) == expected_dtypes
+    # Explicit Dask-on-Ray
+    assert df.equals(ddf.compute(scheduler=ray_dask_get))
+    # Implicit Dask-on-Ray.
+    assert df.equals(ddf.compute())
+
+    # Explicit metadata.
+    df1["two"] = df1["two"].astype(pd.StringDtype())
+    df2["two"] = df2["two"].astype(pd.StringDtype())
+    df = pd.concat([df1, df2])
+    ds = ray.data.from_blocks([df1, df2])
+    if ds_format == "arrow":
+        ds = ds.map_batches(lambda df: df, batch_format="pyarrow", batch_size=None)
+        # After Arrow round-trip both columns come back Arrow-backed via
+        # types_mapper, so the expected df must match that.
+        two_meta_dtype = pd.ArrowDtype(pa.string())
+        df = df.astype({"one": pd.ArrowDtype(pa.int64()), "two": two_meta_dtype})
+    else:
+        two_meta_dtype = pd.StringDtype()
+    ddf = ds.to_dask(
+        meta=pd.DataFrame(
+            {"one": pd.Series(dtype=np.int16), "two": pd.Series(dtype=two_meta_dtype)}
+        ),
+    )
+
+    meta = ddf._meta
+    # Check metadata.
+    assert isinstance(meta, pd.DataFrame)
+    assert meta.empty
+    assert list(meta.columns) == ["one", "two"]
+    assert list(meta.dtypes) == [np.int16, two_meta_dtype]
+
+    # Explicit Dask-on-Ray
+    result = ddf.compute(scheduler=ray_dask_get)
+
+    print("Expected: ", df)
+    print("Result: ", result)
+
+    pd.testing.assert_frame_equal(df, result)
+
+    # Implicit Dask-on-Ray.
+    pd.testing.assert_frame_equal(df, ddf.compute())
+
+    # Test case with blocks which have different schema, where we must
+    # skip the metadata check in order to avoid a Dask metadata mismatch error.
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    df2 = pd.DataFrame({"three": [4, 5, 6], "four": ["e", "f", "g"]})
+    df = pd.concat([df1, df2])
+    ds = ray.data.from_blocks([df1, df2])
+    if ds_format == "arrow":
+        ds = ds.map_batches(lambda df: df, batch_format="pyarrow", batch_size=None)
+
+    ddf = ds.to_dask(verify_meta=False)
+
+    # Explicit Dask-on-Ray
+    result = ddf.compute(scheduler=ray_dask_get)
+
+    print("Expected: ", df)
+    print("Result (1): ", result)
+
+    if ds_format == "arrow":
+        df = df.astype(result.dtypes.to_dict())
+    pd.testing.assert_frame_equal(df, result)
+
+    # Implicit Dask-on-Ray.
+    result = ddf.compute()
+
+    print("Result (2): ", result)
+
+    if ds_format == "arrow":
+        df = df.astype(result.dtypes.to_dict())
+    pd.testing.assert_frame_equal(df, result)
+
+
+def test_to_dask_tensor_column_cast_pandas(ray_start_regular_shared):
+    # Check that tensor column casting occurs when converting a Dataset to a Dask
+    # 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": TensorArray(data)})
+        ds = ray.data.from_pandas(in_df)
+        dtypes = ds.schema().base_schema.types
+        assert len(dtypes) == 1
+        assert isinstance(dtypes[0], TensorDtype)
+        out_df = ds.to_dask().compute()
+        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_to_dask_tensor_column_cast_arrow(ray_start_regular_shared):
+    # Check that tensor column casting occurs when converting a Dataset to a Dask
+    # 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_dask().compute()
+        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
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))