Single-partition Dask executor for cuDF-Polars

python/cudf_polars/cudf_polars/callback.py

@@ -145,6 +145,12 @@ def _callback(
         set_device(device),
         set_memory_resource(memory_resource),
     ):
+        if os.environ.get("CUDF_POLARS_DASK", "OFF").upper() == "ON":
+            # Use experimental Dask executor
+            from cudf_polars.experimental.parallel import evaluate_dask
+
+            return evaluate_dask(ir).to_polars()
+
         return ir.evaluate(cache={}).to_polars()
 
 

python/cudf_polars/cudf_polars/containers/dataframe.py

@@ -226,6 +226,15 @@ def filter(self, mask: Column) -> Self:
         table = plc.stream_compaction.apply_boolean_mask(self.table, mask.obj)
         return type(self).from_table(table, self.column_names).sorted_like(self)
 
+    @classmethod
+    def concat(cls, dfs: list[Self]) -> Self:
+        """Concatenate a list of dataframes."""
+        assert len(dfs), "Cannot concatenate an empty list"
+        if len(dfs) == 1:
+            return dfs[0]
+        table = plc.concatenate.concatenate([df.table for df in dfs])
+        return cls.from_table(table, dfs[0].column_names)
+
     def slice(self, zlice: tuple[int, int] | None) -> Self:
         """
         Slice a dataframe.

python/cudf_polars/cudf_polars/dsl/expressions/sorting.py

@@ -69,7 +69,7 @@ def __init__(
         *by: Expr,
     ) -> None:
         self.dtype = dtype
-        self.options = options
+        self.options = (options[0], tuple(options[1]), tuple(options[2]))
         self.children = (column, *by)
 
     def do_evaluate(

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -1468,13 +1468,20 @@ def __init__(self, schema: Schema, name: str, options: Any, df: IR):
             self.options = (
                 tuple(indices),
                 tuple(pivotees),
-                (variable_name, schema[variable_name]),
-                (value_name, schema[value_name]),
+                variable_name,
+                value_name,
             )
-        self._non_child_args = (name, self.options)
+        self._non_child_args = (schema, name, self.options)
+
+    def get_hashable(self) -> Hashable:
+        """Hashable representation of the node."""
+        schema_hash = tuple(self.schema.items())
+        return (type(self), schema_hash, self.name, str(self.options), *self.children)
 
     @classmethod
-    def do_evaluate(cls, name: str, options: Any, df: DataFrame) -> DataFrame:
+    def do_evaluate(
+        cls, schema: Schema, name: str, options: Any, df: DataFrame
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         if name == "rechunk":
             # No-op in our data model
@@ -1496,8 +1503,8 @@ def do_evaluate(cls, name: str, options: Any, df: DataFrame) -> DataFrame:
             (
                 indices,
                 pivotees,
-                (variable_name, variable_dtype),
-                (value_name, value_dtype),
+                variable_name,
+                value_name,
             ) = options
             npiv = len(pivotees)
             index_columns = [
@@ -1514,15 +1521,18 @@ def do_evaluate(cls, name: str, options: Any, df: DataFrame) -> DataFrame:
                         plc.interop.from_arrow(
                             pa.array(
                                 pivotees,
-                                type=plc.interop.to_arrow(variable_dtype),
+                                type=plc.interop.to_arrow(schema[variable_name]),
                             ),
                         )
                     ]
                 ),
                 df.num_rows,
             ).columns()
             value_column = plc.concatenate.concatenate(
-                [df.column_map[pivotee].astype(value_dtype).obj for pivotee in pivotees]
+                [
+                    df.column_map[pivotee].astype(schema[value_name]).obj
+                    for pivotee in pivotees
+                ]
             )
             return DataFrame(
                 [

python/cudf_polars/cudf_polars/dsl/nodebase.py

@@ -18,6 +18,19 @@
 T = TypeVar("T", bound="Node[Any]")
 
 
+class PartitionInfo:
+    """
+    Partitioning information.
+
+    This class only tracks the partition count (for now).
+    """
+
+    __slots__ = ("count",)
+
+    def __init__(self, count: int):
+        self.count = count
+
+
 class Node(Generic[T]):
     """
     An abstract node type.
@@ -34,9 +47,10 @@ class Node(Generic[T]):
     *children).``
     """
 
-    __slots__ = ("_hash_value", "_repr_value", "children")
+    __slots__ = ("_hash_value", "_repr_value", "_parts_info", "children")
     _hash_value: int
     _repr_value: str
+    _parts_info: PartitionInfo
     children: tuple[T, ...]
     _non_child: ClassVar[tuple[str, ...]] = ()
 
@@ -150,3 +164,17 @@ def __repr__(self) -> str:
             args = ", ".join(f"{arg!r}" for arg in self._ctor_arguments(self.children))
             self._repr_value = f"{type(self).__name__}({args})"
             return self._repr_value
+
+    @property
+    def parts(self) -> PartitionInfo:
+        """Return Partitioning information for this node."""
+        try:
+            return self._parts_info
+        except AttributeError:
+            self._parts_info = self._get_partition_info()
+            return self._parts_info
+
+    def _get_partition_info(self) -> PartitionInfo:
+        from cudf_polars.experimental.parallel import _ir_partition_info_dispatch
+
+        return _ir_partition_info_dispatch(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 depreciation period."""
+
+from __future__ import annotations
+
+__all__: list[str] = []