Skip to content

Commit 81ade4d

Browse files
fix dialect passing
1 parent 2350d62 commit 81ade4d

4 files changed

Lines changed: 40 additions & 17 deletions

File tree

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

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -301,7 +301,9 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
301301
if self.dialect == "polars":
302302
return _lower_cast_to_polars(expr.op, expr.inputs[0])
303303
else:
304-
return _lower_cast_to_substrait(expr.op, expr.inputs[0], dialect=self.dialect)
304+
return _lower_cast_to_substrait(
305+
expr.op, expr.inputs[0], dialect=self.dialect
306+
)
305307

306308

307309
def invert_bytes(byte_string):
@@ -476,7 +478,9 @@ def _lower_cast_to_polars(cast_op: ops.AsTypeOp, arg: expression.Expression):
476478
def _lower_cast_to_substrait(
477479
cast_op: ops.AsTypeOp,
478480
arg: expression.Expression,
479-
dialect: Literal["substrait-datafusion", "substrait-acero"] = "substrait-datafusion",
481+
dialect: Literal[
482+
"substrait-datafusion", "substrait-acero"
483+
] = "substrait-datafusion",
480484
):
481485
if arg.output_type == dtypes.BOOL_DTYPE and cast_op.to_type == dtypes.STRING_DTYPE:
482486
is_true_cond = ops.eq_op.as_expr(arg, expression.const(True))
@@ -504,7 +508,9 @@ def _lower_cast_to_substrait(
504508
elif arg.output_type == dtypes.TIMESTAMP_DTYPE:
505509
cast_expr = cast_op.as_expr(arg)
506510
if dialect == "substrait-datafusion":
507-
replaced_t = string_ops.ReplaceStrOp(pat="T", repl=" ").as_expr(cast_expr)
511+
replaced_t = string_ops.ReplaceStrOp(pat="T", repl=" ").as_expr(
512+
cast_expr
513+
)
508514
return string_ops.ReplaceStrOp(pat="Z", repl="+00").as_expr(replaced_t)
509515
else:
510516
# Acero: native cast (excluded in test)
@@ -513,14 +519,16 @@ def _lower_cast_to_substrait(
513519
return cast_op.as_expr(arg)
514520

515521

516-
class LowerEqNullsMatchRule(op_lowering.OpLoweringRule):
522+
class SubstraitLowerEqNullsMatchRule(op_lowering.OpLoweringRule):
517523
@property
518524
def op(self) -> type[ops.ScalarOp]:
519525
return comparison_ops.EqNullsMatchOp
520526

521527
def lower(self, expr: expression.OpExpression) -> expression.Expression:
522528
assert isinstance(expr.op, comparison_ops.EqNullsMatchOp)
523-
arg1, arg2 = _coerce_comparables(expr.children[0], expr.children[1])
529+
arg1, arg2 = _coerce_comparables(
530+
expr.children[0], expr.children[1], dialect="substrait"
531+
)
524532

525533
# True constant
526534
true_const = expression.const(True)
@@ -594,16 +602,19 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
594602
LowerFloorOp(),
595603
)
596604

605+
597606
def lower_ops_to_polars(root: bigframe_node.BigFrameNode) -> bigframe_node.BigFrameNode:
598607
return op_lowering.lower_ops(root, rules=POLARS_LOWERING_RULES)
599608

600609

601610
def lower_ops_to_substrait(
602611
root: bigframe_node.BigFrameNode,
603-
dialect: Literal["substrait-datafusion", "substrait-acero"] = "substrait-datafusion",
612+
dialect: Literal[
613+
"substrait-datafusion", "substrait-acero"
614+
] = "substrait-datafusion",
604615
) -> bigframe_node.BigFrameNode:
605616
rules = (
606-
LowerEqNullsMatchRule(),
617+
SubstraitLowerEqNullsMatchRule(),
607618
*SUBSTRAIT_LOWER_COMPARISONS,
608619
LowerAsTypeRule(dialect=dialect),
609620
)

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -766,7 +766,9 @@ def compile_join(self, node: nodes.JoinNode):
766766
for left_ex, right_ex in node.conditions:
767767
from bigframes.core.compile import lowering
768768

769-
left_ex, right_ex = lowering._coerce_comparables(left_ex, right_ex)
769+
left_ex, right_ex = lowering._coerce_comparables(
770+
left_ex, right_ex, dialect="polars"
771+
)
770772
left_on.append(self.expr_compiler.compile_expression(left_ex))
771773
right_on.append(self.expr_compiler.compile_expression(right_ex))
772774

@@ -787,7 +789,9 @@ def compile_isin(self, node: nodes.InNode):
787789
right_col = ex.ResolvedDerefOp.from_field(node.right_child.fields[0])
788790
from bigframes.core.compile import lowering
789791

790-
left_ex, right_ex = lowering._coerce_comparables(node.left_col, right_col)
792+
left_ex, right_ex = lowering._coerce_comparables(
793+
node.left_col, right_col, dialect="polars"
794+
)
791795

792796
left_pl_ex = self.expr_compiler.compile_expression(left_ex)
793797
right_pl_ex = self.expr_compiler.compile_expression(right_ex)

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

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@
2929
import bigframes.operations.comparison_ops as comparison_ops
3030
import bigframes.operations.generic_ops as generic_ops
3131
import bigframes.operations.numeric_ops as numeric_ops
32-
import bigframes.operations.struct_ops as struct_ops
3332
import bigframes.operations.string_ops as string_ops
33+
import bigframes.operations.struct_ops as struct_ops
3434
from bigframes.core import agg_expressions, bigframe_node, nodes, rewrite
3535
from bigframes.core.compile import lowering
3636

@@ -44,7 +44,9 @@ def __init__(
4444
self,
4545
duration_type: Literal["interval_day", "int"],
4646
use_precision_types: bool = True,
47-
dialect: Literal["substrait-datafusion", "substrait-acero"] = "substrait-datafusion",
47+
dialect: Literal[
48+
"substrait-datafusion", "substrait-acero"
49+
] = "substrait-datafusion",
4850
):
4951
self._duration_type = duration_type
5052
self._use_precision_types = use_precision_types
@@ -911,18 +913,20 @@ def _compile_replace(
911913
replacement: str,
912914
) -> algebra_pb2.Expression:
913915
pb_expr = algebra_pb2.Expression()
914-
pb_expr.scalar_function.function_reference = 76 # "replace" or "replace_substring"
915-
916+
pb_expr.scalar_function.function_reference = (
917+
76 # "replace" or "replace_substring"
918+
)
919+
916920
pb_expr.scalar_function.arguments.add().value.CopyFrom(str_expr)
917-
921+
918922
search_expr = algebra_pb2.Expression()
919923
search_expr.literal.string = search
920924
pb_expr.scalar_function.arguments.add().value.CopyFrom(search_expr)
921-
925+
922926
replace_expr = algebra_pb2.Expression()
923927
replace_expr.literal.string = replacement
924928
pb_expr.scalar_function.arguments.add().value.CopyFrom(replace_expr)
925-
929+
926930
return pb_expr
927931

928932
@_compile_op.register(struct_ops.StructOp)

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,11 @@ def test_engines_astype_string(scalars_array_value: array_value.ArrayValue, engi
183183
# Acero's Substrait consumer lacks support for string functions like replace/replace_substring
184184
# and precision_time, so we cannot format time_col and timestamp_col inside the Substrait plan.
185185
from bigframes.session.substrait_executor import SubstraitExecutor
186-
if isinstance(engine, SubstraitExecutor) and not engine._compiler._use_precision_types:
186+
187+
if (
188+
isinstance(engine, SubstraitExecutor)
189+
and not engine._compiler._use_precision_types
190+
):
187191
excluded_cols.extend(["time_col", "timestamp_col"])
188192

189193
arr = apply_op(

0 commit comments

Comments
 (0)