Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 56 additions & 0 deletions src/substrait/type_inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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,
Expand Down
22 changes: 14 additions & 8 deletions src/substrait/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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":
Expand All @@ -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)
Expand Down
Loading