Skip to content

feat: Experimental transpilation of unannotated python callables #17419

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
TrevorBergeron wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat: Experimental transpilation of unannotated python callables #17419

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions packages/bigframes/bigframes/_config/experiment_options.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ def __init__(self):
self._semantic_operators: bool = False
self._ai_operators: bool = False
self._sql_compiler: Literal["legacy", "stable", "experimental"] = "stable"
self._enable_python_transpiler: bool = False

@property
def semantic_operators(self) -> bool:
Expand Down Expand Up @@ -166,3 +167,17 @@ def blob_display_height(self, value: Optional[int]):
warnings.warn(msg, category=bfe.ApiDeprecationWarning)

bigframes.options.display.blob_display_height = value

@property
def enable_python_transpiler(self) -> bool:
return self._enable_python_transpiler

@enable_python_transpiler.setter
def enable_python_transpiler(self, value: bool):
if value:
msg = bfe.format_message(
"Python transpiler is an unstable, experimental feature, and not yet fully "
"validated, use at your own risk."
)
warnings.warn(msg, category=bfe.PreviewWarning)
self._enable_python_transpiler = value
79 changes: 73 additions & 6 deletions packages/bigframes/bigframes/dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4692,13 +4692,17 @@ def _prepare_export(
return array_value, id_overrides

def map(self, func, na_action: Optional[str] = None) -> DataFrame:
if not isinstance(func, bigframes.functions.Udf):
from bigframes._config import options

if not isinstance(func, bigframes.functions.Udf) and not (
options.experiments.enable_python_transpiler and callable(func)
):
raise TypeError("the first argument must be callable")

if na_action not in {None, "ignore"}:
raise ValueError(f"na_action={na_action} not supported")

expr = ops.func_to_op(func).as_expr(ex.free_var("input"))
expr = ops.func_to_expr(func).apply(ex.free_var("input"))
if na_action == "ignore":
# True case, predicate, False case
expr = ops.where_op.as_expr(
Expand All @@ -4718,11 +4722,74 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
)
warnings.warn(msg, category=bfe.FunctionAxisOnePreviewWarning)

if not isinstance(func, bigframes.functions.Udf):
from bigframes._config import options

if not isinstance(func, bigframes.functions.Udf) and not (
options.experiments.enable_python_transpiler and callable(func)
):
raise ValueError(
"For axis=1 a BigFrames BigQuery function must be used."
)

if (
not isinstance(func, bigframes.functions.Udf)
and options.experiments.enable_python_transpiler
and callable(func)
):
from bigframes.operations.to_op import CallableExpression

callable_expr = CallableExpression.from_callable(
func, unpack_mode=False
)

# Bind the extra arguments (args and kwargs) starting from parameter 1
bindings = {}
# Positional arguments:
for idx, val in enumerate(args):
param_name = callable_expr.arg_specs[idx + 1].name
bindings[param_name] = val
# Keyword arguments:
for key, val in kwargs.items():
bindings[key] = val

# Bind defaults for other parameters (excluding the first 'row' parameter)
for spec in callable_expr.arg_specs[1:]:
if (
spec.name not in bindings
and spec.default_value is not inspect.Parameter.empty
):
bindings[spec.name] = spec.default_value
Comment on lines +4745 to +4761

@gemini-code-assist gemini-code-assist Bot Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The manual binding of args and kwargs lacks validation for extra positional arguments and unexpected keyword arguments. This can lead to cryptic IndexError exceptions or silent correctness bugs where unexpected keyword arguments accidentally override free variables (columns) in the expression.

                # Bind the extra arguments (args and kwargs) starting from parameter 1
                bindings = {}
                allowed_params = callable_expr.arg_specs[1:]
                allowed_names = {spec.name for spec in allowed_params}
                func_name = getattr(func, "__name__", "callable")

                # Positional arguments:
                for idx, val in enumerate(args):
                    if idx >= len(allowed_params):
                        raise TypeError(
                            f"too many positional arguments for {func_name}: expected {len(allowed_params)}, got {len(args)}"
                        )
                    param_name = allowed_params[idx].name
                    bindings[param_name] = val

                # Keyword arguments:
                for key, val in kwargs.items():
                    if key not in allowed_names:
                        raise TypeError(
                            f"{func_name}() got an unexpected keyword argument '{key}'"
                        )
                    if key in bindings:
                        raise TypeError(
                            f"{func_name}() got multiple values for argument '{key}'"
                        )
                    bindings[key] = val

                # Bind defaults for other parameters (excluding the first 'row' parameter)
                for spec in allowed_params:
                    if (
                        spec.name not in bindings
                        and spec.default_value is not inspect.Parameter.empty
                    ):
                        bindings[spec.name] = spec.default_value


# Wrap all values in bindings as expressions
def to_expr(val):
if isinstance(val, ex.Expression):
return val
return ex.const(val)

bindings = {k: to_expr(v) for k, v in bindings.items()}

# Now bind these variables in the expression!
expr = callable_expr.expr.bind_variables(
bindings, allow_partial_bindings=True
)

# Now bind the remaining free variables to the DataFrame columns:
col_bindings = {}
block = self._get_block()
for col in self.columns:
if col in expr.free_variables:
col_id = block.resolve_label_exact(col)
if col_id is not None:
col_bindings[col] = ex.deref(col_id)

expr = expr.bind_variables(col_bindings)

# Project the expression on the DataFrame block to get a new Series!
block, result_id = self._get_block().project_expr(expr)
from bigframes.series import Series

return Series(block.select_column(result_id))

if func.udf_def.signature.is_row_processor:
# Early check whether the dataframe dtypes are currently supported
# in the bigquery function
Expand Down Expand Up @@ -4777,7 +4844,7 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):

