more tests

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/operations/python_op_maps.py

@@ -32,12 +32,13 @@
     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_null_match_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,
@@ -49,6 +50,8 @@
     ## 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,
@@ -58,6 +61,8 @@
     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,

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
 )

tests/unit/test_series_polars.py

@@ -15,6 +15,7 @@
 import datetime as dt
 import json
 import math
+import operator
 import pathlib
 import re
 import tempfile
@@ -4697,6 +4698,30 @@ def wrapped(x):
     assert_series_equal(bf_result, pd_result, check_dtype=False)
 
 
+@pytest.mark.parametrize(
+    ("ufunc",),
+    [
+        pytest.param(str.upper),
+        pytest.param(str.lower),
+        pytest.param(len),
+    ],
+)
+def test_series_apply_python_string_fns(scalars_dfs, ufunc):
+    scalars_df, scalars_pandas_df = scalars_dfs
+
+    bf_col = scalars_df["string_col"]
+    bf_result = bf_col.apply(ufunc).to_pandas()
+
+    pd_col = scalars_pandas_df["string_col"]
+
+    def wrapped(x):
+        return ufunc(x) if isinstance(x, str) else None
+
+    pd_result = pd_col.apply(wrapped)
+
+    assert_series_equal(bf_result, pd_result, check_dtype=False)
+
+
 @pytest.mark.parametrize(
     ("ufunc",),
     [
@@ -4720,6 +4745,29 @@ def test_combine_series_ufunc(scalars_dfs, ufunc):
     assert_series_equal(bf_result, pd_result, check_dtype=False)
 
 
+@pytest.mark.parametrize(
+    ("func",),
+    [
+        pytest.param(operator.add),
+        pytest.param(operator.truediv),
+    ],
+    ids=[
+        "add",
+        "divide",
+    ],
+)
+def test_combine_series_pyfunc(scalars_dfs, func):
+    scalars_df, scalars_pandas_df = scalars_dfs
+
+    bf_col = scalars_df["int64_col"].dropna()
+    bf_result = bf_col.combine(bf_col, func).to_pandas()
+
+    pd_col = scalars_pandas_df["int64_col"].dropna()
+    pd_result = pd_col.combine(pd_col, func)
+
+    assert_series_equal(bf_result, pd_result, check_dtype=False)
+
+
 def test_combine_scalar_ufunc(scalars_dfs):
     scalars_df, scalars_pandas_df = scalars_dfs