diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index 519217f..2d997b0 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -29,6 +29,7 @@ _join_output_struct, _join_struct_from_schemas, _outer_anchor_binding, + infer_expression_type, infer_plan_schema, join_output_names, ) @@ -451,6 +452,44 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: return build_scoped(resolve) +def _require_boolean_condition( + bound_expression: stee.ExtendedExpression, + schema: stt.NamedStruct, + registry: ExtensionRegistry, + *, + context: str, + hint: str = "", +) -> None: + """Reject a ``filter``/join condition that does not infer to ``boolean``. + + Substrait requires a ``FilterRel``/``JoinRel`` condition to be a boolean + predicate, but ``resolve_expression`` will happily bind any expression -- + including a bare column reference. Such a condition builds a legit-looking + plan that behaves very differently at execution (a non-boolean join + condition degrades to a Cartesian product), so we catch it at build time. + + ``context`` names the relation for the message; ``hint`` adds a caller-specific + suggestion (e.g. how to spell an equi-join). + """ + condition = bound_expression.referred_expr[0].expression + cond_type = infer_expression_type(condition, schema.struct, registry=registry) + kind = cond_type.WhichOneof("kind") + if kind != "bool": + message = ( + f"{context} condition must be a boolean predicate, but got type {kind!r}." + ) + if hint: + message = f"{message} {hint}" + raise ValueError(message) + + +_JOIN_CONDITION_HINT = ( + "A bare column name is a column reference, not a match key -- it does not " + "join on that column. Use a predicate such as col('a') == col('b'); for a " + "Cartesian product use cross_join()." +) + + def filter( plan: PlanOrUnbound, expression: ExtendedExpressionOrUnbound, @@ -462,6 +501,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_expression: stee.ExtendedExpression = resolve_expression( expression, ns, registry ) + _require_boolean_condition(bound_expression, ns, registry, context="filter") return _plan_from( [bound_plan], @@ -635,6 +675,9 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_expression: stee.ExtendedExpression = resolve_expression( expression, ns, registry ) + _require_boolean_condition( + bound_expression, ns, registry, context="join", hint=_JOIN_CONDITION_HINT + ) # The output names must match the columns the join type actually emits # (semi/anti drop a side, mark appends a boolean). @@ -729,6 +772,14 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: if expression is not None else None ) + if bound_expression is not None: + _require_boolean_condition( + bound_expression, + ns, + registry, + context="lateral join", + hint=_JOIN_CONDITION_HINT, + ) # Output names/columns follow the same per-join-type shape as a # regular join (semi/anti drop the right side, mark appends a boolean). @@ -1216,6 +1267,13 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: names=list(left_ns.names) + list(right_ns.names), ) bound_expression = resolve_expression(expression, ns, registry) + _require_boolean_condition( + bound_expression, + ns, + registry, + context="nested loop join", + hint=_JOIN_CONDITION_HINT, + ) out_names = join_output_names( stalg.NestedLoopJoinRel.JoinType.Name(type), diff --git a/tests/builders/plan/test_filter.py b/tests/builders/plan/test_filter.py index 1b7ce1d..8681ffa 100644 --- a/tests/builders/plan/test_filter.py +++ b/tests/builders/plan/test_filter.py @@ -1,8 +1,9 @@ +import pytest import substrait.algebra_pb2 as stalg import substrait.plan_pb2 as stp import substrait.type_pb2 as stt -from substrait.builders.extended_expression import literal +from substrait.builders.extended_expression import column, literal from substrait.builders.plan import default_version, filter, read_named_table from substrait.builders.type import boolean, i64 from substrait.extension_registry import ExtensionRegistry @@ -43,3 +44,19 @@ def test_filter(): ) assert actual == expected + + +def test_filter_non_boolean_condition_raises(): + # A filter condition must be a boolean predicate; a non-boolean expression + # (here the i64 column `id`) is rejected at build time. + table = read_named_table("table", named_struct) + with pytest.raises(ValueError, match="boolean predicate"): + filter(table, column("id"))(registry) + + +def test_filter_bare_boolean_column_is_allowed(): + # Unlike a join, a bare boolean column is the *normal* filter condition + # ("keep rows where this flag is true"), so it must not be rejected. + table = read_named_table("table", named_struct) + plan = filter(table, column("is_applicable"))(registry) + assert plan.relations # builds without error diff --git a/tests/builders/plan/test_join.py b/tests/builders/plan/test_join.py index a044d17..e707a65 100644 --- a/tests/builders/plan/test_join.py +++ b/tests/builders/plan/test_join.py @@ -71,6 +71,28 @@ def test_join(): assert actual == expected +def test_join_non_boolean_condition_raises(): + # A join condition must be a boolean predicate. A bare column reference (the + # pandas ``on="key"`` idiom) binds to a non-boolean column and would otherwise + # build a legit-looking plan that degrades to a Cartesian product at execution. + table = read_named_table("table", named_struct) + table2 = read_named_table("table2", named_struct_2) + with pytest.raises(ValueError, match="boolean predicate"): + join(table, table2, column("id"), stalg.JoinRel.JOIN_TYPE_INNER)(registry) + + +def test_join_boolean_column_condition_is_allowed(): + # A boolean column is a valid (if unusual) join condition -- a filtered + # Cartesian product, expressible in SQL as ``JOIN ... ON `` -- so it + # must not be rejected by the boolean-predicate check. + table = read_named_table("table", named_struct) + table2 = read_named_table("table2", named_struct_2) + plan = join(table, table2, column("is_applicable"), stalg.JoinRel.JOIN_TYPE_INNER)( + registry + ) + assert plan.relations # builds without error + + def _post_field(plan): ref = plan.relations[-1].root.input.join.post_join_filter.selection return ref.direct_reference.struct_field.field diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index a46fbb5..0ecc5d3 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -184,6 +184,16 @@ def test_join_unknown_type_raises(): left.join(right, on=sub.col("x") == sub.col("x"), how="banana") +def test_join_on_bare_column_name_raises(): + # `on` is a boolean match predicate, not a key name. The pandas `on="key"` + # idiom binds to a bare (non-boolean) column reference, which would silently + # build a Cartesian product; it is rejected at build time instead. + left = sub.read_named_table("customers", {"cust_id": sub.i64, "name": sub.string}) + right = sub.read_named_table("orders", {"order_id": sub.i64, "cust_ref": sub.i64}) + with pytest.raises(ValueError, match="boolean predicate"): + left.join(right, on="cust_id", how="inner").to_plan() + + @pytest.mark.parametrize("how, join_type", sorted(_JOIN_TYPES.items())) def test_join_all_types_match_builder(how, join_type): left_ns = named_struct(