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

cudf_polars/experimental/repartition.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Repartitioning Logic."""
+
+from __future__ import annotations
+
+import itertools
+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...N] repartitioning
+    (inclusive). The output partition count is tracked
+    separately 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 rapartitioning to fewer (for now).
+
+    (child,) = ir.children
+    count_in = partition_info[child].count
+    count_out = partition_info[ir].count
+
+    if count_out > count_in:  # pragma: no cover
+        raise NotImplementedError(
+            f"Repartition {count_in} -> {count_out} not supported."
+        )
+
+    key_name = get_key_name(ir)
+    n, remainder = divmod(count_in, count_out)
+    # Spread remainder evenly over the partitions.
+    offsets = [0, *itertools.accumulate(n + (i < remainder) for i in range(count_out))]
+    child_keys = tuple(partition_info[child].keys(child))
+    return {
+        (key_name, i): (_concat, *child_keys[offsets[i] : offsets[i + 1]])
+        for i in range(count_out)
+    }

cudf_polars/experimental/scheduler.py

@@ -0,0 +1,155 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Synchronous task scheduler."""
+
+from __future__ import annotations
+
+from collections import Counter
+from collections.abc import MutableMapping
+from itertools import chain
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from typing_extensions import Unpack
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from typing import TypeAlias
+
+
+Key: TypeAlias = str | tuple[str, Unpack[tuple[int, ...]]]
+Graph: TypeAlias = MutableMapping[Key, Any]
+T_ = TypeVar("T_")
+
+
+# NOTE: This is a slimmed-down version of the single-threaded
+# (synchronous) scheduler in `dask.core`.
+#
+# Key Differences:
+# * We do not allow a task to contain a list of key names.
+#   Keys must be distinct elements of the task.
+# * We do not support nested tasks.
+
+
+def istask(x: Any) -> bool:
+    """Check if x is a callable task."""
+    return isinstance(x, tuple) and bool(x) and callable(x[0])
+
+
+def is_hashable(x: Any) -> bool:
+    """Check if x is hashable."""
+    try:
+        hash(x)
+    except BaseException:
+        return False
+    else:
+        return True
+
+
+def _execute_task(arg: Any, cache: Mapping) -> Any:
+    """Execute a compute task."""
+    if istask(arg):
+        return arg[0](*(_execute_task(a, cache) for a in arg[1:]))
+    elif is_hashable(arg):
+        return cache.get(arg, arg)
+    else:
+        return arg
+
+
+def required_keys(key: Key, graph: Graph) -> list[Key]:
+    """
+    Return the dependencies to extract a key from the graph.
+
+    Parameters
+    ----------
+    key
+        Root key we want to extract.
+    graph
+        The full task graph.
+
+    Returns
+    -------
+    List of other keys needed to extract ``key``.
+    """
+    maybe_task = graph[key]
+    return [
+        k
+        for k in (
+            maybe_task[1:]
+            if istask(maybe_task)
+            else [maybe_task]  # maybe_task might be a key
+        )
+        if is_hashable(k) and k in graph
+    ]
+
+
+def toposort(graph: Graph, dependencies: Mapping[Key, list[Key]]) -> list[Key]:
+    """Return a list of task keys sorted in topological order."""
+    # Stack-based depth-first search traversal. This is based on Tarjan's
+    # algorithm for strongly-connected components
+    # (https://en.wikipedia.org/wiki/Tarjan%27s_strongly_connected_components_algorithm)
+    ordered: list[Key] = []
+    completed: set[Key] = set()
+
+    for key in graph:
+        if key in completed:
+            continue
+        nodes = [key]
+        while nodes:
+            # Keep current node on the stack until all descendants are visited
+            current = nodes[-1]
+            if current in completed:  # pragma: no cover
+                # Already fully traversed descendants of current
+                nodes.pop()
+                continue
+
+            # Add direct descendants of current to nodes stack
+            next_nodes = set(dependencies[current]) - completed
+            if next_nodes:
+                nodes.extend(next_nodes)
+            else:
+                # Current has no more descendants to explore
+                ordered.append(current)
+                completed.add(current)
+                nodes.pop()
+
+    return ordered
+
+
+def synchronous_scheduler(
+    graph: Graph,
+    key: Key,
+    *,
+    cache: MutableMapping | None = None,
+) -> Any:
+    """
+    Execute the task graph for a given key.
+
+    Parameters
+    ----------
+    graph
+        The task graph to execute.
+    key
+        The final output key to extract from the graph.
+    cache
+        Intermediate-data cache.
+
+    Returns
+    -------
+    Executed task-graph result for ``key``.
+    """
+    if key not in graph:  # pragma: no cover
+        raise KeyError(f"{key} is not a key in the graph")
+    if cache is None:
+        cache = {}
+
+    dependencies = {k: required_keys(k, graph) for k in graph}
+    refcount = Counter(chain.from_iterable(dependencies.values()))
+
+    for k in toposort(graph, dependencies):
+        cache[k] = _execute_task(graph[k], cache)
+        for dep in dependencies[k]:
+            refcount[dep] -= 1
+            if refcount[dep] == 0 and dep != key:
+                del cache[dep]
+
+    return cache[key]

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."""
 
@@ -6,16 +6,95 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from cudf_polars.dsl.ir import Select
+from cudf_polars.dsl.ir import HConcat, Select
 from cudf_polars.dsl.traversal import traversal
+from cudf_polars.experimental.base import PartitionInfo
 from cudf_polars.experimental.dispatch import lower_ir_node
+from cudf_polars.experimental.expressions import decompose_expr_graph
+from cudf_polars.experimental.utils import _lower_ir_fallback
 
 if TYPE_CHECKING:
     from collections.abc import MutableMapping
 
     from cudf_polars.dsl.ir import IR
-    from cudf_polars.experimental.base import PartitionInfo
     from cudf_polars.experimental.parallel import LowerIRTransformer
+    from cudf_polars.utils.config import ConfigOptions
+
+
+def decompose_select(
+    select_ir: Select,
+    input_ir: IR,
+    partition_info: MutableMapping[IR, PartitionInfo],
+    config_options: ConfigOptions,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    """
+    Decompose a multi-partition Select operation.
+
+    Parameters
+    ----------
+    select_ir
+        The original Select operation to decompose.
+        This object has not been reconstructed with
+        ``input_ir`` as its child yet.
+    input_ir
+        The lowered child of ``select_ir``. This object
+        will be decomposed into a "partial" selection
+        for each element of  ``select_ir.exprs``.
+    partition_info
+        A mapping from all unique IR nodes to the
+        associated partitioning information.
+    config_options
+        GPUEngine configuration options.
+
+    Returns
+    -------
+    new_ir, partition_info
+        The rewritten Select node, and a mapping from
+        unique nodes in the new graph to associated
+        partitioning information.
+
+    Notes
+    -----
+    This function uses ``decompose_expr_graph`` to further
+    decompose each element of  ``select_ir.exprs``.
+
+    See Also
+    --------
+    decompose_expr_graph
+    """
+    # Collect partial selections
+    selections = []
+    for ne in select_ir.exprs:
+        # Decompose this partial expression
+        new_ne, partial_input_ir, _partition_info = decompose_expr_graph(
+            ne, input_ir, partition_info, config_options
+        )
+        pi = _partition_info[partial_input_ir]
+        partial_input_ir = Select(
+            {ne.name: ne.value.dtype},
+            [new_ne],
+            True,  # noqa: FBT003
+            partial_input_ir,
+        )
+        _partition_info[partial_input_ir] = pi
+        partition_info.update(_partition_info)
+        selections.append(partial_input_ir)
+
+    # Concatenate partial selections
+    new_ir: HConcat | Select
+    if len(selections) > 1:
+        new_ir = HConcat(
+            select_ir.schema,
+            True,  # noqa: FBT003
+            *selections,
+        )
+        partition_info[new_ir] = PartitionInfo(
+            count=max(partition_info[c].count for c in selections)
+        )
+    else:
+        new_ir = selections[0]
+
+    return new_ir, partition_info
 
 
 @lower_ir_node.register(Select)
@@ -27,10 +106,16 @@ def _(
     if pi.count > 1 and not all(
         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."
-        )
+        try:
+            # Try decomposing the underlying expressions
+            return decompose_select(
+                ir, child, partition_info, rec.state["config_options"]
+            )
+        except NotImplementedError:
+            return _lower_ir_fallback(
+                ir, rec, msg="This selection is not supported for multiple partitions."
+            )
+
     new_node = ir.reconstruct([child])
     partition_info[new_node] = pi
     return new_node, partition_info

cudf_polars/experimental/shuffle.py

@@ -0,0 +1,294 @@
+# 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, TypedDict
+
+import pylibcudf as plc
+import rmm.mr
+from rmm.pylibrmm.stream import DEFAULT_STREAM
+
+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 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
+
+    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
+
+
+# Supported shuffle methods
+_SHUFFLE_METHODS = ("rapidsmpf", "tasks")
+
+
+class ShuffleOptions(TypedDict):
+    """RapidsMPF shuffling options."""
+
+    on: Sequence[str]
+    column_names: Sequence[str]
+
+
+# Experimental rapidsmpf shuffler integration
+class RMPFIntegration:  # pragma: no cover
+    """cuDF-Polars protocol for rapidsmpf shuffler."""
+
+    @staticmethod
+    def insert_partition(
+        df: DataFrame,
+        partition_id: int,  # Not currently used
+        partition_count: int,
+        shuffler: Any,
+        options: ShuffleOptions,
+        *other: Any,
+    ) -> None:
+        """Add cudf-polars DataFrame chunks to an RMP shuffler."""
+        from rapidsmpf.shuffler import partition_and_pack
+
+        on = options["on"]
+        assert not other, f"Unexpected arguments: {other}"
+        columns_to_hash = tuple(df.column_names.index(val) for val in on)
+        packed_inputs = partition_and_pack(
+            df.table,
+            columns_to_hash=columns_to_hash,
+            num_partitions=partition_count,
+            stream=DEFAULT_STREAM,
+            device_mr=rmm.mr.get_current_device_resource(),
+        )
+        shuffler.insert_chunks(packed_inputs)
+
+    @staticmethod
+    def extract_partition(
+        partition_id: int,
+        shuffler: Any,
+        options: ShuffleOptions,
+    ) -> DataFrame:
+        """Extract a finished partition from the RMP shuffler."""
+        from rapidsmpf.shuffler import unpack_and_concat
+
+        shuffler.wait_on(partition_id)
+        column_names = options["column_names"]
+        return DataFrame.from_table(
+            unpack_and_concat(
+                shuffler.extract(partition_id),
+                stream=DEFAULT_STREAM,
+                device_mr=rmm.mr.get_current_device_resource(),
+            ),
+            column_names,
+        )
+
+
+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,
+    ) -> 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.
+    """
+    if df.num_rows == 0:
+        # Fast path for empty DataFrame
+        return {i: df for i in range(count)}
+
+    # 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.Scalar.from_py(count, plc.DataType(plc.TypeId.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]:
+    # Extract "shuffle_method" configuration
+    assert ir.config_options.executor.name == "streaming", (
+        "'in-memory' executor not supported in 'generate_ir_tasks'"
+    )
+
+    shuffle_method = ir.config_options.executor.shuffle_method
+
+    # Try using rapidsmpf shuffler if we have "simple" shuffle
+    # keys, and the "shuffle_method" config is set to "rapidsmpf"
+    _keys: list[Col]
+    if shuffle_method in (None, "rapidsmpf") and len(
+        _keys := [ne.value for ne in ir.keys if isinstance(ne.value, Col)]
+    ) == len(ir.keys):  # pragma: no cover
+        shuffle_on = [k.name for k in _keys]
+        try:
+            from rapidsmpf.integrations.dask import rapidsmpf_shuffle_graph
+
+            return rapidsmpf_shuffle_graph(
+                get_key_name(ir.children[0]),
+                get_key_name(ir),
+                partition_info[ir.children[0]].count,
+                partition_info[ir].count,
+                RMPFIntegration,
+                {"on": shuffle_on, "column_names": list(ir.schema.keys())},
+            )
+        except (ImportError, ValueError) as err:
+            # ImportError: rapidsmpf is not installed
+            # ValueError: rapidsmpf couldn't find a distributed client
+            if shuffle_method == "rapidsmpf":
+                # Only raise an error if the user specifically
+                # set the shuffle method to "rapidsmpf"
+                raise ValueError(
+                    "Rapidsmp is not installed correctly or the current "
+                    "Dask cluster does not support rapidsmpf shuffling."
+                ) from err
+
+    # Simple task-based fall-back
+    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,
+    )

