Repository: rapidsai/cudf
Pull request #18405: Automatic single-partition fallback in cudf-polars
Author: login: rjzamora
Labels: name: Python
description: Affects Python cuDF API.
name: 5 - Ready to Merge
description: Testing and reviews complete, ready to merge
name: improvement
description: Improvement / enhancement to an existing function
name: non-breaking
description: Non-breaking change
name: cudf-polars
description: Issues specific to cudf-polars
Created: 2025-04-01T21:07:13Z | Merged: 2025-04-08T16:45:53Z

## Description
These changes allow the "partitioned" cudf-polars executor (currently called `"dask-experimental"`) to automatically collapse (`Repartition`) multi-partition dependencies whenever an unsupported algorithm is encountered. This behavior is probably needed before we can encourage users to test out partitioned-execution (certainly before it can become the default).

## Checklist
- [x] I am familiar with the [Contributing Guidelines](https://github.com/rapidsai/cudf/blob/HEAD/CONTRIBUTING.md).
- [x] New or existing tests cover these changes.
- [ ] The documentation is up to date with these changes.



Diff:
diff --git a/python/cudf_polars/cudf_polars/callback.py b/python/cudf_polars/cudf_polars/callback.py
index 30d7b7d392b..a420e62a1c4 100644
--- a/python/cudf_polars/cudf_polars/callback.py
+++ b/python/cudf_polars/cudf_polars/callback.py
@@ -32,6 +32,7 @@
 
     from cudf_polars.dsl.ir import IR
     from cudf_polars.typing import NodeTraverser
+    from cudf_polars.utils.config import ConfigOptions
 
 __all__: list[str] = ["execute_with_cudf"]
 
@@ -187,6 +188,7 @@ def _callback(
     device: int | None,
     memory_resource: int | None,
     executor: Literal["pylibcudf", "dask-experimental"] | None,
+    config_options: ConfigOptions,
     timer: Timer | None,
 ) -> pl.DataFrame: ...
 
@@ -202,6 +204,7 @@ def _callback(
     device: int | None,
     memory_resource: int | None,
     executor: Literal["pylibcudf", "dask-experimental"] | None,
+    config_options: ConfigOptions,
     timer: Timer | None,
 ) -> tuple[pl.DataFrame, list[tuple[int, int, str]]]: ...
 
@@ -216,6 +219,7 @@ def _callback(
     device: int | None,
     memory_resource: int | None,
     executor: Literal["pylibcudf", "dask-experimental"] | None,
+    config_options: ConfigOptions,
     timer: Timer | None,
 ) -> pl.DataFrame | tuple[pl.DataFrame, list[tuple[int, int, str]]]:
     assert with_columns is None
@@ -238,7 +242,7 @@ def _callback(
         elif executor == "dask-experimental":
             from cudf_polars.experimental.parallel import evaluate_dask
 
-            return evaluate_dask(ir).to_polars()
+            return evaluate_dask(ir, config_options).to_polars()
         else:
             raise ValueError(f"Unknown executor '{executor}'")
 
@@ -313,6 +317,7 @@ def execute_with_cudf(
                         device=device,
                         memory_resource=memory_resource,
                         executor=executor,
+                        config_options=translator.config_options,
                         timer=None,
                     )
                 )
@@ -324,6 +329,7 @@ def execute_with_cudf(
                         device=device,
                         memory_resource=memory_resource,
                         executor=executor,
+                        config_options=translator.config_options,
                         timer=timer,
                     )
                 )
diff --git a/python/cudf_polars/cudf_polars/experimental/base.py b/python/cudf_polars/cudf_polars/experimental/base.py
index 36d0d14d0d7..86f773f8a21 100644
--- a/python/cudf_polars/cudf_polars/experimental/base.py
+++ b/python/cudf_polars/cudf_polars/experimental/base.py
@@ -6,12 +6,9 @@
 
 from typing import TYPE_CHECKING
 
-from cudf_polars.dsl.ir import Union
-
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
-    from cudf_polars.containers import DataFrame
     from cudf_polars.dsl.expr import NamedExpr
     from cudf_polars.dsl.nodebase import Node
 
@@ -42,8 +39,3 @@ def keys(self, node: Node) -> Iterator[tuple[str, int]]:
 def get_key_name(node: Node) -> str:
     """Generate the key name for a Node."""
     return f"{type(node).__name__.lower()}-{hash(node)}"
-
-
-def _concat(*dfs: DataFrame) -> DataFrame:
-    # Concatenate a sequence of DataFrames vertically
-    return Union.do_evaluate(None, *dfs)
diff --git a/python/cudf_polars/cudf_polars/experimental/groupby.py b/python/cudf_polars/cudf_polars/experimental/groupby.py
index f33fbd5c721..c95a0404ac3 100644
--- a/python/cudf_polars/cudf_polars/experimental/groupby.py
+++ b/python/cudf_polars/cudf_polars/experimental/groupby.py
@@ -22,9 +22,10 @@
 )
 from cudf_polars.dsl.ir import GroupBy, Select
 from cudf_polars.dsl.traversal import traversal
-from cudf_polars.experimental.base import PartitionInfo, _concat, get_key_name
+from cudf_polars.experimental.base import PartitionInfo, get_key_name
 from cudf_polars.experimental.dispatch import generate_ir_tasks, lower_ir_node
 from cudf_polars.experimental.shuffle import Shuffle
+from cudf_polars.experimental.utils import _concat, _lower_ir_fallback
 
 if TYPE_CHECKING:
     from collections.abc import Callable, MutableMapping
@@ -148,7 +149,7 @@ def _wrap_unary(select: Expr) -> Expr:
             return selection, aggregations, reductions
         else:
             raise NotImplementedError(
-                "GroupBy does not support multiple partitions "
+                "group_by does not support multiple partitions "
                 f"for this aggregation type:\n{type(expr)}\n"
                 f"Only {_GB_AGG_SUPPORTED} are supported."
             )
