metricflow-sql 0.1.1__py3-none-any.whl

metricflow_sql/compiler/rules/expand_constellations.py

@@ -0,0 +1,790 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import NoReturn
+
+from sqlglot import exp
+from sqlglot.helper import find_new_name
+from sqlglot.optimizer.qualify import qualify
+from sqlglot.optimizer.qualify_columns import quote_identifiers
+from sqlglot.optimizer.scope import find_all_in_scope, find_in_scope, traverse_scope
+from sqlglot.optimizer.simplify import simplify_parens
+
+from metricflow_sql.compiler.catalog.cluster import Cluster
+from metricflow_sql.compiler.catalog.semantic import SemanticCatalog
+from metricflow_sql.compiler.table_arguments import TABLE_ARGUMENTS
+from metricflow_sql.objects.base import resolve_dataset_path
+from metricflow_sql.objects.constellation import Metric, Role
+from metricflow_sql.registry.reference_resolution import resolve_metric
+from metricflow_sql.registry.registry import ConstellationProxy, Registry
+from metricflow_sql.utils.sql import group_all_keys, normalize_group, path_names, table_scopes
+
+logger = logging.getLogger(__name__)
+
+# Query clauses that can't be compiled on constellations yet. HAVING, QUALIFY, DISTINCT,
+# ORDER BY, LIMIT and OFFSET never get here, since isolate_fragments moves them to an outer
+# query first.
+# - JOIN would be pushed down into every branch, like WHERE. It isn't built yet.
+UNSUPPORTED_CLAUSES = {"joins": "JOIN"}
+
+
+def _column(*parts: exp.Identifier) -> exp.Expr:
+    """`accounts, account_facts, x` -> `accounts.account_facts.x`, reusing the parts."""
+    return exp.column(*reversed(parts[:4]), fields=list(parts[4:]) or None, copy=False)
+
+
+def _prefixed(column: exp.Column, name: exp.Identifier) -> exp.Expr:
+    """`bill_customer.birth_year` -> `catalog_sales.bill_customer.birth_year`"""
+    # `name` prefixes every column of an expression.
+    return _column(name.copy(), *path_names(column))
+
+
+def _replace(node: exp.Expr, replacement: exp.Expr) -> exp.Expr:
+    """
+    Returns `replacement` to stand in for `node` in a transform, parenthesized when `node` sits
+    inside a larger expression, e.g. a metric `a + b` in `100 * AGG(metric)` gives `100 * (a + b)`.
+
+    sqlglot doesn't parenthesize a replaced node, so it's always wrapped here, and the rule drops
+    the parentheses that aren't needed once the query is expanded.
+    """
+    if node.parent is not None:
+        return exp.Paren(this=replacement)
+    return replacement
+
+
+@dataclass(frozen=True)
+class Leaf:
+    """
+    An aggregate over a single dataset, computed by that dataset's branch.
+
+    Metrics are unfolded until only leaves are left. In the TPC-DS registry,
+
+        store_return_rate   = AGG(store_return_amount) / AGG(store_sales_amount)
+        store_return_amount = AGG(store_returns.return_amount)
+        store_sales_amount  = AGG(store_sales.sales_amount)
+
+    so `AGG(store_return_rate)` unfolds to two leaves: `AGG(return_amount)` on `store_returns`
+    and `AGG(sales_amount)` on `store_sales`.
+
+    A leaf is named after what it computes: the calculation (`return_amount`), or the metric
+    for an aggregate written inline in it (`catalog_sales_amount = SUM(...)`).
+    """
+
+    dataset: exp.Table
+    expr: exp.Expr
+    name: str
+    # A filter over roles that only this aggregate applies, from a query's
+    # `AGG(metric) FILTER (WHERE ...)`, e.g. `customer.birth_year > 1960`.
+    condition: exp.Expr | None
+
+
+@dataclass
+class ConstellationContext:
+    registry: Registry
+    constellation: ConstellationProxy
+    catalog: str
+    # The name the query's FROM gives the constellation, e.g. `u` in `FROM universe AS u`.
+    # qualify prefixes the constellation's columns with it.
+    alias: str
+    # `FROM universe(min_birth_year => 1985)` reads each branch's dataset with them.
+    arguments: dict[str, exp.Expr]
+    # The schema branches are qualified against before ExpandDatasetsRule expands them.
+    cluster: Cluster
+    # The query's GROUP BY expressions, i.e. the join keys, mapped to the column the branches
+    # output them as, e.g. `item.category` -> `item_category`. Matched structurally, since
+    # qualify writes every use of a GROUP BY expression the same way.
+    keys: dict[exp.Expr, str] = field(default_factory=dict)
+    # Keyed by (dataset, aggregate, filter), so an aggregate shared by several metrics is
+    # computed once.
+    leaves: dict[tuple[exp.Table, exp.Expr, exp.Expr | None], Leaf] = field(default_factory=dict)
+    # Metrics being unfolded right now, to catch a metric that references itself. Keyed by
+    # constellation too, as imported metrics may share a name.
+    expanding: set[tuple[exp.Table, exp.Identifier]] = field(default_factory=set)
+    # Whether the query has no GROUP BY or AGG, so it reads the datasets' rows, stacked.
+    stack: bool = False
+    # The query's dialect, used to qualify branches, normalize metrics and print errors.
+    dialect: str | None = None
+
+
+class ExpandConstellationsRule:
+    """
+    Sqlglot optimizer rule that compiles queries over constellations with drill across.
+
+    A constellation's metrics can read from several datasets, e.g. a return rate divides
+    returns from `store_returns` by sales from `store_sales`. Joining those datasets row by row
+    through a shared role such as the item category would pair every sale with every return
+    in the same category, and inflate both sums (the chasm trap). So the rule compiles:
+
+    1. One branch per dataset the metrics read from: a query that aggregates that dataset by
+       the query's GROUP BY expressions, so it has exactly one row per key, e.g. sales per
+       category. A GROUP BY expression can be any expression over roles, e.g. a birth decade,
+       and each branch computes it from its own assignments, like it does the WHERE.
+    2. A full outer join of the branches on those keys. Since every branch has one row per key,
+       the join can't multiply anything, and being full it keeps keys that only some datasets
+       have.
+    3. The metric arithmetic, on the joined rows. In a metric, `+` and `-` treat NULL as 0 and
+       a division by 0 gives NULL; arithmetic written in the query is left as written.
+
+    HAVING, ORDER BY and LIMIT were already moved to an outer query by isolate_fragments, which
+    also selects any metric they use, so e.g. an ORDER BY metric's dataset joins like any other.
+
+    For example, `SELECT item.category, AGG(store_return_rate) FROM universe GROUP BY 1` becomes
+    return amounts per category from `store_returns`, full outer joined with sales per category
+    from `store_sales`, and the rate is their ratio. Without GROUP BY, each branch is a single
+    row and the join matches them all.
+
+    Without metrics, the constellation stands for all its datasets stacked, so every dataset
+    that provides all the roles the query uses gets a branch with just the keys, and the join
+    merges them. Without GROUP BY either, the query reads the stacked rows themselves, so its
+    projections are the keys and the branches are stacked with UNION ALL instead.
+    """
+
+    def _fail(
+        self, ctx: ConstellationContext, message: str, expr: exp.Expr | None = None
+    ) -> NoReturn:
+        if expr is not None:
+            message = f"{message}: {expr.sql(dialect=ctx.dialect)}"
+        raise ValueError(message)
+
+    def __call__(
+        self, expression: exp.Expr, cluster: Cluster, dialect: str | None = None
+    ) -> exp.Expr:
+        for scope, select, table in table_scopes(expression):
+            if not (registry_and_constellation := cluster.get_constellation(table)):
+                continue
+            registry, constellation = registry_and_constellation
+            # Only a semantic catalog has constellations.
+            if not isinstance(cluster.get_catalog(table.catalog), SemanticCatalog):
+                continue
+
+            ctx = ConstellationContext(
+                registry=registry,
+                constellation=constellation,
+                catalog=table.catalog,
+                alias=table.alias_or_name,
+                arguments=table.meta.get(TABLE_ARGUMENTS, {}),
+                cluster=cluster,
+                dialect=dialect,
+            )
+
+            query = self._expand_scope(select, ctx).transform(simplify_parens, dialect, copy=False)
+
+            select.replace(query)
+
+            if not scope.parent:
+                return query
+
+        return expression
+
+    def _expand_scope(self, select: exp.Select, ctx: ConstellationContext) -> exp.Select:
+        """
+        Rewrites one SELECT over a constellation, following the steps in the class docstring.
+
+        Rewriting the projections registers the leaves, and the leaves decide which branches get
+        built and joined. Without leaves, every dataset that provides the roles the query uses gets a
+        branch.
+        """
+        for arg, clause in UNSUPPORTED_CLAUSES.items():
+            if select.args.get(arg):
+                self._fail(ctx, f"Constellation queries don't support {clause} yet")
+
+        ctx.keys = self._keys(select, ctx)
+        condition = None
+        if isinstance(where := select.args.get("where"), exp.Where):
+            condition = where.this
+            self._check_condition(
+                condition,
+                "WHERE on a constellation can only use roles, use HAVING for metrics",
+                ctx,
+            )
+        ctx.stack = not ctx.keys and not find_in_scope(select, exp.AggFunc)
+        if ctx.stack:
+            # Each projection over roles is read from every dataset's rows, named like a GROUP BY
+            # key, e.g. `customer_birth_year`. _expand_expr rejects the rest, or NULLs metrics.
+            for projection in select.expressions:
+                expr = projection.unalias()
+                paths = self._paths(expr)
+                if (
+                    paths
+                    and not expr.find(exp.Query)
+                    and all(self._role_parts(path, ctx) is not None for path in paths)
+                ):
+                    parts = self._role_parts(expr, ctx)
+                    name = (
+                        "_".join(part.name for part in parts) if parts else projection.alias_or_name
+                    )
+                    ctx.keys.setdefault(expr, find_new_name(set(ctx.keys.values()), name))
+        # The alias is reused as an identifier, so a quoted name like `"Sales"` stays quoted.
+        projections = [
+            exp.alias_(
+                self._expand_expr(projection.unalias(), ctx),
+                alias.copy()
+                if isinstance(alias := projection.args.get("alias"), exp.Identifier)
+                else projection.alias_or_name,
+                copy=False,
+            )
+            for projection in select.expressions
+        ]
+        if ctx.leaves:
+            datasets = list(dict.fromkeys(leaf.dataset for leaf in ctx.leaves.values()))
+        elif ctx.keys:
+            datasets = self._key_datasets(condition, ctx)
+        else:
+            self._fail(ctx, "Constellation queries must aggregate a metric or group by roles")
+
+        if ctx.stack:
+            branches = [self._branch(dataset, condition, ctx) for dataset in datasets]
+            rows = exp.union(*branches, distinct=False, copy=False) if branches[1:] else branches[0]
+        else:
+            rows = self._join_branches(datasets, condition, ctx)
+        joined = rows.subquery(ctx.alias, copy=False)
+        return exp.select(*projections, copy=False).from_(joined, copy=False)
+
+    def _keys(self, select: exp.Select, ctx: ConstellationContext) -> dict[exp.Expr, str]:
+        """
+        Returns the query's GROUP BY expressions, mapped to their column names.
+
+        A role is named after its path, and any other expression after its output name, e.g.
+        `SELECT item.category, customer.birth_year // 10 * 10 AS decade ... GROUP BY 1, 2` gives
+        `{item.category: "item_category", customer.birth_year // 10 * 10: "decade"}`.
+        GROUP BY ALL groups by every projection that doesn't aggregate.
+        """
+        if not isinstance(group := select.args.get("group"), exp.Group):
+            return {}
+        # TODO: support ROLLUP, CUBE and GROUPING SETS. A branch then has one row per key and
+        # grouping set, so a rolled up NULL and a real NULL share a key and the join mixes them
+        # up. It needs GROUPING() carried as part of the key.
+        grouping_sets = normalize_group(group)
+        if len(grouping_sets) != 1:
+            self._fail(
+                ctx, "GROUPING SETS, CUBE and ROLLUP are not supported on constellations yet"
+            )
+        groupings = grouping_sets[0]
+        # A bare GROUP BY ALL; with expressions, ALL is just the default quantifier.
+        if group.args.get("all") and not groupings:
+            groupings = tuple(select.expressions[i].unalias() for i in group_all_keys(select))
+
+        output_names = {
+            projection.unalias(): projection.alias_or_name for projection in select.expressions
+        }
+        keys: dict[exp.Expr, str] = {}
+        for grouping in groupings:
+            # A GROUP BY expression is computed by every branch, so it can only read roles.
+            paths = self._paths(grouping)
+            if (
+                not paths
+                or grouping.find(exp.AggFunc, exp.Window, exp.Query)
+                or any(self._role_parts(path, ctx) is None for path in paths)
+            ):
+                self._fail(ctx, "GROUP BY on a constellation must reference roles", grouping)
+            if grouping not in keys:
+                parts = self._role_parts(grouping, ctx)
+                name = (
+                    "_".join(part.name for part in parts)
+                    if parts
+                    else output_names.get(grouping, "key")
+                )
+                keys[grouping] = find_new_name(set(keys.values()), name)
+        return keys
+
+    def _expand_expr(self, expr: exp.Expr, ctx: ConstellationContext) -> exp.Expr:
+        """
+        Rewrites an expression over the constellation in terms of the joined columns, e.g.
+
+            item.category                 ->  item_category
+            customer.birth_year // 10     ->  decade (for `customer.birth_year // 10 AS decade`)
+            AGG(store_return_rate)        ->  return_amount / sales_amount
+            100 * AGG(store_return_rate)  ->  100 * (return_amount / sales_amount)
+        """
+        # Any other aggregate or window would run over the joined rows, not the datasets' rows.
+        for node in find_all_in_scope(expr, exp.AggFunc, exp.Window):
+            if not isinstance(node, exp.Agg):
+                self._fail(ctx, "Only AGG(metric) can aggregate over a constellation", node)
+
+        def rewrite(node: exp.Expr) -> exp.Expr:
+            if (key := ctx.keys.get(node)) is not None:
+                return exp.column(key, ctx.alias)
+            if self._role_parts(node, ctx) is not None:
+                self._fail(ctx, "Roles used on a constellation must appear in GROUP BY", node)
+            match node:
+                case exp.Filter(
+                    this=exp.Agg(this=exp.Column() as column), expression=exp.Where() as where
+                ) if isinstance(self._field(column, ctx), Metric):
+                    self._check_condition(
+                        where.this, "FILTER on a constellation can only use roles", ctx
+                    )
+                    return _replace(node, self._expand_metric(column.this, ctx, where.this))
+                case exp.Filter():
+                    self._fail(ctx, "FILTER on a constellation must be on AGG(metric)", node)
+                # A FILTER's subqueries were checked with its condition above.
+                case exp.Query():
+                    self._fail(ctx, "Subqueries are not supported on constellations yet", node)
+                case exp.Agg(this=exp.Column() as column) if isinstance(
+                    self._field(column, ctx), Metric
+                ):
+                    return _replace(node, self._expand_metric(column.this, ctx))
+                case exp.Agg():
+                    self._fail(ctx, "AGG on a constellation must reference a metric", node)
+                case exp.Column() if isinstance(self._field(node, ctx), Metric):
+                    # A metric has no value per row, so it's NULL, which lets e.g. a Superset
+                    # virtual dataset list it as a column. Its leaves still pick the datasets.
+                    if ctx.stack:
+                        self._expand_metric(node.this, ctx)
+                        return exp.null()
+                    self._fail(ctx, "Metrics must be aggregated with AGG", node)
+                case exp.Column():
+                    self._fail(
+                        ctx, "Constellation queries can only use roles and AGG(metric)", node
+                    )
+            return node
+
+        return expr.copy().transform(rewrite, copy=False)
+
+    def _expand_metric(
+        self,
+        name: exp.Identifier,
+        ctx: ConstellationContext,
+        condition: exp.Expr | None = None,
+        source: exp.Table | None = None,
+    ) -> exp.Expr:
+        """
+        Unfolds a metric into leaves and returns its expression over them, e.g.
+
+            store_return_rate
+            -> AGG(store_return_amount) / AGG(store_sales_amount)
+            -> AGG(store_returns.return_amount) / AGG(store_sales.sales_amount)
+            -> return_amount / sales_amount
+
+        A metric built from other metrics unfolds them by calling this again. `condition` is
+        passed down to every leaf, so a filtered ratio filters both its sides.
+
+        Its parts are looked up in the queried constellation, then in `source`, the
+        constellation defining the metric that uses them.
+        """
+        metric = resolve_metric(ctx.registry, ctx.constellation, name, source)
+        if metric is None:
+            self._fail(
+                ctx,
+                f"Unknown metric: {self._sql(ctx, ctx.constellation.name)}.{self._sql(ctx, name)}",
+            )
+        key = (metric.constellation, name)
+        if key in ctx.expanding:
+            self._fail(
+                ctx,
+                f"Metric references itself: {self._sql(ctx, metric.constellation)}.{self._sql(ctx, name)}",
+            )
+        ctx.expanding.add(key)
+        expr = metric.expr.copy()
+        # `revenue.accounts.account_count` reads `accounts.account_count` on `revenue`. As in SQL,
+        # a dataset `retail` in the metric's own namespace wins over `retail.store_sales`.
+        reads: dict[int, tuple[exp.Table, tuple[exp.Identifier, ...]]] = {}
+        for path in self._paths(expr):
+            if isinstance(path, exp.Column | exp.Dot) and len(names := path_names(path)) > 1:
+                reads[id(path)] = resolve_dataset_path(
+                    names,
+                    metric.namespace,
+                    lambda name: ctx.registry.get_dataset(name) is not None,
+                )
+        if window := expr.find(exp.Window):
+            self._fail(ctx, "Windows in metrics are not supported yet", window)
+
+        def rewrite(node: exp.Expr) -> exp.Expr:
+            match node:
+                case exp.Filter(this=exp.Agg()):
+                    self._fail(ctx, "FILTER on AGG in a metric is not supported yet", node)
+                # `AGG(dataset.calculation)` is a leaf; `AGG(metric)` composes another metric.
+                case exp.Agg(this=exp.Column() | exp.Dot() as path) if id(path) in reads:
+                    table, parts = reads[id(path)]
+                    agg = exp.Agg(this=_column(*parts))
+                    return self._leaf(table, agg, parts[-1].name, condition, ctx)
+                case exp.Agg(this=exp.Column() as column):
+                    part = self._expand_metric(column.this, ctx, condition, metric.constellation)
+                    return _replace(node, part)
+                case exp.Agg():
+                    self._fail(
+                        ctx, "AGG in a metric must reference a metric or a dataset field", node
+                    )
+                # A FILTER belongs to its aggregate, so the branch computes both.
+                case exp.AggFunc() | exp.Filter():
+                    read = {
+                        reads[id(path)][0] if id(path) in reads else None
+                        for path in self._paths(node)
+                    }
+                    dataset = read.pop() if len(read) == 1 else None
+                    if dataset is None:
+                        self._fail(
+                            ctx, "An aggregate in a metric must read from exactly one dataset", node
+                        )
+                    # Inside the copied metric, so its paths are reused.
+                    local = node.transform(
+                        lambda n: _column(*reads[id(n)][1]) if id(n) in reads else n,
+                        copy=False,
+                    )
+                    return self._leaf(dataset, local, name.name, condition, ctx)
+                # Columns inside aggregates were replaced above, so this one is outside them, e.g.
+                # a role in `AGG(m) * item.category` or a metric missing its AGG.
+                case exp.Column():
+                    self._fail(
+                        ctx, "Columns outside aggregates in metrics are not supported yet", node
+                    )
+                # A leaf is NULL for keys its dataset has no rows for, so like dbt's `safe_add`,
+                # `a + b` becomes `COALESCE(a, 0) + COALESCE(b, 0)`. Nested sums aren't wrapped.
+                case exp.Add() | exp.Sub():
+                    for side in ("this", "expression"):
+                        operand = node.args[side]
+                        if not isinstance(operand.unnest(), exp.Add | exp.Sub):
+                            node.set(
+                                side,
+                                exp.Coalesce(this=operand, expressions=[exp.Literal.number(0)]),
+                            )
+                # Divisions use `a / NULLIF(b, 0)`, so a denominator of 0 gives NULL instead of an
+                # error or infinity, e.g. `store_return_rate` for an item class with store sales of 0.
+                case exp.Div():
+                    node.set("safe", True)
+            return node
+
+        expanded = expr.transform(rewrite, copy=False)
+        ctx.expanding.discard(key)
+        return expanded
+
+    def _leaf(
+        self,
+        dataset: exp.Table,
+        expr: exp.Expr,
+        name: str,
+        condition: exp.Expr | None,
+        ctx: ConstellationContext,
+    ) -> exp.Expr:
+        """
+        Registers an aggregate over `dataset` as a leaf and returns its joined column.
+
+        The same aggregate always gets the same leaf, so `AGG(sales_amount)` on `store_sales` is
+        computed once when both `store_sales_amount` and `store_return_rate` are selected. The
+        leaf is called `name` when that's free, e.g. `sales_amount`, and gets a numbered suffix
+        otherwise, e.g. `sales_amount_2` for a second dataset with a `sales_amount`.
+        """
+        if ctx.registry.get_dataset(dataset) is None:
+            self._fail(ctx, f"Unknown dataset in metric: {self._sql(ctx, dataset)}")
+        leaf_key = (dataset, expr, condition)
+        if (leaf := ctx.leaves.get(leaf_key)) is None:
+            taken = {*ctx.keys.values(), *(other.name for other in ctx.leaves.values())}
+            leaf = Leaf(
+                dataset=dataset, expr=expr, name=find_new_name(taken, name), condition=condition
+            )
+            ctx.leaves[leaf_key] = leaf
+        return exp.column(leaf.name, ctx.alias)
+
+    def _join_branches(
+        self, datasets: list[exp.Table], condition: exp.Expr | None, ctx: ConstellationContext
+    ) -> exp.Select:
+        """
+        Builds one branch per dataset and joins them on the keys (step 2 in the class docstring).
+
+        The result outputs each key once, taken from whichever branch has it, and each leaf from
+        the branch that computes it. Branches are aliased by their dataset's name without its
+        namespace, numbered when datasets of two namespaces share it, e.g. `sales` and `sales_2`.
+        """
+        aliases: dict[exp.Table, str] = {}
+        for dataset in datasets:
+            aliases[dataset] = find_new_name(set(aliases.values()), dataset.name)
+        branches = [
+            self._branch(dataset, condition, ctx).subquery(aliases[dataset], copy=False)
+            for dataset in datasets
+        ]
+        names = list(aliases.values())
+        keys = {name: self._coalesce(name, names) for name in ctx.keys.values()}
+        # A NULL key is a proper group, e.g. sales with no customer, but `=` never matches NULL,
+        # so with several branches it comes out as one row per branch. Grouping by the keys merges
+        # them into one, like GROUP BY does, and MAX picks each column's only non-NULL value.
+        # This assumes leaves are values MAX accepts: Postgres, for one, has no MAX(boolean), so
+        # a metric like BOOL_OR(...) only merges on engines that do.
+        merge = bool(keys) and len(branches) > 1
+        outputs = [
+            exp.alias_(key, name) if len(datasets) > 1 else key for name, key in keys.items()
+        ]
+        for leaf in ctx.leaves.values():
+            column = exp.column(leaf.name, aliases[leaf.dataset])
+            outputs.append(exp.alias_(exp.Max(this=column), leaf.name) if merge else column)
+        joined = exp.select(*outputs, copy=False).from_(branches[0], copy=False)
+
+        for n, branch in enumerate(branches[1:], start=1):
+            # A key can be missing from any earlier branch, so compare with all of them. Without
+            # keys, every branch is a single row, so they all match.
+            conditions = [
+                exp.column(name, names[n]).eq(self._coalesce(name, names[:n])) for name in keys
+            ]
+            on = exp.and_(*conditions) if conditions else exp.true()
+            joined = joined.join(branch, join_type="FULL OUTER", on=on, copy=False)
+
+        if merge:
+            joined = joined.group_by(*keys.values(), copy=False)
+        return joined
+
+    def _key_datasets(
+        self, condition: exp.Expr | None, ctx: ConstellationContext
+    ) -> list[exp.Table]:
+        """
+        Returns the datasets a query without metrics reads its keys from: every dataset that
+        provides all the roles the query uses.
+
+        The constellation stands for all its datasets stacked, so e.g.
+        `SELECT customer.birth_year FROM universe GROUP BY 1` gives the birth years of the
+        customers in `store_sales`, `catalog_sales` and the other datasets with a customer, and
+        `ship_date.year` skips `store_sales`, which has no ship date.
+        """
+        roles = {
+            parts[0]
+            for expr in (*ctx.keys, condition)
+            if expr is not None
+            for path in self._paths(expr)
+            if (parts := self._role_parts(path, ctx)) is not None
+        }
+        datasets = [
+            dataset
+            for dataset in sorted(
+                ctx.constellation.assigned_datasets(), key=lambda name: (name.db, name.name)
+            )
+            if all(ctx.constellation.get_assignment(dataset, role) is not None for role in roles)
+        ]
+        if not datasets:
+            names = ", ".join(sorted(self._sql(ctx, role) for role in roles))
+            self._fail(ctx, f"No dataset provides every role the query uses: {names}")
+        return datasets
+
+    def _branch(
+        self, dataset: exp.Table, condition: exp.Expr | None, ctx: ConstellationContext
+    ) -> exp.Select:
+        """
+        Builds the query for one dataset: its keys and leaves over the rows that pass `condition`,
+        with one row per key, or per row when stacking.
+
+        For example, `store_sales` by `item.category` gives the sales amount per item category,
+        with the category read through `store_sales`' own assignment for the `item` role.
+        """
+        translated = {name: self._translate(expr, dataset, ctx) for expr, name in ctx.keys.items()}
+        projections = [exp.alias_(expr, name) for name, expr in translated.items()]
+        projections += [
+            exp.alias_(self._filtered(leaf, ctx), leaf.name)
+            for leaf in ctx.leaves.values()
+            if leaf.dataset == dataset and not ctx.stack
+        ]
+
+        table = exp.table_(dataset.this, db=dataset.args["db"], catalog=ctx.catalog)
+        table.meta[TABLE_ARGUMENTS] = ctx.arguments
+        branch = exp.select(*projections, copy=False).from_(table, copy=False)
+        if condition is not None:
+            branch.where(self._translate(condition, dataset, ctx), copy=False)
+        # Constant keys like `'store'` don't need grouping, and some engines reject or
+        # misread `GROUP BY <literal>`.
+        groupings = [expr for expr in translated.values() if expr.find(exp.Column)]
+        if groupings and not ctx.stack:
+            branch.set("group", exp.Group(expressions=groupings))
+        # ExpandDatasetsRule expects qualified SQL. Keys are named after the query's output names,
+        # which can match a column of the dataset, so they mustn't be read as alias references.
+        # Its aliases are normalized names, which normalizing again could change, so they're quoted.
+        qualified = qualify(
+            quote_identifiers(branch, dialect=ctx.dialect),
+            schema=ctx.cluster,
+            dialect=ctx.dialect,
+            quote_identifiers=False,
+            expand_alias_refs=False,
+        )
+        # Without leaves or a GROUP BY, e.g. just `'store' AS sales_channel`, the branch has a row
+        # per row of the dataset, which the join would multiply by every other branch's rows. The
+        # DISTINCT goes outside, since ExpandDatasetsRule rebuilds the dataset's select without it.
+        if not ctx.leaves and not groupings and not ctx.stack:
+            return (
+                exp.select("*")
+                .from_(qualified.subquery(dataset.name, copy=False), copy=False)
+                .distinct()
+            )
+        return qualified
+
+    def _translate(self, expr: exp.Expr, dataset: exp.Table, ctx: ConstellationContext) -> exp.Expr:
+        """
+        Rewrites an expression over roles, i.e. a key, the WHERE or a FILTER condition, for
+        `dataset`, replacing each role with the dataset's assignment for it, e.g.
+
+            customer.birth_year on store_sales            ->  customer.birth_year
+            customer.birth_year on catalog_sales          ->  bill_customer.birth_year
+            customer.birth_year // 10 on catalog_sales    ->  bill_customer.birth_year // 10
+            sales_channel on store_sales                  ->  'store'
+
+        since `catalog_sales` assigns the role `customer` as its `bill_customer` relationship,
+        and `store_sales` assigns `sales_channel` as the constant `'store'`.
+        """
+
+        def rewrite(node: exp.Expr) -> exp.Expr:
+            # A subquery's columns are its own tables', e.g. `c.customer`, but not
+            # `universe.customer.birth_year`.
+            correlated = node.find_ancestor(exp.Query) is not None
+            if (correlated and not self._is_outer(node, ctx)) or (
+                parts := self._role_parts(node, ctx)
+            ) is None:
+                return node
+            role, *path = parts
+            if (assignment := ctx.constellation.get_assignment(dataset, role)) is None:
+                name = f"{self._sql(ctx, ctx.constellation.name)}.{self._sql(ctx, role)}"
+                self._fail(ctx, f"Dataset {self._sql(ctx, dataset)} does not provide role {name}")
+            assigned = (
+                # The branch reads the dataset by its name.
+                assignment.expr.transform(
+                    lambda n: _prefixed(n, dataset.this) if isinstance(n, exp.Column) else n
+                )
+                if correlated
+                else assignment.expr.copy()
+            )
+            if not path:
+                # A reference role is a key to follow, not a value to group or filter by.
+                # TODO: let a bare reference role stand for its key, e.g. `customer IN (SELECT ...)`
+                # filtering on `customer.c_customer_sk`, or on `ss_customer_sk` without a join.
+                field = ctx.constellation.get_column(role)
+                if field is not None and field.datatype and field.datatype.is_type(exp.DType.REF):
+                    name = self._sql(ctx, role)
+                    self._fail(
+                        ctx, f"Use an attribute of the entity role {name}, e.g. {name}.<attribute>"
+                    )
+                return _replace(node, assigned)
+            if not isinstance(assigned, exp.Column | exp.Dot):
+                self._fail(
+                    ctx,
+                    f"Role {self._sql(ctx, role)} is not a relationship on dataset "
+                    f"{self._sql(ctx, dataset)}",
+                    assigned,
+                )
+            return exp.Dot.build([assigned, *path])
+
+        return expr.transform(rewrite)
+
+    def _role_parts(
+        self, expr: exp.Expr, ctx: ConstellationContext
+    ) -> tuple[exp.Identifier, ...] | None:
+        """
+        Returns the name path of a role reference, or None if `expr` isn't one.
+
+        For example, `customer.birth_year` gives `("customer", "birth_year")`, and a metric such
+        as `store_sales_amount` gives None.
+        """
+        if not isinstance(expr, exp.Column | exp.Dot):
+            return None
+        # qualify writes role references as `universe.item.category`, so drop the alias.
+        parts = list(path_names(expr))
+        if parts and parts[0].name == ctx.alias:
+            parts = parts[1:]
+        if not parts or not isinstance(ctx.constellation.get_column(parts[0]), Role):
+            return None
+        return tuple(parts)
+
+    def _filtered(self, leaf: Leaf, ctx: ConstellationContext) -> exp.Expr:
+        """
+        Returns a leaf's aggregate with its condition translated for its dataset, e.g.
+        `customer.birth_year > 1960` on `catalog_sales` gives
+        `AGG(sales_amount) FILTER (WHERE bill_customer.birth_year > 1960)`. A filter the metric's
+        own definition has keeps applying, combined with `AND`.
+        """
+        if leaf.condition is None:
+            return leaf.expr
+        condition = self._translate(leaf.condition, leaf.dataset, ctx)
+        if isinstance(leaf.expr, exp.Filter):
+            own = leaf.expr.expression.this
+            return exp.Filter(
+                this=leaf.expr.this.copy(), expression=exp.Where(this=exp.and_(own, condition))
+            )
+        return exp.Filter(this=leaf.expr.copy(), expression=exp.Where(this=condition))
+
+    def _check_condition(
+        self, condition: exp.Expr, message: str, ctx: ConstellationContext
+    ) -> None:
+        """
+        Fails unless `condition` only reads roles, since it's translated into each branch. A
+        metric only has a value per key, after aggregating, and any other column, e.g. one of an
+        outer query, would be read from each branch's own dataset instead.
+        """
+        # A subquery that doesn't read the constellation's rows is pushed down as is, e.g.
+        # `item.brand IN (SELECT brand FROM ...)`, and each branch runs its own copy.
+        for subquery in find_all_in_scope(condition, exp.Query):
+            scopes = traverse_scope(subquery)
+            own = {name for scope in scopes for name in scope.sources}
+            for column in (column for scope in scopes for column in scope.columns):
+                path: exp.Column | exp.Dot = column
+                while isinstance(path.parent, exp.Dot):
+                    path = path.parent
+                # qualify can read an outer role as a field of the subquery's own table, e.g.
+                # `universe.customer.birth_year` as `c.universe.customer.birth_year`.
+                # TODO: an unqualified role, e.g. `item.category` read as `r.item.category`,
+                # can't be told apart from the subquery's own relationship.
+                if ctx.alias in (part.name for part in path.parts[1:-1]):
+                    self._fail(
+                        ctx, "Correlated subqueries are not supported on constellations yet", path
+                    )
+                if column.table not in own and (
+                    column.table != ctx.alias or self._role_parts(path, ctx) is None
+                ):
+                    self._fail(
+                        ctx,
+                        "A correlated subquery on a constellation can only read its roles",
+                        path,
+                    )
+        invalid = find_in_scope(condition, exp.AggFunc, exp.Window) or next(
+            (path for path in self._paths(condition) if self._role_parts(path, ctx) is None), None
+        )
+        if invalid:
+            self._fail(ctx, message, invalid)
+
+    def _is_outer(self, path: exp.Expr, ctx: ConstellationContext) -> bool:
+        """
+        Whether a path in a subquery reads the constellation, e.g. in
+
+            WHERE EXISTS (SELECT 1 FROM c WHERE c.year = universe.customer.birth_year)
+
+        but not in `... FROM (SELECT ...) AS universe WHERE universe.customer = 1`.
+        """
+        if not isinstance(path, exp.Column | exp.Dot):
+            return False
+        names = path_names(path)
+        if not names or names[0].name != ctx.alias:
+            return False
+        select = path.find_ancestor(exp.Select)
+        while select is not None:
+            sources = [select.args.get("from_"), *select.args.get("joins", [])]
+            if any(source and source.this.alias_or_name == ctx.alias for source in sources):
+                return False
+            select = select.find_ancestor(exp.Select)
+        return True
+
+    def _paths(self, expr: exp.Expr) -> list[exp.Expr]:
+        """
+        Returns the column paths in `expr`, each whole: sqlglot parses `customer.birth_year` as
+        dots over the `customer` column, so that column gives the full `customer.birth_year`.
+        Columns of a subquery belong to its own tables, so they're skipped.
+        """
+        paths = []
+        for column in find_all_in_scope(expr, exp.Column):
+            path: exp.Expr = column
+            while isinstance(path.parent, exp.Dot):
+                path = path.parent
+            paths.append(path)
+        return paths
+
+    def _field(self, column: exp.Column, ctx: ConstellationContext) -> Metric | Role | None:
+        """
+        Returns the metric or role a column of the constellation names, e.g. the `Metric` for
+        `store_sales_amount`. A column of any other table gives None.
+        """
+        if column.table and column.table != ctx.alias:
+            return None
+        return ctx.constellation.get_column(column.this)
+
+    def _sql(self, ctx: ConstellationContext, name: exp.Expr) -> str:
+        return name.sql(ctx.registry.dialect)
+
+    def _coalesce(self, name: str, branches: list[str]) -> exp.Expr:
+        """
+        Reads a key from the first of the `branches`, by alias, that has it, since the key may be
+        missing from any of them. A single branch's column is returned as is.
+        """
+        columns = [exp.column(name, branch) for branch in branches]
+        if len(columns) == 1:
+            return columns[0]
+        return exp.Coalesce(this=columns[0], expressions=columns[1:])