cudf-polars-cu13 25.10.0__py3-none-any.whl → 25.12.0__py3-none-any.whl

cudf_polars/dsl/expressions/struct.py

@@ -8,7 +8,9 @@ from __future__ import annotations
 
 from enum import IntEnum, auto
 from io import StringIO
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, cast
+
+import polars as pl
 
 import pylibcudf as plc
 
@@ -18,7 +20,7 @@ from cudf_polars.dsl.expressions.base import ExecutionContext, Expr
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-    from polars.polars import _expr_nodes as pl_expr
+    from polars import polars  # type: ignore[attr-defined]
 
     from cudf_polars.containers import DataFrame, DataType
 
@@ -42,7 +44,7 @@ class StructFunction(Expr):
         )  # https://github.com/pola-rs/polars/pull/23022#issuecomment-2933910958
 
         @classmethod
-        def from_polars(cls, obj: pl_expr.StructFunction) -> Self:
+        def from_polars(cls, obj: polars._expr_nodes.StructFunction) -> Self:
             """Convert from polars' `StructFunction`."""
             try:
                 function, name = str(obj).split(".", maxsplit=1)
@@ -87,11 +89,14 @@ class StructFunction(Expr):
         """Evaluate this expression given a dataframe for context."""
         columns = [child.evaluate(df, context=context) for child in self.children]
         (column,) = columns
