feat: Support some python standard lib callables in apply/combine (#2187)

Author: TrevorBergeron

bigframes/core/compile/polars/compiler.py

@@ -328,6 +328,26 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
             assert isinstance(op, string_ops.StrContainsRegexOp)
             return input.str.contains(pattern=op.pat, literal=False)
 
+        @compile_op.register(string_ops.UpperOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.UpperOp)
+            return input.str.to_uppercase()
+
+        @compile_op.register(string_ops.LowerOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.LowerOp)
+            return input.str.to_lowercase()
+
+        @compile_op.register(string_ops.ArrayLenOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.ArrayLenOp)
+            return input.list.len()
+
+        @compile_op.register(string_ops.StrLenOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.StrLenOp)
+            return input.str.len_chars()
+
         @compile_op.register(string_ops.StartsWithOp)
         def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
             assert isinstance(op, string_ops.StartsWithOp)

bigframes/core/compile/polars/lowering.py

@@ -27,6 +27,7 @@
     generic_ops,
     json_ops,
     numeric_ops,
+    string_ops,
 )
 import bigframes.operations as ops
 
@@ -347,11 +348,28 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
             return ops.coalesce_op.as_expr(new_isin, expression.const(False))
 
 
+class LowerLenOp(op_lowering.OpLoweringRule):
+    @property
+    def op(self) -> type[ops.ScalarOp]:
+        return string_ops.LenOp
+
+    def lower(self, expr: expression.OpExpression) -> expression.Expression:
+        assert isinstance(expr.op, string_ops.LenOp)
+        arg = expr.children[0]
+
+        if dtypes.is_string_like(arg.output_type):
+            return string_ops.StrLenOp().as_expr(arg)
+        elif dtypes.is_array_like(arg.output_type):
+            return string_ops.ArrayLenOp().as_expr(arg)
+        else:
+            raise ValueError(f"Unexpected type: {arg.output_type}")
+
+
 def _coerce_comparables(
     expr1: expression.Expression,
     expr2: expression.Expression,
     *,
-    bools_only: bool = False
+    bools_only: bool = False,
 ):
     if bools_only:
         if (
@@ -446,6 +464,7 @@ def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
     LowerAsTypeRule(),
     LowerInvertOp(),
     LowerIsinOp(),
+    LowerLenOp(),
 )
 
 

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

@@ -29,15 +29,6 @@
     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,
@@ -80,6 +71,78 @@ def sin_op_impl(
     return input.sin()
 
 
+@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.TanOp)
+def tan_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.SinOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.tan()
+
+
+@polars_compiler.register_op(numeric_ops.SinhOp)
+def sinh_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.SinOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.sinh()
+
+
+@polars_compiler.register_op(numeric_ops.CoshOp)
+def cosh_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.CosOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.cosh()
+
+
+@polars_compiler.register_op(numeric_ops.TanhOp)
+def tanh_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.SinOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.tanh()
+
+
+@polars_compiler.register_op(numeric_ops.ArcsinOp)
+def asin_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.ArcsinOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.arcsin()
+
+
+@polars_compiler.register_op(numeric_ops.ArccosOp)
+def acos_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.ArccosOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.arccos()
+
+
+@polars_compiler.register_op(numeric_ops.ArctanOp)
+def atan_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.ArctanOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.arctan()
+
+
 @polars_compiler.register_op(numeric_ops.SqrtOp)
 def sqrt_op_impl(
     compiler: polars_compiler.PolarsExpressionCompiler,

bigframes/operations/python_op_maps.py

@@ -0,0 +1,85 @@
+# 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.
+
+import math
+import operator
+from typing import Optional
+
+import bigframes.operations
+from bigframes.operations import (
+    aggregations,
+    array_ops,
+    bool_ops,
+    comparison_ops,
+    numeric_ops,
+    string_ops,
+)
+
+PYTHON_TO_BIGFRAMES = {
+    ## operators
+    operator.add: numeric_ops.add_op,
+    operator.sub: numeric_ops.sub_op,
+    operator.mul: numeric_ops.mul_op,
+    operator.truediv: numeric_ops.div_op,
+    operator.floordiv: numeric_ops.floordiv_op,
+    operator.mod: numeric_ops.mod_op,
+    operator.pow: numeric_ops.pow_op,
+    operator.pos: numeric_ops.pos_op,
+    operator.neg: numeric_ops.neg_op,
+    operator.abs: numeric_ops.abs_op,
+    operator.eq: comparison_ops.eq_op,
+    operator.ne: comparison_ops.ne_op,
+    operator.gt: comparison_ops.gt_op,
+    operator.lt: comparison_ops.lt_op,
+    operator.ge: comparison_ops.ge_op,
+    operator.le: comparison_ops.le_op,
+    operator.and_: bool_ops.and_op,
+    operator.or_: bool_ops.or_op,
+    operator.xor: bool_ops.xor_op,
+    ## math
+    math.log: numeric_ops.ln_op,
+    math.log10: numeric_ops.log10_op,
+    math.log1p: numeric_ops.log1p_op,
+    math.expm1: numeric_ops.expm1_op,
+    math.sin: numeric_ops.sin_op,
+    math.cos: numeric_ops.cos_op,
+    math.tan: numeric_ops.tan_op,
+    math.sinh: numeric_ops.sinh_op,
+    math.cosh: numeric_ops.cosh_op,
+    math.tanh: numeric_ops.tanh_op,
+    math.asin: numeric_ops.arcsin_op,
+    math.acos: numeric_ops.arccos_op,
+    math.atan: numeric_ops.arctan_op,
+    math.floor: numeric_ops.floor_op,
+    math.ceil: numeric_ops.ceil_op,
+    ## str
+    str.upper: string_ops.upper_op,
+    str.lower: string_ops.lower_op,
+    ## builtins
+    len: string_ops.len_op,
+    abs: numeric_ops.abs_op,
+    pow: numeric_ops.pow_op,
+    ### builtins -- iterable
+    all: array_ops.ArrayReduceOp(aggregations.all_op),
+    any: array_ops.ArrayReduceOp(aggregations.any_op),
+    sum: array_ops.ArrayReduceOp(aggregations.sum_op),
+    min: array_ops.ArrayReduceOp(aggregations.min_op),
+    max: array_ops.ArrayReduceOp(aggregations.max_op),
+}
+
+
+def python_callable_to_op(obj) -> Optional[bigframes.operations.RowOp]:
+    if obj in PYTHON_TO_BIGFRAMES:
+        return PYTHON_TO_BIGFRAMES[obj]
+    return None

bigframes/operations/string_ops.py

@@ -30,6 +30,23 @@
 )
 len_op = LenOp()
 
+## Specialized len ops for compile-time lowering
+StrLenOp = base_ops.create_unary_op(
+    name="strlen",
+    type_signature=op_typing.FixedOutputType(
+        dtypes.is_string_like, dtypes.INT_DTYPE, description="string-like"
+    ),
+)
+str_len_op = StrLenOp()
+
+ArrayLenOp = base_ops.create_unary_op(
+    name="arraylen",
+    type_signature=op_typing.FixedOutputType(
+        dtypes.is_array_like, dtypes.INT_DTYPE, description="array-like"
+    ),
+)
+array_len_op = ArrayLenOp()
+
 ReverseOp = base_ops.create_unary_op(
     name="reverse", type_signature=op_typing.STRING_TRANSFORM
 )