Improve JSON type handling for to_gbq and to_pandas_batches

Author: shuoweil

bigframes/core/blocks.py

@@ -731,6 +731,12 @@ 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)
@@ -746,7 +752,9 @@ def to_pandas_batches(
                 # MyPy doesn't automatically narrow the type of 'dtype' here,
                 # so we add an explicit check.
                 if isinstance(dtype, pd.ArrowDtype):
-                    safe_pa_type = _replace_json_arrow_with_string(dtype.pyarrow_dtype)
+                    safe_pa_type = bigframes.dtypes._replace_json_arrow_with_string(
+                        dtype.pyarrow_dtype
+                    )
                     safe_dtype = pd.ArrowDtype(safe_pa_type)
                     series_map[col] = pd.Series([], dtype=safe_dtype).astype(dtype)
                 else:

bigframes/dtypes.py

@@ -954,6 +954,21 @@ def contains_db_dtypes_json_dtype(dtype):
     return contains_db_dtypes_json_arrow_type(dtype.pyarrow_dtype)
 
 
+def _replace_json_arrow_with_string(pa_type: pa.DataType) -> pa.DataType:
+    """Recursively replace JSONArrowType with string type."""
+    if isinstance(pa_type, db_dtypes.JSONArrowType):
+        return pa.string()
+    if isinstance(pa_type, pa.ListType):
+        return pa.list_(_replace_json_arrow_with_string(pa_type.value_type))
+    if isinstance(pa_type, pa.StructType):
+        new_fields = [
+            field.with_type(_replace_json_arrow_with_string(field.type))
+            for field in pa_type
+        ]
+        return pa.struct(new_fields)
+    return pa_type
+
+
 def warn_on_db_dtypes_json_dtype(dtypes):
     """Warn that the JSON dtype is changing.
 

tests/system/small/test_dataframe_io.py

@@ -376,6 +376,83 @@ 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
+    """
+    import db_dtypes
+
+    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)
+
+    results = list(empty_df_with_json_list.to_pandas_batches())
+
+    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
+
+
+# --- Behavior 2: JSON in Struct ---
+
+
+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.
+
+    Regression test for https://github.com/googleapis/python-bigquery-dataframes/issues/1273
+    """
+    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
+
+
+# --- Behavior 3: Simple JSON ---
+
+
+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
+    """
+    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)
+
+    results = list(empty_df_with_json.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
+
+
 @pytest.mark.parametrize("allow_large_results", (True, False))
 def test_to_pandas_batches_w_page_size_and_max_results(session, allow_large_results):
     """Verify to_pandas_batches() APIs returns the expected page size.