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

cudf_polars/VERSION

@@ -1 +1 @@
-25.02.02
+25.06.00

cudf_polars/callback.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
 
 """Callback for the polars collect function to execute on device."""
@@ -7,11 +7,13 @@ from __future__ import annotations
 
 import contextlib
 import os
+import time
 import warnings
 from functools import cache, partial
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Literal, overload
 
 import nvtx
+from typing_extensions import assert_never
 
 from polars.exceptions import ComputeError, PerformanceWarning
 
@@ -20,6 +22,7 @@ import rmm
 from rmm._cuda import gpu
 
 from cudf_polars.dsl.translate import Translator
+from cudf_polars.utils.timer import Timer
 
 if TYPE_CHECKING:
     from collections.abc import Generator
@@ -29,6 +32,7 @@ if TYPE_CHECKING:
 
     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"]
 
@@ -41,7 +45,7 @@ _SUPPORTED_PREFETCHES = {
 }
 
 
-def _env_get_int(name, default):
+def _env_get_int(name: str, default: int) -> int:
     try:
         return int(os.getenv(name, default))
     except (ValueError, TypeError):  # pragma: no cover
@@ -173,74 +177,72 @@ def set_device(device: int | None) -> Generator[int, None, None]:
         gpu.setDevice(previous)
 
 
+@overload
 def _callback(
     ir: IR,
     with_columns: list[str] | None,
     pyarrow_predicate: str | None,
     n_rows: int | None,
+    should_time: Literal[False],
     *,
-    device: int | None,
-    memory_resource: int | None,
-    executor: Literal["pylibcudf", "dask-experimental"] | None,
-) -> pl.DataFrame:
+    memory_resource: rmm.mr.DeviceMemoryResource | None,
+    config_options: ConfigOptions,
+    timer: Timer | None,
+) -> pl.DataFrame: ...
+
+
+@overload
+def _callback(
+    ir: IR,
+    with_columns: list[str] | None,
+    pyarrow_predicate: str | None,
+    n_rows: int | None,
+    should_time: Literal[True],
+    *,
+    memory_resource: rmm.mr.DeviceMemoryResource | None,
+    config_options: ConfigOptions,
+    timer: Timer | None,
+) -> tuple[pl.DataFrame, list[tuple[int, int, str]]]: ...
+
+
+def _callback(
+    ir: IR,
+    with_columns: list[str] | None,
+    pyarrow_predicate: str | None,
+    n_rows: int | None,
+    should_time: bool,  # noqa: FBT001
+    *,
+    memory_resource: rmm.mr.DeviceMemoryResource | None,
+    config_options: ConfigOptions,
+    timer: Timer | None,
+) -> pl.DataFrame | tuple[pl.DataFrame, list[tuple[int, int, str]]]:
     assert with_columns is None
     assert pyarrow_predicate is None
     assert n_rows is None
+    if timer is not None:
+        assert should_time
     with (
         nvtx.annotate(message="ExecuteIR", domain="cudf_polars"),
         # Device must be set before memory resource is obtained.
-        set_device(device),
+        set_device(config_options.device),
         set_memory_resource(memory_resource),
     ):
