chore: add experimental blob tansform func tuning params (#1422)

Author: GarrettWu

bigframes/blob/_functions.py

@@ -14,7 +14,7 @@
 
 from dataclasses import dataclass
 import inspect
-from typing import Callable, Iterable
+from typing import Callable, Iterable, Union
 
 import google.cloud.bigquery as bigquery
 
@@ -42,12 +42,18 @@ def __init__(
         session: bigframes.session.Session,
         connection: str,
         max_batching_rows: int,
+        container_cpu: Union[float, int],
+        container_memory: str,
     ):
         self._func = func_def.func
         self._requirements = func_def.requirements
         self._session = session
         self._connection = connection
-        self._max_batching_rows = max_batching_rows
+        self._max_batching_rows = (
+            int(max_batching_rows) if max_batching_rows > 1 else max_batching_rows
+        )
+        self._container_cpu = container_cpu
+        self._container_memory = container_memory
 
     def _input_bq_signature(self):
         sig = inspect.signature(self._func)
@@ -72,7 +78,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}, max_batching_rows={self._max_batching_rows})
+OPTIONS (entry_point='{func_name}', runtime_version='python-3.11', packages={packages}, max_batching_rows={self._max_batching_rows}, container_cpu={self._container_cpu}, container_memory='{self._container_memory}')
 AS r\"\"\"
 
 

bigframes/operations/blob.py

@@ -278,7 +278,9 @@ def image_blur(
         *,
         dst: Optional[Union[str, bigframes.series.Series]] = None,
         connection: Optional[str] = None,
-        max_batching_rows: int = 10000,
+        max_batching_rows: int = 8096,
+        container_cpu: Union[float, int] = 0.33,
+        container_memory: str = "512Mi",
     ) -> bigframes.series.Series:
         """Blurs images.
 
@@ -289,7 +291,9 @@ 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.
+            max_batching_rows (int, default 8,096): Max number of rows per batch send to cloud run to execute the function.
+            container_cpu (int or float, default 0.33): number of container CPUs. Possible values are [0.33, 8]. Floats larger than 1 are cast to intergers.
+            container_memory (str, default "512Mi"): container memory size. String of the format <number><unit>. Possible values are from 512Mi to 32Gi.
 
         Returns:
             BigFrames Blob Series
@@ -305,6 +309,8 @@ def image_blur(
                 session=self._block.session,
                 connection=connection,
                 max_batching_rows=max_batching_rows,
+                container_cpu=container_cpu,
+                container_memory=container_memory,
             ).udf()
 
             df["ksize_x"], df["ksize_y"] = ksize
@@ -326,6 +332,8 @@ def image_blur(
             session=self._block.session,
             connection=connection,
             max_batching_rows=max_batching_rows,
+            container_cpu=container_cpu,
+            container_memory=container_memory,
         ).udf()
 
         dst_rt = dst.blob._get_runtime_json_str(mode="RW")
@@ -346,7 +354,9 @@ def image_resize(
         fy: float = 0.0,
         dst: Optional[Union[str, bigframes.series.Series]] = None,
         connection: Optional[str] = None,
-        max_batching_rows: int = 10000,
+        max_batching_rows: int = 8096,
+        container_cpu: Union[float, int] = 0.33,
+        container_memory: str = "512Mi",
     ):
         """Resize images.
 
@@ -359,7 +369,9 @@ 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.
+            max_batching_rows (int, default 8,096): Max number of rows per batch send to cloud run to execute the function.
+            container_cpu (int or float, default 0.33): number of container CPUs. Possible values are [0.33, 8]. Floats larger than 1 are cast to intergers.
+            container_memory (str, default "512Mi"): container memory size. String of the format <number><unit>. Possible values are from 512Mi to 32Gi.
 
         Returns:
             BigFrames Blob Series
@@ -382,6 +394,8 @@ def image_resize(
                 session=self._block.session,
                 connection=connection,
                 max_batching_rows=max_batching_rows,
+                container_cpu=container_cpu,
+                container_memory=container_memory,
             ).udf()
 
             df["dsize_x"], df["dsizye_y"] = dsize
