feat(bigframes): Transpiler supports more string ops

packages/bigframes/bigframes/core/bytecode.py

@@ -87,18 +87,26 @@
     "POP_JUMP_BACKWARD_IF_TRUE",
 }
 
+_JUMP_IF_NONE_OPNAMES = {
+    "POP_JUMP_IF_NONE",
+    "POP_JUMP_FORWARD_IF_NONE",
+    "POP_JUMP_BACKWARD_IF_NONE",
+}
+
+_JUMP_IF_NOT_NONE_OPNAMES = {
+    "POP_JUMP_IF_NOT_NONE",
+    "POP_JUMP_FORWARD_IF_NOT_NONE",
+    "POP_JUMP_BACKWARD_IF_NOT_NONE",
+}
+
 _CONDITIONAL_JUMP_OPNAMES = (
     _JUMP_IF_FALSE_OPNAMES
     | _JUMP_IF_TRUE_OPNAMES
+    | _JUMP_IF_NONE_OPNAMES
+    | _JUMP_IF_NOT_NONE_OPNAMES
     | {
         "JUMP_IF_FALSE_OR_POP",
         "JUMP_IF_TRUE_OR_POP",
-        "POP_JUMP_IF_NONE",
-        "POP_JUMP_IF_NOT_NONE",
-        "POP_JUMP_FORWARD_IF_NONE",
-        "POP_JUMP_FORWARD_IF_NOT_NONE",
-        "POP_JUMP_BACKWARD_IF_NONE",
-        "POP_JUMP_BACKWARD_IF_NOT_NONE",
     }
 )
 
@@ -424,7 +432,10 @@ def _compile_bytecode_to_py_expr(func: Callable) -> expression.Expression:
                         and (inst.arg & 1)
                     )
                     if is_method_lookup:
-                        if isinstance(target, py_exprs.Module):
+                        if isinstance(target, py_exprs.Module) or (
+                            isinstance(target, py_exprs.PyObject)
+                            and isinstance(target.value, type)
+                        ):
                             stack.append(_NULL)
                         else:
                             stack.append(target)
@@ -443,6 +454,66 @@ def _compile_bytecode_to_py_expr(func: Callable) -> expression.Expression:
                         )
                     )
 
+                case "FORMAT_SIMPLE":
+                    if not stack:
+                        raise ValueError("Stack is empty")
+                    value = stack.pop()
+                    stack.append(py_exprs.Call(py_exprs.PyObject(str), (value,)))
+
+                case "CONVERT_VALUE":
+                    flags = inst.arg
+                    assert flags is not None
+                    value = stack.pop()
+                    if flags == 1:
+                        stack.append(py_exprs.Call(py_exprs.PyObject(str), (value,)))
+                    else:
+                        raise NotImplementedError(
+                            "repr() and ascii() conversions are not supported"
+                        )
+
+                case "FORMAT_VALUE":
+                    flags = inst.arg
+                    assert flags is not None
+                    if (flags & 0x04) == 0x04:
+                        stack.pop()
+                        raise NotImplementedError(
+                            "Formatting with specifier is not supported"
+                        )
+
+                    value = stack.pop()
+                    conversion = flags & 0x03
+                    if conversion == 0 or conversion == 1:
+                        stack.append(py_exprs.Call(py_exprs.PyObject(str), (value,)))
+                    else:
+                        raise NotImplementedError(
+                            "repr() and ascii() conversions are not supported"
+                        )
+
+                case "FORMAT_WITH_SPEC":
+                    raise NotImplementedError(
+                        "Formatting with specifier is not supported"
+                    )
+
+                case "BUILD_STRING":
+                    count = inst.arg
+                    assert count is not None
+                    if len(stack) < count:
+                        raise ValueError(
+                            "Stack has fewer elements than BUILD_STRING count"
+                        )
+
+                    if count == 0:
+                        stack.append(py_exprs.PyObject(""))
+                    else:
+                        strings = [stack.pop() for _ in range(count)][::-1]
+                        result = strings[0]
+                        for s in strings[1:]:
+                            result = py_exprs.Call(
+                                py_exprs.PyObject(operator.add),
+                                (result, s),
+                            )
+                        stack.append(result)
+
                 case "COPY":
                     idx = inst.arg
                     if idx is None or idx < 1 or len(stack) < idx:
@@ -533,6 +604,42 @@ def _compile_bytecode_to_py_expr(func: Callable) -> expression.Expression:
                         )
                     )
 
+                case "IS_OP":
+                    if len(stack) < 2:
+                        raise ValueError("Stack has < 2 elements")
+                    right = stack.pop()
+                    left = stack.pop()
+                    invert = inst.arg
+
+                    def is_none_const(expr) -> bool:
+                        if isinstance(expr, py_exprs.PyObject) and expr.value is None:
+                            return True
+                        if (
+                            isinstance(expr, expression.ScalarConstantExpression)
+                            and expr.value is None
+                        ):
+                            return True
+                        return False
+
+                    if is_none_const(right):
+                        op = (
+                            generic_ops.isnull_op
+                            if not invert
+                            else generic_ops.notnull_op
+                        )
+                        stack.append(py_exprs.Call(py_exprs.PyObject(op), (left,)))
+                    elif is_none_const(left):
+                        op = (
+                            generic_ops.isnull_op
+                            if not invert
+                            else generic_ops.notnull_op
+                        )
+                        stack.append(py_exprs.Call(py_exprs.PyObject(op), (right,)))
+                    else:
+                        raise NotImplementedError(
+                            "Identity comparison (is/is not) is only supported for None"
+                        )
+
                 case "COMPARE_OP":
                     if len(stack) < 2:
                         raise ValueError("Stack has < 2 elements")
@@ -728,6 +835,53 @@ def _compile_bytecode_to_py_expr(func: Callable) -> expression.Expression:
                     jumped = True
                     break
 
+                case name if (
+                    name in _JUMP_IF_NONE_OPNAMES or name in _JUMP_IF_NOT_NONE_OPNAMES
+                ):
+                    if not stack:
+                        raise ValueError("Stack is empty")
+                    cond_expr = stack.pop()
+                    cond_bool = py_exprs.Call(
+                        py_exprs.PyObject(generic_ops.isnull_op),
+                        (cond_expr,),
+                    )
+
+                    dest = inst.argval
+                    next_offset = next_offsets.get(inst.offset)
+
+                    if opname in _JUMP_IF_NONE_OPNAMES:
+                        not_cond_bool = py_exprs.Call(
+                            py_exprs.PyObject(operator.not_), (cond_bool,)
+                        )
+                        edge_conditions[(offset, dest)] = py_exprs.Call(
+                            py_exprs.PyObject(operator.and_),
+                            (reach_cond, cond_bool),
+                        )
+                        edge_stacks[(offset, dest)] = stack.copy()
+                        if next_offset is not None:
+                            edge_conditions[(offset, next_offset)] = py_exprs.Call(
+                                py_exprs.PyObject(operator.and_),
+                                (reach_cond, not_cond_bool),
+                            )
+                            edge_stacks[(offset, next_offset)] = stack.copy()
+                    else:  # opname in _JUMP_IF_NOT_NONE_OPNAMES
+                        not_cond_bool = py_exprs.Call(
+                            py_exprs.PyObject(operator.not_), (cond_bool,)
+                        )
+                        edge_conditions[(offset, dest)] = py_exprs.Call(
+                            py_exprs.PyObject(operator.and_),
+                            (reach_cond, not_cond_bool),
+                        )
+                        edge_stacks[(offset, dest)] = stack.copy()
+                        if next_offset is not None:
+                            edge_conditions[(offset, next_offset)] = py_exprs.Call(
+                                py_exprs.PyObject(operator.and_),
+                                (reach_cond, cond_bool),
+                            )
+                            edge_stacks[(offset, next_offset)] = stack.copy()
+                    jumped = True
+                    break
+
                 case name if name in _ALL_JUMP_OPNAMES:
                     raise ValueError(f"Unsupported jump opcode: {opname}")
 

packages/bigframes/bigframes/core/compile/polars/compiler.py

