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

cudf_polars/VERSION

@@ -1 +1 @@
-25.04.00
+25.08.00

cudf_polars/callback.py

@@ -7,12 +7,14 @@ from __future__ import annotations
 
 import contextlib
 import os
+import textwrap
 import time
 import warnings
 from functools import cache, partial
 from typing import TYPE_CHECKING, Literal, overload
 
 import nvtx
+from typing_extensions import assert_never
 
 from polars.exceptions import ComputeError, PerformanceWarning
 
@@ -20,9 +22,10 @@ import pylibcudf
 import rmm
 from rmm._cuda import gpu
 
+from cudf_polars.dsl.tracing import CUDF_POLARS_NVTX_DOMAIN
 from cudf_polars.dsl.translate import Translator
+from cudf_polars.utils.config import _env_get_int, get_total_device_memory
 from cudf_polars.utils.timer import Timer
-from cudf_polars.utils.versions import POLARS_VERSION_LT_125
 
 if TYPE_CHECKING:
     from collections.abc import Generator
@@ -32,6 +35,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"]
 
@@ -44,13 +48,6 @@ _SUPPORTED_PREFETCHES = {
 }
 
 
-def _env_get_int(name, default):
-    try:
-        return int(os.getenv(name, default))
-    except (ValueError, TypeError):  # pragma: no cover
-        return default  # pragma: no cover
-
-
 @cache
 def default_memory_resource(
     device: int,
@@ -101,8 +98,7 @@ def default_memory_resource(
         ):
             raise ComputeError(
                 "GPU engine requested, but incorrect cudf-polars package installed. "
-                "If your system has a CUDA 11 driver, please uninstall `cudf-polars-cu12` "
-                "and install `cudf-polars-cu11`"
+                "cudf-polars requires CUDA 12.0+ to installed."
             ) from None
         else:
             raise
@@ -139,7 +135,11 @@ def set_memory_resource(
         mr = default_memory_resource(
             device=device,
             cuda_managed_memory=bool(
-                _env_get_int("POLARS_GPU_ENABLE_CUDA_MANAGED_MEMORY", default=1) != 0
+                _env_get_int(
+                    "POLARS_GPU_ENABLE_CUDA_MANAGED_MEMORY",
+                    default=1 if get_total_device_memory() is not None else 0,
+                )
+                != 0
             ),
         )
     rmm.mr.set_current_device_resource(mr)
@@ -184,9 +184,8 @@ def _callback(
     n_rows: int | None,
     should_time: Literal[False],
     *,
-    device: int | None,
-    memory_resource: int | None,
-    executor: Literal["pylibcudf", "dask-experimental"] | None,
+    memory_resource: rmm.mr.DeviceMemoryResource | None,
+    config_options: ConfigOptions,
     timer: Timer | None,
 ) -> pl.DataFrame: ...
 
@@ -199,9 +198,8 @@ def _callback(
     n_rows: int | None,
     should_time: Literal[True],
     *,
-    device: int | None,
-    memory_resource: int | None,
-    executor: Literal["pylibcudf", "dask-experimental"] | None,
+    memory_resource: rmm.mr.DeviceMemoryResource | None,
+    config_options: ConfigOptions,
     timer: Timer | None,
 ) -> tuple[pl.DataFrame, list[tuple[int, int, str]]]: ...
 
@@ -213,34 +211,42 @@ def _callback(
     n_rows: int | None,
     should_time: bool,  # noqa: FBT001
     *,
-    device: int | None,
-    memory_resource: int | None,
-    executor: Literal["pylibcudf", "dask-experimental"] | None,
+    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"),
+        nvtx.annotate(message="ExecuteIR", domain=CUDF_POLARS_NVTX_DOMAIN),
         # 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":
+        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 executor == "dask-experimental":
-            from cudf_polars.experimental.parallel import evaluate_dask
+        elif config_options.executor.name == "streaming":
+            from cudf_polars.experimental.parallel import evaluate_streaming
 
-            return evaluate_dask(ir).to_polars()
-        else:
-            raise ValueError(f"Unknown executor '{executor}'")
+            if timer is not None:
+                msg = textwrap.dedent("""\
+                    LazyFrame.profile() is not supported with the streaming executor.
+                    To profile execution with the streaming executor, use:
+
+                    - NVIDIA NSight Systems with the 'streaming' scheduler.
+                    - Dask's built-in profiling tools with the 'distributed' scheduler.
+                    """)
+                raise NotImplementedError(msg)
+
+            return evaluate_streaming(ir, config_options).to_polars()
+        assert_never(f"Unknown executor '{config_options.executor}'")
 
 
 def execute_with_cudf(
@@ -259,7 +265,7 @@ def execute_with_cudf(
         profiling should occur).
 
     config
-        GPUEngine configuration object
+        GPUEngine object. Configuration is available as ``engine.config``.
 
     Raises
     ------
@@ -277,16 +283,22 @@ def execute_with_cudf(
     else:
         start = time.monotonic_ns()
         timer = Timer(start - duration_since_start)
-    device = config.device
+
     memory_resource = config.memory_resource
-    raise_on_fail = config.config.get("raise_on_fail", False)
-    executor = config.config.get("executor", None)
-    with nvtx.annotate(message="ConvertIR", domain="cudf_polars"):
+
+    with nvtx.annotate(message="ConvertIR", domain=CUDF_POLARS_NVTX_DOMAIN):
         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
@@ -301,33 +313,15 @@ def execute_with_cudf(
             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:
-            if POLARS_VERSION_LT_125:  # pragma: no cover
-                nt.set_udf(
-                    partial(
-                        _callback,
-                        ir,
-                        should_time=False,
-                        device=device,
-                        memory_resource=memory_resource,
-                        executor=executor,
-                        timer=None,
-                    )
-                )
-            else:
-                nt.set_udf(
-                    partial(
-                        _callback,
-                        ir,
-                        device=device,
-                        memory_resource=memory_resource,
-                        executor=executor,
-                        timer=timer,
-                    )
+            nt.set_udf(
+                partial(
+                    _callback,
+                    ir,
+                    memory_resource=memory_resource,
+                    config_options=translator.config_options,
+                    timer=timer,
                 )
-
-
-if POLARS_VERSION_LT_125:  # pragma: no cover
-    execute_with_cudf = partial(execute_with_cudf, duration_since_start=None)
+            )

cudf_polars/containers/__init__.py

@@ -1,11 +1,13 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
 """Containers of concrete data."""
 
 from __future__ import annotations
 
-__all__: list[str] = ["Column", "DataFrame"]
+__all__: list[str] = ["Column", "DataFrame", "DataType"]
 
+# dataframe.py & column.py imports DataType, so import in this order to avoid circular import
+from cudf_polars.containers.datatype import DataType  # noqa: I001
 from cudf_polars.containers.column import Column
 from cudf_polars.containers.dataframe import DataFrame

cudf_polars/containers/column.py

@@ -8,6 +8,8 @@ from __future__ import annotations
 import functools
 from typing import TYPE_CHECKING
 
+import polars as pl
+import polars.datatypes.convert
 from polars.exceptions import InvalidOperationError
 
 import pylibcudf as plc
@@ -19,19 +21,39 @@ from pylibcudf.strings.convert.convert_integers import (
 )
 from pylibcudf.traits import is_floating_point
 
+from cudf_polars.containers import DataType
 from cudf_polars.utils import conversion
 from cudf_polars.utils.dtypes import is_order_preserving_cast
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-    import polars as pl
-
-    from cudf_polars.typing import ColumnHeader, ColumnOptions, Slice
+    from cudf_polars.typing import (
+        ColumnHeader,
+        ColumnOptions,
+        DeserializedColumnOptions,
+        Slice,
+    )
 
 __all__: list[str] = ["Column"]
 
 
+def _dtype_short_repr_to_dtype(dtype_str: str) -> pl.DataType:
+    """Convert a Polars dtype short repr to a Polars dtype."""
+    # limitations of dtype_short_repr_to_dtype described in
+    # py-polars/polars/datatypes/convert.py#L299
+    if dtype_str.startswith("list["):
+        stripped = dtype_str.removeprefix("list[").removesuffix("]")
+        return pl.List(_dtype_short_repr_to_dtype(stripped))
+    pl_type = polars.datatypes.convert.dtype_short_repr_to_dtype(dtype_str)
+    if pl_type is None:
+        raise ValueError(f"{dtype_str} was not able to be parsed by Polars.")
+    if isinstance(pl_type, polars.datatypes.DataTypeClass):
+        return pl_type()
+    else:
+        return pl_type
+
+
 class Column:
     """An immutable column with sortedness metadata."""
 
@@ -43,10 +65,12 @@ class Column:
     # Optional name, only ever set by evaluation of NamedExpr nodes
     # The internal evaluation should not care about the name.
     name: str | None
+    dtype: DataType
 
     def __init__(
         self,
         column: plc.Column,
+        dtype: DataType,
         *,
         is_sorted: plc.types.Sorted = plc.types.Sorted.NO,
         order: plc.types.Order = plc.types.Order.ASCENDING,
@@ -56,6 +80,7 @@ class Column:
         self.obj = column
         self.is_scalar = self.size == 1
         self.name = name
+        self.dtype = dtype
         self.set_sorted(is_sorted=is_sorted, order=order, null_order=null_order)
 
     @classmethod
@@ -81,7 +106,23 @@ class Column:
         (plc_column,) = plc.contiguous_split.unpack_from_memoryviews(
             packed_metadata, packed_gpu_data
         ).columns()
-        return cls(plc_column, **header["column_kwargs"])
+        return cls(plc_column, **cls.deserialize_ctor_kwargs(header["column_kwargs"]))
+
+    @staticmethod
+    def deserialize_ctor_kwargs(
+        column_kwargs: ColumnOptions,
+    ) -> DeserializedColumnOptions:
+        """Deserialize the constructor kwargs for a Column."""
+        dtype = DataType(  # pragma: no cover
+            _dtype_short_repr_to_dtype(column_kwargs["dtype"])
+        )
+        return {
+            "is_sorted": column_kwargs["is_sorted"],
+            "order": column_kwargs["order"],
+            "null_order": column_kwargs["null_order"],
+            "name": column_kwargs["name"],
+            "dtype": dtype,
+        }
 
     def serialize(
         self,
@@ -105,17 +146,21 @@ class Column:
             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 = {
+        header: ColumnHeader = {
+            "column_kwargs": self.serialize_ctor_kwargs(),
+            "frame_count": 2,
+        }
+        return header, packed.release()
+
+    def serialize_ctor_kwargs(self) -> ColumnOptions:
+        """Serialize the constructor kwargs for self."""
+        return {
             "is_sorted": self.is_sorted,
             "order": self.order,
             "null_order": self.null_order,
             "name": self.name,
+            "dtype": pl.polars.dtype_str_repr(self.dtype.polars),
         }
-        header: ColumnHeader = {
-            "column_kwargs": column_kwargs,
-            "frame_count": 2,
-        }
-        return header, packed.release()
 
     @functools.cached_property
     def obj_scalar(self) -> plc.Scalar:
@@ -172,12 +217,51 @@ class Column:
         return type(self)(
             self.obj,
             name=self.name,
+            dtype=self.dtype,
             is_sorted=like.is_sorted,
             order=like.order,
             null_order=like.null_order,
         )
 
-    def astype(self, dtype: plc.DataType) -> Column:
+    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.size <= 1 or self.size == self.null_count:
+            return True
+        if self.is_sorted == plc.types.Sorted.YES:
+            return self.order == order and (
+                self.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: DataType) -> Column:
         """
         Cast the column to as the requested dtype.
 
@@ -200,14 +284,18 @@ class Column:
         This only produces a copy if the requested dtype doesn't match
         the current one.
         """
-        if self.obj.type() == dtype:
+        plc_dtype = dtype.plc
+        if self.obj.type() == plc_dtype:
             return self
 
-        if dtype.id() == plc.TypeId.STRING or self.obj.type().id() == plc.TypeId.STRING:
-            return Column(self._handle_string_cast(dtype))
+        if (
+            plc_dtype.id() == plc.TypeId.STRING
+            or self.obj.type().id() == plc.TypeId.STRING
+        ):
+            return Column(self._handle_string_cast(plc_dtype), dtype=dtype)
         else:
-            result = Column(plc.unary.cast(self.obj, dtype))
-            if is_order_preserving_cast(self.obj.type(), dtype):
+            result = Column(plc.unary.cast(self.obj, plc_dtype), dtype=dtype)
+            if is_order_preserving_cast(self.obj.type(), plc_dtype):
                 return result.sorted_like(self)
             return result
 
@@ -220,24 +308,20 @@ class Column:
         else:
             if is_floating_point(dtype):
                 floats = is_float(self.obj)
-                if not plc.interop.to_arrow(
-                    plc.reduce.reduce(
-                        floats,
-                        plc.aggregation.all(),
-                        plc.DataType(plc.TypeId.BOOL8),
-                    )
-                ).as_py():
+                if not plc.reduce.reduce(
+                    floats,
+                    plc.aggregation.all(),
+                    plc.DataType(plc.TypeId.BOOL8),
+                ).to_py():
                     raise InvalidOperationError("Conversion from `str` failed.")
                 return to_floats(self.obj, dtype)
             else:
                 integers = is_integer(self.obj)
-                if not plc.interop.to_arrow(
-                    plc.reduce.reduce(
-                        integers,
-                        plc.aggregation.all(),
-                        plc.DataType(plc.TypeId.BOOL8),
-                    )
-                ).as_py():
+                if not plc.reduce.reduce(
+                    integers,
+                    plc.aggregation.all(),
+                    plc.DataType(plc.TypeId.BOOL8),
+                ).to_py():
                     raise InvalidOperationError("Conversion from `str` failed.")
                 return to_integers(self.obj, dtype)
 
@@ -323,6 +407,7 @@ class Column:
             order=self.order,
             null_order=self.null_order,
             name=self.name,
+            dtype=self.dtype,
         )
 
     def mask_nans(self) -> Self:
@@ -330,7 +415,7 @@ class Column:
         if plc.traits.is_floating_point(self.obj.type()):
             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))
+            result = type(self)(self.obj.with_mask(mask, new_count), self.dtype)
             if old_count == new_count:
                 return result.sorted_like(self)
             return result
@@ -339,14 +424,12 @@ class Column:
     @functools.cached_property
     def nan_count(self) -> int:
         """Return the number of NaN values in the column."""
-        if plc.traits.is_floating_point(self.obj.type()):
-            return plc.interop.to_arrow(
-                plc.reduce.reduce(
-                    plc.unary.is_nan(self.obj),
-                    plc.aggregation.sum(),
-                    plc.types.SIZE_TYPE,
-                )
-            ).as_py()
+        if self.size > 0 and plc.traits.is_floating_point(self.obj.type()):
+            return plc.reduce.reduce(
+                plc.unary.is_nan(self.obj),
+                plc.aggregation.sum(),
+                plc.types.SIZE_TYPE,
+            ).to_py()
         return 0
 
     @property
@@ -380,4 +463,4 @@ class Column:
             conversion.from_polars_slice(zlice, num_rows=self.size),
         )
         (column,) = table.columns()
-        return type(self)(column, name=self.name).sorted_like(self)
+        return type(self)(column, name=self.name, dtype=self.dtype).sorted_like(self)