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

cudf_polars/experimental/spilling.py

@@ -0,0 +1,151 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Spilling in multi-partition Dask execution using RAPIDSMPF."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from dask.sizeof import sizeof
+from distributed import get_worker
+from rapidsmpf.buffer.buffer import MemoryType
+from rapidsmpf.integrations.dask.core import get_worker_context
+from rapidsmpf.integrations.dask.spilling import SpillableWrapper
+
+from cudf_polars.containers import DataFrame
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, MutableMapping
+    from typing import Any
+
+    from cudf_polars.utils.config import ConfigOptions
+
+
+def wrap_arg(obj: Any) -> Any:
+    """
+    Make `obj` spillable if it is a DataFrame.
+
+    Parameters
+    ----------
+    obj
+        The object to be wrapped (if it is a DataFrame).
+
+    Returns
+    -------
+    A SpillableWrapper if obj is a DataFrame, otherwise the original object.
+    """
+    if isinstance(obj, DataFrame):
+        return SpillableWrapper(on_device=obj)
+    return obj
+
+
+def unwrap_arg(obj: Any) -> Any:
+    """
+    Unwraps a SpillableWrapper to retrieve the original object.
+
+    Parameters
+    ----------
+    obj
+        The object to be unwrapped.
+
+    Returns
+    -------
+    The unwrapped obj is a SpillableWrapper, otherwise the original object.
+    """
+    if isinstance(obj, SpillableWrapper):
+        return obj.unspill()
+    return obj
+
+
+def wrap_func_spillable(
+    func: Callable,
+    *,
+    make_func_output_spillable: bool,
+    target_partition_size: int,
+) -> Callable:
+    """
+    Wraps a function to handle spillable DataFrames.
+
+    Parameters
+    ----------
+    func
+        The function to be wrapped.
+    make_func_output_spillable
+        Whether to wrap the function's output in a SpillableWrapper.
+    target_partition_size
+        Target byte size for IO tasks.
+
+    Returns
+    -------
+    A wrapped function that processes spillable DataFrames.
+    """
+
+    def wrapper(*args: Any) -> Any:
+        # Make headroom before executing the task
+        headroom = 0
+        probable_io_task = True
+        for arg in args:
+            if isinstance(arg, SpillableWrapper):
+                if arg.mem_type() == MemoryType.HOST:
+                    headroom += sizeof(arg._on_host)
+                probable_io_task = False
+        if probable_io_task:
+            # Likely an IO task - Assume we need target_partition_size
+            headroom = target_partition_size
+        if headroom > 128_000_000:  # Don't waste time on smaller data
+            ctx = get_worker_context(get_worker())
+            with ctx.lock:
+                ctx.br.spill_manager.spill_to_make_headroom(headroom=headroom)
+
+        ret: Any = func(*(unwrap_arg(arg) for arg in args))
+        if make_func_output_spillable:
+            ret = wrap_arg(ret)
+        return ret
+
+    return wrapper
+
+
+def wrap_dataframe_in_spillable(
+    graph: MutableMapping[Any, Any],
+    ignore_key: str | tuple[str, int],
+    config_options: ConfigOptions,
+) -> MutableMapping[Any, Any]:
+    """
+    Wraps functions within a task graph to handle spillable DataFrames.
+
+    Only supports flat task graphs where each DataFrame can be found in the
+    outermost level. Currently, this is true for all cudf-polars task graphs.
+
+    Parameters
+    ----------
+    graph
+        Task graph.
+    ignore_key
+        The key to ignore when wrapping function, typically the key of the
+        output node.
+    config_options
+        GPUEngine configuration options.
+
+    Returns
+    -------
+    A new task graph with wrapped functions.
+    """
+    assert config_options.executor.name == "streaming", (
+        "'in-memory' executor not supported in 'wrap_dataframe_in_spillable'"
+    )
+    target_partition_size = config_options.executor.target_partition_size
+
+    ret = {}
+    for key, task in graph.items():
+        assert isinstance(task, tuple)
+        ret[key] = tuple(
+            wrap_func_spillable(
+                a,
+                make_func_output_spillable=key != ignore_key,
+                target_partition_size=target_partition_size,
+            )
+            if callable(a)
+            else a
+            for a in task
+        )
+    return ret

cudf_polars/experimental/utils.py

