From 6afedef3573d69fe7606e050eae1242f878e3158 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Mon, 7 Sep 2026 23:52:33 +0300 Subject: [PATCH 1/7] 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) From 95a2e9a027469b2d38ac4d70341d105c8170082a Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Sun, 20 Sep 2026 10:00:09 +0300 Subject: [PATCH 2/7] docs: say why a projected read keeps a single-field struct The comment rested on substrait-java doing the same, which says nothing about this code a year from now, and it left the nested case unstated. A relation's schema is a struct, so a mask that selects one field at the top level has nothing to unwrap into; nested levels follow that, and maintain_singular_struct is never read. --- src/substrait/type_inference.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 7451008..cf5ea8b 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -747,10 +747,11 @@ 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. + # The mask selects fields. Where the spec would unwrap a + # single-field selection, this keeps the struct: at this level a + # relation's schema is a struct, and nested levels follow it, so + # maintain_singular_struct is never read. The projection runs + # before the emit mapping below, which indexes its output. struct = _project_read_struct(struct, rel.read.projection.select) elif rel_type == "filter": (common, struct) = ( From f3fe47027fef586d2dccd6acd8990b3098564754 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 22 Sep 2026 10:17:54 +0300 Subject: [PATCH 3/7] fix: keep an emitting read's filters unanchored and a select-less mask whole A read's filter binds against its base schema, so an emit reshapes the read's output as a projection does and the filter keeps offset-based outer references. A mask with no select selects every field, as an absent projection does. --- src/substrait/type_inference.py | 2 +- src/substrait/utils/__init__.py | 10 +++++++--- tests/test_read_projection.py | 7 +++++++ tests/test_utils.py | 10 ++++++---- 4 files changed, 21 insertions(+), 8 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index cf5ea8b..0cf449f 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -746,7 +746,7 @@ 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"): + if rel.read.projection.HasField("select"): # The mask selects fields. Where the spec would unwrap a # single-field selection, this keeps the struct: at this level a # relation's schema is a struct, and nested levels follow it, so diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index adc0d96..9da34b1 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -607,11 +607,15 @@ 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. Join conditions and projected read filters may use a - # different row with no relation to anchor (binding None -> left as-is). + # own output. Join conditions and the filters of a projected or + # emitting read 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") + projected_read = rel_type == "read" and ( + node.HasField("projection") + or node.common.WhichOneof("emit_kind") == "emit" + ) for name, expr in _iter_named_direct_expressions(node): if single_input is not None: binding = single_input diff --git a/tests/test_read_projection.py b/tests/test_read_projection.py index d16f9ab..f92b43a 100644 --- a/tests/test_read_projection.py +++ b/tests/test_read_projection.py @@ -60,6 +60,13 @@ def test_read_without_projection_keeps_the_schema(): assert infer_rel_schema(_read(schema)) == schema +def test_read_projection_without_select_keeps_the_schema(): + schema = _struct(i64(nullable=False), string(), variation=7) + rel = _read(schema) + rel.read.projection.maintain_singular_struct = True + assert infer_rel_schema(rel) == 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)) diff --git a/tests/test_utils.py b/tests/test_utils.py index fe83f74..cb27a2f 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -526,10 +526,12 @@ 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): +@pytest.mark.parametrize("reshaped", [None, "projection", "emit"]) +def test_convert_read_filter_uses_unprojected_scope(filter_field, reshaped): read = _read("o", ncols=3) - if projected: + if reshaped == "emit": + read.read.common.emit.output_mapping.append(0) + elif reshaped == "projection": read.read.projection.CopyFrom( stalg.Expression.MaskExpression( select=stalg.Expression.MaskExpression.StructSelect( @@ -549,7 +551,7 @@ def test_convert_read_filter_uses_unprojected_scope(filter_field, projected): out_read.read, filter_field ).subquery.set_predicate.tuples.filter.condition.selection assert ref.direct_reference.struct_field.field == 2 - if projected: + if reshaped: 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 From cc691656c70d54ae533d59de3d2d540ae27f74fd Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 22 Sep 2026 10:28:55 +0300 Subject: [PATCH 4/7] fix: anchor the filters of a read whose mask has no select A mask with no select keeps every field, so the read's row is its base schema and its filters can anchor to it, as for an unprojected read. --- src/substrait/utils/__init__.py | 2 +- tests/test_utils.py | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index 9da34b1..5169abc 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -613,7 +613,7 @@ def convert_rel(rel, scope): 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") + node.projection.HasField("select") or node.common.WhichOneof("emit_kind") == "emit" ) for name, expr in _iter_named_direct_expressions(node): diff --git a/tests/test_utils.py b/tests/test_utils.py index cb27a2f..6ed7b02 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -526,10 +526,14 @@ def _outer_refs(plan: stplan.Plan): @pytest.mark.parametrize("filter_field", ["filter", "best_effort_filter"]) -@pytest.mark.parametrize("reshaped", [None, "projection", "emit"]) +@pytest.mark.parametrize("reshaped", [None, "no_select", "projection", "emit"]) def test_convert_read_filter_uses_unprojected_scope(filter_field, reshaped): read = _read("o", ncols=3) - if reshaped == "emit": + if reshaped == "no_select": + # A mask with no select keeps every field, so the read's row is unchanged. + read.read.projection.maintain_singular_struct = True + reshaped = None + elif reshaped == "emit": read.read.common.emit.output_mapping.append(0) elif reshaped == "projection": read.read.projection.CopyFrom( From fdce125e08a737644efc62cfc3e9403a4c3e186c Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 22 Sep 2026 11:08:53 +0300 Subject: [PATCH 5/7] fix: drop a projected struct's variation and print a projected read's names A type variation describes the unprojected row, so a pruned struct keeps only its nullability, as in substrait-java and the validator. The plan printer resolves field names in the order the read's mask outputs them. --- src/substrait/type_inference.py | 18 ++++++------ src/substrait/utils/__init__.py | 6 +++- src/substrait/utils/display.py | 24 +++++++++++++++- tests/test_display.py | 51 ++++++++++++++++++++++++++++++++- tests/test_read_projection.py | 10 +++---- 5 files changed, 92 insertions(+), 17 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 0cf449f..fbb15e6 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -687,7 +687,11 @@ def _set_output_struct(op_name: str, inputs: list) -> stt.Type.Struct: 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.""" + """Project fields in mask order, as substrait-java does. + + The result keeps the struct's nullability but not its type variation, which + describes the unprojected row's layout. + """ fields = [] for item in select.struct_items: if not 0 <= item.field < len(struct.types): @@ -699,11 +703,7 @@ def _project_read_struct( 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 + return stt.Type.Struct(types=fields, nullability=struct.nullability) def _project_read_type( @@ -722,8 +722,10 @@ def _project_read_type( 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. + # The spec leaves this open: it unwraps a single-element list selection by + # default and says nothing of map keys. As substrait-java and the validator + # do, a list or map selection keeps the wrapper and its type, and a child + # mask projects the element or the value. selection = getattr(select, kind) if selection.HasField("child"): member = "type" if kind == "list" else "value" diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index 5169abc..a088d48 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -475,6 +475,10 @@ def _plan_has_steps_out(plan: stplan.Plan) -> bool: # condition scope names columns the output drops and has no anchorable relation. _JOIN_COMBINED_SCOPED_FIELDS = frozenset({"expression", "residual_expression"}) +# A read's filters bind against its base schema, before the projection and the +# emit that reshape its output row. +_READ_FILTER_FIELDS = frozenset({"filter", "best_effort_filter"}) + def _is_reducing_join(node) -> bool: """Whether a join relation-variant ``node`` emits only one side (semi/anti), so @@ -621,7 +625,7 @@ def convert_rel(rel, scope): binding = single_input elif reducing and name in _JOIN_COMBINED_SCOPED_FIELDS: binding = None - elif projected_read and name in ("filter", "best_effort_filter"): + elif projected_read and name in _READ_FILTER_FIELDS: binding = None else: binding = rel diff --git a/src/substrait/utils/display.py b/src/substrait/utils/display.py index e1e1841..994eced 100644 --- a/src/substrait/utils/display.py +++ b/src/substrait/utils/display.py @@ -5,10 +5,32 @@ in a readable format using indentation, -> characters, and colors. """ +import itertools + import substrait.algebra_pb2 as stalg import substrait.plan_pb2 as stp import substrait.type_pb2 as stt +from substrait.utils import type_num_names + + +def _read_output_names(read: stalg.ReadRel) -> list: + """The base schema's names in the order the read's projection outputs them. + + Each top-level field owns a block of the depth-first names, and a mask + selects whole blocks in its order; a child mask does not prune a block. + """ + names = list(read.base_schema.names) + if not read.projection.HasField("select"): + return names + lengths = [type_num_names(t) for t in read.base_schema.struct.types] + starts = [0, *itertools.accumulate(lengths)] + return [ + name + for item in read.projection.select.struct_items + for name in names[starts[item.field] : starts[item.field + 1]] + ] + # ANSI color codes class Colors: @@ -202,7 +224,7 @@ def _stream_read_rel(self, read: stalg.ReadRel, stream, depth: int): if read.HasField("base_schema"): # Capture schema names for field resolution - self.schema_names = list(read.base_schema.names) + self.schema_names = _read_output_names(read) if self.show_metadata: stream.write( f"{self._get_indent_with_arrow(depth + 1)}{self._color('schema:', Colors.BLUE)} {self._color(self.schema_names, Colors.YELLOW)}\n" diff --git a/tests/test_display.py b/tests/test_display.py index 2bba08a..4652e05 100644 --- a/tests/test_display.py +++ b/tests/test_display.py @@ -1,8 +1,10 @@ +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.plan import fetch, read_named_table, virtual_table -from substrait.builders.type import boolean, i64 +from substrait.builders.type import boolean, i64, string from substrait.extension_registry import ExtensionRegistry from substrait.utils.display import PlanPrinter @@ -53,3 +55,50 @@ def test_stringify_fetch_unset_offset_and_count(): out = _printer().stringify_plan(plan) assert "fetch: offset=0, count=all" in out + + +def test_stringify_resolves_names_through_a_read_projection(): + schema = stt.NamedStruct( + names=["id", "txt", "flag"], + struct=stt.Type.Struct( + types=[i64(nullable=False), string(), boolean()], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), + ) + read = stalg.ReadRel( + base_schema=schema, + named_table=stalg.ReadRel.NamedTable(names=["t"]), + projection=stalg.Expression.MaskExpression( + select=stalg.Expression.MaskExpression.StructSelect( + struct_items=[stalg.Expression.MaskExpression.StructItem(field=2)] + ), + maintain_singular_struct=True, + ), + ) + condition = 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(), + ) + ) + plan = stp.Plan( + relations=[ + stp.PlanRel( + root=stalg.RelRoot( + input=stalg.Rel( + filter=stalg.FilterRel( + input=stalg.Rel(read=read), condition=condition + ) + ), + names=["flag"], + ) + ) + ] + ) + + out = _printer().stringify_plan(plan) + + assert "field: flag" in out + assert "field: id" not in out diff --git a/tests/test_read_projection.py b/tests/test_read_projection.py index f92b43a..11f92f1 100644 --- a/tests/test_read_projection.py +++ b/tests/test_read_projection.py @@ -49,9 +49,8 @@ def test_read_projection_selects_fields_in_mask_order(fields): 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 - ) + # The variation describes the unprojected row, so the projection drops it. + assert infer_rel_schema(rel) == _struct(*(schema.types[i] for i in fields)) assert rel.SerializeToString() == before @@ -128,7 +127,6 @@ def test_read_projection_keeps_nested_structure_and_root_names(): boolean(nullable=False), i64(nullable=False), nullable=NULL, - variation=8, ) ), i32(), @@ -161,7 +159,7 @@ def test_read_projection_prunes_collection_children(kind): ) expected = stt.Type( list=stt.Type.List( - type=stt.Type(struct=_struct(string(), nullable=NULL, variation=9)), + type=stt.Type(struct=_struct(string(), nullable=NULL)), nullability=NULL, type_variation_reference=10, ) @@ -181,7 +179,7 @@ def test_read_projection_prunes_collection_children(kind): expected = stt.Type( map=stt.Type.Map( key=string(nullable=False), - value=stt.Type(struct=_struct(string(), nullable=NULL, variation=9)), + value=stt.Type(struct=_struct(string(), nullable=NULL)), nullability=NULL, type_variation_reference=10, ) From 28450e85606e1fa95da12f4217958e591d8c9a08 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 22 Sep 2026 11:27:56 +0300 Subject: [PATCH 6/7] fix: skip read mask items outside the schema when printing names The printer degrades on a malformed plan rather than raising, as its field-name lookup already does. --- src/substrait/utils/display.py | 1 + tests/test_display.py | 42 ++++++++++++++++++++++++++-------- 2 files changed, 33 insertions(+), 10 deletions(-) diff --git a/src/substrait/utils/display.py b/src/substrait/utils/display.py index 994eced..b9b76cb 100644 --- a/src/substrait/utils/display.py +++ b/src/substrait/utils/display.py @@ -28,6 +28,7 @@ def _read_output_names(read: stalg.ReadRel) -> list: return [ name for item in read.projection.select.struct_items + if 0 <= item.field < len(lengths) for name in names[starts[item.field] : starts[item.field + 1]] ] diff --git a/tests/test_display.py b/tests/test_display.py index 4652e05..7bb1366 100644 --- a/tests/test_display.py +++ b/tests/test_display.py @@ -57,24 +57,29 @@ def test_stringify_fetch_unset_offset_and_count(): assert "fetch: offset=0, count=all" in out -def test_stringify_resolves_names_through_a_read_projection(): - schema = stt.NamedStruct( - names=["id", "txt", "flag"], - struct=stt.Type.Struct( - types=[i64(nullable=False), string(), boolean()], - nullability=stt.Type.NULLABILITY_REQUIRED, +def _projected_read(*fields: int) -> stalg.ReadRel: + return stalg.ReadRel( + base_schema=stt.NamedStruct( + names=["id", "txt", "flag"], + struct=stt.Type.Struct( + types=[i64(nullable=False), string(), boolean()], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), ), - ) - read = stalg.ReadRel( - base_schema=schema, named_table=stalg.ReadRel.NamedTable(names=["t"]), projection=stalg.Expression.MaskExpression( select=stalg.Expression.MaskExpression.StructSelect( - struct_items=[stalg.Expression.MaskExpression.StructItem(field=2)] + struct_items=[ + stalg.Expression.MaskExpression.StructItem(field=f) for f in fields + ] ), maintain_singular_struct=True, ), ) + + +def test_stringify_resolves_names_through_a_read_projection(): + read = _projected_read(2) condition = stalg.Expression( selection=stalg.Expression.FieldReference( direct_reference=stalg.Expression.ReferenceSegment( @@ -102,3 +107,20 @@ def test_stringify_resolves_names_through_a_read_projection(): assert "field: flag" in out assert "field: id" not in out + + +def test_stringify_tolerates_a_read_mask_out_of_range(): + plan = stp.Plan( + relations=[ + stp.PlanRel( + root=stalg.RelRoot( + input=stalg.Rel(read=_projected_read(5, -1, 2)), names=["flag"] + ) + ) + ] + ) + printer = _printer() + + printer.stringify_plan(plan) + + assert printer.schema_names == ["flag"] From ae7f2991660e514ce0bf26592144d1228a9ea5c5 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 22 Sep 2026 11:49:16 +0300 Subject: [PATCH 7/7] fix: apply a read's emit when printing its names Honoring the mask without the emit mislabels reads the printer used to get right by ignoring both, such as a mask [2, 0] with emit [1]. --- src/substrait/utils/display.py | 25 ++++++++++++-------- tests/test_display.py | 43 ++++++++++++++++++++++------------ 2 files changed, 43 insertions(+), 25 deletions(-) diff --git a/src/substrait/utils/display.py b/src/substrait/utils/display.py index b9b76cb..b15d433 100644 --- a/src/substrait/utils/display.py +++ b/src/substrait/utils/display.py @@ -15,22 +15,27 @@ def _read_output_names(read: stalg.ReadRel) -> list: - """The base schema's names in the order the read's projection outputs them. + """The base schema's names in the order the read outputs them. - Each top-level field owns a block of the depth-first names, and a mask - selects whole blocks in its order; a child mask does not prune a block. + Each top-level field owns a block of the depth-first names. The mask selects + whole blocks in its order, then the emit picks among those; a child mask does + not prune a block. Indices outside the schema are skipped. """ names = list(read.base_schema.names) - if not read.projection.HasField("select"): + projected = read.projection.HasField("select") + emitted = read.common.WhichOneof("emit_kind") == "emit" + if not projected and not emitted: return names lengths = [type_num_names(t) for t in read.base_schema.struct.types] starts = [0, *itertools.accumulate(lengths)] - return [ - name - for item in read.projection.select.struct_items - if 0 <= item.field < len(lengths) - for name in names[starts[item.field] : starts[item.field + 1]] - ] + blocks = [names[starts[i] : starts[i + 1]] for i in range(len(lengths))] + if projected: + fields = [item.field for item in read.projection.select.struct_items] + blocks = [blocks[i] for i in fields if 0 <= i < len(blocks)] + if emitted: + mapping = read.common.emit.output_mapping + blocks = [blocks[i] for i in mapping if 0 <= i < len(blocks)] + return [name for block in blocks for name in block] # ANSI color codes diff --git a/tests/test_display.py b/tests/test_display.py index 7bb1366..de476a9 100644 --- a/tests/test_display.py +++ b/tests/test_display.py @@ -1,3 +1,4 @@ +import pytest import substrait.algebra_pb2 as stalg import substrait.plan_pb2 as stp import substrait.type_pb2 as stt @@ -57,8 +58,8 @@ def test_stringify_fetch_unset_offset_and_count(): assert "fetch: offset=0, count=all" in out -def _projected_read(*fields: int) -> stalg.ReadRel: - return stalg.ReadRel( +def _projected_read(*fields: int, emit=None) -> stalg.ReadRel: + read = stalg.ReadRel( base_schema=stt.NamedStruct( names=["id", "txt", "flag"], struct=stt.Type.Struct( @@ -76,10 +77,26 @@ def _projected_read(*fields: int) -> stalg.ReadRel: maintain_singular_struct=True, ), ) - - -def test_stringify_resolves_names_through_a_read_projection(): - read = _projected_read(2) + if emit is not None: + read.common.emit.output_mapping.extend(emit) + return read + + +# The field a filter above the read names as field 0. +@pytest.mark.parametrize( + "mask, emit, name", + [ + ([2], None, "flag"), + ([2, 0], [1], "id"), + ([2, 0], [0], "flag"), + ([1, 2], [1], "flag"), + (None, [2], "flag"), + ], +) +def test_stringify_resolves_names_through_a_read_projection(mask, emit, name): + read = _projected_read(*(mask or []), emit=emit) + if mask is None: + read.ClearField("projection") condition = stalg.Expression( selection=stalg.Expression.FieldReference( direct_reference=stalg.Expression.ReferenceSegment( @@ -105,22 +122,18 @@ def test_stringify_resolves_names_through_a_read_projection(): out = _printer().stringify_plan(plan) - assert "field: flag" in out - assert "field: id" not in out + assert f"field: {name}\n" in out -def test_stringify_tolerates_a_read_mask_out_of_range(): +def test_stringify_tolerates_read_indices_out_of_range(): + read = _projected_read(5, -1, 2, 0, emit=[3, -1, 1]) plan = stp.Plan( relations=[ - stp.PlanRel( - root=stalg.RelRoot( - input=stalg.Rel(read=_projected_read(5, -1, 2)), names=["flag"] - ) - ) + stp.PlanRel(root=stalg.RelRoot(input=stalg.Rel(read=read), names=["id"])) ] ) printer = _printer() printer.stringify_plan(plan) - assert printer.schema_names == ["flag"] + assert printer.schema_names == ["id"]