diff --git a/README.md b/README.md index 775ed5b..afaa3eb 100644 --- a/README.md +++ b/README.md @@ -251,6 +251,30 @@ relations { } ``` +## Deeply nested expressions +Protobuf parsers, including the one in the Python `protobuf` library, reject messages nested more +than 100 levels deep, and each expression level takes several message levels. Substrait works +around this with [detached expressions](https://substrait.io/serialization/binary_serialization/#bounded-expression-nesting): +deep subtrees are moved to a table on the `PlanRel` or `ExtendedExpression` and referenced by +ordinal. + +```python +from substrait.proto import Plan +from substrait.utils.detached import detach, resolve_detached + +data = detach(plan).SerializeToString() # parses under the default limits +plan = resolve_detached(Plan.FromString(data)) # ordinary expression trees again +``` + +`detach` returns a plan with nothing to detach as is. `infer_plan_schema`, +`infer_extended_expression_schema`, the plan printer, and the builders resolve detached expressions +on their own. Functions that take a bare `Rel` or `Expression` cannot, since the table lives on the +containing message, so resolve the `Plan` or `ExtendedExpression` first. + +The builders work on the inlined tree, and the `protobuf` library copies messages passed to a +constructor through its parser, so building on a detached plan works only while the plan still fits +the limit once inlined. Detach the finished plan instead, just before serializing it. + ## Produce a Substrait Plan with Ibis Let's use an existing Substrait producer, [Ibis](https://ibis-project.org), to provide an example using Python Substrait as the consumer. diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index ce0bf9b..74de4ba 100644 --- a/src/substrait/builders/extended_expression.py +++ b/src/substrait/builders/extended_expression.py @@ -26,6 +26,7 @@ top_level_names, type_num_names, ) +from substrait.utils.detached import resolve_for_building # Monotonic source of unique RelCommon.rel_anchor values within a single build. # Not reset between builds by default; reset at each top-level materialization @@ -162,10 +163,13 @@ def resolve_expression( whichever build produced it, so the collector re-derives them from the durable ``(urn, name)`` identities and the expression is rewritten to match -- the expression-level counterpart of ``builders.plan._bind``. Unchanged when the - numbering already agrees, as it does for anything this build resolved. + numbering already agrees, as it does for anything this build resolved. Its + detached subtrees are inlined, since their tables do not travel with the + expression into whatever is built from it. """ if not isinstance(expression, stee.ExtendedExpression): - return expression(base_schema, registry) + return resolve_for_building(expression(base_schema, registry)) + expression = resolve_for_building(expression) collector = current_collector() if collector is None: return expression @@ -1077,6 +1081,9 @@ def _inner_rel(query, registry: ExtensionRegistry, base_schema) -> stalg.Rel: plan = query(registry) if callable(query) else query finally: outer_schemas.reset(token) + # Only the Rel travels into the Expression.Subquery; the PlanRel's detached + # tables stay behind with the plan, so inline its ordinals first. + plan = resolve_for_building(plan) # Fold the inner plan's extensions into this build before lifting its Rel: only # the Rel travels into the Expression.Subquery, so the declarations its function # references name stay behind with the plan. A plan built elsewhere therefore diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index 6525fd2..bfb2cc9 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -40,6 +40,7 @@ remap_function_references, type_num_names, ) +from substrait.utils.detached import resolve_for_building from substrait.version import substrait_version UnboundPlan = Callable[[ExtensionRegistry], stp.Plan] @@ -73,9 +74,11 @@ def _bind(plan: PlanOrUnbound, registry: ExtensionRegistry) -> stp.Plan: resolver writes them). Every builder binds its inputs through here, so this is the single point at - which a foreign plan's anchor space is reconciled with ours. + which a foreign plan's anchor space is reconciled with ours. It is also where a + foreign plan's detached subtrees are inlined: their ordinals index tables on + the plan's own ``PlanRel``, which does not survive being built upon. """ - bound = plan if isinstance(plan, stp.Plan) else plan(registry) + bound = resolve_for_building(plan if isinstance(plan, stp.Plan) else plan(registry)) collector = current_collector() if collector is None: return bound diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index b5a071f..b19eaef 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -7,6 +7,7 @@ import substrait.type_pb2 as stt from substrait.utils import child_rels, iter_plan_rels, plan_subtrees, rel_anchor_of +from substrait.utils.detached import resolve_present_tables class _SubtreeScope: @@ -796,6 +797,13 @@ def infer_expression_type( return stt.Type(string=ecv.current_timezone) else: raise Exception(f"Unknown execution context variable {variable}") + elif rex_type == "detached_expression_ordinal": + raise ValueError( + f"detached expression ordinal {expression.detached_expression_ordinal} " + "can only be resolved against its containing PlanRel or " + "ExtendedExpression; call substrait.utils.detached.resolve_detached " + "on that message first" + ) else: raise Exception(f"Unknown rex_type {rex_type}") @@ -803,6 +811,7 @@ def infer_expression_type( def infer_extended_expression_schema( ee: stee.ExtendedExpression, *, registry=None ) -> stt.Type.Struct: + ee = resolve_present_tables(ee) exprs = [e for e in ee.referred_expr] types = [ @@ -1394,6 +1403,7 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None, subtrees=()) -> stt.Type. def infer_plan_schema(plan: stp.Plan, *, registry=None) -> stt.NamedStruct: + plan = resolve_present_tables(plan) # A Plan carries its shared subtrees in-band as the leading ``rel`` entries of # ``relations`` (the query root is the trailing ``root`` entry), so a # ReferenceRel anywhere in the tree resolves against them by ordinal. Wrap them diff --git a/src/substrait/utils/detached.py b/src/substrait/utils/detached.py new file mode 100644 index 0000000..6d47a89 --- /dev/null +++ b/src/substrait/utils/detached.py @@ -0,0 +1,431 @@ +"""Reading and writing detached subtrees (spec v0.105.0). + +Protobuf parsers cap how deeply messages may nest: upb (the Python runtime), +protobuf-java and C++ all reject a message nested more than 100 levels below the +one being parsed. Every expression level costs several message levels -- a nested +struct is ``Expression`` -> ``Nested`` -> ``Struct`` -> ``Expression`` -- so a +plan with an expression only ~33 levels deep can already be unreadable. + +The spec's workaround is `bounded expression nesting`_: a producer moves subtrees +into a table on the containing ``PlanRel`` or ``ExtendedExpression`` -- expressions +into ``detached_expressions`` -- and leaves an ordinal reference in their place +(``detached_expression_ordinal``). Detachment is encoding-only; resolving the +ordinals recovers the same logical tree. + +* :func:`detach` is the producer side: it cuts expressions that would otherwise + exceed a nesting budget, so the serialized message parses under the default + limits. A message within budget is returned as is. +* :func:`resolve_detached` is the consumer side: it inlines every ordinal, + validating the spec's rules, so the rest of this library -- schema inference, + display, the builders -- sees ordinary trees. + +Both walk the message through its protobuf descriptor, so every relation and +expression kind is covered, and both are iterative, so trees deeper than Python's +recursion limit are fine. Neither looks inside a ``google.protobuf.Any``: its +payload is opaque bytes here. + +.. _bounded expression nesting: + https://substrait.io/serialization/binary_serialization/#bounded-expression-nesting +""" + +from dataclasses import dataclass + +import substrait.algebra_pb2 as stalg +import substrait.extended_expression_pb2 as stee +import substrait.plan_pb2 as stplan +import substrait.type_pb2 as stt +from google.protobuf.message import Message + +# The default parse recursion limit of protobuf-java, C++ and upb: the deepest +# submessage may sit this many levels below the message being parsed. +DEFAULT_MAX_DEPTH = 100 + +_ANY = "google.protobuf.Any" + + +@dataclass(frozen=True) +class _Table: + """One kind of detachable subtree: its message type, the oneof member that + holds an ordinal in place of a subtree, and the container field holding the + detached entries. Each table is resolved on its own: an ordinal of one kind + only ever indexes the table of that kind.""" + + label: str + message: type + oneof: str + ordinal: str + field: str + + @property + def full_name(self) -> str: + return self.message.DESCRIPTOR.full_name + + @property + def tag(self) -> bytes: + """The ordinal field's wire tag (a varint, wire type 0).""" + value = self.message.DESCRIPTOR.fields_by_name[self.ordinal].number << 3 + out = bytearray() + while value > 0x7F: + out.append(value & 0x7F | 0x80) + value >>= 7 + out.append(value) + return bytes(out) + + +_EXPRESSIONS = _Table( + "expression", + stalg.Expression, + "rex_type", + "detached_expression_ordinal", + "detached_expressions", +) +_TABLES = (_EXPRESSIONS,) +_TABLE_BY_MESSAGE = {table.full_name: table for table in _TABLES} +_TABLE_FIELDS = frozenset(table.field for table in _TABLES) +_TAGS = tuple(table.tag for table in _TABLES) + +# Stands in for a resolved entry that has to keep its position. Nothing can +# validly reference it, so its content only has to be a valid expression. +_PLACEHOLDER = stalg.Expression( + literal=stalg.Expression.Literal( + null=stt.Type(bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_NULLABLE)) + ) +) + + +def _declaration_index(item) -> int: + return item[0].index + + +def _children(msg: Message, *, skip_tables: bool = False): + """Yield ``(child, levels)`` for each set submessage of ``msg``. + + ``levels`` is how many message levels below ``msg`` the child sits on the + wire: 1, or 2 for a map value (the map entry is a message of its own). + ``google.protobuf.Any`` payloads are opaque and not descended into, and + ``skip_tables`` leaves out a container's own detached tables. + + Fields come in declaration order, not field-number order (which + ``ListFields`` uses), so the cuts :func:`detach` makes are numbered the way + substrait-java numbers them. + """ + if msg.DESCRIPTOR.full_name == _ANY: + return + for field, value in sorted(msg.ListFields(), key=_declaration_index): + if field.message_type is None or (skip_tables and field.name in _TABLE_FIELDS): + continue + if isinstance(value, Message): + yield value, 1 + elif hasattr(value, "values"): + # A map (repeated fields have no .values()); only message-valued maps + # hold submessages. + for item in value.values(): + if isinstance(item, Message): + yield item, 2 + else: + for item in value: + yield item, 1 + + +def _find(root: Message, predicate, *, skip_tables: bool = False) -> bool: + """Whether ``predicate`` holds for ``root`` or any submessage below it.""" + stack = [root] + while stack: + msg = stack.pop() + if predicate(msg): + return True + stack.extend( + child + for child, _ in _children(msg, skip_tables=skip_tables and msg is root) + ) + return False + + +def _containers(msg): + """The ``(container, depth)`` pairs that own detached tables in ``msg``.""" + if isinstance(msg, stplan.Plan): + return [(pr, 1) for pr in msg.relations] + if isinstance(msg, (stplan.PlanRel, stee.ExtendedExpression)): + return [(msg, 0)] + raise TypeError( + "detached subtrees live on a Plan, PlanRel or ExtendedExpression, " + f"not a {type(msg).__name__}" + ) + + +def _copy(msg): + out = type(msg)() + out.CopyFrom(msg) + return out + + +def _ordinal_table(msg: Message): + """The table ``msg`` refers into, if it is an ordinal reference, else None.""" + table = _TABLE_BY_MESSAGE.get(msg.DESCRIPTOR.full_name) + if table is not None and msg.WhichOneof(table.oneof) == table.ordinal: + return table + return None + + +def _is_any(msg: Message) -> bool: + return msg.DESCRIPTOR.full_name == _ANY + + +def _holds_ordinals(container) -> bool: + """Whether ``container``'s relation or referred expressions hold an ordinal. + + Screened on the wire bytes first: serialization is canonical, so if no + ordinal field's tag occurs in them there is no ordinal, and that is decided + in native code. Only when a tag's bytes do occur -- an ordinal, or the same + bytes inside some other value -- does the descriptor walk settle it. + """ + data = container.SerializeToString() + if not any(tag in data for tag in _TAGS): + return False + return _find(container, _ordinal_table, skip_tables=True) + + +def resolve_detached(msg): + """``msg`` with every detached ordinal inlined. + + ``msg`` is a ``Plan`` (each of its ``PlanRel`` is resolved against its own + tables), a ``PlanRel`` or an ``ExtendedExpression``. Resolution starts from + the relation or referred expressions and follows ordinals transitively, + enforcing the spec's rules on that reachable portion; a violation raises + ``ValueError``: + + 1. an ordinal indexes its own container's table of its kind; + 2. each entry is referenced at most once; + 3. references are acyclic; + 4. an entry is not itself an ordinal (its children may be). + + Unreachable entries are neither resolved nor validated. + + A container's tables are then cleared -- unless it carries a + ``google.protobuf.Any`` payload. An extension payload may reference entries + the relation tree does not, and since it is opaque here its references + cannot be followed, so those entries keep their ordinals, as the spec + requires. The entries resolved here are a different matter: each entry is + referenced at most once, so no payload can validly refer to one of them. The + trailing run of them is dropped, which renumbers nothing, and any left in + place are replaced by a null literal, so none still points into that run. + + ``msg`` itself is returned when it holds no ordinals, which is the common + case, so callers need not check first. Otherwise the result is a copy. + """ + if not any(_holds_ordinals(container) for container, _ in _containers(msg)): + return msg + out = _copy(msg) + for container, _ in _containers(out): + _resolve_in_place(container) + return out + + +def _has_entries(msg) -> bool: + return any( + getattr(container, table.field) + for container, _ in _containers(msg) + for table in _TABLES + ) + + +def resolve_present_tables(msg): + """:func:`resolve_detached`, skipped when ``msg`` has no detached entries. + + For this library's own hot paths: the builders and schema inference run it + on every step, on messages that almost never carry a table. Skipping the + check there leaves an ordinal with no table unvalidated, but nothing can + consume one: schema inference rejects it. + """ + return resolve_detached(msg) if _has_entries(msg) else msg + + +def resolve_for_building(msg): + """:func:`resolve_present_tables`, for a message a builder is about to build on. + + A builder lifts the relation tree or the expression into a new container, + and the tables do not travel with it. Entries that outlive resolution -- kept + for a ``google.protobuf.Any`` payload that may reference them -- would be + lost, leaving the payload's ordinals pointing at nothing or, inside a + subquery, at the enclosing container's entries, so that raises + ``ValueError`` instead. + """ + out = resolve_present_tables(msg) + if _has_entries(out): + raise ValueError( + "cannot build on a message whose detached entries an extension " + "payload may still reference: they would not travel with it" + ) + return out + + +def _resolve_in_place(container) -> None: + tables = {table: getattr(container, table.field) for table in _TABLES} + used = {table: set() for table in _TABLES} + has_any = False + stack = [child for child, _ in _children(container, skip_tables=True)] + while stack: + node = stack.pop() + if _is_any(node): + has_any = True + continue + table = _ordinal_table(node) + if table is not None: + entries = tables[table] + ordinal = getattr(node, table.ordinal) + if ordinal >= len(entries): + raise ValueError( + f"detached {table.label} ordinal {ordinal} is out of range: the " + f"container has {len(entries)} detached {table.label}s" + ) + if ordinal in used[table]: + # With every entry inlined at most once, a cycle shows up here + # too: following it back reaches an entry already in use. + raise ValueError( + f"detached {table.label} {ordinal} is referenced more than " + "once, or by a cycle" + ) + entry = entries[ordinal] + if _ordinal_table(entry) is table: + raise ValueError( + f"detached {table.label} {ordinal} is itself an ordinal " + f"({getattr(entry, table.ordinal)})" + ) + used[table].add(ordinal) + node.CopyFrom(entry) + stack.extend(child for child, _ in _children(node)) + + # An Any in an entry nothing here reached may still be live: a payload can + # reference that entry. + has_any = has_any or any( + _find(entries[i], _is_any) + for table, entries in tables.items() + for i in range(len(entries)) + if i not in used[table] + ) + for table, entries in tables.items(): + if not has_any: + container.ClearField(table.field) + continue + while entries and len(entries) - 1 in used[table]: + del entries[-1] + for i in used[table]: + if i < len(entries): + entries[i].CopyFrom(_PLACEHOLDER) + + +def detach(msg, *, max_depth: int = DEFAULT_MAX_DEPTH): + """``msg`` with its expressions nesting at most ``max_depth`` levels deep. + + ``msg`` is a ``Plan``, ``PlanRel`` or ``ExtendedExpression``; depth is counted + in protobuf message levels below ``msg`` itself, the way parsers count it, and + the default matches their default limit. A message with nothing to detach is + returned itself, so its serialization is unchanged; otherwise the result is a + copy. + + The walk goes top-down, through the existing entries first and then the + relation or referred expressions, each in field declaration order -- the + order substrait-java uses, so both produce the same encoding. An expression + whose subtree would end below ``max_depth``, and which itself sits deeper than + half the budget (and than the container's entries), is moved to the end of + its container's ``detached_expressions`` and replaced by its ordinal; the + moved entry is then processed from its new position, right under the + container. Cutting at half the budget leaves chunks about half the budget + deep rather than one entry per level past the limit. Entries come out in + pre-order, so a child's ordinal is larger than its parent's. Entries already + in the table keep their positions, so existing ordinals still find their + targets. + + Only expressions are detached. Nesting inside literals, types, field + reference paths, masks or relation trees (see ``ReferenceRel``), and inside + ``google.protobuf.Any`` payloads, is left as is, so such a message can still + exceed ``max_depth``: the result is best effort, since the limit belongs to + the consumer, and some raise it. + + ``max_depth=0`` detaches every expression that can be, which is handy for + testing a consumer against the detached form. + """ + # Cutting works on a copy. Copying is native and cheap next to the walk, and + # this way the walk runs only once. + out = _copy(msg) + cut = False + for container, base in _containers(out): + cut = _detach_in_place(container, base, max_depth) or cut + return out if cut else msg + + +def _detach_in_place(container, base: int, max_depth: int) -> bool: + """Cut ``container``'s over-deep subtrees; whether anything was cut.""" + tables = {table: getattr(container, table.field) for table in _TABLES} + entry_depth = base + 1 + # Only subtrees below the entry roots are cut, so an entry root is never + # detached again however small the budget. + threshold = max(max_depth // 2, entry_depth) + + # Index the tree once: ``nodes[i]`` is a submessage and ``children[i]`` its + # ``(index, levels)`` pairs. Holding every wrapper keeps the tree stable while + # it is rewritten. The existing entries come first, then the relation tree. + nodes, children = [], [] + + def add(msg): + nodes.append(msg) + children.append([]) + return len(nodes) - 1 + + roots = [ + (add(entry), entry_depth) for entries in tables.values() for entry in entries + ] + roots += [ + (add(child), base + levels) + for child, levels in _children(container, skip_tables=True) + ] + stack = [i for i, _ in roots] + while stack: + i = stack.pop() + for child, levels in _children(nodes[i]): + j = add(child) + children[i].append((j, levels)) + stack.append(j) + + # Children are indexed after their parent, so a reverse sweep sees them first. + height = [0] * len(nodes) + for i in reversed(range(len(nodes))): + height[i] = max((levels + height[j] for j, levels in children[i]), default=0) + + # Choose the cuts top-down. A cut only changes depths below it, so every + # choice can be made against the subtree heights of the uncut tree. Each cut's + # subtree is walked straight after it is chosen, so ``cuts`` is in pre-order + # and each table's ordinals are handed out in that order. + cuts = [] + next_ordinal = {table: len(entries) for table, entries in tables.items()} + stack = [(i, depth, True) for i, depth in reversed(roots)] + while stack: + i, depth, is_root = stack.pop() + node = nodes[i] + table = _TABLE_BY_MESSAGE.get(node.DESCRIPTOR.full_name) + if ( + table is not None + and not is_root + and depth > threshold + # An ordinal is already as shallow as a subtree gets; detaching it + # would make an entry that is itself an ordinal. + and _ordinal_table(node) is None + ): + if depth + height[i] <= max_depth: + continue + cuts.append((i, table, next_ordinal[table])) + next_ordinal[table] += 1 + depth = entry_depth + stack.extend((j, depth + levels, False) for j, levels in reversed(children[i])) + + # Materialize the deepest cuts first, across tables, so each entry copies only + # its own chunk: the cuts below it are already ordinals. + for table, entries in tables.items(): + for _ in range(next_ordinal[table] - len(entries)): + entries.add() + for i, table, ordinal in reversed(cuts): + node = nodes[i] + tables[table][ordinal].CopyFrom(node) + node.Clear() + setattr(node, table.ordinal, ordinal) + return bool(cuts) diff --git a/src/substrait/utils/display.py b/src/substrait/utils/display.py index 535ec27..908867a 100644 --- a/src/substrait/utils/display.py +++ b/src/substrait/utils/display.py @@ -12,6 +12,7 @@ import substrait.type_pb2 as stt from substrait.utils import type_num_names +from substrait.utils.detached import resolve_detached def _read_output_names(read: stalg.ReadRel) -> list: @@ -156,12 +157,25 @@ def _stream_plan(self, plan: stp.Plan, stream, depth: int): def _stream_relation(self, rel: stp.PlanRel, stream, depth: int): """Print a plan relation concisely""" + try: + rel = resolve_detached(rel) + malformed = False + except ValueError: + # Show a malformed relation as is: its ordinals print as such, and its + # detached table follows the tree. + malformed = True if rel.HasField("root"): self._stream_rel_root(rel.root, stream, depth) elif rel.HasField("rel"): self._stream_rel(rel.rel, stream, depth) + if malformed: + indent = " " * (depth * self.indent_size) + for i, entry in enumerate(rel.detached_expressions): + stream.write(f"{indent}$detached[{i}]:\n") + self._stream_expression(entry, stream, depth + 1) + def _stream_rel_root(self, root: stalg.RelRoot, stream, depth: int): """Print a relation root concisely""" indent = " " * (depth * self.indent_size) @@ -492,6 +506,9 @@ def _stream_expression(self, expression: stalg.Expression, stream, depth: int): self._stream_if_then(expression.if_then, stream, depth) elif expression.HasField("window_function"): self._stream_window_function(expression.window_function, stream, depth) + elif expression.HasField("detached_expression_ordinal"): + ordinal = expression.detached_expression_ordinal + stream.write(f"{indent}$detached[{ordinal}]\n") else: stream.write(f"{indent}\n") @@ -682,6 +699,8 @@ def _get_function_argument_string(self, arg) -> str: # For nested scalar functions, we'll handle them specially in the main printing # Return a placeholder that indicates it needs recursive expansion return "" + elif arg.value.HasField("detached_expression_ordinal"): + return f"$detached[{arg.value.detached_expression_ordinal}]" else: return "" elif arg.HasField("enum"): @@ -733,6 +752,9 @@ def _stream_function_argument(self, arg, stream, depth: int): stream.write(f"{indent}field: root\n") elif arg.value.HasField("scalar_function"): self._stream_scalar_function(arg.value.scalar_function, stream, depth) + elif arg.value.HasField("detached_expression_ordinal"): + ordinal = arg.value.detached_expression_ordinal + stream.write(f"{indent}$detached[{ordinal}]\n") else: stream.write(f"{indent}\n") elif arg.HasField("enum"): diff --git a/tests/test_detached.py b/tests/test_detached.py new file mode 100644 index 0000000..1e7a633 --- /dev/null +++ b/tests/test_detached.py @@ -0,0 +1,548 @@ +"""Tests for detached expression subtrees (spec v0.105.0, bounded expression nesting).""" + +import pytest +import substrait.algebra_pb2 as stalg +import substrait.extended_expression_pb2 as stee +import substrait.plan_pb2 as stp +import substrait.type_pb2 as stt +from google.protobuf import any_pb2, text_format +from google.protobuf.message import DecodeError, Message + +import substrait.dataframe as sub +from substrait.builders.extended_expression import ( + column, + in_predicate, + scalar_subquery, + set_comparison, + set_predicate, +) +from substrait.builders.plan import project, read_named_table, select +from substrait.builders.type import fp64, named_struct, struct +from substrait.extension_registry import ExtensionRegistry +from substrait.type_inference import ( + infer_expression_type, + infer_extended_expression_schema, + infer_plan_schema, +) +from substrait.utils.detached import ( + detach, + resolve_detached, +) +from substrait.utils.display import PlanPrinter + +registry = ExtensionRegistry(load_default_extensions=True) + +# The spec's examples, site/examples/proto-textformat/{plan_rel,extended_expression}. +SPEC_PLAN_REL = """ +root { + input { + project { + common { direct {} } + input { + read { + common { direct {} } + base_schema { + names: "value" + struct { + types { i64 { nullability: NULLABILITY_REQUIRED } } + nullability: NULLABILITY_REQUIRED + } + } + named_table { names: "example" } + } + } + expressions { detached_expression_ordinal: 0 } + } + } + names: "value" + names: "answer" +} +detached_expressions { literal { i64: 42 } } +""" + +SPEC_EXTENDED_EXPRESSION = """ +version { minor_number: 104 } +base_schema { struct { nullability: NULLABILITY_REQUIRED } } +referred_expr { + expression { detached_expression_ordinal: 0 } + output_names: "result" + output_names: "answer" +} +detached_expressions { + nested { struct { fields { detached_expression_ordinal: 1 } } } +} +detached_expressions { literal { i64: 42 } } +""" + + +def _plan_rel(textproto: str) -> stp.PlanRel: + return text_format.Parse(textproto, stp.PlanRel()) + + +def _plan(textproto: str) -> stp.Plan: + plan = stp.Plan() + plan.relations.add().CopyFrom(_plan_rel(textproto)) + return plan + + +def _deep_plan(levels: int) -> stp.Plan: + """A plan projecting a struct literal nested ``levels`` structs deep. + + Built by mutating in place: under upb, a message constructor copies its + arguments through the parser and so is itself bound by the depth limit.""" + plan = stp.Plan() + root = plan.relations.add().root + root.names.append("x") + project = root.input.project + project.input.read.named_table.names.append("t") + project.input.read.base_schema.struct.nullability = stt.Type.NULLABILITY_REQUIRED + expr = project.expressions.add() + for _ in range(levels): + expr = expr.nested.struct.fields.add() + expr.literal.i64 = 42 + return plan + + +def _submessages(msg): + """Yield ``(child, levels)``: a map value sits two levels down, under its entry.""" + for field, value in msg.ListFields(): + if isinstance(value, Message): + yield value, 1 + elif hasattr(value, "values"): + yield from ((v, 2) for v in value.values() if isinstance(v, Message)) + elif field.message_type is not None: + yield from ((v, 1) for v in value) + + +def _height(msg) -> int: + """How many message levels the deepest submessage sits below ``msg``.""" + deepest, stack = 0, [(msg, 0)] + while stack: + node, depth = stack.pop() + deepest = max(deepest, depth) + stack.extend((child, depth + levels) for child, levels in _submessages(node)) + return deepest + + +def _ordinals(msg): + """Every detached expression ordinal in ``msg``, in no particular order.""" + found, stack = [], [msg] + while stack: + node = stack.pop() + if isinstance(node, stalg.Expression) and ( + node.WhichOneof("rex_type") == "detached_expression_ordinal" + ): + found.append(node.detached_expression_ordinal) + stack.extend(child for child, _ in _submessages(node)) + return found + + +# --------------------------------------------------------------------------- +# Resolving +# --------------------------------------------------------------------------- + + +def test_resolve_spec_plan_rel_example(): + resolved = resolve_detached(_plan_rel(SPEC_PLAN_REL)) + assert resolved.root.input.project.expressions[0].literal.i64 == 42 + assert len(resolved.detached_expressions) == 0 + + +def test_resolve_spec_extended_expression_example(): + ee = text_format.Parse(SPEC_EXTENDED_EXPRESSION, stee.ExtendedExpression()) + resolved = resolve_detached(ee) + expression = resolved.referred_expr[0].expression + assert expression.nested.struct.fields[0].literal.i64 == 42 + assert len(resolved.detached_expressions) == 0 + + +def test_resolve_does_not_modify_input(): + plan = _plan(SPEC_PLAN_REL) + before = plan.SerializeToString() + resolve_detached(plan) + assert plan.SerializeToString() == before + + +def test_resolve_without_detached_expressions_returns_input(): + plan = _deep_plan(3) + assert resolve_detached(plan) is plan + + +def test_resolve_ignores_unreachable_entries(): + # Unreachable entries need not be valid: this one is itself an ordinal. + plan_rel = _plan_rel( + SPEC_PLAN_REL + "detached_expressions { detached_expression_ordinal: 0 }" + ) + resolved = resolve_detached(plan_rel) + assert resolved.root.input.project.expressions[0].literal.i64 == 42 + + +@pytest.mark.parametrize( + "table, message", + [ + ("", "out of range"), + ( + "detached_expressions { detached_expression_ordinal: 0 }", + "itself an ordinal", + ), + ( + # Entry 0 points back at itself through a nested struct. + ( + "detached_expressions {" + " nested { struct { fields { detached_expression_ordinal: 0 } } } }" + ), + "more than once, or by a cycle", + ), + ], +) +def test_resolve_rejects_invalid_references(table, message): + plan_rel = _plan_rel( + SPEC_PLAN_REL.replace("detached_expressions { literal { i64: 42 } }", table) + ) + with pytest.raises(ValueError, match=message): + resolve_detached(plan_rel) + + +def test_resolve_rejects_reused_entry(): + plan_rel = _plan_rel(SPEC_PLAN_REL) + project = plan_rel.root.input.project + project.expressions.add().detached_expression_ordinal = 0 + with pytest.raises(ValueError, match="more than once"): + resolve_detached(plan_rel) + + +def test_resolve_does_not_cross_plan_rels(): + # Ordinals are local to their PlanRel: the second one has no table of its own. + plan = _plan(SPEC_PLAN_REL) + second = plan.relations.add() + second.CopyFrom(plan.relations[0]) + second.ClearField("detached_expressions") + with pytest.raises(ValueError, match="out of range"): + resolve_detached(plan) + + +def _with_any_payload(plan_rel: stp.PlanRel) -> stp.PlanRel: + plan_rel.root.input.project.common.advanced_extension.enhancement.CopyFrom( + any_pb2.Any(type_url="example.com/Detail", value=b"\x08\x00") + ) + return plan_rel + + +def test_resolve_keeps_entries_an_any_payload_may_reference(): + # Entry 1 is unreachable from the relation tree, so the opaque payload may + # reference it: it must keep its ordinal. Entry 2, resolved here, is dropped. + plan_rel = _with_any_payload( + _plan_rel( + SPEC_PLAN_REL.replace( + "detached_expression_ordinal: 0", "detached_expression_ordinal: 2" + ) + + "detached_expressions { literal { i64: 7 } }" + "detached_expressions { literal { i64: 42 } }" + ) + ) + resolved = resolve_detached(plan_rel) + assert resolved.root.input.project.expressions[0].literal.i64 == 42 + assert [e.literal.i64 for e in resolved.detached_expressions] == [42, 7] + + +def test_resolve_with_any_payload_drops_entries_it_resolved(): + resolved = resolve_detached(_with_any_payload(_plan_rel(SPEC_PLAN_REL))) + assert len(resolved.detached_expressions) == 0 + + +def test_round_trips_with_any_payload_do_not_grow_the_table(): + plan = _deep_plan(200) + _with_any_payload(plan.relations[0]) + sizes = [] + for _ in range(3): + detached = detach(plan) + sizes.append(len(detached.relations[0].detached_expressions)) + plan = resolve_detached(stp.Plan.FromString(detached.SerializeToString())) + assert len(plan.relations[0].detached_expressions) == 0 + assert sizes == [sizes[0]] * 3 + + +def test_resolve_screens_out_ordinal_tag_bytes_elsewhere(): + # 144 encodes as the bytes 0x90 0x01, the same as the ordinal field's tag. + plan = _deep_plan(1) + plan.relations[0].root.input.project.expressions[0].nested.struct.fields[ + 0 + ].literal.i64 = 144 + assert resolve_detached(plan) is plan + + +def test_resolve_rejects_unsupported_message(): + with pytest.raises(TypeError, match="Plan, PlanRel or ExtendedExpression"): + resolve_detached(stalg.Expression()) + + +# --------------------------------------------------------------------------- +# Detaching +# --------------------------------------------------------------------------- + + +def test_detach_within_budget_returns_input(): + plan = _deep_plan(5) + assert detach(plan) is plan + + +def test_deep_plan_round_trips_through_bytes(): + plan = _deep_plan(1000) + with pytest.raises(DecodeError): + stp.Plan.FromString(plan.SerializeToString()) + + detached = detach(plan) + assert _height(detached) <= 100 + parsed = stp.Plan.FromString(detached.SerializeToString()) + assert resolve_detached(parsed) == plan + + +# The relation tree around the expression is 8 levels deep; detaching only +# shortens expressions, so smaller budgets are best effort. +@pytest.mark.parametrize("max_depth", [0, 8, 20, 64]) +def test_detach_respects_custom_budget(max_depth): + plan = _deep_plan(60) + detached = detach(plan, max_depth=max_depth) + if max_depth: + assert _height(detached) <= max_depth + assert resolve_detached(detached) == plan + + +def test_detach_orders_entries_in_pre_order(): + detached = detach(_deep_plan(200)) + table = detached.relations[0].detached_expressions + assert len(table) > 1 + # Each entry holds exactly the next cut down the chain. + for ordinal, entry in enumerate(table[:-1]): + assert _ordinals(entry) == [ordinal + 1] + assert _ordinals(table[-1]) == [] + + +def test_detach_keeps_existing_entries(): + plan = _plan(SPEC_PLAN_REL) + detached = detach(plan, max_depth=0) + table = detached.relations[0].detached_expressions + assert table[0].literal.i64 == 42 + assert resolve_detached(detached) == resolve_detached(plan) + + +def test_detach_extended_expression(): + ee = text_format.Parse(SPEC_EXTENDED_EXPRESSION, stee.ExtendedExpression()) + resolved = resolve_detached(ee) + detached = detach(resolved, max_depth=0) + expression = detached.referred_expr[0].expression + assert expression.WhichOneof("rex_type") == "detached_expression_ordinal" + assert resolve_detached(detached) == resolved + + +def _dataframe_plans(): + t = sub.read_named_table( + "t", {"id": sub.i64, "region": sub.string, "amount": sub.fp64} + ) + return [ + t.filter((sub.col("id") > 1) & (sub.col("amount") < 10.0)).to_plan(), + t.select((sub.col("amount") * 2 + 1).alias("x")).to_plan(), + t.group_by("region").agg(sub.f.avg(sub.col("amount")).alias("a")).to_plan(), + ] + + +@pytest.mark.parametrize("plan", _dataframe_plans()) +def test_fully_detached_dataframe_plans_round_trip(plan): + detached = detach(plan, max_depth=0) + assert detached.relations[-1].detached_expressions + parsed = stp.Plan.FromString(detached.SerializeToString()) + assert resolve_detached(parsed) == plan + + +def test_detach_rejects_unsupported_message(): + with pytest.raises(TypeError, match="Plan, PlanRel or ExtendedExpression"): + detach(stalg.Rel()) + + +# --------------------------------------------------------------------------- +# Consumers accept detached input +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("plan", _dataframe_plans()) +def test_infer_plan_schema_resolves_detached_expressions(plan): + detached = detach(plan, max_depth=0) + assert infer_plan_schema(detached, registry=registry) == infer_plan_schema( + plan, registry=registry + ) + + +def test_infer_extended_expression_schema_resolves_detached_expressions(): + ee = text_format.Parse(SPEC_EXTENDED_EXPRESSION, stee.ExtendedExpression()) + schema = infer_extended_expression_schema(ee) + assert schema.types[0].struct.types[0].WhichOneof("kind") == "i64" + + +def test_builders_inline_a_detached_input_plan(): + plan = _dataframe_plans()[1] + detached = detach(plan, max_depth=0) + built = select(detached, expressions=[column("x")])(registry) + assert not any(pr.detached_expressions for pr in built.relations) + assert _ordinals(built) == [] + assert built == select(plan, expressions=[column("x")])(registry) + + +def test_display_resolves_detached_expressions(): + plan = _dataframe_plans()[1] + detached = detach(plan, max_depth=0) + printer = PlanPrinter() + assert printer.stringify_plan(detached) == printer.stringify_plan(plan) + + +_SUBQUERY_BUILDERS = { + "scalar_subquery": scalar_subquery, + "set_predicate": lambda query: set_predicate( + query, stalg.Expression.Subquery.SetPredicate.PREDICATE_OP_EXISTS + ), + "in_predicate": lambda query: in_predicate([column("v")], query), + "set_comparison": lambda query: set_comparison( + column("v"), + query, + stalg.Expression.Subquery.SetComparison.REDUCTION_OP_ANY, + stalg.Expression.Subquery.SetComparison.COMPARISON_OP_EQ, + ), +} + + +@pytest.mark.parametrize("builder", _SUBQUERY_BUILDERS.values(), ids=_SUBQUERY_BUILDERS) +def test_subquery_builders_inline_a_detached_plan(builder): + inner = _dataframe_plans()[1] + outer_schema = named_struct( + names=["v"], struct=struct(types=[fp64(nullable=False)], nullable=False) + ) + + def build(query): + return project( + read_named_table("o", outer_schema), expressions=[builder(query)] + )(registry) + + built = build(detach(inner, max_depth=0)) + assert _ordinals(built) == [] + assert built == build(inner) + + +def test_expression_type_of_an_ordinal_names_the_fix(): + expression = stalg.Expression(detached_expression_ordinal=3) + with pytest.raises(ValueError, match="ordinal 3 .*resolve_detached"): + infer_expression_type(expression, stt.Type.Struct()) + + +def test_display_shows_a_malformed_detached_plan_as_is(): + plan = _plan( + SPEC_PLAN_REL.replace( + "detached_expression_ordinal: 0", "detached_expression_ordinal: 5" + ) + ) + assert "$detached[5]" in PlanPrinter().stringify_plan(plan) + + +def test_resolve_blanks_resolved_entries_kept_in_place(): + # Entry 0 (resolved) holds ordinal 2 (resolved, trailing, dropped); entry 1 is + # unreached and kept for the payload. Entry 0 must not keep pointing at 2. + plan_rel = _with_any_payload( + _plan_rel( + SPEC_PLAN_REL.replace( + "detached_expressions { literal { i64: 42 } }", + "detached_expressions {" + " nested { struct { fields { detached_expression_ordinal: 2 } } } }" + "detached_expressions { literal { i64: 7 } }" + "detached_expressions { literal { i64: 42 } }", + ) + ) + ) + resolved = resolve_detached(plan_rel) + table = resolved.detached_expressions + assert len(table) == 2 + assert _ordinals(table[0]) == [] + assert table[1].literal.i64 == 7 + + +def test_detach_walks_fields_in_declaration_order(): + # SwitchExpression declares ``match`` (field 3) before ``ifs`` (1) and + # ``else`` (2); substrait-java numbers cuts in declaration order. + plan = _deep_plan(0) + expr = plan.relations[0].root.input.project.expressions[0] + expr.Clear() + switch = expr.switch_expression + switch.match.literal.i64 = 1 + clause = switch.ifs.add() + clause.then.literal.i64 = 2 + getattr(switch, "else").literal.i64 = 3 + clause.SetInParent() + getattr(clause, "if").i64 = 1 + + table = detach(plan, max_depth=0).relations[0].detached_expressions + assert [e.literal.i64 for e in table[1:]] == [1, 2, 3] + + +def test_detach_processes_existing_entries_first(): + plan = _plan( + SPEC_PLAN_REL.replace( + "detached_expressions { literal { i64: 42 } }", + "detached_expressions { nested { struct { fields { literal { i64: 42 } } } } }", + ) + ) + project = plan.relations[0].root.input.project + project.expressions.add().nested.struct.fields.add().literal.i64 = 7 + + table = detach(plan, max_depth=0).relations[0].detached_expressions + # Entry 0 existed; its child is cut first, then the relation tree's. + assert table[1].literal.i64 == 42 + assert [e.literal.i64 for e in table[2:] if e.HasField("literal")] == [7] + + +def test_detach_with_nothing_to_cut_returns_input(): + plan = _deep_plan(200) + plan.relations[0].root.input.project.ClearField("expressions") + assert detach(plan, max_depth=10) is plan + + +def test_builders_inline_a_detached_plan_from_a_callable(): + plan = _dataframe_plans()[1] + detached = detach(plan, max_depth=0) + built = select(lambda registry: detached, expressions=[column("x")])(registry) + assert _ordinals(built) == [] + assert built == select(plan, expressions=[column("x")])(registry) + + +def test_builders_refuse_entries_a_payload_may_reference(): + plan = _plan(SPEC_PLAN_REL + "detached_expressions { literal { i64: 7 } }") + _with_any_payload(plan.relations[0]) + with pytest.raises(ValueError, match="would not travel with it"): + select(plan, expressions=[column("value")])(registry) + + +def test_builders_reject_an_ordinal_without_a_table(): + plan = _plan(SPEC_PLAN_REL) + plan.relations[0].ClearField("detached_expressions") + with pytest.raises(ValueError, match="resolve_detached"): + select(plan, expressions=[column("value")])(registry) + + +def test_display_shows_ordinal_function_arguments(): + plan = _dataframe_plans()[1] + detached = detach(plan, max_depth=0) + # Corrupt the table so the printer shows the plan unresolved. + detached.relations[0].detached_expressions.add().detached_expression_ordinal = 0 + detached.relations[ + 0 + ].root.input.project.expressions.add().detached_expression_ordinal = 99 + output = PlanPrinter().stringify_plan(detached) + assert "" not in output + assert "$detached[0]:" in output + + +def test_display_resolves_each_plan_rel_on_its_own(): + plan = _plan(SPEC_PLAN_REL) + malformed = plan.relations.add() + malformed.CopyFrom(plan.relations[0]) + malformed.root.input.project.expressions[0].detached_expression_ordinal = 5 + output = PlanPrinter().stringify_plan(plan) + # The first PlanRel resolves; only the second prints its ordinal and table. + assert "$detached[5]" in output + assert output.count("$detached[0]:") == 1