@@ -0,0 +1,100 @@
+# 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 itertools import chain
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.expr import Col
+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.expr import Expr
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+    from cudf_polars.utils.config import ConfigOptions
+
+
+def _concat(*dfs: DataFrame) -> DataFrame:
+    # Concatenate a sequence of DataFrames vertically
+    return Union.do_evaluate(None, *dfs)
+
+
+def _fallback_inform(msg: str, config_options: ConfigOptions) -> None:
+    """Inform the user of single-partition fallback."""
+    # Satisfy type checking
+    assert config_options.executor.name == "streaming", (
+        "'in-memory' executor not supported in '_fallback_inform'"
+    )
+
+    match fallback_mode := config_options.executor.fallback_mode:
+        case "warn":
+            warnings.warn(msg, stacklevel=2)
+        case "raise":
+            raise NotImplementedError(msg)
+        case "silent":
+            pass
+        case _:  # pragma: no cover; Should never get here.
+            raise ValueError(
+                f"{fallback_mode} is not a supported 'fallback_mode' "
+                "option. Please use 'warn', 'raise', or 'silent'."
+            )
+
+
+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"
+        _fallback_inform(msg, rec.state["config_options"])
+
+    # Reconstruct and return
+    new_node = ir.reconstruct(children)
+    partition_info[new_node] = PartitionInfo(count=1)
+    return new_node, partition_info
+
+
+def _leaf_column_names(expr: Expr) -> tuple[str, ...]:
+    """Find the leaf column names of an expression."""
+    if expr.children:
+        return tuple(
+            chain.from_iterable(_leaf_column_names(child) for child in expr.children)
+        )
+    elif isinstance(expr, Col):
+        return (expr.name,)
+    else:
+        return ()

cudf_polars/testing/asserts.py

@@ -5,24 +5,30 @@
 
 from __future__ import annotations
 
+from pathlib import Path
 from typing import TYPE_CHECKING
 
+import polars as pl
 from polars import GPUEngine
 from polars.testing.asserts import assert_frame_equal
 
 from cudf_polars.dsl.translate import Translator
 
 if TYPE_CHECKING:
-    import polars as pl
-
     from cudf_polars.typing import OptimizationArgs
 
-__all__: list[str] = ["assert_gpu_result_equal", "assert_ir_translation_raises"]
 
+__all__: list[str] = [
+    "assert_gpu_result_equal",
+    "assert_ir_translation_raises",
+    "assert_sink_ir_translation_raises",
+    "assert_sink_result_equal",
+]
 
 # Will be overriden by `conftest.py` with the value from the `--executor`
