cudf-polars-cu12 25.2.2__py3-none-any.whl → 25.6.0__py3-none-any.whl

cudf_polars/experimental/join.py

@@ -0,0 +1,353 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Parallel Join Logic."""
+
+from __future__ import annotations
+
+import operator
+from functools import reduce
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import ConditionalJoin, Join
+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.repartition import Repartition
+from cudf_polars.experimental.shuffle import Shuffle, _partition_dataframe
+from cudf_polars.experimental.utils import _concat, _fallback_inform, _lower_ir_fallback
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.dsl.expr import NamedExpr
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.parallel import LowerIRTransformer
+    from cudf_polars.utils.config import ConfigOptions
+
+
+def _maybe_shuffle_frame(
+    frame: IR,
+    on: tuple[NamedExpr, ...],
+    partition_info: MutableMapping[IR, PartitionInfo],
+    config_options: ConfigOptions,
+    output_count: int,
+) -> IR:
+    # Shuffle `frame` if it isn't already shuffled.
+    if (
+        partition_info[frame].partitioned_on == on
+        and partition_info[frame].count == output_count
+    ):
+        # Already shuffled
+        return frame
+    else:
+        # Insert new Shuffle node
+        frame = Shuffle(
+            frame.schema,
+            on,
+            config_options,
+            frame,
+        )
+        partition_info[frame] = PartitionInfo(
+            count=output_count,
+            partitioned_on=on,
+        )
+        return frame
+
+
+def _make_hash_join(
+    ir: Join,
+    output_count: int,
+    partition_info: MutableMapping[IR, PartitionInfo],
+    left: IR,
+    right: IR,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Shuffle left and right dataframes (if necessary)
+    new_left = _maybe_shuffle_frame(
+        left,
+        ir.left_on,
+        partition_info,
+        ir.config_options,
+        output_count,
+    )
+    new_right = _maybe_shuffle_frame(
+        right,
+        ir.right_on,
+        partition_info,
+        ir.config_options,
+        output_count,
+    )
+    if left != new_left or right != new_right:
+        ir = ir.reconstruct([new_left, new_right])
+    left = new_left
+    right = new_right
+
+    # Record new partitioning info
+    partitioned_on: tuple[NamedExpr, ...] = ()
+    if ir.left_on == ir.right_on or (ir.options[0] in ("Left", "Semi", "Anti")):
+        partitioned_on = ir.left_on
+    elif ir.options[0] == "Right":
+        partitioned_on = ir.right_on
+    partition_info[ir] = PartitionInfo(
+        count=output_count,
+        partitioned_on=partitioned_on,
+    )
+
+    return ir, partition_info
+
+
+def _should_bcast_join(
+    ir: Join,
+    left: IR,
+    right: IR,
+    partition_info: MutableMapping[IR, PartitionInfo],
+    output_count: int,
+) -> bool:
+    # Decide if a broadcast join is appropriate.
+    if partition_info[left].count >= partition_info[right].count:
+        small_count = partition_info[right].count
+        large = left
+        large_on = ir.left_on
+    else:
+        small_count = partition_info[left].count
+        large = right
+        large_on = ir.right_on
+
+    # Avoid the broadcast if the "large" table is already shuffled
+    large_shuffled = (
+        partition_info[large].partitioned_on == large_on
+        and partition_info[large].count == output_count
+    )
+
+    # Broadcast-Join Criteria:
+    # 1. Large dataframe isn't already shuffled
+    # 2. Small dataframe has 8 partitions (or fewer).
+    #    TODO: Make this value/heuristic configurable).
+    #    We may want to account for the number of workers.
+    # 3. The "kind" of join is compatible with a broadcast join
+    assert ir.config_options.executor.name == "streaming", (
+        "'in-memory' executor not supported in 'generate_ir_tasks'"
+    )
+
+    return (
+        not large_shuffled
+        and small_count <= ir.config_options.executor.broadcast_join_limit
+        and (
+            ir.options[0] == "Inner"
+            or (ir.options[0] in ("Left", "Semi", "Anti") and large == left)
+            or (ir.options[0] == "Right" and large == right)
+        )
+    )
+
+
+def _make_bcast_join(
+    ir: Join,
+    output_count: int,
+    partition_info: MutableMapping[IR, PartitionInfo],
+    left: IR,
+    right: IR,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    if ir.options[0] != "Inner":
+        left_count = partition_info[left].count
+        right_count = partition_info[right].count
+
+        # Shuffle the smaller table (if necessary) - Notes:
+        # - We need to shuffle the smaller table if
+        #   (1) we are not doing an "inner" join,
+        #   and (2) the small table contains multiple
+        #   partitions.
+        # - We cannot simply join a large-table partition
+        #   to each small-table partition, and then
+        #   concatenate the partial-join results, because
+        #   a non-"inner" join does NOT commute with
+        #   concatenation.
+        # - In some cases, we can perform the partial joins
+        #   sequentially. However, we are starting with a
+        #   catch-all algorithm that works for all cases.
+        if left_count >= right_count:
+            right = _maybe_shuffle_frame(
+                right,
+                ir.right_on,
+                partition_info,
+                ir.config_options,
+                right_count,
+            )
+        else:
+            left = _maybe_shuffle_frame(
+                left,
+                ir.left_on,
+                partition_info,
+                ir.config_options,
+                left_count,
+            )
+
+    new_node = ir.reconstruct([left, right])
+    partition_info[new_node] = PartitionInfo(count=output_count)
+    return new_node, partition_info
+
+
+@lower_ir_node.register(ConditionalJoin)
+def _(
+    ir: ConditionalJoin, rec: LowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    if ir.options[2]:  # pragma: no cover
+        return _lower_ir_fallback(
+            ir,
+            rec,
+            msg="Slice not supported in ConditionalJoin for multiple partitions.",
+        )
+
+    # Lower children
+    left, right = ir.children
+    left, pi_left = rec(left)
+    right, pi_right = rec(right)
+
+    # Fallback to single partition on the smaller table
+    left_count = pi_left[left].count
+    right_count = pi_right[right].count
+    output_count = max(left_count, right_count)
+    fallback_msg = "ConditionalJoin not supported for multiple partitions."
+    if left_count < right_count:
+        if left_count > 1:
+            left = Repartition(left.schema, left)
+            pi_left[left] = PartitionInfo(count=1)
+            _fallback_inform(fallback_msg, rec.state["config_options"])
+    elif right_count > 1:
+        right = Repartition(left.schema, right)
+        pi_right[right] = PartitionInfo(count=1)
+        _fallback_inform(fallback_msg, rec.state["config_options"])
+
+    # Reconstruct and return
+    new_node = ir.reconstruct([left, right])
+    partition_info = reduce(operator.or_, (pi_left, pi_right))
+    partition_info[new_node] = PartitionInfo(count=output_count)
+    return new_node, partition_info
+
+
+@lower_ir_node.register(Join)
+def _(
+    ir: Join, rec: LowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Lower children
+    children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
+    partition_info = reduce(operator.or_, _partition_info)
+
+    left, right = children
+    output_count = max(partition_info[left].count, partition_info[right].count)
+    if output_count == 1:
+        new_node = ir.reconstruct(children)
+        partition_info[new_node] = PartitionInfo(count=1)
+        return new_node, partition_info
+    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
+        return _make_bcast_join(
+            ir,
+            output_count,
+            partition_info,
+            left,
+            right,
+        )
+    else:
+        # Create a hash join
+        return _make_hash_join(
+            ir,
+            output_count,
+            partition_info,
+            left,
+            right,
+        )
+
+
+@generate_ir_tasks.register(Join)
+def _(
+    ir: Join, partition_info: MutableMapping[IR, PartitionInfo]
+) -> MutableMapping[Any, Any]:
+    left, right = ir.children
+    output_count = partition_info[ir].count
+
+    left_partitioned = (
+        partition_info[left].partitioned_on == ir.left_on
+        and partition_info[left].count == output_count
+    )
+    right_partitioned = (
+        partition_info[right].partitioned_on == ir.right_on
+        and partition_info[right].count == output_count
+    )
+
+    if output_count == 1 or (left_partitioned and right_partitioned):
+        # Partition-wise join
+        left_name = get_key_name(left)
+        right_name = get_key_name(right)
+        return {
+            key: (
+                ir.do_evaluate,
+                *ir._non_child_args,
+                (left_name, i),
+                (right_name, i),
+            )
+            for i, key in enumerate(partition_info[ir].keys(ir))
+        }
+    else:
+        # Broadcast join
+        left_parts = partition_info[left]
+        right_parts = partition_info[right]
+        if left_parts.count >= right_parts.count:
+            small_side = "Right"
+            small_name = get_key_name(right)
+            small_size = partition_info[right].count
+            large_name = get_key_name(left)
+            large_on = ir.left_on
+        else:
+            small_side = "Left"
+            small_name = get_key_name(left)
+            small_size = partition_info[left].count
+            large_name = get_key_name(right)
+            large_on = ir.right_on
+
+        graph: MutableMapping[Any, Any] = {}
+
+        out_name = get_key_name(ir)
+        out_size = partition_info[ir].count
+        split_name = f"split-{out_name}"
+        getit_name = f"getit-{out_name}"
+        inter_name = f"inter-{out_name}"
+
+        for part_out in range(out_size):
+            if ir.options[0] != "Inner":
+                graph[(split_name, part_out)] = (
+                    _partition_dataframe,
+                    (large_name, part_out),
+                    large_on,
+                    small_size,
+                )
+
+            _concat_list = []
+            for j in range(small_size):
+                left_key: tuple[str, int] | tuple[str, int, int]
+                if ir.options[0] != "Inner":
+                    left_key = (getit_name, part_out, j)
+                    graph[left_key] = (operator.getitem, (split_name, part_out), j)
+                else:
+                    left_key = (large_name, part_out)
+                join_children = [left_key, (small_name, j)]
+                if small_side == "Left":
+                    join_children.reverse()
+
+                inter_key = (inter_name, part_out, j)
+                graph[(inter_name, part_out, j)] = (
+                    ir.do_evaluate,
+                    ir.left_on,
+                    ir.right_on,
+                    ir.options,
+                    *join_children,
+                )
+                _concat_list.append(inter_key)
+            if len(_concat_list) == 1:
+                graph[(out_name, part_out)] = graph.pop(_concat_list[0])
+            else:
+                graph[(out_name, part_out)] = (_concat, *_concat_list)
+
+        return graph

cudf_polars/experimental/parallel.py

@@ -1,59 +1,61 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
-"""Multi-partition Dask execution."""
+"""Multi-partition evaluation."""
 
 from __future__ import annotations
 
 import itertools
 import operator
-from functools import reduce
+from functools import partial, reduce
 from typing import TYPE_CHECKING, Any
 
+import cudf_polars.experimental.distinct
+import cudf_polars.experimental.groupby
 import cudf_polars.experimental.io
-import cudf_polars.experimental.select  # noqa: F401
-from cudf_polars.dsl.ir import IR, Cache, Filter, HStack, Projection, Select, Union
+import cudf_polars.experimental.join
+import cudf_polars.experimental.select
+import cudf_polars.experimental.shuffle
+import cudf_polars.experimental.sort  # noqa: F401
+from cudf_polars.dsl.ir import (
+    IR,
+    Cache,
+    Filter,
+    HConcat,
+    HStack,
+    MapFunction,
+    Projection,
+    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
 
 if TYPE_CHECKING:
     from collections.abc import MutableMapping
+    from typing import Any
 
     from cudf_polars.containers import DataFrame
     from cudf_polars.experimental.dispatch import LowerIRTransformer
+    from cudf_polars.utils.config import ConfigOptions
 
 
 @lower_ir_node.register(IR)
-def _(ir: IR, rec: LowerIRTransformer) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+def _(
+    ir: IR, rec: LowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:  # pragma: no cover
     # Default logic - Requires single partition
-
-    if len(ir.children) == 0:
-        # Default leaf node has single partition
-        return ir, {
-            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
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     """
     Rewrite an IR graph and extract partitioning information.
 
