feat: leaf expression markers — recovery reads authored leaves, never guesses

src/boring_semantic_layer/serialization/__init__.py

@@ -52,6 +52,77 @@ def do_import():
 # to_tagged
 # ---------------------------------------------------------------------------
 
+#: Marker tag stamped on each leaf model's AUTHORED table expression inside
+#: the lowered payload. Recovery reads the marked subtree back verbatim
+#: instead of re-deriving the leaf from the lowered plan (base-relation
+#: walking / join splitting), which discarded authored shaping and could not
+#: see through aggregation lowering at all. metadata: {"tag": BSL_LEAF_TAG,
+#: "leaf": <model name>}.
+BSL_LEAF_TAG = "__bsl_leaf__"
+
+
+def _collect_leaf_tables(op, out=None):
+    """Map each leaf SemanticTableOp's name to its authored table op.
+
+    Walks the SEMANTIC op tree (source/left/right chains; join wrappers
+    descend into their _source_join). First declaration wins on a duplicated
+    name — an ambiguous name cannot be marked meaningfully and falls back to
+    heuristic recovery.
+    """
+    from .. import ops as bsl_ops
+
+    if out is None:
+        out = {}
+    if op is None:
+        return out
+    if isinstance(op, bsl_ops.SemanticTableOp):
+        source_join = getattr(op, "_source_join", None)
+        if source_join is not None:
+            return _collect_leaf_tables(source_join, out)
+        name = getattr(op, "name", None)
+        table = getattr(op, "table", None)
+        if table is not None and hasattr(table, "op"):
+            table = table.op()
+        if name and table is not None and name not in out:
+            out[name] = table
+        return out
+    for attr in ("source", "left", "right"):
+        child = getattr(op, attr, None)
+        if child is not None:
+            _collect_leaf_tables(child, out)
+    return out
+
+
+def _mark_leaf_tables(xorq_table, leaf_tables):
+    """Wrap every occurrence of an authored leaf table in a marker tag.
+
+    The lowered expression embeds each leaf's table op structurally, so a
+    node-equality rewrite finds them wherever lowering placed them — under
+    rename projections, under pre-aggregation legs, under a query's
+    Aggregate. A leaf whose table was itself rewritten by lowering (so no
+    node matches) is simply left unmarked and recovers via the heuristics.
+    """
+    from .._xorq import replace_nodes
+
+    by_op = {}
+    for name, table in leaf_tables.items():
+        by_op.setdefault(table, name)
+    if not by_op:
+        return xorq_table
+
+    def replacer(node, _kwargs):
+        # Plain-callable replacers must recreate the node from _kwargs
+        # themselves — that is how child substitutions propagate upward
+        # (see ibis graph._coerce_replacer; Pattern/Mapping replacers get
+        # this for free, callables do not).
+        rebuilt = node.__recreate__(_kwargs) if _kwargs else node
+        name = by_op.get(node)
+        if name is None:
+            return rebuilt
+        return rebuilt.to_expr().hashing_tag(tag=BSL_LEAF_TAG, leaf=name).op()
+
+    return replace_nodes(replacer, xorq_table).to_expr()
+
 
 def to_tagged(semantic_expr, aggregate_cache_storage=None):
     """Tag a BSL expression with serialized metadata.
@@ -112,6 +183,7 @@ def extract_path_from_view(table_name):
             return node
 
         xorq_table = replace_nodes(replace_read_parquet, xorq_table).to_expr()
+        xorq_table = _mark_leaf_tables(xorq_table, _collect_leaf_tables(op))
 
         metadata = extract_op_tree(op, context)
         tag_data = {k: freeze(v) for k, v in metadata.items()}

src/boring_semantic_layer/serialization/reconstruct.py

@@ -200,8 +200,13 @@ def _reconstruct_table():
             _source_join=join_op,
         )
 
+    marked_leaf = _find_marked_leaf(xorq_expr, metadata.get("name"))
     return bsl_expr.SemanticModel(
-        table=_strip_internal_join_temps(_reconstruct_table()),
+        table=(
+            marked_leaf
+            if marked_leaf is not None
+            else _strip_internal_join_temps(_reconstruct_table())
+        ),
         dimensions=dimensions,
         measures=measures,
         calc_measures=calc_measures,
@@ -359,6 +364,30 @@ def _reconstruct_limit(metadata: dict, xorq_expr, source, context: BSLSerializat
     return source.limit(n=int(metadata.get("n", 0)), offset=int(metadata.get("offset", 0)))
 
 
+def _find_marked_leaf(xorq_expr, name):
+    """Return the AUTHORED table expression for leaf *name*, if the payload
+    carries a leaf marker (BSL_LEAF_TAG, written by to_tagged since the
+    leaf-payload change). This is the lossless recovery path: the marked
+    subtree is exactly what the author passed to to_semantic_table —
+    shaped views, renames, even aggregate-grain rollups — so no heuristic
+    re-derivation from the lowered plan is needed. Returns None when the
+    payload predates markers or the leaf was not markable (unnamed model,
+    or its table op was rewritten by lowering)."""
+    if not name:
+        return None
+    from .._xorq import Tag, walk_nodes
+
+    try:
+        for tag_op in walk_nodes((Tag,), xorq_expr):
+            metadata = getattr(tag_op, "metadata", None) or {}
+            if metadata.get("tag") == "__bsl_leaf__" and metadata.get("leaf") == name:
+                parent = tag_op.parent
+                return parent.to_expr() if hasattr(parent, "to_expr") else parent
+    except Exception:
+        return None
+    return None
+
+
 def _strip_internal_join_temps(expr):
     """Invert BSL's temporary join-key renames on a recovered leaf table.
 

