feat: support bpd.Series(json_data, dtype="json") (#1882)

* feat: support bpd.Series(json_data, dtype="json")

* undo dtypes._dtype_from_string and apply json only

Author: chelsea-lin

bigframes/core/indexes/base.py

@@ -86,6 +86,8 @@ def __new__(
             pd_df = pandas.DataFrame(index=data)
             block = df.DataFrame(pd_df, session=session)._block
         else:
+            if isinstance(dtype, str) and dtype.lower() == "json":
+                dtype = bigframes.dtypes.JSON_DTYPE
             pd_index = pandas.Index(data=data, dtype=dtype, name=name)
             pd_df = pandas.DataFrame(index=pd_index)
             block = df.DataFrame(pd_df, session=session)._block

bigframes/dataframe.py

@@ -196,6 +196,9 @@ def __init__(
                 block = block.multi_apply_unary_op(ops.AsTypeOp(to_type=bf_dtype))
 
         else:
+            if isinstance(dtype, str) and dtype.lower() == "json":
+                dtype = bigframes.dtypes.JSON_DTYPE
+
             import bigframes.pandas
 
             pd_dataframe = pandas.DataFrame(

bigframes/operations/base.py

@@ -121,6 +121,8 @@ def __init__(
                 bf_dtype = bigframes.dtypes.bigframes_type(dtype)
                 block = block.multi_apply_unary_op(ops.AsTypeOp(to_type=bf_dtype))
         else:
+            if isinstance(dtype, str) and dtype.lower() == "json":
+                dtype = bigframes.dtypes.JSON_DTYPE
             pd_series = pd.Series(
                 data=data,
                 index=index,  # type:ignore

tests/system/small/test_dataframe.py

@@ -180,6 +180,26 @@ def test_df_construct_from_dict():
     )
 
 
+@pytest.mark.parametrize(
+    ("json_type"),
+    [
+        pytest.param(dtypes.JSON_DTYPE),
+        pytest.param("json"),
+    ],
+)
+def test_df_construct_w_json_dtype(json_type):
+    data = [
+        "1",
+        "false",
+        '["a", {"b": 1}, null]',
+        None,
+    ]
+    df = dataframe.DataFrame({"json_col": data}, dtype=json_type)
+
+    assert df["json_col"].dtype == dtypes.JSON_DTYPE
+    assert df["json_col"][1] == "false"
+
+
 def test_df_construct_inline_respects_location(reset_default_session_and_location):
     # Note: This starts a thread-local session.
     with bpd.option_context("bigquery.location", "europe-west1"):

tests/system/small/test_index.py

@@ -18,6 +18,7 @@
 import pandas as pd
 import pytest
 
+from bigframes import dtypes
 import bigframes.pandas as bpd
 from bigframes.testing.utils import assert_pandas_index_equal_ignore_index_type
 
@@ -61,6 +62,26 @@ def test_index_construct_from_index():
     pd.testing.assert_index_equal(bf_result, pd_result)
 
 
+@pytest.mark.parametrize(
+    ("json_type"),
+    [
+        pytest.param(dtypes.JSON_DTYPE),
+        pytest.param("json"),
+    ],
+)
+def test_index_construct_w_json_dtype(json_type):
+    data = [
+        "1",
+        "false",
+        '["a", {"b": 1}, null]',
+        None,
+    ]
+    index = bpd.Index(data, dtype=json_type)
+
+    assert index.dtype == dtypes.JSON_DTYPE
+    assert index[1] == "false"
+
+
 def test_get_index(scalars_df_index, scalars_pandas_df_index):
     index = scalars_df_index.index
     bf_result = index.to_pandas()

tests/system/small/test_series.py

@@ -326,7 +326,14 @@ def test_series_construct_local_unordered_has_sequential_index(unordered_session
     pd.testing.assert_index_equal(series.index.to_pandas(), expected)
 
 
-def test_series_construct_w_dtype_for_json():
+@pytest.mark.parametrize(
+    ("json_type"),
+    [
+        pytest.param(dtypes.JSON_DTYPE),
+        pytest.param("json"),
+    ],
+)
+def test_series_construct_w_json_dtype(json_type):
     data = [
         "1",
         '"str"',
@@ -335,8 +342,9 @@ def test_series_construct_w_dtype_for_json():
         None,
         '{"a": {"b": [1, 2, 3], "c": true}}',
     ]
-    s = bigframes.pandas.Series(data, dtype=dtypes.JSON_DTYPE)
+    s = bigframes.pandas.Series(data, dtype=json_type)
 
+    assert s.dtype == dtypes.JSON_DTYPE
     assert s[0] == "1"
     assert s[1] == '"str"'
     assert s[2] == "false"