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

cudf_polars/experimental/io.py

@@ -22,14 +22,16 @@ if TYPE_CHECKING:
     from cudf_polars.dsl.expr import NamedExpr
     from cudf_polars.experimental.dispatch import LowerIRTransformer
     from cudf_polars.typing import Schema
+    from cudf_polars.utils.config import ConfigOptions
 
 
 @lower_ir_node.register(DataFrameScan)
 def _(
     ir: DataFrameScan, rec: LowerIRTransformer
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
-    rows_per_partition = ir.config_options.get("executor_options", {}).get(
-        "max_rows_per_partition", 1_000_000
+    rows_per_partition = ir.config_options.get(
+        "executor_options.max_rows_per_partition",
+        default=1_000_000,
     )
 
     nrows = max(ir.df.shape()[0], 1)
@@ -91,8 +93,10 @@ class ScanPartitionPlan:
         """Extract the partitioning plan of a Scan operation."""
         if ir.typ == "parquet":
             # TODO: Use system info to set default blocksize
-            parallel_options = ir.config_options.get("executor_options", {})
-            blocksize: int = parallel_options.get("parquet_blocksize", 1024**3)
+            blocksize: int = ir.config_options.get(
+                "executor_options.parquet_blocksize",
+                default=1024**3,
+            )
             stats = _sample_pq_statistics(ir)
             file_size = sum(float(stats[column]) for column in ir.schema)
             if file_size > 0:
@@ -168,7 +172,7 @@ class SplitScan(IR):
         schema: Schema,
         typ: str,
         reader_options: dict[str, Any],
-        config_options: dict[str, Any],
+        config_options: ConfigOptions,
         paths: list[str],
         with_columns: list[str] | None,
         skip_rows: int,
@@ -243,7 +247,7 @@ def _sample_pq_statistics(ir: Scan) -> dict[str, float]:
 
     # Use average total_uncompressed_size of three files
     # TODO: Use plc.io.parquet_metadata.read_parquet_metadata
-    n_sample = 3
+    n_sample = min(3, len(ir.paths))
     column_sizes = {}
     ds = pa_ds.dataset(random.sample(ir.paths, n_sample), format="parquet")
     for i, frag in enumerate(ds.get_fragments()):
@@ -270,11 +274,9 @@ def _(
         paths = list(ir.paths)
         if plan.flavor == ScanPartitionFlavor.SPLIT_FILES:
             # Disable chunked reader when splitting files
-            config_options = ir.config_options.copy()
-            config_options["parquet_options"] = config_options.get(
-                "parquet_options", {}
-            ).copy()
-            config_options["parquet_options"]["chunked"] = False
+            config_options = ir.config_options.set(
+                name="parquet_options.chunked", value=False
+            )
 
             slices: list[SplitScan] = []
             for path in paths:

cudf_polars/experimental/join.py

@@ -0,0 +1,318 @@
+# 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 Join
+from cudf_polars.experimental.base import PartitionInfo, _concat, get_key_name
+from cudf_polars.experimental.dispatch import generate_ir_tasks, lower_ir_node
+from cudf_polars.experimental.shuffle import Shuffle, _partition_dataframe
+
+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
+    return (
+        not large_shuffled
+        and small_count
+        <= ir.config_options.get(
+            # Maximum number of "small"-table partitions to bcast
+            "executor_options.broadcast_join_limit",
+            default=16,
+        )
+        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(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":
+        raise NotImplementedError(
+            "Cross join not support for multiple partitions."
+        )  # pragma: no cover
+
+    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}"
+        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):
+                join_children = [
+                    (
+                        (
+                            operator.getitem,
+                            (split_name, part_out),
+                            j,
+                        )
+                        if ir.options[0] != "Inner"
+                        else (large_name, part_out)
+                    ),
+                    (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,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
 """Multi-partition Dask execution."""
 
@@ -7,10 +7,13 @@ from __future__ import annotations
 import itertools
 import operator
 from functools import reduce
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
+import cudf_polars.experimental.groupby
 import cudf_polars.experimental.io
-import cudf_polars.experimental.select  # noqa: F401
+import cudf_polars.experimental.join
+import cudf_polars.experimental.select
+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
@@ -22,10 +25,38 @@ from cudf_polars.experimental.dispatch import (
 if TYPE_CHECKING:
     from collections.abc import MutableMapping
 
+    from distributed import Client
+
     from cudf_polars.containers import DataFrame
     from cudf_polars.experimental.dispatch import LowerIRTransformer
 
 
+class SerializerManager:
+    """Manager to ensure ensure serializer is only registered once."""
+
+    _serializer_registered: bool = False
+    _client_run_executed: ClassVar[set[str]] = set()
+
+    @classmethod
+    def register_serialize(cls) -> None:
+        """Register Dask/cudf-polars serializers in calling process."""
+        if not cls._serializer_registered:
+            from cudf_polars.experimental.dask_serialize import register
+
+            register()
+            cls._serializer_registered = True
+
+    @classmethod
+    def run_on_cluster(cls, client: Client) -> None:
+        """Run serializer registration on the workers and scheduler."""
+        if (
+            client.id not in cls._client_run_executed
+        ):  # pragma: no cover; Only executes with Distributed scheduler
+            client.run(cls.register_serialize)
+            client.run_on_scheduler(cls.register_serialize)
+            cls._client_run_executed.add(client.id)
+
+
 @lower_ir_node.register(IR)
 def _(ir: IR, rec: LowerIRTransformer) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     # Default logic - Requires single partition
@@ -119,18 +150,38 @@ 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 get_client():
+    """Get appropriate Dask client or scheduler."""
+    SerializerManager.register_serialize()
+
+    try:  # pragma: no cover; block depends on executor type and Distributed cluster
+        from distributed import get_client
+
+        client = get_client()
+        SerializerManager.run_on_cluster(client)
+    except (
+        ImportError,
+        ValueError,
+    ):  # pragma: no cover; block depends on Dask local scheduler
+        from dask import get
+
+        return get
+    else:  # pragma: no cover; block depends on executor type and Distributed cluster
+        return client.get
+
+
 def evaluate_dask(ir: IR) -> DataFrame:
     """Evaluate an IR graph with Dask."""
-    from dask import get
-
     ir, partition_info = lower_ir_graph(ir)
 
+    get = get_client()
+
     graph, key = task_graph(ir, partition_info)
     return get(graph, key)
 

cudf_polars/experimental/shuffle.py

@@ -0,0 +1,194 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Shuffle Logic."""
+
+from __future__ import annotations
+
+import operator
+from typing import TYPE_CHECKING, Any
+
+import pyarrow as pa
+
+import pylibcudf as plc
+
+from cudf_polars.containers import DataFrame
+from cudf_polars.dsl.ir import IR
+from cudf_polars.experimental.base import _concat, get_key_name
+from cudf_polars.experimental.dispatch import generate_ir_tasks, lower_ir_node
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.dsl.expr import NamedExpr
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+    from cudf_polars.experimental.parallel import PartitionInfo
+    from cudf_polars.typing import Schema
+    from cudf_polars.utils.config import ConfigOptions
+
+
+class Shuffle(IR):
+    """
+    Shuffle multi-partition data.
+
+    Notes
+    -----
+    Only hash-based partitioning is supported (for now).
+    """
+
+    __slots__ = ("config_options", "keys")
+    _non_child = ("schema", "keys", "config_options")
+    keys: tuple[NamedExpr, ...]
+    """Keys to shuffle on."""
+    config_options: ConfigOptions
+    """Configuration options."""
+
+    def __init__(
+        self,
+        schema: Schema,
+        keys: tuple[NamedExpr, ...],
+        config_options: ConfigOptions,
+        df: IR,
+    ):
+        self.schema = schema
+        self.keys = keys
+        self.config_options = config_options
+        self._non_child_args = (schema, keys, config_options)
+        self.children = (df,)
+
+    @classmethod
+    def do_evaluate(
+        cls,
+        schema: Schema,
+        keys: tuple[NamedExpr, ...],
+        config_options: ConfigOptions,
+        df: DataFrame,
+    ):  # pragma: no cover
+        """Evaluate and return a dataframe."""
+        # Single-partition Shuffle evaluation is a no-op
+        return df
+
+
+def _partition_dataframe(
+    df: DataFrame,
+    keys: tuple[NamedExpr, ...],
+    count: int,
+) -> dict[int, DataFrame]:
+    """
+    Partition an input DataFrame for shuffling.
+
+    Notes
+    -----
+    This utility only supports hash partitioning (for now).
+
+    Parameters
+    ----------
+    df
+        DataFrame to partition.
+    keys
+        Shuffle key(s).
+    count
+        Total number of output partitions.
+
+    Returns
+    -------
+    A dictionary mapping between int partition indices and
+    DataFrame fragments.
+    """
+    # Hash the specified keys to calculate the output
+    # partition for each row
+    partition_map = plc.binaryop.binary_operation(
+        plc.hashing.murmurhash3_x86_32(
+            DataFrame([expr.evaluate(df) for expr in keys]).table
+        ),
+        plc.interop.from_arrow(pa.scalar(count, type="uint32")),
+        plc.binaryop.BinaryOperator.PYMOD,
+        plc.types.DataType(plc.types.TypeId.UINT32),
+    )
+
+    # Apply partitioning
+    t, offsets = plc.partitioning.partition(
+        df.table,
+        partition_map,
+        count,
+    )
+
+    # Split and return the partitioned result
+    return {
+        i: DataFrame.from_table(
+            split,
+            df.column_names,
+        )
+        for i, split in enumerate(plc.copying.split(t, offsets[1:-1]))
+    }
+
+
+def _simple_shuffle_graph(
+    name_in: str,
+    name_out: str,
+    keys: tuple[NamedExpr, ...],
+    count_in: int,
+    count_out: int,
+) -> MutableMapping[Any, Any]:
+    """Make a simple all-to-all shuffle graph."""
+    split_name = f"split-{name_out}"
+    inter_name = f"inter-{name_out}"
+
+    graph: MutableMapping[Any, Any] = {}
+    for part_out in range(count_out):
+        _concat_list = []
+        for part_in in range(count_in):
+            graph[(split_name, part_in)] = (
+                _partition_dataframe,
+                (name_in, part_in),
+                keys,
+                count_out,
+            )
+            _concat_list.append((inter_name, part_out, part_in))
+            graph[_concat_list[-1]] = (
+                operator.getitem,
+                (split_name, part_in),
+                part_out,
+            )
+        graph[(name_out, part_out)] = (_concat, *_concat_list)
+    return graph
+
+
+@lower_ir_node.register(Shuffle)
+def _(
+    ir: Shuffle, rec: LowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Simple lower_ir_node handling for the default hash-based shuffle.
+    # More-complex logic (e.g. joining and sorting) should
+    # be handled separately.
+    from cudf_polars.experimental.parallel import PartitionInfo
+
+    (child,) = ir.children
+
+    new_child, pi = rec(child)
+    if pi[new_child].count == 1 or ir.keys == pi[new_child].partitioned_on:
+        # Already shuffled
+        return new_child, pi
+    new_node = ir.reconstruct([new_child])
+    pi[new_node] = PartitionInfo(
+        # Default shuffle preserves partition count
+        count=pi[new_child].count,
+        # Add partitioned_on info
+        partitioned_on=ir.keys,
+    )
+    return new_node, pi
+
+
+@generate_ir_tasks.register(Shuffle)
+def _(
+    ir: Shuffle, partition_info: MutableMapping[IR, PartitionInfo]
+) -> MutableMapping[Any, Any]:
+    # Use a simple all-to-all shuffle graph.
+
+    # TODO: Optionally use rapidsmp.
+    return _simple_shuffle_graph(
+        get_key_name(ir.children[0]),
+        get_key_name(ir),
+        ir.keys,
+        partition_info[ir.children[0]].count,
+        partition_info[ir].count,
+    )