# Apply the function
result_series = rows_as_json_series._apply_nary_op(
ops.func_to_op(func),
ops.func_to_expr(func).expr.op,
list(args),
)

Expand Down Expand Up @@ -4837,8 +4904,8 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):

series_list = [self[col] for col in self.columns]
op_list = series_list[1:] + list(args)
result_series = series_list[0]._apply_nary_op(
ops.func_to_op(func), op_list
result_series = series_list[0]._apply_callable_expr(
ops.func_to_expr(func), op_list
)
result_series.name = None

Expand Down
4 changes: 2 additions & 2 deletions packages/bigframes/bigframes/operations/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,7 @@
timestamp_add_op,
timestamp_sub_op,
)
from bigframes.operations.to_op import func_to_op
from bigframes.operations.to_op import func_to_expr

__all__ = [
# Base ops
Expand Down Expand Up @@ -437,7 +437,7 @@
"AIScore",
"AISimilarity",
# Helper functions
"func_to_op",
"func_to_expr",
# Numpy ops mapping
"NUMPY_TO_BINOP",
"NUMPY_TO_OP",
Expand Down
178 changes: 166 additions & 12 deletions packages/bigframes/bigframes/operations/to_op.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,31 +11,185 @@
# 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.
from __future__ import annotations

import dataclasses
import inspect
import typing

import bigframes.core.expression as ex
import bigframes.core.identifiers as ids
import bigframes.dtypes as dtypes
from bigframes._config import options
from bigframes.functions import Udf
from bigframes.functions.udf_def import BigqueryUdf, PythonUdf
from bigframes.operations import base_ops, remote_function_ops
from bigframes.operations import remote_function_ops


@dataclasses.dataclass(frozen=True)
class ArgumentSpec:
"""
Information about a single argument to a function
"""

name: str
default_value: typing.Any
is_varargs: bool


def func_to_op(op) -> base_ops.NaryOp:
@dataclasses.dataclass(frozen=True)
class CallableExpression(ex.Expression):
"""
Convert various bigframes, python functions into bigframes operations.
Encodes a calling convention and an expression to bind arguments to.
"""

expr: ex.Expression
arg_specs: typing.Sequence[ArgumentSpec]

@classmethod
def from_callable(
cls, func: typing.Callable, unpack_mode: bool = False
) -> CallableExpression:
sig = inspect.signature(func)
arg_specs = []
for name, param in sig.parameters.items():
is_varargs = param.kind == inspect.Parameter.VAR_POSITIONAL
arg_specs.append(
ArgumentSpec(
name=name,
default_value=param.default,
is_varargs=is_varargs,
)
)

from bigframes.core.bytecode import dis_to_expr

expr = dis_to_expr(func, unpack_mode=unpack_mode)
return cls(expr=expr, arg_specs=arg_specs)

def apply(self, *args, **kwargs) -> ex.Expression:
"""
Apply the arguments to the expression.

All args are expected to be column references, or scalars.
"""
bindings = {}
pos_idx = 0

def to_expr(val):
if isinstance(val, ex.Expression):
return val
return ex.const(val)

for spec in self.arg_specs:
if spec.is_varargs:
raise NotImplementedError(
"varargs in compiled python functions is not supported"
)

This should handle anything that might be passed to eg map, combine, other pandas methods that take a function.
if pos_idx < len(args):
bindings[spec.name] = to_expr(args[pos_idx])
pos_idx += 1
elif spec.name in kwargs:
bindings[spec.name] = to_expr(kwargs[spec.name])
elif spec.default_value is not inspect.Parameter.empty:
bindings[spec.name] = to_expr(spec.default_value)
else:
raise TypeError(f"missing required argument: '{spec.name}'")

It should raise a TypeError if the object is not a supported type.
if pos_idx < len(args):
raise TypeError(
f"too many positional arguments: expected {len(self.arg_specs)}, got {len(args)}"
)

Args:
op: The object to convert.
return self.expr.bind_variables(bindings)
Comment on lines +76 to +105

@gemini-code-assist gemini-code-assist Bot Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

CallableExpression.apply does not validate unexpected keyword arguments or check if a keyword argument has already been supplied as a positional argument. This can lead to silent correctness bugs where keyword arguments are ignored or unexpected arguments are accepted.

        bindings = {}
        pos_idx = 0
        arg_names = {spec.name for spec in self.arg_specs}

        # Validate unexpected keyword arguments
        for key in kwargs:
            if key not in arg_names:
                raise TypeError(f"got an unexpected keyword argument '{key}'")

        def to_expr(val):
            if isinstance(val, ex.Expression):
                return val
            return ex.const(val)

        for spec in self.arg_specs:
            if spec.is_varargs:
                raise NotImplementedError(
                    "varargs in compiled python functions is not supported"
                )

            if pos_idx < len(args):
                if spec.name in kwargs:
                    raise TypeError(f"got multiple values for keyword argument '{spec.name}'")
                bindings[spec.name] = to_expr(args[pos_idx])
                pos_idx += 1
            elif spec.name in kwargs:
                bindings[spec.name] = to_expr(kwargs[spec.name])
            elif spec.default_value is not inspect.Parameter.empty:
                bindings[spec.name] = to_expr(spec.default_value)
            else:
                raise TypeError(f"missing required argument: '{spec.name}'")

        if pos_idx < len(args):
            raise TypeError(
                f"too many positional arguments: expected {len(self.arg_specs)}, got {len(args)}"
            )

        return self.expr.bind_variables(bindings)


Returns:
A bigframes operations.
@property
def column_references(self) -> typing.Tuple[ids.ColumnId, ...]:
return self.expr.column_references

@property
def free_variables(self) -> typing.Tuple[typing.Hashable, ...]:
return self.expr.free_variables

@property
def is_const(self) -> bool:
return self.expr.is_const

@property
def is_resolved(self) -> bool:
return False

@property
def output_type(self) -> dtypes.ExpressionType:
raise ValueError(
"CallableExpression does not have a fixed output type until arguments are applied."
)

def bind_refs(
self,
bindings: typing.Mapping[ids.ColumnId, ex.Expression],
allow_partial_bindings: bool = False,
) -> CallableExpression:
return dataclasses.replace(
self,
expr=self.expr.bind_refs(
bindings, allow_partial_bindings=allow_partial_bindings
),
)

def bind_variables(
self,
bindings: typing.Mapping[typing.Hashable, ex.Expression],
allow_partial_bindings: bool = False,
) -> CallableExpression:
arg_names = {spec.name for spec in self.arg_specs}
filtered_bindings = {k: v for k, v in bindings.items() if k not in arg_names}
return dataclasses.replace(
self,
expr=self.expr.bind_variables(
filtered_bindings, allow_partial_bindings=allow_partial_bindings
),
)

def transform_children(
self, t: typing.Callable[[ex.Expression], ex.Expression]
) -> ex.Expression:
new_expr = t(self.expr)
if new_expr != self.expr:
return dataclasses.replace(self, expr=new_expr)
return self


def func_to_expr(op, unpack_mode: bool = False) -> CallableExpression:
"""
Convert various bigframes, python functions into bigframes CallableExpression.
"""
# TODO(b/517578802): Handle numpy ufuncs, builtin functions, etc.
if isinstance(op, Udf):
if isinstance(op.udf_def, BigqueryUdf):
return remote_function_ops.RemoteFunctionOp(function_def=op.udf_def)
bq_op = remote_function_ops.RemoteFunctionOp(function_def=op.udf_def)
elif isinstance(op.udf_def, PythonUdf):
return remote_function_ops.PythonUdfOp(function_def=op.udf_def)
bq_op = remote_function_ops.PythonUdfOp(function_def=op.udf_def)
else:
raise TypeError(f"Unsupported UDF definition: {op.udf_def}")

inputs_expr = tuple(
ex.free_var(arg.name) for arg in op.udf_def.signature.inputs
)
expr = ex.OpExpression(bq_op, inputs_expr)

arg_specs = [
ArgumentSpec(
name=arg.name,
default_value=inspect.Parameter.empty,
is_varargs=False,
)
for arg in op.udf_def.signature.inputs
]
return CallableExpression(expr=expr, arg_specs=arg_specs)

elif options.experiments.enable_python_transpiler and callable(op):
return CallableExpression.from_callable(op, unpack_mode=unpack_mode)

else:
raise TypeError(f"Unsupported function type: {op}")
Loading
Loading