-        if executor is None or executor == "pylibcudf":
-            return ir.evaluate(cache={}).to_polars()
-        elif executor == "dask-experimental":
-            from cudf_polars.experimental.parallel import evaluate_dask
-
-            return evaluate_dask(ir).to_polars()
-        else:
-            raise ValueError(f"Unknown executor '{executor}'")
-
-
-def validate_config_options(config: dict) -> None:
-    """
-    Validate the configuration options for the GPU engine.
-
-    Parameters
-    ----------
-    config
-        Configuration options to validate.
-
-    Raises
-    ------
-    ValueError
-        If the configuration contains unsupported options.
-    """
-    if unsupported := (
-        config.keys()
-        - {"raise_on_fail", "parquet_options", "executor", "executor_options"}
-    ):
-        raise ValueError(
-            f"Engine configuration contains unsupported settings: {unsupported}"
-        )
-    assert {"chunked", "chunk_read_limit", "pass_read_limit"}.issuperset(
-        config.get("parquet_options", {})
-    )
-
-    # Validate executor_options
-    executor = config.get("executor", "pylibcudf")
-    if executor == "dask-experimental":
-        unsupported = config.get("executor_options", {}).keys() - {
-            "max_rows_per_partition",
-            "parquet_blocksize",
-        }
-    else:
-        unsupported = config.get("executor_options", {}).keys()
-    if unsupported:
-        raise ValueError(f"Unsupported executor_options for {executor}: {unsupported}")
-
-
-def execute_with_cudf(nt: NodeTraverser, *, config: GPUEngine) -> None:
+        if config_options.executor.name == "in-memory":
+            df = ir.evaluate(cache={}, timer=timer).to_polars()
+            if timer is None:
+                return df
+            else:
+                return df, timer.timings
+        elif config_options.executor.name == "streaming":
+            from cudf_polars.experimental.parallel import evaluate_streaming
+
+            return evaluate_streaming(ir, config_options).to_polars()
+        assert_never(f"Unknown executor '{config_options.executor}'")
+
+
+def execute_with_cudf(
+    nt: NodeTraverser, duration_since_start: int | None, *, config: GPUEngine
+) -> None:
     """
     A post optimization callback that attempts to execute the plan with cudf.
 
@@ -249,8 +251,12 @@ def execute_with_cudf(nt: NodeTraverser, *, config: GPUEngine) -> None:
     nt
         NodeTraverser
 
+    duration_since_start
+        Time since the user started executing the query (or None if no
+        profiling should occur).
+
     config
-        GPUEngine configuration object
+        GPUEngine object. Configuration is available as ``engine.config``.
 
     Raises
     ------
@@ -263,16 +269,27 @@ def execute_with_cudf(nt: NodeTraverser, *, config: GPUEngine) -> None:
     -----
     The NodeTraverser is mutated if the libcudf executor can handle the plan.
     """
-    device = config.device
+    if duration_since_start is None:
+        timer = None
+    else:
+        start = time.monotonic_ns()
+        timer = Timer(start - duration_since_start)
+
     memory_resource = config.memory_resource
-    raise_on_fail = config.config.get("raise_on_fail", False)
-    executor = config.config.get("executor", None)
-    validate_config_options(config.config)
 
     with nvtx.annotate(message="ConvertIR", domain="cudf_polars"):
         translator = Translator(nt, config)
         ir = translator.translate_ir()
         ir_translation_errors = translator.errors
+        if timer is not None:
+            timer.store(start, time.monotonic_ns(), "gpu-ir-translation")
+
+        if (
+            memory_resource is None
+            and translator.config_options.executor.name == "streaming"
+            and translator.config_options.executor.scheduler == "distributed"
+        ):  # pragma: no cover; Requires distributed cluster
+            memory_resource = rmm.mr.get_current_device_resource()
         if len(ir_translation_errors):
             # TODO: Display these errors in user-friendly way.
             # tracked in https://github.com/rapidsai/cudf/issues/17051
@@ -287,15 +304,15 @@ def execute_with_cudf(nt: NodeTraverser, *, config: GPUEngine) -> None:
             exception = NotImplementedError(error_message, unique_errors)
             if bool(int(os.environ.get("POLARS_VERBOSE", 0))):
                 warnings.warn(error_message, PerformanceWarning, stacklevel=2)
-            if raise_on_fail:
+            if translator.config_options.raise_on_fail:
                 raise exception
         else:
             nt.set_udf(
                 partial(
                     _callback,
                     ir,
-                    device=device,
                     memory_resource=memory_resource,
-                    executor=executor,
+                    config_options=translator.config_options,
+                    timer=timer,
                 )
             )

cudf_polars/containers/column.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
 
 """A column, with some properties."""
@@ -19,6 +19,7 @@ from pylibcudf.strings.convert.convert_integers import (
 )
 from pylibcudf.traits import is_floating_point
 
+from cudf_polars.utils import conversion
 from cudf_polars.utils.dtypes import is_order_preserving_cast
 
 if TYPE_CHECKING:
@@ -26,6 +27,8 @@ if TYPE_CHECKING:
 
     import polars as pl
 
+    from cudf_polars.typing import ColumnHeader, ColumnOptions, Slice
+
 __all__: list[str] = ["Column"]
 
 
@@ -51,10 +54,69 @@ class Column:
         name: str | None = None,
     ):
         self.obj = column
-        self.is_scalar = self.obj.size() == 1
+        self.is_scalar = self.size == 1
         self.name = name
         self.set_sorted(is_sorted=is_sorted, order=order, null_order=null_order)
 
+    @classmethod
+    def deserialize(
+        cls, header: ColumnHeader, frames: tuple[memoryview, plc.gpumemoryview]
+    ) -> Self:
+        """
+        Create a Column from a serialized representation returned by `.serialize()`.
+
+        Parameters
+        ----------
+        header
+            The (unpickled) metadata required to reconstruct the object.
+        frames
+            Two-tuple of frames (a memoryview and a gpumemoryview).
+
+        Returns
+        -------
+        Column
+            The deserialized Column.
+        """
+        packed_metadata, packed_gpu_data = frames
+        (plc_column,) = plc.contiguous_split.unpack_from_memoryviews(
+            packed_metadata, packed_gpu_data
+        ).columns()
+        return cls(plc_column, **header["column_kwargs"])
+
+    def serialize(
+        self,
+    ) -> tuple[ColumnHeader, tuple[memoryview, plc.gpumemoryview]]:
+        """
+        Serialize the Column into header and frames.
+
+        Follows the Dask serialization scheme with a picklable header (dict) and
+        a tuple of frames (in this case a contiguous host and device buffer).
+
+        To enable dask support, dask serializers must be registered
+
+            >>> from cudf_polars.experimental.dask_serialize import register
+            >>> register()
+
+        Returns
+        -------
+        header
+            A dict containing any picklable metadata required to reconstruct the object.
+        frames
+            Two-tuple of frames suitable for passing to `plc.contiguous_split.unpack_from_memoryviews`
+        """
+        packed = plc.contiguous_split.pack(plc.Table([self.obj]))
+        column_kwargs: ColumnOptions = {
+            "is_sorted": self.is_sorted,
+            "order": self.order,
+            "null_order": self.null_order,
+            "name": self.name,
+        }
+        header: ColumnHeader = {
+            "column_kwargs": column_kwargs,
+            "frame_count": 2,
+        }
+        return header, packed.release()
+
     @functools.cached_property
     def obj_scalar(self) -> plc.Scalar:
         """
@@ -70,9 +132,7 @@ class Column:
             If the column is not length-1.
         """
         if not self.is_scalar:
-            raise ValueError(
-                f"Cannot convert a column of length {self.obj.size()} to scalar"
-            )
+            raise ValueError(f"Cannot convert a column of length {self.size} to scalar")
         return plc.copying.get_element(self.obj, 0)
 
     def rename(self, name: str | None, /) -> Self:
@@ -117,6 +177,44 @@ class Column:
             null_order=like.null_order,
         )
 
+    def check_sorted(
+        self,
+        *,
+        order: plc.types.Order,
+        null_order: plc.types.NullOrder,
+    ) -> bool:
+        """
+        Check if the column is sorted.
+
+        Parameters
+        ----------
+        order
+            The requested sort order.
+        null_order
+            Where nulls sort to.
+
+        Returns
+        -------
+        True if the column is sorted, false otherwise.
+
+        Notes
+        -----
+        If the sortedness flag is not set, this launches a kernel to
+        check sortedness.
+        """
+        if self.obj.size() <= 1 or self.obj.size() == self.obj.null_count():
+            return True
+        if self.is_sorted == plc.types.Sorted.YES:
+            return self.order == order and (
+                self.obj.null_count() == 0 or self.null_order == null_order
+            )
+        if plc.sorting.is_sorted(plc.Table([self.obj]), [order], [null_order]):
+            self.sorted = plc.types.Sorted.YES
+            self.order = order
+            self.null_order = null_order
+            return True
+        return False
+
     def astype(self, dtype: plc.DataType) -> Column:
         """
         Cast the column to as the requested dtype.
@@ -242,7 +340,7 @@ class Column:
         -------
         Self with metadata set.
         """