@@ -198,10 +199,14 @@ def _(
         return single_part_node, partition_info
 
     # Check group-by keys
-    if not all(expr.is_pointwise for expr in traversal([e.value for e in ir.keys])):
-        raise NotImplementedError(
-            f"GroupBy does not support multiple partitions for keys:\n{ir.keys}"
-        )  # pragma: no cover
+    if not all(
+        expr.is_pointwise for expr in traversal([e.value for e in ir.keys])
+    ):  # pragma: no cover
+        return _lower_ir_fallback(
+            ir,
+            rec,
+            msg="group_by does not support multiple partitions for non-pointwise keys.",
+        )
 
     # Check if we are dealing with any high-cardinality columns
     post_aggregation_count = 1  # Default tree reduction
@@ -226,9 +231,14 @@ def _(
         )
 
     # Decompose the aggregation requests into three distinct phases
-    selection_exprs, piecewise_exprs, reduction_exprs = combine(
-        *(decompose(agg.name, agg.value) for agg in ir.agg_requests)
-    )
+    try:
+        selection_exprs, piecewise_exprs, reduction_exprs = combine(
+            *(decompose(agg.name, agg.value) for agg in ir.agg_requests)
+        )
+    except NotImplementedError:
+        return _lower_ir_fallback(
+            ir, rec, msg="Failed to decompose groupby aggs for multiple partitions."
+        )
 
     # Partition-wise groupby operation
     pwise_schema = {k.name: k.value.dtype for k in ir.keys} | {
@@ -250,8 +260,10 @@ def _(
     gb_inter: GroupBy | Shuffle = gb_pwise
     if post_aggregation_count > 1:
         if ir.maintain_order:  # pragma: no cover
-            raise NotImplementedError(
-                "maintain_order not supported for multiple output partitions."
+            return _lower_ir_fallback(
+                ir,
+                rec,
+                msg="maintain_order not supported for multiple output partitions.",
             )
 
         gb_inter = Shuffle(
diff --git a/python/cudf_polars/cudf_polars/experimental/join.py b/python/cudf_polars/cudf_polars/experimental/join.py
index 58879bb8122..40d9c51bb24 100644
--- a/python/cudf_polars/cudf_polars/experimental/join.py
+++ b/python/cudf_polars/cudf_polars/experimental/join.py
@@ -9,9 +9,10 @@
 from typing import TYPE_CHECKING, Any
 
 from cudf_polars.dsl.ir import Join
-from cudf_polars.experimental.base import PartitionInfo, _concat, get_key_name
+from cudf_polars.experimental.base import PartitionInfo, get_key_name
 from cudf_polars.experimental.dispatch import generate_ir_tasks, lower_ir_node
 from cudf_polars.experimental.shuffle import Shuffle, _partition_dataframe
+from cudf_polars.experimental.utils import _concat, _lower_ir_fallback
 
 if TYPE_CHECKING:
     from collections.abc import MutableMapping
@@ -197,10 +198,10 @@ def _(
         new_node = ir.reconstruct(children)
         partition_info[new_node] = PartitionInfo(count=1)
         return new_node, partition_info
-    elif ir.options[0] == "Cross":
-        raise NotImplementedError(
-            "Cross join not support for multiple partitions."
-        )  # pragma: no cover
+    elif ir.options[0] == "Cross":  # pragma: no cover
+        return _lower_ir_fallback(
+            ir, rec, msg="Cross join not support for multiple partitions."
+        )
 
     if _should_bcast_join(ir, left, right, partition_info, output_count):
         # Create a broadcast join
diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py
index 819913ccc0f..2170183bbb1 100644
--- a/python/cudf_polars/cudf_polars/experimental/parallel.py
+++ b/python/cudf_polars/cudf_polars/experimental/parallel.py
@@ -16,11 +16,13 @@
 import cudf_polars.experimental.shuffle  # noqa: F401
 from cudf_polars.dsl.ir import IR, Cache, Filter, HStack, Projection, Select, Union
 from cudf_polars.dsl.traversal import CachingVisitor, traversal
-from cudf_polars.experimental.base import PartitionInfo, _concat, get_key_name
+from cudf_polars.experimental.base import PartitionInfo, get_key_name
 from cudf_polars.experimental.dispatch import (
     generate_ir_tasks,
     lower_ir_node,
 )
+from cudf_polars.experimental.utils import _concat, _lower_ir_fallback
+from cudf_polars.utils.config import ConfigOptions
 
 if TYPE_CHECKING:
     from collections.abc import MutableMapping
@@ -68,24 +70,14 @@ def _(ir: IR, rec: LowerIRTransformer) -> tuple[IR, MutableMapping[IR, Partition
             ir: PartitionInfo(count=1)
         }  # pragma: no cover; Missed by pylibcudf executor
 
-    # Lower children
-    children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
-    partition_info = reduce(operator.or_, _partition_info)
-
-    # Check that child partitioning is supported
-    if any(partition_info[c].count > 1 for c in children):
-        raise NotImplementedError(
-            f"Class {type(ir)} does not support multiple partitions."
-        )  # pragma: no cover
-
-    # Return reconstructed node and partition-info dict
-    partition = PartitionInfo(count=1)
-    new_node = ir.reconstruct(children)
-    partition_info[new_node] = partition
-    return new_node, partition_info
+    return _lower_ir_fallback(
+        ir, rec, msg=f"Class {type(ir)} does not support multiple partitions."
+    )
 
 
-def lower_ir_graph(ir: IR) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+def lower_ir_graph(
+    ir: IR, config_options: ConfigOptions | None = None
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     """
     Rewrite an IR graph and extract partitioning information.
 
@@ -93,6 +85,8 @@ def lower_ir_graph(ir: IR) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     ----------
     ir
         Root of the graph to rewrite.
+    config_options
+        GPUEngine configuration options.
 
     Returns
     -------
@@ -109,7 +103,8 @@ def lower_ir_graph(ir: IR) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     --------
     lower_ir_node
     """
-    mapper = CachingVisitor(lower_ir_node)
+    config_options = config_options or ConfigOptions({})
+    mapper = CachingVisitor(lower_ir_node, state={"config_options": config_options})
     return mapper(ir)
 
 
@@ -180,9 +175,9 @@ def get_client() -> Any:
         return client.get
 
 
-def evaluate_dask(ir: IR) -> DataFrame:
-    """Evaluate an IR graph with Dask."""
-    ir, partition_info = lower_ir_graph(ir)
+def evaluate_dask(ir: IR, config_options: ConfigOptions | None = None) -> DataFrame:
+    """Evaluate an IR graph with partitioning."""
+    ir, partition_info = lower_ir_graph(ir, config_options)
 
     get = get_client()
 
@@ -223,18 +218,16 @@ def _(
 def _(
     ir: Union, rec: LowerIRTransformer
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Check zlice
+    if ir.zlice is not None:  # pragma: no cover
+        return _lower_ir_fallback(
+            ir, rec, msg="zlice is not supported for multiple partitions."
+        )
+
     # Lower children
     children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
     partition_info = reduce(operator.or_, _partition_info)
 
-    # Check zlice
-    if ir.zlice is not None:  # pragma: no cover
-        if any(p[c].count > 1 for p, c in zip(children, _partition_info, strict=False)):
-            raise NotImplementedError("zlice is not supported for multiple partitions.")
-        new_node = ir.reconstruct(children)
-        partition_info[new_node] = PartitionInfo(count=1)
-        return new_node, partition_info
-
     # Partition count is the sum of all child partitions
     count = sum(partition_info[c].count for c in children)
 
@@ -268,10 +261,12 @@ def _lower_ir_pwise(
     counts = {partition_info[c].count for c in children}
 
     # Check that child partitioning is supported
-    if len(counts) > 1:
-        raise NotImplementedError(
-            f"Class {type(ir)} does not support unbalanced partitions."
-        )  # pragma: no cover
+    if len(counts) > 1:  # pragma: no cover
+        return _lower_ir_fallback(
+            ir,
+            rec,
+            msg=f"Class {type(ir)} does not support children with mismatched partition counts.",
+        )
 
     # Return reconstructed node and partition-info dict
     partition = PartitionInfo(count=max(counts))
diff --git a/python/cudf_polars/cudf_polars/experimental/repartition.py b/python/cudf_polars/cudf_polars/experimental/repartition.py
new file mode 100644
index 00000000000..c56466842ce
--- /dev/null
+++ b/python/cudf_polars/cudf_polars/experimental/repartition.py
@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Repartitioning Logic."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import IR
+from cudf_polars.experimental.base import get_key_name
+from cudf_polars.experimental.dispatch import generate_ir_tasks
+from cudf_polars.experimental.utils import _concat
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.experimental.parallel import PartitionInfo
+    from cudf_polars.typing import Schema
+
+
+class Repartition(IR):
+    """
+    Repartition a DataFrame.
+
+    Notes
+    -----
+    Repartitioning means that we are not modifying any
+    data, nor are we reordering or shuffling rows. We
+    are only changing the overall partition count. For
+    now, we only support an N -> 1 repartitioning. The
+    new partition count is tracked outside the object
+    (using PartitionInfo).
+    """
+
+    __slots__ = ()
+    _non_child = ("schema",)
+
+    def __init__(self, schema: Schema, df: IR):
+        self.schema = schema
+        self._non_child_args = ()
+        self.children = (df,)
+
+
+@generate_ir_tasks.register(Repartition)
+def _(
+    ir: Repartition, partition_info: MutableMapping[IR, PartitionInfo]
+) -> MutableMapping[Any, Any]:
+    # Repartition an IR node.
+    # Only supports N -> 1 rapartitioning (for now).
+
+    if partition_info[ir].count > 1:  # pragma: no cover
+        raise NotImplementedError(
+            f"Repartition -> {partition_info[ir].count} not supported."
+        )
+
+    (child,) = ir.children
+    key_name = get_key_name(ir)
+    child_name = get_key_name(child)
+    return {
+        (key_name, 0): (
+            _concat,
+            *((child_name, idx) for idx in range(partition_info[child].count)),
+        )
+    }
diff --git a/python/cudf_polars/cudf_polars/experimental/select.py b/python/cudf_polars/cudf_polars/experimental/select.py
index 5f79384b569..74e5506bad5 100644
--- a/python/cudf_polars/cudf_polars/experimental/select.py
+++ b/python/cudf_polars/cudf_polars/experimental/select.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 """Parallel Select Logic."""
 
@@ -9,6 +9,7 @@
 from cudf_polars.dsl.ir import Select
 from cudf_polars.dsl.traversal import traversal
 from cudf_polars.experimental.dispatch import lower_ir_node
+from cudf_polars.experimental.utils import _lower_ir_fallback
 
 if TYPE_CHECKING:
     from collections.abc import MutableMapping
@@ -28,8 +29,8 @@ def _(
         expr.is_pointwise for expr in traversal([e.value for e in ir.exprs])
     ):
         # TODO: Handle non-pointwise expressions.
-        raise NotImplementedError(
-            f"Selection {ir} does not support multiple partitions."
+        return _lower_ir_fallback(
+            ir, rec, msg="This selection not support for multiple partitions."
         )
     new_node = ir.reconstruct([child])
     partition_info[new_node] = pi
diff --git a/python/cudf_polars/cudf_polars/experimental/shuffle.py b/python/cudf_polars/cudf_polars/experimental/shuffle.py
index ccbf0bbe357..e0dcedf8921 100644
--- a/python/cudf_polars/cudf_polars/experimental/shuffle.py
+++ b/python/cudf_polars/cudf_polars/experimental/shuffle.py
@@ -16,8 +16,9 @@
 from cudf_polars.containers import DataFrame
 from cudf_polars.dsl.expr import Col
 from cudf_polars.dsl.ir import IR
-from cudf_polars.experimental.base import _concat, get_key_name
+from cudf_polars.experimental.base import get_key_name
 from cudf_polars.experimental.dispatch import generate_ir_tasks, lower_ir_node
+from cudf_polars.experimental.utils import _concat
 
 if TYPE_CHECKING:
     from collections.abc import MutableMapping, Sequence
diff --git a/python/cudf_polars/cudf_polars/experimental/utils.py b/python/cudf_polars/cudf_polars/experimental/utils.py
new file mode 100644
index 00000000000..a647cecbef5
--- /dev/null
+++ b/python/cudf_polars/cudf_polars/experimental/utils.py
@@ -0,0 +1,76 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-partition utilities."""
+
+from __future__ import annotations
+
+import operator
+import warnings
+from functools import reduce
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.ir import Union
+from cudf_polars.experimental.base import PartitionInfo
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.containers import DataFrame
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+
+
+def _concat(*dfs: DataFrame) -> DataFrame:
+    # Concatenate a sequence of DataFrames vertically
+    return Union.do_evaluate(None, *dfs)
+
+
+def _lower_ir_fallback(
+    ir: IR,
+    rec: LowerIRTransformer,
+    *,
+    msg: str | None = None,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Catch-all single-partition lowering logic.
+    # If any children contain multiple partitions,
+    # those children will be collapsed with `Repartition`.
+    from cudf_polars.experimental.repartition import Repartition
+
+    # Lower children
+    lowered_children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
+    partition_info = reduce(operator.or_, _partition_info)
+
+    # Ensure all children are single-partitioned
+    children = []
+    fallback = False
+    for c in lowered_children:
+        child = c
+        if partition_info[c].count > 1:
+            # Fall-back logic
+            fallback = True
+            child = Repartition(child.schema, child)
+            partition_info[child] = PartitionInfo(count=1)
+        children.append(child)
+
+    if fallback and msg:
+        # Warn/raise the user if any children were collapsed
+        # and the "fallback_mode" configuration is not "silent"
+        match fallback_mode := rec.state["config_options"].get(
+            "executor_options.fallback_mode", default="warn"
+        ):
+            case "warn":
+                warnings.warn(msg, stacklevel=2)
+            case "raise":
+                raise NotImplementedError(msg)
+            case "silent":
+                pass
+            case _:
+                raise ValueError(
+                    f"{fallback_mode} is not a supported 'fallback_mode' "
+                    "option. Please use 'warn', 'raise', or 'silent'."
+                )
+
+    # Reconstruct and return
+    new_node = ir.reconstruct(children)
+    partition_info[new_node] = PartitionInfo(count=1)
+    return new_node, partition_info
diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py
index 4e508749da1..aadc8347496 100644
--- a/python/cudf_polars/cudf_polars/utils/config.py
+++ b/python/cudf_polars/cudf_polars/utils/config.py
@@ -124,6 +124,7 @@ def validate(config: dict) -> None:
         executor = config.get("executor", "pylibcudf")
         if executor == "dask-experimental":
             unsupported = config.get("executor_options", {}).keys() - {
+                "fallback_mode",
                 "max_rows_per_partition",
                 "parquet_blocksize",
                 "cardinality_factor",
diff --git a/python/cudf_polars/tests/experimental/test_groupby.py b/python/cudf_polars/tests/experimental/test_groupby.py
index b0c8b6d66fa..61ac70f734e 100644
--- a/python/cudf_polars/tests/experimental/test_groupby.py
+++ b/python/cudf_polars/tests/experimental/test_groupby.py
@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+import contextlib
+
 import pytest
 
 import polars as pl
@@ -80,12 +82,31 @@ def test_groupby_agg_config_options(df, op, keys):
     assert_gpu_result_equal(q, engine=engine, check_row_order=False)
 
 
-def test_groupby_raises(df, engine):
+@pytest.mark.parametrize("fallback_mode", ["silent", "raise", "warn", "foo"])
+def test_groupby_fallback(df, engine, fallback_mode):
+    engine = pl.GPUEngine(
+        raise_on_fail=True,
+        executor="dask-experimental",
+        executor_options={
+            "fallback_mode": fallback_mode,
+            "max_rows_per_partition": 4,
+        },
+    )
+    match = "Failed to decompose groupby aggs"
+
     q = df.group_by("y").median()
-    with pytest.raises(
-        pl.exceptions.ComputeError,
-        match="NotImplementedError",
-    ):
+
+    if fallback_mode == "silent":
+        ctx = contextlib.nullcontext()
+    elif fallback_mode == "raise":
+        ctx = pytest.raises(pl.exceptions.ComputeError, match=match)
+    elif fallback_mode == "foo":
+        ctx = pytest.raises(
+            pl.exceptions.ComputeError, match="not a supported 'fallback_mode' option"
+        )
+    else:
+        ctx = pytest.warns(UserWarning, match=match)
+    with ctx:
         assert_gpu_result_equal(q, engine=engine, check_row_order=False)
 
 
diff --git a/python/cudf_polars/tests/experimental/test_select.py b/python/cudf_polars/tests/experimental/test_select.py
index 7dfe6ead148..47982e9d1fe 100644
--- a/python/cudf_polars/tests/experimental/test_select.py
+++ b/python/cudf_polars/tests/experimental/test_select.py
@@ -1,8 +1,10 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
 
+import contextlib
+
 import pytest
 
 import polars as pl
@@ -36,15 +38,34 @@ def test_select(df, engine):
     assert_gpu_result_equal(query, engine=engine)
 
 
-def test_select_reduce_raises(df, engine):
+@pytest.mark.parametrize("fallback_mode", ["silent", "raise", "warn", "foo"])
+def test_select_reduce_fallback(df, fallback_mode):
+    engine = pl.GPUEngine(
+        raise_on_fail=True,
+        executor="dask-experimental",
+        executor_options={
+            "fallback_mode": fallback_mode,
+            "max_rows_per_partition": 3,
+        },
+    )
+    match = "This selection not support for multiple partitions."
+
     query = df.select(
         (pl.col("a") + pl.col("b")).max(),
         (pl.col("a") * 2 + pl.col("b")).alias("d").mean(),
     )
-    with pytest.raises(
-        pl.exceptions.ComputeError,
-        match="NotImplementedError",
-    ):
+
+    if fallback_mode == "silent":
+        ctx = contextlib.nullcontext()
+    elif fallback_mode == "raise":
+        ctx = pytest.raises(pl.exceptions.ComputeError, match=match)
+    elif fallback_mode == "foo":
+        ctx = pytest.raises(
+            pl.exceptions.ComputeError, match="not a supported 'fallback_mode' option"
+        )
+    else:
+        ctx = pytest.warns(UserWarning, match=match)
+    with ctx:
         assert_gpu_result_equal(query, engine=engine)
 
 


Reviews:
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-01T21:08:34Z
user: login: TomAugspurger
body: This ended up looking pretty nice, thanks.

I think having a global configuration for how to handle the fallback would be valuable.

```python
engine = pl.GPUEngine(executor="partitioned", fallback="warn"/True/False)
```

or something like that. Then we can configure it to `False` when testing that something definitely doesn't fall back, and users can set it to `True` if they get tired of the warnings.


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T12:51:35Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T14:11:14Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T14:18:49Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T17:03:15Z
user: login: TomAugspurger
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T19:40:24Z
user: login: wence-
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T11:43:54Z
user: login: wence-
body: Some small suggestions. I think this looks pretty good
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T13:41:17Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-04T02:39:26Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-04T14:37:59Z
user: login: TomAugspurger
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T16:30:54Z
user: login: Matt711
body: Thanks @rjzamora, I had a few non-blocking questions.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T22:14:14Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T03:05:02Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T03:05:29Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T03:06:15Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T03:15:30Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T13:44:52Z
user: login: Matt711
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-08T13:46:07Z
user: login: wence-
body: Tiny nits. Thanks Rick
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-08T14:43:18Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T14:46:27Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T14:49:34Z
user: login: rjzamora
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T14:53:33Z

Review comments:
diff_hunk: @@ -42,8 +39,3 @@ def keys(self, node: Node) -> Iterator[tuple[str, int]]:
 def get_key_name(node: Node) -> str:
     """Generate the key name for a Node."""
     return f"{type(node).__name__.lower()}-{hash(node)}"
-
-
-def _concat(*dfs: DataFrame) -> DataFrame:
-    # Concatenate a sequence of DataFrames vertically
-    return Union.do_evaluate(None, *dfs)
path: python/cudf_polars/cudf_polars/experimental/base.py
user: login: rjzamora
body: We don't need to move `_concat` to `utils.py` - It just "felt" like it belonged there.
created_at: 2025-04-01T21:08:34Z
author_association: MEMBER
diff_hunk: @@ -42,8 +39,3 @@ def keys(self, node: Node) -> Iterator[tuple[str, int]]:
 def get_key_name(node: Node) -> str:
     """Generate the key name for a Node."""
     return f"{type(node).__name__.lower()}-{hash(node)}"
-
-
-def _concat(*dfs: DataFrame) -> DataFrame:
-    # Concatenate a sequence of DataFrames vertically
-    return Union.do_evaluate(None, *dfs)
path: python/cudf_polars/cudf_polars/experimental/base.py
user: login: TomAugspurger
body: Agreed!
created_at: 2025-04-01T21:32:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -198,10 +199,14 @@ def _(
         return single_part_node, partition_info
 
     # Check group-by keys
-    if not all(expr.is_pointwise for expr in traversal([e.value for e in ir.keys])):
-        raise NotImplementedError(
-            f"GroupBy does not support multiple partitions for keys:\n{ir.keys}"
-        )  # pragma: no cover
+    if not all(
+        expr.is_pointwise for expr in traversal([e.value for e in ir.keys])
+    ):  # pragma: no cover
+        return _lower_ir_fallback(
+            ir,
+            rec,
+            msg=f"GroupBy does not support multiple partitions for keys:\n{ir.keys}",
path: python/cudf_polars/cudf_polars/experimental/groupby.py
user: login: TomAugspurger
body: ```suggestion
            msg=f"group_by does not support multiple partitions for keys:\n{ir.keys}",
```

to match the public Polars API?
created_at: 2025-04-02T11:56:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -226,9 +231,14 @@ def _(
         )
 
     # Decompose the aggregation requests into three distinct phases
-    selection_exprs, piecewise_exprs, reduction_exprs = combine(
-        *(decompose(agg.name, agg.value) for agg in ir.agg_requests)
-    )
+    try:
path: python/cudf_polars/cudf_polars/experimental/groupby.py
user: login: TomAugspurger
body: I haven't looked at combine / decompose yet, but just confirming: is a `try / except NotImplemented` needed here? Or can we maybe push this fallback down into `decompose`?

My question is mainly motivated by my aversion to using exceptions for control flow. But perhaps it's the best option in this case.
created_at: 2025-04-02T11:57:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -229,11 +218,9 @@ def _(
 
     # Check zlice
     if ir.zlice is not None:  # pragma: no cover
-        if any(p[c].count > 1 for p, c in zip(children, _partition_info, strict=False)):
-            raise NotImplementedError("zlice is not supported for multiple partitions.")
-        new_node = ir.reconstruct(children)
-        partition_info[new_node] = PartitionInfo(count=1)
-        return new_node, partition_info
+        return _lower_ir_fallback(
path: python/cudf_polars/cudf_polars/experimental/parallel.py
user: login: TomAugspurger
body: I see the `if any(p[c].count > 1...)` condition was removed, and we're always using `_lower_ir_fallback`. Is that deliberate?

IIUC, the outcome is the same either way (a new node with `PartitionInfo(count=1)`. Only, won't we emit an unnecessary warning in the case where all the children partition infos have `count=1`?
created_at: 2025-04-02T12:50:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -226,9 +231,14 @@ def _(
         )
 
     # Decompose the aggregation requests into three distinct phases
-    selection_exprs, piecewise_exprs, reduction_exprs = combine(
-        *(decompose(agg.name, agg.value) for agg in ir.agg_requests)
-    )
+    try:
path: python/cudf_polars/cudf_polars/experimental/groupby.py
user: login: rjzamora
body: I can understand your aversion to these kinds of exceptions.

I did want to trigger the `_lower_ir_fallback` within `decompose`, but that function requires `IR` and `LowerIRTransformer` arguments that aren't available within decompose. I don't think we want to plumb `ir` and `rec` into `decompose`.

Maybe the try/except approach is a bit cleaner if we define a custom `MultiPartitionSupport` exception for this?
created_at: 2025-04-02T14:11:14Z
author_association: MEMBER
diff_hunk: @@ -229,11 +218,9 @@ def _(
 
     # Check zlice
     if ir.zlice is not None:  # pragma: no cover
-        if any(p[c].count > 1 for p, c in zip(children, _partition_info, strict=False)):
-            raise NotImplementedError("zlice is not supported for multiple partitions.")
-        new_node = ir.reconstruct(children)
-        partition_info[new_node] = PartitionInfo(count=1)
-        return new_node, partition_info
+        return _lower_ir_fallback(
path: python/cudf_polars/cudf_polars/experimental/parallel.py
user: login: rjzamora
body: >I see the if any(p[c].count > 1...) condition was removed, and we're always using _lower_ir_fallback. Is that deliberate?

Yes, the if `any(p[c].count > 1...)` condition is exactly the case that `_lower_ir_fallback` will collapse the child(ren) and emit a warning. This means that the entire `if ir.zlice is not None:` block will behave exactly the same as `_lower_ir_fallback` in the case that `any(p[c].count > 1` is False. Therefore, there is really no reason for us to check that condition here.

Note that we should also move the `if ir.zlice is not None:` block to the top of this function, since we don't need to lower the children here if we are just calling `_lower_ir_fallback` anyway.
created_at: 2025-04-02T14:18:49Z
author_association: MEMBER
diff_hunk: @@ -109,7 +103,8 @@ def lower_ir_graph(ir: IR) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     --------
     lower_ir_node
     """
-    mapper = CachingVisitor(lower_ir_node)
+    config_options = config_options or ConfigOptions({})
+    mapper = CachingVisitor(lower_ir_node, state={"config_options": config_options})
path: python/cudf_polars/cudf_polars/experimental/parallel.py
user: login: rjzamora
body: cc @wence- 

In order to make the fallback behavior "configurable", we probably want to be able to access the `GPUEngine` options for any `IR` node without needing to explicitly pass those options into the IR object at translation time.

This seems like one reasonable way to do this?
created_at: 2025-04-02T17:03:14Z
author_association: MEMBER
diff_hunk: @@ -226,9 +231,14 @@ def _(
         )
 
     # Decompose the aggregation requests into three distinct phases
-    selection_exprs, piecewise_exprs, reduction_exprs = combine(
-        *(decompose(agg.name, agg.value) for agg in ir.agg_requests)
-    )
+    try:
path: python/cudf_polars/cudf_polars/experimental/groupby.py
user: login: TomAugspurger
body: Makes sense, thanks.

> Maybe the try/except approach is a bit cleaner if we define a custom MultiPartitionSupport exception for this?

I think they're about equivalent. Hopefully we're the only ones raising `NotImplemented` here...
created_at: 2025-04-02T19:40:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -109,7 +103,8 @@ def lower_ir_graph(ir: IR) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     --------
     lower_ir_node
     """
-    mapper = CachingVisitor(lower_ir_node)
+    config_options = config_options or ConfigOptions({})
+    mapper = CachingVisitor(lower_ir_node, state={"config_options": config_options})
path: python/cudf_polars/cudf_polars/experimental/parallel.py
user: login: wence-
body: Yes, this is what the state is for.
created_at: 2025-04-03T11:43:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Repartitioning Logic."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import IR
+from cudf_polars.experimental.base import get_key_name
+from cudf_polars.experimental.dispatch import generate_ir_tasks
+from cudf_polars.experimental.utils import _concat
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.experimental.parallel import PartitionInfo
+    from cudf_polars.typing import Schema
+
+
+class Repartition(IR):
+    """Repartition a DataFrame."""
+
+    __slots__ = ()
+    _non_child = ("schema",)
+
+    def __init__(self, schema: Schema, df: IR):
+        self.schema = schema
+        self._non_child_args = ()
+        self.children = (df,)
path: python/cudf_polars/cudf_polars/experimental/repartition.py
user: login: wence-
body: *question*: Right now we're using `Repartition` as a tag to mean "concatenate to a single partition". Should we already accept a `num_partitions` argument (which we raise if it's not `1` for)
created_at: 2025-04-03T13:03:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Repartitioning Logic."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import IR
+from cudf_polars.experimental.base import get_key_name
+from cudf_polars.experimental.dispatch import generate_ir_tasks
+from cudf_polars.experimental.utils import _concat
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.experimental.parallel import PartitionInfo
+    from cudf_polars.typing import Schema
+
+
+class Repartition(IR):
+    """Repartition a DataFrame."""
+
+    __slots__ = ()
+    _non_child = ("schema",)
+
+    def __init__(self, schema: Schema, df: IR):
+        self.schema = schema
+        self._non_child_args = ()
+        self.children = (df,)
+
+
+@generate_ir_tasks.register(Repartition)
+def _(
+    ir: Repartition, partition_info: MutableMapping[IR, PartitionInfo]
+) -> MutableMapping[Any, Any]:
+    # Repartition an IR node.
+    # Only supports N -> 1 rapartitioning (for now).
+
+    if partition_info[ir].count > 1:  # pragma: no cover
+        raise NotImplementedError(
+            f"Repartition -> {partition_info[ir].count} not supported."
+        )
+
+    (child,) = ir.children
+    key_name = get_key_name(ir)
+    child_name = get_key_name(child)
+    return {
+        (key_name, 0): (
+            _concat,
+            *[(child_name, idx) for idx in range(partition_info[child].count)],
path: python/cudf_polars/cudf_polars/experimental/repartition.py
user: login: wence-
body: ```suggestion
            *((child_name, idx) for idx in range(partition_info[child].count)),
```
No need to cons the list before splatting the generator
created_at: 2025-04-03T13:04:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-partition utilities."""
+
+from __future__ import annotations
+
+import operator
+import warnings
+from functools import reduce
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.ir import Union
+from cudf_polars.experimental.base import PartitionInfo
+from cudf_polars.utils.config import ConfigOptions
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.containers import DataFrame
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+
+
+def _concat(*dfs: DataFrame) -> DataFrame:
+    # Concatenate a sequence of DataFrames vertically
+    return Union.do_evaluate(None, *dfs)
+
+
+def _lower_ir_fallback(
+    ir: IR,
+    rec: LowerIRTransformer,
+    *,
+    msg: str | None = None,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Catch-all single-partition lowering logic.
+    # If any children contain multiple partitions,
+    # those children will be collapsed with `Repartition`.
+    from cudf_polars.experimental.repartition import Repartition
+
+    # Lower children
+    _children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
+    partition_info = reduce(operator.or_, _partition_info)
+
+    # Ensure all children are single-partitioned
+    children = []
+    fallback = False
+    for c in _children:
+        child = c
+        if partition_info[c].count > 1:
+            # Fall-back logic
+            fallback = True
+            child = Repartition(child.schema, child)
+            partition_info[child] = PartitionInfo(count=1)
+        children.append(child)
+
+    if fallback and msg:
+        # Warn/raise the user if any children were collapsed
+        # and the "fallback_mode" configuration is not "silent"
+        fallback_mode = rec.state.get("config_options", ConfigOptions({})).get(
+            "executor_options.fallback_mode", default="warn"
+        )
path: python/cudf_polars/cudf_polars/experimental/utils.py
user: login: wence-
body: Can we arrange that the state _always_ has config_options so that this is simpler.
created_at: 2025-04-03T13:34:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-partition utilities."""
+
+from __future__ import annotations
+
+import operator
+import warnings
+from functools import reduce
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.ir import Union
+from cudf_polars.experimental.base import PartitionInfo
+from cudf_polars.utils.config import ConfigOptions
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.containers import DataFrame
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+
+
+def _concat(*dfs: DataFrame) -> DataFrame:
+    # Concatenate a sequence of DataFrames vertically
+    return Union.do_evaluate(None, *dfs)
+
+
+def _lower_ir_fallback(
+    ir: IR,
+    rec: LowerIRTransformer,
+    *,
+    msg: str | None = None,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Catch-all single-partition lowering logic.
+    # If any children contain multiple partitions,
+    # those children will be collapsed with `Repartition`.
+    from cudf_polars.experimental.repartition import Repartition
+
+    # Lower children
+    _children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
+    partition_info = reduce(operator.or_, _partition_info)
+
+    # Ensure all children are single-partitioned
+    children = []
+    fallback = False
+    for c in _children:
+        child = c
+        if partition_info[c].count > 1:
+            # Fall-back logic
+            fallback = True
+            child = Repartition(child.schema, child)
+            partition_info[child] = PartitionInfo(count=1)
+        children.append(child)
+
+    if fallback and msg:
+        # Warn/raise the user if any children were collapsed
+        # and the "fallback_mode" configuration is not "silent"
+        fallback_mode = rec.state.get("config_options", ConfigOptions({})).get(
+            "executor_options.fallback_mode", default="warn"
+        )
+        if fallback_mode == "warn":
+            warnings.warn(msg, stacklevel=2)
+        elif fallback_mode == "raise":
+            raise NotImplementedError(msg)
+        elif fallback_mode != "silent":  # pragma: no cover
+            raise ValueError(
+                f"{fallback_mode} is not a supported 'fallback_mode' option. "
+                "Please use 'warn', 'raise', or 'silent'."
+            )
path: python/cudf_polars/cudf_polars/experimental/utils.py
user: login: wence-
body: Please test this.
created_at: 2025-04-03T13:34:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Repartitioning Logic."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import IR
+from cudf_polars.experimental.base import get_key_name
+from cudf_polars.experimental.dispatch import generate_ir_tasks
+from cudf_polars.experimental.utils import _concat
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.experimental.parallel import PartitionInfo
+    from cudf_polars.typing import Schema
+
+
+class Repartition(IR):
+    """Repartition a DataFrame."""
+
+    __slots__ = ()
+    _non_child = ("schema",)
+
+    def __init__(self, schema: Schema, df: IR):
+        self.schema = schema
+        self._non_child_args = ()
+        self.children = (df,)
path: python/cudf_polars/cudf_polars/experimental/repartition.py
user: login: rjzamora
body: So far, we have avoided attaching explicit partition-count information to the IR objects themselves. Instead, we refer to the extracted PartitionInfo object to decide how many output partitions this is generating.
created_at: 2025-04-04T02:39:26Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-partition utilities."""
+
+from __future__ import annotations
+
+import operator
+import warnings
+from functools import reduce
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.ir import Union
+from cudf_polars.experimental.base import PartitionInfo
+from cudf_polars.utils.config import ConfigOptions
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.containers import DataFrame
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+
+
+def _concat(*dfs: DataFrame) -> DataFrame:
+    # Concatenate a sequence of DataFrames vertically
+    return Union.do_evaluate(None, *dfs)
+
+
+def _lower_ir_fallback(
+    ir: IR,
+    rec: LowerIRTransformer,
+    *,
+    msg: str | None = None,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Catch-all single-partition lowering logic.
+    # If any children contain multiple partitions,
+    # those children will be collapsed with `Repartition`.
+    from cudf_polars.experimental.repartition import Repartition
+
+    # Lower children
+    _children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
+    partition_info = reduce(operator.or_, _partition_info)
+
+    # Ensure all children are single-partitioned
+    children = []
+    fallback = False
+    for c in _children:
+        child = c
+        if partition_info[c].count > 1:
+            # Fall-back logic
+            fallback = True
+            child = Repartition(child.schema, child)
+            partition_info[child] = PartitionInfo(count=1)
+        children.append(child)
+
+    if fallback and msg:
+        # Warn/raise the user if any children were collapsed
+        # and the "fallback_mode" configuration is not "silent"
+        fallback_mode = rec.state.get("config_options", ConfigOptions({})).get(
+            "executor_options.fallback_mode", default="warn"
+        )
path: python/cudf_polars/cudf_polars/experimental/utils.py
user: login: rjzamora
body: Good point. The lowering process is kicked off in `lower_ir_graph`, where we **always** attach `config_options`. So, we don't need this default logic.
created_at: 2025-04-04T14:37:59Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Repartitioning Logic."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import IR
+from cudf_polars.experimental.base import get_key_name
+from cudf_polars.experimental.dispatch import generate_ir_tasks
+from cudf_polars.experimental.utils import _concat
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.experimental.parallel import PartitionInfo
+    from cudf_polars.typing import Schema
+
+
+class Repartition(IR):
path: python/cudf_polars/cudf_polars/experimental/repartition.py
user: login: Matt711
body: Is the `Repartition` the right name here? I think something like `Recombine` or `Recollect` sounds closer what this node does.
created_at: 2025-04-07T21:57:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-partition utilities."""
+
+from __future__ import annotations
+
+import operator
+import warnings
+from functools import reduce
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.ir import Union
+from cudf_polars.experimental.base import PartitionInfo
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.containers import DataFrame
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+
+
+def _concat(*dfs: DataFrame) -> DataFrame:
+    # Concatenate a sequence of DataFrames vertically
+    return Union.do_evaluate(None, *dfs)
+
+
+def _lower_ir_fallback(
+    ir: IR,
+    rec: LowerIRTransformer,
+    *,
+    msg: str | None = None,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Catch-all single-partition lowering logic.
+    # If any children contain multiple partitions,
+    # those children will be collapsed with `Repartition`.
+    from cudf_polars.experimental.repartition import Repartition
+
+    # Lower children
+    _children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
path: python/cudf_polars/cudf_polars/experimental/utils.py
user: login: Matt711
body: Just because there's `children` too
```suggestion
    lowered_children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
```
created_at: 2025-04-07T21:58:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-partition utilities."""
+
+from __future__ import annotations
+
+import operator
+import warnings
+from functools import reduce
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.ir import Union
+from cudf_polars.experimental.base import PartitionInfo
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.containers import DataFrame
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+
+
+def _concat(*dfs: DataFrame) -> DataFrame:
+    # Concatenate a sequence of DataFrames vertically
+    return Union.do_evaluate(None, *dfs)
+
+
+def _lower_ir_fallback(
+    ir: IR,
+    rec: LowerIRTransformer,
+    *,
+    msg: str | None = None,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Catch-all single-partition lowering logic.
+    # If any children contain multiple partitions,
+    # those children will be collapsed with `Repartition`.
+    from cudf_polars.experimental.repartition import Repartition
+
+    # Lower children
+    _children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
+    partition_info = reduce(operator.or_, _partition_info)
+
+    # Ensure all children are single-partitioned
+    children = []
+    fallback = False
+    for c in _children:
+        child = c
+        if partition_info[c].count > 1:
+            # Fall-back logic
+            fallback = True
+            child = Repartition(child.schema, child)
+            partition_info[child] = PartitionInfo(count=1)
+        children.append(child)
+
+    if fallback and msg:
+        # Warn/raise the user if any children were collapsed
+        # and the "fallback_mode" configuration is not "silent"
+        fallback_mode = rec.state["config_options"].get(
+            "executor_options.fallback_mode", default="warn"
+        )
+        if fallback_mode == "warn":
+            warnings.warn(msg, stacklevel=2)
+        elif fallback_mode == "raise":
+            raise NotImplementedError(msg)
+        elif fallback_mode != "silent":
+            raise ValueError(
+                f"{fallback_mode} is not a supported 'fallback_mode' option. "
+                "Please use 'warn', 'raise', or 'silent'."
+            )
path: python/cudf_polars/cudf_polars/experimental/utils.py
user: login: Matt711
body: Just a observation: Do you know if python 3.10 is our minimum supported version? This block made me think since `match-case` was added in python 3.10, we maybe could start using it. I'm curious what folks think it it versus `if-else`. 
created_at: 2025-04-07T21:58:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -80,12 +82,31 @@ def test_groupby_agg_config_options(df, op, keys):
     assert_gpu_result_equal(q, engine=engine, check_row_order=False)
 
 
-def test_groupby_raises(df, engine):
+@pytest.mark.parametrize("fallback_mode", ["silent", "raise", "warn", "foo"])
path: python/cudf_polars/tests/experimental/test_groupby.py
user: login: Matt711
body: In our CI we run the cudf-polars tests 3x. See https://github.com/rapidsai/cudf/blob/branch-25.06/ci/run_cudf_polars_pytests.sh#L15

For the second run, when we pass `--executor dask-experimental`. Do we already fallback to non-partitioned cudf-polars if the test has unsupported operations. Eg. What happens when we run conditional join tests with `dask-experimental`?
created_at: 2025-04-07T22:07:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Repartitioning Logic."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import IR
+from cudf_polars.experimental.base import get_key_name
+from cudf_polars.experimental.dispatch import generate_ir_tasks
+from cudf_polars.experimental.utils import _concat
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.experimental.parallel import PartitionInfo
+    from cudf_polars.typing import Schema
+
+
+class Repartition(IR):
path: python/cudf_polars/cudf_polars/experimental/repartition.py
user: login: rjzamora
body: This is borrowing from the dask definition of "repartition", which means we are not changing the data, nor are we re-ordering or shuffling rows. All we are doing is changing the partition count. For now, we only use this to collapse down to one partition. However, it *could* potentially be used for other repartitioning purposes in the future (not sure how realistic that is).
created_at: 2025-04-08T03:05:02Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-partition utilities."""
+
+from __future__ import annotations
+
+import operator
+import warnings
+from functools import reduce
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.ir import Union
+from cudf_polars.experimental.base import PartitionInfo
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.containers import DataFrame
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+
+
+def _concat(*dfs: DataFrame) -> DataFrame:
+    # Concatenate a sequence of DataFrames vertically
+    return Union.do_evaluate(None, *dfs)
+
+
+def _lower_ir_fallback(
+    ir: IR,
+    rec: LowerIRTransformer,
+    *,
+    msg: str | None = None,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Catch-all single-partition lowering logic.
+    # If any children contain multiple partitions,
+    # those children will be collapsed with `Repartition`.
+    from cudf_polars.experimental.repartition import Repartition
+
+    # Lower children
+    _children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
path: python/cudf_polars/cudf_polars/experimental/utils.py
user: login: rjzamora
body: Yeah, seems reasonable
created_at: 2025-04-08T03:05:29Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-partition utilities."""
+
+from __future__ import annotations
+
+import operator
+import warnings
+from functools import reduce
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.ir import Union
+from cudf_polars.experimental.base import PartitionInfo
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.containers import DataFrame
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+
+
+def _concat(*dfs: DataFrame) -> DataFrame:
+    # Concatenate a sequence of DataFrames vertically
+    return Union.do_evaluate(None, *dfs)
+
+
+def _lower_ir_fallback(
+    ir: IR,
+    rec: LowerIRTransformer,
+    *,
+    msg: str | None = None,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Catch-all single-partition lowering logic.
+    # If any children contain multiple partitions,
+    # those children will be collapsed with `Repartition`.
+    from cudf_polars.experimental.repartition import Repartition
+
+    # Lower children
+    _children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
+    partition_info = reduce(operator.or_, _partition_info)
+
+    # Ensure all children are single-partitioned
+    children = []
+    fallback = False
+    for c in _children:
+        child = c
+        if partition_info[c].count > 1:
+            # Fall-back logic
+            fallback = True
+            child = Repartition(child.schema, child)
+            partition_info[child] = PartitionInfo(count=1)
+        children.append(child)
+
+    if fallback and msg:
+        # Warn/raise the user if any children were collapsed
+        # and the "fallback_mode" configuration is not "silent"
+        fallback_mode = rec.state["config_options"].get(
+            "executor_options.fallback_mode", default="warn"
+        )
+        if fallback_mode == "warn":
+            warnings.warn(msg, stacklevel=2)
+        elif fallback_mode == "raise":
+            raise NotImplementedError(msg)
+        elif fallback_mode != "silent":
+            raise ValueError(
+                f"{fallback_mode} is not a supported 'fallback_mode' option. "
+                "Please use 'warn', 'raise', or 'silent'."
+            )
path: python/cudf_polars/cudf_polars/experimental/utils.py
user: login: rjzamora
body: Sure - I don't mind adopting match-case.
created_at: 2025-04-08T03:06:14Z
author_association: MEMBER
diff_hunk: @@ -80,12 +82,31 @@ def test_groupby_agg_config_options(df, op, keys):
     assert_gpu_result_equal(q, engine=engine, check_row_order=False)
 
 
-def test_groupby_raises(df, engine):
+@pytest.mark.parametrize("fallback_mode", ["silent", "raise", "warn", "foo"])
path: python/cudf_polars/tests/experimental/test_groupby.py
user: login: rjzamora
body: You are right that we run the tests three times:

1. "dask-experimental" executor (all tests)
2. "pylibcudf" executor(all tests)
3. "dask-experimental" + Distributed (only experimental tests)

The "dask-experimental" executor (i.e. "partitioned" cudf-polars) works for all features in cudf-polars as long as there is only one partition. All tests outside the `tests/experimental` directory should work without the need for fallback, because they will only use a single-partition anyway. The tests inside `tests/experimental` are specifically designed to test multiple partitions.

Keep in mind that the "fallback" mechanism does not actually switch from the "dask-experimental" executor to the "pylibcudf" executor. Rather, we simply collapse multi-partition DataFrames so that the single-partition algorithm will produce correct results.
created_at: 2025-04-08T03:15:30Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Repartitioning Logic."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import IR
+from cudf_polars.experimental.base import get_key_name
+from cudf_polars.experimental.dispatch import generate_ir_tasks
+from cudf_polars.experimental.utils import _concat
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.experimental.parallel import PartitionInfo
+    from cudf_polars.typing import Schema
+
+
+class Repartition(IR):
path: python/cudf_polars/cudf_polars/experimental/repartition.py
user: login: rjzamora
body: Added notes to the `Repartition` doc-string for now.
created_at: 2025-04-08T13:44:51Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Repartitioning Logic."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import IR
+from cudf_polars.experimental.base import get_key_name
+from cudf_polars.experimental.dispatch import generate_ir_tasks
+from cudf_polars.experimental.utils import _concat
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.experimental.parallel import PartitionInfo
+    from cudf_polars.typing import Schema
+
+
+class Repartition(IR):
path: python/cudf_polars/cudf_polars/experimental/repartition.py
user: login: Matt711
body: > All we are doing is changing the partition count. For now, we only use this to collapse down to one partition. However, it could potentially be used for other repartitioning purposes in the future (not sure how realistic that is).

Thanks, that makes sense
created_at: 2025-04-08T13:46:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -197,10 +198,10 @@ def _(
         new_node = ir.reconstruct(children)
         partition_info[new_node] = PartitionInfo(count=1)
         return new_node, partition_info
-    elif ir.options[0] == "Cross":
-        raise NotImplementedError(
-            "Cross join not support for multiple partitions."
-        )  # pragma: no cover
+    elif ir.options[0] == "Cross":  # pragma: no cover
+        return _lower_ir_fallback(
+            ir, rec, msg="Cross join not support for multiple partitions."
+        )
path: python/cudf_polars/cudf_polars/experimental/join.py
user: login: wence-
body: Can we be slightly smarter here? If one of the two tables in a cross-join is single-partition I think we can do the partition-wise cross-join like a broadcast join. That is probably worth it.
created_at: 2025-04-08T13:46:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -268,10 +261,10 @@ def _lower_ir_pwise(
     counts = {partition_info[c].count for c in children}
 
     # Check that child partitioning is supported
-    if len(counts) > 1:
-        raise NotImplementedError(
-            f"Class {type(ir)} does not support unbalanced partitions."
-        )  # pragma: no cover
+    if len(counts) > 1:  # pragma: no cover
+        return _lower_ir_fallback(
+            ir, rec, msg=f"Class {type(ir)} does not support unbalanced partitions."
path: python/cudf_polars/cudf_polars/experimental/parallel.py
user: login: wence-
body: *nit*: "unbalanced" is an odd word here. Do you mean "multiple"?
created_at: 2025-04-08T14:38:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -229,11 +218,9 @@ def _(
 
     # Check zlice
     if ir.zlice is not None:  # pragma: no cover
-        if any(p[c].count > 1 for p, c in zip(children, _partition_info, strict=False)):
-            raise NotImplementedError("zlice is not supported for multiple partitions.")
-        new_node = ir.reconstruct(children)
-        partition_info[new_node] = PartitionInfo(count=1)
-        return new_node, partition_info
+        return _lower_ir_fallback(
path: python/cudf_polars/cudf_polars/experimental/parallel.py
user: login: wence-
body: Just to note. I'd love to find a way to do the minimum amount of work for slices like this (which might appears as `head` or `tail` for example).
created_at: 2025-04-08T14:42:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -197,10 +198,10 @@ def _(
         new_node = ir.reconstruct(children)
         partition_info[new_node] = PartitionInfo(count=1)
         return new_node, partition_info
-    elif ir.options[0] == "Cross":
-        raise NotImplementedError(
-            "Cross join not support for multiple partitions."
-        )  # pragma: no cover
+    elif ir.options[0] == "Cross":  # pragma: no cover
+        return _lower_ir_fallback(
+            ir, rec, msg="Cross join not support for multiple partitions."
+        )
path: python/cudf_polars/cudf_polars/experimental/join.py
user: login: rjzamora
body: Good point - I'm making a list of smaller feature requests to file as follow-up issues. I'll raise one to add broadcast-join support for `"Cross"`.
created_at: 2025-04-08T14:46:27Z
author_association: MEMBER
diff_hunk: @@ -229,11 +218,9 @@ def _(
 
     # Check zlice
     if ir.zlice is not None:  # pragma: no cover
-        if any(p[c].count > 1 for p, c in zip(children, _partition_info, strict=False)):
-            raise NotImplementedError("zlice is not supported for multiple partitions.")
-        new_node = ir.reconstruct(children)
-        partition_info[new_node] = PartitionInfo(count=1)
-        return new_node, partition_info
+        return _lower_ir_fallback(
path: python/cudf_polars/cudf_polars/experimental/parallel.py
user: login: rjzamora
body: Right - We should definitely take advantage of `zlice` options that translate to a head or tail. Especially after a sort. This will definitely be another follow-up issue.
created_at: 2025-04-08T14:49:34Z
author_association: MEMBER
diff_hunk: @@ -268,10 +261,10 @@ def _lower_ir_pwise(
     counts = {partition_info[c].count for c in children}
 
     # Check that child partitioning is supported
-    if len(counts) > 1:
-        raise NotImplementedError(
-            f"Class {type(ir)} does not support unbalanced partitions."
-        )  # pragma: no cover
+    if len(counts) > 1:  # pragma: no cover
+        return _lower_ir_fallback(
+            ir, rec, msg=f"Class {type(ir)} does not support unbalanced partitions."
path: python/cudf_polars/cudf_polars/experimental/parallel.py
user: login: rjzamora
body: This is actually meant to detect the case that we have multiple children, and those children have an unaligned partition count. I'm not sure if/when this might happen, but I'm guarding against it for now.  Maybe ~"unaligned" partitions is better?~ (I'll just say "does not support children with mismatched partition counts" to be as clear as possible.)
created_at: 2025-04-08T14:53:32Z
author_association: MEMBER

PR comments:
user: login: copy-pr-bot[bot]
created_at: 2025-04-01T21:07:17Z
body: Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message [here](https://docs.gha-runners.nvidia.com/cpr/auto-sync).

author_association: NONE
user: login: rjzamora
created_at: 2025-04-01T21:07:43Z
body: /ok to test
author_association: MEMBER
user: login: rjzamora
created_at: 2025-04-02T18:44:57Z
body: /ok to test
author_association: MEMBER
user: login: rjzamora
created_at: 2025-04-04T15:53:45Z
body: Thanks for the preliminary reviews @TomAugspurger & @wence- ! I made a few revisions.
author_association: MEMBER
user: login: rjzamora
created_at: 2025-04-08T16:45:47Z
body: /merge
author_association: MEMBER