feat: Experimental transpilation of unannotated python callables by TrevorBergeron · Pull Request #17419 · googleapis/google-cloud-python

TrevorBergeron · 2026-06-11T00:14:33Z

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 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> 🦕

gemini-code-assist

Code Review

This pull request introduces an experimental Python transpiler feature, allowing standard Python callables to be transpiled and executed within DataFrame.map, DataFrame.apply, Series.apply, and Series.combine. It replaces the func_to_op helper with func_to_expr, which generates a new CallableExpression to manage argument binding. Feedback on the changes highlights critical validation gaps in both DataFrame.apply and CallableExpression.apply where missing checks for unexpected keyword arguments, duplicate arguments, or extra positional arguments could lead to silent correctness bugs or cryptic errors.

tswast · 2026-06-16T19:21:43Z

System test looks like it might indicate a real error, but probably unrelated to this change:

___________________ test_read_parquet_gcs[bigquery_wildcard] ___________________
[gw19] linux -- Python 3.12.12 /tmpfs/src/github/google-cloud-python/packages/bigframes/.nox/system/bin/python

session = 
scalars_dfs = (          bool_col                                          bytes_col  \
rowindex                                    ...7             True  ...  0 days 00:00:00.000004
8            False  ...         5 days 00:00:00

[9 rows x 14 columns])
gcs_folder = 'gs://bigframes-dev-testing/bigframes_tests_system_20260616000152_914011/'
engine = 'bigquery', filename = '*.parquet'

    @pytest.mark.parametrize(
        ("engine", "filename"),
        (
            pytest.param(
                "auto",
                "000000000000.parquet",
                id="auto",
            ),
            pytest.param(
                "pyarrow",
                "000000000000.parquet",
                id="pyarrow",
            ),
            pytest.param(
                "bigquery",
                "000000000000.parquet",
                id="bigquery",
            ),
            pytest.param(
                "bigquery",
                "*.parquet",
                id="bigquery_wildcard",
            ),
            pytest.param(
                "auto",
                "*.parquet",
                id="auto_wildcard",
                marks=pytest.mark.xfail(
                    raises=ValueError,
                ),
            ),
        ),
    )
    def test_read_parquet_gcs(
        session: bigframes.Session, scalars_dfs, gcs_folder, engine, filename
    ):
        scalars_df, _ = scalars_dfs
        # Include wildcard so that multiple files can be written/read if > 1 GB.
        # https://cloud.google.com/bigquery/docs/exporting-data#exporting_data_into_one_or_more_files
        write_path = gcs_folder + test_read_parquet_gcs.__name__ + "*.parquet"
        read_path = gcs_folder + test_read_parquet_gcs.__name__ + filename
    
        df_in: bigframes.dataframe.DataFrame = scalars_df.copy()
        # GEOGRAPHY not supported in parquet export.
        df_in = df_in.drop(columns="geography_col")
        # Make sure we can also serialize the order.
        df_write = df_in.reset_index(drop=False)
        df_write.index.name = f"ordering_id_{random.randrange(1_000_000)}"
        df_write.to_parquet(write_path, index=True)
    
        df_out = (
            session.read_parquet(read_path, engine=engine)
            # Restore order.
>           .set_index(df_write.index.name)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
            .sort_index()
            # Restore index.
            .set_index(typing.cast(str, df_in.index.name))
        )

tests/system/small/test_session.py:1916: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
bigframes/core/logging/log_adapter.py:183: in wrapper
    return method(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self =     ordering_id_430420  rowindex  bool_col  \
0                          4     False   
2                     ...                  432000000000  
10                                  432000000000  
...

[36 rows x 15 columns]
keys = ('ordering_id_463089',), append = False, drop = True

    def set_index(
        self,
        keys: typing.Union[blocks.Label, typing.Sequence[blocks.Label]],
        append: bool = False,
        drop: bool = True,
    ) -> DataFrame:
        if not utils.is_list_like(keys):
            keys = typing.cast(typing.Sequence[blocks.Label], (keys,))
        else:
            keys = typing.cast(typing.Sequence[blocks.Label], tuple(keys))
        col_ids = [self._resolve_label_exact(key) for key in keys]
        missing = [keys[i] for i in range(len(col_ids)) if col_ids[i] is None]
        if len(missing) > 0:
>           raise KeyError(f"None of {missing} are in the columns")
E           KeyError: "None of ['ordering_id_463089'] are in the columns"

bigframes/dataframe.py:2419: KeyError

CC oncall @sycai for visibility

tswast

Very cool! A few comments.

tswast · 2026-06-16T19:23:11Z

+                "Python transpiler is an unstable, experimental feature, and not yet fully "
+                "validated, use at your own risk."
+            )
+            warnings.warn(msg, category=bfe.PreviewWarning)


