feat: guard polar engine

Author: shuoweil

bigframes/core/compile/polars/compiler.py

@@ -429,7 +429,13 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
         @compile_op.register(json_ops.JSONDecode)
         def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
             assert isinstance(op, json_ops.JSONDecode)
-            return input.str.json_decode(_DTYPE_MAPPING[op.to_type])
+            if op.safe:
+                # Polars does not support safe JSON decoding (returning null on failure).
+                # Fallback to BigQuery execution.
+                raise NotImplementedError(
+                    "Safe JSON decoding is not supported in Polars executor."
+                )
+            return input.str.json_decode(_bigframes_dtype_to_polars_dtype(op.to_type))
 
         @compile_op.register(arr_ops.ToArrayOp)
         def _(self, op: ops.ToArrayOp, *inputs: pl.Expr) -> pl.Expr:

bigframes/core/compile/polars/lowering.py

@@ -391,7 +391,7 @@ def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
         return arg
 
     if arg.output_type == dtypes.JSON_DTYPE:
-        return json_ops.JSONDecode(cast_op.to_type).as_expr(arg)
+        return json_ops.JSONDecode(cast_op.to_type, safe=cast_op.safe).as_expr(arg)
     if (
         arg.output_type == dtypes.STRING_DTYPE
         and cast_op.to_type == dtypes.DATETIME_DTYPE

bigframes/operations/json_ops.py

@@ -220,6 +220,7 @@ def output_type(self, *input_types):
 class JSONDecode(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "json_decode"
     to_type: dtypes.Dtype
+    safe: bool = False
 
     def output_type(self, *input_types):
         input_type = input_types[0]

bigframes/session/polars_executor.py

@@ -34,6 +34,7 @@
     numeric_ops,
     string_ops,
 )
+import bigframes.operations.json_ops as json_ops
 from bigframes.session import executor, semi_executor
 
 if TYPE_CHECKING:
@@ -94,6 +95,7 @@
     string_ops.EndsWithOp,
     string_ops.StrContainsOp,
     string_ops.StrContainsRegexOp,
+    json_ops.JSONDecode,
 )
 _COMPATIBLE_AGG_OPS = (
     agg_ops.SizeOp,