testcase update

Author: shuoweil

bigframes/core/blocks.py

@@ -711,12 +711,6 @@ def to_pandas_batches(
         # To reduce the number of edge cases to consider when working with the
         # results of this, always return at least one DataFrame. See:
         # b/428918844.
-        empty_val = pd.DataFrame(
-            {
-                col: pd.Series([], dtype=self.expr.get_column_type(col))
-                for col in itertools.chain(self.value_columns, self.index_columns)
-            }
-        )
         series_map = {}
         for col in itertools.chain(self.value_columns, self.index_columns):
             dtype = self.expr.get_column_type(col)
@@ -733,10 +727,8 @@ def to_pandas_batches(
                     safe_dtype = pd.ArrowDtype(safe_pa_type)
                     series_map[col] = pd.Series([], dtype=safe_dtype).astype(dtype)
                 else:
-                    # This branch should ideally not be reached if
-                    # contains_db_dtypes_json_dtype is accurate,
-                    # but it's here for MyPy's sake.
-                    series_map[col] = pd.Series([], dtype=dtype)
+                    # Fallback for other types that might error
+                    series_map[col] = pd.Series([], dtype="object").astype(dtype)
         empty_val = pd.DataFrame(series_map)
         dfs = map(
             lambda a: a[0],

tests/system/small/test_dataframe_io.py

@@ -376,75 +376,52 @@ def test_to_pandas_batches_w_empty_dataframe(session):
     pandas.testing.assert_series_equal(results[0].dtypes, empty.dtypes)
 
 
-def test_to_pandas_batches_w_empty_dataframe_json_in_list(session):
-    """Tests to_pandas_batches() with an empty DataFrame containing a list of JSON.
-
-    Regression test for https://github.com/googleapis/python-bigquery-dataframes/issues/1273
+def test_to_pandas_batches_preserves_dtypes_for_populated_nested_json(session):
+    """Verifies to_pandas_batches() preserves dtypes for nested JSON."""
+    # This SQL query only tests the POPULATED case.
+    sql = """
+        SELECT
+            0 AS id,
+            [JSON '{"a":1}', JSON '{"b":2}'] AS json_array,
+            STRUCT(JSON '{"x":1}' AS json_field, 'test' AS str_field) AS json_struct
     """
-    import db_dtypes
+    df = session.read_gbq(sql, index_col="id")
 
-    json_list_dtype = pd.ArrowDtype(pa.list_(db_dtypes.JSONArrowType()))
-    empty_df_with_json_list = bpd.DataFrame(
-        {
-            "idx": pd.Series([], dtype="Int64"),
-            "json_list_col": pd.Series([], dtype=json_list_dtype),
-        },
-        session=session,
-    ).set_index("idx", drop=True)
+    batches = list(df.to_pandas_batches())
 
-    results = list(empty_df_with_json_list.to_pandas_batches())
+    # Check that we processed the row
+    assert sum(len(b) for b in batches) == 1
 
-    assert len(results) == 1
-    assert list(results[0].columns) == ["json_list_col"]
-    assert results[0].dtypes["json_list_col"] == json_list_dtype
-    assert len(results[0]) == 0
+    # Check dtypes on the resulting batch
+    assert isinstance(batches[0].dtypes["json_array"], pd.ArrowDtype)
+    assert isinstance(batches[0].dtypes["json_array"].pyarrow_dtype, pa.ListType)
+    assert isinstance(batches[0].dtypes["json_struct"], pd.ArrowDtype)
+    assert isinstance(batches[0].dtypes["json_struct"].pyarrow_dtype, pa.StructType)
 
 
-def test_to_pandas_batches_w_empty_dataframe_json_in_struct(session):
-    """Tests to_pandas_batches() with an empty DataFrame containing a struct of JSON.
+def test_to_pandas_batches_should_not_error_on_empty_nested_json(session):
+    """Verify to_pandas_batches() works with empty nested JSON types.
 
-    Regression test for https://github.com/googleapis/python-bigquery-dataframes/issues/1273
+    Regression test for PyArrow limitation with empty JSON arrays.
     """
-    import db_dtypes
-
-    json_struct_dtype = pd.ArrowDtype(
-        pa.struct([("json_field", db_dtypes.JSONArrowType())])
-    )
-    empty_df_with_json_struct = bpd.DataFrame(
-        {
-            "idx": pd.Series([], dtype="Int64"),
-            "json_struct_col": pd.Series([], dtype=json_struct_dtype),
-        },
-        session=session,
-    ).set_index("idx", drop=True)
-
-    results = list(empty_df_with_json_struct.to_pandas_batches())
-
-    assert len(results) == 1
-    assert list(results[0].columns) == ["json_struct_col"]
-    assert results[0].dtypes["json_struct_col"] == json_struct_dtype
-    assert len(results[0]) == 0
-
-
-def test_to_pandas_batches_w_empty_dataframe_simple_json(session):
-    """Tests to_pandas_batches() with an empty DataFrame containing a simple JSON column.
-
-    Regression test for https://github.com/googleapis/python-bigquery-dataframes/issues/1273
+    # This SQL query is MINIMAL and tests only the EMPTY regression case.
+    sql = """
+        SELECT
+            1 AS id,
+            [] AS json_array,
+            STRUCT(NULL AS json_field, 'test2' AS str_field) AS json_struct
     """
-    empty_df_with_json = bpd.DataFrame(
-        {
-            "idx": pd.Series([], dtype="Int64"),
-            "json_col": pd.Series([], dtype=dtypes.JSON_DTYPE),
-        },
-        session=session,
-    ).set_index("idx", drop=True)
+    df = session.read_gbq(sql, index_col="id")
 
-    results = list(empty_df_with_json.to_pandas_batches())
+    # The main point of this test is that this line does not raise an error.
+    batches = list(df.to_pandas_batches())
 
-    assert len(results) == 1
-    assert list(results[0].columns) == ["json_col"]
-    assert results[0].dtypes["json_col"] == dtypes.JSON_DTYPE
-    assert len(results[0]) == 0
+    # Verify the row was actually processed and not just skipped
+    assert sum(len(b) for b in batches) == 1
+
+    # Verify dtypes are still correct, even with empty data
+    assert isinstance(batches[0].dtypes["json_array"], pd.ArrowDtype)
+    assert isinstance(batches[0].dtypes["json_struct"], pd.ArrowDtype)
 
 
 @pytest.mark.parametrize("allow_large_results", (True, False))