@@ -61,6 +63,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
     -------
@@ -77,7 +81,7 @@ def lower_ir_graph(ir: IR) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     --------
     lower_ir_node
     """
-    mapper = CachingVisitor(lower_ir_node)
+    mapper = CachingVisitor(lower_ir_node, state={"config_options": config_options})
     return mapper(ir)
 
 
@@ -119,48 +123,118 @@ def task_graph(
     key_name = get_key_name(ir)
     partition_count = partition_info[ir].count
     if partition_count > 1:
-        graph[key_name] = (_concat, list(partition_info[ir].keys(ir)))
+        graph[key_name] = (_concat, *partition_info[ir].keys(ir))
         return graph, key_name
     else:
         return graph, (key_name, 0)
 
 
-def evaluate_dask(ir: IR) -> DataFrame:
-    """Evaluate an IR graph with Dask."""
-    from dask import get
+# The true type signature for get_scheduler() needs an overload. Not worth it.
+
+
+def get_scheduler(config_options: ConfigOptions) -> Any:
+    """Get appropriate task scheduler."""
+    assert config_options.executor.name == "streaming", (
+        "'in-memory' executor not supported in 'generate_ir_tasks'"
+    )
+
+    scheduler = config_options.executor.scheduler
+
+    if (
+        scheduler == "distributed"
+    ):  # pragma: no cover; block depends on executor type and Distributed cluster
+        from distributed import get_client
+
+        from cudf_polars.experimental.dask_registers import DaskRegisterManager
+
+        client = get_client()
+        DaskRegisterManager.register_once()
+        DaskRegisterManager.run_on_cluster(client)
+        return client.get
+    elif scheduler == "synchronous":
+        from cudf_polars.experimental.scheduler import synchronous_scheduler
+
+        return synchronous_scheduler
+    else:  # pragma: no cover
+        raise ValueError(f"{scheduler} not a supported scheduler option.")
+
+
+def post_process_task_graph(
+    graph: MutableMapping[Any, Any],
+    key: str | tuple[str, int],
+    config_options: ConfigOptions,
+) -> MutableMapping[Any, Any]:
+    """
+    Post-process the task graph.
+
+    Parameters
+    ----------
+    graph
+        Task graph to post-process.
+    key
+        Output key for the graph.
+    config_options
+        GPUEngine configuration options.
+
+    Returns
+    -------
+    graph
+        A Dask-compatible task graph.
+    """
+    assert config_options.executor.name == "streaming", (
+        "'in-memory' executor not supported in 'post_process_task_graph'"
+    )
+
+    if config_options.executor.rapidsmpf_spill:  # pragma: no cover
+        from cudf_polars.experimental.spilling import wrap_dataframe_in_spillable
+
+        return wrap_dataframe_in_spillable(
+            graph, ignore_key=key, config_options=config_options
+        )
+    return graph
+
+
+def evaluate_streaming(ir: IR, config_options: ConfigOptions) -> DataFrame:
+    """
+    Evaluate an IR graph with partitioning.
+
+    Parameters
+    ----------
+    ir
+        Logical plan to evaluate.
+    config_options
+        GPUEngine configuration options.
 
