Merge branch 'main' into routine-cleanup

Author: jialuoo

bigframes/blob/_functions.py

@@ -24,6 +24,8 @@
 # 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:
     import bigframes._importing

bigframes/core/compile/polars/__init__.py

@@ -493,9 +493,9 @@ def compile_agg_op(
             if isinstance(op, agg_ops.MedianOp):
                 return pl.median(*inputs)
             if isinstance(op, agg_ops.AllOp):
-                return pl.all(*inputs)
+                return pl.col(inputs).cast(pl.Boolean).all()
             if isinstance(op, agg_ops.AnyOp):
-                return pl.any(*inputs)  # type: ignore
+                return pl.col(inputs).cast(pl.Boolean).any()
             if isinstance(op, agg_ops.NuniqueOp):
                 return pl.col(*inputs).drop_nulls().n_unique()
             if isinstance(op, agg_ops.MinOp):

bigframes/core/compile/polars/compiler.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/polars/operations/numeric_ops.py

@@ -0,0 +1,48 @@
+# 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.generic_ops.
+
+Please keep implementations in sequential order by op name.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import bigframes_vendored.constants
+
+import bigframes.core.compile.polars.compiler as polars_compiler
+from bigframes.operations import struct_ops
+
+if TYPE_CHECKING:
+    import polars as pl
+
+
+@polars_compiler.register_op(struct_ops.StructFieldOp)
+def struct_field_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: struct_ops.StructFieldOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    if isinstance(op.name_or_index, str):
+        name = op.name_or_index
+    else:
+        raise NotImplementedError(
+            "Referencing a struct field by number not implemented in polars compiler. "
+            f"{bigframes_vendored.constants.FEEDBACK_LINK}"
+        )
+
+    return input.struct.field(name)

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

@@ -40,14 +40,9 @@ def apply_window_if_present(
         # Unbound grouping window.
         order_by = None
     elif window.is_range_bounded:
-        # Note that, when the window is range-bounded, we only need one ordering key.
-        # There are two reasons:
-        # 1. Manipulating null positions requires more than one ordering key, which
-        #  is forbidden by SQL window syntax for range rolling.
-        # 2. Pandas does not allow range rolling on timeseries with nulls.
-        order_by = get_window_order_by((window.ordering[0],), override_null_order=False)
+        order_by = get_window_order_by((window.ordering[0],))
     else:
-        order_by = get_window_order_by(window.ordering, override_null_order=True)
+        order_by = get_window_order_by(window.ordering)
 
     order = sge.Order(expressions=order_by) if order_by else None
 
@@ -102,7 +97,15 @@ def get_window_order_by(
     ordering: typing.Tuple[ordering_spec.OrderingExpression, ...],
     override_null_order: bool = False,
 ) -> typing.Optional[tuple[sge.Ordered, ...]]:
-    """Returns the SQL order by clause for a window specification."""
+    """Returns the SQL order by clause for a window specification.
+    Args:
+        ordering (Tuple[ordering_spec.OrderingExpression, ...]):
+            A tuple of ordering specification objects.
+        override_null_order (bool):
+            If True, overrides BigQuery's default null ordering behavior, which
+            is sometimes incompatible with ordered aggregations. The generated SQL
+            will include extra expressions to correctly enforce NULL FIRST/LAST.
+    """
     if not ordering:
         return None
 
@@ -115,8 +118,6 @@ def get_window_order_by(
         nulls_first = not ordering_spec_item.na_last
 
         if override_null_order:
-            # Bigquery SQL considers NULLS to be "smallest" values, but we need
-            # to override in these cases.
             is_null_expr = sge.Is(this=expr, expression=sge.Null())
             if nulls_first and desc:
                 order_by.append(

bigframes/core/compile/sqlglot/aggregations/windows.py

@@ -299,6 +299,8 @@ def compile_random_sample(
     def compile_aggregate(
         self, node: nodes.AggregateNode, child: ir.SQLGlotIR
     ) -> ir.SQLGlotIR:
+        # The BigQuery ordered aggregation cannot support for NULL FIRST/LAST,
+        # so we need to add extra expressions to enforce the null ordering.
         ordering_cols = windows.get_window_order_by(
             node.order_by, override_null_order=True
         )

bigframes/core/compile/sqlglot/compiler.py

@@ -23,6 +23,7 @@
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+register_ternary_op = scalar_compiler.scalar_op_compiler.register_ternary_op
 
 
 @register_unary_op(ops.AsTypeOp, pass_op=True)
@@ -66,6 +67,18 @@ def _(expr: TypedExpr, op: ops.AsTypeOp) -> sge.Expression:
     return _cast(sg_expr, sg_to_type, op.safe)
 
 
+@register_ternary_op(ops.clip_op)
+def _(
+    original: TypedExpr,
+    lower: TypedExpr,
+    upper: TypedExpr,
+) -> sge.Expression:
+    return sge.Greatest(
+        this=sge.Least(this=original.expr, expressions=[upper.expr]),
+        expressions=[lower.expr],
+    )
+
+
 @register_unary_op(ops.hash_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("FARM_FINGERPRINT", expr.expr)
@@ -94,6 +107,13 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.Not(this=sge.Is(this=expr.expr, expression=sge.Null()))
 
 
+@register_ternary_op(ops.where_op)
+def _(
+    original: TypedExpr, condition: TypedExpr, replacement: TypedExpr
+) -> sge.Expression:
+    return sge.If(this=condition.expr, true=original.expr, false=replacement.expr)
+
+
 # Helper functions
 def _cast_to_json(expr: TypedExpr, op: ops.AsTypeOp) -> sge.Expression:
     from_type = expr.dtype

bigframes/core/compile/sqlglot/expressions/generic_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,
             )
         ],