Support cudf as a DataFrame backend

dask_expr/_collection.py

@@ -1088,9 +1088,10 @@ def read_parquet(
     aggregate_files=None,
     parquet_file_extension=(".parq", ".parquet", ".pq"),
     filesystem="fsspec",
+    engine=None,
     **kwargs,
 ):
-    from dask_expr.io.parquet import ReadParquet
+    from dask_expr.io.parquet import ReadParquet, _set_parquet_engine
 
     if not isinstance(path, str):
         path = stringify_path(path)
@@ -1113,6 +1114,7 @@ def read_parquet(
             aggregate_files=aggregate_files,
             parquet_file_extension=parquet_file_extension,
             filesystem=filesystem,
+            engine=_set_parquet_engine(engine),
             kwargs=kwargs,
         )
     )

dask_expr/_expr.py

@@ -54,6 +54,18 @@ def __init__(self, *args, **kwargs):
                 operands.append(type(self)._defaults[parameter])
         assert not kwargs
         self.operands = operands
+        if self._required_attribute:
+            dep = next(iter(self.dependencies()))._meta
+            if not hasattr(dep, self._required_attribute):
+                # Raise a ValueError instead of AttributeError to
+                # avoid infinite recursion
+                raise ValueError(f"{dep} has no attribute {self._required_attribute}")
+
+    @property
+    def _required_attribute(self) -> str:
+        # Specify if the first `dependency` must support
+        # a specific attribute for valid behavior.
+        return None
 
     @functools.cached_property
     def ndim(self):
@@ -941,6 +953,12 @@ class Blockwise(Expr):
     _keyword_only = []
     _projection_passthrough = False
 
+    @property
+    def _required_attribute(self):
+        if isinstance(self.operation, type(M.method_caller)):
+            return self.operation.method
+        return None
+
     @functools.cached_property
     def _meta(self):
         args = [op._meta if isinstance(op, Expr) else op for op in self._args]
@@ -1027,7 +1045,13 @@ def _combine_similar(self, root: Expr):
         # Push projections back up through `_projection_passthrough`
         # operations if it reduces the number of unique expression nodes.
         if self._projection_passthrough and isinstance(self.frame, Projection):
-            common = type(self)(self.frame.frame, *self.operands[1:])
+            try:
+                common = type(self)(self.frame.frame, *self.operands[1:])
+            except ValueError:
+                # May have encountered a problem with `_required_attribute`.
+                # (There is no guarentee that the same method will exist for
+                # both a Series and DataFrame)
+                return None
             projection = self.frame.operand("columns")
             push_up_projection = False
             for op in self._find_similar_operations(root, ignore=self._parameters):

dask_expr/_reductions.py

@@ -488,6 +488,7 @@ class NBytes(Reduction):
     # Only supported for Series objects
     reduction_chunk = lambda ser: ser.nbytes
     reduction_aggregate = sum
+    _required_attribute = "nbytes"
 
 
 class Var(Reduction):
@@ -519,7 +520,7 @@ def aggregate_kwargs(self):
     @classmethod
     def reduction_chunk(cls, x, skipna=True, numeric_only=False):
         kwargs = {"numeric_only": numeric_only} if is_dataframe_like(x) else {}
-        if skipna:
+        if skipna or numeric_only:
             n = x.count(**kwargs)
             kwargs["skipna"] = skipna
             avg = x.mean(**kwargs)
@@ -529,6 +530,11 @@ def reduction_chunk(cls, x, skipna=True, numeric_only=False):
             n = len(x)
             kwargs["skipna"] = skipna
             avg = x.sum(**kwargs) / n
+        if numeric_only:
+            # Workaround for cudf bug
+            # (see: https://github.com/rapidsai/cudf/issues/13731)
+            x = x.select_dtypes("number")
+            n = n.loc[x.columns]
         m2 = ((x - avg) ** 2).sum(**kwargs)
         return n, avg, m2
 

dask_expr/io/parquet.py

@@ -24,7 +24,7 @@
 from dask.dataframe.io.parquet.utils import _split_user_options
 from dask.dataframe.io.utils import _is_local_fs
 from dask.delayed import delayed
