From e50fb96e8cd77012f3ed06049e69d1e15763f8ac Mon Sep 17 00:00:00 2001 From: Dimitri Yatsenko Date: Sat, 17 Oct 2026 19:07:06 -0500 Subject: [PATCH 1/3] test(diagram): move the expand strategies into their own module The Hypothesis strategies that draw model schemas, databases, seeds and conditions move from test_expand_spec.py to expand_strategies.py, so that the planner's tests can draw from them too. No test changes. --- tests/unit/expand_strategies.py | 310 ++++++++++++++++++++++++++++++++ tests/unit/test_expand_spec.py | 293 +----------------------------- 2 files changed, 311 insertions(+), 292 deletions(-) create mode 100644 tests/unit/expand_strategies.py diff --git a/tests/unit/expand_strategies.py b/tests/unit/expand_strategies.py new file mode 100644 index 000000000..7e7c42034 --- /dev/null +++ b/tests/unit/expand_strategies.py @@ -0,0 +1,310 @@ +""" +Hypothesis strategies for the ``Diagram.expand`` tests. + +They draw schemas in the reference model's format (``expand_model.py``), +consistent databases over small integer domains, seed rows, and restriction +conditions. ``test_expand_spec.py`` checks the model against the +specification with them, and ``test_expand_plan.py`` checks the planner against +the model. +""" + +from __future__ import annotations + +import itertools +from dataclasses import dataclass, field + +from hypothesis import assume +from hypothesis import strategies as st + +from .expand_model import DB, FK, Not, Relation, RowGraph, Schema, TableSpec, all_rows, q_acyclic + +#: Attribute values are drawn from this many small integers, so that keys collide. +DOMAIN = 3 +#: Tier prefixes of DataJoint table names: manual, lookup, imported, computed. +TIER_PREFIXES = ("", "#", "_", "__") + + +# ============================================================================= +# Generators +# ============================================================================= + + +class _TableBuilder: + """Collects the attributes and FKs of one table while it is drawn.""" + + def __init__(self, name: str, master: str | None = None) -> None: + self.name = name + self.master = master + self.pk: list[str] = [] + self.secondary: list[str] = [] + self.fks: list[FK] = [] + + def has(self, attr: str) -> bool: + return attr in self.pk or attr in self.secondary + + def add_fk(self, parent: TableSpec, primary: bool, renames: dict, nullable: bool = False) -> None: + # An FK adds only the attributes not declared yet, as in declare.py, so + # unrenamed attributes shared with an earlier FK are reused. + attr_map = {} + for attr in parent.pk: + name = renames.get(attr, attr) + attr_map[name] = attr + if not self.has(name): + (self.pk if primary else self.secondary).append(name) + if any(f.parent == parent.name and dict(f.attr_map) == attr_map for f in self.fks): + return # an identical FK adds nothing + self.fks.append(FK(self.name, parent.name, attr_map, nullable=nullable and not primary)) + + def spec(self) -> TableSpec: + return TableSpec(self.name, tuple(self.pk), tuple(self.secondary), self.master) + + +def _draw_renames(draw, parent: TableSpec, fresh: itertools.count, keep: tuple = ()) -> dict: + """Rename none or some of the parent's key attributes, never those in ``keep``.""" + candidates = [a for a in parent.pk if a not in keep] + if not candidates or not draw(st.booleans()): + return {} + renamed = draw(st.sets(st.sampled_from(candidates), min_size=1)) + return {a: f"{a}_r{next(fresh)}" for a in sorted(renamed)} + + +def _descendants(fks: list[FK], names: set[str]) -> set[str]: + """Tables that reference any of ``names``, directly or through other tables.""" + found, frontier = set(), set(names) + while frontier: + frontier = {fk.child for fk in fks if fk.parent in frontier} - found - names + found |= frontier + return found + + +def _draw_fks(draw, tables, fks, fresh, n, solid, group=None) -> list: + """ + Draw ``n`` FKs from the table being built to tables declared earlier. + + A uniform choice of parent rarely makes the shapes that matter, so the + parent is drawn from one of several pools: every table; tables whose key + shares some but not all of its attributes with this table's non-null + attributes ``solid`` (diamonds and, for a nullable FK, partial NULLs); and, + for a part, its own group (renamed references to the master or a sibling), + the parts of other groups, and the tables downstream of its group (group + cycles). + """ + drawn = [] + solid, known = set(solid), set(solid) + for _ in range(n): + kind = draw(st.sampled_from(["primary", "secondary", "nullable"])) + sharing = [t for t in tables if set(t.pk) & solid and not set(t.pk) <= known] + pools = [tables, sharing] + if kind == "nullable": + pools.append(sharing) # twice as likely, to make partial NULLs + if group is not None: + master = group[0] + pools.append([t for t in tables if t.name in group]) + pools.append([t for t in tables if t.master and t.master != master]) + downstream = _descendants(fks, set(group)) + pools.append([t for t in tables if t.name in downstream]) + parent = draw(st.sampled_from(draw(st.sampled_from([p for p in pools if p])))) + renames = _draw_renames(draw, parent, fresh) + attrs = {renames.get(a, a) for a in parent.pk} + if kind != "nullable": + solid |= attrs - known + known |= attrs + drawn.append((parent, kind == "primary", renames, kind == "nullable")) + return drawn + + +def _finish(draw, builder: _TableBuilder, tables: list[TableSpec], fks: list, own_pk: str, force_own: bool) -> TableSpec: + # Primary FKs first, then the table's own key attribute, then secondary FKs, + # as a definition lists them above and below the --- line. + for parent, primary, renames, nullable in fks: + if primary: + builder.add_fk(parent, True, renames) + # Without its own attribute, a table's key is its parents' (one-to-one, as in S1). + if force_own or not builder.pk or draw(st.sampled_from([True, True, False])): + builder.pk.append(own_pk) + for parent, primary, renames, nullable in fks: + if not primary: + builder.add_fk(parent, False, renames, nullable) + # Sometimes add a namesake: an attribute that shares a name with another + # table's key attribute but belongs to no FK. + namesakes = sorted({a for t in tables for a in t.pk if not builder.has(a)}) + if namesakes and draw(st.booleans()): + builder.secondary.append(draw(st.sampled_from(namesakes))) + return builder.spec() + + +@st.composite +def schemas(draw, max_tables: int = 7, owned_only: bool = False) -> Schema: + """ + Draw a schema: tables in declaration order, so the FK graph is a DAG. + + - Each step declares a plain table, a master, or a part of a master declared + earlier, so parts can be added late, after other tables. + - A master has 0-3 parts. A part is owned through ``-> master`` in its + primary key, a renamed one, or a primary-key FK to an owned sibling with + or without renames (#1429). Unless ``owned_only``, a part may also be not + owned: it references the master only below the line, or through a + nullable FK, or through a part that is not owned, or not at all. + - Other FKs are primary or secondary, nullable, renamed and parallel, to any + table declared earlier. Diamonds that share attributes, part -> part across + groups, references inside a group that link instances, and group cycles + all occur. + - Some tables get a namesake attribute that belongs to no FK. + """ + fresh = itertools.count() + tables: list[TableSpec] = [] + fks: list[FK] = [] + parts_of: dict[str, list[str]] = {} + for i in range(draw(st.integers(min(2, max_tables), max_tables))): + by_name = {t.name: t for t in tables} + open_masters = [m for m, parts in parts_of.items() if len(parts) < 3] + kind = draw(st.sampled_from(["plain", "master"] + (["part", "part"] if open_masters else []))) + if kind != "part": + builder = _TableBuilder(f"{draw(st.sampled_from(TIER_PREFIXES))}t{i}") + extra = _draw_fks(draw, tables, fks, fresh, draw(st.integers(0, 3)) if tables else 0, set()) + spec = _finish(draw, builder, tables, extra, own_pk=f"k{i}", force_own=False) + if kind == "master": + parts_of[spec.name] = [] + else: + master = draw(st.sampled_from(open_masters)) + mspec = by_name[master] + builder = _TableBuilder(f"{master}__p{len(parts_of[master])}", master=master) + owned = Schema(tables, fks).owned_parts.get(master, ()) + not_owned = [p for p in parts_of[master] if p not in owned] + owner_modes = ["direct", "renamed"] + (["sibling"] if owned else []) + ways_not_owned = ["secondary", "nullable", "missing"] + (["through a part not owned"] if not_owned else []) + mode = draw(st.sampled_from(owner_modes + ([] if owned_only else ["not owned"]))) + if mode == "not owned": + mode = draw(st.sampled_from(ways_not_owned)) + below_line = [] # a reference to the master below the line + if mode == "direct": + builder.add_fk(mspec, True, {}) + elif mode == "renamed": + renamed = draw(st.sets(st.sampled_from(mspec.pk), min_size=1)) + builder.add_fk(mspec, True, {a: f"{a}_r{next(fresh)}" for a in sorted(renamed)}) + elif mode == "sibling": + sibling = by_name[draw(st.sampled_from(owned))] + builder.add_fk(sibling, True, _draw_renames(draw, sibling, fresh)) + elif mode == "through a part not owned": + other = by_name[draw(st.sampled_from(not_owned))] + builder.add_fk(other, True, _draw_renames(draw, other, fresh)) + elif mode in ("secondary", "nullable"): + below_line = [(mspec, False, {}, mode == "nullable")] + solid = set(builder.pk) | {a for o in below_line if not o[3] for a in o[0].pk} + group = [master] + parts_of[master] + extra = _draw_fks(draw, tables, fks, fresh, draw(st.integers(0, 3)), solid, group) + # With the master referenced below the line, the part needs a key attribute of its own. + spec = _finish(draw, builder, tables, below_line + extra, own_pk=f"k{i}", force_own=bool(below_line)) + parts_of[master].append(spec.name) + tables.append(spec) + fks.extend(builder.fks) + schema = Schema(tables, fks) + if owned_only: + assert all(t.name in schema.group for t in schema.tables if t.master), schema + return schema + + +def _draw_row(draw, schema: Schema, db: DB, table: TableSpec) -> dict | None: + """Draw one row whose FK values name existing parent rows or are NULL; None if no row fits.""" + row = {} + nullable = schema.nullable[table.name] + for fk in schema.fks_of(table.name): + pairs = list(fk.attr_map.items()) + if any(row.get(c, 0) is None for c, _ in pairs): + continue # an earlier FK set a shared attribute to NULL, so this FK references nothing + matches = [p for p in db[fk.parent].values() if all(c not in row or row[c] == p[a] for c, a in pairs)] + can_be_null = any(c in nullable and c not in row for c, _ in pairs) + options = len(matches) + can_be_null + if not options: + return None + choice = draw(st.integers(0, options - 1)) + if choice < len(matches): + row.update({c: matches[choice][a] for c, a in pairs}) + else: + # NULL in the attributes this FK introduced; shared key attributes keep values (partial NULL). + row.update({c: None for c, _ in pairs if c in nullable and c not in row}) + for attr in table.heading: + if attr not in row: + row[attr] = draw(st.integers(0, DOMAIN - 1)) + return row + + +@st.composite +def databases(draw, schema: Schema, max_rows: int = 4) -> DB: + """Draw a consistent database: tables filled in declaration order over small integer domains.""" + db = {} + for table in schema.tables: + rows = {} + for _ in range(draw(st.integers(1, max_rows))): + row = _draw_row(draw, schema, db, table) + if row is not None: + rows.setdefault(tuple(row[a] for a in table.pk), row) + db[table.name] = rows + return db + + +@dataclass +class Case: + """A drawn schema and database, with the row graph built once.""" + + schema: Schema + db: DB + graph: RowGraph = field(repr=False) + + def __repr__(self) -> str: + rows = {t: list(r.values()) for t, r in self.db.items() if r} + return f"Case(\nschema={self.schema!r},\ndb={rows!r})" + + +@st.composite +def cases(draw, max_tables: int = 7, max_rows: int = 4, acyclic: bool = True, max_total: int | None = None) -> Case: + """Draw a schema and a database. With ``acyclic``, schemas whose Q has a cycle are discarded.""" + schema = draw(schemas(max_tables)) + if acyclic: + assume(q_acyclic(schema)) + db = draw(databases(schema, max_rows)) + if max_total is not None: + assume(len(all_rows(db)) <= max_total) + return Case(schema, db, RowGraph(schema, db)) + + +def seeds(rows) -> st.SearchStrategy: + """Draw a set of seed rows from ``rows``.""" + rows = sorted(rows) + return st.frozensets(st.sampled_from(rows)) if rows else st.just(frozenset()) + + +# ============================================================================= +# Conditions +# ============================================================================= + + +VALUES = st.sampled_from([0, 1, 2, None]) + + +def key_dicts(names: list[str]) -> st.SearchStrategy: + """A dict over some of ``names``, with values that may be NULL.""" + return st.lists(st.sampled_from(names), unique=True, max_size=3).flatmap( + lambda keys: st.fixed_dictionaries({k: VALUES for k in keys}) + ) + + +def relations(names: list[str]) -> st.SearchStrategy: + """A query expression used as a condition: a few rows over some of ``names``.""" + + def over(heading): + rows = st.frozensets(st.tuples(*[VALUES] * len(heading)), max_size=4) if heading else st.just(frozenset({()})) + return rows.map(lambda r: Relation(tuple(heading), r)) + + return st.lists(st.sampled_from(names), unique=True, max_size=3).flatmap(over) + + +def conditions(names: list[str]) -> st.SearchStrategy: + """A dict, a list of dicts, a query expression (Relation) or a Not, over ``names``.""" + return st.one_of( + key_dicts(names), + st.lists(key_dicts(names), max_size=2), + relations(names), + relations(names).map(Not), + key_dicts(names).map(Not), + ) diff --git a/tests/unit/test_expand_spec.py b/tests/unit/test_expand_spec.py index e1cd31bf1..43322d859 100644 --- a/tests/unit/test_expand_spec.py +++ b/tests/unit/test_expand_spec.py @@ -50,7 +50,6 @@ from __future__ import annotations import itertools -from dataclasses import dataclass, field import networkx as nx import pytest @@ -61,15 +60,11 @@ from datajoint.errors import DataJointError from .expand_model import ( - DB, FK, - Not, PartIntegrityError, - Relation, RowGraph, Schema, TableSpec, - all_rows, apply_delete, copy_allowed, dependency_graph, @@ -91,262 +86,7 @@ subdatabase, virtual_upstream, ) - -#: Attribute values are drawn from this many small integers, so that keys collide. -DOMAIN = 3 -#: Tier prefixes of DataJoint table names: manual, lookup, imported, computed. -TIER_PREFIXES = ("", "#", "_", "__") - - -# ============================================================================= -# Generators -# ============================================================================= - - -class _TableBuilder: - """Collects the attributes and FKs of one table while it is drawn.""" - - def __init__(self, name: str, master: str | None = None) -> None: - self.name = name - self.master = master - self.pk: list[str] = [] - self.secondary: list[str] = [] - self.fks: list[FK] = [] - - def has(self, attr: str) -> bool: - return attr in self.pk or attr in self.secondary - - def add_fk(self, parent: TableSpec, primary: bool, renames: dict, nullable: bool = False) -> None: - # An FK adds only the attributes not declared yet, as in declare.py, so - # unrenamed attributes shared with an earlier FK are reused. - attr_map = {} - for attr in parent.pk: - name = renames.get(attr, attr) - attr_map[name] = attr - if not self.has(name): - (self.pk if primary else self.secondary).append(name) - if any(f.parent == parent.name and dict(f.attr_map) == attr_map for f in self.fks): - return # an identical FK adds nothing - self.fks.append(FK(self.name, parent.name, attr_map, nullable=nullable and not primary)) - - def spec(self) -> TableSpec: - return TableSpec(self.name, tuple(self.pk), tuple(self.secondary), self.master) - - -def _draw_renames(draw, parent: TableSpec, fresh: itertools.count, keep: tuple = ()) -> dict: - """Rename none or some of the parent's key attributes, never those in ``keep``.""" - candidates = [a for a in parent.pk if a not in keep] - if not candidates or not draw(st.booleans()): - return {} - renamed = draw(st.sets(st.sampled_from(candidates), min_size=1)) - return {a: f"{a}_r{next(fresh)}" for a in sorted(renamed)} - - -def _descendants(fks: list[FK], names: set[str]) -> set[str]: - """Tables that reference any of ``names``, directly or through other tables.""" - found, frontier = set(), set(names) - while frontier: - frontier = {fk.child for fk in fks if fk.parent in frontier} - found - names - found |= frontier - return found - - -def _draw_fks(draw, tables, fks, fresh, n, solid, group=None) -> list: - """ - Draw ``n`` FKs from the table being built to tables declared earlier. - - A uniform choice of parent rarely makes the shapes that matter, so the - parent is drawn from one of several pools: every table; tables whose key - shares some but not all of its attributes with this table's non-null - attributes ``solid`` (diamonds and, for a nullable FK, partial NULLs); and, - for a part, its own group (renamed references to the master or a sibling), - the parts of other groups, and the tables downstream of its group (group - cycles). - """ - drawn = [] - solid, known = set(solid), set(solid) - for _ in range(n): - kind = draw(st.sampled_from(["primary", "secondary", "nullable"])) - sharing = [t for t in tables if set(t.pk) & solid and not set(t.pk) <= known] - pools = [tables, sharing] - if kind == "nullable": - pools.append(sharing) # twice as likely, to make partial NULLs - if group is not None: - master = group[0] - pools.append([t for t in tables if t.name in group]) - pools.append([t for t in tables if t.master and t.master != master]) - downstream = _descendants(fks, set(group)) - pools.append([t for t in tables if t.name in downstream]) - parent = draw(st.sampled_from(draw(st.sampled_from([p for p in pools if p])))) - renames = _draw_renames(draw, parent, fresh) - attrs = {renames.get(a, a) for a in parent.pk} - if kind != "nullable": - solid |= attrs - known - known |= attrs - drawn.append((parent, kind == "primary", renames, kind == "nullable")) - return drawn - - -def _finish(draw, builder: _TableBuilder, tables: list[TableSpec], fks: list, own_pk: str, force_own: bool) -> TableSpec: - # Primary FKs first, then the table's own key attribute, then secondary FKs, - # as a definition lists them above and below the --- line. - for parent, primary, renames, nullable in fks: - if primary: - builder.add_fk(parent, True, renames) - # Without its own attribute, a table's key is its parents' (one-to-one, as in S1). - if force_own or not builder.pk or draw(st.sampled_from([True, True, False])): - builder.pk.append(own_pk) - for parent, primary, renames, nullable in fks: - if not primary: - builder.add_fk(parent, False, renames, nullable) - # Sometimes add a namesake: an attribute that shares a name with another - # table's key attribute but belongs to no FK. - namesakes = sorted({a for t in tables for a in t.pk if not builder.has(a)}) - if namesakes and draw(st.booleans()): - builder.secondary.append(draw(st.sampled_from(namesakes))) - return builder.spec() - - -@st.composite -def schemas(draw, max_tables: int = 7, owned_only: bool = False) -> Schema: - """ - Draw a schema: tables in declaration order, so the FK graph is a DAG. - - - Each step declares a plain table, a master, or a part of a master declared - earlier, so parts can be added late, after other tables. - - A master has 0-3 parts. A part is owned through ``-> master`` in its - primary key, a renamed one, or a primary-key FK to an owned sibling with - or without renames (#1429). Unless ``owned_only``, a part may also be not - owned: it references the master only below the line, or through a - nullable FK, or through a part that is not owned, or not at all. - - Other FKs are primary or secondary, nullable, renamed and parallel, to any - table declared earlier. Diamonds that share attributes, part -> part across - groups, references inside a group that link instances, and group cycles - all occur. - - Some tables get a namesake attribute that belongs to no FK. - """ - fresh = itertools.count() - tables: list[TableSpec] = [] - fks: list[FK] = [] - parts_of: dict[str, list[str]] = {} - for i in range(draw(st.integers(min(2, max_tables), max_tables))): - by_name = {t.name: t for t in tables} - open_masters = [m for m, parts in parts_of.items() if len(parts) < 3] - kind = draw(st.sampled_from(["plain", "master"] + (["part", "part"] if open_masters else []))) - if kind != "part": - builder = _TableBuilder(f"{draw(st.sampled_from(TIER_PREFIXES))}t{i}") - extra = _draw_fks(draw, tables, fks, fresh, draw(st.integers(0, 3)) if tables else 0, set()) - spec = _finish(draw, builder, tables, extra, own_pk=f"k{i}", force_own=False) - if kind == "master": - parts_of[spec.name] = [] - else: - master = draw(st.sampled_from(open_masters)) - mspec = by_name[master] - builder = _TableBuilder(f"{master}__p{len(parts_of[master])}", master=master) - owned = Schema(tables, fks).owned_parts.get(master, ()) - not_owned = [p for p in parts_of[master] if p not in owned] - owner_modes = ["direct", "renamed"] + (["sibling"] if owned else []) - ways_not_owned = ["secondary", "nullable", "missing"] + (["through a part not owned"] if not_owned else []) - mode = draw(st.sampled_from(owner_modes + ([] if owned_only else ["not owned"]))) - if mode == "not owned": - mode = draw(st.sampled_from(ways_not_owned)) - below_line = [] # a reference to the master below the line - if mode == "direct": - builder.add_fk(mspec, True, {}) - elif mode == "renamed": - renamed = draw(st.sets(st.sampled_from(mspec.pk), min_size=1)) - builder.add_fk(mspec, True, {a: f"{a}_r{next(fresh)}" for a in sorted(renamed)}) - elif mode == "sibling": - sibling = by_name[draw(st.sampled_from(owned))] - builder.add_fk(sibling, True, _draw_renames(draw, sibling, fresh)) - elif mode == "through a part not owned": - other = by_name[draw(st.sampled_from(not_owned))] - builder.add_fk(other, True, _draw_renames(draw, other, fresh)) - elif mode in ("secondary", "nullable"): - below_line = [(mspec, False, {}, mode == "nullable")] - solid = set(builder.pk) | {a for o in below_line if not o[3] for a in o[0].pk} - group = [master] + parts_of[master] - extra = _draw_fks(draw, tables, fks, fresh, draw(st.integers(0, 3)), solid, group) - # With the master referenced below the line, the part needs a key attribute of its own. - spec = _finish(draw, builder, tables, below_line + extra, own_pk=f"k{i}", force_own=bool(below_line)) - parts_of[master].append(spec.name) - tables.append(spec) - fks.extend(builder.fks) - schema = Schema(tables, fks) - if owned_only: - assert all(t.name in schema.group for t in schema.tables if t.master), schema - return schema - - -def _draw_row(draw, schema: Schema, db: DB, table: TableSpec) -> dict | None: - """Draw one row whose FK values name existing parent rows or are NULL; None if no row fits.""" - row = {} - nullable = schema.nullable[table.name] - for fk in schema.fks_of(table.name): - pairs = list(fk.attr_map.items()) - if any(row.get(c, 0) is None for c, _ in pairs): - continue # an earlier FK set a shared attribute to NULL, so this FK references nothing - matches = [p for p in db[fk.parent].values() if all(c not in row or row[c] == p[a] for c, a in pairs)] - can_be_null = any(c in nullable and c not in row for c, _ in pairs) - options = len(matches) + can_be_null - if not options: - return None - choice = draw(st.integers(0, options - 1)) - if choice < len(matches): - row.update({c: matches[choice][a] for c, a in pairs}) - else: - # NULL in the attributes this FK introduced; shared key attributes keep values (partial NULL). - row.update({c: None for c, _ in pairs if c in nullable and c not in row}) - for attr in table.heading: - if attr not in row: - row[attr] = draw(st.integers(0, DOMAIN - 1)) - return row - - -@st.composite -def databases(draw, schema: Schema, max_rows: int = 4) -> DB: - """Draw a consistent database: tables filled in declaration order over small integer domains.""" - db = {} - for table in schema.tables: - rows = {} - for _ in range(draw(st.integers(1, max_rows))): - row = _draw_row(draw, schema, db, table) - if row is not None: - rows.setdefault(tuple(row[a] for a in table.pk), row) - db[table.name] = rows - return db - - -@dataclass -class Case: - """A drawn schema and database, with the row graph built once.""" - - schema: Schema - db: DB - graph: RowGraph = field(repr=False) - - def __repr__(self) -> str: - rows = {t: list(r.values()) for t, r in self.db.items() if r} - return f"Case(\nschema={self.schema!r},\ndb={rows!r})" - - -@st.composite -def cases(draw, max_tables: int = 7, max_rows: int = 4, acyclic: bool = True, max_total: int | None = None) -> Case: - """Draw a schema and a database. With ``acyclic``, schemas whose Q has a cycle are discarded.""" - schema = draw(schemas(max_tables)) - if acyclic: - assume(q_acyclic(schema)) - db = draw(databases(schema, max_rows)) - if max_total is not None: - assume(len(all_rows(db)) <= max_total) - return Case(schema, db, RowGraph(schema, db)) - - -def seeds(rows) -> st.SearchStrategy: - """Draw a set of seed rows from ``rows``.""" - rows = sorted(rows) - return st.frozensets(st.sampled_from(rows)) if rows else st.just(frozenset()) - +from .expand_strategies import Case, cases, conditions, key_dicts, schemas, seeds #: Tiny databases for the brute-force minimality checks (at most 2^8 supersets). tiny_cases = cases(max_tables=4, max_rows=2, max_total=8) @@ -750,37 +490,6 @@ def test_every_fk_inside_a_group_carries_a_master_key(schema): # ============================================================================= -VALUES = st.sampled_from([0, 1, 2, None]) - - -def key_dicts(names: list[str]) -> st.SearchStrategy: - """A dict over some of ``names``, with values that may be NULL.""" - return st.lists(st.sampled_from(names), unique=True, max_size=3).flatmap( - lambda keys: st.fixed_dictionaries({k: VALUES for k in keys}) - ) - - -def relations(names: list[str]) -> st.SearchStrategy: - """A query expression used as a condition: a few rows over some of ``names``.""" - - def over(heading): - rows = st.frozensets(st.tuples(*[VALUES] * len(heading)), max_size=4) if heading else st.just(frozenset({()})) - return rows.map(lambda r: Relation(tuple(heading), r)) - - return st.lists(st.sampled_from(names), unique=True, max_size=3).flatmap(over) - - -def conditions(names: list[str]) -> st.SearchStrategy: - """A dict, a list of dicts, a query expression (Relation) or a Not, over ``names``.""" - return st.one_of( - key_dicts(names), - st.lists(key_dicts(names), max_size=2), - relations(names), - relations(names).map(Not), - key_dicts(names).map(Not), - ) - - def _edge_and_names(case: Case, data, copyable: bool): fks = [fk for fk in case.schema.fks if fk.child != fk.parent] if copyable: From 7da8c05150ebba334dfc44cbbe0b169ffbe08f7e Mon Sep 17 00:00:00 2001 From: Dimitri Yatsenko Date: Sat, 17 Oct 2026 19:09:11 -0500 Subject: [PATCH 2/3] feat(diagram): plan expand as an intermediate form Add datajoint.expand, the planner behind Diagram.expand (diagram.md, "Traversal specification", Computation): - TraversalSchema, a schema description built from connection.dependencies and the headings: tables, primary keys, nullable attributes, foreign keys with their attr_map, and each master's owned parts with their owning attributes (A1). - The plan: Seed, Literal, Restrict, Lift, Reexpand and Or nodes that name the restrictions they use by handle. plan_expand makes one pass over the group quotient for down and up (both is the up pass seeded with the down pass), copies a condition only under R1's rule, relabels literal keys under the Seeds rule, and marks the groups whose foreign keys link instances to iterate. - Two evaluators: RowEvaluator on rows held in memory, and QueryCompiler, which compiles each node to a FreeTable with one condition and reads rows only to find the master keys of a group that iterates. tests/unit/test_expand_plan.py (L0') checks the plan against the reference model: the description agrees with the model, the plan evaluated in memory selects D(S), U(S) and U(D(S)) for every seed kind, plan structure for the catalog shapes, and a delete simulation. The simulation leaves DB - D(S) when no lift is evaluated at delete time, and when the nodes the spec says to materialize are materialized; with the plan's own flags it is a strict xfail until delete-time materialization lands. --- src/datajoint/expand.py | 1339 ++++++++++++++++++++++++++++++++ tests/unit/test_expand_plan.py | 1059 +++++++++++++++++++++++++ 2 files changed, 2398 insertions(+) create mode 100644 src/datajoint/expand.py create mode 100644 tests/unit/test_expand_plan.py diff --git a/src/datajoint/expand.py b/src/datajoint/expand.py new file mode 100644 index 000000000..4781f45f4 --- /dev/null +++ b/src/datajoint/expand.py @@ -0,0 +1,1339 @@ +""" +The traversal planner behind ``Diagram.expand``. + +``Diagram.expand(seed, direction)`` selects, in each table it reaches, the rows +that the traversal specification defines: ``D(S)``, the seed rows and every row +that depends on them (``"down"``); ``U(S)``, the seed rows and every row they +depend on (``"up"``); and ``U(D(S))`` (``"both"``). The specification is the +"Traversal specification" section of the Diagram specification +(``reference/specs/diagram.md`` in datajoint-docs); its shape catalog, +``reference/specs/traversal-shapes.md``, works through one example per schema +shape. + +The rows are computed in three steps, as the specification's "Computation" +section describes them: + +1. A **schema description** (:class:`TraversalSchema`) holds the tables, their + primary keys and nullable attributes, the foreign keys with their attribute + maps, and each master's owned parts with their owning attributes (A1). It is + plain data: :meth:`TraversalSchema.from_dependencies` builds it from + ``connection.dependencies`` and the table headings, and tests build it by hand. +2. The **planner** (:func:`plan_expand`) builds a :class:`Plan`: an intermediate + form, not SQL, whose nodes are :class:`Seed`, :class:`Literal`, + :class:`Restrict`, :class:`Lift`, :class:`Reexpand` and :class:`Or`. A node + refers to the restrictions it uses by a handle, the index of another node. The + planner makes one pass over the group quotient (R1, the edge rule, and R2, the + group rule) and marks the groups that iterate. It reads no rows. +3. An **evaluator** gives the plan its rows. :func:`compile_plan` turns it into + query expressions; :class:`RowEvaluator` evaluates it on rows held in memory, + which is how the test suite compares it with the definitions. + +Compilation reads rows only for the groups that iterate. A ``dj.Top`` seed is +replaced by the keys it selects before planning (:func:`seed_from_expression`). +Materialization before a delete is not done here; ``Table.delete`` does not use +this module yet. +""" + +from __future__ import annotations + +import copy +import inspect +import re +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any, Callable, Iterable, Mapping + +import networkx as nx + +from .condition import AndList, make_condition +from .dependencies import extract_master +from .errors import DataJointError + +#: The directions of ``Diagram.expand``. +DIRECTIONS = ("down", "up", "both") + + +# ============================================================================= +# Schema description +# ============================================================================= + + +@dataclass(frozen=True) +class TableInfo: + """ + One table of a schema description. + + Parameters + ---------- + name : str + Table name: the full table name when built from a connection. + primary_key : tuple[str, ...] + Primary-key attributes, in order. + attributes : tuple[str, ...], optional + All attributes in heading order, the primary key included. Defaults to the + primary key. + nullable : frozenset[str], optional + Attributes that can hold NULL. + master : str, optional + The master's name if the table is a part table (by name, as + ``extract_master`` reads it). Whether the master owns it is decided by + :class:`TraversalSchema`. + """ + + name: str + primary_key: tuple[str, ...] + attributes: tuple[str, ...] = () + nullable: frozenset = frozenset() + master: str | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "primary_key", tuple(self.primary_key)) + object.__setattr__(self, "attributes", tuple(self.attributes) or self.primary_key) + object.__setattr__(self, "nullable", frozenset(self.nullable)) + + +@dataclass(frozen=True) +class ForeignKey: + """ + One foreign key: the child references the parent's primary key. + + Parameters + ---------- + child : str + The referencing table. + parent : str + The referenced table. + attr_map : mapping or tuple of pairs + Each child attribute of the foreign key and the parent primary-key + attribute it references. Stored as a tuple of ``(child, parent)`` pairs. + """ + + child: str + parent: str + attr_map: tuple = () + + def __post_init__(self) -> None: + pairs = self.attr_map.items() if isinstance(self.attr_map, Mapping) else self.attr_map + object.__setattr__(self, "attr_map", tuple((c, p) for c, p in pairs)) + + @property + def mapping(self) -> dict[str, str]: + """The attribute map as a dict, child attribute to parent attribute.""" + return dict(self.attr_map) + + @property + def attributes(self) -> tuple[str, ...]: + """The child's attributes in this foreign key.""" + return tuple(c for c, _ in self.attr_map) + + @property + def renamed(self) -> bool: + """True if any attribute has a different name in the child than in the parent.""" + return any(c != p for c, p in self.attr_map) + + +class TraversalSchema: + """ + The tables and foreign keys a traversal needs, with the owned parts (A1). + + Parameters + ---------- + tables : iterable of TableInfo + The tables. Their order breaks ties wherever the topological order leaves + a choice, so that plans are reproducible. + foreign_keys : iterable of ForeignKey + The foreign keys between the tables. Each child's foreign keys are listed + in the order of its definition (see Notes). + + Attributes + ---------- + groups : dict[str, tuple[str, ...]] + For each master that owns at least one part, its owned parts, parents + before children. + group_of : dict[str, str] + The master of every table in a group, the master itself included. + owning : dict[str, dict[str, str]] + For each table in a group, the attribute of that table that holds each + attribute of the master's primary key: its owning attributes. The master + maps its key onto itself. + + Notes + ----- + A part is owned by its master if a foreign key in the part's primary key (all + its attributes are primary-key attributes of the part) references the master, + or an owned part of the same master, renamed or not. Its owning attributes + are the attributes that chain maps onto the master's key. If a part has + several owning references, an unrenamed reference to the master wins, and + otherwise the first one in the part's definition (diagram.md, A1). + + The same rule decides ownership in ``datajoint.dependencies.owned_parts`` + (#1620), which finds the owned parts in the dependency graph. This class + also needs the owning attributes and that tie-break, so it computes both from + the foreign keys directly. + """ + + def __init__(self, tables: Iterable[TableInfo], foreign_keys: Iterable[ForeignKey]) -> None: + self.tables: dict[str, TableInfo] = {} + for table in tables: + if table.name in self.tables: + raise DataJointError(f"Table {table.name} is described twice.") + self.tables[table.name] = table + self.foreign_keys = tuple(foreign_keys) + self._position = {name: i for i, name in enumerate(self.tables)} + self._fks_of = defaultdict(list) # child -> its foreign keys, in definition order + self._fks_to = defaultdict(list) # parent -> the foreign keys that reference it + for fk in self.foreign_keys: + if fk.child not in self.tables or fk.parent not in self.tables: + raise DataJointError(f"Foreign key {fk.child} -> {fk.parent} names a table that is not described.") + self._fks_of[fk.child].append(fk) + self._fks_to[fk.parent].append(fk) + self.groups: dict[str, tuple[str, ...]] = {} + self.group_of: dict[str, str] = {} + self.owning: dict[str, dict[str, str]] = {} + self._find_owned_parts() + + def __repr__(self) -> str: + lines = ["TraversalSchema("] + for t in self.tables.values(): + extra = f", nullable={sorted(t.nullable)}" if t.nullable else "" + extra += f", master={t.master!r}" if t.master else "" + lines.append(f" TableInfo({t.name!r}, {t.primary_key!r}, {t.attributes!r}{extra}),") + for fk in self.foreign_keys: + lines.append(f" ForeignKey({fk.child!r}, {fk.parent!r}, {fk.mapping!r}),") + return "\n".join(lines) + "\n)" + + # --- construction --------------------------------------------------------- + + @classmethod + def from_dependencies(cls, dependencies: nx.MultiDiGraph, names: Iterable[str], heading: Callable) -> TraversalSchema: + """ + Describe the tables ``names`` of a dependency graph. + + Parameters + ---------- + dependencies : networkx.MultiDiGraph + ``connection.dependencies``: one edge per foreign key, parent to + child, with its ``attr_map``. + names : iterable of str + Full names of the tables to describe. Foreign keys between them are + included; foreign keys to other tables are left out. + heading : callable + Returns the ``Heading`` of a table, given its full name. Attribute + order and nullability come from it. + + Returns + ------- + TraversalSchema + + Notes + ----- + The database does not record the order in which a definition lists its + foreign keys. Each foreign key adds the attributes the table does not + have yet, in definition order, so the foreign keys are ordered here by the + heading positions of their attributes. This order is used only to break + the tie between two renamed owning references of one part. + """ + names = sorted(set(names)) + tables = [] + position = {} + for name in names: + h = heading(name) + attributes = tuple(h.names) + position[name] = {a: i for i, a in enumerate(attributes)} + master = extract_master(name) + tables.append( + TableInfo( + name, + tuple(h.primary_key), + attributes, + frozenset(a for a in attributes if h.attributes[a].nullable), + master if master in names else None, + ) + ) + selected = set(names) + fks = [ + ForeignKey(child, parent, props["attr_map"]) + for parent, child, props in dependencies.edges(data=True) + if parent in selected and child in selected + ] + fks.sort(key=lambda fk: (fk.child, sorted(position[fk.child].get(a, len(position[fk.child])) for a in fk.attributes))) + return cls(tables, fks) + + def _find_owned_parts(self) -> None: + # Parts are taken parents first, so that a part's references to other + # parts are decided before the part itself. + parts_of = defaultdict(list) + for t in self.tables.values(): + if t.master is not None and t.master in self.tables: + parts_of[t.master].append(t.name) + for master, parts in parts_of.items(): + key = self.tables[master].primary_key + owning = {master: {a: a for a in key}} + graph = nx.DiGraph() + graph.add_nodes_from(parts) + graph.add_edges_from((fk.parent, fk.child) for p in parts for fk in self._fks_of[p] if fk.parent in parts) + for part in nx.lexicographical_topological_sort(graph, key=self._position.get): + chains = [fk for fk in self._fks_of[part] if fk.parent in owning and self.is_primary(fk)] + if not chains: + continue + unrenamed = [fk for fk in chains if fk.parent == master and not fk.renamed] + fk = (unrenamed or chains)[0] + by_parent = {p: c for c, p in fk.attr_map} + owning[part] = {a: by_parent[owning[fk.parent][a]] for a in key} + owned = tuple(p for p in nx.lexicographical_topological_sort(graph, key=self._position.get) if p in owning) + if owned: + self.groups[master] = owned + self.owning.update(owning) + self.group_of.update({name: master for name in owning}) + + # --- facts about tables and foreign keys ----------------------------------- + + def fks_of(self, child: str) -> list[ForeignKey]: + """The foreign keys of ``child``, in definition order.""" + return list(self._fks_of[child]) + + def fks_to(self, parent: str) -> list[ForeignKey]: + """The foreign keys that reference ``parent``.""" + return list(self._fks_to[parent]) + + def is_primary(self, fk: ForeignKey) -> bool: + """True if all the foreign key's attributes are in the child's primary key.""" + return set(fk.attributes) <= set(self.tables[fk.child].primary_key) + + def owning_attributes(self, table: str) -> tuple[str, ...] | None: + """The table's owning attributes, in the order of the master's key; None outside a group.""" + master = self.group_of.get(table) + if master is None: + return None + return tuple(self.owning[table][a] for a in self.tables[master].primary_key) + + def intra_group_fks(self, master: str) -> list[ForeignKey]: + """The foreign keys between two different tables of the group of ``master``.""" + return [ + fk + for fk in self.foreign_keys + if fk.child != fk.parent and self.group_of.get(fk.child) == master and self.group_of.get(fk.parent) == master + ] + + def preserves_instance(self, fk: ForeignKey) -> bool: + """ + True if a foreign key inside a group maps the referencing table's owning + attributes onto the referenced table's owning attributes. + + Such a foreign key stays inside one master-part instance. One that does + not links different instances, and its group iterates (diagram.md, + "Groups that iterate"). + """ + master = self.group_of[fk.child] + child, parent, mapping = self.owning[fk.child], self.owning[fk.parent], fk.mapping + return all(mapping.get(child[a]) == parent[a] for a in self.tables[master].primary_key) + + def links(self, master: str) -> tuple[ForeignKey, ...]: + """The foreign keys inside the group of ``master`` that link different instances.""" + return tuple(fk for fk in self.intra_group_fks(master) if not self.preserves_instance(fk)) + + def copy_allowed(self, fk: ForeignKey, attributes: frozenset | None) -> bool: + """ + R1's rule for copying a parent's condition to the child unchanged. + + The attributes the condition names are known and are among the foreign + key's attributes, the foreign key is not renamed, and none of its columns + is nullable. + """ + return ( + attributes is not None + and attributes <= set(fk.attributes) + and not fk.renamed + and not set(fk.attributes) & self.tables[fk.child].nullable + ) + + def relabel_down_allowed(self, fk: ForeignKey, keys: Iterable[Mapping]) -> bool: + """ + The Seeds rule for relabeling a parent's literal keys for the child. + + Every key names, among the parent's attributes, only attributes the + foreign key carries; and if the foreign key has a nullable column, every + key names all of the foreign key's attributes, none of them NULL. + """ + keys = list(keys) + parent = set(self.tables[fk.parent].attributes) + carried = {p for _, p in fk.attr_map} + if any((set(key) & parent) - carried for key in keys): + return False + if not set(fk.attributes) & self.tables[fk.child].nullable: + return True + return all(carried <= set(key) and all(key[a] is not None for a in carried) for key in keys) + + # --- reach and order -------------------------------------------------------- + + def reach(self, tables: Iterable[str], direction: str) -> frozenset: + """ + The table-level reach of ``tables``: the node set of an expanded diagram. + + The tables reached by following foreign keys in the direction of travel, + together with the master and all owned parts of every group that has a + reached member. For ``"both"``, the upward reach of the downward reach. + """ + if direction == "both": + return self.reach(self.reach(tables, "down"), "up") + _check_direction(direction) + seen = set(tables) + stack = list(seen) + while stack: + t = stack.pop() + nxt = {fk.child for fk in self._fks_to[t]} if direction == "down" else {fk.parent for fk in self._fks_of[t]} + master = self.group_of.get(t) + if master is not None: + nxt |= {master, *self.groups[master]} + for n in nxt - seen: + seen.add(n) + stack.append(n) + return frozenset(seen) + + def quotient_order(self, names: Iterable[str]) -> list[str]: + """ + The nodes of the group quotient Q among ``names``, in topological order. + + A group is represented by its master. Raises ``DataJointError`` if Q has a + cycle, which the declaration check refuses (#1619). + """ + names = sorted(set(names), key=lambda n: self._position[n]) + node = {t: self.group_of.get(t, t) for t in names} + q = nx.DiGraph() + q.add_nodes_from(node[t] for t in names) + for fk in self.foreign_keys: + if fk.parent in node and fk.child in node: + u, v = node[fk.parent], node[fk.child] + if u != v or fk.parent == fk.child: + q.add_edge(u, v) + try: + return list(nx.lexicographical_topological_sort(q, key=self._position.get)) + except nx.NetworkXUnfeasible: + cycle = [u for u, _ in nx.find_cycle(q)] + first = cycle.index(min(cycle, key=lambda n: self._position[n])) + cycle = cycle[first:] + cycle[:first] + raise DataJointError( + "Cannot traverse tables whose master-part groups form a cycle: " + f"{' -> '.join(cycle + cycle[:1])}. Each group is named by its master." + ) from None + + def order(self, names: Iterable[str]) -> list[str]: + """``names`` in topological order, with each master's owned parts right after it.""" + names = set(names) + result = [] + for q in self.quotient_order(names): + result.append(q) + result.extend(p for p in self.groups.get(q, ()) if p in names) + return result + + +def _check_direction(direction: str) -> None: + if direction not in ("down", "up"): + raise ValueError(f"direction must be 'down', 'up' or 'both', not {direction!r}") + + +# ============================================================================= +# The plan +# ============================================================================= +# Every node selects rows of one table, ``node.table``, by a condition. A node +# that uses another node's rows names it by its handle, its index in +# ``Plan.nodes``. The attribute maps of Restrict, Lift and Reexpand pair each +# attribute of the node's table with the attribute of the source's table whose +# values it must match. + + +@dataclass(frozen=True) +class Seed: + """ + The seed: one table restricted by a condition. + + Parameters + ---------- + table : str + The seed table. + condition : object + The restriction, opaque to the planner. ``seed_from_expression`` stores + a query expression's compiled conditions; tests store model conditions. + attributes : frozenset[str] or None + The attributes the condition names, which decide whether R1 may copy it. + None if they cannot be read from the condition, which is then never copied. + """ + + table: str + condition: Any = field(compare=False) + attributes: frozenset | None = None + + +@dataclass(frozen=True) +class Literal: + """ + Literal keys of a table. + + Parameters + ---------- + table : str + The table. + keys : tuple[dict, ...] or None + Key dicts, combined with OR; an empty tuple selects no rows. None marks + the master keys of a group that iterates, which the evaluator finds by + repeating the group step (see ``GroupStep``). + read : bool + True if the keys were read from the table's own rows, as ``dj.Top`` keys + are. Only such keys cross a foreign key upward by relabeling. + """ + + table: str + keys: tuple | None + read: bool = False + + +@dataclass(frozen=True) +class Restrict: + """ + R1 across one foreign key: the rows of ``table`` that match rows of ``src``. + + Downward, ``table`` is the child and ``attr_map`` is the foreign key's map; + upward, ``table`` is the parent and the map is reversed. Compiled, it is + ``table & src_rows.proj()``. + + With ``copy=True``, the child is restricted by ``src``'s own condition, + unchanged, as R1 allows when the condition names only attributes of an + unrenamed foreign key with no nullable column. + """ + + table: str + src: int + attr_map: tuple + copy: bool = False + + +@dataclass(frozen=True) +class Lift: + """ + R2's lift: the master rows of the part rows that ``src`` selects. + + ``table`` is the master and ``attr_map`` pairs each attribute of the master's + key with the part's owning attribute that holds it. Compiled, it is + ``master & part_rows.proj(m1="o1", ...)``. + """ + + table: str + src: int + attr_map: tuple + + +@dataclass(frozen=True) +class Reexpand: + """ + R2's re-expansion: every row of the part ``table`` whose master ``src`` selects. + + ``attr_map`` pairs each owning attribute of the part with the attribute of + the master's key it holds. Compiled, it is ``part & master_rows.proj(o1="m1", ...)``. + """ + + table: str + src: int + attr_map: tuple + + +@dataclass(frozen=True) +class Or: + """The union of the rows of several nodes of one table.""" + + table: str + terms: tuple + + +#: The node types of a plan. +NODE_TYPES = (Seed, Literal, Restrict, Lift, Reexpand, Or) +_MATCHING = (Restrict, Lift, Reexpand) + + +@dataclass(frozen=True) +class GroupStep: + """ + The visit of one group in one pass of a plan. + + Parameters + ---------- + master : str + The group's master. + parts : tuple[str, ...] + Its owned parts. + direction : str + The pass: ``"down"`` or ``"up"``. + entry : int + The master's rows after one step: the seed on the master, R1 into the + master, and the lift of R1 into each part and of the seed on each part. + links : tuple[ForeignKey, ...] + The foreign keys inside the group that link different instances. The + group iterates if there is one. + keys : int, optional + For a group that iterates: the ``Literal`` with ``keys=None`` that stands + for its master keys. The master's restriction is this node. + step : int, optional + For a group that iterates: the master rows reached in one more group step + from the instances of ``keys``, across ``links``. The evaluator binds + ``keys`` to the keys of ``entry`` and adds the keys of ``step`` until no + new key appears. + """ + + master: str + parts: tuple + direction: str + entry: int + links: tuple = () + keys: int | None = None + step: int | None = None + + @property + def iterate(self) -> bool: + """Whether the group needs iterative evaluation.""" + return bool(self.links) + + +class Plan: + """ + The plan of one expansion. + + Attributes + ---------- + schema : TraversalSchema + The schema the plan was built for. + direction : str + ``"down"``, ``"up"`` or ``"both"``. + seed : int + Handle of the seed node. + nodes : list + The nodes; a handle is an index into this list. + restriction : dict[str, int] + For each table of the expanded diagram, the handle of its rows. + order : tuple[str, ...] + The tables of the expanded diagram in topological order, owned parts + right after their master. + groups : list[GroupStep] + The groups visited, in the order of the passes. + """ + + def __init__(self, schema: TraversalSchema, direction: str) -> None: + self.schema = schema + self.direction = direction + self.seed: int = -1 # set by plan_expand + self.nodes: list = [] + self.restriction: dict[str, int] = {} + self.order: tuple = () + self.groups: list[GroupStep] = [] + self._index: dict = {} + + def add(self, node) -> int: + """Add a node and return its handle; an identical node already present is reused.""" + key = _node_key(node) + if key is not None and key in self._index: + return self._index[key] + self.nodes.append(node) + handle = len(self.nodes) - 1 + if key is not None: + self._index[key] = handle + return handle + + def union(self, table: str, terms: Iterable[int]) -> int: + """The handle of the union of ``terms``: the term itself if there is one, no rows if none.""" + flat = [] + for t in terms: + node = self.nodes[t] + for term in node.terms if isinstance(node, Or) else (t,): + if term not in flat: + flat.append(term) + if not flat: + return self.add(Literal(table, ())) + if len(flat) == 1: + return flat[0] + return self.add(Or(table, tuple(flat))) + + def attributes(self, handle: int) -> frozenset | None: + """ + The attributes the condition of a node names, for R1's copy rule. + + A Seed's are given with it, and a Literal's are the names in its keys. A + Restrict, Lift or Reexpand names the attributes it matches; a copy names + those of the condition it copies; an Or names those of all its terms. + None means unknown. + """ + node = self.nodes[handle] + if isinstance(node, Seed): + return node.attributes + if isinstance(node, Literal): + if node.keys is None: + return frozenset(self.schema.tables[node.table].primary_key) + return frozenset(a for key in node.keys for a in key) + if isinstance(node, Restrict) and node.copy: + return self.attributes(node.src) + if isinstance(node, _MATCHING): + return frozenset(a for a, _ in node.attr_map) + names = [self.attributes(t) for t in node.terms] + known = [n for n in names if n is not None] + return frozenset().union(*known) if len(known) == len(names) else None + + def iterating_groups(self) -> list[GroupStep]: + """The groups that need iterative evaluation.""" + return [g for g in self.groups if g.iterate] + + def describe(self) -> str: + """The plan, one node per line, then each table's restriction.""" + lines = [] + for h, node in enumerate(self.nodes): + if isinstance(node, Seed): + text = f"Seed({node.table}; names {_names(node.attributes)})" + elif isinstance(node, Literal): + keys = "the keys of an iterating group" if node.keys is None else f"{len(node.keys)} keys" + text = f"Literal({node.table}; {keys}{', read' if node.read else ''})" + elif isinstance(node, Or): + text = f"Or({node.table}; {', '.join(f'#{t}' for t in node.terms)})" + else: + kind = "copy of" if getattr(node, "copy", False) else "by" + pairs = ", ".join(a if a == b else f"{a}={b}" for a, b in node.attr_map) + text = f"{type(node).__name__}({node.table} {kind} #{node.src}; {pairs})" + lines.append(f"#{h} {text}") + for g in self.groups: + if g.iterate: + links = "; ".join(f"{fk.child} -> {fk.parent}" for fk in g.links) + lines.append(f"{g.direction}: {g.master} iterates over {links}: #{g.keys} from #{g.entry}, step #{g.step}") + lines.extend(f"{t} = #{self.restriction[t]}" for t in self.order) + return "\n".join(lines) + + def __repr__(self) -> str: + return f"Plan({self.direction})\n{self.describe()}" + + +def _names(attributes) -> str: + return "unknown attributes" if attributes is None else "{" + ", ".join(sorted(attributes)) + "}" + + +def _node_key(node): + """A hashable key for reusing identical nodes; None for nodes that are never shared.""" + if isinstance(node, Seed): + return None + if isinstance(node, Literal): + if node.keys is None: + return None # each iterating group has its own keys + return (Literal, node.table, tuple(tuple(sorted(k.items(), key=lambda kv: kv[0])) for k in node.keys), node.read) + return node + + +# ============================================================================= +# The planner +# ============================================================================= + + +def plan_expand(schema: TraversalSchema, seed, direction: str = "down") -> Plan: + """ + Plan ``expand(seed, direction)``. + + Parameters + ---------- + schema : TraversalSchema + The schema; its group quotient must be acyclic. + seed : Seed or Literal + The seed node. + direction : {"down", "up", "both"} + ``"down"`` plans ``D(S)``, ``"up"`` plans ``U(S)``, and ``"both"`` plans + ``U(D(S))``: the downward pass, then an upward pass seeded with every + table's rows from it. + + Returns + ------- + Plan + + Notes + ----- + diagram.md, "One pass over the group quotient": each node of Q is visited + once, in topological order for ``down`` and in reverse order for ``up``. A + plain table gets its seed and R1 from each restricted neighbor on the + incoming side. A group's master gets the seed on the master, R1 into the + master, and the lift of R1 into each part and of the seed on each part; each + part then gets the re-expansion of the master. A group with a foreign key + that links instances is marked to iterate. + """ + if direction not in DIRECTIONS: + raise ValueError(f"direction must be 'down', 'up' or 'both', not {direction!r}") + if seed.table not in schema.tables: + raise DataJointError(f"The seed table {seed.table} is not in the schema description.") + plan = Plan(schema, direction) + plan.seed = plan.add(seed) + seeds = {seed.table: [plan.seed]} + if direction == "both": + down = _plan_pass(plan, seeds, "down") + plan.restriction = _plan_pass(plan, {t: [h] for t, h in down.items()}, "up") + else: + plan.restriction = _plan_pass(plan, seeds, direction) + plan.order = tuple(schema.order(plan.restriction)) + return plan + + +def _plan_pass(plan: Plan, seeds: dict, direction: str) -> dict: + """One pass over Q; returns the handle of each reached table's rows.""" + schema = plan.schema + order = schema.quotient_order(schema.reach(seeds, direction)) + if direction == "up": + order.reverse() + restriction: dict[str, int] = {} + + def incoming(table): + # The seed on the table and R1 from each restricted neighbor outside its group. + terms = list(seeds.get(table, ())) + group = schema.group_of.get(table) + for fk in schema.fks_of(table) if direction == "down" else schema.fks_to(table): + other = fk.parent if direction == "down" else fk.child + if other in restriction and (group is None or schema.group_of.get(other) != group): + terms.append(_cross(plan, fk, restriction[other], direction)) + return terms + + for q in order: + parts = schema.groups.get(q) + if parts is None: + restriction[q] = plan.union(q, incoming(q)) + continue + # R2: the master's rows are its own incoming rows and the lift of each part's. + terms = incoming(q) + for part in parts: + part_terms = incoming(part) + if part_terms: + terms.append(plan.add(Lift(q, plan.union(part, part_terms), _lift_map(schema, part)))) + entry = plan.union(q, terms) + links = schema.links(q) + keys = step = None + if links: + # A foreign key that links instances: one more group step from the + # instances of the current master keys, across those foreign keys. + keys = plan.add(Literal(q, None, read=True)) + instance = {q: keys} + instance.update({p: plan.add(Reexpand(p, keys, _reexpand_map(schema, p))) for p in parts}) + reached = [] + for fk in links: + source, target = (fk.parent, fk.child) if direction == "down" else (fk.child, fk.parent) + term = _cross(plan, fk, instance[source], direction) + reached.append(term if target == q else plan.add(Lift(q, term, _lift_map(schema, target)))) + step = plan.union(q, reached) + restriction[q] = entry if keys is None else keys + for part in parts: + restriction[part] = plan.add(Reexpand(part, restriction[q], _reexpand_map(schema, part))) + plan.groups.append(GroupStep(q, parts, direction, entry, links, keys, step)) + return restriction + + +def _cross(plan: Plan, fk: ForeignKey, src: int, direction: str) -> int: + """R1 across ``fk`` from the restricted table ``src``, in the direction of travel.""" + schema = plan.schema + node = plan.nodes[src] + literal = isinstance(node, Literal) and node.keys is not None + if direction == "down": + if literal: + if schema.relabel_down_allowed(fk, node.keys): + return plan.add(Literal(fk.child, tuple(_relabel_down(fk, node.keys)))) + elif schema.copy_allowed(fk, plan.attributes(src)): + return plan.add(Restrict(fk.child, src, tuple((c, c) for c in fk.attributes), copy=True)) + return plan.add(Restrict(fk.child, src, fk.attr_map)) + # Upward, R1 is always a restriction: a condition on the child describes child + # rows. Only keys read from the child itself may be relabeled. + if literal and node.read and all(set(fk.attributes) <= set(key) for key in node.keys): + return plan.add(Literal(fk.parent, tuple(_relabel_up(fk, node.keys)))) + return plan.add(Restrict(fk.parent, src, tuple((p, c) for c, p in fk.attr_map))) + + +def _lift_map(schema: TraversalSchema, part: str) -> tuple: + owning = schema.owning[part] + return tuple((a, owning[a]) for a in schema.tables[schema.group_of[part]].primary_key) + + +def _reexpand_map(schema: TraversalSchema, part: str) -> tuple: + return tuple((o, m) for m, o in _lift_map(schema, part)) + + +def _relabel_down(fk: ForeignKey, keys) -> list[dict]: + by_parent = {p: c for c, p in fk.attr_map} + return [{by_parent[a]: v for a, v in key.items() if a in by_parent} for key in keys] + + +def _relabel_up(fk: ForeignKey, keys) -> list[dict]: + mapping = fk.mapping + return [{mapping[c]: v for c, v in key.items() if c in mapping} for key in keys] + + +# ============================================================================= +# Evaluation +# ============================================================================= + + +class _Evaluator: + """ + Shared evaluation logic: memoized node values and the groups that iterate. + + A value that depends on the keys of a group that iterates is computed once + those keys are known. The keys are found by repeating the group step: they + start as the keys of the group's ``entry``, and the keys of its ``step``, + evaluated from the current keys, are added until no new key appears. + """ + + def __init__(self, plan: Plan) -> None: + self.plan = plan + self.schema = plan.schema + self._memo: dict[int, Any] = {} + self._depends: dict[int, frozenset] = {} + self._keys: dict[int, set] = {} # handle of a group's Literal(keys=None) -> final set of key tuples + self._trial: dict[int, set] = {} # the same, while the group step repeats + self._groups = {g.keys: g for g in plan.groups if g.keys is not None} + #: For each group that iterates, ``(direction, master)``: the master keys after each round. + self.rounds: dict[tuple, list] = {} + + def value(self, handle: int): + """The value of a node: rows in memory, or a query expression.""" + if handle in self._memo: + return self._memo[handle] + depends = self._placeholders(handle) + for p in depends: + if p not in self._keys and p not in self._trial: + self._iterate(self._groups[p]) + result = self._compute(handle) + if all(p in self._keys for p in depends): + self._memo[handle] = result + return result + + def table(self, name: str): + """The value of a table of the expanded diagram.""" + return self.value(self.plan.restriction[name]) + + def result(self) -> dict: + """The value of every table of the expanded diagram, in topological order.""" + return {name: self.table(name) for name in self.plan.order} + + def group_keys(self, handle: int) -> list[dict]: + """The current master keys of a group that iterates, as key dicts.""" + keys = self._trial[handle] if handle in self._trial else self._keys[handle] + pk = self.schema.tables[self.plan.nodes[handle].table].primary_key + return [dict(zip(pk, k)) for k in sorted(keys, key=repr)] + + def _placeholders(self, handle: int) -> frozenset: + """The keys of the groups that iterate that a node depends on.""" + if handle not in self._depends: + node = self.plan.nodes[handle] + if isinstance(node, Literal) and node.keys is None: + found = frozenset({handle}) + elif isinstance(node, _MATCHING): + found = self._placeholders(node.src) + elif isinstance(node, Or): + found = frozenset().union(*(self._placeholders(t) for t in node.terms)) + else: + found = frozenset() + self._depends[handle] = found + return self._depends[handle] + + def _iterate(self, group: GroupStep) -> None: + handle, step = group.keys, group.step + assert handle is not None and step is not None, "only a group that iterates is iterated" + keys = set(self._key_tuples(group.entry)) + history = [set(keys)] + while True: + self._trial[handle] = keys + new = set(self._key_tuples(step)) - keys + if not new: + break + keys = keys | new + history.append(set(keys)) + del self._trial[handle] + self._keys[handle] = keys + self.rounds[(group.direction, group.master)] = history + + def _compute(self, handle: int): # pragma: no cover - implemented by subclasses + raise NotImplementedError + + def _key_tuples(self, handle: int): # pragma: no cover - implemented by subclasses + raise NotImplementedError + + +def _dict_matches(key: Mapping, row: Mapping, heading) -> bool: + """A key dict as a restriction: keys that are not attributes are ignored, and None matches NULL.""" + return all(row[k] == v for k, v in key.items() if k in heading) + + +class RowEvaluator(_Evaluator): + """ + Evaluate a plan on rows held in memory. + + Parameters + ---------- + plan : Plan + The plan. + db : dict + ``{table: {pk_tuple: row_dict}}``, where ``pk_tuple`` lists the + primary-key values in key order; the format of the reference model in the + test suite. + seed_match : callable, optional + ``seed_match(condition, row, table_info) -> bool`` decides whether a row + passes the seed's condition. By default the condition is a key dict or a + list of key dicts. + fixed : dict[int, set], optional + Nodes whose rows are given, as sets of primary-key tuples: the delete + simulation passes the nodes it materialized. + keys : dict[int, set], optional + The master keys of groups that iterate, found already, as returned by + ``group_key_sets``. + + Notes + ----- + A value is a frozenset of primary-key tuples. A restriction by a projection + matches the attributes it pairs, and a row with a NULL among them matches + nothing (A2). A copy evaluates the copied condition on the child's own rows, + so that the copy rule is tested rather than assumed. + """ + + def __init__(self, plan: Plan, db: Mapping, seed_match: Callable | None = None, fixed=None, keys=None) -> None: + super().__init__(plan) + self.db = db + self.seed_match = seed_match or (lambda cond, row, table: _literal_matches(cond, row, table.attributes)) + self.fixed = dict(fixed or {}) + self._keys.update({h: set(k) for h, k in (keys or {}).items()}) + self._projected: dict[int, frozenset] = {} + + def group_key_sets(self) -> dict: + """The master keys found for each group that iterates, to reuse on another database state.""" + for g in self.plan.iterating_groups(): + if g.keys not in self._keys: + self._iterate(g) + return {h: set(k) for h, k in self._keys.items()} + + def _key_tuples(self, handle: int): + return self.value(handle) + + def _compute(self, handle: int) -> frozenset: + if handle in self.fixed: + return frozenset(self.fixed[handle]) + node = self.plan.nodes[handle] + rows = self.db[node.table] + if isinstance(node, Or): + return frozenset().union(*(self.value(t) for t in node.terms)) + if isinstance(node, _MATCHING) and not getattr(node, "copy", False): + projected = self._projection(handle) + names = [a for a, _ in node.attr_map] + return frozenset(k for k, row in rows.items() if _value_in(tuple(row[a] for a in names), projected)) + return frozenset(k for k, row in rows.items() if self.holds(handle, node.table, row)) + + def _projection(self, handle: int) -> frozenset: + """The source rows of a Restrict, Lift or Reexpand, projected onto the attributes it matches.""" + if handle in self._projected: + return self._projected[handle] + node = self.plan.nodes[handle] + table = self.plan.nodes[node.src].table + rows, pk = self.db[table], self.schema.tables[table].primary_key + values = set() + for k in self.value(node.src): + # A materialized source can name rows deleted since; their keys still hold the key attributes. + row = rows.get(k) or dict(zip(pk, k)) + if all(s in row for _, s in node.attr_map): + values.add(tuple(row[s] for _, s in node.attr_map)) + projected = frozenset(values) + if all(p in self._keys for p in self._placeholders(node.src)): + self._projected[handle] = projected + return projected + + def holds(self, handle: int, table: str, row: Mapping) -> bool: + """Whether ``row`` of ``table`` passes the condition of node ``handle``.""" + info = self.schema.tables[table] + node = self.plan.nodes[handle] + if handle in self.fixed: + # A materialized node: its keys, as a restriction of this table. + pk = self.schema.tables[node.table].primary_key + return any(_dict_matches(dict(zip(pk, k)), row, info.attributes) for k in self.fixed[handle]) + if isinstance(node, Seed): + return self.seed_match(node.condition, row, info) + if isinstance(node, Literal): + keys = node.keys if node.keys is not None else self.group_keys(handle) + return any(_dict_matches(key, row, info.attributes) for key in keys) + if isinstance(node, Restrict) and node.copy: + return self.holds(node.src, table, row) + if isinstance(node, _MATCHING): + # A restriction by a query expression matches on the shared attributes. + shared = [i for i, (a, _) in enumerate(node.attr_map) if a in info.attributes] + projected = self._projection(handle) + if not shared: + return bool(projected) + value = tuple(row[node.attr_map[i][0]] for i in shared) + return _value_in(value, frozenset(tuple(p[i] for i in shared) for p in projected)) + return any(self.holds(t, table, row) for t in node.terms) + + +def _value_in(value: tuple, values: frozenset) -> bool: + """SQL's ``(a, b) IN (...)``: a NULL on either side matches nothing.""" + return None not in value and value in values + + +def _literal_matches(cond, row: Mapping, heading) -> bool: + if isinstance(cond, Mapping): + return _dict_matches(cond, row, heading) + return any(_literal_matches(c, row, heading) for c in cond) + + +def evaluate_rows(plan: Plan, db: Mapping, seed_match: Callable | None = None) -> dict: + """ + Evaluate a plan on rows held in memory. + + Returns + ------- + dict[str, frozenset] + For each table of the expanded diagram, the primary-key tuples of its rows. + """ + return RowEvaluator(plan, db, seed_match).result() + + +class QueryCompiler(_Evaluator): + """ + Compile a plan to query expressions. + + Parameters + ---------- + plan : Plan + A plan whose table names are full table names. + connection : Connection + The database connection. + tables : dict, optional + ``FreeTable`` objects to reuse, by full table name, so that headings are + loaded once. + + Notes + ----- + Each node is compiled to a ``FreeTable`` of its table with one condition + appended. A Restrict, Lift or Reexpand compiles to + ``(a1, a2) IN (SELECT s1, s2 FROM )``, which is how + ``table & src.proj(a1="s1", a2="s2")`` compiles; matching by position keeps + other attributes of the source out of the match, also when a rename reuses a + name the source holds (parallel foreign keys). A copy appends the copied + condition, compiled for the child. A seed's conditions are already compiled + and are appended as they are: compiling them again would escape ``%`` twice + and reject the column names of a subquery (#1610). A Literal is a list of + key dicts, and an Or joins its terms' conditions with OR. + + A group that iterates is the only place where compilation reads rows: the + master keys are fetched after each group step. + """ + + def __init__(self, plan: Plan, connection, tables: dict | None = None) -> None: + super().__init__(plan) + self.connection = connection + self._tables = dict(tables or {}) + + def free_table(self, name: str): + """The unrestricted table ``name``.""" + if name not in self._tables: + from .table import FreeTable + + self._tables[name] = FreeTable(self.connection, name) + return self._tables[name] + + def _key_tuples(self, handle: int): + pk = self.schema.tables[self.plan.nodes[handle].table].primary_key + return {tuple(row[a] for a in pk) for row in self.value(handle).keys()} + + def _compute(self, handle: int): + node = self.plan.nodes[handle] + table = self.free_table(node.table) + return _append_condition(table, self.condition(handle, table), self.plan.attributes(handle)) + + def condition(self, handle: int, target) -> str | bool: + """The SQL condition that restricts ``target`` by the condition of node ``handle``.""" + node = self.plan.nodes[handle] + if isinstance(node, Seed): + items = list(node.condition) + if not items: + return True + if any(item is False for item in items): + return False + return "(" + ")AND(".join(str(item) for item in items) + ")" + if isinstance(node, Literal): + keys = node.keys if node.keys is not None else self.group_keys(handle) + return make_condition(target, [dict(k) for k in keys], set()) + if isinstance(node, Restrict) and node.copy: + return self.condition(node.src, target) + if isinstance(node, _MATCHING): + quote = self.connection.adapter.quote_identifier + source = self.value(node.src) + fields = ", ".join(quote(a) for a, _ in node.attr_map) + return f"({fields}) IN ({source.make_sql([s for _, s in node.attr_map])})" + parts = [self.condition(t, target) for t in node.terms] + if any(p is True for p in parts): + return True + parts = [p for p in parts if p is not False] + if not parts: + return False + return "(" + " OR ".join(f"({p})" for p in parts) + ")" + + +def _append_condition(expr, condition, attributes): + """ + Return ``expr`` restricted by an SQL condition that is already compiled. + + ``QueryExpression.restrict`` compiles its argument, so a compiled condition + is appended here directly. ``expr`` is a base table, so no subquery is needed. + """ + if condition is True: + return expr + result = copy.copy(expr) + result._restriction = AndList(expr.restriction) + result._restriction.append(condition) + result._restriction_attributes = set(expr.restriction_attributes) | (set(attributes or ()) & set(expr.heading.names)) + return result + + +def compile_plan(plan: Plan, connection, tables: dict | None = None) -> dict: + """ + Compile a plan to query expressions. + + Returns + ------- + dict[str, QueryExpression] + For each table of the expanded diagram, in topological order, the table + restricted to its rows. + """ + return QueryCompiler(plan, connection, tables).result() + + +# ============================================================================= +# Seeds and the dependency graph +# ============================================================================= + +_SELECT = re.compile(r"\bselect\b", re.IGNORECASE) + + +def _has_subquery(condition) -> bool: + """Whether a compiled condition contains a subquery (as in 2.3.5, #1610).""" + if not isinstance(condition, str): + return False + text = re.sub(r"(\\\")|(\\\')", "", condition) # escaped quotes + text = re.sub(r"'[^']*'|\"[^\"]*\"", "", text) # quoted text and identifiers + return _SELECT.search(text) is not None + + +def seed_from_expression(expr) -> Seed | Literal: + """ + The seed node for a query expression. + + A seed is one table with any restriction (diagram.md, "Seeds"). The + restriction may use other queries, but a projection, a join, an aggregation + or a union is refused as a seed. A ``dj.Top`` restriction is replaced by the + primary keys it selects. + + Parameters + ---------- + expr : QueryExpression or Table class + The seed. + + Returns + ------- + Seed or Literal + A Literal for a ``dj.Top`` seed; otherwise a Seed holding the compiled + conditions. Their attributes are ``restriction_attributes``, unless a + condition contains a subquery: the attributes a subquery names cannot be + told from those of the seed table in compiled text, so such a seed is + never copied. + + Raises + ------ + DataJointError + If the expression is not one table with a restriction. + """ + from .expression import Aggregation, QueryExpression, Union + + if inspect.isclass(expr) and issubclass(expr, QueryExpression): + expr = expr() + if not isinstance(expr, QueryExpression): + raise DataJointError(f"The seed of Diagram.expand must be a table or a restricted table, not {type(expr).__name__}.") + node, top = expr, False + while True: + if isinstance(node, Union): + raise DataJointError("A union cannot seed Diagram.expand. Expand each operand and add the diagrams.") + if isinstance(node, Aggregation): + raise DataJointError("An aggregation cannot seed Diagram.expand. Restrict a table by it instead.") + if len(node.support) != 1 or node._joins: + raise DataJointError( + "A join cannot seed Diagram.expand. Seed one table and restrict it by the other operand instead." + ) + top = top or node._top is not None + if isinstance(node.support[0], str): + break + node = node.support[0] # a subquery made around a dj.Top + name = node.support[0] + from .table import FreeTable + + base = FreeTable(expr.connection, name).heading + if expr.heading.new_attributes or expr.heading.names != base.names or expr.primary_key != base.primary_key: + raise DataJointError( + "A projection cannot seed Diagram.expand. Seed the table with the same restriction, without the projection." + ) + if top: + return Literal(name, tuple(expr.keys()), read=True) + if node is not expr: # pragma: no cover - only a dj.Top makes a subquery around a restriction + raise DataJointError("The seed of Diagram.expand must be one table with a restriction.") + items = tuple(expr.restriction) + attributes = None if any(_has_subquery(c) for c in items) else frozenset(expr.restriction_attributes) + return Seed(name, items, attributes) + + +def graph_reach(graph: nx.MultiDiGraph, table: str, direction: str) -> set[str]: + """ + The tables a traversal of the dependency graph can reach, with every part. + + Like ``TraversalSchema.reach``, with each master grouped with all its parts + by name, owned or not. The result contains the exact reach, and the tables a + schema description needs: describe these, then plan on the description. + """ + parts = defaultdict(set) + for node in graph.nodes: + master = extract_master(node) + if master is not None and master in graph: + parts[master].add(node) + + def closure(tables, step): + seen = set(tables) + stack = list(seen) + while stack: + t = stack.pop() + owner = extract_master(t) + master = owner if owner is not None and owner in parts else t + for n in (set(step(t)) | {master} | parts.get(master, set())) - seen: + seen.add(n) + stack.append(n) + return seen + + if direction == "both": + return closure(closure({table}, graph.successors), graph.predecessors) + _check_direction(direction) + return closure({table}, graph.successors if direction == "down" else graph.predecessors) + + +def load_dependencies(connection, direction: str, schema_names: Iterable[str] = ()): + """ + Load every schema a traversal in ``direction`` can reach (diagram.md, "Unloaded schemas"). + + Starting from the activated schemas and ``schema_names``, schemas that + reference them (``down``) or that they reference (``up``) are added until no + new one appears. For ``both``, the downward closure comes first, then the + upward closure of the result. The dependency graph is reloaded only if a + schema is missing from it. + + Returns + ------- + Dependencies + ``connection.dependencies``. + """ + adapter = connection.adapter + known = set(connection.schemas) | set(schema_names) + for step in ("down", "up") if direction == "both" else (direction,): + find = adapter.find_downstream_schemas_sql if step == "down" else adapter.find_upstream_schemas_sql + while True: + listed = ", ".join(adapter.quote_string(s) for s in known) + new = {row[0] for row in connection.query(find(listed))} - known + if not new: + break + known |= new + dependencies = connection.dependencies + if not (dependencies._loaded and known <= dependencies._loaded_schemas): + dependencies.load(force=True, schema_names=known) + return dependencies diff --git a/tests/unit/test_expand_plan.py b/tests/unit/test_expand_plan.py new file mode 100644 index 000000000..2f278730b --- /dev/null +++ b/tests/unit/test_expand_plan.py @@ -0,0 +1,1059 @@ +""" +Tests of the traversal planner (``datajoint.expand``) against the reference model. + +The planner builds a plan, an intermediate form, from a schema description and +a seed (the "Computation" section of the traversal specification, ``diagram.md`` +in datajoint-docs). This module checks it without a database: + +- **The schema description** agrees with the reference model: owned parts, + owning attributes, groups that iterate, R1's copy rule, the relabeling rules, + the table-level reach and the group quotient. +- **The plan, evaluated on rows held in memory**, selects exactly the model's + ``D(S)``, ``U(S)`` and ``U(D(S))``, for random schemas, databases and seeds of + every kind (a condition, literal key dicts, keys read from the seed table). + Copies are evaluated as copies, by applying the copied condition to the + child's rows, so the copy rule is tested and not assumed. +- **A delete simulation** runs the plan in reverse topological order, evaluating + each table's restriction on the database as it is at that moment, with the + materialized nodes replaced by the rows they selected before the first + DELETE, and refusing a DELETE that would leave a reference to a deleted row + (``ON DELETE RESTRICT``). It must leave ``DB - D(S)``. +- **The plan's structure** for the shapes of the catalog (``traversal-shapes.md``): + where a condition is copied and where it is restricted, the lifts, and which + groups iterate. + +``HYPOTHESIS_PROFILE`` selects the number of examples (``tests/conftest.py``). +""" + +from __future__ import annotations + +import copy + +import pytest +from hypothesis import assume, event, example, find, given, settings +from hypothesis import strategies as st + +from datajoint import dependencies +from datajoint.errors import DataJointError +from datajoint.expand import ( + ForeignKey, + Lift, + Literal, + Or, + Plan, + Reexpand, + Restrict, + RowEvaluator, + Seed, + TableInfo, + TraversalSchema, + evaluate_rows, + plan_expand, +) + +from .expand_model import ( + FK, + RowGraph, + Schema, + TableSpec, + apply_delete, + condition_attributes, + copy_allowed, + dependency_graph, + fk_value, + full_table_name, + iterates, + make_db, + q_acyclic, + relabel_down_allowed, + restrict, + satisfies, + table_reach, +) +from .expand_strategies import Case, cases, conditions, key_dicts, schemas + +DIRECTIONS = ["down", "up", "both"] + + +# ============================================================================= +# Helpers +# ============================================================================= + + +def describe(schema: Schema) -> TraversalSchema: + """The planner's schema description of a model schema.""" + tables = [TableInfo(t.name, t.pk, t.heading, schema.nullable[t.name], t.master) for t in schema.tables] + return TraversalSchema(tables, [ForeignKey(fk.child, fk.parent, fk.attr_map) for fk in schema.fks]) + + +def model_match(cond, row, table: TableInfo) -> bool: + """A seed condition in the model's format (dict, list, Relation, Not), as a restriction.""" + return satisfies(row, table.attributes, cond) + + +def by_table(rows) -> dict: + """``{table: {pk, ...}}`` from model row identifiers.""" + result = {} + for t, k in rows: + result.setdefault(t, set()).add(k) + return result + + +def expected(case, seed_rows, direction): + op = {"down": case.graph.D, "up": case.graph.U, "both": case.graph.both}[direction] + return by_table(op(seed_rows)) + + +def evaluated(plan: Plan, db) -> dict: + """The plan's rows in memory, leaving out tables with none.""" + return {t: set(rows) for t, rows in evaluate_rows(plan, db, model_match).items() if rows} + + +@st.composite +def seed_nodes(draw, case, parts_first: bool = False): + """ + Draw a seed table and a seed node of one of the kinds the planner handles. + + - a condition (a dict, a list of dicts, a query expression, a Not), whose + attributes are those the model's ``condition_attributes`` reads; + - literal key dicts (a materialized dict or list of dicts); + - keys read from rows of the seed table, over any of its attributes; + - no restriction. + + With ``parts_first``, the seed table is an owned part when the schema has one, + so that the seed lifts to its master. Returns the seed node and the rows it + selects. + """ + schema = case.schema + parts = [t.name for t in schema.tables if t.master and t.name in schema.group] + table = draw(st.sampled_from(parts if parts_first and parts else [t.name for t in schema.tables]), label="seed table") + heading = schema.table(table).heading + names = list(heading) * 2 + sorted({a for t in schema.tables for a in t.heading}) + ["unrelated"] + kind = draw(st.sampled_from(["condition", "literal", "read", "all"]), label="seed kind") + if kind == "condition": + cond = draw(conditions(names), label="condition") + seed = Seed(table, cond, frozenset(condition_attributes(cond))) + elif kind == "literal": + cond = draw(st.lists(key_dicts(names), max_size=3), label="keys") + seed = Literal(table, tuple(cond)) + elif kind == "read": + rows = sorted(case.db[table]) + picked = draw(st.lists(st.sampled_from(rows), unique=True, max_size=3) if rows else st.just([])) + attrs = draw(st.sets(st.sampled_from(heading)), label="attributes read") + cond = [{a: case.db[table][k][a] for a in sorted(attrs)} for k in picked] + seed = Literal(table, tuple(cond), read=True) + else: + cond = {} + seed = Seed(table, cond, frozenset()) + return seed, restrict(schema, case.db, table, cond) + + +@st.composite +def seeded_cases(draw, parts_first: bool = False): + """A drawn case with a seed node: ``(case, seed, seed_rows)``.""" + case = draw(cases()) + return (case, *draw(seed_nodes(case, parts_first))) + + +def lifts_evaluated_at_delete(plan: Plan) -> set: + """ + The Lift nodes that a delete evaluates when it reaches a table. + + A group that iterates has its master keys found before the first DELETE, so + the lifts behind them are not evaluated again. + """ + found, seen = set(), set() + stack = list(plan.restriction.values()) + while stack: + h = stack.pop() + if h in seen: + continue + seen.add(h) + node = plan.nodes[h] + if isinstance(node, Lift): + found.add(h) + if isinstance(node, (Restrict, Lift, Reexpand)): + stack.append(node.src) + elif isinstance(node, Or): + stack.extend(node.terms) + return found + + +class RestrictViolation(Exception): + """A DELETE would leave a row that references a deleted row (``ON DELETE RESTRICT``).""" + + +def simulate_delete(plan: Plan, model: Schema, db, match_for=lambda db: model_match, materialize=()): + """ + Run a ``down`` plan as ``Table.delete`` runs a cascade, on a copy of ``db``. + + Before the first DELETE, the nodes in ``materialize`` are replaced by the rows + they select, and the master keys of the groups that iterate are found (as + compiling the plan does). Then each table, leaves first, deletes the rows its + restriction selects on the database as it is at that moment. A DELETE that + would leave a row referencing a deleted row raises ``RestrictViolation``. + + ``match_for(db)`` returns the seed matcher for a database state, so that a + seed whose condition reads other tables reads them as they are. + """ + first = RowEvaluator(plan, db, match_for(db)) + keys = first.group_key_sets() + fixed = {h: first.value(h) for h in materialize} + current = copy.deepcopy(db) + for table in reversed(plan.order): + doomed = RowEvaluator(plan, current, match_for(current), fixed=fixed, keys=keys).table(table) + for fk in model.fks: + if fk.parent != table: + continue + for k, row in current[fk.child].items(): + if fk.child == table and k in doomed: + continue + if fk_value(model, fk, row) in doomed: + raise RestrictViolation(f"deleting {table} {fk_value(model, fk, row)} leaves {fk.child} {k}") + current[table] = {k: r for k, r in current[table].items() if k not in doomed} + return current + + +def spec_materialization(plan: Plan) -> set: + """ + What diagram.md ("Materialization happens only at delete time") materializes: + every lift, and the seed unless its restriction is literal keys. + """ + flagged = set(lifts_evaluated_at_delete(plan)) + if isinstance(plan.nodes[plan.seed], Seed): + flagged.add(plan.seed) + return flagged + + +# ============================================================================= +# The schema description agrees with the model +# ============================================================================= + + +@given(schema=schemas()) +def test_description_agrees_with_the_model(schema): + """Owned parts, owning attributes, the groups that iterate, and the group quotient.""" + d = describe(schema) + model_groups = {m: parts for m, parts in schema.owned_parts.items() if parts} + assert {m: set(p) for m, p in d.groups.items()} == {m: set(p) for m, p in model_groups.items()} + for table in d.group_of: + assert d.owning_attributes(table) == schema.owning_attributes(table), table + for master in d.groups: + assert bool(d.links(master)) == iterates(schema, master) + if q_acyclic(schema): + order = d.order(d.tables) + position = {t: i for i, t in enumerate(order)} + assert all(position[fk.parent] < position[fk.child] for fk in schema.fks) + for master, parts in d.groups.items(): + assert order[position[master] : position[master] + 1 + len(parts)] == [master, *parts] + else: + with pytest.raises(DataJointError, match="cycle"): + d.quotient_order(d.tables) + + +@given(schema=schemas()) +def test_reach_is_the_table_level_reach(schema): + """The node set of an expanded diagram, in each direction.""" + d = describe(schema) + for t in schema.tables: + for direction in DIRECTIONS: + assert d.reach({t.name}, direction) == table_reach(schema, {t.name}, direction) + + +@given(case=cases(), data=st.data()) +def test_copy_and_relabel_rules_agree_with_the_model(case, data): + """R1's copy rule and the Seeds rule for relabeling keys, as the model states them.""" + d = describe(case.schema) + fks = [fk for fk in case.schema.fks if fk.child != fk.parent] + assume(fks) + fk = data.draw(st.sampled_from(fks)) + names = list(fk.attr_map) * 3 + sorted({a for t in case.schema.tables for a in t.heading}) + ["unrelated"] + cond = data.draw(conditions(names)) + keys = data.draw(st.lists(key_dicts(names), max_size=3)) + ours = ForeignKey(fk.child, fk.parent, fk.attr_map) + assert d.copy_allowed(ours, frozenset(condition_attributes(cond))) == copy_allowed(case.schema, fk, cond) + assert d.relabel_down_allowed(ours, keys) == relabel_down_allowed(case.schema, fk, keys) + + +class _Heading: + """The parts of a ``Heading`` that ``TraversalSchema.from_dependencies`` reads.""" + + class _Attribute: + def __init__(self, nullable): + self.nullable = nullable + + def __init__(self, spec: TableSpec, nullable): + self.names = list(spec.heading) + self.primary_key = list(spec.pk) + self.attributes = {a: self._Attribute(a in nullable) for a in spec.heading} + + +@pytest.mark.parametrize("quote", ["`", '"']) +@given(schema=schemas()) +def test_description_from_the_dependency_graph(quote, schema): + """ + Built from ``connection.dependencies``, the description has the same groups + and owning attributes; the order of the foreign keys comes from the headings. + """ + graph = dependency_graph(schema, quote) + name = {t.name: full_table_name(t.name, quote) for t in schema.tables} + headings = {name[t.name]: _Heading(t, schema.nullable[t.name]) for t in schema.tables} + d = TraversalSchema.from_dependencies(graph, headings, headings.__getitem__) + expected = describe(schema) + assert {m: set(p) for m, p in d.groups.items()} == {name[m]: {name[p] for p in ps} for m, ps in expected.groups.items()} + for t in expected.group_of: + assert d.owning_attributes(name[t]) == expected.owning_attributes(t) + assert {n: d.tables[n].nullable for n in d.tables} == {name[t]: v.nullable for t, v in expected.tables.items()} + + +@given(schema=schemas()) +def test_owned_parts_agree_with_dependencies(schema): + """ + ``datajoint.dependencies.owned_parts`` (#1620) finds the same owned parts in + the dependency graph. Skipped until that function exists on this branch. + """ + owned_parts = getattr(dependencies, "owned_parts", None) + if owned_parts is None: + pytest.skip("datajoint.dependencies.owned_parts is not on this branch (#1620)") + graph = dependency_graph(schema) + d = describe(schema) + assert owned_parts(graph) == {full_table_name(p) for parts in d.groups.values() for p in parts} + + +# ============================================================================= +# The plan evaluated in memory equals the definitions +# ============================================================================= + + +@pytest.mark.parametrize("direction", DIRECTIONS) +@given(case=cases(), data=st.data()) +def test_plan_equals_the_model(direction, case, data): + """ + The plan, evaluated on rows, selects the model's D(S), U(S) or U(D(S)), and + its tables are the table-level reach of the seed table. + """ + seed, seed_rows = data.draw(seed_nodes(case), label="seed") + plan = plan_expand(describe(case.schema), seed, direction) + assert set(plan.order) == table_reach(case.schema, {seed.table}, direction) + assert evaluated(plan, case.db) == expected(case, seed_rows, direction) + + +@given(case=cases(), data=st.data()) +def test_seed_kinds_give_the_same_rows(case, data): + """``plan(T & r)`` and ``plan(T & (T & r).keys())`` select the same rows (K3).""" + seed, seed_rows = data.draw(seed_nodes(case), label="seed") + pk = case.schema.table(seed.table).pk + keys = Literal(seed.table, tuple(dict(zip(pk, k)) for _, k in sorted(seed_rows)), read=True) + d = describe(case.schema) + for direction in DIRECTIONS: + assert evaluated(plan_expand(d, seed, direction), case.db) == evaluated(plan_expand(d, keys, direction), case.db) + + +@pytest.mark.parametrize("direction", ["down", "up"]) +@given(case=cases(), data=st.data()) +def test_groups_iterate_by_the_model_rule(direction, case, data): + """A plan marks a group to iterate exactly when the model's rule does.""" + seed, _ = data.draw(seed_nodes(case), label="seed") + plan = plan_expand(describe(case.schema), seed, direction) + for group in plan.groups: + assert group.iterate == iterates(case.schema, group.master) + assert (group.keys is None) == (not group.iterate) + assert plan.restriction[group.master] == (group.keys if group.iterate else group.entry) + + +def _plans(seeded, direction): + case, seed, _ = seeded + return plan_expand(describe(case.schema), seed, direction) + + +def _has(kind, direction="down", **fields): + def predicate(seeded): + plan = _plans(seeded, direction) + return any( + type(n) is kind and all(getattr(n, k) == v for k, v in fields.items()) and h != plan.seed + for h, n in enumerate(plan.nodes) + ) + + return predicate + + +PLAN_FEATURES = { + "a copy": _has(Restrict, copy=True), + "a restriction across a renamed foreign key": lambda x: _has(Restrict, copy=False)(x) + and any(not all(a == b for a, b in n.attr_map) for n in _plans(x, "down").nodes if type(n) is Restrict), + "keys relabeled downward": _has(Literal, read=False), + "keys relabeled upward": _has(Literal, direction="up", read=False), + "a lift": _has(Lift), + "a group that iterates downward": lambda x: bool(_plans(x, "down").iterating_groups()), + "a group that iterates upward": lambda x: bool(_plans(x, "up").iterating_groups()), +} + + +@pytest.mark.parametrize("feature", sorted(PLAN_FEATURES)) +def test_drawn_plans_have_every_node_kind(feature): + """The drawn schemas and seeds reach every kind of plan node the properties rely on.""" + find(seeded_cases(), PLAN_FEATURES[feature], settings=settings(max_examples=5000, database=None, derandomize=True)) + + +# ============================================================================= +# The catalog shapes: plan structure and rows +# ============================================================================= + + +class Shape: + """A catalog shape: its model schema and rows, and its description.""" + + def __init__(self, tables, fks, data): + self.model = Schema(tables, fks) + self.db = make_db(self.model, data) + self.schema = describe(self.model) + + def plan(self, table, cond, direction, attributes=None): + attrs = frozenset(condition_attributes(cond)) if attributes is None else attributes + return plan_expand(self.schema, Seed(table, cond, attrs), direction) + + def check(self, plan, expected_rows): + """The plan's rows equal the catalog's (and so the model's).""" + rows = {t: {k if isinstance(k, tuple) else (k,) for k in ks} for t, ks in expected_rows.items() if ks} + assert evaluated(plan, self.db) == rows + assert set(plan.order) == set(expected_rows) + + +def node(plan: Plan, table: str): + """The node that holds the rows of ``table``.""" + return plan.nodes[plan.restriction[table]] + + +def terms(plan: Plan, table: str) -> list: + """The terms of the union that holds the rows of ``table``.""" + n = node(plan, table) + return [plan.nodes[t] for t in n.terms] if isinstance(n, Or) else [n] + + +def s1(): + return Shape( + [TableSpec("subject", ("subject_id",)), TableSpec("session", ("subject_id", "session_id"))], + [FK("session", "subject", {"subject_id": "subject_id"})], + {"subject": [(1,), (2,)], "session": [(1, 1), (1, 2), (2, 1)]}, + ) + + +def test_S1_plan(): + """S1: downward the condition is copied; upward the parent is restricted.""" + s = s1() + down = s.plan("subject", {"subject_id": 1}, "down") + assert node(down, "session") == Restrict("session", down.seed, (("subject_id", "subject_id"),), copy=True) + s.check(down, {"subject": [1], "session": [(1, 1), (1, 2)]}) + up = s.plan("session", {"session_id": 1}, "up") + assert node(up, "subject") == Restrict("subject", up.seed, (("subject_id", "subject_id"),)) + s.check(up, {"session": [(1, 1), (2, 1)], "subject": [1, 2]}) + + +def test_S2_plan(): + """S2: a secondary FK copies downward and restricts upward.""" + s = Shape( + [TableSpec("#rig", ("rig_id",)), TableSpec("session", ("session_id",), ("rig_id",))], + [FK("session", "#rig", {"rig_id": "rig_id"})], + {"#rig": [(1,), (2,), (3,)], "session": [(1, 1), (2, 1), (3, 2)]}, + ) + down = s.plan("#rig", {"rig_id": 1}, "down") + assert node(down, "session").copy + s.check(down, {"#rig": [1], "session": [1, 2]}) + up = s.plan("session", {"session_id": 3}, "up") + assert node(up, "#rig") == Restrict("#rig", up.seed, (("rig_id", "rig_id"),)) + s.check(up, {"session": [3], "#rig": [2]}) + + +def s3(): + return Shape( + [ + TableSpec("subject", ("subject_id",)), + TableSpec("surgery", ("subject_id", "surgery_id")), + TableSpec("session", ("subject_id", "session_id"), ("surgery_id",)), + ], + [ + FK("surgery", "subject", {"subject_id": "subject_id"}), + FK("session", "subject", {"subject_id": "subject_id"}), + FK("session", "surgery", {"subject_id": "subject_id", "surgery_id": "surgery_id"}, nullable=True), + ], + { + "subject": [(1,), (2,)], + "surgery": [(1, 1), (1, 2), (2, 1)], + "session": [(1, 1, 1), (1, 2, None), (1, 3, 2), (2, 1, 1)], + }, + ) + + +def test_S3_plan(): + """S3: a foreign key with a nullable column is crossed by restriction, never copied or relabeled.""" + s = s3() + down = s.plan("surgery", {"subject_id": 1}, "down") + pairs = (("subject_id", "subject_id"), ("surgery_id", "surgery_id")) + assert node(down, "session") == Restrict("session", down.seed, pairs) + s.check(down, {"surgery": [(1, 1), (1, 2)], "session": [(1, 1), (1, 3)]}) + up = s.plan("session", {"subject_id": 1, "session_id": 2}, "up") + s.check(up, {"session": [(1, 2)], "subject": [1], "surgery": []}) + # Literal keys that name the whole foreign key, none NULL, are relabeled; a NULL is not. + fk = s.schema.fks_of("session")[1] + assert s.schema.relabel_down_allowed(fk, [{"subject_id": 1, "surgery_id": 2}]) + assert not s.schema.relabel_down_allowed(fk, [{"subject_id": 1, "surgery_id": None}]) + assert not s.schema.relabel_down_allowed(fk, [{"subject_id": 1}]) + keys = plan_expand(s.schema, Literal("surgery", ({"subject_id": 1, "surgery_id": 2},)), "down") + assert node(keys, "session") == Literal("session", ({"subject_id": 1, "surgery_id": 2},)) + + +def test_S4_plan(): + """S4: a renamed foreign key is crossed through its attr_map.""" + s = Shape( + [TableSpec("neuron", ("neuron_id",)), TableSpec("input", ("source_id", "input_id"), ("neuron_id",))], + [FK("input", "neuron", {"source_id": "neuron_id"})], + {"neuron": [(1,), (2,), (3,)], "input": [(1, 1, 2), (1, 2, 3), (2, 1, 1), (3, 1, 1)]}, + ) + down = s.plan("neuron", {"neuron_id": 1}, "down") + assert node(down, "input") == Restrict("input", down.seed, (("source_id", "neuron_id"),)) + s.check(down, {"neuron": [1], "input": [(1, 1), (1, 2)]}) + up = s.plan("input", {"source_id": 2, "input_id": 1}, "up") + assert node(up, "neuron") == Restrict("neuron", up.seed, (("neuron_id", "source_id"),)) + s.check(up, {"input": [(2, 1)], "neuron": [2]}) + + +def test_S5_plan(): + """S5: parallel foreign keys are two terms of a union; the unrenamed one is a copy.""" + s = Shape( + [TableSpec("neuron", ("neuron_id",)), TableSpec("synapse", ("neuron_id", "target_id"))], + [FK("synapse", "neuron", {"neuron_id": "neuron_id"}), FK("synapse", "neuron", {"target_id": "neuron_id"})], + {"neuron": [(1,), (2,), (3,)], "synapse": [(1, 2), (2, 3), (3, 1)]}, + ) + down = s.plan("neuron", {"neuron_id": 1}, "down") + assert terms(down, "synapse") == [ + Restrict("synapse", down.seed, (("neuron_id", "neuron_id"),), copy=True), + Restrict("synapse", down.seed, (("target_id", "neuron_id"),)), + ] + s.check(down, {"neuron": [1], "synapse": [(1, 2), (3, 1)]}) + s.check(s.plan("neuron", {"neuron_id": 1}, "both"), {"neuron": [1, 2, 3], "synapse": [(1, 2), (3, 1)]}) + up = s.plan("synapse", {"neuron_id": 1, "target_id": 2}, "up") + assert len(terms(up, "neuron")) == 2 + s.check(up, {"synapse": [(1, 2)], "neuron": [1, 2]}) + + +def s7(): + return Shape( + [ + TableSpec("lab", ("lab_id",)), + TableSpec("rig", ("rig_id",), ("lab_id",)), + TableSpec("session", ("lab_id", "session_id"), ("rig_id",)), + ], + [ + FK("rig", "lab", {"lab_id": "lab_id"}), + FK("session", "lab", {"lab_id": "lab_id"}), + FK("session", "rig", {"rig_id": "rig_id"}), + ], + {"lab": [(1,), (2,)], "rig": [(10, 1), (20, 2)], "session": [(1, 1, 20), (2, 1, 10), (2, 2, 20)]}, + ) + + +def test_S7_plan(): + """S7: Rig's copied condition names lab_id, which Session's FK to Rig does not carry, so Rig restricts Session.""" + s = s7() + down = s.plan("lab", {"lab_id": 1}, "down") + rig = down.restriction["rig"] + assert node(down, "rig") == Restrict("rig", down.seed, (("lab_id", "lab_id"),), copy=True) + assert terms(down, "session") == [ + Restrict("session", down.seed, (("lab_id", "lab_id"),), copy=True), + Restrict("session", rig, (("rig_id", "rig_id"),)), + ] + s.check(down, {"lab": [1], "rig": [10], "session": [(1, 1), (2, 1)]}) + s.check(s.plan("session", {"lab_id": 2, "session_id": 1}, "up"), {"session": [(2, 1)], "rig": [10], "lab": [1, 2]}) + s.check(s.plan("rig", {"rig_id": 10}, "both"), {"rig": [10], "session": [(2, 1)], "lab": [1, 2]}) + + +def g1(): + return Shape( + [ + TableSpec("session", ("session_id",)), + TableSpec("session__trial", ("session_id", "trial_id"), master="session"), + TableSpec("review", ("session_id", "review_id")), + ], + [ + FK("session__trial", "session", {"session_id": "session_id"}), + FK("review", "session", {"session_id": "session_id"}), + ], + { + "session": [(1,), (2,), (3,)], + "session__trial": [(1, 1), (1, 2), (2, 1), (3, 1)], + "review": [(1, 1), (2, 1)], + }, + ) + + +def test_G1_plan(): + """G1: the part lifts to its master, and the master re-expands to all its part rows.""" + s = g1() + down = s.plan("session__trial", {"trial_id": 2}, "down") + master = down.restriction["session"] + assert node(down, "session") == Lift("session", down.seed, (("session_id", "session_id"),)) + assert node(down, "session__trial") == Reexpand("session__trial", master, (("session_id", "session_id"),)) + assert node(down, "review") == Restrict("review", master, (("session_id", "session_id"),), copy=True) + assert [g.iterate for g in down.groups] == [False] + s.check(down, {"session__trial": [(1, 1), (1, 2)], "session": [1], "review": [(1, 1)]}) + up = s.plan("session__trial", {"trial_id": 2}, "up") + assert node(up, "session") == Lift("session", up.seed, (("session_id", "session_id"),)) + s.check(up, {"session__trial": [(1, 1), (1, 2)], "session": [1]}) + + +def test_G2_plan(): + """G2: the lift is one step on the master key, whichever chain owns the part; no iteration.""" + s = Shape( + [ + TableSpec("ext", ("ext_id",)), + TableSpec("master", ("master_id",)), + TableSpec("master__q", ("master_id", "q_id"), ("ext_id",), master="master"), + TableSpec("master__p", ("master_id", "p_id"), ("q_id", "ext_id"), master="master"), + ], + [ + FK("master__q", "master", {"master_id": "master_id"}), + FK("master__q", "ext", {"ext_id": "ext_id"}), + FK("master__p", "master", {"master_id": "master_id"}), + FK("master__p", "master__q", {"master_id": "master_id", "q_id": "q_id"}), + FK("master__p", "ext", {"ext_id": "ext_id"}), + ], + { + "ext": [(1,), (2,)], + "master": [(1,), (2,), (3,)], + "master__q": [(1, 10, 1), (2, 20, 2), (3, 30, 2)], + "master__p": [(1, 100, 10, 2), (2, 200, 20, 1), (3, 300, 30, 2)], + }, + ) + down = s.plan("ext", {"ext_id": 1}, "down") + assert [type(t) for t in terms(down, "master")] == [Lift, Lift] + assert not down.iterating_groups() + s.check(down, {"ext": [1], "master__q": [(1, 10), (2, 20)], "master__p": [(1, 100), (2, 200)], "master": [1, 2]}) + up = s.plan("master__p", {"master_id": 1, "p_id": 100}, "up") + s.check(up, {"master__p": [(1, 100)], "master": [1], "master__q": [(1, 10)], "ext": [1, 2]}) + + +def test_G3_plan(): + """G3: only the owning reference lifts; the second, renamed FK to the master links instances, so the group iterates.""" + s = Shape( + [ + TableSpec("session", ("session_id",)), + TableSpec("session__link", ("session_id", "link_id"), ("other_session",), master="session"), + ], + [ + FK("session__link", "session", {"session_id": "session_id"}), + FK("session__link", "session", {"other_session": "session_id"}), + ], + {"session": [(1,), (2,), (3,)], "session__link": [(1, 1, 2), (3, 1, 1)]}, + ) + down = s.plan("session__link", {"session_id": 1}, "down") + (group,) = down.groups + assert group.iterate and [fk.mapping for fk in group.links] == [{"other_session": "session_id"}] + assert down.nodes[group.entry] == Lift("session", down.seed, (("session_id", "session_id"),)) + s.check(down, {"session__link": [(1, 1), (3, 1)], "session": [1, 3]}) + s.check(s.plan("session__link", {"session_id": 1}, "up"), {"session__link": [(1, 1)], "session": [1, 2]}) + + +def g4(): + return Shape( + [ + TableSpec("run", ("run_id",)), + TableSpec("run__step", ("run_id", "step_id"), master="run"), + TableSpec("run__carry", ("run_id", "carry_id"), ("prev_run", "prev_step"), master="run"), + ], + [ + FK("run__step", "run", {"run_id": "run_id"}), + FK("run__carry", "run", {"run_id": "run_id"}), + FK("run__carry", "run__step", {"prev_run": "run_id", "prev_step": "step_id"}), + ], + { + "run": [(1,), (2,), (3,), (4,)], + "run__step": [(1, 1), (2, 1), (3, 1), (4, 1)], + "run__carry": [(2, 1, 1, 1), (3, 1, 2, 1)], + }, + ) + + +def test_G4_plan(): + """G4: the renamed sibling FK links instances; the master keys grow {1}, {1, 2}, {1, 2, 3}.""" + s = g4() + down = s.plan("run", {"run_id": 1}, "down") + (group,) = down.iterating_groups() + assert [(fk.child, fk.parent) for fk in group.links] == [("run__carry", "run__step")] + assert down.restriction["run"] == group.keys + assert down.nodes[group.keys] == Literal("run", None, read=True) + # One group step: carry rows that reference a step row of the current runs, lifted to their run. + step = down.nodes[group.step] + assert isinstance(step, Lift) and down.nodes[step.src].attr_map == (("prev_run", "run_id"), ("prev_step", "step_id")) + evaluator = RowEvaluator(down, s.db, model_match) + assert {t: set(r) for t, r in evaluator.result().items()} == { + "run": {(1,), (2,), (3,)}, + "run__step": {(1, 1), (2, 1), (3, 1)}, + "run__carry": {(2, 1), (3, 1)}, + } + assert evaluator.rounds[("down", "run")] == [{(1,)}, {(1,), (2,)}, {(1,), (2,), (3,)}] + up = s.plan("run", {"run_id": 2}, "up") + evaluator = RowEvaluator(up, s.db, model_match) + evaluator.result() + assert evaluator.rounds[("up", "run")] == [{(2,)}, {(1,), (2,)}] + s.check(up, {"run": [1, 2], "run__step": [(1, 1), (2, 1)], "run__carry": [(2, 1)]}) + s.check( + s.plan("run", {"run_id": 2}, "both"), + {"run": [1, 2, 3], "run__step": [(1, 1), (2, 1), (3, 1)], "run__carry": [(2, 1), (3, 1)]}, + ) + # Both passes iterate, each with its own keys. + both = s.plan("run", {"run_id": 2}, "both") + assert [(g.direction, g.iterate) for g in both.groups] == [("down", True), ("up", True)] + + +def test_G5_plan(): + """G5: a trace reaches the parts of an ancestor master and their parents; a node may hold zero rows.""" + s = Shape( + [ + TableSpec("ephys_source", ("ephys_id",)), + TableSpec("imaging_source", ("imaging_id",)), + TableSpec("recording", ("recording_id",)), + TableSpec("recording__ephys", ("recording_id",), ("ephys_id",), master="recording"), + TableSpec("recording__imaging", ("recording_id",), ("imaging_id",), master="recording"), + TableSpec("analysis", ("recording_id",)), + ], + [ + FK("recording__ephys", "recording", {"recording_id": "recording_id"}), + FK("recording__ephys", "ephys_source", {"ephys_id": "ephys_id"}), + FK("recording__imaging", "recording", {"recording_id": "recording_id"}), + FK("recording__imaging", "imaging_source", {"imaging_id": "imaging_id"}), + FK("analysis", "recording", {"recording_id": "recording_id"}), + ], + { + "ephys_source": [(1,), (2,)], + "imaging_source": [(1,)], + "recording": [(1,), (2,), (3,)], + "recording__ephys": [(1, 1), (2, 2)], + "recording__imaging": [(3, 1)], + "analysis": [(1,), (2,), (3,)], + }, + ) + up = s.plan("analysis", {"recording_id": 1}, "up") + assert node(up, "ephys_source") == Restrict( + "ephys_source", up.restriction["recording__ephys"], (("ephys_id", "ephys_id"),) + ) + expected_up = { + "analysis": [1], + "recording": [1], + "recording__ephys": [1], + "recording__imaging": [], + "ephys_source": [1], + "imaging_source": [], + } + s.check(up, expected_up) + down = s.plan("ephys_source", {"ephys_id": 2}, "down") + assert isinstance(node(down, "recording"), Lift) + s.check( + down, + {"ephys_source": [2], "recording__ephys": [2], "recording": [2], "recording__imaging": [], "analysis": [2]}, + ) + + +def test_G6_plan(): + """G6: a part its master does not own is a plain table: no group, no lift.""" + s = Shape( + [TableSpec("probe", ("probe_id",)), TableSpec("probe__channel", ("channel_id",), ("probe_id",), master="probe")], + [FK("probe__channel", "probe", {"probe_id": "probe_id"})], + {"probe": [(1,), (2,)], "probe__channel": [(1, 1), (2, 1), (3, 2)]}, + ) + assert s.schema.groups == {} + down = s.plan("probe__channel", {"channel_id": 1}, "down") + assert not down.groups + s.check(down, {"probe__channel": [1]}) + s.check(s.plan("probe__channel", {"channel_id": 1}, "up"), {"probe__channel": [1], "probe": [1]}) + s.check(s.plan("probe", {"probe_id": 1}, "down"), {"probe": [1], "probe__channel": [1, 2]}) + + +def g7(): + return Shape( + [ + TableSpec("master", ("master_id",)), + TableSpec("master__part_a", ("master_id", "part_a_id"), master="master"), + TableSpec("master__part_b", ("src_master", "src_part", "part_b_id"), master="master"), + TableSpec("report", ("master_id", "report_id")), + ], + [ + FK("master__part_a", "master", {"master_id": "master_id"}), + FK("master__part_b", "master__part_a", {"src_master": "master_id", "src_part": "part_a_id"}), + FK("report", "master", {"master_id": "master_id"}), + ], + { + "master": [(1,), (2,), (3,)], + "master__part_a": [(1, 10), (1, 11), (2, 20)], + "master__part_b": [(1, 10, 100), (1, 11, 110), (2, 20, 200)], + "report": [(1, 1), (2, 1)], + }, + ) + + +def test_G7_plan(): + """G7: the lift renames the owning attribute onto the master's key; no iteration.""" + s = g7() + down = s.plan("master__part_b", {"part_b_id": 100}, "down") + assert node(down, "master") == Lift("master", down.seed, (("master_id", "src_master"),)) + assert node(down, "master__part_b") == Reexpand( + "master__part_b", down.restriction["master"], (("src_master", "master_id"),) + ) + assert not down.iterating_groups() + s.check( + down, + { + "master__part_b": [(1, 10, 100), (1, 11, 110)], + "master": [1], + "master__part_a": [(1, 10), (1, 11)], + "report": [(1, 1)], + }, + ) + + +def k_shape(): + return Shape( + [ + TableSpec("subject", ("subject_id",), ("species",)), + TableSpec("session", ("subject_id", "session_id"), ("quality",)), + TableSpec("trial", ("subject_id", "session_id", "trial_id")), + TableSpec("flag", ("flag_id",), ("subject_id",)), + ], + [ + FK("session", "subject", {"subject_id": "subject_id"}), + FK("trial", "session", {"subject_id": "subject_id", "session_id": "session_id"}), + FK("flag", "subject", {"subject_id": "subject_id"}, nullable=True), + ], + { + "subject": [(1, "mouse"), (2, "rat"), (3, "mouse"), (4, "rat")], + "session": [(1, 1, 0), (1, 2, 1), (2, 1, 1), (2, 2, 0), (3, 1, 1)], + "trial": [(1, 1, 1), (1, 2, 1), (2, 1, 1), (2, 2, 1), (3, 1, 1)], + "flag": [(1, 1), (2, 3)], + }, + ) + + +class SessionsOfQuality: + """``Subject & (Session & "quality = ")`` as a live condition: it reads Session as it is.""" + + def __init__(self, quality): + self.quality = quality + + +def k1_match_for(db): + def match(cond, row, table): + if isinstance(cond, SessionsOfQuality): + return any(s["subject_id"] == row["subject_id"] and s["quality"] == cond.quality for s in db["session"].values()) + return model_match(cond, row, table) + + return match + + +K1_DOWN = { + "subject": [1, 2], + "session": [(1, 1), (1, 2), (2, 1), (2, 2)], + "trial": [(1, 1, 1), (1, 2, 1), (2, 1, 1), (2, 2, 1)], + "flag": [1], +} + + +def test_K1_plan(): + """ + K1: a seed restricted by a descendant. Its condition, a query expression, + names the whole heading of Session, so it is not copied; nor is the seed + copied across Flag's nullable foreign key. + """ + s = k_shape() + seed = Seed("subject", SessionsOfQuality(0), frozenset({"subject_id", "session_id", "quality"})) + down = plan_expand(s.schema, seed, "down") + assert node(down, "session") == Restrict("session", down.seed, (("subject_id", "subject_id"),)) + assert node(down, "flag") == Restrict("flag", down.seed, (("subject_id", "subject_id"),)) + rows = {t: {k if isinstance(k, tuple) else (k,) for k in ks} for t, ks in K1_DOWN.items()} + assert {t: set(r) for t, r in RowEvaluator(down, s.db, k1_match_for(s.db)).result().items() if r} == rows + # Deleting with the seed materialized first removes exactly these rows. + after = simulate_delete(down, s.model, s.db, k1_match_for, materialize={down.seed}) + expected_rows = {(t, k) for t, ks in rows.items() for k in ks} + assert after == apply_delete(s.db, expected_rows) + + +def test_K4_plan(): + """K4: a dj.Top seed is replaced by its keys, which cross downward by relabeling.""" + s = k_shape() + keys = ({"subject_id": 1, "session_id": 2}, {"subject_id": 2, "session_id": 1}) + down = plan_expand(s.schema, Literal("session", keys, read=True), "down") + assert node(down, "trial") == Literal("trial", keys) + s.check(down, {"session": [(1, 2), (2, 1)], "trial": [(1, 2, 1), (2, 1, 1)]}) + # Upward, keys read from the table itself that hold the whole foreign key are relabeled too. + up = plan_expand(s.schema, Literal("session", keys, read=True), "up") + assert node(up, "subject") == Literal("subject", ({"subject_id": 1}, {"subject_id": 2})) + s.check(up, {"session": [(1, 2), (2, 1)], "subject": [1, 2]}) + + +def test_K5_empty_seeds(): + """K5: a seed that selects no rows expands to nodes with no rows.""" + s = k_shape() + for seed in (Literal("session", ()), Seed("session", [], frozenset()), Seed("session", {"subject_id": 4}, None)): + s.check(plan_expand(s.schema, seed, "down"), {"session": [], "trial": []}) + s.check(plan_expand(s.schema, seed, "up"), {"session": [], "subject": []}) + + +def test_user_keys_cross_upward_by_restriction(): + """K5: a user-supplied key dict may name rows that do not exist, so it is not relabeled upward.""" + s = k_shape() + up = plan_expand(s.schema, Literal("session", ({"subject_id": 4},)), "up") + assert isinstance(node(up, "subject"), Restrict) + s.check(up, {"session": [], "subject": []}) + + +# ============================================================================= +# The delete simulation +# ============================================================================= + + +@given(seeded=seeded_cases()) +def test_delete_simulation_without_live_lifts(seeded): + """ + Without a lift evaluated at delete time, deleting table by table, leaves + first, with live evaluation, removes exactly D(S) and never breaks a + foreign key. + + The model's seed conditions read no table, and a parent is deleted after its + children, so every other restriction still sees the rows it reads. + """ + case, seed, seed_rows = seeded + plan = plan_expand(describe(case.schema), seed, "down") + assume(not lifts_evaluated_at_delete(plan)) + after = simulate_delete(plan, case.schema, case.db) + assert after == apply_delete(case.db, case.graph.D(seed_rows)) + + +@pytest.mark.parametrize("parts_first", [False, True], ids=["any seed", "part seed"]) +@given(data=st.data()) +def test_delete_simulation_with_the_spec_materialization(parts_first, data): + """ + Materializing what diagram.md says to materialize before the first DELETE + (every lift, and the seed unless it is literal keys) makes the delete remove + exactly D(S), and it never breaks a foreign key. + """ + case, seed, seed_rows = data.draw(seeded_cases(parts_first), label="case and seed") + plan = plan_expand(describe(case.schema), seed, "down") + event("a lift is evaluated at delete time" if lifts_evaluated_at_delete(plan) else "no lift at delete time") + after = simulate_delete(plan, case.schema, case.db, materialize=spec_materialization(plan)) + assert after == apply_delete(case.db, case.graph.D(seed_rows)) + + +NEEDS_E3 = ( + "E3: Table.delete materializes the lifts and live seeds of the plan before the first DELETE " + "(diagram.md, 'Materialization happens only at delete time'); the E2a plan carries no materialization flags" +) + + +def _plan_materialization(plan: Plan) -> set: + # The flags E3 adds to the plan; none yet. + return set(getattr(plan, "materialize", ())) + + +def _g1_seeded(): + """G1's schema with the seed Session.Trial & {"trial_id": 2}, as ``seeded_cases`` draws them.""" + s = g1() + cond = {"trial_id": 2} + return ( + Case(s.model, s.db, RowGraph(s.model, s.db)), + Seed("session__trial", cond, frozenset(cond)), + restrict(s.model, s.db, "session__trial", cond), + ) + + +@pytest.mark.xfail(strict=True, reason=NEEDS_E3) +@given(seeded=seeded_cases(parts_first=True)) +@example(seeded=_g1_seeded()) +def test_delete_simulation(seeded): + """The L0′ delete property: the plan's own materialization flags make the delete remove exactly D(S).""" + case, seed, seed_rows = seeded + plan = plan_expand(describe(case.schema), seed, "down") + after = simulate_delete(plan, case.schema, case.db, materialize=_plan_materialization(plan)) + assert after == apply_delete(case.db, case.graph.D(seed_rows)) + + +# ============================================================================= +# The delete simulation on catalog shapes +# ============================================================================= + + +@pytest.mark.parametrize( + "shape, table, cond", + [ + pytest.param(g1, "session__trial", {"trial_id": 2}, id="G1"), + pytest.param(g7, "master__part_b", {"part_b_id": 100}, id="G7"), + ], +) +def test_delete_needs_the_lift_materialized(shape, table, cond): + """ + The part rows are deleted before their master, so a lift evaluated at the + master's DELETE finds none and the master stays; materialized first, the + delete removes D(S). + """ + s = shape() + plan = s.plan(table, cond, "down") + seed_rows = restrict(s.model, s.db, table, cond) + expected_db = apply_delete(s.db, RowGraph(s.model, s.db).D(seed_rows)) + assert simulate_delete(plan, s.model, s.db) != expected_db + assert simulate_delete(plan, s.model, s.db, materialize=spec_materialization(plan)) == expected_db + + +@pytest.mark.xfail(strict=True, reason=NEEDS_E3) +@pytest.mark.parametrize( + "shape, table, cond", + [ + pytest.param(g1, "session__trial", {"trial_id": 2}, id="G1"), + pytest.param(g7, "master__part_b", {"part_b_id": 100}, id="G7"), + pytest.param(k_shape, "subject", SessionsOfQuality(0), id="K1"), + ], +) +def test_delete_with_the_plan_flags(shape, table, cond): + """With the materialization flags of the plan itself, the delete removes D(S) (#1496).""" + s = shape() + attrs = frozenset({"subject_id", "session_id", "quality"}) if isinstance(cond, SessionsOfQuality) else None + plan = s.plan(table, cond, "down", attributes=attrs) + match_for = k1_match_for if isinstance(cond, SessionsOfQuality) else (lambda db: model_match) + seed_rows = {(table, k) for k, row in s.db[table].items() if match_for(s.db)(cond, row, s.schema.tables[table])} + expected_db = apply_delete(s.db, RowGraph(s.model, s.db).D(seed_rows)) + assert simulate_delete(plan, s.model, s.db, match_for, materialize=_plan_materialization(plan)) == expected_db + + +# ============================================================================= +# Plans and errors +# ============================================================================= + + +def test_unknown_direction_and_seed_table(): + s = s1() + with pytest.raises(ValueError, match="direction"): + s.plan("subject", {}, "sideways") + with pytest.raises(DataJointError, match="not in the schema"): + plan_expand(s.schema, Seed("nowhere", {}, frozenset()), "down") + + +def test_group_cycle_is_refused(): + """X2: a group whose part references a table that depends on its master cannot be ordered.""" + model = Schema( + [ + TableSpec("session", ("session_id",)), + TableSpec("review", ("session_id",)), + TableSpec("session__note", ("session_id",), master="session"), + ], + [ + FK("review", "session", {"session_id": "session_id"}), + FK("session__note", "session", {"session_id": "session_id"}), + FK("session__note", "review", {"session_id": "session_id"}), + ], + ) + with pytest.raises(DataJointError, match="cycle: session -> review -> session"): + plan_expand(describe(model), Seed("session", {}, frozenset()), "down") + + +def test_describe_lists_nodes_and_tables(): + s = g4() + text = s.plan("run", {"run_id": 1}, "down").describe() + assert "Seed(run; names {run_id})" in text + assert "down: run iterates over run__carry -> run__step" in text + assert "run__carry = #" in text From b2d703a5e516533c7aee0ac7663ca9d4b9c59842 Mon Sep 17 00:00:00 2001 From: Dimitri Yatsenko Date: Sat, 17 Oct 2026 19:09:13 -0500 Subject: [PATCH 3/3] feat(diagram): add Diagram.expand on the compiled plan Diagram.expand(seed, direction) returns a diagram whose nodes are the table-level reach of the seed and whose d[T], counts() and iteration give each table's rows as a restricted query expression, compiled from the plan. It loads the schemas reachable in the direction of travel, refuses projections, joins, aggregations and unions as seeds, and replaces a dj.Top seed by its keys. cascade() and trace() are unchanged. tests/integration/test_expand_planner.py checks S1, S3, S7, G1, G4, G7 and K1, K3-K6 of the shape catalog on MySQL and PostgreSQL. --- src/datajoint/diagram.py | 150 +++++++++ tests/integration/test_expand_planner.py | 382 +++++++++++++++++++++++ 2 files changed, 532 insertions(+) create mode 100644 tests/integration/test_expand_planner.py diff --git a/src/datajoint/diagram.py b/src/datajoint/diagram.py index 6ff60d061..df818daea 100644 --- a/src/datajoint/diagram.py +++ b/src/datajoint/diagram.py @@ -203,6 +203,12 @@ class Diagram(nx.MultiDiGraph): # noqa: C901 dj.Diagram(schema).draw() """ + # Set by Diagram.expand: the direction, the plan, and each node's rows as a + # query expression, in topological order. None on other diagrams. + _direction = None + _plan = None + _expansion = None + def __init__(self, source, context=None) -> None: if isinstance(source, Diagram): # copy constructor @@ -417,6 +423,126 @@ def __mul__(self, arg) -> "Diagram": self.nodes_to_show.intersection_update(arg.nodes_to_show) return self + @classmethod + def expand(cls, seed, direction="down"): + """ + Expand a seed into the rows that depend on it, that it depends on, or both. + + The rows are those the traversal specification defines (the "Traversal + specification" section of the Diagram specification, ``diagram.md``): + + - ``"down"``: ``D(S)``, the seed rows and every row that depends on them, + through foreign keys and through master-part groups; + - ``"up"``: ``U(S)``, the seed rows and every row they depend on, with + the whole master-part instance of every row reached; + - ``"both"``: ``U(D(S))``, the smallest region around the seed that is + consistent on its own. + + The traversal is planned once over the group quotient (``datajoint.expand``) + and compiled to one restricted query expression per table. No rows are + read, except to find the master keys of a group whose foreign keys link + one master-part instance to another, and the keys of a ``dj.Top`` seed. + + Parameters + ---------- + seed : QueryExpression + One table with any restriction, e.g. ``Session & {"subject_id": 1}``. + The restriction may use other queries. Projections, joins, + aggregations and unions are refused as seeds. + direction : str, optional + ``"down"`` (default), ``"up"``, or ``"both"``. + + Returns + ------- + Diagram + A diagram whose nodes are the tables the expansion reaches (the + table-level reach: a node may hold zero rows). ``d[T]`` returns table + ``T`` restricted to its rows, ``counts()`` counts them, and iteration + yields each node's restricted table in topological order. + + Raises + ------ + DataJointError + If the seed is not one table with a restriction. + ValueError + If ``direction`` is not ``"down"``, ``"up"``, or ``"both"``. + + Examples + -------- + >>> dj.Diagram.expand(Subject & {"subject_id": 1}, "down").counts() + >>> dj.Diagram.expand(Session & key, "up")[Subject] + """ + from .expand import ( + TraversalSchema, + compile_plan, + graph_reach, + load_dependencies, + plan_expand, + seed_from_expression, + ) + from .table import FreeTable + + if direction not in ("down", "up", "both"): + raise ValueError(f"direction must be 'down', 'up', or 'both', got {direction!r}") + if inspect.isclass(seed) and issubclass(seed, Table): + seed = seed() + seed_node = seed_from_expression(seed) + conn = seed.connection + database, _ = conn.adapter.split_full_table_name(seed_node.table) + dependencies = load_dependencies(conn, direction, {database}) + # Describe the tables the traversal can reach, then plan on the description. + names = graph_reach(dependencies, seed_node.table, direction) + tables = {name: FreeTable(conn, name) for name in names} + schema = TraversalSchema.from_dependencies(dependencies, names, lambda name: tables[name].heading) + plan = plan_expand(schema, seed_node, direction) + expansion = compile_plan(plan, conn, tables) + + result = cls.__new__(cls) + nx.MultiDiGraph.__init__(result, dependencies) + result.remove_nodes_from(set(result.nodes()) - set(plan.order)) + result._connection = conn + result.context = {} + result._cascade_restrictions = {} + result._restrict_conditions = {} + result._restriction_attrs = {} + result.nodes_to_show = set(plan.order) + result._expanded_nodes = set(plan.order) + result._direction = direction + result._plan = plan + result._expansion = expansion + return result + + def _expansion_node(self, key): + """ + Resolve ``d[key]`` on an expanded diagram to a full table name. + + ``key`` is a table class or instance, a full table name, or a class name + (``"Session"``, or ``"Session.Trial"`` for a part). + """ + from .utils import from_camel_case + + if isinstance(key, type) and issubclass(key, Table) or isinstance(key, Table): + name = key.full_table_name + elif isinstance(key, str) and ("`" in key or '"' in key): + name = key + elif isinstance(key, str): + try: + tail = "__".join(from_camel_case(part) for part in key.split(".")) + except DataJointError: + tail = None + split = self._connection.adapter.split_full_table_name + matches = [n for n in self._expansion if split(n)[1].lstrip("#_") == tail] + if len(matches) > 1: + raise DataJointError( + f"Ambiguous table reference {key!r}: matches {', '.join(matches)}. Use a full table name." + ) + name = matches[0] if matches else key + else: + raise DataJointError(f"d[...] expects a Table class, Table instance, or string; got {type(key).__name__}.") + if name not in self._expansion: + raise DataJointError(f"Table {name} is not a node of this diagram.") + return name + @classmethod def cascade(cls, table_expr, part_integrity="enforce"): """ @@ -664,6 +790,10 @@ def __getitem__(self, key): >>> trace[Session].fetch1("session_date") # class index >>> trace["my_schema.Session"].to_dicts() # string index → FreeTable """ + # Expanded diagrams: each node's rows, in every direction. + if self._expansion is not None: + return self._expansion[self._expansion_node(key)] + # Non-trace diagrams: defer to networkx adjacency lookup so existing # `diagram[node_name]` patterns (used in diagram algebra, diagram tests) # keep working. @@ -716,6 +846,8 @@ def _restricted_table(self, node): """ from .table import FreeTable + if self._expansion is not None and node in self._expansion: + return self._expansion[node] ft = FreeTable(self._connection, node) restrictions = (self._cascade_restrictions or self._restrict_conditions).get(node, []) if not restrictions: @@ -1117,6 +1249,12 @@ def counts(self): dict[str, int] Mapping of full table name to affected row count. """ + if self._expansion is not None: + result = {} + for ft in self: + result[ft.full_table_name] = len(ft) + logger.info("{table} ({count} tuples)".format(table=ft.full_table_name, count=result[ft.full_table_name])) + return result restrictions = self._cascade_restrictions or self._restrict_conditions if not restrictions: raise DataJointError( @@ -1137,7 +1275,14 @@ def __iter__(self): Yields restricted ``FreeTable`` objects when cascade or restrict conditions have been applied, unrestricted ``FreeTable`` otherwise. + On an expanded diagram, the order is the plan's: each master's owned + parts right after it. """ + if self._expansion is not None: + for node in self._expansion: + if node in self.nodes_to_show: + yield self._expansion[node] + return for node in topo_sort(self): if node in self.nodes_to_show: yield self._restricted_table(node) @@ -1149,6 +1294,11 @@ def __reversed__(self): Same as ``__iter__`` but reversed — useful for cascading deletes and drops. """ + if self._expansion is not None: + for node in reversed(list(self._expansion)): + if node in self.nodes_to_show: + yield self._expansion[node] + return for node in reversed(topo_sort(self)): if node in self.nodes_to_show: yield self._restricted_table(node) diff --git a/tests/integration/test_expand_planner.py b/tests/integration/test_expand_planner.py new file mode 100644 index 000000000..dd8e07039 --- /dev/null +++ b/tests/integration/test_expand_planner.py @@ -0,0 +1,382 @@ +""" +Integration tests for ``Diagram.expand`` on MySQL and PostgreSQL. + +A handful of shapes from the Traversal Shape Catalog (``traversal-shapes.md`` in +datajoint-docs), declared through DataJoint: the rows ``Diagram.expand`` selects +in each table equal the rows the catalog states. The full catalog, the +differential tests on random schemas, and the delete tests come later. +""" + +import time + +import pytest + +import datajoint as dj +from datajoint.errors import DataJointError + + +@pytest.fixture(scope="function") +def schema_by_backend(connection_by_backend, db_creds_by_backend): + """A fresh schema for each test.""" + backend = db_creds_by_backend["backend"] + name = f"djtest_expand_{backend}_{str(int(time.time() * 1000))[-8:]}"[:64] + drop = f"DROP DATABASE IF EXISTS {connection_by_backend.adapter.quote_identifier(name)}" + connection_by_backend.query(drop) + yield dj.Schema(name, connection=connection_by_backend) + if connection_by_backend.is_connected: + connection_by_backend.query(drop) + + +def keys(expr) -> set: + """The primary keys of an expression's rows, as tuples.""" + pk = expr.primary_key + return {tuple(row[a] for a in pk) for row in expr.keys()} + + +def check(diagram, expected: dict) -> None: + """The diagram's nodes are the tables of ``expected``, and each holds the rows listed.""" + names = {table.full_table_name: table for table in expected} + assert diagram.nodes_to_show == set(names) + for table, rows in expected.items(): + assert keys(diagram[table]) == {r if isinstance(r, tuple) else (r,) for r in rows}, table.__name__ + assert diagram.counts() == {table.full_table_name: len(rows) for table, rows in expected.items()} + + +def test_S1_identity_fk(schema_by_backend): + """S1: the condition is copied to the child; upward the parent is restricted.""" + + @schema_by_backend + class Subject(dj.Manual): + definition = """ + subject_id : int32 + """ + + @schema_by_backend + class Session(dj.Manual): + definition = """ + -> Subject + session_id : int32 + """ + + Subject.insert([(1,), (2,)]) + Session.insert([(1, 1), (1, 2), (2, 1)]) + + down = dj.Diagram.expand(Subject & {"subject_id": 1}, "down") + check(down, {Subject: [1], Session: [(1, 1), (1, 2)]}) + assert "IN (" not in down[Session].where_clause() # a copy, not a subquery + check(dj.Diagram.expand(Session & {"session_id": 1}, "up"), {Session: [(1, 1), (2, 1)], Subject: [1, 2]}) + # Iteration is in topological order, parents first; reversed, leaves first. + assert [t.full_table_name for t in down] == [Subject.full_table_name, Session.full_table_name] + assert [t.full_table_name for t in reversed(down)] == [Session.full_table_name, Subject.full_table_name] + # A table can be named by its class name or its full name. + assert keys(down["Session"]) == keys(down[Session.full_table_name]) == {(1, 1), (1, 2)} + + +def test_S3_nullable_fk_sharing_a_primary_key_column(schema_by_backend): + """S3: a foreign key with a NULL column references nothing (A2), in both directions.""" + + @schema_by_backend + class Subject(dj.Manual): + definition = """ + subject_id : int32 + """ + + @schema_by_backend + class Surgery(dj.Manual): + definition = """ + -> Subject + surgery_id : int32 + """ + + @schema_by_backend + class Session(dj.Manual): + definition = """ + -> Subject + session_id : int32 + --- + -> [nullable] Surgery + """ + + Subject.insert([(1,), (2,)]) + Surgery.insert([(1, 1), (1, 2), (2, 1)]) + Session.insert([(1, 1, 1), (1, 2, None), (1, 3, 2), (2, 1, 1)]) + + check( + dj.Diagram.expand(Surgery & {"subject_id": 1}, "down"), + {Surgery: [(1, 1), (1, 2)], Session: [(1, 1), (1, 3)]}, + ) + check( + dj.Diagram.expand(Session & {"subject_id": 1, "session_id": 2}, "up"), + {Session: [(1, 2)], Subject: [1], Surgery: []}, + ) + + +def test_S7_diamond_reached_through_a_secondary_fk(schema_by_backend): + """S7: Session (2, 1) reaches the seed only through its rig.""" + + @schema_by_backend + class Lab(dj.Manual): + definition = """ + lab_id : int32 + """ + + @schema_by_backend + class Rig(dj.Manual): + definition = """ + rig_id : int32 + --- + -> Lab + """ + + @schema_by_backend + class Session(dj.Manual): + definition = """ + -> Lab + session_id : int32 + --- + -> Rig + """ + + Lab.insert([(1,), (2,)]) + Rig.insert([(10, 1), (20, 2)]) + Session.insert([(1, 1, 20), (2, 1, 10), (2, 2, 20)]) + + check(dj.Diagram.expand(Lab & {"lab_id": 1}, "down"), {Lab: [1], Rig: [10], Session: [(1, 1), (2, 1)]}) + check( + dj.Diagram.expand(Session & {"lab_id": 2, "session_id": 1}, "up"), + {Session: [(2, 1)], Rig: [10], Lab: [1, 2]}, + ) + check(dj.Diagram.expand(Rig & {"rig_id": 10}, "both"), {Rig: [10], Session: [(2, 1)], Lab: [1, 2]}) + + +def test_G1_part_to_master_lift(schema_by_backend): + """G1: the lift is existential; upward, the instance is read whole.""" + + @schema_by_backend + class Session(dj.Manual): + definition = """ + session_id : int32 + """ + + class Trial(dj.Part): + definition = """ + -> master + trial_id : int32 + """ + + @schema_by_backend + class Review(dj.Manual): + definition = """ + -> Session + review_id : int32 + """ + + Session.insert([(1,), (2,), (3,)]) + Session.Trial.insert([(1, 1), (1, 2), (2, 1), (3, 1)]) + Review.insert([(1, 1), (2, 1)]) + + seed = Session.Trial & {"trial_id": 2} + down = dj.Diagram.expand(seed, "down") + check(down, {Session.Trial: [(1, 1), (1, 2)], Session: [1], Review: [(1, 1)]}) + assert keys(down["Session.Trial"]) == {(1, 1), (1, 2)} + check(dj.Diagram.expand(seed, "up"), {Session.Trial: [(1, 1), (1, 2)], Session: [1]}) + + +def test_G4_group_step_iterates(schema_by_backend): + """G4: a renamed sibling FK links instances; the master keys grow {1}, {1, 2}, {1, 2, 3}.""" + + @schema_by_backend + class Run(dj.Manual): + definition = """ + run_id : int32 + """ + + class Step(dj.Part): + definition = """ + -> master + step_id : int32 + """ + + class Carry(dj.Part): + definition = """ + -> master + carry_id : int32 + --- + -> Run.Step.proj(prev_run="run_id", prev_step="step_id") + """ + + Run.insert([(1,), (2,), (3,), (4,)]) + Run.Step.insert([(1, 1), (2, 1), (3, 1), (4, 1)]) + Run.Carry.insert([(2, 1, 1, 1), (3, 1, 2, 1)]) + + down = dj.Diagram.expand(Run & {"run_id": 1}, "down") + check(down, {Run: [1, 2, 3], Run.Step: [(1, 1), (2, 1), (3, 1)], Run.Carry: [(2, 1), (3, 1)]}) + assert [g.iterate for g in down._plan.groups] == [True] + check( + dj.Diagram.expand(Run & {"run_id": 2}, "down"), + {Run: [2, 3], Run.Step: [(2, 1), (3, 1)], Run.Carry: [(2, 1), (3, 1)]}, + ) + check( + dj.Diagram.expand(Run & {"run_id": 2}, "up"), + {Run: [1, 2], Run.Step: [(1, 1), (2, 1)], Run.Carry: [(2, 1)]}, + ) + check( + dj.Diagram.expand(Run & {"run_id": 2}, "both"), + {Run: [1, 2, 3], Run.Step: [(1, 1), (2, 1), (3, 1)], Run.Carry: [(2, 1), (3, 1)]}, + ) + + +def test_G7_owned_part_of_part_with_renamed_attributes(schema_by_backend): + """G7: the lift renames the owning attribute src_master onto the master's key.""" + + @schema_by_backend + class Master(dj.Manual): + definition = """ + master_id : int32 + """ + + class PartA(dj.Part): + definition = """ + -> master + part_a_id : int32 + """ + + class PartB(dj.Part): + definition = """ + -> Master.PartA.proj(src_master="master_id", src_part="part_a_id") + part_b_id : int32 + """ + + @schema_by_backend + class Report(dj.Manual): + definition = """ + -> Master + report_id : int32 + """ + + Master.insert([(1,), (2,), (3,)]) + Master.PartA.insert([(1, 10), (1, 11), (2, 20)]) + Master.PartB.insert([(1, 10, 100), (1, 11, 110), (2, 20, 200)]) + Report.insert([(1, 1), (2, 1)]) + + seed = Master.PartB & {"part_b_id": 100} + instance = {Master.PartB: [(1, 10, 100), (1, 11, 110)], Master: [1], Master.PartA: [(1, 10), (1, 11)]} + check(dj.Diagram.expand(seed, "down"), {**instance, Report: [(1, 1)]}) + check(dj.Diagram.expand(seed, "up"), instance) + + +@pytest.fixture +def k_tables(schema_by_backend): + """The schema shared by the seed shapes K1 to K6.""" + + @schema_by_backend + class Subject(dj.Manual): + definition = """ + subject_id : int32 + --- + species : varchar(16) + """ + + @schema_by_backend + class Session(dj.Manual): + definition = """ + -> Subject + session_id : int32 + --- + quality : int32 + """ + + @schema_by_backend + class Trial(dj.Manual): + definition = """ + -> Session + trial_id : int32 + """ + + @schema_by_backend + class Flag(dj.Manual): + definition = """ + flag_id : int32 + --- + -> [nullable] Subject + """ + + Subject.insert([(1, "mouse"), (2, "rat"), (3, "mouse"), (4, "rat")]) + Session.insert([(1, 1, 0), (1, 2, 1), (2, 1, 1), (2, 2, 0), (3, 1, 1)]) + Trial.insert([(1, 1, 1), (1, 2, 1), (2, 1, 1), (2, 2, 1), (3, 1, 1)]) + Flag.insert([(1, 1), (2, 3)]) + return Subject, Session, Trial, Flag + + +def test_K1_seed_restricted_by_a_descendant(k_tables): + """K1: down takes all sessions of the seed subjects, not only those of quality 0.""" + Subject, Session, Trial, Flag = k_tables + seed = Subject & (Session & "quality = 0") + check( + dj.Diagram.expand(seed, "down"), + { + Subject: [1, 2], + Session: [(1, 1), (1, 2), (2, 1), (2, 2)], + Trial: [(1, 1, 1), (1, 2, 1), (2, 1, 1), (2, 2, 1)], + Flag: [1], + }, + ) + check(dj.Diagram.expand(seed, "up"), {Subject: [1, 2]}) + + +def test_K3_seed_kinds_give_the_same_rows(k_tables): + """K3: a dict, a list of dicts, a string, a query expression, an AndList and a Not select the same rows.""" + Subject, Session, Trial, Flag = k_tables + for seed in ( + Session & {"subject_id": 2}, + Session & [{"subject_id": 2, "session_id": 1}, {"subject_id": 2, "session_id": 2}], + Session & "subject_id = 2", + Session & (Subject & {"species": "rat"}), + Session & dj.AndList([{"subject_id": 2}, "quality >= 0"]), + Session - Flag, + ): + check(dj.Diagram.expand(seed, "down"), {Session: [(2, 1), (2, 2)], Trial: [(2, 1, 1), (2, 2, 1)]}) + check(dj.Diagram.expand(seed, "up"), {Session: [(2, 1), (2, 2)], Subject: [2]}) + + +def test_K4_top_seed(k_tables): + """K4: a dj.Top seed is replaced by the keys it selects.""" + Subject, Session, Trial, Flag = k_tables + seed = Session & "quality = 1" & dj.Top(2) + check(dj.Diagram.expand(seed, "down"), {Session: [(1, 2), (2, 1)], Trial: [(1, 2, 1), (2, 1, 1)]}) + check(dj.Diagram.expand(seed, "up"), {Session: [(1, 2), (2, 1)], Subject: [1, 2]}) + + +def test_K5_empty_seed_and_unknown_node(k_tables): + """K5: an empty seed gives nodes with zero rows; indexing a table outside the diagram raises.""" + Subject, Session, Trial, Flag = k_tables + down = dj.Diagram.expand(Session & {"subject_id": 4}, "down") + check(down, {Session: [], Trial: []}) + with pytest.raises(DataJointError, match="not a node"): + down[Flag] + + +def test_K6_unrestricted_seed(k_tables): + """K6: an unrestricted seed selects every row; upward, subject 4 is not reached.""" + Subject, Session, Trial, Flag = k_tables + sessions = [(1, 1), (1, 2), (2, 1), (2, 2), (3, 1)] + trials = [(1, 1, 1), (1, 2, 1), (2, 1, 1), (2, 2, 1), (3, 1, 1)] + for seed in (Session, Session()): + check(dj.Diagram.expand(seed, "down"), {Session: sessions, Trial: trials}) + check(dj.Diagram.expand(seed, "up"), {Session: sessions, Subject: [1, 2, 3]}) + + +def test_refused_seeds(k_tables): + """A projection, a join, an aggregation and a union are refused as seeds; the direction is checked.""" + Subject, Session, Trial, Flag = k_tables + for seed, kind in ( + (Session.proj(), "projection"), + (Session * Trial, "join"), + (Subject.aggr(Session, n="count(*)"), "aggregation"), + ((Session & "quality = 0").proj() + (Session & "quality = 1").proj(), "union"), + (dj.U("subject_id") & Session, "projection"), + ): + with pytest.raises(DataJointError, match=f"(?i){kind}"): + dj.Diagram.expand(seed, "down") + with pytest.raises(ValueError, match="direction"): + dj.Diagram.expand(Session, "sideways")