src/boring_semantic_layer/tests/test_xorq_join_leaf_recovery.py

@@ -199,17 +199,11 @@ def test_shaped_join_leaves_round_trip():
     assert list(result["transactions.gross_purchases"]) == [0.0, 11.0]
 
 
-@pytest.mark.xfail(
-    strict=True,
-    reason="Known gap: a lowered QUERY entry over a shaped view still base-walks "
-    "— BSL's query-time injected mutates and authored shaping are "
-    "indistinguishable in the lowered tree. Needs a lowering-time base-boundary "
-    "marker. MODEL entries (what the star-schema doctrine catalogs) are covered "
-    "by the tests above.",
-)
-def test_query_entry_still_digs_under_the_aggregate():
-    """Recovery of a tagged aggregate over a SHAPED view should replay the
-    query against the shaped base — today the base walk discards the shaping."""
+def test_query_entry_recovers_the_shaped_base():
+    """Recovery of a tagged aggregate over a SHAPED view replays the query
+    against the shaped base: the leaf marker survives under the Aggregate,
+    so recovery no longer needs to guess where lowering ends and authored
+    shaping begins (this was a strict xfail before leaf markers)."""
     accounts_view, _ = _seam_free_tables()
     model = (
         to_semantic_table(accounts_view, name="accounts")
@@ -333,3 +327,169 @@ def test_join_many_aggregate_query_fails_loud_not_silently_wrong(plain_tables):
 
     with pytest.raises(ValueError, match="Round-trip could not recover"):
         from_tagged(to_tagged(query))
+
+
+# ---------------------------------------------------------------------------
+# Leaf markers: shaped and aggregate-grain leaves round-trip losslessly
+# ---------------------------------------------------------------------------
+
+
+def test_grouped_grain_fact_round_trips():
+    """A fact view PRE-AGGREGATED to a coarser grain is a legal model leaf.
+
+    Dimensional-modeling aggregate fact tables: roll transactions to account
+    grain as a deferred xorq view, then a thin semantic layer of simple sums
+    on top. Before leaf markers this failed round-trip — the base walk dug
+    under the Aggregate and the rollup columns vanished.
+    """
+    con = xo.connect()
+    raw_txn = con.register(
+        pd.DataFrame(
+            {
+                "transaction_id": [1, 2, 3, 4],
+                "account_id": [1, 1, 2, 3],
+                "amount": [5.0, 6.0, 7.0, 8.0],
+            }
+        ),
+        "grain_transactions",
+    )
+    raw_accounts = con.register(
+        pd.DataFrame({"account_id": [1, 2, 3], "customer_id": [10, 10, 20]}),
+        "grain_accounts",
+    )
+    txn_by_account = raw_txn.group_by("account_id").aggregate(
+        txn_count=xo._.count(), txn_amount=xo._.amount.sum()
+    )
+
+    accounts = to_semantic_table(raw_accounts, name="accounts").with_dimensions(
+        account_id=lambda t: t.account_id,
+        customer_id=lambda t: t.customer_id,
+    )
+    fact = to_semantic_table(txn_by_account, name="txn_rollup").with_measures(
+        total_amount=lambda t: t.txn_amount.sum(),
+        total_txns=lambda t: t.txn_count.sum(),
+    )
+    model = accounts.join_one(fact, on=lambda a, f: a.account_id == f.account_id)
+
+    recovered = from_tagged(to_tagged(model))
+    result = (
+        recovered.group_by("accounts.customer_id")
+        .aggregate("txn_rollup.total_amount", "txn_rollup.total_txns")
+        .to_untagged()
+        .execute()
+        .sort_values("accounts.customer_id")
+        .reset_index(drop=True)
+    )
+    assert list(result["txn_rollup.total_amount"]) == [18.0, 8.0]
+    assert list(result["txn_rollup.total_txns"]) == [3, 1]
+
+
+def test_single_table_model_over_ibis_aggregate_round_trips():
+    """to_semantic_table(table.group_by(...).aggregate(...)) — a semantic
+    model DIRECTLY over an ibis/xorq aggregate query — is a legal model.
+
+    The aggregate view fixes the grain (account level here); the semantic
+    layer names simple reductions over the rollup columns. Everything stays
+    deferred; the marker carries the authored aggregate through round-trip.
+    """
+    con = xo.connect()
+    raw_txn = con.register(
+        pd.DataFrame(
+            {
+                "transaction_id": [1, 2, 3, 4],
+                "account_id": [1, 1, 2, 3],
+                "amount": [5.0, 6.0, 7.0, 8.0],
+            }
+        ),
+        "st_transactions",
+    )
+    rollup = raw_txn.group_by("account_id").aggregate(
+        txn_count=xo._.count(), txn_amount=xo._.amount.sum()
+    )
+    model = (
+        to_semantic_table(rollup, name="account_rollup")
+        .with_dimensions(account_id=lambda t: t.account_id)
+        .with_measures(
+            accounts_with_txns=lambda t: t.count(),
+            total_amount=lambda t: t.txn_amount.sum(),
+            max_account_amount=lambda t: t.txn_amount.max(),
+        )
+    )
+
+    recovered = from_tagged(to_tagged(model))
+    result = (
+        recovered.aggregate("accounts_with_txns", "total_amount", "max_account_amount")
+        .to_untagged()
+        .execute()
+    )
+    assert list(result["accounts_with_txns"]) == [3]
+    assert list(result["total_amount"]) == [26.0]
+    assert list(result["max_account_amount"]) == [11.0]
+
+
+def test_grouped_grain_model_survives_disk_round_trip(tmp_path):
+    """The aggregate-grain model survives xorq build -> load_expr ->
+    ls.builder — the full catalog path."""
+    pytest.importorskip("duckdb")
+    db = str(tmp_path / "wh.ddb")
+    seed = xo.duckdb.connect(db)
+    seed.create_table(
+        "transactions",
+        pd.DataFrame(
+            {
+                "transaction_id": [1, 2, 3, 4],
+                "account_id": [1, 1, 2, 3],
+                "amount": [5.0, 6.0, 7.0, 8.0],
+            }
+        ),
+    )
+    del seed
+    con = xo.duckdb.connect(db)
+    rollup = (
+        con.table("transactions")
+        .group_by("account_id")
+        .aggregate(txn_amount=xo._.amount.sum())
+    )
+    model = to_semantic_table(rollup, name="account_rollup").with_measures(
+        total_amount=lambda t: t.txn_amount.sum(),
+    )
+    tagged = to_tagged(model)
+
+    path = xo.build_expr(tagged, builds_dir=str(tmp_path / "builds"))
+    loaded = xo.load_expr(path)
+    recovered = loaded.ls.builder
+    result = recovered.aggregate("total_amount").to_untagged().execute()
+    assert list(result["total_amount"]) == [26.0]
+
+
+def test_unmarked_payload_falls_back_to_heuristics():
+    """Payloads written before leaf markers (no __bsl_leaf__ tags) must keep
+    recovering via the heuristic paths — strip the marker tags off a fresh
+    payload to simulate one."""
+    from boring_semantic_layer._xorq import Tag, walk_nodes
+    from boring_semantic_layer.serialization import BSL_LEAF_TAG
+
+    accounts_view, _txn_view = _seam_free_tables()
+    accounts = to_semantic_table(accounts_view, name="accounts").with_measures(
+        open_account_count=lambda t: t.is_open.sum(),
+    )
+    tagged = to_tagged(accounts)
+
+    from xorq.common.utils.graph_utils import replace_nodes
+
+    def strip(node, _kwargs):
+        rebuilt = node.__recreate__(_kwargs) if _kwargs else node
+        if isinstance(rebuilt, Tag) and (rebuilt.metadata or {}).get("tag") == BSL_LEAF_TAG:
+            return rebuilt.parent
+        return rebuilt
+
+    stripped = replace_nodes(strip, tagged).to_expr()
+    assert not [
+        t
+        for t in walk_nodes((Tag,), stripped)
+        if (getattr(t, "metadata", {}) or {}).get("tag") == BSL_LEAF_TAG
+    ]
+    recovered = from_tagged(stripped)
+    result = recovered.aggregate("open_account_count").to_untagged().execute()
+    # heuristic per-row-chain preservation still recovers the shaped leaf
+    assert list(result["open_account_count"]) == [2]