fix: execute semantic nest lambdas in aggregate

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()}

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,

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:
     """