+        # Type checker doesn't know polars only calls StructFunction with struct types
         if self.name == StructFunction.Name.FieldByName:
             field_index = next(
                 (
                     i
-                    for i, field in enumerate(self.children[0].dtype.polars.fields)
+                    for i, field in enumerate(
+                        cast(pl.Struct, self.children[0].dtype.polars_type).fields
+                    )
                     if field.name == self.options[0]
                 ),
                 None,
@@ -109,7 +114,12 @@ class StructFunction(Expr):
             table = plc.Table(column.obj.children())
             metadata = plc.io.TableWithMetadata(
                 table,
-                [(field.name, []) for field in self.children[0].dtype.polars.fields],
+                [
+                    (field.name, [])
+                    for field in cast(
+                        pl.Struct, self.children[0].dtype.polars_type
+                    ).fields
+                ],
             )
             options = (
                 plc.io.json.JsonWriterOptions.builder(target, table)
@@ -120,9 +130,11 @@ class StructFunction(Expr):
                 .utf8_escaped(val=False)
                 .build()
             )
-            plc.io.json.write_json(options)
+            plc.io.json.write_json(options, stream=df.stream)
             return Column(
-                plc.Column.from_iterable_of_py(buff.getvalue().split()),
+                plc.Column.from_iterable_of_py(
+                    buff.getvalue().split(), stream=df.stream
+                ),
                 dtype=self.dtype,
             )
         elif self.name in {

cudf_polars/dsl/expressions/ternary.py

@@ -41,9 +41,15 @@ class Ternary(Expr):
         when, then, otherwise = (
             child.evaluate(df, context=context) for child in self.children
         )
-        then_obj = then.obj_scalar if then.is_scalar else then.obj
-        otherwise_obj = otherwise.obj_scalar if otherwise.is_scalar else otherwise.obj
+        then_obj = then.obj_scalar(stream=df.stream) if then.is_scalar else then.obj
+        otherwise_obj = (
+            otherwise.obj_scalar(stream=df.stream)
+            if otherwise.is_scalar
+            else otherwise.obj
+        )
         return Column(
-            plc.copying.copy_if_else(then_obj, otherwise_obj, when.obj),
+            plc.copying.copy_if_else(
+                then_obj, otherwise_obj, when.obj, stream=df.stream
+            ),
             dtype=self.dtype,
         )

cudf_polars/dsl/expressions/unary.py

@@ -15,7 +15,6 @@ from cudf_polars.containers import Column
 from cudf_polars.dsl.expressions.base import ExecutionContext, Expr
 from cudf_polars.dsl.expressions.literal import Literal
 from cudf_polars.utils import dtypes
-from cudf_polars.utils.versions import POLARS_VERSION_LT_129
 
 if TYPE_CHECKING:
     from cudf_polars.containers import DataFrame, DataType
@@ -33,7 +32,7 @@ class Cast(Expr):
         self.dtype = dtype
         self.children = (value,)
         self.is_pointwise = True
-        if not dtypes.can_cast(value.dtype.plc, self.dtype.plc):
+        if not dtypes.can_cast(value.dtype.plc_type, self.dtype.plc_type):
             raise NotImplementedError(
                 f"Can't cast {value.dtype.id().name} to {self.dtype.id().name}"
             )
@@ -44,7 +43,7 @@ class Cast(Expr):
         """Evaluate this expression given a dataframe for context."""
         (child,) = self.children
         column = child.evaluate(df, context=context)
-        return column.astype(self.dtype)
+        return column.astype(self.dtype, stream=df.stream)
 
 
 class Len(Expr):
@@ -61,8 +60,9 @@ class Len(Expr):
         """Evaluate this expression given a dataframe for context."""
         return Column(
             plc.Column.from_scalar(
-                plc.Scalar.from_py(df.num_rows, self.dtype.plc),
+                plc.Scalar.from_py(df.num_rows, self.dtype.plc_type, stream=df.stream),
                 1,
+                stream=df.stream,
             ),
             dtype=self.dtype,
         )
@@ -150,7 +150,7 @@ class UnaryFunction(Expr):
         )
 
         if self.name not in UnaryFunction._supported_fns:
-            raise NotImplementedError(f"Unary function {name=}")
+            raise NotImplementedError(f"Unary function {name=}")  # pragma: no cover
         if self.name in UnaryFunction._supported_cum_aggs:
             (reverse,) = self.options
             if reverse:
@@ -174,26 +174,25 @@ class UnaryFunction(Expr):
         """Evaluate this expression given a dataframe for context."""
         if self.name == "mask_nans":
             (child,) = self.children
-            return child.evaluate(df, context=context).mask_nans()
+            return child.evaluate(df, context=context).mask_nans(stream=df.stream)
         if self.name == "null_count":
             (column,) = (child.evaluate(df, context=context) for child in self.children)
             return Column(
                 plc.Column.from_scalar(
-                    plc.Scalar.from_py(column.null_count, self.dtype.plc),
+                    plc.Scalar.from_py(
+                        column.null_count, self.dtype.plc_type, stream=df.stream
+                    ),
                     1,
+                    stream=df.stream,
                 ),
                 dtype=self.dtype,
             )
+        arg: plc.Column | plc.Scalar
         if self.name == "round":
-            round_mode = "half_away_from_zero"
-            if POLARS_VERSION_LT_129:
-                (decimal_places,) = self.options  # pragma: no cover
-            else:
-                # pragma: no cover
-                (
-                    decimal_places,
-                    round_mode,
-                ) = self.options
+            (
+                decimal_places,
+                round_mode,
+            ) = self.options
             (values,) = (child.evaluate(df, context=context) for child in self.children)
             return Column(
                 plc.round.round(
@@ -204,6 +203,7 @@ class UnaryFunction(Expr):
                         if round_mode == "half_to_even"
                         else plc.round.RoundingMethod.HALF_UP
                     ),
+                    stream=df.stream,
                 ),
                 dtype=self.dtype,
             ).sorted_like(values)  # pragma: no cover
@@ -215,30 +215,31 @@ class UnaryFunction(Expr):
             keep = plc.stream_compaction.DuplicateKeepOption.KEEP_ANY
             if values.is_sorted:
                 maintain_order = True
-                result = plc.stream_compaction.unique(
+                (compacted,) = plc.stream_compaction.unique(
                     plc.Table([values.obj]),
                     [0],
                     keep,
                     plc.types.NullEquality.EQUAL,
-                )
+                    stream=df.stream,
+                ).columns()
             else:
                 distinct = (
                     plc.stream_compaction.stable_distinct
                     if maintain_order
                     else plc.stream_compaction.distinct
                 )
