diff --git a/packages/bigframes/bigframes/core/bytecode.py b/packages/bigframes/bigframes/core/bytecode.py index d62cd8d5136c..f657ce707ea6 100644 --- a/packages/bigframes/bigframes/core/bytecode.py +++ b/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}") diff --git a/packages/bigframes/bigframes/core/compile/polars/compiler.py b/packages/bigframes/bigframes/core/compile/polars/compiler.py index ccb3d8ef25bb..65b217602b96 100644 --- a/packages/bigframes/bigframes/core/compile/polars/compiler.py +++ b/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) diff --git a/packages/bigframes/bigframes/core/py_expressions.py b/packages/bigframes/bigframes/core/py_expressions.py index 6aa7d45ec791..da9376770508 100644 --- a/packages/bigframes/bigframes/core/py_expressions.py +++ b/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: diff --git a/packages/bigframes/bigframes/operations/python_op_maps.py b/packages/bigframes/bigframes/operations/python_op_maps.py index 4917b1cf8892..b4c58e14c7b0 100644 --- a/packages/bigframes/bigframes/operations/python_op_maps.py +++ b/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, } diff --git a/packages/bigframes/tests/unit/test_py_udf.py b/packages/bigframes/tests/unit/test_py_udf.py index 59f46d65eb2a..dcd6fa286583 100644 --- a/packages/bigframes/tests/unit/test_py_udf.py +++ b/packages/bigframes/tests/unit/test_py_udf.py @@ -24,6 +24,7 @@ import bigframes import bigframes.core.global_session import bigframes.pandas as bpd +from bigframes.core.bytecode import py_to_expression from bigframes.testing.utils import ( assert_frame_equal, assert_series_equal, @@ -569,3 +570,192 @@ def foo(row): pd_result = pd_df.apply(foo, axis=1).astype("Int64") assert_series_equal(bf_result, pd_result) + + +def test_series_apply_fstrings(session): + pd_series = pd.Series(["apple", "banana", None], dtype="string") + bf_series = bpd.Series(pd_series, session=session) + + def format_udf(x): + if x is None: + return "Null value" + return f"Fruit: {x}!" + + bf_result = bf_series.apply(format_udf).to_pandas() + pd_result = pd.Series( + ["Fruit: apple!", "Fruit: banana!", "Null value"], dtype="string" + ) + assert_series_equal(bf_result, pd_result, check_dtype=False) + + +def test_series_apply_nullity_jumps(session): + pd_series = pd.Series([10, None, 20], dtype="Int64") + bf_series = bpd.Series(pd_series, session=session) + + def nullity_udf(x): + if x is None: + return "Absent" + if x is not None: + return "Present" + return "Unknown" + + bf_result = bf_series.apply(nullity_udf).to_pandas() + pd_result = pd.Series(["Present", "Absent", "Present"], dtype="string") + assert_series_equal(bf_result, pd_result, check_dtype=False) + + +def test_series_apply_string_ops(session): + pd_series = pd.Series(["hello world", "BigFrames", "123a"], dtype="string") + bf_series = bpd.Series(pd_series, session=session) + + def str_udf(x): + if x is None: + return None + return x.upper() + " " + x.lower() + " " + str.upper(x) + " " + x.capitalize() + + bf_result = bf_series.apply(str_udf).to_pandas() + pd_result = pd.Series( + [ + "HELLO WORLD hello world HELLO WORLD Hello world", + "BIGFRAMES bigframes BIGFRAMES Bigframes", + "123A 123a 123A 123a", + ], + dtype="string", + ) + assert_series_equal(bf_result, pd_result, check_dtype=False) + + +def test_series_apply_string_predicates(session): + pd_series = pd.Series( + ["hello world", "abc123", "123", "HELLO!", None], dtype="string" + ) + bf_series = bpd.Series(pd_series, session=session) + + def predicates_udf(x): + if x is None: + return None + return f"{x.islower()}_{x.isupper()}" + + bf_result = bf_series.apply(predicates_udf).to_pandas() + pd_result = pd.Series( + ["True_False", "True_False", "False_False", "False_True", None], dtype="string" + ) + assert_series_equal(bf_result, pd_result, check_dtype=False) + + +def test_fstring_multiple_placeholders(session): + pd_series = pd.Series(["apple", "banana"], dtype="string") + bf_series = bpd.Series(pd_series, session=session) + + def multiple_placeholders(x): + return f"{x} and {x.upper()}!" + + bf_res = bf_series.apply(multiple_placeholders).to_pandas() + pd_res = pd.Series(["apple and APPLE!", "banana and BANANA!"], dtype="string") + assert_series_equal(bf_res, pd_res, check_dtype=False) + + +def test_fstring_empty(session): + pd_series = pd.Series(["apple", "banana"], dtype="string") + bf_series = bpd.Series(pd_series, session=session) + + def empty_fstring(x): + return "" + + bf_res = bf_series.apply(empty_fstring).to_pandas() + pd_res = pd.Series(["", ""], dtype="string") + assert_series_equal(bf_res, pd_res, check_dtype=False) + + +def test_fstring_consecutive_placeholders(session): + pd_series = pd.Series(["apple", "banana"], dtype="string") + bf_series = bpd.Series(pd_series, session=session) + + def consecutive_placeholders(x): + return f"{x}{x.upper()}" + + bf_res = bf_series.apply(consecutive_placeholders).to_pandas() + pd_res = pd.Series(["appleAPPLE", "bananaBANANA"], dtype="string") + assert_series_equal(bf_res, pd_res, check_dtype=False) + + +def test_fstring_with_specifier_raises(): + def format_with_spec(x): + return f"{x:2d}" + + with pytest.raises( + NotImplementedError, match="Formatting with specifier is not supported" + ): + py_to_expression(format_with_spec) + + +def test_fstring_with_repr_raises(): + def format_with_repr(x): + return f"{x!r}" + + with pytest.raises( + NotImplementedError, + match="repr\\(\\) and ascii\\(\\) conversions are not supported", + ): + py_to_expression(format_with_repr) + + +def test_fstring_with_ascii_raises(): + def format_with_ascii(x): + return f"{x!a}" + + with pytest.raises( + NotImplementedError, + match="repr\\(\\) and ascii\\(\\) conversions are not supported", + ): + py_to_expression(format_with_ascii) + + +def test_identity_unsupported_raises(): + def is_true_udf(x): + return x is True + + with pytest.raises( + NotImplementedError, + match="Identity comparison \\(is/is not\\) is only supported for None", + ): + py_to_expression(is_true_udf) + + +def test_fstring_int_input(session): + pd_int_series = pd.Series([10, 20, None], dtype="Int64") + bf_int_series = bpd.Series(pd_int_series, session=session) + bf_res_int = bf_int_series.apply(lambda x: f"val: {x}").to_pandas() + pd_res_int = pd.Series(["val: 10", "val: 20", None], dtype="string") + assert_series_equal(bf_res_int, pd_res_int, check_dtype=False) + + +def test_fstring_float_input(session): + pd_float_series = pd.Series([1.5, 2.75], dtype="Float64") + bf_float_series = bpd.Series(pd_float_series, session=session) + bf_res_float = bf_float_series.apply(lambda x: f"val: {x}").to_pandas() + pd_res_float = pd.Series(["val: 1.5", "val: 2.75"], dtype="string") + assert_series_equal(bf_res_float, pd_res_float, check_dtype=False) + + +def test_fstring_bool_input(session): + pd_bool_series = pd.Series([True, False], dtype="boolean") + bf_bool_series = bpd.Series(pd_bool_series, session=session) + bf_res_bool = bf_bool_series.apply(lambda x: f"val: {x}").to_pandas() + pd_res_bool = pd.Series(["val: True", "val: False"], dtype="string") + assert_series_equal(bf_res_bool, pd_res_bool, check_dtype=False) + + +def test_fstring_list_input_raises(session): + array_pa_type = pa.list_(pa.int64()) + pd_series = pd.Series( + pa.array([[10, 20]], array_pa_type), + dtype=pd.ArrowDtype(array_pa_type), + ) + bf_series = bpd.Series(pd_series, session=session) + + def udf_with_list(x): + return f"list: {x}" + + with pytest.raises((TypeError, ValueError)): + bf_series.apply(udf_with_list)