@@ -404,6 +418,8 @@ def image_resize(
             session=self._block.session,
             connection=connection,
             max_batching_rows=max_batching_rows,
+            container_cpu=container_cpu,
+            container_memory=container_memory,
         ).udf()
 
         dst_rt = dst.blob._get_runtime_json_str(mode="RW")
@@ -425,7 +441,9 @@ def image_normalize(
         norm_type: str = "l2",
         dst: Optional[Union[str, bigframes.series.Series]] = None,
         connection: Optional[str] = None,
-        max_batching_rows: int = 10000,
+        max_batching_rows: int = 8096,
+        container_cpu: Union[float, int] = 0.33,
+        container_memory: str = "512Mi",
     ) -> bigframes.series.Series:
         """Normalize images.
 
@@ -438,7 +456,9 @@ 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.
+            max_batching_rows (int, default 8,096): Max number of rows per batch send to cloud run to execute the function.
+            container_cpu (int or float, default 0.33): number of container CPUs. Possible values are [0.33, 8]. Floats larger than 1 are cast to intergers.
+            container_memory (str, default "512Mi"): container memory size. String of the format <number><unit>. Possible values are from 512Mi to 32Gi.
 
         Returns:
             BigFrames Blob Series
@@ -454,6 +474,8 @@ def image_normalize(
                 session=self._block.session,
                 connection=connection,
                 max_batching_rows=max_batching_rows,
+                container_cpu=container_cpu,
+                container_memory=container_memory,
             ).udf()
 
             df["alpha"] = alpha
@@ -477,6 +499,8 @@ def image_normalize(
             session=self._block.session,
             connection=connection,
             max_batching_rows=max_batching_rows,
+            container_cpu=container_cpu,
+            container_memory=container_memory,
         ).udf()
 
         dst_rt = dst.blob._get_runtime_json_str(mode="RW")
@@ -495,7 +519,9 @@ def pdf_extract(
         self,
         *,
         connection: Optional[str] = None,
-        max_batching_rows: int = 10000,
+        max_batching_rows: int = 8096,
+        container_cpu: Union[float, int] = 0.33,
+        container_memory: str = "512Mi",
     ) -> bigframes.series.Series:
         """Extracts and chunks text from PDF URLs and saves the text as
            arrays of string.
@@ -508,8 +534,10 @@ 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
+            max_batching_rows (int, default 8,096): Max number of rows per batch
                 send to cloud run to execute the function.
+            container_cpu (int or float, default 0.33): number of container CPUs. Possible values are [0.33, 8]. Floats larger than 1 are cast to intergers.
+            container_memory (str, default "512Mi"): container memory size. String of the format <number><unit>. Possible values are from 512Mi to 32Gi.
 
         Returns:
             bigframes.series.Series: conatins all text from a pdf file
@@ -524,6 +552,8 @@ def pdf_extract(
             session=self._block.session,
             connection=connection,
             max_batching_rows=max_batching_rows,
+            container_cpu=container_cpu,
+            container_memory=container_memory,
         ).udf()
 
         src_rt = self._get_runtime_json_str(mode="R")
@@ -536,7 +566,9 @@ def pdf_chunk(
         connection: Optional[str] = None,
         chunk_size: int = 1000,
         overlap_size: int = 200,
-        max_batching_rows: int = 10000,
+        max_batching_rows: int = 8096,
+        container_cpu: Union[float, int] = 0.33,
+        container_memory: str = "512Mi",
     ) -> bigframes.series.Series:
         """Extracts and chunks text from PDF URLs and saves the text as
            arrays of strings.
@@ -554,8 +586,10 @@ 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
+            max_batching_rows (int, default 8,096): Max number of rows per batch
                 send to cloud run to execute the function.
+            container_cpu (int or float, default 0.33): number of container CPUs. Possible values are [0.33, 8]. Floats larger than 1 are cast to intergers.
+            container_memory (str, default "512Mi"): container memory size. String of the format <number><unit>. Possible values are from 512Mi to 32Gi.
 
         Returns:
             bigframe.series.Series of array[str], where each string is a
@@ -579,6 +613,8 @@ def pdf_chunk(
             session=self._block.session,
             connection=connection,
             max_batching_rows=max_batching_rows,
+            container_cpu=container_cpu,
+            container_memory=container_memory,
         ).udf()
 
         src_rt = self._get_runtime_json_str(mode="R")