Polars: DataFrame Serialization

ci/test_wheel_cudf_polars.sh

@@ -17,11 +17,11 @@ rapids-logger "Installing cudf_polars and its dependencies"
 # generate constraints (possibly pinning to oldest support versions of dependencies)
 rapids-generate-pip-constraints py_test_cudf_polars ./constraints.txt
 
-# echo to expand wildcard before adding `[test]` requires for pip
+# echo to expand wildcard before adding `[test,experimental]` requires for pip
 python -m pip install \
     -v \
     --constraint ./constraints.txt \
-    "$(echo ./dist/cudf_polars_${RAPIDS_PY_CUDA_SUFFIX}*.whl)[test]" \
+    "$(echo ./dist/cudf_polars_${RAPIDS_PY_CUDA_SUFFIX}*.whl)[test,experimental]" \
     "$(echo ./dist/libcudf_${RAPIDS_PY_CUDA_SUFFIX}*.whl)" \
     "$(echo ./dist/pylibcudf_${RAPIDS_PY_CUDA_SUFFIX}*.whl)"
 

dependencies.yaml

@@ -41,9 +41,11 @@ files:
       - test_cpp
       - test_python_common
       - test_python_cudf
+      - test_python_cudf_common
       - test_python_dask_cudf
       - test_python_pylibcudf
       - test_python_cudf_pandas
+      - test_python_cudf_polars
   test_static_build:
     output: none
     includes:
@@ -59,6 +61,7 @@ files:
       - cuda_version
       - py_version
       - test_python_common
+      - test_python_cudf_common
       - test_python_cudf
       - test_python_cudf_pandas
   test_python_cudf:
@@ -67,13 +70,15 @@ files:
       - cuda_version
       - py_version
       - test_python_common
+      - test_python_cudf_common
       - test_python_cudf
   test_python_other:
     output: none
     includes:
       - cuda_version
       - py_version
       - test_python_common
+      - test_python_cudf_common
       - test_python_dask_cudf
   test_java:
     output: none
@@ -152,6 +157,7 @@ files:
       key: test
     includes:
       - test_python_common
+      - test_python_cudf_common
       - test_python_cudf
   py_build_libcudf:
     output: pyproject
@@ -216,6 +222,7 @@ files:
       key: test
     includes:
       - test_python_common
+      - test_python_cudf_common
       - test_python_pylibcudf
   py_test_pandas_cudf:
     output: pyproject
@@ -248,6 +255,14 @@ files:
     includes:
       - run_cudf_polars
       - depends_on_pylibcudf
+  py_run_cudf_polars_experimental:
+    output: pyproject
+    pyproject_dir: python/cudf_polars
+    extras:
+      table: project.optional-dependencies
+      key: experimental
+    includes:
+      - run_cudf_polars_experimental
   py_test_cudf_polars:
     output: pyproject
     pyproject_dir: python/cudf_polars
@@ -256,6 +271,7 @@ files:
       key: test
     includes:
       - test_python_common
+      - test_python_cudf_polars
   py_build_dask_cudf:
     output: pyproject
     pyproject_dir: python/dask_cudf
@@ -281,6 +297,7 @@ files:
       key: test
     includes:
       - test_python_common
+      - test_python_cudf_common
       - test_python_dask_cudf
   py_build_cudf_kafka:
     output: pyproject
@@ -313,6 +330,7 @@ files:
       key: test
     includes:
       - test_python_common
+      - test_python_cudf_common
   py_build_custreamz:
     output: pyproject
     pyproject_dir: python/custreamz
@@ -337,6 +355,7 @@ files:
       key: test
     includes:
       - test_python_common
+      - test_python_cudf_common
 channels:
   - rapidsai
   - rapidsai-nightly
@@ -730,6 +749,11 @@ dependencies:
       - output_types: [conda, requirements, pyproject]
         packages:
           - polars>=1.11,<1.14
