Some fixes for unary functions (#16719)

Correctly handle `pow` and `log` by translating to binary expressions
when we observe the node.

Upgrade our minimum supported polars version (so that we see all these
function names from the rust IR).

Also tighten check for which groupby-aggs are supported when the
expression contains a unary function.

dependencies.yaml

@@ -631,7 +631,7 @@ dependencies:
     common:
       - output_types: [conda, requirements, pyproject]
         packages:
-          - polars>=1.3
+          - polars>=1.6
   run_dask_cudf:
     common:
       - output_types: [conda, requirements, pyproject]

python/cudf_polars/cudf_polars/dsl/expr.py

@@ -1003,6 +1003,7 @@ class UnaryFunction(Expr):
     _non_child = ("dtype", "name", "options")
     children: tuple[Expr, ...]
 
+    # Note: log, and pow are handled via translation to binops
     _OP_MAPPING: ClassVar[dict[str, plc.unary.UnaryOperator]] = {
         "sin": plc.unary.UnaryOperator.SIN,
         "cos": plc.unary.UnaryOperator.COS,
@@ -1017,7 +1018,6 @@ class UnaryFunction(Expr):
         "arccosh": plc.unary.UnaryOperator.ARCCOSH,
         "arctanh": plc.unary.UnaryOperator.ARCTANH,
         "exp": plc.unary.UnaryOperator.EXP,
-        "log": plc.unary.UnaryOperator.LOG,
         "sqrt": plc.unary.UnaryOperator.SQRT,
         "cbrt": plc.unary.UnaryOperator.CBRT,
         "ceil": plc.unary.UnaryOperator.CEIL,
@@ -1034,7 +1034,6 @@ class UnaryFunction(Expr):
             "round",
             "set_sorted",
             "unique",
-            "pow",
         }
     )
     _supported_cum_aggs = frozenset(
@@ -1169,21 +1168,6 @@ def do_evaluate(
                 )
                 arg = evaluated.obj_scalar if evaluated.is_scalar else evaluated.obj
             return Column(plc.replace.replace_nulls(column.obj, arg))
-        elif self.name == "pow":
-            (base, exponent) = (
-                c.evaluate(df, context=context, mapping=mapping) for c in self.children
-            )
-            base_obj = (
-                base.obj_scalar
-                if (base.is_scalar and not exponent.is_scalar)
-                else base.obj
-            )
-            exponent_obj = exponent.obj_scalar if exponent.is_scalar else exponent.obj
-            return Column(
-                plc.binaryop.binary_operation(
-                    base_obj, exponent_obj, plc.binaryop.BinaryOperator.POW, self.dtype
-                )
-            )
         elif self.name in self._OP_MAPPING:
             column = self.children[0].evaluate(df, context=context, mapping=mapping)
             if column.obj.type().id() != self.dtype.id():
@@ -1241,6 +1225,8 @@ def do_evaluate(
 
     def collect_agg(self, *, depth: int) -> AggInfo:
         """Collect information about aggregations in groupbys."""
+        if self.name in {"unique", "drop_nulls"} | self._supported_cum_aggs:
+            raise NotImplementedError(f"{self.name} in groupby")
         if depth == 1:
             # inside aggregation, need to pre-evaluate, groupby
             # construction has checked that we don't have nested aggs,

python/cudf_polars/cudf_polars/dsl/translate.py

@@ -76,25 +76,11 @@ def _translate_ir(
 def _(
     node: pl_ir.PythonScan, visitor: NodeTraverser, schema: dict[str, plc.DataType]
 ) -> ir.IR:
-    if visitor.version()[0] == 1:
-        # https://github.com/pola-rs/polars/pull/17939
-        # Versioning can be dropped once polars 1.4 is lowest
-        # supported version.
-        scan_fn, with_columns, source_type, predicate, nrows = node.options
-        options = (scan_fn, with_columns, source_type, nrows)
-        predicate = (
-            translate_named_expr(visitor, n=predicate)
-            if predicate is not None
-            else None
-        )
-    else:  # pragma: no cover; CI tests 1.4
-        # version == 0
-        options = node.options
-        predicate = (
-            translate_named_expr(visitor, n=node.predicate)
-            if node.predicate is not None
-            else None
-        )
+    scan_fn, with_columns, source_type, predicate, nrows = node.options
+    options = (scan_fn, with_columns, source_type, nrows)
+    predicate = (
+        translate_named_expr(visitor, n=predicate) if predicate is not None else None
+    )
     return ir.PythonScan(schema, options, predicate)
 
 
@@ -115,13 +101,8 @@ def _(
         n_rows = -1  # All rows
         skip_rows = 0  # Don't skip
     else:
-        if visitor.version() >= (1, 0):
-            # Polars 1.4 n_rows property is (skip, nrows)
-            skip_rows, n_rows = n_rows
-        else:  # pragma: no cover; CI tests 1.4
-            # Polars 1.3 n_rows property is integer, skip rows was
-            # always zero because it was not pushed down to reader.
-            skip_rows = 0
+        # TODO: with versioning, rename on the rust side
+        skip_rows, n_rows = n_rows
 
     row_index = file_options.row_index
     return ir.Scan(
@@ -445,12 +426,19 @@ def _(node: pl_expr.Function, visitor: NodeTraverser, dtype: plc.DataType) -> ex
             *(translate_expr(visitor, n=n) for n in node.input),
         )
     elif isinstance(name, str):
-        return expr.UnaryFunction(
-            dtype,
-            name,
-            options,
-            *(translate_expr(visitor, n=n) for n in node.input),
-        )
+        children = (translate_expr(visitor, n=n) for n in node.input)
+        if name == "log":
+            (base,) = options
+            (child,) = children
+            return expr.BinOp(
+                dtype,
+                plc.binaryop.BinaryOperator.LOG_BASE,
+                child,
+                expr.Literal(dtype, pa.scalar(base, type=plc.interop.to_arrow(dtype))),
+            )
+        elif name == "pow":
+            return expr.BinOp(dtype, plc.binaryop.BinaryOperator.POW, *children)
+        return expr.UnaryFunction(dtype, name, options, *children)
     raise NotImplementedError(
         f"No handler for Expr function node with {name=}"
     )  # pragma: no cover; polars raises on the rust side for now

python/cudf_polars/cudf_polars/utils/versions.py

@@ -12,11 +12,11 @@
 
 POLARS_VERSION = parse(__version__)
 
-POLARS_VERSION_GE_13 = POLARS_VERSION >= parse("1.3")
-POLARS_VERSION_GT_13 = POLARS_VERSION > parse("1.3")
-POLARS_VERSION_LT_13 = POLARS_VERSION < parse("1.3")
+POLARS_VERSION_GE_16 = POLARS_VERSION >= parse("1.6")
+POLARS_VERSION_GT_16 = POLARS_VERSION > parse("1.6")
+POLARS_VERSION_LT_16 = POLARS_VERSION < parse("1.6")
 
-if POLARS_VERSION < parse("1.3"):
+if POLARS_VERSION_LT_16:
     raise ImportError(
-        "cudf_polars requires py-polars v1.3 or greater."
+        "cudf_polars requires py-polars v1.6 or greater."
     )  # pragma: no cover

python/cudf_polars/pyproject.toml

@@ -20,7 +20,7 @@ license = { text = "Apache 2.0" }
 requires-python = ">=3.9"
 dependencies = [
     "cudf==24.8.*,>=0.0.0a0",
-    "polars>=1.3",
+    "polars>=1.6",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 classifiers = [
     "Intended Audience :: Developers",

python/cudf_polars/tests/expressions/test_numeric_unaryops.py

@@ -24,8 +24,7 @@
         "arcsinh",
         "arccosh",
         "arctanh",
-        # "exp", Missing rust side impl
-        # "log", Missing rust side impl
+        "exp",
         "sqrt",
         "cbrt",
         "ceil",
@@ -78,3 +77,15 @@ def test_pow(ldf, base_literal, exponent_literal):
     q = ldf.select(base.pow(exponent))
 
     assert_gpu_result_equal(q, check_exact=False)
+
+
+@pytest.mark.parametrize("natural", [True, False])
+def test_log(ldf, natural):
+    if natural:
+        expr = pl.col("a").log()
+    else:
+        expr = pl.col("a").log(10)
+
+    q = ldf.select(expr)
+
+    assert_gpu_result_equal(q, check_exact=False)

python/cudf_polars/tests/test_groupby.py

@@ -178,3 +178,12 @@ def test_groupby_literal_in_agg(df, key, expr):
     # so just sort by the group key
     q = df.group_by(key).agg(expr).sort(key, maintain_order=True)
     assert_gpu_result_equal(q)
+
+
+@pytest.mark.parametrize(
+    "expr",
+    [pl.col("int").unique(), pl.col("int").drop_nulls(), pl.col("int").cum_max()],
+)
+def test_groupby_unary_non_pointwise_raises(df, expr):
+    q = df.group_by("key1").agg(expr)
+    assert_ir_translation_raises(q, NotImplementedError)