[data] Integrate Ray Dataset with Daft Dataframe (#51531)

jaychia · richardliaw · web-flow · commit ff1ecbe7d01f · 2025-03-26T11:30:53.000-07:00
## Why are these changes needed? This adds nice integrations between Daft and Ray Data along with: Unit test for data round-trip (including more interesting types such as a Tensor) Documentation Implementation on Ray Dataset for conversion to/from Daft Dataframes ## Related issue number Re-opening #50630 which was borked by a bad rebase. ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [x] I've included any doc changes needed for https://docs.ray.io/en/master/. - [x] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Jay Chia <17691182+jaychia@users.noreply.github.com> Signed-off-by: Richard Liaw <rliaw@berkeley.edu> Co-authored-by: Richard Liaw <rliaw@berkeley.edu>
diff --git a/doc/source/data/api/input_output.rst b/doc/source/data/api/input_output.rst
@@ -233,6 +233,16 @@ ClickHouse
 
    read_clickhouse
 
+Daft
+----
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_daft
+   Dataset.to_daft
+
 Dask
 ----
 
diff --git a/doc/source/data/comparisons.rst b/doc/source/data/comparisons.rst
@@ -19,11 +19,11 @@ How does Ray Data compare to other solutions for offline inference?
 
     For a more detailed performance comparison between Ray Data and Sagemaker Batch Transform, see `Offline Batch Inference: Comparing Ray, Apache Spark, and SageMaker <https://www.anyscale.com/blog/offline-batch-inference-comparing-ray-apache-spark-and-sagemaker>`_.
 
-.. dropdown:: Distributed Data Processing Frameworks: Apache Spark
+.. dropdown:: Distributed Data Processing Frameworks: Apache Spark and Daft
 
-    Ray Data handles many of the same batch processing workloads as `Apache Spark <https://spark.apache.org/>`_, but with a streaming paradigm that is better suited for GPU workloads for deep learning inference.
+    Ray Data handles many of the same batch processing workloads as `Apache Spark <https://spark.apache.org/>`_ and `Daft <https://www.getdaft.io>`_, but with a streaming paradigm that is better suited for GPU workloads for deep learning inference.
 
-    Ray Data doesn't have a SQL interface and isn't meant as a replacement for generic ETL pipelines like Spark.
+    However, Ray Data doesn't have a SQL interface and isn't meant as a replacement for generic ETL pipelines like Spark and Daft.
 
     For a more detailed performance comparison between Ray Data and Apache Spark, see `Offline Batch Inference: Comparing Ray, Apache Spark, and SageMaker <https://www.anyscale.com/blog/offline-batch-inference-comparing-ray-apache-spark-and-sagemaker>`_.
 
diff --git a/doc/source/data/loading-data.rst b/doc/source/data/loading-data.rst
@@ -469,7 +469,7 @@ Ray Data interoperates with libraries like pandas, NumPy, and Arrow.
 Loading data from distributed DataFrame libraries
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Ray Data interoperates with distributed data processing frameworks like
+Ray Data interoperates with distributed data processing frameworks like `Daft <https://www.getdaft.io>`_,
 :ref:`Dask <dask-on-ray>`, :ref:`Spark <spark-on-ray>`, :ref:`Modin <modin-on-ray>`, and
 :ref:`Mars <mars-on-ray>`.
 
@@ -480,6 +480,30 @@ Ray Data interoperates with distributed data processing frameworks like
 
 .. tab-set::
 
+    .. tab-item:: Daft
+
+        To create a :class:`~ray.data.dataset.Dataset` from a `Daft DataFrame <https://www.getdaft.io/projects/docs/en/stable/api_docs/dataframe.html>`_, call
+        :func:`~ray.data.from_daft`. This function executes the Daft dataframe and constructs a ``Dataset`` backed by the resultant arrow data produced
+        by your Daft query.
+
+        .. testcode::
+
+            import daft
+            import ray
+
+            ray.init()
+
+            df = daft.from_pydict({"int_col": [i for i in range(10000)], "str_col": [str(i) for i in range(10000)]})
+            ds = ray.data.from_daft(df)
+
+            ds.show(3)
+
+        .. testoutput::
+
+            {'int_col': 0, 'str_col': '0'}
+            {'int_col': 1, 'str_col': '1'}
+            {'int_col': 2, 'str_col': '2'}
+
     .. tab-item:: Dask
 
         To create a :class:`~ray.data.dataset.Dataset` from a
diff --git a/doc/source/data/saving-data.rst b/doc/source/data/saving-data.rst
@@ -208,12 +208,25 @@ on the head node.
 Converting Datasets to distributed DataFrames
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Ray Data interoperates with distributed data processing frameworks like
+Ray Data interoperates with distributed data processing frameworks like `Daft <https://www.getdaft.io>`_,
 :ref:`Dask <dask-on-ray>`, :ref:`Spark <spark-on-ray>`, :ref:`Modin <modin-on-ray>`, and
 :ref:`Mars <mars-on-ray>`.
 
 .. tab-set::
 
+    .. tab-item:: Daft
+
+        To convert a :class:`~ray.data.dataset.Dataset` to a `Daft Dataframe <https://www.getdaft.io/projects/docs/en/stable/api_docs/dataframe.html>`_, call
+        :meth:`Dataset.to_daft() <ray.data.Dataset.to_daft>`.
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+            df = ds.to_daft()
+
     .. tab-item:: Dask
 
         To convert a :class:`~ray.data.dataset.Dataset` to a
@@ -228,6 +241,35 @@ Ray Data interoperates with distributed data processing frameworks like
 
             df = ds.to_dask()
 
+            df
+
+        .. testoutput::
+
+            ╭───────────────────┬──────────────────┬───────────────────┬──────────────────┬────────╮
+            │ sepal length (cm) ┆ sepal width (cm) ┆ petal length (cm) ┆ petal width (cm) ┆ target │
+            │ ---               ┆ ---              ┆ ---               ┆ ---              ┆ ---    │
+            │ Float64           ┆ Float64          ┆ Float64           ┆ Float64          ┆ Int64  │
+            ╞═══════════════════╪══════════════════╪═══════════════════╪══════════════════╪════════╡
+            │ 5.1               ┆ 3.5              ┆ 1.4               ┆ 0.2              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 4.9               ┆ 3                ┆ 1.4               ┆ 0.2              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 4.7               ┆ 3.2              ┆ 1.3               ┆ 0.2              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 4.6               ┆ 3.1              ┆ 1.5               ┆ 0.2              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 5                 ┆ 3.6              ┆ 1.4               ┆ 0.2              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 5.4               ┆ 3.9              ┆ 1.7               ┆ 0.4              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 4.6               ┆ 3.4              ┆ 1.4               ┆ 0.3              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 5                 ┆ 3.4              ┆ 1.5               ┆ 0.2              ┆ 0      │
+            ╰───────────────────┴──────────────────┴───────────────────┴──────────────────┴────────╯
+
+            (Showing first 8 of 150 rows)
+
+
     .. tab-item:: Spark
 
         To convert a :class:`~ray.data.dataset.Dataset` to a `Spark DataFrame
diff --git a/doc/source/images/daft.webp b/doc/source/images/daft.webp
diff --git a/doc/source/ray-more-libs/index.rst b/doc/source/ray-more-libs/index.rst
@@ -27,6 +27,7 @@ Ray has a variety of additional integrations with ecosystem libraries.
 - :ref:`spark-on-ray`
 - :ref:`mars-on-ray`
 - :ref:`modin-on-ray`
+- `daft <https://www.getdaft.io>`_
 
 .. _air-ecosystem-map:
 
diff --git a/doc/source/ray-overview/ray-libraries.rst b/doc/source/ray-overview/ray-libraries.rst
@@ -81,6 +81,27 @@ more information.
             Classy Vision Integration
 
 
+    .. grid-item-card::
+
+        .. figure:: ../images/daft.webp
+            :class: card-figure
+
+        .. div::
+
+            .. image:: https://img.shields.io/github/stars/daft/daft?style=social)]
+                :target: https://github.com/Eventual-Inc/Daft
+
+            Daft is a data engine that supports SQL and Python DataFrames for data processing and analytics natively on your Ray clusters.
+
+        +++
+        .. button-link:: https://www.getdaft.io
+            :color: primary
+            :outline:
+            :expand:
+
+            Daft Integration
+
+
     .. grid-item-card::
 
         .. figure:: ../images/dask.png
diff --git a/python/ray/data/BUILD b/python/ray/data/BUILD
@@ -1265,3 +1265,17 @@ py_test(
         "//:ray_lib",
     ],
 )
