Skip to content

Commit 7d2bc21

Browse files
TrevorBergerontswast
TrevorBergeron
and
tswast
authored
feat(bigframes): Transpiler supports more string ops (#17693)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/google-cloud-python/issues) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕 --------- Co-authored-by: Tim Sweña <swast@google.com>
1 parent 654a96f commit 7d2bc21

5 files changed

Lines changed: 426 additions & 7 deletions

File tree

‎packages/bigframes/bigframes/core/bytecode.py‎

Lines changed: 161 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -87,18 +87,26 @@
8787
"POP_JUMP_BACKWARD_IF_TRUE",
8888
}
8989

90+
_JUMP_IF_NONE_OPNAMES = {
91+
"POP_JUMP_IF_NONE",
92+
"POP_JUMP_FORWARD_IF_NONE",
93+
"POP_JUMP_BACKWARD_IF_NONE",
94+
}
95+
96+
_JUMP_IF_NOT_NONE_OPNAMES = {
97+
"POP_JUMP_IF_NOT_NONE",
98+
"POP_JUMP_FORWARD_IF_NOT_NONE",
99+
"POP_JUMP_BACKWARD_IF_NOT_NONE",
100+
}
101+
90102
_CONDITIONAL_JUMP_OPNAMES = (
91103
_JUMP_IF_FALSE_OPNAMES
92104
| _JUMP_IF_TRUE_OPNAMES
105+
| _JUMP_IF_NONE_OPNAMES
106+
| _JUMP_IF_NOT_NONE_OPNAMES
93107
| {
94108
"JUMP_IF_FALSE_OR_POP",
95109
"JUMP_IF_TRUE_OR_POP",
96-
"POP_JUMP_IF_NONE",
97-
"POP_JUMP_IF_NOT_NONE",
98-
"POP_JUMP_FORWARD_IF_NONE",
99-
"POP_JUMP_FORWARD_IF_NOT_NONE",
100-
"POP_JUMP_BACKWARD_IF_NONE",
101-
"POP_JUMP_BACKWARD_IF_NOT_NONE",
102110
}
103111
)
104112

@@ -424,7 +432,10 @@ def _compile_bytecode_to_py_expr(func: Callable) -> expression.Expression:
424432
and (inst.arg & 1)
425433
)
426434
if is_method_lookup:
427-
if isinstance(target, py_exprs.Module):
435+
if isinstance(target, py_exprs.Module) or (
436+
isinstance(target, py_exprs.PyObject)
437+
and isinstance(target.value, type)
438+
):
428439
stack.append(_NULL)
429440
else:
430441
stack.append(target)
@@ -443,6 +454,66 @@ def _compile_bytecode_to_py_expr(func: Callable) -> expression.Expression:
443454
)
444455
)
445456

457+
case "FORMAT_SIMPLE":
458+
if not stack:
459+
raise ValueError("Stack is empty")
460+
value = stack.pop()
461+
stack.append(py_exprs.Call(py_exprs.PyObject(str), (value,)))
462+
463+
case "CONVERT_VALUE":
464+
flags = inst.arg
465+
assert flags is not None
466+
value = stack.pop()
467+
if flags == 1:
468+
stack.append(py_exprs.Call(py_exprs.PyObject(str), (value,)))
469+
else:
470+
raise NotImplementedError(
471+
"repr() and ascii() conversions are not supported"
472+
)
473+
474+
case "FORMAT_VALUE":
475+
flags = inst.arg
476+
assert flags is not None
477+
if (flags & 0x04) == 0x04:
478+
stack.pop()
479+
raise NotImplementedError(
480+
"Formatting with specifier is not supported"
481+
)
482+
483+
value = stack.pop()
484+
conversion = flags & 0x03
485+
if conversion == 0 or conversion == 1:
486+
stack.append(py_exprs.Call(py_exprs.PyObject(str), (value,)))
487+
else:
488+
raise NotImplementedError(
489+
"repr() and ascii() conversions are not supported"
490+
)
491+
492+
case "FORMAT_WITH_SPEC":
493+
raise NotImplementedError(
494+
"Formatting with specifier is not supported"
495+
)
496+
497+
case "BUILD_STRING":
498+
count = inst.arg
499+
assert count is not None
500+
if len(stack) < count:
501+
raise ValueError(
502+
"Stack has fewer elements than BUILD_STRING count"
503+
)
504+
505+
if count == 0:
506+
stack.append(py_exprs.PyObject(""))
507+
else:
508+
strings = [stack.pop() for _ in range(count)][::-1]
509+
result = strings[0]
510+
for s in strings[1:]:
511+
result = py_exprs.Call(
512+
py_exprs.PyObject(operator.add),
513+
(result, s),
514+
)
515+
stack.append(result)
516+
446517
case "COPY":
447518
idx = inst.arg
448519
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:
533604
)
534605
)
535606

