increase test coverage

Author: chelsea-lin

bigframes/core/utils.py

@@ -245,17 +245,17 @@ def replace_timedeltas_with_micros(dataframe: pd.DataFrame) -> List[str]:
     return updated_columns
 
 
-def search_for_nested_json_type(arrow_type: pa.DataType) -> bool:
+def _search_for_nested_json_type(arrow_type: pa.DataType) -> bool:
     """
-    Searches recursively for pa.JsonType within a PyArrow DataType.
+    Searches recursively for JSON array type within a PyArrow DataType.
     """
-    if isinstance(arrow_type, pa.JsonType):
+    if arrow_type == dtypes.JSON_ARROW_TYPE:
         return True
     if pa.types.is_list(arrow_type):
-        return search_for_nested_json_type(arrow_type.value_type)
+        return _search_for_nested_json_type(arrow_type.value_type)
     if pa.types.is_struct(arrow_type):
         return any(
-            search_for_nested_json_type(field.type) for field in arrow_type.fields
+            _search_for_nested_json_type(field.type) for field in arrow_type.fields
         )
     return False
 
@@ -274,7 +274,9 @@ def replace_json_with_string(dataframe: pd.DataFrame) -> List[str]:
         if column_type == dtypes.JSON_DTYPE:
             dataframe[col] = dataframe[col].astype(dtypes.STRING_DTYPE)
             updated_columns.append(col)
-        elif isinstance(column_type, pd.ArrowDtype):
+        elif isinstance(column_type, pd.ArrowDtype) and _search_for_nested_json_type(
+            column_type.pyarrow_dtype
+        ):
             raise NotImplementedError(
                 f"Nested JSON types, found in column `{col}`: `{column_type}`', "
                 f"are currently unsupported for upload. {constants.FEEDBACK_LINK}"
@@ -283,9 +285,11 @@ def replace_json_with_string(dataframe: pd.DataFrame) -> List[str]:
     if dataframe.index.dtype == dtypes.JSON_DTYPE:
         dataframe.index = dataframe.index.astype(dtypes.STRING_DTYPE)
         updated_columns.append(dataframe.index.name)
-    elif isinstance(column_type, pd.ArrowDtype):
+    elif isinstance(
+        dataframe.index.dtype, pd.ArrowDtype
+    ) and _search_for_nested_json_type(dataframe.index.dtype.pyarrow_dtype):
         raise NotImplementedError(
-            f"Nested JSON types, found in column `{col}`: `{column_type}`', "
+            f"Nested JSON types, found in the index: `{dataframe.index.dtype}`', "
             f"are currently unsupported for upload. {constants.FEEDBACK_LINK}"
         )
 

tests/system/small/test_session.py

@@ -840,19 +840,21 @@ def test_read_pandas_json_index(session, write_engine):
 )
 def test_read_pandas_w_nested_json(session, write_engine):
     data = [
-        {"json_field": "1"},
-        {"json_field": None},
-        {"json_field": '["1","3","5"]'},
-        {"json_field": '{"a":1,"b":["x","y"],"c":{"x":[],"z":false}}'},
+        [{"json_field": "1"}],
+        [{"json_field": None}],
+        [{"json_field": '["1","3","5"]'}],
+        [{"json_field": '{"a":1,"b":["x","y"],"c":{"x":[],"z":false}}'}],
     ]
     # PyArrow currently lacks support for creating structs or lists containing extension types.
     # See issue: https://github.com/apache/arrow/issues/45262
-    pa_array = pa.array(data, type=pa.struct([("name", pa.string())]))
+    pa_array = pa.array(data, type=pa.list_(pa.struct([("name", pa.string())])))
     pd_s = pd.Series(
         arrays.ArrowExtensionArray(pa_array),  # type: ignore
-        dtype=pd.ArrowDtype(pa.struct([("name", bigframes.dtypes.JSON_ARROW_TYPE)])),
+        dtype=pd.ArrowDtype(
+            pa.list_(pa.struct([("name", bigframes.dtypes.JSON_ARROW_TYPE)]))
+        ),
     )
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(NotImplementedError, match="Nested JSON types, found in column"):
         # Until b/401630655 is resolved, json not compatible with allow_large_results=False
         session.read_pandas(pd_s, write_engine=write_engine).to_pandas(
             allow_large_results=True
@@ -868,19 +870,21 @@ def test_read_pandas_w_nested_json(session, write_engine):
 )
 def test_read_pandas_w_nested_json_index(session, write_engine):
     data = [
-        {"json_field": "1"},
-        {"json_field": None},
-        {"json_field": '["1","3","5"]'},
-        {"json_field": '{"a":1,"b":["x","y"],"c":{"x":[],"z":false}}'},
+        [{"json_field": "1"}],
+        [{"json_field": None}],
+        [{"json_field": '["1","3","5"]'}],
+        [{"json_field": '{"a":1,"b":["x","y"],"c":{"x":[],"z":false}}'}],
     ]
     # PyArrow currently lacks support for creating structs or lists containing extension types.
     # See issue: https://github.com/apache/arrow/issues/45262
-    pa_array = pa.array(data, type=pa.struct([("name", pa.string())]))
+    pa_array = pa.array(data, type=pa.list_(pa.struct([("name", pa.string())])))
     pd_idx: pd.Index = pd.Index(
         arrays.ArrowExtensionArray(pa_array),  # type: ignore
-        dtype=pd.ArrowDtype(pa.struct([("name", bigframes.dtypes.JSON_ARROW_TYPE)])),
+        dtype=pd.ArrowDtype(
+            pa.list_(pa.struct([("name", bigframes.dtypes.JSON_ARROW_TYPE)]))
+        ),
     )
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(NotImplementedError, match="Nested JSON types, found in the index"):
         # Until b/401630655 is resolved, json not compatible with allow_large_results=False
         session.read_pandas(pd_idx, write_engine=write_engine).to_pandas(
             allow_large_results=True