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)