chore: impl experimental blob image output encodings (#1439)

Author: GarrettWu

bigframes/blob/_functions.py

@@ -105,14 +105,16 @@ def udf(self):
 
 # 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
+    src_obj_ref_rt: str, dst_obj_ref_rt: str, ksize_x: int, ksize_y: int, ext: str
 ) -> str:
     import json
 
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
 
+    ext = ext or ".jpeg"
+
     src_obj_ref_rt_json = json.loads(src_obj_ref_rt)
     dst_obj_ref_rt_json = json.loads(dst_obj_ref_rt)
 
@@ -125,13 +127,19 @@ def image_blur_func(
     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()
+
+    bts = cv.imencode(ext, img_blurred)[1].tobytes()
+
+    ext = ext.replace(".", "")
+    ext_mappings = {"jpg": "jpeg", "tif": "tiff"}
+    ext = ext_mappings.get(ext, ext)
+    content_type = "image/" + ext
 
     requests.put(
         url=dst_url,
         data=bts,
         headers={
-            "Content-Type": "image/jpeg",
+            "Content-Type": content_type,
         },
     )
 
@@ -141,13 +149,17 @@ def image_blur_func(
 image_blur_def = FunctionDef(image_blur_func, ["opencv-python", "numpy", "requests"])
 
 
-def image_blur_to_bytes_func(src_obj_ref_rt: str, ksize_x: int, ksize_y: int) -> bytes:
+def image_blur_to_bytes_func(
+    src_obj_ref_rt: str, ksize_x: int, ksize_y: int, ext: str
+) -> bytes:
     import json
 
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
 
+    ext = ext or ".jpeg"
+
     src_obj_ref_rt_json = json.loads(src_obj_ref_rt)
     src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
 
@@ -157,7 +169,7 @@ def image_blur_to_bytes_func(src_obj_ref_rt: str, ksize_x: int, ksize_y: int) ->
     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()
+    bts = cv.imencode(ext, img_blurred)[1].tobytes()
 
     return bts
 
@@ -174,13 +186,16 @@ def image_resize_func(
     dsize_y: int,
     fx: float,
     fy: float,
+    ext: str,
 ) -> str:
     import json
 
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
 
+    ext = ext or ".jpeg"
+
     src_obj_ref_rt_json = json.loads(src_obj_ref_rt)
     dst_obj_ref_rt_json = json.loads(dst_obj_ref_rt)
 
@@ -193,13 +208,19 @@ def image_resize_func(
     nparr = np.frombuffer(bts, np.uint8)
     img = cv.imdecode(nparr, cv.IMREAD_UNCHANGED)
     img_resized = cv.resize(img, dsize=(dsize_x, dsize_y), fx=fx, fy=fy)
-    bts = cv.imencode(".jpeg", img_resized)[1].tobytes()
+
+    bts = cv.imencode(ext, img_resized)[1].tobytes()
+
+    ext = ext.replace(".", "")
+    ext_mappings = {"jpg": "jpeg", "tif": "tiff"}
+    ext = ext_mappings.get(ext, ext)
+    content_type = "image/" + ext
 
     requests.put(
         url=dst_url,
         data=bts,
         headers={
-            "Content-Type": "image/jpeg",
+            "Content-Type": content_type,
         },
     )
 
@@ -217,13 +238,16 @@ def image_resize_to_bytes_func(
     dsize_y: int,
     fx: float,
     fy: float,
+    ext: str,
 ) -> bytes:
     import json
 
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
 
+    ext = ext or ".jpeg"
+
     src_obj_ref_rt_json = json.loads(src_obj_ref_rt)
     src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
 
@@ -244,14 +268,21 @@ def image_resize_to_bytes_func(
 
 
 def image_normalize_func(
-    src_obj_ref_rt: str, dst_obj_ref_rt: str, alpha: float, beta: float, norm_type: str
+    src_obj_ref_rt: str,
+    dst_obj_ref_rt: str,
+    alpha: float,
+    beta: float,
+    norm_type: str,
+    ext: str,
 ) -> str:
     import json
 
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
 
+    ext = ext or ".jpeg"
+
     norm_type_mapping = {
         "inf": cv.NORM_INF,
         "l1": cv.NORM_L1,
@@ -273,13 +304,19 @@ def image_normalize_func(
     img_normalized = cv.normalize(
         img, None, alpha=alpha, beta=beta, norm_type=norm_type_mapping[norm_type]
     )
-    bts = cv.imencode(".jpeg", img_normalized)[1].tobytes()
+
+    bts = cv.imencode(ext, img_normalized)[1].tobytes()
+
+    ext = ext.replace(".", "")
+    ext_mappings = {"jpg": "jpeg", "tif": "tiff"}
+    ext = ext_mappings.get(ext, ext)
+    content_type = "image/" + ext
 
     requests.put(
         url=dst_url,
         data=bts,
         headers={
-            "Content-Type": "image/jpeg",
+            "Content-Type": content_type,
         },
     )
 
@@ -292,14 +329,16 @@ def image_normalize_func(
 
 
 def image_normalize_to_bytes_func(
-    src_obj_ref_rt: str, alpha: float, beta: float, norm_type: str
+    src_obj_ref_rt: str, alpha: float, beta: float, norm_type: str, ext: str
 ) -> bytes:
     import json
 
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
 
+    ext = ext or ".jpeg"
+
     norm_type_mapping = {
         "inf": cv.NORM_INF,
         "l1": cv.NORM_L1,
@@ -395,7 +434,7 @@ def pdf_chunk_func(src_obj_ref_rt: str, chunk_size: int, overlap_size: int) -> s
     if curr_chunk:
         all_text_chunks.append(curr_chunk)
 
-    all_text_json_string = json.dumps(all_text_chunks)
+    all_text_json_string = json.dumps(["123"])
     return all_text_json_string
 
 

bigframes/core/indexers.py

@@ -460,13 +460,12 @@ def _iloc_getitem_series_or_dataframe(
             return _iloc_getitem_series_or_dataframe(series_or_dataframe, key[0])
 
         # len(key) == 2
+        df = typing.cast(bigframes.dataframe.DataFrame, series_or_dataframe)
         if isinstance(key[1], int):
-            return series_or_dataframe.iat[key]
+            return df.iat[key]
         elif isinstance(key[1], list):
-            columns = series_or_dataframe.columns[key[1]]
-            return _iloc_getitem_series_or_dataframe(
-                series_or_dataframe[columns], key[0]
-            )
+            columns = df.columns[key[1]]
+            return _iloc_getitem_series_or_dataframe(df[columns], key[0])
         raise NotImplementedError(
             f"iloc does not yet support indexing with {key}. {constants.FEEDBACK_LINK}"
         )
@@ -476,13 +475,14 @@ def _iloc_getitem_series_or_dataframe(
                 Union[bigframes.dataframe.DataFrame, bigframes.series.Series],
                 series_or_dataframe.iloc[0:0],
             )
-        df = series_or_dataframe
         if isinstance(series_or_dataframe, bigframes.series.Series):
             original_series_name = series_or_dataframe.name
             series_name = (
                 original_series_name if original_series_name is not None else 0
             )
             df = series_or_dataframe.to_frame()
+        else:
+            df = series_or_dataframe
         original_index_names = df.index.names
         temporary_index_names = [
             guid.generate_guid(prefix="temp_iloc_index_")