chore: add experimental blob.image_blur function (#1256)

* chore: add experimental blob.image_blur function

* apply to obj_ref

* docs

* fix mypy

Author: GarrettWu

bigframes/blob/_functions.py

@@ -0,0 +1,130 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+import inspect
+from typing import Callable, Iterable
+
+import google.cloud.bigquery as bigquery
+
+import bigframes
+import bigframes.session._io.bigquery as bf_io_bigquery
+
+_PYTHON_TO_BQ_TYPES = {int: "INT64", float: "FLOAT64", str: "STRING", bytes: "BYTES"}
+
+
+@dataclass(frozen=True)
+class FunctionDef:
+    """Definition of a Python UDF."""
+
+    func: Callable  # function body
+    requirements: Iterable[str]  # required packages
+
+
+# TODO(garrettwu): migrate to bigframes UDF when it is available
+class TransformFunction:
+    """Simple transform function class to deal with Python UDF."""
+
+    def __init__(
+        self, func_def: FunctionDef, session: bigframes.Session, connection: str
+    ):
+        self._func = func_def.func
+        self._requirements = func_def.requirements
+        self._session = session
+        self._connection = connection
+
+    def _input_bq_signature(self):
+        sig = inspect.signature(self._func)
+        inputs = []
+        for k, v in sig.parameters.items():
+            inputs.append(f"{k} {_PYTHON_TO_BQ_TYPES[v.annotation]}")
+        return ", ".join(inputs)
+
+    def _output_bq_type(self):
+        sig = inspect.signature(self._func)
+        return _PYTHON_TO_BQ_TYPES[sig.return_annotation]
+
+    def _create_udf(self):
+        """Create Python UDF in BQ. Return name of the UDF."""
+        udf_name = str(self._session._loader._storage_manager._random_table())
+
+        func_body = inspect.getsource(self._func)
+        func_name = self._func.__name__
+        packages = str(list(self._requirements))
+
+        sql = f"""
+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})
+AS r\"\"\"
+
+
+{func_body}
+
+
+\"\"\"
+        """
+
+        bf_io_bigquery.start_query_with_client(
+            self._session.bqclient,
+            sql,
+            job_config=bigquery.QueryJobConfig(),
+            metrics=self._session._metrics,
+        )
+
+        return udf_name
+
+    def udf(self):
+        """Create and return the UDF object."""
+        udf_name = self._create_udf()
+        return self._session.read_gbq_function(udf_name)
+
+
+# Blur images. Takes ObjectRefRuntime as JSON string. Outputs ObjectRefRuntime JSON string.
+def image_blur_func(
+    src_obj_ref_rt: str, dst_obj_ref_rt: str, ksize_x: int, ksize_y: int
+) -> str:
+    import json
+
+    import cv2 as cv  # type: ignore
+    import numpy as np
+    import requests
+
+    src_obj_ref_rt_json = json.loads(src_obj_ref_rt)
+    dst_obj_ref_rt_json = json.loads(dst_obj_ref_rt)
+
+    src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
+    dst_url = dst_obj_ref_rt_json["access_urls"]["write_url"]
+
+    response = requests.get(src_url)
+    bts = response.content
+
+    nparr = np.frombuffer(bts, np.uint8)
+    img = cv.imdecode(nparr, cv.IMREAD_UNCHANGED)
+    img_blurred = cv.blur(img, ksize=(ksize_x, ksize_y))
+    bts = cv.imencode(".jpeg", img_blurred)[1].tobytes()
+
+    requests.put(
+        url=dst_url,
+        data=bts,
+        headers={
+            "Content-Type": "image/jpeg",
+        },
+    )
+
+    return dst_obj_ref_rt
+
+
+image_blur_def = FunctionDef(image_blur_func, ["opencv-python", "numpy", "requests"])

bigframes/core/compile/scalar_op_compiler.py

@@ -1210,6 +1210,11 @@ def json_extract_string_array_op_impl(
     return json_extract_string_array(json_obj=x, json_path=op.json_path)
 
 
+@scalar_op_compiler.register_unary_op(ops.ToJSONString)
+def to_json_string_op_impl(json_obj: ibis_types.Value):
+    return to_json_string(json_obj=json_obj)
+
+
 # Blob Ops
 @scalar_op_compiler.register_unary_op(ops.obj_fetch_metadata_op)
 def obj_fetch_metadata_op_impl(obj_ref: ibis_types.Value):
@@ -1909,6 +1914,13 @@ def json_extract_string_array(  # type: ignore[empty-body]
     """Extracts a JSON array and converts it to a SQL ARRAY of STRINGs."""
 
 
+@ibis_udf.scalar.builtin(name="to_json_string")
+def to_json_string(  # type: ignore[empty-body]
+    json_obj: ibis_dtypes.JSON,
+) -> ibis_dtypes.String:
+    """Convert JSON to STRING."""
+
+
 @ibis_udf.scalar.builtin(name="ML.DISTANCE")
 def vector_distance(vector1, vector2, type: str) -> ibis_dtypes.Float64:  # type: ignore[empty-body]
     """Computes the distance between two vectors using specified type ("EUCLIDEAN", "MANHATTAN", or "COSINE")"""

bigframes/operations/__init__.py

@@ -740,6 +740,23 @@ def output_type(self, *input_types):
         )
 
 
+@dataclasses.dataclass(frozen=True)
+class ToJSONString(UnaryOp):
+    name: typing.ClassVar[str] = "to_json_string"
+
+    def output_type(self, *input_types):
+        input_type = input_types[0]
+        if not dtypes.is_json_like(input_type):
+            raise TypeError(
+                "Input type must be an valid JSON object or JSON-formatted string type."
+                + f" Received type: {input_type}"
+            )
+        return dtypes.STRING_DTYPE
+
+
+to_json_string_op = ToJSONString()
+
+
 ## Blob Ops
 @dataclasses.dataclass(frozen=True)
 class ObjGetAccessUrl(UnaryOp):

bigframes/operations/blob.py

@@ -14,9 +14,12 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import IPython.display as ipy_display
 import requests
 
+from bigframes import clients
 from bigframes.operations import base
 import bigframes.operations as ops
 import bigframes.series
@@ -66,3 +69,51 @@ def display(self, n: int = 3):
             read_url = str(read_url).strip('"')
             response = requests.get(read_url)
             ipy_display.display(ipy_display.Image(response.content))
+
+    def image_blur(
+        self,
+        ksize: tuple[int, int],
+        *,
+        dst: bigframes.series.Series,
+        connection: Optional[str] = None,
+    ) -> bigframes.series.Series:
+        """Blurs images.
+
+        .. note::
+            BigFrames Blob is still under experiments. It may not work and subject to change in the future.
+
+        Args:
+            ksize (tuple(int, int)): Kernel size.
+            dst (bigframes.series.Series): Destination blob series.
+            connection (str or None, default None): BQ connection used for internet transactions. If None, uses default connection of the session.
+
+        Returns:
+            JSON: Runtime info of the Blob.
+        """
+        import bigframes.blob._functions as blob_func
+
+        connection = connection or self._block.session._bq_connection
+        connection = clients.resolve_full_bq_connection_name(
+            connection,
+            default_project=self._block.session._project,
+            default_location=self._block.session._location,
+        )
+
+        image_blur_udf = blob_func.TransformFunction(
+            blob_func.image_blur_def,
+            session=self._block.session,
+            connection=connection,
+        ).udf()
+
+        src_rt = bigframes.series.Series(self._block)._apply_unary_op(
+            ops.ObjGetAccessUrl(mode="R")
+        )
+        dst_rt = dst._apply_unary_op(ops.ObjGetAccessUrl(mode="RW"))
+
+        src_rt = src_rt._apply_unary_op(ops.to_json_string_op)
+        dst_rt = dst_rt._apply_unary_op(ops.to_json_string_op)
+
+        df = src_rt.to_frame().join(dst_rt.to_frame(), how="outer")
+        df["ksize_x"], df["ksize_y"] = ksize
+
+        return df.apply(image_blur_udf, axis=1)