@@ -454,6 +454,54 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
             else:
                 return pl.any_horizontal(*(input.str.ends_with(pat) for pat in op.pat))
 
+        @compile_op.register(string_ops.CapitalizeOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.CapitalizeOp)
+            return (
+                input.str.slice(0, 1).str.to_uppercase()
+                + input.str.slice(1).str.to_lowercase()
+            )
+
+        @compile_op.register(string_ops.IsAlnumOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.IsAlnumOp)
+            return input.str.contains(r"^[a-zA-Z0-9]+$")
+
+        @compile_op.register(string_ops.IsAlphaOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.IsAlphaOp)
+            return input.str.contains(r"^[a-zA-Z]+$")
+
+        @compile_op.register(string_ops.IsDigitOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.IsDigitOp)
+            return input.str.contains(r"^[0-9]+$")
+
+        @compile_op.register(string_ops.IsSpaceOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.IsSpaceOp)
+            return input.str.contains(r"^\s+$")
+
+        @compile_op.register(string_ops.IsDecimalOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.IsDecimalOp)
+            return input.str.contains(r"^[0-9]+$")
+
+        @compile_op.register(string_ops.IsNumericOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.IsNumericOp)
+            return input.str.contains(r"^[0-9]+$")
+
+        @compile_op.register(string_ops.IsLowerOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.IsLowerOp)
+            return input.str.contains(r"[a-z]") & ~input.str.contains(r"[A-Z]")
+
+        @compile_op.register(string_ops.IsUpperOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.IsUpperOp)
+            return input.str.contains(r"[A-Z]") & ~input.str.contains(r"[a-z]")
+
         @compile_op.register(freq_ops.FloorDtOp)
         def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
             assert isinstance(op, freq_ops.FloorDtOp)

packages/bigframes/bigframes/core/py_expressions.py

@@ -547,6 +547,13 @@ def resolve_call(
             if fn in _CALLABLE_TO_OP:
                 op = _CALLABLE_TO_OP[fn]
                 return OpExpression(op, call.inputs)
+        elif isinstance(callable.input, PyObject) and isinstance(
+            callable.input.value, type
+        ):
+            fn = getattr(callable.input.value, attr, None)
+            if fn in python_op_maps.PYTHON_TO_BIGFRAMES:
+                op = python_op_maps.PYTHON_TO_BIGFRAMES[fn]
+                return OpExpression(op, call.inputs)
         else:
             # Method call on an expression (e.g. df.col.sum() or s.mean())
             try:

packages/bigframes/bigframes/operations/python_op_maps.py

@@ -69,6 +69,15 @@
     ## str
     str.upper: string_ops.upper_op,
     str.lower: string_ops.lower_op,
+    str.isalnum: string_ops.isalnum_op,
+    str.isalpha: string_ops.isalpha_op,
+    str.isdecimal: string_ops.isdecimal_op,
+    str.isdigit: string_ops.isdigit_op,
+    str.isnumeric: string_ops.isnumeric_op,
+    str.isspace: string_ops.isspace_op,
+    str.islower: string_ops.islower_op,
+    str.isupper: string_ops.isupper_op,
+    str.capitalize: string_ops.capitalize_op,
     ## builtins
     len: string_ops.len_op,
     abs: numeric_ops.abs_op,
@@ -103,4 +112,15 @@ def python_callable_to_op(obj) -> Optional[bigframes.operations.RowOp]:
     "isna": generic_ops.isnull_op,
     "notnull": generic_ops.notnull_op,
     "notna": generic_ops.notnull_op,
+    "upper": string_ops.upper_op,
+    "lower": string_ops.lower_op,
+    "isalnum": string_ops.isalnum_op,
+    "isalpha": string_ops.isalpha_op,
+    "isdecimal": string_ops.isdecimal_op,
+    "isdigit": string_ops.isdigit_op,
+    "isnumeric": string_ops.isnumeric_op,
+    "isspace": string_ops.isspace_op,
+    "islower": string_ops.islower_op,
+    "isupper": string_ops.isupper_op,
+    "capitalize": string_ops.capitalize_op,
 }