-from dask.utils import apply, natural_sort_key
+from dask.utils import apply, natural_sort_key, typename
 from fsspec.utils import stringify_path
 
 from dask_expr._expr import (
@@ -157,7 +157,6 @@ def _layer(self):
 def to_parquet(
     df,
     path,
-    engine="pyarrow",
     compression="snappy",
     write_index=True,
     append=False,
@@ -177,6 +176,7 @@ def to_parquet(
     from dask_expr._collection import new_collection
     from dask_expr.io.parquet import NONE_LABEL, ToParquet
 
+    engine = _set_parquet_engine(meta=df._meta)
     compute_kwargs = compute_kwargs or {}
 
     partition_on = partition_on or []
@@ -391,6 +391,7 @@ class ReadParquet(PartitionsFiltered, BlockwiseIO):
         "aggregate_files",
         "parquet_file_extension",
         "filesystem",
+        "engine",
         "kwargs",
         "_partitions",
         "_series",
@@ -409,6 +410,7 @@ class ReadParquet(PartitionsFiltered, BlockwiseIO):
         "aggregate_files": None,
         "parquet_file_extension": (".parq", ".parquet", ".pq"),
         "filesystem": "fsspec",
+        "engine": "pyarrow",
         "kwargs": None,
         "_partitions": None,
         "_series": False,
@@ -417,7 +419,10 @@ class ReadParquet(PartitionsFiltered, BlockwiseIO):
 
     @property
     def engine(self):
-        return get_engine("pyarrow")
+        _engine = self.operand("engine")
+        if isinstance(_engine, str):
+            return get_engine(_engine)
+        return _engine
 
     @property
     def columns(self):
@@ -681,6 +686,20 @@ def _update_length_statistics(self):
 #
 
 
+def _set_parquet_engine(engine=None, meta=None):
+    # Use `engine` or `meta` input to set the parquet engine
+    if engine is None:
+        if (
+            meta is not None and typename(meta).split(".")[0] == "cudf"
+        ) or dask.config.get("dataframe.backend", "pandas") == "cudf":
+            from dask_cudf.io.parquet import CudfEngine
+
+            engine = CudfEngine
+        else:
+            engine = "pyarrow"
+    return engine
+
+
 def _align_statistics(parts, statistics):
     # Make sure parts and statistics are aligned
     # (if statistics is not empty)

dask_expr/io/tests/test_io.py

@@ -1,20 +1,25 @@
+import importlib
 import os
 
 import dask.dataframe as dd
-import pandas as pd
 import pytest
+from dask import config
 from dask.dataframe.utils import assert_eq
 
 from dask_expr import from_dask_dataframe, from_pandas, optimize, read_csv, read_parquet
 from dask_expr._expr import Expr, Lengths, Literal, Replace
 from dask_expr._reductions import Len
 from dask_expr.io import ReadParquet
 
+# Import backend DataFrame library to test
+BACKEND = config.get("dataframe.backend", "pandas")
+lib = importlib.import_module(BACKEND)
+
 
 def _make_file(dir, format="parquet", df=None):
     fn = os.path.join(str(dir), f"myfile.{format}")
     if df is None:
-        df = pd.DataFrame({c: range(10) for c in "abcde"})
+        df = lib.DataFrame({c: range(10) for c in "abcde"})
     if format == "csv":
         df.to_csv(fn)
     elif format == "parquet":
@@ -83,7 +88,7 @@ def test_io_fusion(tmpdir, fmt):
 
 
 def test_predicate_pushdown(tmpdir):
-    original = pd.DataFrame(
+    original = lib.DataFrame(
         {
             "a": [1, 2, 3, 4, 5] * 10,
             "b": [0, 1, 2, 3, 4] * 10,
@@ -106,11 +111,11 @@ def test_predicate_pushdown(tmpdir):
     y_result = y.compute()
     assert y_result.name == "b"
     assert len(y_result) == 6
-    assert all(y_result == 4)
+    assert (y_result == 4).all()
 
 
 def test_predicate_pushdown_compound(tmpdir):
-    pdf = pd.DataFrame(
+    pdf = lib.DataFrame(
         {
             "a": [1, 2, 3, 4, 5] * 10,
             "b": [0, 1, 2, 3, 4] * 10,
@@ -134,15 +139,18 @@ def test_predicate_pushdown_compound(tmpdir):
     )
 
     # Test OR
-    x = df[(df.a == 5) | (df.c > 20)][df.b != 0]["b"]
+    x = df[(df.a == 5) | (df.c > 20)]
+    x = x[x.b != 0]["b"]
     y = optimize(x, fuse=False)
     assert isinstance(y.expr, ReadParquet)
     filters = [set(y.filters[0]), set(y.filters[1])]
     assert {("c", ">", 20), ("b", "!=", 0)} in filters
     assert {("a", "==", 5), ("b", "!=", 0)} in filters
+    expect = pdf[(pdf.a == 5) | (pdf.c > 20)]
+    expect = expect[expect.b != 0]["b"]
     assert_eq(
         y,
-        pdf[(pdf.a == 5) | (pdf.c > 20)][pdf.b != 0]["b"],
+        expect,
         check_index=False,
     )
 
@@ -158,7 +166,7 @@ def test_predicate_pushdown_compound(tmpdir):
 
 @pytest.mark.parametrize("fmt", ["parquet", "csv", "pandas"])
 def test_io_culling(tmpdir, fmt):
-    pdf = pd.DataFrame({c: range(10) for c in "abcde"})
+    pdf = lib.DataFrame({c: range(10) for c in "abcde"})
     if fmt == "parquet":
         dd.from_pandas(pdf, 2).to_parquet(tmpdir)
         df = read_parquet(tmpdir)
@@ -191,23 +199,24 @@ def _check_culling(expr, partitions):
 
 @pytest.mark.parametrize("sort", [True, False])
 def test_from_pandas(sort):
-    pdf = pd.DataFrame({"x": [1, 4, 3, 2, 0, 5]})
+    pdf = lib.DataFrame({"x": [1, 4, 3, 2, 0, 5]})
     df = from_pandas(pdf, npartitions=2, sort=sort)
 
     assert df.divisions == (0, 3, 5) if sort else (None,) * 3
     assert_eq(df, pdf)
 
 
 def test_from_pandas_immutable():
-    pdf = pd.DataFrame({"x": [1, 2, 3, 4]})
+    pdf = lib.DataFrame({"x": [1, 2, 3, 4]})
     expected = pdf.copy()
     df = from_pandas(pdf)
     pdf["z"] = 100
     assert_eq(df, expected)
 
 
 def test_parquet_complex_filters(tmpdir):
-    df = read_parquet(_make_file(tmpdir))
+    with config.set({"dataframe.backend": BACKEND}):
+        df = read_parquet(_make_file(tmpdir))
     pdf = df.compute()
     got = df["a"][df["b"] > df["b"].mean()]
     expect = pdf["a"][pdf["b"] > pdf["b"].mean()]
@@ -247,7 +256,7 @@ def test_from_dask_dataframe(optimize):
 
 @pytest.mark.parametrize("optimize", [True, False])
 def test_to_dask_dataframe(optimize):
-    pdf = pd.DataFrame({"x": [1, 4, 3, 2, 0, 5]})
+    pdf = lib.DataFrame({"x": [1, 4, 3, 2, 0, 5]})
     df = from_pandas(pdf, npartitions=2)
     ddf = df.to_dask_dataframe(optimize=optimize)
     assert isinstance(ddf, dd.DataFrame)
@@ -256,7 +265,7 @@ def test_to_dask_dataframe(optimize):
 
 @pytest.mark.parametrize("write_metadata_file", [True, False])
 def test_to_parquet(tmpdir, write_metadata_file):
-    pdf = pd.DataFrame({"x": [1, 4, 3, 2, 0, 5]})
+    pdf = lib.DataFrame({"x": [1, 4, 3, 2, 0, 5]})
     df = from_pandas(pdf, npartitions=2)
 
     # Check basic parquet round trip
@@ -277,7 +286,7 @@ def test_to_parquet(tmpdir, write_metadata_file):
 
 
 def test_combine_similar(tmpdir):
-    pdf = pd.DataFrame(
+    pdf = lib.DataFrame(
         {"x": [0, 1, 2, 3] * 4, "y": range(16), "z": [None, 1, 2, 3] * 4}
     )
     fn = _make_file(tmpdir, format="parquet", df=pdf)