diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index b0648d65ba..cc74b72716 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -4170,8 +4170,10 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs): # to the applied function should be a Series, not a scalar. if utils.get_axis_number(axis) == 1: - msg = bfe.format_message("axis=1 scenario is in preview.") - warnings.warn(msg, category=bfe.PreviewWarning) + msg = bfe.format_message( + "DataFrame.apply with parameter axis=1 scenario is in preview." + ) + warnings.warn(msg, category=bfe.FunctionAxisOnePreviewWarning) if not hasattr(func, "bigframes_bigquery_function"): raise ValueError( diff --git a/bigframes/exceptions.py b/bigframes/exceptions.py index 6197481253..8924295c29 100644 --- a/bigframes/exceptions.py +++ b/bigframes/exceptions.py @@ -95,6 +95,10 @@ class ObsoleteVersionWarning(Warning): """The BigFrames version is too old.""" +class FunctionAxisOnePreviewWarning(PreviewWarning): + """Remote Function and Managed UDF with axis=1 preview.""" + + def format_message(message: str, fill: bool = True): """Formats a warning message with ANSI color codes for the warning color. diff --git a/bigframes/operations/blob.py b/bigframes/operations/blob.py index ab529c1312..369e472ce2 100644 --- a/bigframes/operations/blob.py +++ b/bigframes/operations/blob.py @@ -16,6 +16,7 @@ import os from typing import cast, Optional, Union +import warnings import IPython.display as ipy_display import pandas as pd @@ -23,6 +24,7 @@ from bigframes import clients import bigframes.dataframe +import bigframes.exceptions as bfe from bigframes.operations import base import bigframes.operations as ops import bigframes.series @@ -166,6 +168,30 @@ def _get_runtime( return s._apply_unary_op(ops.ObjGetAccessUrl(mode=mode)) + def _df_apply_udf( + self, df: bigframes.dataframe.DataFrame, udf + ) -> bigframes.series.Series: + # Catch and rethrow function axis=1 warning to be more user-friendly. + with warnings.catch_warnings(record=True) as catched_warnings: + s = df.apply(udf, axis=1) + for w in catched_warnings: + if isinstance(w.message, bfe.FunctionAxisOnePreviewWarning): + warnings.warn( + "Blob Functions use bigframes DataFrame Managed function with axis=1 senario, which is a preview feature.", + category=w.category, + stacklevel=2, + ) + else: + warnings.warn_explicit( + message=w.message, + category=w.category, + filename=w.filename, + lineno=w.lineno, + source=w.source, + ) + + return s + def read_url(self) -> bigframes.series.Series: """Retrieve the read URL of the Blob. @@ -335,7 +361,7 @@ def image_blur( df["ksize_x"], df["ksize_y"] = ksize df["ext"] = ext # type: ignore - res = df.apply(image_blur_udf, axis=1) + res = self._df_apply_udf(df, image_blur_udf) return res @@ -364,7 +390,7 @@ def image_blur( df["ksize_x"], df["ksize_y"] = ksize df["ext"] = ext # type: ignore - res = df.apply(image_blur_udf, axis=1) + res = self._df_apply_udf(df, image_blur_udf) res.cache() # to execute the udf return dst @@ -430,7 +456,7 @@ def image_resize( df["dsize_x"], df["dsizye_y"] = dsize df["fx"], df["fy"] = fx, fy df["ext"] = ext # type: ignore - res = df.apply(image_resize_udf, axis=1) + res = self._df_apply_udf(df, image_resize_udf) return res @@ -460,7 +486,7 @@ def image_resize( df["fx"], df["fy"] = fx, fy df["ext"] = ext # type: ignore - res = df.apply(image_resize_udf, axis=1) + res = self._df_apply_udf(df, image_resize_udf) res.cache() # to execute the udf return dst @@ -520,7 +546,7 @@ def image_normalize( df["beta"] = beta df["norm_type"] = norm_type df["ext"] = ext # type: ignore - res = df.apply(image_normalize_udf, axis=1) + res = self._df_apply_udf(df, image_normalize_udf) return res @@ -551,7 +577,7 @@ def image_normalize( df["norm_type"] = norm_type df["ext"] = ext # type: ignore - res = df.apply(image_normalize_udf, axis=1) + res = self._df_apply_udf(df, image_normalize_udf) res.cache() # to execute the udf return dst @@ -661,7 +687,7 @@ def pdf_chunk( df["chunk_size"] = chunk_size df["overlap_size"] = overlap_size - res = df.apply(pdf_chunk_udf, axis=1) + res = self._df_apply_udf(df, pdf_chunk_udf) res_array = bbq.json_extract_string_array(res) return res_array