+  run_cudf_polars_experimental:
+    common:
+      - output_types: [conda, requirements, pyproject]
+        packages:
+          - rapids-dask-dependency==24.12.*,>=0.0.0a0
   run_dask_cudf:
     common:
       - output_types: [conda, requirements, pyproject]
@@ -779,6 +803,7 @@ dependencies:
           - pytest<8
           - pytest-cov
           - pytest-xdist
+  test_python_cudf_common:
     specific:
       # Define additional constraints for testing with oldest dependencies.
       - output_types: [conda, requirements]
@@ -884,6 +909,11 @@ dependencies:
               - pyarrow==14.0.1
           - matrix:
             packages:
+  test_python_cudf_polars:
+    common:
+      - output_types: [conda, requirements, pyproject]
+        packages:
+          - *numpy
   depends_on_libcudf:
     common:
       - output_types: conda

python/cudf_polars/cudf_polars/containers/dataframe.py

@@ -5,8 +5,9 @@
 
 from __future__ import annotations
 
+import pickle
 from functools import cached_property
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import pyarrow as pa
 
@@ -147,6 +148,75 @@ def from_table(cls, table: plc.Table, names: Sequence[str]) -> Self:
             Column(c, name=name) for c, name in zip(table.columns(), names, strict=True)
         )
 
+    @classmethod
+    def deserialize(
+        cls, header: Mapping[str, Any], frames: tuple[memoryview, plc.gpumemoryview]
+    ) -> Self:
+        """
+        Create a DataFrame from a serialized representation returned by `.serialize()`.
+
+        Parameters
+        ----------
+        header
+            The (unpickled) metadata required to reconstruct the object.
+        frames
+            Two-tuple of frames (a memoryview and a gpumemoryview).
+
+        Returns
+        -------
+        DataFrame
+            The deserialized DataFrame.
+        """
+        packed_metadata, packed_gpu_data = frames
+        table = plc.contiguous_split.unpack_from_memoryviews(
+            packed_metadata, packed_gpu_data
+        )
+        return cls(
+            Column(c, **kw)
+            for c, kw in zip(table.columns(), header["columns_kwargs"], strict=True)
+        )
+
+    def serialize(
+        self,
+    ) -> tuple[Mapping[str, Any], tuple[memoryview, plc.gpumemoryview]]:
+        """
+        Serialize the table into header and frames.
+
+        Follows the Dask serialization scheme with a picklable header (dict) and
+        a tuple of frames (in this case a contiguous host and device buffer).
+
+        To enable dask support, dask serializers must be registered
+
+        >>> from cudf_polars.experimental.dask_serialize import register
+        >>> register()
+
+        Returns
+        -------
+        header
+            A dict containing any picklable metadata required to reconstruct the object.
+        frames
+            Two-tuple of frames suitable for passing to `unpack_from_memoryviews`
+        """
+        packed = plc.contiguous_split.pack(self.table)
+
+        # Keyword arguments for `Column.__init__`.
+        columns_kwargs = [
+            {
+                "is_sorted": col.is_sorted,
+                "order": col.order,
+                "null_order": col.null_order,
+                "name": col.name,
+            }
+            for col in self.columns
+        ]
+        header = {
+            "columns_kwargs": columns_kwargs,
+            # Dask Distributed uses "type-serialized" to dispatch deserialization
+            "type-serialized": pickle.dumps(type(self)),
+            "frame_count": 2,
+        }
+        return header, packed.release()
+
     def sorted_like(
         self, like: DataFrame, /, *, subset: Set[str] | None = None
     ) -> Self:

python/cudf_polars/cudf_polars/experimental/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Experimental features, which can change without any deprecation period."""
+
+from __future__ import annotations
+
+__all__: list[str] = []

python/cudf_polars/cudf_polars/experimental/dask_serialize.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Dask serialization."""
+
+from __future__ import annotations
+
+from distributed.protocol import dask_deserialize, dask_serialize
+from distributed.protocol.cuda import cuda_deserialize, cuda_serialize
+from distributed.utils import log_errors
+
+import pylibcudf as plc
+import rmm
+
+from cudf_polars.containers import DataFrame
+
+__all__ = ["register"]
+
+
+def register() -> None:
+    """Register dask serialization routines for DataFrames."""
+
+    @cuda_serialize.register(DataFrame)
+    def _(x: DataFrame):
+        with log_errors():
+            header, frames = x.serialize()
+            return header, list(frames)  # Dask expect a list of frames
+
+    @cuda_deserialize.register(DataFrame)
+    def _(header, frames):
+        with log_errors():
+            assert len(frames) == 2
+            return DataFrame.deserialize(header, tuple(frames))
+
+    @dask_serialize.register(DataFrame)
+    def _(x: DataFrame):
+        with log_errors():
+            header, (metadata, gpudata) = x.serialize()
+
+            # For robustness, we check that the gpu data is contiguous
+            cai = gpudata.__cuda_array_interface__
+            assert len(cai["shape"]) == 1
+            assert cai["strides"] is None or cai["strides"] == (1,)
+            assert cai["typestr"] == "|u1"
+            nbytes = cai["shape"][0]
+
+            # Copy the gpudata to host memory
+            gpudata_on_host = memoryview(
+                rmm.DeviceBuffer(ptr=gpudata.ptr, size=nbytes).copy_to_host()
+            )
+            return header, (metadata, gpudata_on_host)
+
+    @dask_deserialize.register(DataFrame)
+    def _(header, frames) -> DataFrame:
+        with log_errors():
+            assert len(frames) == 2
+            # Copy the second frame (the gpudata in host memory) back to the gpu
+            frames = frames[0], plc.gpumemoryview(rmm.DeviceBuffer.to_device(frames[1]))
+            return DataFrame.deserialize(header, frames)

python/cudf_polars/pyproject.toml

@@ -35,10 +35,14 @@ classifiers = [
 
 [project.optional-dependencies]
 test = [
+    "numpy>=1.23,<3.0a0",
     "pytest-cov",
     "pytest-xdist",
     "pytest<8",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
+experimental = [
+    "rapids-dask-dependency==24.12.*,>=0.0.0a0",
+] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [project.urls]
 Homepage = "https://github.com/rapidsai/cudf"

python/cudf_polars/tests/containers/test_dataframe.py

@@ -3,9 +3,11 @@
 
 from __future__ import annotations
 
+import pyarrow as pa
 import pytest
 
 import polars as pl
+from polars.testing.asserts import assert_frame_equal
 
 import pylibcudf as plc
 
@@ -161,3 +163,24 @@ def test_empty_name_roundtrips_overlap():
 def test_empty_name_roundtrips_no_overlap():
     df = pl.LazyFrame({"": [1, 2, 3], "b": [4, 5, 6]})
     assert_gpu_result_equal(df)
+
+
+@pytest.mark.parametrize(
+    "arrow_tbl",
+    [
+        pa.table([]),
+        pa.table({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}),
+        pa.table({"a": [1, 2, 3]}),
+        pa.table({"a": [1], "b": [2], "c": [3]}),
+        pa.table({"a": ["a", "bb", "ccc"]}),
+        pa.table({"a": [1, 2, None], "b": [None, 3, 4]}),
+    ],
+)
+def test_serialization_roundtrip(arrow_tbl):
+    plc_tbl = plc.interop.from_arrow(arrow_tbl)
+    df = DataFrame.from_table(plc_tbl, names=arrow_tbl.column_names)
+
+    header, frames = df.serialize()
+    res = DataFrame.deserialize(header, frames)
+
+    assert_frame_equal(df.to_polars(), res.to_polars())

python/cudf_polars/tests/experimental/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Testing experimental features."""
+
+from __future__ import annotations
+
+__all__: list[str] = []