-                result = distinct(
+                (compacted,) = distinct(
                     plc.Table([values.obj]),
                     [0],
                     keep,
                     plc.types.NullEquality.EQUAL,
                     plc.types.NanEquality.ALL_EQUAL,
-                )
-            (column,) = result.columns()
-            result = Column(column, dtype=self.dtype)
+                    stream=df.stream,
+                ).columns()
+            column = Column(compacted, dtype=self.dtype)
             if maintain_order:
-                result = result.sorted_like(values)
-            return result
+                column = column.sorted_like(values)
+            return column
         elif self.name == "set_sorted":
             (column,) = (child.evaluate(df, context=context) for child in self.children)
             (asc,) = self.options
@@ -250,9 +251,11 @@ class UnaryFunction(Expr):
             null_order = plc.types.NullOrder.BEFORE
             if column.null_count > 0 and (n := column.size) > 1:
                 # PERF: This invokes four stream synchronisations!
-                has_nulls_first = not plc.copying.get_element(column.obj, 0).is_valid()
+                has_nulls_first = not plc.copying.get_element(
+                    column.obj, 0, stream=df.stream
+                ).is_valid()
                 has_nulls_last = not plc.copying.get_element(
-                    column.obj, n - 1
+                    column.obj, n - 1, stream=df.stream
                 ).is_valid()
                 if (order == plc.types.Order.DESCENDING and has_nulls_first) or (
                     order == plc.types.Order.ASCENDING and has_nulls_last
@@ -280,30 +283,43 @@ class UnaryFunction(Expr):
                     counts_table,
                     [plc.types.Order.DESCENDING],
                     [plc.types.NullOrder.BEFORE],
+                    stream=df.stream,
                 )
                 counts_table = plc.copying.gather(
-                    counts_table, sort_indices, plc.copying.OutOfBoundsPolicy.DONT_CHECK
+                    counts_table,
+                    sort_indices,
+                    plc.copying.OutOfBoundsPolicy.DONT_CHECK,
+                    stream=df.stream,
                 )
                 keys_table = plc.copying.gather(
-                    keys_table, sort_indices, plc.copying.OutOfBoundsPolicy.DONT_CHECK
+                    keys_table,
+                    sort_indices,
+                    plc.copying.OutOfBoundsPolicy.DONT_CHECK,
+                    stream=df.stream,
                 )
             keys_col = keys_table.columns()[0]
             counts_col = counts_table.columns()[0]
             if normalize:
                 total_counts = plc.reduce.reduce(
-                    counts_col, plc.aggregation.sum(), plc.DataType(plc.TypeId.UINT64)
+                    counts_col,
+                    plc.aggregation.sum(),
+                    plc.DataType(plc.TypeId.UINT64),
+                    stream=df.stream,
                 )
                 counts_col = plc.binaryop.binary_operation(
                     counts_col,
                     total_counts,
                     plc.binaryop.BinaryOperator.DIV,
                     plc.DataType(plc.TypeId.FLOAT64),
+                    stream=df.stream,
                 )
             elif counts_col.type().id() == plc.TypeId.INT32:
