diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 8c431a6..fbb15e6 100644 --- a/src/substrait/type_inference.py +++ b/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, diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index c41782a..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 @@ -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) diff --git a/src/substrait/utils/display.py b/src/substrait/utils/display.py index e1e1841..b15d433 100644 --- a/src/substrait/utils/display.py +++ b/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" diff --git a/tests/test_display.py b/tests/test_display.py index 2bba08a..de476a9 100644 --- a/tests/test_display.py +++ b/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"] diff --git a/tests/test_read_projection.py b/tests/test_read_projection.py new file mode 100644 index 0000000..11f92f1 --- /dev/null +++ b/tests/test_read_projection.py @@ -0,0 +1,304 @@ +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() + + # 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 + + +def test_read_without_projection_keeps_the_schema(): + schema = _struct(i64(nullable=False), string(), variation=7) + 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)) + 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, + ) + ), + 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)), + 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)), + 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..6ed7b02 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -525,6 +525,65 @@ def _outer_refs(plan: stplan.Plan): ] +@pytest.mark.parametrize("filter_field", ["filter", "best_effort_filter"]) +@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 == "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( + 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 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 + 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)