fix testcase

Author: shuoweil

bigframes/blob/_functions.py

@@ -215,6 +215,8 @@ def image_blur_func(
             timeout=30,
         )
 
+        result_dict["content"] = dst_obj_ref_rt_json["objectref"]["uri"]
+
     except Exception as e:
         result_dict["status"] = str(e)
 
@@ -233,8 +235,7 @@ def image_blur_to_bytes_func(
     import base64
     import json
 
-    status = ""
-    content = b""
+    result_dict = {"status": "", "content": ""}
 
     try:
         import cv2 as cv  # type: ignore
@@ -258,11 +259,12 @@ def image_blur_to_bytes_func(
         img_blurred = cv.blur(img, ksize=(ksize_x, ksize_y))
         content = cv.imencode(ext, img_blurred)[1].tobytes()
 
+        encoded_content = base64.b64encode(content).decode("utf-8")
+        result_dict["content"] = encoded_content
+
     except Exception as e:
-        status = str(e)
+        result_dict["status"] = str(e)
 
-    encoded_content = base64.b64encode(content).decode("utf-8")
-    result_dict = {"status": status, "content": encoded_content}
     if verbose:
         return json.dumps(result_dict)
     else:
@@ -328,6 +330,8 @@ def image_resize_func(
             timeout=30,
         )
 
+        result_dict["content"] = dst_obj_ref_rt_json["objectref"]["uri"]
+
     except Exception as e:
         result_dict["status"] = str(e)
 
@@ -354,8 +358,7 @@ def image_resize_to_bytes_func(
     import base64
     import json
 
-    status = ""
-    content = b""
+    result_dict = {"status": "", "content": ""}
 
     try:
         import cv2 as cv  # type: ignore
@@ -379,11 +382,12 @@ def image_resize_to_bytes_func(
         img_resized = cv.resize(img, dsize=(dsize_x, dsize_y), fx=fx, fy=fy)
         content = cv.imencode(".jpeg", img_resized)[1].tobytes()
 
+        encoded_content = base64.b64encode(content).decode("utf-8")
+        result_dict["content"] = encoded_content
+
     except Exception as e:
-        status = str(e)
+        result_dict["status"] = str(e)
 
-    encoded_content = base64.b64encode(content).decode("utf-8")
-    result_dict = {"status": status, "content": encoded_content}
     if verbose:
         return json.dumps(result_dict)
     else:
@@ -457,6 +461,8 @@ def image_normalize_func(
             timeout=30,
         )
 
+        result_dict["content"] = dst_obj_ref_rt_json["objectref"]["uri"]
+
     except Exception as e:
         result_dict["status"] = str(e)
 

bigframes/dataframe.py

@@ -95,12 +95,8 @@
 
     import bigframes.session
 
-    SingleItemValue = Union[
-        bigframes.series.Series, int, float, str, pandas.Timedelta, Callable
-    ]
-    MultiItemValue = Union[
-        "DataFrame", Sequence[int | float | str | pandas.Timedelta | Callable]
-    ]
+    SingleItemValue = Union[bigframes.series.Series, int, float, str, Callable]
+    MultiItemValue = Union["DataFrame", Sequence[Union[int, float, str, Callable]]]
 
 LevelType = typing.Hashable
 LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]]

bigframes/operations/blob.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import os
+import typing
 from typing import cast, Literal, Optional, Union
 import warnings
 
@@ -371,7 +372,10 @@ def image_blur(
         container_cpu: Union[float, int] = 0.33,
         container_memory: str = "512Mi",
         verbose: bool = False,
-    ) -> bigframes.series.Series:
+    ) -> typing.Union[
+        bigframes.series.Series,
+        typing.Tuple[bigframes.series.Series, bigframes.series.Series],
+    ]:
         """Blurs images.
 
         Args:
@@ -474,11 +478,8 @@ def image_blur(
             )
             content_series = res._apply_unary_op(ops.JSONValue(json_path="$.content"))
             dst_blobs = content_series.str.to_blob(connection=connection)
-            results_df = bpd.DataFrame(
-                {"status": blurred_status_series, "content": dst_blobs}
-            )
-            results_struct = bbq.struct(results_df).rename("blurred_results")
-            return results_struct
+
+            return dst_blobs, blurred_status_series
         else:
             return res.str.to_blob(connection=connection)
 
@@ -631,7 +632,10 @@ def image_normalize(
         container_cpu: Union[float, int] = 0.33,
         container_memory: str = "512Mi",
         verbose: bool = False,
-    ) -> bigframes.series.Series:
+    ) -> typing.Union[
+        bigframes.series.Series,
+        typing.Tuple[bigframes.series.Series, bigframes.series.Series],
+    ]:
         """Normalize images.
 
         Args:
@@ -740,14 +744,8 @@ def image_normalize(
             )
             content_series = res._apply_unary_op(ops.JSONValue(json_path="$.content"))
             dst_blobs = content_series.str.to_blob(connection=connection)
-            results_df = bpd.DataFrame(
-                {
-                    "status": normalized_status_series,
-                    "content": dst_blobs,
-                }
-            )
-            results_struct = bbq.struct(results_df).rename("normalized_results")
-            return results_struct
+
+            return dst_blobs, normalized_status_series
         else:
             return res.str.to_blob(connection=connection)
 

tests/system/large/blob/test_function.py

@@ -120,6 +120,7 @@ def test_blob_image_blur_to_series(
         content_series = actual_exploded["content"]
         # Content should be blob objects for GCS destination
         assert hasattr(content_series, "blob")
+        assert not content_series.blob.size().isna().any()
 
     else:
         expected_df = pd.DataFrame(
@@ -131,14 +132,13 @@ def test_blob_image_blur_to_series(
             }
         )
         pd.testing.assert_frame_equal(
-            actual.struct.explode().to_pandas(),
+            actual.blob.to_frame().to_pandas(),
             expected_df,
             check_dtype=False,
             check_index_type=False,
         )
-
-    # verify the files exist
-    assert not actual.blob.size().isna().any()
+        # verify the files exist
+        assert not actual.blob.size().isna().any()
 
 
 @pytest.mark.parametrize("verbose", [True, False])
@@ -168,6 +168,7 @@ def test_blob_image_blur_to_folder(
         content_series = actual_exploded["content"]
         # Content should be blob objects for GCS destination
         assert hasattr(content_series, "blob")
+        assert not content_series.blob.size().isna().any()
 
     else:
         expected_df = pd.DataFrame(
@@ -179,14 +180,13 @@ def test_blob_image_blur_to_folder(
             }
         )
         pd.testing.assert_frame_equal(
-            actual.struct.explode().to_pandas(),
+            actual.blob.to_frame().to_pandas(),
             expected_df,
             check_dtype=False,
             check_index_type=False,
         )
-
-    # verify the files exist
-    assert not actual.blob.size().isna().any()
+        # verify the files exist
+        assert not actual.blob.size().isna().any()
 
 
 @pytest.mark.parametrize("verbose", [True, False])
@@ -248,6 +248,7 @@ def test_blob_image_resize_to_series(
         content_series = actual_exploded["content"]
         # Content should be blob objects for GCS destination
         assert hasattr(content_series, "blob")
+        assert not content_series.blob.size().isna().any()
 
     else:
         expected_df = pd.DataFrame(
@@ -259,14 +260,13 @@ def test_blob_image_resize_to_series(
             }
         )
         pd.testing.assert_frame_equal(
-            actual.struct.explode().to_pandas(),
+            actual.blob.to_frame().to_pandas(),
             expected_df,
             check_dtype=False,
             check_index_type=False,
         )
-
-    # verify the files exist
-    assert not actual.blob.size().isna().any()
+        # verify the files exist
+        assert not actual.blob.size().isna().any()
 
 
 @pytest.mark.parametrize("verbose", [True, False])
@@ -297,6 +297,7 @@ def test_blob_image_resize_to_folder(
         content_series = actual_exploded["content"]
         # Content should be blob objects for GCS destination
         assert hasattr(content_series, "blob")
+        assert not content_series.blob.size().isna().any()
 
     else:
         expected_df = pd.DataFrame(
@@ -308,14 +309,13 @@ def test_blob_image_resize_to_folder(
             }
         )
         pd.testing.assert_frame_equal(
-            actual.struct.explode().to_pandas(),
+            actual.blob.to_frame().to_pandas(),
             expected_df,
             check_dtype=False,
             check_index_type=False,
         )
-
-    # verify the files exist
-    assert not actual.blob.size().isna().any()
+        # verify the files exist
+        assert not actual.blob.size().isna().any()
 
 
 @pytest.mark.parametrize("verbose", [True, False])
@@ -368,7 +368,6 @@ def test_blob_image_normalize_to_series(
     )
 
     if verbose:
-
         assert hasattr(actual, "struct")
         actual_exploded = actual.struct.explode()
         assert "status" in actual_exploded.columns
@@ -380,6 +379,7 @@ def test_blob_image_normalize_to_series(
         content_series = actual_exploded["content"]
         # Content should be blob objects for GCS destination
         assert hasattr(content_series, "blob")
+        assert not content_series.blob.size().isna().any()
 
     else:
         expected_df = pd.DataFrame(
@@ -391,7 +391,7 @@ def test_blob_image_normalize_to_series(
             }
         )
         pd.testing.assert_frame_equal(
-            actual.struct.explode().to_pandas(),
+            actual.blob.to_frame().to_pandas(),
             expected_df,
             check_dtype=False,
             check_index_type=False,
@@ -431,6 +431,7 @@ def test_blob_image_normalize_to_folder(
         content_series = actual_exploded["content"]
         # Content should be blob objects for GCS destination
         assert hasattr(content_series, "blob")
+        assert not content_series.blob.size().isna().any()
 
     else:
         expected_df = pd.DataFrame(
@@ -442,14 +443,14 @@ def test_blob_image_normalize_to_folder(
             }
         )
         pd.testing.assert_frame_equal(
-            actual.struct.explode().to_pandas(),
+            actual.blob.to_frame().to_pandas(),
             expected_df,
             check_dtype=False,
             check_index_type=False,
         )
 
-    # verify the files exist
-    assert not actual.blob.size().isna().any()
+        # verify the files exist
+        assert not actual.blob.size().isna().any()
 
 
 @pytest.mark.parametrize("verbose", [True, False])