-    ir, partition_info = lower_ir_graph(ir)
+    Returns
+    -------
+    A cudf-polars DataFrame object.
+    """
+    ir, partition_info = lower_ir_graph(ir, config_options)
 
     graph, key = task_graph(ir, partition_info)
-    return get(graph, key)
+
+    graph = post_process_task_graph(graph, key, config_options)
+
+    return get_scheduler(config_options)(graph, key)
 
 
 @generate_ir_tasks.register(IR)
 def _(
     ir: IR, partition_info: MutableMapping[IR, PartitionInfo]
 ) -> MutableMapping[Any, Any]:
-    # Single-partition default behavior.
-    # This is used by `generate_ir_tasks` for all unregistered IR sub-types.
-    if partition_info[ir].count > 1:
-        raise NotImplementedError(
-            f"Failed to generate multiple output tasks for {ir}."
-        )  # pragma: no cover
-
-    child_names = []
-    for child in ir.children:
-        child_names.append(get_key_name(child))
-        if partition_info[child].count > 1:
-            raise NotImplementedError(
-                f"Failed to generate tasks for {ir} with child {child}."
-            )  # pragma: no cover
-
-    key_name = get_key_name(ir)
+    # Generate pointwise (embarrassingly-parallel) tasks by default
+    child_names = [get_key_name(c) for c in ir.children]
+    bcast_child = [partition_info[c].count == 1 for c in ir.children]
     return {
-        (key_name, 0): (
+        key: (
             ir.do_evaluate,
             *ir._non_child_args,
-            *((child_name, 0) for child_name in child_names),
+            *[
+                (child_name, 0 if bcast_child[j] else i)
+                for j, child_name in enumerate(child_names)
+            ],
         )
+        for i, key in enumerate(partition_info[ir].keys(ir))
     }
 
 
@@ -168,18 +242,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)
 
@@ -202,8 +274,22 @@ def _(
     }
 
 
+@lower_ir_node.register(MapFunction)
+def _(
+    ir: MapFunction, rec: LowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Allow pointwise operations
+    if ir.name in ("rename", "explode"):
+        return _lower_ir_pwise(ir, rec)
+
+    # Fallback for everything else
+    return _lower_ir_fallback(
+        ir, rec, msg=f"{ir.name} is not supported for multiple partitions."
+    )
+
+
 def _lower_ir_pwise(
-    ir: IR, rec: LowerIRTransformer
+    ir: IR, rec: LowerIRTransformer, *, preserve_partitioning: bool = False
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     # Lower a partition-wise (i.e. embarrassingly-parallel) IR node
 
@@ -213,41 +299,28 @@ 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.",
+        )
+
+    # Preserve child partition_info if possible
+    if preserve_partitioning and len(children) == 1:
+        partition = partition_info[children[0]]
+    else:
+        partition = PartitionInfo(count=max(counts))
 
     # Return reconstructed node and partition-info dict
-    partition = PartitionInfo(count=max(counts))
     new_node = ir.reconstruct(children)
     partition_info[new_node] = partition
     return new_node, partition_info
 
 
-lower_ir_node.register(Projection, _lower_ir_pwise)
+_lower_ir_pwise_preserve = partial(_lower_ir_pwise, preserve_partitioning=True)
+lower_ir_node.register(Projection, _lower_ir_pwise_preserve)
+lower_ir_node.register(Filter, _lower_ir_pwise_preserve)
 lower_ir_node.register(Cache, _lower_ir_pwise)
-lower_ir_node.register(Filter, _lower_ir_pwise)
 lower_ir_node.register(HStack, _lower_ir_pwise)
-
-
-def _generate_ir_tasks_pwise(
-    ir: IR, partition_info: MutableMapping[IR, PartitionInfo]
-) -> MutableMapping[Any, Any]:
-    # Generate partition-wise (i.e. embarrassingly-parallel) tasks
-    child_names = [get_key_name(c) for c in ir.children]
-    return {
-        key: (
-            ir.do_evaluate,
-            *ir._non_child_args,
-            *[(child_name, i) for child_name in child_names],
-        )
-        for i, key in enumerate(partition_info[ir].keys(ir))
-    }
-
-
-generate_ir_tasks.register(Projection, _generate_ir_tasks_pwise)
-generate_ir_tasks.register(Cache, _generate_ir_tasks_pwise)
-generate_ir_tasks.register(Filter, _generate_ir_tasks_pwise)
-generate_ir_tasks.register(HStack, _generate_ir_tasks_pwise)
-generate_ir_tasks.register(Select, _generate_ir_tasks_pwise)
+lower_ir_node.register(HConcat, _lower_ir_pwise)