-# command-line argument
-Executor = None
+# and `--scheduler` command-line arguments
+DEFAULT_EXECUTOR = "in-memory"
+DEFAULT_SCHEDULER = "synchronous"
 
 
 def assert_gpu_result_equal(
@@ -89,7 +95,14 @@ def assert_gpu_result_equal(
         If GPU collection failed in some way.
     """
     if engine is None:
-        engine = GPUEngine(raise_on_fail=True, executor=executor or Executor)
+        executor = executor or DEFAULT_EXECUTOR
+        engine = GPUEngine(
+            raise_on_fail=True,
+            executor=executor,
+            executor_options=(
+                {"scheduler": DEFAULT_SCHEDULER} if executor == "streaming" else {}
+            ),
+        )
 
     final_polars_collect_kwargs, final_cudf_collect_kwargs = _process_kwargs(
         collect_kwargs, polars_collect_kwargs, cudf_collect_kwargs
@@ -236,3 +249,130 @@ def assert_collect_raises(
     else:
         if cudf_except != ():
             raise AssertionError(f"GPU execution DID NOT RAISE {cudf_except}")
+
+
+def _resolve_sink_format(path: Path) -> str:
+    """Returns valid sink format for assert utilities."""
+    suffix = path.suffix.lower()
+    supported_ext = {
+        ".csv": "csv",
+        ".pq": "parquet",
+        ".parquet": "parquet",
+        ".json": "ndjson",
+        ".ndjson": "ndjson",
+    }
+    if suffix not in supported_ext:
+        raise ValueError(f"Unsupported file format: {suffix}")
+    return supported_ext[suffix]
+
+
+def assert_sink_result_equal(
+    lazydf: pl.LazyFrame,
+    path: str | Path,
+    *,
+    engine: str | GPUEngine | None = None,
+    read_kwargs: dict | None = None,
+    write_kwargs: dict | None = None,
+    executor: str | None = None,
+) -> None:
+    """
+    Assert that writing a LazyFrame via sink produces the same output.
+
+    Parameters
+    ----------
+    lazydf
+        The LazyFrame to sink.
+    path
+        The file path to use. Suffix must be one of:
+        '.csv', '.parquet', '.pq', '.json', '.ndjson'.
+    engine
+        The GPU engine to use for the sink operation.
+    read_kwargs
+        Optional keyword arguments to pass to the corresponding `pl.read_*` function.
+    write_kwargs
+        Optional keyword arguments to pass to the corresponding `sink_*` function.
+    executor
+        The executor configuration to pass to `GPUEngine`. If not specified
+        uses the module level `Executor` attribute.
+
+    Raises
+    ------
+    AssertionError
+        If the outputs from CPU and GPU sink differ.
+    ValueError
+        If the file extension is not one of the supported formats.
+    """
+    if engine is None:
+        executor = executor or DEFAULT_EXECUTOR
+        engine = GPUEngine(
+            raise_on_fail=True,
+            executor=executor,
+            executor_options=(
+                {"scheduler": DEFAULT_SCHEDULER} if executor == "streaming" else {}
+            ),
+        )
+    path = Path(path)
+    read_kwargs = read_kwargs or {}
+    write_kwargs = write_kwargs or {}
+
+    fmt = _resolve_sink_format(path)
+
+    cpu_path = path.with_name(f"{path.stem}_cpu{path.suffix}")
+    gpu_path = path.with_name(f"{path.stem}_gpu{path.suffix}")
+
+    sink_fn = getattr(lazydf, f"sink_{fmt}")
+    read_fn = getattr(pl, f"read_{fmt}")
+
+    sink_fn(cpu_path, **write_kwargs)
+    sink_fn(gpu_path, engine=engine, **write_kwargs)
+
+    expected = read_fn(cpu_path, **read_kwargs)
+    result = read_fn(gpu_path, **read_kwargs)
+
+    assert_frame_equal(expected, result)
+
+
+def assert_sink_ir_translation_raises(
+    lazydf: pl.LazyFrame,
+    path: str | Path,
+    write_kwargs: dict,
+    *exceptions: type[Exception],
+) -> None:
+    """
+    Assert that translation of a sink query raises an exception.
+
+    Parameters
+    ----------
+    lazydf
+        The LazyFrame to sink.
+    path
+        The file path. Must have one of the supported suffixes.
+    write_kwargs
+        Keyword arguments to pass to the `sink_*` method.
+    *exceptions
+        One or more expected exception types that should be raised during translation.
+
+    Raises
+    ------
+    AssertionError
+        If translation does not raise any of the expected exceptions.
+        If an exception occurs before translation begins.
+    ValueError
+        If the file extension is not one of the supported formats.
+    """
+    path = Path(path)
+    fmt = _resolve_sink_format(path)
+
+    try:
+        lazy_sink = getattr(lazydf, f"sink_{fmt}")(
+            path,
+            engine="gpu",
+            lazy=True,
+            **write_kwargs,
+        )
+    except Exception as e:
+        raise AssertionError(
+            f"Sink function raised an exception before translation: {e}"
+        ) from e
+
+    assert_ir_translation_raises(lazy_sink, *exceptions)

cudf_polars/testing/io.py

@@ -0,0 +1,72 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""IO testing utilities."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Literal
+
+    import polars as pl
+
+__all__: list[str] = ["make_partitioned_source"]
+
+
+def make_partitioned_source(
+    df: pl.DataFrame,
+    path: str | Path,
+    fmt: Literal["csv", "ndjson", "parquet", "chunked_parquet"],
+    *,
+    n_files: int = 1,
+    row_group_size: int | None = None,
+    write_kwargs: dict | None = None,
+) -> None:
+    """
+    Write the Polars DataFrame to one or more files of the desired format.
+
+    Parameters
+    ----------
+    df : polars.DataFrame
+        The input DataFrame to write.
+    path : str | pathlib.Path
+        The base path to write the file(s) to.
+    fmt : Literal["csv", "ndjson", "parquet", "chunked_parquet"]
+        The format to write in.
+    n_files : int, default 1
+        If greater than 1, splits the data into multiple files.
+    row_group_size : optional, int
+        Only used for Parquet. Specifies the row group size per file.
+    write_kwargs : dict, optional
+        Additional keyword arguments to pass to the write_* functions.
+    """
+    path = Path(path)
+    write_kwargs = write_kwargs or {}
+
+    def write(part: pl.DataFrame, file_path: Path) -> None:
+        match fmt:
+            case "csv":
+                part.write_csv(file_path, **write_kwargs)
+            case "ndjson":
+                part.write_ndjson(file_path, **write_kwargs)
+            case "parquet" | "chunked_parquet":
+                part.write_parquet(
+                    file_path,
+                    row_group_size=row_group_size or (len(part) // 2),
+                    **write_kwargs,
+                )
+            case _:
+                raise ValueError(f"Unsupported format: {fmt}")
+
+    if n_files == 1:
+        if path.is_dir():
+            path = path / f"part.0.{fmt}"
+        write(df, path)
+    else:
+        stride = len(df) // n_files
+        for i, part in enumerate(df.iter_slices(stride)):
+            file_path = path / f"part.{i}.{fmt}"
+            write(part, file_path)