-        if self.obj.size() <= 1:
+        if self.size <= 1:
             is_sorted = plc.types.Sorted.YES
         self.is_sorted = is_sorted
         self.order = order
@@ -268,7 +366,7 @@ class Column:
     def mask_nans(self) -> Self:
         """Return a shallow copy of self with nans masked out."""
         if plc.traits.is_floating_point(self.obj.type()):
-            old_count = self.obj.null_count()
+            old_count = self.null_count
             mask, new_count = plc.transform.nans_to_nulls(self.obj)
             result = type(self)(self.obj.with_mask(mask, new_count))
             if old_count == new_count:
@@ -288,3 +386,36 @@ class Column:
                 )
             ).as_py()
         return 0
+
+    @property
+    def size(self) -> int:
+        """Return the size of the column."""
+        return self.obj.size()
+
+    @property
+    def null_count(self) -> int:
+        """Return the number of Null values in the column."""
+        return self.obj.null_count()
+
+    def slice(self, zlice: Slice | None) -> Self:
+        """
+        Slice a column.
+
+        Parameters
+        ----------
+        zlice
+            optional, tuple of start and length, negative values of start
+            treated as for python indexing. If not provided, returns self.
+
+        Returns
+        -------
+        New column (if zlice is not None) otherwise self (if it is)
+        """
+        if zlice is None:
+            return self
+        (table,) = plc.copying.slice(
+            plc.Table([self.obj]),
+            conversion.from_polars_slice(zlice, num_rows=self.size),
+        )
+        (column,) = table.columns()
+        return type(self)(column, name=self.name).sorted_like(self)

cudf_polars/containers/dataframe.py

@@ -1,27 +1,26 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
 """A dataframe, with some properties."""
 
 from __future__ import annotations
 
-import pickle
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, cast
-
-import pyarrow as pa
+from typing import TYPE_CHECKING, cast
 
 import polars as pl
 
 import pylibcudf as plc
 
 from cudf_polars.containers import Column
-from cudf_polars.utils import dtypes
+from cudf_polars.utils import conversion
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Mapping, Sequence, Set
 
-    from typing_extensions import Self
+    from typing_extensions import Any, Self
+
+    from cudf_polars.typing import ColumnOptions, DataFrameHeader, Slice
 
 
 __all__: list[str] = ["DataFrame"]
@@ -107,17 +106,12 @@ class DataFrame:
         -------
         New dataframe representing the input.
         """
-        table = df.to_arrow()
-        schema = table.schema
-        for i, field in enumerate(schema):
-            schema = schema.set(
-                i, pa.field(field.name, dtypes.downcast_arrow_lists(field.type))
-            )
-        # No-op if the schema is unchanged.
-        d_table = plc.interop.from_arrow(table.cast(schema))
+        plc_table = plc.Table(df)
         return cls(
-            Column(column).copy_metadata(h_col)
-            for column, h_col in zip(d_table.columns(), df.iter_columns(), strict=True)
+            Column(d_col, name=name).copy_metadata(h_col)
+            for d_col, h_col, name in zip(
+                plc_table.columns(), df.iter_columns(), df.columns, strict=True
+            )
         )
 
     @classmethod
@@ -150,7 +144,7 @@ class DataFrame:
 
     @classmethod
     def deserialize(
-        cls, header: Mapping[str, Any], frames: tuple[memoryview, plc.gpumemoryview]
+        cls, header: DataFrameHeader, frames: tuple[memoryview, plc.gpumemoryview]
     ) -> Self:
         """
         Create a DataFrame from a serialized representation returned by `.serialize()`.