607+
case "IS_OP":
608+
if len(stack) < 2:
609+
raise ValueError("Stack has < 2 elements")
610+
right = stack.pop()
611+
left = stack.pop()
612+
invert = inst.arg
613+
614+
def is_none_const(expr) -> bool:
615+
if isinstance(expr, py_exprs.PyObject) and expr.value is None:
616+
return True
617+
if (
618+
isinstance(expr, expression.ScalarConstantExpression)
619+
and expr.value is None
620+
):
621+
return True
622+
return False
623+
624+
if is_none_const(right):
625+
op = (
626+
generic_ops.isnull_op
627+
if not invert
628+
else generic_ops.notnull_op
629+
)
630+
stack.append(py_exprs.Call(py_exprs.PyObject(op), (left,)))
631+
elif is_none_const(left):
632+
op = (
633+
generic_ops.isnull_op
634+
if not invert
635+
else generic_ops.notnull_op
636+
)
637+
stack.append(py_exprs.Call(py_exprs.PyObject(op), (right,)))
638+
else:
639+
raise NotImplementedError(
640+
"Identity comparison (is/is not) is only supported for None"
641+
)
642+
536643
case "COMPARE_OP":
537644
if len(stack) < 2:
538645
raise ValueError("Stack has < 2 elements")
@@ -728,6 +835,53 @@ def _compile_bytecode_to_py_expr(func: Callable) -> expression.Expression:
728835
jumped = True
729836
break
730837

838+
case name if (
839+
name in _JUMP_IF_NONE_OPNAMES or name in _JUMP_IF_NOT_NONE_OPNAMES
840+
):
841+
if not stack:
842+
raise ValueError("Stack is empty")
843+
cond_expr = stack.pop()
844+
cond_bool = py_exprs.Call(
845+
py_exprs.PyObject(generic_ops.isnull_op),
846+
(cond_expr,),
847+
)
848+
849+
dest = inst.argval
850+
next_offset = next_offsets.get(inst.offset)
851+
852+
if opname in _JUMP_IF_NONE_OPNAMES:
853+
not_cond_bool = py_exprs.Call(
854+
py_exprs.PyObject(operator.not_), (cond_bool,)
855+
)
856+
edge_conditions[(offset, dest)] = py_exprs.Call(
857+
py_exprs.PyObject(operator.and_),
858+
(reach_cond, cond_bool),
859+
)
860+
edge_stacks[(offset, dest)] = stack.copy()
861+
if next_offset is not None:
862+
edge_conditions[(offset, next_offset)] = py_exprs.Call(
863+
py_exprs.PyObject(operator.and_),
864+
(reach_cond, not_cond_bool),
865+
)
866+
edge_stacks[(offset, next_offset)] = stack.copy()
867+
else: # opname in _JUMP_IF_NOT_NONE_OPNAMES
868+
not_cond_bool = py_exprs.Call(
869+
py_exprs.PyObject(operator.not_), (cond_bool,)
870+
)
871+
edge_conditions[(offset, dest)] = py_exprs.Call(
872+
py_exprs.PyObject(operator.and_),
873+
(reach_cond, not_cond_bool),
874+
)
875+
edge_stacks[(offset, dest)] = stack.copy()
876+
if next_offset is not None:
877+
edge_conditions[(offset, next_offset)] = py_exprs.Call(
878+
py_exprs.PyObject(operator.and_),
879+
(reach_cond, cond_bool),
880+
)
881+
edge_stacks[(offset, next_offset)] = stack.copy()
882+
jumped = True
883+
break
884+
731885
case name if name in _ALL_JUMP_OPNAMES:
732886
raise ValueError(f"Unsupported jump opcode: {opname}")
733887

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

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -454,6 +454,54 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
454454
else:
455455
return pl.any_horizontal(*(input.str.ends_with(pat) for pat in op.pat))
456456

