diff --git a/src/boring_semantic_layer/serialization/__init__.py b/src/boring_semantic_layer/serialization/__init__.py index eea30bbe..f89c98be 100644 --- a/src/boring_semantic_layer/serialization/__init__.py +++ b/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": }. +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()} diff --git a/src/boring_semantic_layer/serialization/reconstruct.py b/src/boring_semantic_layer/serialization/reconstruct.py index f026b4c8..bdaa169e 100644 --- a/src/boring_semantic_layer/serialization/reconstruct.py +++ b/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. diff --git a/src/boring_semantic_layer/tests/test_xorq_join_leaf_recovery.py b/src/boring_semantic_layer/tests/test_xorq_join_leaf_recovery.py index ce9438a3..815505bb 100644 --- a/src/boring_semantic_layer/tests/test_xorq_join_leaf_recovery.py +++ b/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]