Add multi-partition DataFrameScan support to cuDF-Polars

python/cudf_polars/cudf_polars/callback.py

@@ -217,14 +217,16 @@ def validate_config_options(config: dict) -> None:
         If the configuration contains unsupported options.
     """
     if unsupported := (
-        config.keys() - {"raise_on_fail", "parquet_options", "executor"}
+        config.keys()
+        - {"raise_on_fail", "parquet_options", "parallel_options", "executor"}
     ):
         raise ValueError(
             f"Engine configuration contains unsupported settings: {unsupported}"
         )
     assert {"chunked", "chunk_read_limit", "pass_read_limit"}.issuperset(
         config.get("parquet_options", {})
     )
+    assert {"num_rows_threshold"}.issuperset(config.get("parallel_options", {}))
 
 
 def execute_with_cudf(nt: NodeTraverser, *, config: GPUEngine) -> None:

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -683,26 +683,30 @@ class DataFrameScan(IR):
     This typically arises from ``q.collect().lazy()``
     """
 
-    __slots__ = ("df", "projection", "predicate")
-    _non_child = ("schema", "df", "projection", "predicate")
+    __slots__ = ("df", "projection", "predicate", "config_options")
+    _non_child = ("schema", "df", "projection", "predicate", "config_options")
     df: Any
     """Polars LazyFrame object."""
     projection: tuple[str, ...] | None
     """List of columns to project out."""
     predicate: expr.NamedExpr | None
     """Mask to apply."""
+    config_options: dict[str, Any]
+    """GPU-specific configuration options"""
 
     def __init__(
         self,
         schema: Schema,
         df: Any,
         projection: Sequence[str] | None,
         predicate: expr.NamedExpr | None,
+        config_options: dict[str, Any],
     ):
         self.schema = schema
         self.df = df
         self.projection = tuple(projection) if projection is not None else None
         self.predicate = predicate
+        self.config_options = config_options
         self._non_child_args = (schema, df, self.projection, predicate)
         self.children = ()
 
