feat: implement cos, sin, and log operations for polars compiler (#2170)

* feat: implement cos, sin, and log operations for polars compiler

* fix domain for log

* update snapshot

* revert sqrt change

* revert sqrt change

Author: tswast

bigframes/core/compile/polars/__init__.py

@@ -24,6 +24,7 @@
 # polars shouldn't be needed at import time, as register is a no-op if polars
 # isn't installed.
 import bigframes.core.compile.polars.operations.generic_ops  # noqa: F401
+import bigframes.core.compile.polars.operations.numeric_ops  # noqa: F401
 import bigframes.core.compile.polars.operations.struct_ops  # noqa: F401
 
 try:

bigframes/core/compile/polars/operations/numeric_ops.py

@@ -0,0 +1,91 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+BigFrames -> Polars compilation for the operations in bigframes.operations.numeric_ops.
+
+Please keep implementations in sequential order by op name.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import bigframes.core.compile.polars.compiler as polars_compiler
+from bigframes.operations import numeric_ops
+
+if TYPE_CHECKING:
+    import polars as pl
+
+
+@polars_compiler.register_op(numeric_ops.CosOp)
+def cos_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.CosOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.cos()
+
+
+@polars_compiler.register_op(numeric_ops.LnOp)
+def ln_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.LnOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    import polars as pl
+
+    return pl.when(input <= 0).then(float("nan")).otherwise(input.log())
+
+
+@polars_compiler.register_op(numeric_ops.Log10Op)
+def log10_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.Log10Op,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    import polars as pl
+
+    return pl.when(input <= 0).then(float("nan")).otherwise(input.log(base=10))
+
+
+@polars_compiler.register_op(numeric_ops.Log1pOp)
+def log1p_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.Log1pOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    import polars as pl
+
+    return pl.when(input <= -1).then(float("nan")).otherwise((input + 1).log())
+
+
+@polars_compiler.register_op(numeric_ops.SinOp)
+def sin_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.SinOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.sin()
+
+
+@polars_compiler.register_op(numeric_ops.SqrtOp)
+def sqrt_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.SqrtOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    import polars as pl
+
+    return pl.when(input < 0).then(float("nan")).otherwise(input.sqrt())

bigframes/core/compile/sqlglot/expressions/numeric_ops.py

@@ -158,7 +158,7 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.Case(
         ifs=[
             sge.If(
-                this=expr.expr < sge.convert(0),
+                this=expr.expr <= sge.convert(0),
                 true=constants._NAN,
             )
         ],
@@ -171,7 +171,7 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.Case(
         ifs=[
             sge.If(
-                this=expr.expr < sge.convert(0),
+                this=expr.expr <= sge.convert(0),
                 true=constants._NAN,
             )
         ],
@@ -184,7 +184,7 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.Case(
         ifs=[
             sge.If(
-                this=expr.expr < sge.convert(-1),
+                this=expr.expr <= sge.convert(-1),
                 true=constants._NAN,
             )
         ],

tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_ln/out.sql

@@ -5,7 +5,7 @@ WITH `bfcte_0` AS (
 ), `bfcte_1` AS (
   SELECT
     *,
-    CASE WHEN `bfcol_0` < 0 THEN CAST('NaN' AS FLOAT64) ELSE LN(`bfcol_0`) END AS `bfcol_1`
+    CASE WHEN `bfcol_0` <= 0 THEN CAST('NaN' AS FLOAT64) ELSE LN(`bfcol_0`) END AS `bfcol_1`
   FROM `bfcte_0`
 )
 SELECT

tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_log10/out.sql

@@ -5,7 +5,7 @@ WITH `bfcte_0` AS (
 ), `bfcte_1` AS (
   SELECT
     *,
-    CASE WHEN `bfcol_0` < 0 THEN CAST('NaN' AS FLOAT64) ELSE LOG(10, `bfcol_0`) END AS `bfcol_1`
+    CASE WHEN `bfcol_0` <= 0 THEN CAST('NaN' AS FLOAT64) ELSE LOG(10, `bfcol_0`) END AS `bfcol_1`
   FROM `bfcte_0`
 )
 SELECT

tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_log1p/out.sql

@@ -5,7 +5,7 @@ WITH `bfcte_0` AS (
 ), `bfcte_1` AS (
   SELECT
     *,
-    CASE WHEN `bfcol_0` < -1 THEN CAST('NaN' AS FLOAT64) ELSE LN(1 + `bfcol_0`) END AS `bfcol_1`
+    CASE WHEN `bfcol_0` <= -1 THEN CAST('NaN' AS FLOAT64) ELSE LN(1 + `bfcol_0`) END AS `bfcol_1`
   FROM `bfcte_0`
 )
 SELECT

tests/unit/test_series_polars.py

@@ -4622,20 +4622,15 @@ def test_apply_lambda(scalars_dfs, col, lambda_):
     )
 
 
-@pytest.mark.skip(
-    reason="NotImplementedError: Polars compiler hasn't implemented log()"
-)
 @pytest.mark.parametrize(
     ("ufunc",),
     [
-        pytest.param(numpy.log),
-        pytest.param(numpy.sqrt),
-        pytest.param(numpy.sin),
-    ],
-    ids=[
-        "log",
-        "sqrt",
-        "sin",
+        pytest.param(numpy.cos, id="cos"),
+        pytest.param(numpy.log, id="log"),
+        pytest.param(numpy.log10, id="log10"),
+        pytest.param(numpy.log1p, id="log1p"),
+        pytest.param(numpy.sqrt, id="sqrt"),
+        pytest.param(numpy.sin, id="sin"),
     ],
 )
 def test_apply_numpy_ufunc(scalars_dfs, ufunc):