@@ -178,7 +172,7 @@ class DataFrame:
 
     def serialize(
         self,
-    ) -> tuple[Mapping[str, Any], tuple[memoryview, plc.gpumemoryview]]:
+    ) -> tuple[DataFrameHeader, tuple[memoryview, plc.gpumemoryview]]:
         """
         Serialize the table into header and frames.
 
@@ -187,20 +181,20 @@ class DataFrame:
 
         To enable dask support, dask serializers must be registered
 
-        >>> from cudf_polars.experimental.dask_serialize import register
-        >>> register()
+            >>> from cudf_polars.experimental.dask_serialize import register
+            >>> register()
 
         Returns
         -------
         header
             A dict containing any picklable metadata required to reconstruct the object.
         frames
-            Two-tuple of frames suitable for passing to `unpack_from_memoryviews`
+            Two-tuple of frames suitable for passing to `plc.contiguous_split.unpack_from_memoryviews`
         """
         packed = plc.contiguous_split.pack(self.table)
 
         # Keyword arguments for `Column.__init__`.
-        columns_kwargs = [
+        columns_kwargs: list[ColumnOptions] = [
             {
                 "is_sorted": col.is_sorted,
                 "order": col.order,
@@ -209,10 +203,8 @@ class DataFrame:
             }
             for col in self.columns
         ]
-        header = {
+        header: DataFrameHeader = {
             "columns_kwargs": columns_kwargs,
-            # Dask Distributed uses "type-serialized" to dispatch deserialization
-            "type-serialized": pickle.dumps(type(self)),
             "frame_count": 2,
         }
         return header, packed.release()
@@ -247,7 +239,9 @@ class DataFrame:
             for c, other in zip(self.columns, like.columns, strict=True)
         )
 
-    def with_columns(self, columns: Iterable[Column], *, replace_only=False) -> Self:
+    def with_columns(
+        self, columns: Iterable[Column], *, replace_only: bool = False
+    ) -> Self:
         """
         Return a new dataframe with extra columns.
 
@@ -276,7 +270,7 @@ class DataFrame:
         """Drop columns by name."""
         return type(self)(column for column in self.columns if column.name not in names)
 
-    def select(self, names: Sequence[str]) -> Self:
+    def select(self, names: Sequence[str] | Mapping[str, Any]) -> Self:
         """Select columns by name returning DataFrame."""
         try:
             return type(self)(self.column_map[name] for name in names)
@@ -296,7 +290,7 @@ class DataFrame:
         table = plc.stream_compaction.apply_boolean_mask(self.table, mask.obj)
         return type(self).from_table(table, self.column_names).sorted_like(self)
 
-    def slice(self, zlice: tuple[int, int] | None) -> Self:
+    def slice(self, zlice: Slice | None) -> Self:
         """
         Slice a dataframe.
 
@@ -312,14 +306,7 @@ class DataFrame:
         """
         if zlice is None:
             return self
-        start, length = zlice
-        if start < 0:
-            start += self.num_rows
-        # Polars implementation wraps negative start by num_rows, then
-        # adds length to start to get the end, then clamps both to
-        # [0, num_rows)
-        end = start + length
-        start = max(min(start, self.num_rows), 0)
-        end = max(min(end, self.num_rows), 0)
-        (table,) = plc.copying.slice(self.table, [start, end])
+        (table,) = plc.copying.slice(
+            self.table, conversion.from_polars_slice(zlice, num_rows=self.num_rows)
+        )
         return type(self).from_table(table, self.column_names).sorted_like(self)

cudf_polars/dsl/expr.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
 # TODO: remove need for this
 # ruff: noqa: D101
@@ -30,6 +30,7 @@ from cudf_polars.dsl.expressions.datetime import TemporalFunction
 from cudf_polars.dsl.expressions.literal import Literal, LiteralColumn
 from cudf_polars.dsl.expressions.rolling import GroupedRollingWindow, RollingWindow
 from cudf_polars.dsl.expressions.selection import Filter, Gather
+from cudf_polars.dsl.expressions.slicing import Slice
 from cudf_polars.dsl.expressions.sorting import Sort, SortBy
 from cudf_polars.dsl.expressions.string import StringFunction
 from cudf_polars.dsl.expressions.ternary import Ternary
@@ -53,6 +54,7 @@ __all__ = [
     "LiteralColumn",
     "NamedExpr",
     "RollingWindow",
+    "Slice",
     "Sort",
     "SortBy",
     "StringFunction",