@@ -714,7 +718,14 @@ def get_hashable(self) -> Hashable:
         not stable across runs, or repeat instances of the same equal dataframes.
         """
         schema_hash = tuple(self.schema.items())
-        return (type(self), schema_hash, id(self.df), self.projection, self.predicate)
+        return (
+            type(self),
+            schema_hash,
+            id(self.df),
+            self.projection,
+            self.predicate,
+            json.dumps(self.config_options),
+        )
 
     @classmethod
     def do_evaluate(

python/cudf_polars/cudf_polars/dsl/translate.py

@@ -263,6 +263,7 @@ def _(
         translate_named_expr(translator, n=node.selection)
         if node.selection is not None
         else None,
+        translator.config.config.copy(),
     )
 
 

python/cudf_polars/cudf_polars/experimental/io.py

@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Parallel IO Logic."""
+
+from __future__ import annotations
+
+import math
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import DataFrameScan
+from cudf_polars.experimental.parallel import (
+    PartitionInfo,
+    generate_ir_tasks,
+    get_key_name,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.parallel import LowerIRTransformer
+
+
+##
+## DataFrameScan
+##
+
+
+class ParDataFrameScan(DataFrameScan):
+    """Parallel DataFrameScan."""
+
+    @property
+    def _max_n_rows(self) -> int:
+        """Row-count threshold for splitting a DataFrame."""
+        parallel_options = self.config_options.get("parallel_options", {})
+        return parallel_options.get("num_rows_threshold", 1_000_000)
+
+    @cached_property
+    def _count(self) -> int:
+        """Partition count."""
+        total_rows = max(self.df.shape()[0], 1)
+        return math.ceil(total_rows / self._max_n_rows)
+
+    def _tasks(self) -> MutableMapping[Any, Any]:
+        """Task graph."""
+        total_rows = max(self.df.shape()[0], 1)
+        stride = math.ceil(total_rows / self._count)
+        key_name = get_key_name(self)
+        return {
+            (key_name, i): (
+                self.do_evaluate,
+                self.schema,
+                self.df.slice(offset, stride),
+                self.projection,
+                self.predicate,
+            )
+            for i, offset in enumerate(range(0, total_rows, stride))
+        }
+
+
+def lower_dataframescan_node(
+    ir: DataFrameScan, rec: LowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    """Rewrite a Scan node with proper partitioning."""
+    new_node = ParDataFrameScan(
+        ir.schema,
+        ir.df,
+        ir.projection,
+        ir.predicate,
+        ir.config_options,
+    )
+    return new_node, {new_node: PartitionInfo(count=new_node._count)}
+
+
+@generate_ir_tasks.register(ParDataFrameScan)
+def _(
+    ir: ParDataFrameScan, partition_info: MutableMapping[IR, PartitionInfo]
+) -> MutableMapping[Any, Any]:
+    assert (
+        partition_info[ir].count == ir._count
+    ), "Inconsistent ParDataFrameScan partitioning."
+    return ir._tasks()

python/cudf_polars/cudf_polars/experimental/parallel.py

@@ -8,11 +8,11 @@
 from functools import reduce, singledispatch
 from typing import TYPE_CHECKING, Any
 
-from cudf_polars.dsl.ir import IR
+from cudf_polars.dsl.ir import IR, DataFrameScan, Union
 from cudf_polars.dsl.traversal import traversal
 
 if TYPE_CHECKING:
-    from collections.abc import MutableMapping
+    from collections.abc import MutableMapping, Sequence
     from typing import TypeAlias
 
     from cudf_polars.containers import DataFrame
@@ -223,7 +223,14 @@ def task_graph(
         operator.or_,
         (generate_ir_tasks(node, partition_info) for node in traversal(ir)),
     )
-    return graph, (get_key_name(ir), 0)
+
+    key_name = get_key_name(ir)
+    partition_count = partition_info[ir].count
+    if partition_count > 1:
+        graph[key_name] = (_concat, [(key_name, i) for i in range(partition_count)])
+        return graph, key_name
+    else:
+        return graph, (key_name, 0)
 
 
 def evaluate_dask(ir: IR) -> DataFrame:
@@ -234,3 +241,22 @@ def evaluate_dask(ir: IR) -> DataFrame:
 
     graph, key = task_graph(ir, partition_info)
     return get(graph, key)
+
+
+def _concat(dfs: Sequence[DataFrame]) -> DataFrame:
+    # Concatenate a sequence of DataFrames vertically
+    return Union.do_evaluate(None, *dfs)
+
+
+##
+## DataFrameScan
+##
+
+
+@lower_ir_node.register(DataFrameScan)
+def _(
+    ir: DataFrameScan, rec: LowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    import cudf_polars.experimental.io as _io
+
+    return _io.lower_dataframescan_node(ir, rec)

python/cudf_polars/tests/dsl/test_traversal.py

@@ -116,7 +116,11 @@ def test_rewrite_ir_node():
     def replace_df(node, rec):
         if isinstance(node, ir.DataFrameScan):
             return ir.DataFrameScan(
-                node.schema, new_df._df, node.projection, node.predicate
+                node.schema,
+                new_df._df,
+                node.projection,
+                node.predicate,
+                node.config_options,
             )
         return reuse_if_unchanged(node, rec)
 
@@ -144,7 +148,11 @@ def test_rewrite_scan_node(tmp_path):
     def replace_scan(node, rec):
         if isinstance(node, ir.Scan):
             return ir.DataFrameScan(
-                node.schema, right._df, node.with_columns, node.predicate
+                node.schema,
+                right._df,
+                node.with_columns,
+                node.predicate,
+                node.config_options,
             )
         return reuse_if_unchanged(node, rec)
 

python/cudf_polars/tests/experimental/test_dataframescan.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import pytest
+
+import polars as pl
+
+from cudf_polars import Translator
+from cudf_polars.experimental.parallel import lower_ir_graph
+from cudf_polars.testing.asserts import assert_gpu_result_equal
+
+
+@pytest.fixture(scope="module")
+def df():
+    return pl.LazyFrame(
+        {
+            "x": range(30_000),
+            "y": ["cat", "dog", "fish"] * 10_000,
+            "z": [1.0, 2.0, 3.0, 4.0, 5.0] * 6_000,
+        }
+    )
+
+
+@pytest.mark.parametrize("num_rows_threshold", [1_000, 1_000_000])
+def test_parallel_dataframescan(df, num_rows_threshold):
+    total_row_count = len(df.collect())
+    engine = pl.GPUEngine(
+        raise_on_fail=True,
+        parallel_options={"num_rows_threshold": num_rows_threshold},
+        executor="dask-experimental",
+    )
+    assert_gpu_result_equal(df, engine=engine)
+
+    # Check partitioning
+    qir = Translator(df._ldf.visit(), engine).translate_ir()
+    ir, info = lower_ir_graph(qir)
+    count = info[ir].count
+    if num_rows_threshold < total_row_count:
+        assert count > 1
+    else:
+        assert count == 1