Merge remote-tracking branch 'origin/main' into refactor-isnull-op

Author: tswast

CHANGELOG.md

@@ -4,6 +4,31 @@
 
 [1]: https://pypi.org/project/bigframes/#history
 
+## [2.14.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v2.13.0...v2.14.0) (2025-08-05)
+
+
+### Features
+
+* Dynamic table width for better display across devices (https://github.com/googleapis/python-bigquery-dataframes/issues/1948) ([a6d30ae](https://github.com/googleapis/python-bigquery-dataframes/commit/a6d30ae3f4358925c999c53b558c1ecd3ee03e6c)) ([a6d30ae](https://github.com/googleapis/python-bigquery-dataframes/commit/a6d30ae3f4358925c999c53b558c1ecd3ee03e6c))
+* Retry AI/ML jobs that fail more often ([#1965](https://github.com/googleapis/python-bigquery-dataframes/issues/1965)) ([25bde9f](https://github.com/googleapis/python-bigquery-dataframes/commit/25bde9f9b89112db0efcc119bf29b6d1f3896c33))
+* Support series input in managed function ([#1920](https://github.com/googleapis/python-bigquery-dataframes/issues/1920)) ([62a189f](https://github.com/googleapis/python-bigquery-dataframes/commit/62a189f4d69f6c05fe348a1acd1fbac364fa60b9))
+
+
+### Bug Fixes
+
+* Enhance type error messages for bigframes functions ([#1958](https://github.com/googleapis/python-bigquery-dataframes/issues/1958)) ([770918e](https://github.com/googleapis/python-bigquery-dataframes/commit/770918e998bf1fde7a656e8f8a0ff0a8c68509f2))
+
+
+### Performance Improvements
+
+* Use promote_offsets for consistent row number generation for index.get_loc  ([#1957](https://github.com/googleapis/python-bigquery-dataframes/issues/1957)) ([c67a25a](https://github.com/googleapis/python-bigquery-dataframes/commit/c67a25a879ab2a35ca9053a81c9c85b5660206ae))
+
+
+### Documentation
+
+* Add code snippet for storing dataframes to a CSV file ([#1943](https://github.com/googleapis/python-bigquery-dataframes/issues/1943)) ([a511e09](https://github.com/googleapis/python-bigquery-dataframes/commit/a511e09e6924d2e8302af2eb4a602c6b9e5d2d72))
+* Add code snippet for storing dataframes to a CSV file ([#1953](https://github.com/googleapis/python-bigquery-dataframes/issues/1953)) ([a298a02](https://github.com/googleapis/python-bigquery-dataframes/commit/a298a02b451f03ca200fe0756b9a7b57e3d1bf0e))
+
 ## [2.13.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v2.12.0...v2.13.0) (2025-07-25)
 
 

bigframes/exceptions.py

@@ -103,6 +103,13 @@ class FunctionAxisOnePreviewWarning(PreviewWarning):
     """Remote Function and Managed UDF with axis=1 preview."""
 
 
+class FunctionPackageVersionWarning(PreviewWarning):
+    """
+    Managed UDF package versions for Numpy, Pandas, and Pyarrow may not
+    precisely match users' local environment or the exact versions specified.
+    """
+
+
 def format_message(message: str, fill: bool = True):
     """Formats a warning message with ANSI color codes for the warning color.
 

bigframes/functions/_function_client.py

@@ -19,7 +19,6 @@
 import logging
 import os
 import random
-import re
 import shutil
 import string
 import tempfile
@@ -247,7 +246,7 @@ def provision_bq_managed_function(
         # Augment user package requirements with any internal package
         # requirements.
         packages = _utils._get_updated_package_requirements(
-            packages, is_row_processor, capture_references
+            packages, is_row_processor, capture_references, ignore_package_version=True
         )
         if packages:
             managed_function_options["packages"] = packages
@@ -270,28 +269,6 @@ def provision_bq_managed_function(
         )
 
         udf_name = func.__name__
-        if capture_references:
-            # This code path ensures that if the udf body contains any
-            # references to variables and/or imports outside the body, they are
-            # captured as well.
-            import cloudpickle
-
-            pickled = cloudpickle.dumps(func)
-            udf_code = textwrap.dedent(
-                f"""
-                import cloudpickle
-                {udf_name} = cloudpickle.loads({pickled})
-            """
-            )
-        else:
-            # This code path ensures that if the udf body is self contained,
-            # i.e. there are no references to variables or imports outside the
-            # body.
-            udf_code = textwrap.dedent(inspect.getsource(func))
-            match = re.search(r"^def ", udf_code, flags=re.MULTILINE)
-            if match is None:
-                raise ValueError("The UDF is not defined correctly.")
-            udf_code = udf_code[match.start() :]
 
         with_connection_clause = (
             (
@@ -301,6 +278,13 @@ def provision_bq_managed_function(
             else ""
         )
 
+        # Generate the complete Python code block for the managed Python UDF,
+        # including the user's function, necessary imports, and the BigQuery
+        # handler wrapper.
+        python_code_block = bff_template.generate_managed_function_code(
+            func, udf_name, is_row_processor, capture_references
+        )
+
         create_function_ddl = (
             textwrap.dedent(
                 f"""
@@ -311,13 +295,11 @@ def provision_bq_managed_function(
                 OPTIONS ({managed_function_options_str})
                 AS r'''
                 __UDF_PLACE_HOLDER__
-                def bigframes_handler(*args):
-                    return {udf_name}(*args)
                 '''
             """
             )
             .strip()
-            .replace("__UDF_PLACE_HOLDER__", udf_code)
+            .replace("__UDF_PLACE_HOLDER__", python_code_block)
         )
 
         self._ensure_dataset_exists()

bigframes/functions/_function_session.py

@@ -847,15 +847,15 @@ def wrapper(func):
             if output_type:
                 py_sig = py_sig.replace(return_annotation=output_type)
 
-            udf_sig = udf_def.UdfSignature.from_py_signature(py_sig)
-
             # The function will actually be receiving a pandas Series, but allow
             # both BigQuery DataFrames and pandas object types for compatibility.
             is_row_processor = False
             if new_sig := _convert_row_processor_sig(py_sig):
                 py_sig = new_sig
                 is_row_processor = True
 
+            udf_sig = udf_def.UdfSignature.from_py_signature(py_sig)
+
             managed_function_client = _function_client.FunctionClient(
                 dataset_ref.project,
                 bq_location,

bigframes/functions/_utils.py

@@ -18,6 +18,7 @@
 import sys
 import typing
 from typing import cast, Optional, Set
+import warnings
 
 import cloudpickle
 import google.api_core.exceptions
@@ -26,6 +27,7 @@
 import pandas
 import pyarrow
 
+import bigframes.exceptions as bfe
 import bigframes.formatting_helpers as bf_formatting
 from bigframes.functions import function_typing
 
@@ -61,21 +63,40 @@ def get_remote_function_locations(bq_location):
 
 
 def _get_updated_package_requirements(
-    package_requirements=None, is_row_processor=False, capture_references=True
+    package_requirements=None,
+    is_row_processor=False,
+    capture_references=True,
+    ignore_package_version=False,
 ):
     requirements = []
     if capture_references:
         requirements.append(f"cloudpickle=={cloudpickle.__version__}")
 
     if is_row_processor:
-        # bigframes function will send an entire row of data as json, which
-        # would be converted to a pandas series and processed Ensure numpy
-        # versions match to avoid unpickling problems. See internal issue
-        # b/347934471.
-        requirements.append(f"numpy=={numpy.__version__}")
-        requirements.append(f"pandas=={pandas.__version__}")
-        requirements.append(f"pyarrow=={pyarrow.__version__}")
-
+        if ignore_package_version:
+            # TODO(jialuo): Add back the version after b/410924784 is resolved.
+            # Due to current limitations on the packages version in Python UDFs,
+            # we use `ignore_package_version` to optionally omit the version for
+            # managed functions only.
+            msg = bfe.format_message(
+                "numpy, pandas, and pyarrow versions in the function execution"
+                " environment may not precisely match your local environment."
+            )
+            warnings.warn(msg, category=bfe.FunctionPackageVersionWarning)
+            requirements.append("pandas")
+            requirements.append("pyarrow")
+            requirements.append("numpy")
+        else:
+            # bigframes function will send an entire row of data as json, which
+            # would be converted to a pandas series and processed Ensure numpy
+            # versions match to avoid unpickling problems. See internal issue
+            # b/347934471.
+            requirements.append(f"pandas=={pandas.__version__}")
+            requirements.append(f"pyarrow=={pyarrow.__version__}")
+            requirements.append(f"numpy=={numpy.__version__}")
+
+    # TODO(b/435023957): Fix the issue of potential duplicate package versions
+    # when `package_requirements` also contains `pandas/pyarrow/numpy`.
     if package_requirements:
         requirements.extend(package_requirements)
 

bigframes/functions/function_template.py

@@ -17,6 +17,7 @@
 import inspect
 import logging
 import os
+import re
 import textwrap
 from typing import Tuple
 
@@ -291,3 +292,55 @@ def generate_cloud_function_main_code(
     logger.debug(f"Wrote {os.path.abspath(main_py)}:\n{open(main_py).read()}")
 
     return handler_func_name
+
+
+def generate_managed_function_code(
+    def_,
+    udf_name: str,
+    is_row_processor: bool,
+    capture_references: bool,
+) -> str:
+    """Generates the Python code block for managed Python UDF."""
+
+    if capture_references:
+        # This code path ensures that if the udf body contains any
+        # references to variables and/or imports outside the body, they are
+        # captured as well.
+        import cloudpickle
+
+        pickled = cloudpickle.dumps(def_)
+        func_code = textwrap.dedent(
+            f"""
+            import cloudpickle
+            {udf_name} = cloudpickle.loads({pickled})
+        """
+        )
+    else:
+        # This code path ensures that if the udf body is self contained,
+        # i.e. there are no references to variables or imports outside the
+        # body.
+        func_code = textwrap.dedent(inspect.getsource(def_))
+        match = re.search(r"^def ", func_code, flags=re.MULTILINE)
+        if match is None:
+            raise ValueError("The UDF is not defined correctly.")
+        func_code = func_code[match.start() :]
+
+    if is_row_processor:
+        udf_code = textwrap.dedent(inspect.getsource(get_pd_series))
+        udf_code = udf_code[udf_code.index("def") :]
+        bigframes_handler_code = textwrap.dedent(
+            f"""def bigframes_handler(str_arg):
+                return {udf_name}({get_pd_series.__name__}(str_arg))"""
+        )
+    else:
+        udf_code = ""
+        bigframes_handler_code = textwrap.dedent(
+            f"""def bigframes_handler(*args):
+                return {udf_name}(*args)"""
+        )
+
+    udf_code_block = textwrap.dedent(
+        f"{udf_code}\n{func_code}\n{bigframes_handler_code}"
+    )
+
+    return udf_code_block

bigframes/session/__init__.py

@@ -40,6 +40,7 @@
 import weakref
 
 import bigframes_vendored.constants as constants
+import bigframes_vendored.google_cloud_bigquery.retry as third_party_gcb_retry
 import bigframes_vendored.ibis.backends.bigquery as ibis_bigquery  # noqa
 import bigframes_vendored.pandas.io.gbq as third_party_pandas_gbq
 import bigframes_vendored.pandas.io.parquet as third_party_pandas_parquet
@@ -2051,6 +2052,7 @@ def _start_query_ml_ddl(
             project=None,
             timeout=None,
             query_with_job=True,
+            job_retry=third_party_gcb_retry.DEFAULT_ML_JOB_RETRY,
         )
         return iterator, query_job
 

bigframes/session/_io/bigquery/__init__.py

@@ -24,8 +24,10 @@
 import typing
 from typing import Dict, Iterable, Literal, Mapping, Optional, overload, Tuple, Union
 
+import bigframes_vendored.google_cloud_bigquery.retry as third_party_gcb_retry
 import bigframes_vendored.pandas.io.gbq as third_party_pandas_gbq
 import google.api_core.exceptions
+import google.api_core.retry
 import google.cloud.bigquery as bigquery
 
 from bigframes.core import log_adapter
@@ -245,7 +247,7 @@ def start_query_with_client(
     location: Optional[str],
     project: Optional[str],
     timeout: Optional[float],
-    metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
+    metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
     query_with_job: Literal[True],
 ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]:
     ...
@@ -260,8 +262,40 @@ def start_query_with_client(
     location: Optional[str],
     project: Optional[str],
     timeout: Optional[float],
-    metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
+    metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
+    query_with_job: Literal[False],
+) -> Tuple[bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
+    ...
+
+
+@overload
+def start_query_with_client(
+    bq_client: bigquery.Client,
+    sql: str,
+    *,
+    job_config: bigquery.QueryJobConfig,
+    location: Optional[str],
+    project: Optional[str],
+    timeout: Optional[float],
+    metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
+    query_with_job: Literal[True],
+    job_retry: google.api_core.retry.Retry,
+) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]:
+    ...
+
+
+@overload
+def start_query_with_client(
+    bq_client: bigquery.Client,
+    sql: str,
+    *,
+    job_config: bigquery.QueryJobConfig,
+    location: Optional[str],
+    project: Optional[str],
+    timeout: Optional[float],
+    metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
     query_with_job: Literal[False],
+    job_retry: google.api_core.retry.Retry,
 ) -> Tuple[bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
     ...
 
@@ -276,6 +310,11 @@ def start_query_with_client(
     timeout: Optional[float] = None,
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
     query_with_job: bool = True,
+    # TODO(tswast): We can stop providing our own default once we use a
+    # google-cloud-bigquery version with
+    # https://github.com/googleapis/python-bigquery/pull/2256 merged, likely
+    # version 3.36.0 or later.
+    job_retry: google.api_core.retry.Retry = third_party_gcb_retry.DEFAULT_JOB_RETRY,
 ) -> Tuple[bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
     """
     Starts query job and waits for results.
@@ -292,6 +331,7 @@ def start_query_with_client(
                 location=location,
                 project=project,
                 api_timeout=timeout,
+                job_retry=job_retry,
             )
             if metrics is not None:
                 metrics.count_job_stats(row_iterator=results_iterator)
@@ -303,6 +343,7 @@ def start_query_with_client(
             location=location,
             project=project,
             timeout=timeout,
+            job_retry=job_retry,
         )
     except google.api_core.exceptions.Forbidden as ex:
         if "Drive credentials" in ex.message:

bigframes/version.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "2.13.0"
+__version__ = "2.14.0"
 
 # {x-release-please-start-date}
-__release_date__ = "2025-07-25"
+__release_date__ = "2025-08-05"
 # {x-release-please-end}