Nit: I like to make custom exceptions the subclass from PreviewWarning for more explicit opt-in, but probably overkill in retrospect.

tswast · 2026-06-16T19:24:04Z

+
+    name: str
+    default_value: typing.Any
+    is_varargs: bool


Do we want to track keyword-only and/or kwargs dictionary separately? Or maybe that's not really inferrable from the Python AST?

Yeah, switched to track all 5 kinds of params (pos+kw, pos, kw, posvar, kwvar). Most of these aren't relevant quite yet, as only the axis=1 case in our API even takes anything besides positional, and it has its own code paths, but will allow us to generalize later.

tswast · 2026-06-16T19:27:43Z

+                )
+            )
+
+        from bigframes.core.bytecode import dis_to_expr


"dis" is hard for me to understand without context. Could we use a more descriptive name?

renamed to py_to_expression

tswast · 2026-06-16T19:29:26Z

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


These bind and transform functions are pretty similar to the other implementations, right? Maybe we can do some sort of mixin class to implement these?

IDK about specific impl (mixin or otherwise), but yeah, not that have expanded from 3 expression subclasses (op, literal, reference) to several more, things are getting redundant.

The logic is actually a bit custom per expression type, but it could maybe just be derived from dataclass field metadata, as most of it amounts to descending through fields (known statically) that subclass expression and returning/replacing them.

Created internal issue 525095335 to track a refactor here.

tswast · 2026-06-16T19:32:21Z

+                for col in self.columns:
+                    if col in expr.free_variables:
+                        col_id = block.resolve_label_exact(col)


Would we expect any trouble from mixing the high-level (dataframe columns) representation and lower level (block column labels) representation?

Also, this looks relatively familiar, such as in our bbq.sql_scalar implementation. Perhaps there's some shared utilities we can refactor for this column mapping logic?

The idea here is to cross that label to id boundary. The python function references series attributes which are labels, but we want to map those to unambiguous block ids.

The implementation here is a bit messy though, refactored it in new revision as apply_to_block_rows

tswast · 2026-06-16T19:33:42Z

