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

@@ -275,6 +275,43 @@ 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
+):
+    exprs = [
+        ops.AsTypeOp(to_type=bigframes.dtypes.JSON_DTYPE).as_expr(
+            expression.deref("label")
+        ),
+    ]
+    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(