feat: Add isin local execution to hybrid engine

bigframes/core/compile/polars/compiler.py

@@ -513,6 +513,30 @@ def compile_join(self, node: nodes.JoinNode):
  left, right, node.type, left_on, right_on, node.joins_nulls
  )
 
+ @compile_node.register
+ def compile_isin(self, node: nodes.InNode):
+ left = self.compile_node(node.left_child)
+ right = self.compile_node(node.right_child).unique(node.right_col.id.sql)
+ right = right.with_columns(pl.lit(True).alias(node.indicator_col.sql))
+
+ left_ex, right_ex = lowering._coerce_comparables(node.left_col, node.right_col)
+
+ left_pl_ex = self.expr_compiler.compile_expression(left_ex)
+ right_pl_ex = self.expr_compiler.compile_expression(right_ex)
+
+ joined = left.join(
+ right,
+ how="left",
+ left_on=left_pl_ex,
+ right_on=right_pl_ex,
+ # Note: join_nulls renamed to nulls_equal for polars 1.24
+ join_nulls=node.joins_nulls, # type: ignore
+ coalesce=False,
+ )
+ passthrough = [pl.col(id) for id in left.columns]
+ indicator = pl.col(node.indicator_col.sql).fill_null(False)
+ return joined.select((*passthrough, indicator))
+
  def _ordered_join(
  self,
  left_frame: pl.LazyFrame,

bigframes/core/rewrite/schema_binding.py

@@ -65,6 +65,16 @@ def bind_schema_to_node(
  node,
  conditions=conditions,
  )
+ if isinstance(node, nodes.InNode):
+ return dataclasses.replace(
+ node,
+ left_col=ex.ResolvedDerefOp.from_field(
+ node.left_child.field_by_id[node.left_col.id]
+ ),
+ right_col=ex.ResolvedDerefOp.from_field(
+ node.right_child.field_by_id[node.right_col.id]
+ ),
+ )
 
  if isinstance(node, nodes.AggregateNode):
  aggregations = []

bigframes/session/polars_executor.py

@@ -39,6 +39,7 @@
  nodes.FilterNode,
  nodes.ConcatNode,
  nodes.JoinNode,
+ nodes.InNode,
 )
 
 _COMPATIBLE_SCALAR_OPS = (

tests/system/small/engines/test_join.py

@@ -88,3 +88,22 @@ def test_engines_cross_join(
  result, _ = scalars_array_value.relational_join(scalars_array_value, type="cross")
 
  assert_equivalence_execution(result.node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize(
+ ("left_key", "right_key"),
+ [
+ ("int64_col", "float64_col"),
+ ("float64_col", "int64_col"),
+ ("int64_too", "int64_col"),
+ ],
+)
+def test_engines_isin(
+ scalars_array_value: array_value.ArrayValue, engine, left_key, right_key
+):
+ result, _ = scalars_array_value.isin(
+ scalars_array_value, lcol=left_key, rcol=right_key
+ )
+
+ assert_equivalence_execution(result.node, REFERENCE_ENGINE, engine)