From 1e7fa98d237e1661710bb144538f6c2be3618829 Mon Sep 17 00:00:00 2001 From: Hussain Sultan Date: Sat, 11 Jul 2026 08:48:01 -0400 Subject: [PATCH] fix: execute semantic nest lambdas in aggregate --- src/boring_semantic_layer/expr.py | 12 ++ src/boring_semantic_layer/ops.py | 150 ++++++++++++++++++ .../tests/test_malloy_inspired.py | 82 ++++++++++ 3 files changed, 244 insertions(+) diff --git a/src/boring_semantic_layer/expr.py b/src/boring_semantic_layer/expr.py index cbe402b3..2c9af93a 100644 --- a/src/boring_semantic_layer/expr.py +++ b/src/boring_semantic_layer/expr.py @@ -1384,6 +1384,18 @@ def nest_agg(ibis_tbl): f"got {type(result).__module__}.{type(result).__name__}", ) + # Keep the semantic lambda available to the aggregate compiler. + # Treating this callable like an ordinary measure causes + # _make_base_measure() to invoke it with ColumnScope, which is + # correct for scalar measures but loses the BSL query API used by + # lambdas such as + # + # lambda t: t.group_by("carrier").aggregate("flight_count") + # + # SemanticAggregateOp lowers that query at the combined outer + + # inner grain and collects its rows. Raw ibis/xorq nest lambdas + # continue through the callable above unchanged. + nest_agg.__bsl_semantic_nest__ = fn return nest_agg nest_aggs = {name: make_nest_agg(fn) for name, fn in nest.items()} diff --git a/src/boring_semantic_layer/ops.py b/src/boring_semantic_layer/ops.py index 1bc870f4..c83e626c 100644 --- a/src/boring_semantic_layer/ops.py +++ b/src/boring_semantic_layer/ops.py @@ -2783,6 +2783,10 @@ def required_columns(self) -> dict[str, set[str]]: return combined.to_dict() def to_untagged(self): + semantic_nest_result = self._lower_semantic_nests() + if semantic_nest_result is not None: + return semantic_nest_result + all_roots = _find_all_root_models(self.source) def find_join_in_tree(node): @@ -2980,6 +2984,152 @@ def collect_filters_to_join(node): is_post_agg=is_post_agg, ) + def _lower_semantic_nests(self): + """Lower nest lambdas which return a BSL aggregate pipeline. + + A nest is correlated with this aggregate's grouping keys. The inner + query therefore has to run at ``outer keys + inner keys`` grain before + its rows are collected into an array of structs. Invoking the lambda + as a scalar measure cannot express that correlation and, historically, + also passed a :class:`ColumnScope` where the BSL query API was expected. + + Return ``None`` when every nest lambda is a raw ibis/xorq lambda; that + preserves the original lightweight struct-collect implementation. + """ + from .expr import SemanticTable + + marked: dict[str, Any] = {} + regular: dict[str, Any] = {} + for name, wrapped in self.aggs.items(): + fn = _unwrap(wrapped) + semantic_fn = getattr(fn, "__bsl_semantic_nest__", None) + if semantic_fn is None: + regular[name] = fn + continue + + # group_by().aggregate() stores the SemanticGroupByOp as the + # aggregate source. Nested pipelines should start from the same + # ungrouped semantic source, not from that bookkeeping wrapper. + base_source = ( + self.source.source + if isinstance(self.source, SemanticGroupByOp) + else self.source + ) + try: + nested = semantic_fn(SemanticTable(base_source)) + except (AttributeError, TypeError, NotImplementedError): + # The established raw-table form (notably + # ``t.group_by(["a", "b"])``) is intentionally not valid BSL + # syntax. Leave it on the old path. + regular[name] = fn + continue + if not isinstance(nested, SemanticTable): + regular[name] = fn + continue + marked[name] = nested.op() + + if not marked: + return None + + # Compile the outer query without the nest measures. Reusing a normal + # SemanticAggregateOp keeps joins, filters, calculated measures, and + # fan-out-safe aggregation on their existing paths. + outer_op = SemanticAggregateOp( + source=self.source, + keys=self.keys, + aggs=regular, + nested_columns=(), + ) + result = outer_op.to_untagged() + + for output_name, pipeline_op in marked.items(): + inner_agg, order_keys, limit_spec, predicates = self._split_nest_pipeline( + pipeline_op, output_name + ) + fine_keys = tuple(dict.fromkeys((*self.keys, *inner_agg.keys))) + fine_op = SemanticAggregateOp( + source=inner_agg.source, + keys=fine_keys, + aggs={name: _unwrap(fn) for name, fn in inner_agg.aggs.items()}, + nested_columns=(), + ) + fine = fine_op.to_untagged() + + # Filters above the inner aggregate are HAVING predicates and must + # run at the fine grain, before collection. + for predicate in reversed(predicates): + fine = fine.filter(_resolve_expr(_unwrap(predicate), ColumnScope(_tbl=fine))) + + struct_cols = tuple(dict.fromkeys((*inner_agg.keys, *inner_agg.aggs.keys()))) + if not struct_cols: + raise TypeError( + f"Nest lambda for {output_name!r} must produce at least one column" + ) + struct_values = {col: fine[col] for col in struct_cols} + first_col = next(iter(struct_values.values())) + if "xorq.vendor.ibis" in type(first_col).__module__: + from ._xorq import ibis as ibis_mod + else: + ibis_mod = ibis + + collect_kwargs = {} + if order_keys: + collect_kwargs["order_by"] = [ + self._resolve_nest_order_key(key, fine) for key in order_keys + ] + collected_expr = ibis_mod.struct(struct_values).collect(**collect_kwargs) + if limit_spec is not None: + n, offset = limit_spec + collected_expr = collected_expr[offset : offset + n] + + if self.keys: + part = fine.group_by([fine[key] for key in self.keys]).aggregate( + **{output_name: collected_expr} + ) + from .nested_compile import join_tables + + result = join_tables(self.keys, [result, part]) + else: + part = fine.aggregate(**{output_name: collected_expr}) + result = result.cross_join(part) + + wanted = [*self.keys, *self.aggs.keys()] + return result.select(*wanted) + + @staticmethod + def _split_nest_pipeline(pipeline_op, output_name): + """Return inner aggregate plus post-aggregate pipeline modifiers.""" + order_keys: tuple[Any, ...] = () + limit_spec: tuple[int, int] | None = None + predicates: list[Any] = [] + current = pipeline_op + while not isinstance(current, SemanticAggregateOp): + if isinstance(current, SemanticLimitOp): + if limit_spec is not None: + raise TypeError(f"Nest lambda for {output_name!r} has multiple limits") + limit_spec = (current.n, current.offset) + elif isinstance(current, SemanticOrderByOp): + if order_keys: + raise TypeError( + f"Nest lambda for {output_name!r} has multiple order_by steps" + ) + order_keys = current.keys + elif isinstance(current, SemanticFilterOp): + predicates.append(current.predicate) + else: + raise TypeError( + f"Nest lambda for {output_name!r} must return an aggregate pipeline, " + f"got {type(current).__name__}" + ) + current = current.source + return current, order_keys, limit_spec, predicates + + @staticmethod + def _resolve_nest_order_key(key, table): + if isinstance(key, str): + return table[key] + return _resolve_expr(_unwrap(key), ColumnScope(_tbl=table)) + def _to_untagged_with_preagg( self, all_roots: list, diff --git a/src/boring_semantic_layer/tests/test_malloy_inspired.py b/src/boring_semantic_layer/tests/test_malloy_inspired.py index 94e74949..0fa18ae9 100644 --- a/src/boring_semantic_layer/tests/test_malloy_inspired.py +++ b/src/boring_semantic_layer/tests/test_malloy_inspired.py @@ -740,6 +740,88 @@ def test_nested_group_by_with_xorq(self): nested_codes = {item["code"] for item in co_row["data"]} assert nested_codes == {"DEN", "ASE"} + def test_semantic_nest_lambda_executes_at_correlated_grain(self): + """A semantic nest pipeline receives BSL, not ColumnScope, semantics.""" + con = ibis.duckdb.connect(":memory:") + flights_tbl = con.create_table( + "semantic_nest_flights", + pd.DataFrame( + { + "origin": ["NYC", "NYC", "NYC", "LAX"], + "carrier": ["AA", "AA", "DL", "UA"], + } + ), + ) + flights = ( + to_semantic_table(flights_tbl, name="flights") + .with_dimensions( + origin=lambda t: t.origin, + carrier=lambda t: t.carrier, + ) + .with_measures(flight_count=lambda t: t.count()) + ) + + result = ( + flights.group_by("origin") + .aggregate( + "flight_count", + nest={ + "by_carrier": lambda t: t.group_by("carrier").aggregate( + "flight_count" + ) + }, + ) + .execute() + .set_index("origin") + ) + + assert result.loc["NYC", "flight_count"] == 3 + assert { + row["carrier"]: row["flight_count"] + for row in result.loc["NYC", "by_carrier"] + } == {"AA": 2, "DL": 1} + assert result.loc["LAX", "by_carrier"] == [{"carrier": "UA", "flight_count": 1}] + + def test_semantic_nest_preserves_inner_order_and_limit(self): + con = ibis.duckdb.connect(":memory:") + flights_tbl = con.create_table( + "ordered_semantic_nest_flights", + pd.DataFrame( + { + "origin": ["NYC"] * 6, + "carrier": ["AA", "AA", "AA", "DL", "DL", "UA"], + } + ), + ) + flights = ( + to_semantic_table(flights_tbl) + .with_dimensions( + origin=lambda t: t.origin, + carrier=lambda t: t.carrier, + ) + .with_measures(flight_count=lambda t: t.count()) + ) + + result = ( + flights.group_by("origin") + .aggregate( + nest={ + "top_carriers": lambda t: ( + t.group_by("carrier") + .aggregate("flight_count") + .order_by(lambda t: t.flight_count.desc()) + .limit(2) + ) + } + ) + .execute() + ) + + assert result.loc[0, "top_carriers"] == [ + {"carrier": "AA", "flight_count": 3}, + {"carrier": "DL", "flight_count": 2}, + ] + class TestCrossJoinAggregation: """