cudf_polars/experimental/sort.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Sorting Logic."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.ir import Sort
+from cudf_polars.experimental.base import PartitionInfo
+from cudf_polars.experimental.dispatch import lower_ir_node
+from cudf_polars.experimental.repartition import Repartition
+from cudf_polars.experimental.utils import _lower_ir_fallback
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+
+
+@lower_ir_node.register(Sort)
+def _(
+    ir: Sort, rec: LowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # Special handling for slicing
+    # (May be a top- or bottom-k operation)
+
+    if ir.zlice is not None and ir.zlice[0] < 1:
+        # TODO: Handle large slices (e.g. 1m+ rows)
+        from cudf_polars.experimental.parallel import _lower_ir_pwise
+
+        # Sort input partitions
+        new_node, partition_info = _lower_ir_pwise(ir, rec)
+        if partition_info[new_node].count > 1:
+            # Collapse down to single partition
+            inter = Repartition(new_node.schema, new_node)
+            partition_info[inter] = PartitionInfo(count=1)
+            # Sort reduced partition
+            new_node = ir.reconstruct([inter])
+            partition_info[new_node] = PartitionInfo(count=1)
+        return new_node, partition_info
+
+    # Fallback
+    return _lower_ir_fallback(ir, rec, msg="Sort does not support multiple partitions.")