tests: add engine tests for casting to json

Author: chelsea-lin

tests/system/small/engines/conftest.py

@@ -73,6 +73,13 @@ def scalars_array_value(
     return ArrayValue.from_managed(managed_data_source, fake_session)
 
 
+@pytest.fixture(scope="module")
+def nested_array_value(
+    nested_data_source: local_data.ManagedArrowTable, fake_session: bigframes.Session
+):
+    return ArrayValue.from_managed(nested_data_source, fake_session)
+
+
 @pytest.fixture(scope="module")
 def zero_row_source() -> local_data.ManagedArrowTable:
     return local_data.ManagedArrowTable.from_pandas(pd.DataFrame({"a": [], "b": []}))

tests/system/small/engines/test_generic_ops.py

@@ -14,6 +14,8 @@
 
 import re
 
+import pandas as pd
+import pyarrow as pa
 import pytest
 
 from bigframes.core import array_value, expression
@@ -275,6 +277,58 @@ def test_engines_astype_from_json(scalars_array_value: array_value.ArrayValue, e
     assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_astype_to_json(scalars_array_value: array_value.ArrayValue, engine):
+    exprs = [
+        ops.AsTypeOp(to_type=bigframes.dtypes.JSON_DTYPE).as_expr(
+            expression.deref("int64_col")
+        ),
+        ops.AsTypeOp(to_type=bigframes.dtypes.JSON_DTYPE).as_expr(
+            # Use a const since float to json has precision issues
+            expression.const(5.2, bigframes.dtypes.FLOAT_DTYPE)
+        ),
+        ops.AsTypeOp(to_type=bigframes.dtypes.JSON_DTYPE).as_expr(
+            expression.deref("bool_col")
+        ),
+        ops.AsTypeOp(to_type=bigframes.dtypes.JSON_DTYPE).as_expr(
+            # Use a const since "str_col" has special chars.
+            expression.const('"hello world"', bigframes.dtypes.STRING_DTYPE)
+        ),
+    ]
+    arr, _ = scalars_array_value.compute_values(exprs)
+
+    assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_astype_struct_to_json(
+    nested_array_value: array_value.ArrayValue, engine
+):
+    json_data = [
+        {"version": 1, "project": "pandas"},
+        {"version": 2, "project": "numpy"},
+    ]
+    exprs = [
+        # ops.AsTypeOp(to_type=bigframes.dtypes.JSON_DTYPE).as_expr(
+        #     expression.deref("label")
+        # ),
+        # ops.AsTypeOp(to_type=bigframes.dtypes.JSON_DTYPE).as_expr(
+        #     expression.deref("address")
+        # ),
+        ops.AsTypeOp(to_type=bigframes.dtypes.JSON_DTYPE).as_expr(
+            expression.const(
+                json_data,
+                pd.ArrowDtype(
+                    pa.struct([("version", pa.int64()), ("project", pa.string())])
+                ),
+            )
+        ),
+    ]
+    arr, _ = nested_array_value.compute_values(exprs)
+
+    assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
+
+
 @pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
 def test_engines_astype_timedelta(scalars_array_value: array_value.ArrayValue, engine):
     arr = apply_op(