feat: Support some python standard lib callables in apply/combine

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
 )