+
+py_test(
+    name = "test_daft",
+    size = "small",
+    srcs = ["tests/test_daft.py"],
+    tags = [
+        "exclusive",
+        "team:data",
+    ],
+    deps = [
+        ":conftest",
+        "//:ray_lib",
+    ],
+)
diff --git a/python/ray/data/__init__.py b/python/ray/data/__init__.py
@@ -29,6 +29,7 @@
     from_arrow,
     from_arrow_refs,
     from_blocks,
+    from_daft,
     from_dask,
     from_huggingface,
     from_items,
@@ -124,6 +125,7 @@
     "ReadTask",
     "RowBasedFileDatasink",
     "Schema",
+    "from_daft",
     "from_dask",
     "from_items",
     "from_arrow",
diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py
@@ -113,6 +113,7 @@
 from ray.widgets.util import repr_with_fallback
 
 if TYPE_CHECKING:
+    import daft
     import dask
     import mars
     import modin
@@ -4775,6 +4776,24 @@ def to_tf(
             additional_type_spec=additional_type_spec,
         )
 
+    @ConsumptionAPI(pattern="Time complexity:")
+    @PublicAPI(api_group=IOC_API_GROUP)
+    def to_daft(self) -> "daft.DataFrame":
+        """Convert this :class:`~ray.data.Dataset` into a
+        `Daft DataFrame <https://www.getdaft.io/projects/docs/en/stable/api_docs/dataframe.html>`_.
+
+        This will convert all the data inside the Ray Dataset into a Daft DataFrame in a zero-copy way
+        (using Arrow as the intermediate data format).
+
+        Time complexity: O(dataset size / parallelism)
+
+        Returns:
+            A `Daft DataFrame`_ created from this dataset.
+        """
+        import daft
+
+        return daft.from_ray_dataset(self)
+
     @ConsumptionAPI(pattern="Time complexity:")
     @PublicAPI(api_group=IOC_API_GROUP)
     def to_dask(
diff --git a/python/ray/data/read_api.py b/python/ray/data/read_api.py
@@ -93,6 +93,7 @@
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 
 if TYPE_CHECKING:
+    import daft
     import dask
     import datasets
     import mars
@@ -2583,6 +2584,22 @@ def read_hudi(
     )
 
 
+@PublicAPI
+def from_daft(df: "daft.DataFrame") -> Dataset:
+    """Create a :class:`~ray.data.Dataset` from a `Daft DataFrame <https://www.getdaft.io/projects/docs/en/stable/api_docs/dataframe.html>`_.
+
+    Args:
+        df: A Daft DataFrame
+
+    Returns:
+        A :class:`~ray.data.Dataset` holding rows read from the DataFrame.
+    """
+    # NOTE: Today this returns a MaterializedDataset. We should also integrate Daft such that we can stream object references into a Ray
+    # dataset. Unfortunately this is very tricky today because of the way Ray Datasources are implemented with a fully-materialized `list`
+    # of ReadTasks, rather than an iterator which can lazily return these tasks.
+    return df.to_ray_dataset()
+
+
 @PublicAPI
 def from_dask(df: "dask.dataframe.DataFrame") -> MaterializedDataset:
     """Create a :class:`~ray.data.Dataset` from a
diff --git a/python/ray/data/tests/test_daft.py b/python/ray/data/tests/test_daft.py
@@ -0,0 +1,52 @@
+import sys
+
+import pytest
+
+import ray
+
+
+@pytest.fixture(scope="module")
+def ray_start(request):
+    try:
+        yield ray.init(num_cpus=16)
+    finally:
+        ray.shutdown()
+
+
+def test_daft_round_trip(ray_start):
+    import pandas as pd
+    import numpy as np
+    import daft
+
+    data = {
+        "int_col": list(range(128)),
+        "str_col": [str(i) for i in range(128)],
+        "nested_list_col": [[i] * 3 for i in range(128)],
+        "tensor_col": [np.array([[i] * 3] * 3) for i in range(128)],
+    }
+    df = daft.from_pydict(data)
+    ds = ray.data.from_daft(df)
+    pd.testing.assert_frame_equal(ds.to_pandas(), df.to_pandas())
+
+    df2 = ds.to_daft()
+    df_pandas = df.to_pandas()
+    df2_pandas = df2.to_pandas()
+
+    for c in data.keys():
+        # NOTE: tensor behavior on round-trip is different because Ray Data provides
+        # Daft with more information about a column being a fixed-shape-tensor.
+        #
+        # Hence the Pandas representation of `df1` is "just" an object column, but
+        # `df2` knows that this is actually a numpy fixed shaped tensor column
+        if c == "tensor_col":
+            np.testing.assert_equal(
+                np.array(list(df_pandas[c])), df2_pandas[c].to_numpy()
+            )
+        else:
+            pd.testing.assert_series_equal(df_pandas[c], df2_pandas[c])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))
diff --git a/python/requirements/ml/data-requirements.txt b/python/requirements/ml/data-requirements.txt
@@ -1,6 +1,7 @@
 # Used by CI for datasets and docs.
 # https://github.com/ray-project/ray/pull/29448#discussion_r1006256498
 
+getdaft==0.4.3
 dask[complete]==2022.10.2; python_version < '3.12'
 distributed==2022.10.2; python_version < '3.12'
 dask[complete]==2024.6.0; python_version >= '3.12'