@@ -4692,13 +4692,17 @@ def _prepare_export(
        return array_value, id_overrides

    def map(self, func, na_action: Optional[str] = None) -> DataFrame:


I'd love to see an example in the docstrings, especially if this is compatible with the polars engine and thus would make for a relatively speedy flakeless doctest.

Added doctests

tswast · 2026-06-18T20:13:05Z

some doctest failures look real:

_ [doctest] bigframes.extensions.core.series_accessor.BigQuerySeriesAccessor.array_length _
[gw4] linux -- Python 3.12.12 /tmpfs/src/github/google-cloud-python/packages/bigframes/.nox/doctest/bin/python
403             >>> ps = pd.Series([[1, 2, 8, 3], [], [3, 4]])
404             >>> ps.bigquery.array_length()
405             0    4
406             1    0
407             2    2
408             dtype: Int64
409 
410         You can also apply this function directly to Series using `apply`.
411 
412             >>> s.apply(bbq.array_length, by_row=False)
UNEXPECTED EXCEPTION: ValueError('Unsupported opcode: COPY_FREE_VARS')
Traceback (most recent call last):
  File "/usr/local/lib/python3.12/doctest.py", line 1368, in __run
    exec(compile(example.source, filename, "single",
  File "", line 1, in 
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/core/logging/log_adapter.py", line 183, in wrapper
    return method(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/series.py", line 2086, in apply
    result_series = self._apply_callable_expr(ops.func_to_expr(func), args)
                                              ^^^^^^^^^^^^^^^^^^^^^^
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/operations/to_op.py", line 194, in func_to_expr
    return CallableExpression.from_callable(op)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/operations/to_op.py", line 98, in from_callable
    expr = py_to_expression(func)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/core/bytecode.py", line 259, in py_to_expression
    py_expr = _compile_bytecode_to_py_expr(func)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/core/bytecode.py", line 247, in _compile_bytecode_to_py_expr
    raise ValueError(f"Unsupported opcode: {opname}")
ValueError: Unsupported opcode: COPY_FREE_VARS
/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/extensions/core/series_accessor.py:412: UnexpectedException
_ [doctest] bigframes.operations.googlesql.global_namespace.array.array_length _
[gw4] linux -- Python 3.12.12 /tmpfs/src/github/google-cloud-python/packages/bigframes/.nox/doctest/bin/python
703         >>> ps = pd.Series([[1, 2, 8, 3], [], [3, 4]])
704         >>> ps.bigquery.array_length()
705         0    4
706         1    0
707         2    2
708         dtype: Int64
709 
710     You can also apply this function directly to Series using `apply`.
711 
712         >>> s.apply(bbq.array_length, by_row=False)
UNEXPECTED EXCEPTION: ValueError('Unsupported opcode: COPY_FREE_VARS')
Traceback (most recent call last):
  File "/usr/local/lib/python3.12/doctest.py", line 1368, in __run
    exec(compile(example.source, filename, "single",
  File "", line 1, in 
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/core/logging/log_adapter.py", line 183, in wrapper
    return method(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/series.py", line 2086, in apply
    result_series = self._apply_callable_expr(ops.func_to_expr(func), args)
                                              ^^^^^^^^^^^^^^^^^^^^^^
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/operations/to_op.py", line 194, in func_to_expr
    return CallableExpression.from_callable(op)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/operations/to_op.py", line 98, in from_callable
    expr = py_to_expression(func)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/core/bytecode.py", line 259, in py_to_expression
    py_expr = _compile_bytecode_to_py_expr(func)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/core/bytecode.py", line 247, in _compile_bytecode_to_py_expr
    raise ValueError(f"Unsupported opcode: {opname}")
ValueError: Unsupported opcode: COPY_FREE_VARS
/tmpfs/src/github/google-cloud-python/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py:712: UnexpectedException

Possible solution: remove the "session" scope from https://github.com/googleapis/google-cloud-python/blob/main/packages/bigframes/conftest.py

gemini-code-assist Bot reviewed Jun 11, 2026

View reviewed changes

Comment thread packages/bigframes/bigframes/dataframe.py Outdated

Comment thread packages/bigframes/bigframes/operations/to_op.py Outdated

parthea assigned TrevorBergeron Jun 11, 2026

TrevorBergeron requested a review from tswast June 16, 2026 00:22

TrevorBergeron marked this pull request as ready for review June 16, 2026 00:22

TrevorBergeron requested review from a team as code owners June 16, 2026 00:22

tswast reviewed Jun 16, 2026

View reviewed changes

TrevorBergeron added 9 commits June 18, 2026 17:39

feat: Experimental transpilation of unannotated python callables

2dd2003

refactor, more tests

b11f211

ruff

cbe405b

fix test for python versions

0f42849

fixes

2488374

rename dis_to_expr

8f42da1

fix names

910027e

doctests

a61ae8b

refactor CallableExpression

c117c50

TrevorBergeron force-pushed the tbergeron_transpile_path branch from 3d5fd04 to c117c50 Compare June 18, 2026 17:39

fix bytecode tests

d2b980c

TrevorBergeron requested a review from tswast June 18, 2026 17:48

TrevorBergeron added 2 commits June 18, 2026 18:26

revert global fixture

7d52cd4

fix mypy

91f65dc

		@@ -4692,13 +4692,17 @@ def _prepare_export(
		return array_value, id_overrides

		def map(self, func, na_action: Optional[str] = None) -> DataFrame:

Conversation

TrevorBergeron commented Jun 11, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

tswast commented Jun 16, 2026

Uh oh!

tswast left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tswast commented Jun 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants