diff --git a/bigframes/core/indexes/base.py b/bigframes/core/indexes/base.py index f653b8700b..ea34a465c1 100644 --- a/bigframes/core/indexes/base.py +++ b/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 diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 9f42e70923..69b251fd5b 100644 --- a/bigframes/dataframe.py +++ b/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( diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index c316d28321..f2bbcb3320 100644 --- a/bigframes/operations/base.py +++ b/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 diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index a2b78fcfef..4e74fe020f 100644 --- a/tests/system/small/test_dataframe.py +++ b/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"): diff --git a/tests/system/small/test_index.py b/tests/system/small/test_index.py index c7e316a9d2..2b2364d3bc 100644 --- a/tests/system/small/test_index.py +++ b/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() diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index d513b0e780..a69c6b945b 100644 --- a/tests/system/small/test_series.py +++ b/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"