casting fixes

Author: TrevorBergeron

packages/bigframes/bigframes/core/compile/lowering.py

@@ -290,7 +290,7 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
 
 @dataclasses.dataclass
 class LowerAsTypeRule(op_lowering.OpLoweringRule):
-    dialect: Literal["polars", "substrait"]
+    dialect: Literal["polars", "substrait-datafusion", "substrait-acero"]
 
     @property
     def op(self) -> type[ops.ScalarOp]:
@@ -300,8 +300,8 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
         assert isinstance(expr.op, ops.AsTypeOp)
         if self.dialect == "polars":
             return _lower_cast_to_polars(expr.op, expr.inputs[0])
-        elif self.dialect == "substrait":
-            return _lower_cast_to_substrait(expr.op, expr.inputs[0])
+        else:
+            return _lower_cast_to_substrait(expr.op, expr.inputs[0], dialect=self.dialect)
 
 
 def invert_bytes(byte_string):
@@ -473,7 +473,11 @@ def _lower_cast_to_polars(cast_op: ops.AsTypeOp, arg: expression.Expression):
     return cast_op.as_expr(arg)
 
 
-def _lower_cast_to_substrait(cast_op: ops.AsTypeOp, arg: expression.Expression):
+def _lower_cast_to_substrait(
+    cast_op: ops.AsTypeOp,
+    arg: expression.Expression,
+    dialect: Literal["substrait-datafusion", "substrait-acero"] = "substrait-datafusion",
+):
     if arg.output_type == dtypes.BOOL_DTYPE and cast_op.to_type == dtypes.STRING_DTYPE:
         is_true_cond = ops.eq_op.as_expr(arg, expression.const(True))
         is_false_cond = ops.eq_op.as_expr(arg, expression.const(False))
@@ -483,6 +487,29 @@ def _lower_cast_to_substrait(cast_op: ops.AsTypeOp, arg: expression.Expression):
             is_false_cond,
             expression.const("False"),
         )
+
+    if cast_op.to_type == dtypes.STRING_DTYPE:
+        if arg.output_type == dtypes.DATETIME_DTYPE:
+            cast_expr = cast_op.as_expr(arg)
+            if dialect == "substrait-datafusion":
+                return string_ops.ReplaceStrOp(pat="T", repl=" ").as_expr(cast_expr)
+            else:
+                # Acero: let it cast natively, compiler will intercept and cast to precision_timestamp(0)
+                return cast_expr
+
+        elif arg.output_type == dtypes.TIME_DTYPE:
+            # Both engines use native cast (Acero is excluded in test)
+            return cast_op.as_expr(arg)
+
+        elif arg.output_type == dtypes.TIMESTAMP_DTYPE:
+            cast_expr = cast_op.as_expr(arg)
+            if dialect == "substrait-datafusion":
+                replaced_t = string_ops.ReplaceStrOp(pat="T", repl=" ").as_expr(cast_expr)
+                return string_ops.ReplaceStrOp(pat="Z", repl="+00").as_expr(replaced_t)
+            else:
+                # Acero: native cast (excluded in test)
+                return cast_expr
+
     return cast_op.as_expr(arg)
 
 
@@ -567,21 +594,20 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
     LowerFloorOp(),
 )
 
-SUBSTRAIT_LOWERING_RULES = (
-    LowerEqNullsMatchRule(),
-    *SUBSTRAIT_LOWER_COMPARISONS,
-    LowerAsTypeRule(dialect="substrait"),
-)
-
-
 def lower_ops_to_polars(root: bigframe_node.BigFrameNode) -> bigframe_node.BigFrameNode:
     return op_lowering.lower_ops(root, rules=POLARS_LOWERING_RULES)
 
 
 def lower_ops_to_substrait(
     root: bigframe_node.BigFrameNode,
+    dialect: Literal["substrait-datafusion", "substrait-acero"] = "substrait-datafusion",
 ) -> bigframe_node.BigFrameNode:
-    return op_lowering.lower_ops(root, rules=SUBSTRAIT_LOWERING_RULES)
+    rules = (
+        LowerEqNullsMatchRule(),
+        *SUBSTRAIT_LOWER_COMPARISONS,
+        LowerAsTypeRule(dialect=dialect),
+    )
+    return op_lowering.lower_ops(root, rules=rules)
 
 
 def _numeric_to_timedelta(expr: expression.Expression) -> expression.Expression:

packages/bigframes/bigframes/core/compile/substrait/compiler.py

@@ -30,6 +30,7 @@
 import bigframes.operations.generic_ops as generic_ops
 import bigframes.operations.numeric_ops as numeric_ops
 import bigframes.operations.struct_ops as struct_ops
+import bigframes.operations.string_ops as string_ops
 from bigframes.core import agg_expressions, bigframe_node, nodes, rewrite
 from bigframes.core.compile import lowering
 
@@ -43,9 +44,11 @@ def __init__(
         self,
         duration_type: Literal["interval_day", "int"],
         use_precision_types: bool = True,
+        dialect: Literal["substrait-datafusion", "substrait-acero"] = "substrait-datafusion",
     ):
         self._duration_type = duration_type
         self._use_precision_types = use_precision_types
+        self._dialect = dialect
 
     def compile(self, plan: bigframe_node.BigFrameNode) -> Optional[bytes]:
         """
@@ -56,7 +59,7 @@ def compile(self, plan: bigframe_node.BigFrameNode) -> Optional[bytes]:
 
         # Need to bind types in before lowering
         plan = rewrite.bind_schema_to_tree(plan)
-        plan = lowering.lower_ops_to_substrait(plan)
+        plan = lowering.lower_ops_to_substrait(plan, dialect=self._dialect)
         pb_rel = self._compile_node(plan)
 
         pb_plan = plan_pb2.Plan()
@@ -73,10 +76,26 @@ def compile(self, plan: bigframe_node.BigFrameNode) -> Optional[bytes]:
                 )
             )
 
-        for name, anchor in self._EXTENSIONS.items():
+        # Determine extensions dynamically based on execution engine
+        extensions = dict(self._EXTENSIONS)
+        if self._use_precision_types:
+            # DataFusion supports standard "replace"
+            extensions["replace"] = 76
+        else:
+            # Acero expects Arrow simple extension function URN for "replace_substring"
+            extensions["replace_substring"] = 76
+
+            # Register Arrow simple extension URN at anchor 1
+            arrow_uri = pb_plan.extension_urns.add()
+            arrow_uri.extension_urn_anchor = 1
+            arrow_uri.urn = "urn:arrow:substrait_simple_extension_function"
+
+        for name, anchor in extensions.items():
             ext = pb_plan.extensions.add()
             ext.extension_function.function_anchor = anchor
             ext.extension_function.name = name
+            if name == "replace_substring" and not self._use_precision_types:
+                ext.extension_function.extension_urn_reference = 1
 
         return pb_plan.SerializeToString()
 
@@ -848,8 +867,64 @@ def _compile_astype(
         child: nodes.BigFrameNode,
     ) -> algebra_pb2.Expression:
         arg_expr = self._compile_expression(inputs[0], child)
+        from_dtype = self._get_expression_dtype(inputs[0], child)
+
+        if op.to_type == dtypes.STRING_DTYPE:
+            if from_dtype == dtypes.DATETIME_DTYPE:
+                # This is only reached for Acero (dialect="substrait-acero"),
+                # because DataFusion was lowered to ReplaceStrOp in lowering.py!
+                if not self._use_precision_types:
+                    # Cast to precision_timestamp with precision 0 first, then to string
+                    second_ts_expr = self._compile_cast_with_type_dict(
+                        arg_expr, {"precision_timestamp": {"precision": 0}}
+                    )
+                    return self._compile_cast(second_ts_expr, dtypes.STRING_DTYPE)
+
         return self._compile_cast(arg_expr, op.to_type)
 
+    @_compile_op.register(string_ops.ReplaceStrOp)
+    def _compile_replace_str(
+        self,
+        op: string_ops.ReplaceStrOp,
+        inputs: Sequence[ex.Expression],
+        child: nodes.BigFrameNode,
+    ) -> algebra_pb2.Expression:
+        arg_expr = self._compile_expression(inputs[0], child)
+        return self._compile_replace(arg_expr, op.pat, op.repl)
+
+    def _compile_cast_with_type_dict(
+        self, input_expr: algebra_pb2.Expression, type_dict: Dict[str, Any]
+    ) -> algebra_pb2.Expression:
+        pb_expr = algebra_pb2.Expression()
+        cast = pb_expr.cast
+        cast.input.CopyFrom(input_expr)
+        json_format.ParseDict(type_dict, cast.type)
+        cast.failure_behavior = (
+            algebra_pb2.Expression.Cast.FAILURE_BEHAVIOR_THROW_EXCEPTION
+        )
+        return pb_expr
+
+    def _compile_replace(
+        self,
+        str_expr: algebra_pb2.Expression,
+        search: str,
+        replacement: str,
+    ) -> algebra_pb2.Expression:
+        pb_expr = algebra_pb2.Expression()
+        pb_expr.scalar_function.function_reference = 76  # "replace" or "replace_substring"
+        
+        pb_expr.scalar_function.arguments.add().value.CopyFrom(str_expr)
+        
+        search_expr = algebra_pb2.Expression()
+        search_expr.literal.string = search
+        pb_expr.scalar_function.arguments.add().value.CopyFrom(search_expr)
+        
+        replace_expr = algebra_pb2.Expression()
+        replace_expr.literal.string = replacement
+        pb_expr.scalar_function.arguments.add().value.CopyFrom(replace_expr)
+        
+        return pb_expr
+
     @_compile_op.register(struct_ops.StructOp)
     def _compile_struct_op(
         self,

packages/bigframes/bigframes/session/substrait_executor.py

@@ -115,13 +115,18 @@ def default_for_engine(cls, engine_name: str) -> SubstraitExecutor:
             return cls(
                 AceroSubstraitConsumer(),
                 substrait_compiler.SubstraitCompiler(
-                    duration_type="int", use_precision_types=False
+                    duration_type="int",
+                    use_precision_types=False,
+                    dialect="substrait-acero",
                 ),
             )
         elif engine_name == "datafusion":
             return cls(
                 DataFusionSubstraitConsumer(),
-                substrait_compiler.SubstraitCompiler(duration_type="int"),
+                substrait_compiler.SubstraitCompiler(
+                    duration_type="int",
+                    dialect="substrait-datafusion",
+                ),
             )
         else:
             raise ValueError(f"Unknown engine: {engine_name}")

packages/bigframes/tests/system/small/engines/test_generic_ops.py

@@ -178,10 +178,18 @@ def test_engines_astype_bool(scalars_array_value: array_value.ArrayValue, engine
 )
 def test_engines_astype_string(scalars_array_value: array_value.ArrayValue, engine):
     # floats work slightly different with trailing zeroes rn
+    excluded_cols = ["float64_col"]
+
+    # Acero's Substrait consumer lacks support for string functions like replace/replace_substring
+    # and precision_time, so we cannot format time_col and timestamp_col inside the Substrait plan.
+    from bigframes.session.substrait_executor import SubstraitExecutor
+    if isinstance(engine, SubstraitExecutor) and not engine._compiler._use_precision_types:
+        excluded_cols.extend(["time_col", "timestamp_col"])
+
     arr = apply_op(
         scalars_array_value,
         ops.AsTypeOp(to_type=bigframes.dtypes.STRING_DTYPE),
-        excluded_cols=["float64_col"],
+        excluded_cols=excluded_cols,
     )
 
     assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)