chore: add max_batching_rows to experimental blob transform functions (#1407)

Author: GarrettWu

bigframes/blob/_functions.py

@@ -37,12 +37,17 @@ class TransformFunction:
     """Simple transform function class to deal with Python UDF."""
 
     def __init__(
-        self, func_def: FunctionDef, session: bigframes.session.Session, connection: str
+        self,
+        func_def: FunctionDef,
+        session: bigframes.session.Session,
+        connection: str,
+        max_batching_rows: int,
     ):
         self._func = func_def.func
         self._requirements = func_def.requirements
         self._session = session
         self._connection = connection
+        self._max_batching_rows = max_batching_rows
 
     def _input_bq_signature(self):
         sig = inspect.signature(self._func)
@@ -67,7 +72,7 @@ def _create_udf(self):
 CREATE OR REPLACE FUNCTION `{udf_name}`({self._input_bq_signature()})
 RETURNS {self._output_bq_type()} LANGUAGE python
 WITH CONNECTION `{self._connection}`
-OPTIONS (entry_point='{func_name}', runtime_version='python-3.11', packages={packages})
+OPTIONS (entry_point='{func_name}', runtime_version='python-3.11', packages={packages}, max_batching_rows={self._max_batching_rows})
 AS r\"\"\"
 
 

bigframes/operations/blob.py

@@ -278,6 +278,7 @@ def image_blur(
         *,
         dst: Optional[Union[str, bigframes.series.Series]] = None,
         connection: Optional[str] = None,
+        max_batching_rows: int = 10000,
     ) -> bigframes.series.Series:
         """Blurs images.
 
@@ -288,6 +289,7 @@ def image_blur(
             ksize (tuple(int, int)): Kernel size.
             dst (str or bigframes.series.Series or None, default None): Destination GCS folder str or blob series. If None, output to BQ as bytes.
             connection (str or None, default None): BQ connection used for function internet transactions, and the output blob if "dst" is str. If None, uses default connection of the session.
+            max_batching_rows (int, default 10,000): Max number of rows per batch send to cloud run to execute the function.
 
         Returns:
             BigFrames Blob Series
@@ -302,6 +304,7 @@ def image_blur(
                 blob_func.image_blur_to_bytes_def,
                 session=self._block.session,
                 connection=connection,
+                max_batching_rows=max_batching_rows,
             ).udf()
 
             df["ksize_x"], df["ksize_y"] = ksize
@@ -322,6 +325,7 @@ def image_blur(
             blob_func.image_blur_def,
             session=self._block.session,
             connection=connection,
+            max_batching_rows=max_batching_rows,
         ).udf()
 
         dst_rt = dst.blob._get_runtime_json_str(mode="RW")
@@ -342,6 +346,7 @@ def image_resize(
         fy: float = 0.0,
         dst: Optional[Union[str, bigframes.series.Series]] = None,
         connection: Optional[str] = None,
+        max_batching_rows: int = 10000,
     ):
         """Resize images.
 
@@ -354,6 +359,7 @@ def image_resize(
             fy (float, defalut 0.0): scale factor along the vertical axis. If set to 0.0, dsize parameter determines the output size.
             dst (str or bigframes.series.Series or None, default None): Destination GCS folder str or blob series. If None, output to BQ as bytes.
             connection (str or None, default None): BQ connection used for function internet transactions, and the output blob if "dst" is str. If None, uses default connection of the session.
+            max_batching_rows (int, default 10,000): Max number of rows per batch send to cloud run to execute the function.
 
         Returns:
             BigFrames Blob Series
@@ -375,6 +381,7 @@ def image_resize(
                 blob_func.image_resize_to_bytes_def,
                 session=self._block.session,
                 connection=connection,
+                max_batching_rows=max_batching_rows,
             ).udf()
 
             df["dsize_x"], df["dsizye_y"] = dsize
@@ -396,6 +403,7 @@ def image_resize(
             blob_func.image_resize_def,
             session=self._block.session,
             connection=connection,
+            max_batching_rows=max_batching_rows,
         ).udf()
 
         dst_rt = dst.blob._get_runtime_json_str(mode="RW")
@@ -417,6 +425,7 @@ def image_normalize(
         norm_type: str = "l2",
         dst: Optional[Union[str, bigframes.series.Series]] = None,
         connection: Optional[str] = None,
+        max_batching_rows: int = 10000,
     ) -> bigframes.series.Series:
         """Normalize images.
 
@@ -429,6 +438,7 @@ def image_normalize(
             norm_type (str, default "l2"): Normalization type. Accepted values are "inf", "l1", "l2" and "minmax".
             dst (str or bigframes.series.Series or None, default None): Destination GCS folder str or blob series. If None, output to BQ as bytes.
             connection (str or None, default None): BQ connection used for function internet transactions, and the output blob if "dst" is str. If None, uses default connection of the session.
+            max_batching_rows (int, default 10,000): Max number of rows per batch send to cloud run to execute the function.
 
         Returns:
             BigFrames Blob Series
@@ -443,6 +453,7 @@ def image_normalize(
                 blob_func.image_normalize_to_bytes_def,
                 session=self._block.session,
                 connection=connection,
+                max_batching_rows=max_batching_rows,
             ).udf()
 
             df["alpha"] = alpha
@@ -465,6 +476,7 @@ def image_normalize(
             blob_func.image_normalize_def,
             session=self._block.session,
             connection=connection,
+            max_batching_rows=max_batching_rows,
         ).udf()
 
         dst_rt = dst.blob._get_runtime_json_str(mode="RW")
@@ -480,7 +492,10 @@ def image_normalize(
         return dst
 
     def pdf_extract(
-        self, *, connection: Optional[str] = None
+        self,
+        *,
+        connection: Optional[str] = None,
+        max_batching_rows: int = 10000,
     ) -> bigframes.series.Series:
         """Extracts and chunks text from PDF URLs and saves the text as
            arrays of string.
@@ -493,6 +508,8 @@ def pdf_extract(
             connection (str or None, default None): BQ connection used for
                 function internet transactions, and the output blob if "dst"
                 is str. If None, uses default connection of the session.
+            max_batching_rows (int, default 10,000): Max number of rows per batch
+                send to cloud run to execute the function.
 
         Returns:
             bigframes.series.Series: conatins all text from a pdf file
@@ -502,14 +519,15 @@ def pdf_extract(
 
         connection = self._resolve_connection(connection)
 
-        pdf_chunk_udf = blob_func.TransformFunction(
+        pdf_extract_udf = blob_func.TransformFunction(
             blob_func.pdf_extract_def,
             session=self._block.session,
             connection=connection,
+            max_batching_rows=max_batching_rows,
         ).udf()
 
         src_rt = self._get_runtime_json_str(mode="R")
-        res = src_rt.apply(pdf_chunk_udf)
+        res = src_rt.apply(pdf_extract_udf)
         return res
 
     def pdf_chunk(
@@ -518,6 +536,7 @@ def pdf_chunk(
         connection: Optional[str] = None,
         chunk_size: int = 1000,
         overlap_size: int = 200,
+        max_batching_rows: int = 10000,
     ) -> bigframes.series.Series:
         """Extracts and chunks text from PDF URLs and saves the text as
            arrays of strings.
@@ -535,6 +554,8 @@ def pdf_chunk(
             overlap_size (int, default 200): the number of overlapping characters
                 between consective chunks. The helps to ensure context is
                 perserved across chunk boundaries.
+            max_batching_rows (int, default 10,000): Max number of rows per batch
+                send to cloud run to execute the function.
 
         Returns:
             bigframe.series.Series of array[str], where each string is a
@@ -557,6 +578,7 @@ def pdf_chunk(
             blob_func.pdf_chunk_def,
             session=self._block.session,
             connection=connection,
+            max_batching_rows=max_batching_rows,
         ).udf()
 
         src_rt = self._get_runtime_json_str(mode="R")