-                counts_col = plc.unary.cast(counts_col, plc.DataType(plc.TypeId.UINT32))
+                counts_col = plc.unary.cast(
+                    counts_col, plc.DataType(plc.TypeId.UINT32), stream=df.stream
+                )
 
             plc_column = plc.Column(
-                self.dtype.plc,
+                self.dtype.plc_type,
                 counts_col.size(),
                 None,
                 None,
@@ -318,7 +334,7 @@ class UnaryFunction(Expr):
                 return column
             return Column(
                 plc.stream_compaction.drop_nulls(
-                    plc.Table([column.obj]), [0], 1
+                    plc.Table([column.obj]), [0], 1, stream=df.stream
                 ).columns()[0],
                 dtype=self.dtype,
             )
@@ -328,19 +344,31 @@ class UnaryFunction(Expr):
                 return column
             fill_value = self.children[1]
             if isinstance(fill_value, Literal):
-                arg = plc.Scalar.from_py(fill_value.value, fill_value.dtype.plc)
+                arg = plc.Scalar.from_py(
+                    fill_value.value, fill_value.dtype.plc_type, stream=df.stream
+                )
             else:
                 evaluated = fill_value.evaluate(df, context=context)
-                arg = evaluated.obj_scalar if evaluated.is_scalar else evaluated.obj
+                arg = (
+                    evaluated.obj_scalar(stream=df.stream)
+                    if evaluated.is_scalar
+                    else evaluated.obj
+                )
             if isinstance(arg, plc.Scalar) and dtypes.can_cast(
-                column.dtype.plc, arg.type()
+                column.dtype.plc_type, arg.type()
             ):  # pragma: no cover
                 arg = (
-                    Column(plc.Column.from_scalar(arg, 1), dtype=fill_value.dtype)
-                    .astype(column.dtype)
-                    .obj.to_scalar()
+                    Column(
+                        plc.Column.from_scalar(arg, 1, stream=df.stream),
+                        dtype=fill_value.dtype,
+                    )
+                    .astype(column.dtype, stream=df.stream)
+                    .obj.to_scalar(stream=df.stream)
                 )
-            return Column(plc.replace.replace_nulls(column.obj, arg), dtype=self.dtype)
+            return Column(
+                plc.replace.replace_nulls(column.obj, arg, stream=df.stream),
+                dtype=self.dtype,
+            )
         elif self.name == "fill_null_with_strategy":
             column = self.children[0].evaluate(df, context=context)
             strategy, limit = self.options
@@ -352,6 +380,8 @@ class UnaryFunction(Expr):
                 )
             ):
                 return column
+
+            replacement: plc.replace.ReplacePolicy | plc.Scalar
             if strategy == "forward":
                 replacement = plc.replace.ReplacePolicy.PRECEDING
             elif strategy == "backward":
@@ -360,37 +390,49 @@ class UnaryFunction(Expr):
                 replacement = plc.reduce.reduce(
                     column.obj,
                     plc.aggregation.min(),
-                    column.dtype.plc,
+                    column.dtype.plc_type,
+                    stream=df.stream,
                 )
             elif strategy == "max":
                 replacement = plc.reduce.reduce(
                     column.obj,
                     plc.aggregation.max(),
-                    column.dtype.plc,
+                    column.dtype.plc_type,
+                    stream=df.stream,
                 )
             elif strategy == "mean":
                 replacement = plc.reduce.reduce(
                     column.obj,
                     plc.aggregation.mean(),
                     plc.DataType(plc.TypeId.FLOAT64),
+                    stream=df.stream,
                 )
             elif strategy == "zero":
-                replacement = plc.scalar.Scalar.from_py(0, dtype=column.dtype.plc)
+                replacement = plc.scalar.Scalar.from_py(
+                    0, dtype=column.dtype.plc_type, stream=df.stream
+                )
             elif strategy == "one":
-                replacement = plc.scalar.Scalar.from_py(1, dtype=column.dtype.plc)
+                replacement = plc.scalar.Scalar.from_py(
+                    1, dtype=column.dtype.plc_type, stream=df.stream
+                )
             else:
                 assert_never(strategy)  # pragma: no cover
 
             if strategy == "mean":
                 return Column(
                     plc.replace.replace_nulls(
-                        plc.unary.cast(column.obj, plc.DataType(plc.TypeId.FLOAT64)),
+                        plc.unary.cast(
+                            column.obj,
+                            plc.DataType(plc.TypeId.FLOAT64),
+                            stream=df.stream,
+                        ),
                         replacement,
+                        stream=df.stream,
                     ),
                     dtype=self.dtype,
-                ).astype(self.dtype)
+                ).astype(self.dtype, stream=df.stream)
             return Column(
-                plc.replace.replace_nulls(column.obj, replacement),
+                plc.replace.replace_nulls(column.obj, replacement, stream=df.stream),
                 dtype=self.dtype,
             )
         elif self.name == "as_struct":
