From 6afedef3573d69fe7606e050eae1242f878e3158 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Mon, 7 Sep 2026 23:52:33 +0300 Subject: [PATCH] fix!: apply ReadRel projection during schema inference `infer_rel_schema` currently returns the full `base_schema` for a masked read. A mask selecting one field can therefore report all input fields, and emit mappings are evaluated against the wrong columns. Apply the read projection before emit, following the [Read Operator in spec v0.99.0](https://github.com/substrait-io/substrait/blob/v0.99.0/site/docs/relations/logical_relations.md#read-operator). Recursively project struct fields, list elements, and map values while preserving their nullability and container types. Keep explicit `RelRoot.names` and the original input schema unchanged. For projected reads, leave correlations inside `filter` and `best_effort_filter` offset-based: [Read Filtering](https://github.com/substrait-io/substrait/blob/v0.99.0/site/docs/relations/logical_relations.md#read-filtering) explicitly places those predicates before projection. Mask unwrapping is a compatibility choice. The [unwrapping section](https://github.com/substrait-io/substrait/blob/v0.99.0/site/docs/expressions/field_references.md#unwrapping-behavior) describes singleton unwrapping by default but leaves its serialization as TBD. This change follows Java's type projector: it preserves structs and containers at every level regardless of `maintain_singular_struct`. For example, selecting only the string from a nested `struct` keeps a nested `struct`, even with the flag false. The validator instead requires the flag for a single-field read projection. Struct fields also follow Java's mask order, including reordered and repeated selections. The same spec page still lists column reordering as an open question, so preserving mask order and duplicates remains a compatibility choice. Closes #264 BREAKING CHANGE: ReadRel projection masks now change inferred output schemas and the indices seen by emit and parent relations. Invalid mask field indices and selector/type mismatches now raise errors instead of being ignored. --- src/substrait/type_inference.py | 56 ++++++ src/substrait/utils/__init__.py | 22 ++- tests/test_read_projection.py | 299 ++++++++++++++++++++++++++++++++ tests/test_utils.py | 53 ++++++ 4 files changed, 422 insertions(+), 8 deletions(-) create mode 100644 tests/test_read_projection.py diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 8c431a6..7451008 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -684,6 +684,56 @@ def _set_output_struct(op_name: str, inputs: list) -> stt.Type.Struct: return stt.Type.Struct(types=types, nullability=primary.nullability) +def _project_read_struct( + struct: stt.Type.Struct, select: stalg.Expression.MaskExpression.StructSelect +) -> stt.Type.Struct: + """Project fields in mask order, matching Java and keeping struct metadata.""" + fields = [] + for item in select.struct_items: + if not 0 <= item.field < len(struct.types): + raise ValueError( + f"Read projection field index {item.field} is out of range " + f"for a struct with {len(struct.types)} fields" + ) + field = struct.types[item.field] + if item.HasField("child"): + field = _project_read_type(field, item.child) + fields.append(field) + result = stt.Type.Struct() + result.CopyFrom(struct) + del result.types[:] + result.types.extend(fields) + return result + + +def _project_read_type( + field: stt.Type, select: stalg.Expression.MaskExpression.Select +) -> stt.Type: + kind = select.WhichOneof("type") + if kind is None: + raise ValueError("Read projection child selection must have a type") + field_kind = field.WhichOneof("kind") + if field_kind != kind: + raise ValueError( + f"Read projection {kind} selection requires a {kind} type, got {field_kind}" + ) + result = stt.Type() + result.CopyFrom(field) + if kind == "struct": + result.struct.CopyFrom(_project_read_struct(field.struct, select.struct)) + else: + # List positions and map keys filter values without changing their type. + # Child masks project the list element or map value, retaining the wrapper. + selection = getattr(select, kind) + if selection.HasField("child"): + member = "type" if kind == "list" else "value" + child_type = getattr(getattr(field, kind), member) + getattr(getattr(result, kind), member).CopyFrom( + _project_read_type(child_type, selection.child) + ) + return result + + def infer_rel_schema(rel: stalg.Rel, *, registry=None, subtrees=()) -> stt.Type.Struct: """Infer a relation's output struct. @@ -696,6 +746,12 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None, subtrees=()) -> stt.Type. if rel_type == "read": (common, struct) = (rel.read.common, rel.read.base_schema.struct) + if rel.read.HasField("projection"): + # Spec v0.99.0 describes default unwrapping but leaves its + # serialization open. Match Java by keeping structs and containers + # at every level, regardless of maintain_singular_struct. + # Apply the projection before the common emit mapping below. + struct = _project_read_struct(struct, rel.read.projection.select) elif rel_type == "filter": (common, struct) = ( rel.filter.common, diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index c41782a..adc0d96 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -509,7 +509,10 @@ def to_id_based_outer_references(plan: stplan.Plan) -> stplan.Plan: This is the shared-subtree / DAG case that offset-based ``steps_out`` cannot address unambiguously. * a ``post_join_filter``, or a leaf host's own filter, exposes the **host's** - output row, so the host is anchored. + output row, so the host is anchored, except for projected reads below. + * a ``ReadRel``'s ``filter`` / ``best_effort_filter`` uses its base schema. + When the read has a projection, its output need not carry that row, so + references into these filters remain offset-based. * a join *condition* / ``residual_expression`` exposes the **combined** left+right row; the join's own output equals that row for a non-reducing join, so the join is anchored. For a *reducing* join (semi/anti) the two differ and no relation @@ -584,11 +587,11 @@ def convert_expr(expr, scope, binding): f"{len(scope)} enclosing query scope(s)" ) target = scope[-steps] - # None marks a combined-inputs scope with no anchorable relation - # (a reducing join's condition). A lateral join's rel_anchor is - # reserved for its right input's left-row reference, so it cannot - # double as the output-row anchor a correlation here would need. - # Both are left offset-based (spec-valid, read by inference). + # None marks a scope with no anchorable relation, such as a + # reducing join's condition or a projected read's filter. + # A lateral join's rel_anchor is reserved for its right input's + # left-row reference, so it cannot double as an output anchor. + # These references stay offset-based. if target is not None and not _binding_is_lateral_join(target): oref.rel_reference = anchor_for(target) elif rex == "subquery": @@ -604,15 +607,18 @@ def convert_rel(rel, scope): if node is not None: # The relation whose output row a subquery here would see one level up: # a single-input host exposes its input; a leaf or multi-input host its - # own output -- except a reducing join's combined-inputs-scoped fields, - # whose scope no relation's output carries (binding None -> left as-is). + # own output. Join conditions and projected read filters may use a + # different row with no relation to anchor (binding None -> left as-is). single_input = _child_rel(*children[0]) if len(children) == 1 else None reducing = single_input is None and _is_reducing_join(node) + projected_read = rel_type == "read" and node.HasField("projection") for name, expr in _iter_named_direct_expressions(node): if single_input is not None: binding = single_input elif reducing and name in _JOIN_COMBINED_SCOPED_FIELDS: binding = None + elif projected_read and name in ("filter", "best_effort_filter"): + binding = None else: binding = rel convert_expr(expr, scope, binding) diff --git a/tests/test_read_projection.py b/tests/test_read_projection.py new file mode 100644 index 0000000..d16f9ab --- /dev/null +++ b/tests/test_read_projection.py @@ -0,0 +1,299 @@ +import pytest +import substrait.algebra_pb2 as stalg +import substrait.plan_pb2 as stp +import substrait.type_pb2 as stt + +from substrait.builders.type import boolean, i32, i64, string +from substrait.type_inference import infer_plan_schema, infer_rel_schema + +MASK = stalg.Expression.MaskExpression +REQ = stt.Type.NULLABILITY_REQUIRED +NULL = stt.Type.NULLABILITY_NULLABLE + + +def _struct(*fields, nullable=REQ, variation=0): + return stt.Type.Struct( + types=fields, nullability=nullable, type_variation_reference=variation + ) + + +def _select(*fields): + return MASK.StructSelect( + struct_items=[ + field + if isinstance(field, MASK.StructItem) + else MASK.StructItem(field=field) + for field in fields + ] + ) + + +def _read(schema, select=None, *, names=(), emit=None): + read = stalg.ReadRel( + base_schema=stt.NamedStruct(names=names, struct=schema), + named_table=stalg.ReadRel.NamedTable(names=["t"]), + ) + if select is not None: + read.projection.CopyFrom(MASK(select=select, maintain_singular_struct=True)) + if emit is not None: + read.common.emit.output_mapping.extend(emit) + read.common.emit.SetInParent() + return stalg.Rel(read=read) + + +@pytest.mark.parametrize("fields", [[], [2], [2, 0], [2, 0, 2]]) +def test_read_projection_selects_fields_in_mask_order(fields): + schema = _struct( + i64(nullable=False), string(), boolean(nullable=False), variation=7 + ) + rel = _read(schema, _select(*fields), names=["id", "text", "flag"]) + before = rel.SerializeToString() + + assert infer_rel_schema(rel) == _struct( + *(schema.types[i] for i in fields), variation=7 + ) + assert rel.SerializeToString() == before + + +def test_read_without_projection_keeps_the_schema(): + schema = _struct(i64(nullable=False), string(), variation=7) + assert infer_rel_schema(_read(schema)) == schema + + +@pytest.mark.parametrize("maintain", [False, True]) +def test_single_field_read_projection_preserves_the_row_struct(maintain): + rel = _read(_struct(i64(), string(), boolean(nullable=False)), _select(2)) + rel.read.projection.maintain_singular_struct = maintain + assert infer_rel_schema(rel) == _struct(boolean(nullable=False)) + + +@pytest.mark.parametrize("maintain", [False, True]) +def test_single_field_read_projection_preserves_nested_structs(maintain): + rel = _read( + _struct(i64(), stt.Type(struct=_struct(i64(), string(), nullable=NULL))), + _select(MASK.StructItem(field=1, child=MASK.Select(struct=_select(1)))), + ) + rel.read.projection.maintain_singular_struct = maintain + + assert infer_rel_schema(rel) == _struct( + stt.Type(struct=_struct(string(), nullable=NULL)) + ) + + +def test_read_projection_precedes_emit(): + schema = _struct(i64(nullable=False), string(), boolean(nullable=False)) + rel = _read(schema, _select(2, 0), emit=[1, 0, 1]) + assert infer_rel_schema(rel) == _struct( + i64(nullable=False), boolean(nullable=False), i64(nullable=False) + ) + + +def test_read_emit_cannot_index_a_field_removed_by_projection(): + rel = _read(_struct(i64(), string(), boolean()), _select(2), emit=[1]) + with pytest.raises(IndexError): + infer_rel_schema(rel) + + +def test_read_projection_keeps_nested_structure_and_root_names(): + inner = _struct( + i64(nullable=False), + string(), + boolean(nullable=False), + nullable=NULL, + variation=8, + ) + rel = _read( + _struct(i32(), stt.Type(struct=inner), string()), + _select(MASK.StructItem(field=1, child=MASK.Select(struct=_select(2, 0))), 0), + names=["unused", "original_struct", "id", "text", "flag", "other"], + ) + root_names = ["renamed_struct", "renamed_flag", "renamed_id", "renamed_scalar"] + plan = stp.Plan( + relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=root_names))] + ) + before = plan.SerializeToString() + + assert infer_plan_schema(plan) == stt.NamedStruct( + names=root_names, + struct=_struct( + stt.Type( + struct=_struct( + boolean(nullable=False), + i64(nullable=False), + nullable=NULL, + variation=8, + ) + ), + i32(), + ), + ) + assert plan.SerializeToString() == before + + +@pytest.mark.parametrize("kind", ["list", "map"]) +def test_read_projection_prunes_collection_children(kind): + element = stt.Type( + struct=_struct(i64(nullable=False), string(), nullable=NULL, variation=9) + ) + child = MASK.Select(struct=_select(1)) + if kind == "list": + field = stt.Type( + list=stt.Type.List( + type=element, nullability=NULL, type_variation_reference=10 + ) + ) + selection = MASK.Select( + list=MASK.ListSelect( + selection=[ + MASK.ListSelect.ListSelectItem( + item=MASK.ListSelect.ListSelectItem.ListElement(field=0) + ) + ], + child=child, + ) + ) + expected = stt.Type( + list=stt.Type.List( + type=stt.Type(struct=_struct(string(), nullable=NULL, variation=9)), + nullability=NULL, + type_variation_reference=10, + ) + ) + else: + field = stt.Type( + map=stt.Type.Map( + key=string(nullable=False), + value=element, + nullability=NULL, + type_variation_reference=10, + ) + ) + selection = MASK.Select( + map=MASK.MapSelect(key=MASK.MapSelect.MapKey(map_key="k"), child=child) + ) + expected = stt.Type( + map=stt.Type.Map( + key=string(nullable=False), + value=stt.Type(struct=_struct(string(), nullable=NULL, variation=9)), + nullability=NULL, + type_variation_reference=10, + ) + ) + rel = _read(_struct(field), _select(MASK.StructItem(field=0, child=selection))) + before = rel.SerializeToString() + assert infer_rel_schema(rel) == _struct(expected) + assert rel.SerializeToString() == before + + +def test_read_projection_recurses_through_nested_collections(): + # list>>> -> the same wrappers, struct. + def wrapped(inner): + return stt.Type( + list=stt.Type.List( + type=stt.Type( + map=stt.Type.Map( + key=string(nullable=False), + value=stt.Type(list=stt.Type.List(type=inner, nullability=REQ)), + nullability=NULL, + ) + ), + nullability=NULL, + ) + ) + + selection = MASK.Select( + list=MASK.ListSelect( + child=MASK.Select( + map=MASK.MapSelect( + child=MASK.Select( + list=MASK.ListSelect( + child=MASK.Select(struct=_select(1)), + ) + ) + ), + ) + ) + ) + rel = _read( + _struct(wrapped(stt.Type(struct=_struct(i64(), string())))), + _select(MASK.StructItem(field=0, child=selection)), + ) + assert infer_rel_schema(rel) == _struct(wrapped(stt.Type(struct=_struct(string())))) + + +@pytest.mark.parametrize("kind", ["list", "map"]) +def test_read_collection_selection_without_child_keeps_its_type(kind): + if kind == "list": + field = stt.Type(list=stt.Type.List(type=i64(), nullability=NULL)) + child = MASK.Select( + list=MASK.ListSelect( + selection=[ + MASK.ListSelect.ListSelectItem( + slice=MASK.ListSelect.ListSelectItem.ListSlice(start=1, end=3) + ) + ] + ) + ) + else: + field = stt.Type( + map=stt.Type.Map(key=string(nullable=False), value=i64(), nullability=REQ) + ) + child = MASK.Select( + map=MASK.MapSelect( + expression=MASK.MapSelect.MapKeyExpression(map_key_expression="k*") + ) + ) + rel = _read( + _struct(string(), field), _select(MASK.StructItem(field=1, child=child)) + ) + assert infer_rel_schema(rel) == _struct(field) + + +@pytest.mark.parametrize("index", [-1, 2]) +@pytest.mark.parametrize("nested", [False, True]) +def test_read_projection_rejects_invalid_struct_indices(index, nested): + schema = _struct(i64(), string()) + select = _select(index) + if nested: + schema = _struct(stt.Type(struct=schema)) + select = _select(MASK.StructItem(field=0, child=MASK.Select(struct=select))) + with pytest.raises(ValueError, match=f"field index {index}"): + infer_rel_schema(_read(schema, select)) + + +@pytest.mark.parametrize( + "child", + [ + MASK.Select(struct=_select(0)), + MASK.Select(list=MASK.ListSelect()), + MASK.Select(map=MASK.MapSelect()), + MASK.Select(), + ], +) +def test_read_projection_rejects_inapplicable_child_selection(child): + rel = _read(_struct(i64()), _select(MASK.StructItem(field=0, child=child))) + with pytest.raises(ValueError, match="Read projection"): + infer_rel_schema(rel) + + +def test_project_above_masked_read_uses_projected_indices(): + rel = _read(_struct(i64(nullable=False), string(), boolean()), _select(1, 0)) + project = stalg.Rel( + project=stalg.ProjectRel( + input=rel, + expressions=[ + stalg.Expression( + selection=stalg.Expression.FieldReference( + direct_reference=stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField( + field=0 + ) + ), + root_reference=stalg.Expression.FieldReference.RootReference(), + ) + ) + ], + common=stalg.RelCommon(emit=stalg.RelCommon.Emit(output_mapping=[2])), + ) + ) + assert infer_rel_schema(project) == _struct(string()) diff --git a/tests/test_utils.py b/tests/test_utils.py index f56e21b..fe83f74 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -525,6 +525,59 @@ def _outer_refs(plan: stplan.Plan): ] +@pytest.mark.parametrize("filter_field", ["filter", "best_effort_filter"]) +@pytest.mark.parametrize("projected", [False, True]) +def test_convert_read_filter_uses_unprojected_scope(filter_field, projected): + read = _read("o", ncols=3) + if projected: + read.read.projection.CopyFrom( + stalg.Expression.MaskExpression( + select=stalg.Expression.MaskExpression.StructSelect( + struct_items=[stalg.Expression.MaskExpression.StructItem(field=0)] + ), + maintain_singular_struct=True, + ) + ) + getattr(read.read, filter_field).CopyFrom( + _exists(_filter(_read("i"), _outer(1, field=2))) + ) + plan = _plan(read) + before = plan.SerializeToString() + out = to_id_based_outer_references(plan) + out_read = out.relations[-1].root.input + ref = getattr( + out_read.read, filter_field + ).subquery.set_predicate.tuples.filter.condition.selection + assert ref.direct_reference.struct_field.field == 2 + if projected: + assert rel_anchor_of(out_read) is None + assert ref.outer_reference.WhichOneof("outer_reference_type") == "steps_out" + assert ref.outer_reference.steps_out == 1 + else: + assert ref.outer_reference.WhichOneof("outer_reference_type") == "rel_reference" + assert ref.outer_reference.rel_reference == rel_anchor_of(out_read) + assert plan.SerializeToString() == before + + +def test_convert_filter_above_projected_read_anchors_read_output(): + read = _read("o", ncols=3) + read.read.projection.CopyFrom( + stalg.Expression.MaskExpression( + select=stalg.Expression.MaskExpression.StructSelect( + struct_items=[stalg.Expression.MaskExpression.StructItem(field=2)] + ), + maintain_singular_struct=True, + ) + ) + out = to_id_based_outer_references( + _plan(_filter(read, _exists(_filter(_read("i"), _outer(1))))) + ) + host = out.relations[-1].root.input.filter + ref = host.condition.subquery.set_predicate.tuples.filter.condition.selection.outer_reference + assert ref.WhichOneof("outer_reference_type") == "rel_reference" + assert ref.rel_reference == rel_anchor_of(host.input) + + def test_convert_correlated_exists_stamps_anchor_and_rewrites(): plan = _plan(_filter(_read("o"), _exists(_filter(_read("i"), _outer(1))))) out = to_id_based_outer_references(plan)