fix!: apply ReadRel projection during schema inference

src/substrait/type_inference.py

@@ -684,6 +684,58 @@ 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, 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):
+            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)
+    return stt.Type.Struct(types=fields, nullability=struct.nullability)
+
+
+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:
+        # 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"
+            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 +748,13 @@ 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.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
+            # 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) = (
             rel.filter.common,

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
@@ -509,7 +513,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 +591,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 +611,22 @@ 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 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.projection.HasField("select")
+                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
                 elif reducing and name in _JOIN_COMBINED_SCOPED_FIELDS:
                     binding = None
+                elif projected_read and name in _READ_FILTER_FIELDS:
+                    binding = None
                 else:
                     binding = rel
                 convert_expr(expr, scope, binding)

src/substrait/utils/display.py

@@ -5,10 +5,38 @@
 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 outputs them.
+
+    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)
+    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)]
+    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
 class Colors:
@@ -202,7 +230,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"

tests/test_display.py

@@ -1,8 +1,11 @@
+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.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 +56,84 @@ def test_stringify_fetch_unset_offset_and_count():
     out = _printer().stringify_plan(plan)
 
     assert "fetch: offset=0, count=all" in out
+
+
+def _projected_read(*fields: int, emit=None) -> stalg.ReadRel:
+    read = 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,
+            ),
+        ),
+        named_table=stalg.ReadRel.NamedTable(names=["t"]),
+        projection=stalg.Expression.MaskExpression(
+            select=stalg.Expression.MaskExpression.StructSelect(
+                struct_items=[
+                    stalg.Expression.MaskExpression.StructItem(field=f) for f in fields
+                ]
+            ),
+            maintain_singular_struct=True,
+        ),
+    )
+    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(
+                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 f"field: {name}\n" in out
+
+
+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=read), names=["id"]))
+        ]
+    )
+    printer = _printer()
+
+    printer.stringify_plan(plan)
+
+    assert printer.schema_names == ["id"]