@@ -399,7 +441,7 @@ class UnaryFunction(Expr):
             ]
             return Column(
                 plc.Column(
-                    data_type=self.dtype.plc,
+                    data_type=self.dtype.plc_type,
                     size=children[0].size(),
                     data=None,
                     mask=None,
@@ -432,19 +474,24 @@ class UnaryFunction(Expr):
                 plc.types.NullPolicy.EXCLUDE,
                 plc.types.NullOrder.BEFORE if descending else plc.types.NullOrder.AFTER,
                 percentage=False,
+                stream=df.stream,
             )
 
             # Min/Max/Dense/Ordinal -> IDX_DTYPE
             # See https://github.com/pola-rs/polars/blob/main/crates/polars-ops/src/series/ops/rank.rs
             if method_str in {"min", "max", "dense", "ordinal"}:
-                dest = self.dtype.plc.id()
+                dest = self.dtype.plc_type.id()
                 src = ranked.type().id()
                 if dest == plc.TypeId.UINT32 and src != plc.TypeId.UINT32:
-                    ranked = plc.unary.cast(ranked, plc.DataType(plc.TypeId.UINT32))
+                    ranked = plc.unary.cast(
+                        ranked, plc.DataType(plc.TypeId.UINT32), stream=df.stream
+                    )
                 elif (
                     dest == plc.TypeId.UINT64 and src != plc.TypeId.UINT64
                 ):  # pragma: no cover
-                    ranked = plc.unary.cast(ranked, plc.DataType(plc.TypeId.UINT64))
+                    ranked = plc.unary.cast(
+                        ranked, plc.DataType(plc.TypeId.UINT64), stream=df.stream
+                    )
 
             return Column(ranked, dtype=self.dtype)
         elif self.name == "top_k":
@@ -459,23 +506,26 @@ class UnaryFunction(Expr):
                     plc.types.Order.ASCENDING
                     if reverse
                     else plc.types.Order.DESCENDING,
+                    stream=df.stream,
                 ),
                 dtype=self.dtype,
             )
         elif self.name in self._OP_MAPPING:
             column = self.children[0].evaluate(df, context=context)
-            if column.dtype.plc.id() != self.dtype.id():
-                arg = plc.unary.cast(column.obj, self.dtype.plc)
+            if column.dtype.plc_type.id() != self.dtype.id():
+                arg = plc.unary.cast(column.obj, self.dtype.plc_type, stream=df.stream)
             else:
                 arg = column.obj
             return Column(
-                plc.unary.unary_operation(arg, self._OP_MAPPING[self.name]),
+                plc.unary.unary_operation(
+                    arg, self._OP_MAPPING[self.name], stream=df.stream
+                ),
                 dtype=self.dtype,
             )
         elif self.name in UnaryFunction._supported_cum_aggs:
             column = self.children[0].evaluate(df, context=context)
             plc_col = column.obj
-            col_type = column.dtype.plc
+            col_type = column.dtype.plc_type
             # cum_sum casts
             # Int8, UInt8, Int16, UInt16 -> Int64 for overflow prevention
             # Bool -> UInt32
@@ -496,9 +546,16 @@ class UnaryFunction(Expr):
                 and plc.traits.is_integral(col_type)
                 and plc.types.size_of(col_type) <= 4
             ):
-                plc_col = plc.unary.cast(plc_col, plc.DataType(plc.TypeId.INT64))
-            elif self.name == "cum_sum" and column.dtype.plc.id() == plc.TypeId.BOOL8:
-                plc_col = plc.unary.cast(plc_col, plc.DataType(plc.TypeId.UINT32))
+                plc_col = plc.unary.cast(
+                    plc_col, plc.DataType(plc.TypeId.INT64), stream=df.stream
+                )
+            elif (
+                self.name == "cum_sum"
+                and column.dtype.plc_type.id() == plc.TypeId.BOOL8
+            ):
+                plc_col = plc.unary.cast(
+                    plc_col, plc.DataType(plc.TypeId.UINT32), stream=df.stream
+                )
             if self.name == "cum_sum":
                 agg = plc.aggregation.sum()
             elif self.name == "cum_prod":
@@ -509,7 +566,9 @@ class UnaryFunction(Expr):
                 agg = plc.aggregation.max()
 
             return Column(
-                plc.reduce.scan(plc_col, agg, plc.reduce.ScanType.INCLUSIVE),
+                plc.reduce.scan(
+                    plc_col, agg, plc.reduce.ScanType.INCLUSIVE, stream=df.stream
+                ),
                 dtype=self.dtype,
             )
         raise NotImplementedError(