457+
@compile_op.register(string_ops.CapitalizeOp)
458+
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
459+
assert isinstance(op, string_ops.CapitalizeOp)
460+
return (
461+
input.str.slice(0, 1).str.to_uppercase()
462+
+ input.str.slice(1).str.to_lowercase()
463+
)
464+
465+
@compile_op.register(string_ops.IsAlnumOp)
466+
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
467+
assert isinstance(op, string_ops.IsAlnumOp)
468+
return input.str.contains(r"^[a-zA-Z0-9]+$")
469+
470+
@compile_op.register(string_ops.IsAlphaOp)
471+
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
472+
assert isinstance(op, string_ops.IsAlphaOp)
473+
return input.str.contains(r"^[a-zA-Z]+$")
474+
475+
@compile_op.register(string_ops.IsDigitOp)
476+
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
477+
assert isinstance(op, string_ops.IsDigitOp)
478+
return input.str.contains(r"^[0-9]+$")
479+
480+
@compile_op.register(string_ops.IsSpaceOp)
481+
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
482+
assert isinstance(op, string_ops.IsSpaceOp)
483+
return input.str.contains(r"^\s+$")
484+
485+
@compile_op.register(string_ops.IsDecimalOp)
486+
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
487+
assert isinstance(op, string_ops.IsDecimalOp)
488+
return input.str.contains(r"^[0-9]+$")
489+
490+
@compile_op.register(string_ops.IsNumericOp)
491+
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
492+
assert isinstance(op, string_ops.IsNumericOp)
493+
return input.str.contains(r"^[0-9]+$")
494+
495+
@compile_op.register(string_ops.IsLowerOp)
496+
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
497+
assert isinstance(op, string_ops.IsLowerOp)
498+
return input.str.contains(r"[a-z]") & ~input.str.contains(r"[A-Z]")
499+
500+
@compile_op.register(string_ops.IsUpperOp)
501+
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
502+
assert isinstance(op, string_ops.IsUpperOp)
503+
return input.str.contains(r"[A-Z]") & ~input.str.contains(r"[a-z]")
504+
457505
@compile_op.register(freq_ops.FloorDtOp)
458506
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
459507
assert isinstance(op, freq_ops.FloorDtOp)

‎packages/bigframes/bigframes/core/py_expressions.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -547,6 +547,13 @@ def resolve_call(
547547
if fn in _CALLABLE_TO_OP:
548548
op = _CALLABLE_TO_OP[fn]
549549
return OpExpression(op, call.inputs)
550+
elif isinstance(callable.input, PyObject) and isinstance(
551+
callable.input.value, type
552+
):
553+
fn = getattr(callable.input.value, attr, None)
554+
if fn in python_op_maps.PYTHON_TO_BIGFRAMES:
555+
op = python_op_maps.PYTHON_TO_BIGFRAMES[fn]
556+
return OpExpression(op, call.inputs)
550557
else:
551558
# Method call on an expression (e.g. df.col.sum() or s.mean())
552559
try:

‎packages/bigframes/bigframes/operations/python_op_maps.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,15 @@
6969
## str
7070
str.upper: string_ops.upper_op,
7171
str.lower: string_ops.lower_op,
72+
str.isalnum: string_ops.isalnum_op,
73+
str.isalpha: string_ops.isalpha_op,
74+
str.isdecimal: string_ops.isdecimal_op,
75+
str.isdigit: string_ops.isdigit_op,
76+
str.isnumeric: string_ops.isnumeric_op,
77+
str.isspace: string_ops.isspace_op,
78+
str.islower: string_ops.islower_op,
79+
str.isupper: string_ops.isupper_op,
80+
str.capitalize: string_ops.capitalize_op,
7281
## builtins
7382
len: string_ops.len_op,
7483
abs: numeric_ops.abs_op,
@@ -103,4 +112,15 @@ def python_callable_to_op(obj) -> Optional[bigframes.operations.RowOp]:
103112
"isna": generic_ops.isnull_op,
104113
"notnull": generic_ops.notnull_op,
105114
"notna": generic_ops.notnull_op,
115+
"upper": string_ops.upper_op,
116+
"lower": string_ops.lower_op,
117+
"isalnum": string_ops.isalnum_op,
118+
"isalpha": string_ops.isalpha_op,
119+
"isdecimal": string_ops.isdecimal_op,
120+
"isdigit": string_ops.isdigit_op,
121+
"isnumeric": string_ops.isnumeric_op,
122+
"isspace": string_ops.isspace_op,
123+
"islower": string_ops.islower_op,
124+
"isupper": string_ops.isupper_op,
125+
"capitalize": string_ops.capitalize_op,
106126
}

0 commit comments

Comments
 (0)