chore: rewrite function axis=1 warning in blob transform functions (#1538)

Author: GarrettWu

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(

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.
 

bigframes/operations/blob.py

@@ -16,13 +16,15 @@
 
 import os
 from typing import cast, Optional, Union
+import warnings
 
 import IPython.display as ipy_display
 import pandas as pd
 import requests
 
 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.
 
@@ -346,7 +372,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
 
@@ -375,7 +401,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
 
         self._add_to_cleanup_set(image_blur_udf)
@@ -443,7 +469,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
 
@@ -473,7 +499,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
 
         self._add_to_cleanup_set(image_resize_udf)
@@ -535,7 +561,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
 
@@ -566,7 +592,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
 
         self._add_to_cleanup_set(image_normalize_udf)
@@ -681,7 +707,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)