From 39bf0c48871551c16be87806a2e6128189ae1838 Mon Sep 17 00:00:00 2001 From: mday-io Date: Tue, 3 Nov 2026 22:33:27 +0000 Subject: [PATCH] fix(metrics): infer measure table for aggregates with ORDER BY A metric such as ARRAY_AGG(t.x ORDER BY t.x) failed to load with "Could not infer a measures table". sqlglot parses the aggregate's argument as Order(this=, expressions=[...]), and is_measure did not climb through the Order node. Climb it via its `this` arg only, so columns that appear just in the ORDER BY clause are never the measure. Partially addresses #6169. Signed-off-by: mday-io --- sqlmesh/core/metric/definition.py | 4 ++++ tests/core/metric/test_metric.py | 8 ++++++++ tests/core/metric/test_rewriter.py | 28 +++++++++++++++++++++++++++- 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/sqlmesh/core/metric/definition.py b/sqlmesh/core/metric/definition.py index 6119a883ed..fc523298c8 100644 --- a/sqlmesh/core/metric/definition.py +++ b/sqlmesh/core/metric/definition.py @@ -199,6 +199,10 @@ def is_measure(node: exp.Expr) -> bool: return is_measure(parent) if isinstance(parent, (exp.Binary, exp.Paren, exp.Distinct)): return is_measure(parent) + # AGG(x ORDER BY y) is parsed as AGG(Order(this=x, expressions=[Ordered(y)])). + # Only the Order's `this` can be the measure, not the ORDER BY expressions. + if isinstance(parent, exp.Order) and node.arg_key == "this": + return is_measure(parent) return False for node in expression.walk(): diff --git a/tests/core/metric/test_metric.py b/tests/core/metric/test_metric.py index 51c97fbe3d..f15c4c8032 100644 --- a/tests/core/metric/test_metric.py +++ b/tests/core/metric/test_metric.py @@ -154,3 +154,11 @@ def test_get_measure_and_dim_tables(): "a", ("b",), ) + assert _get_measure_and_dim_tables(d.parse_one("ARRAY_AGG(a.x ORDER BY a.x)")) == ("a", ()) + assert _get_measure_and_dim_tables(d.parse_one("ARRAY_AGG(a.x ORDER BY b.y)")) == ("a", ("b",)) + assert _get_measure_and_dim_tables( + d.parse_one("ARRAY_AGG(IF(c.z = 'dim', a.x, NULL) ORDER BY b.y)") + ) == ("a", ("b", "c")) + # A column that only appears in the ORDER BY clause cannot be the measure. + with pytest.raises(ConfigError, match="Could not infer a measures table"): + _get_measure_and_dim_tables(d.parse_one("ARRAY_AGG(1 ORDER BY b.y)")) diff --git a/tests/core/metric/test_rewriter.py b/tests/core/metric/test_rewriter.py index 9fe8169508..d2ce821c2d 100644 --- a/tests/core/metric/test_rewriter.py +++ b/tests/core/metric/test_rewriter.py @@ -1,7 +1,9 @@ from sqlglot import parse_one -from sqlmesh.core.metric import rewrite +from sqlmesh.core import dialect as d +from sqlmesh.core.metric import expand_metrics, load_metric_ddl, rewrite from sqlmesh.core.reference import ReferenceGraph +from sqlmesh.utils import UniqueKeyDict def test_rewrite(sushi_context_pre_scheduling, assert_exp_eq): @@ -222,3 +224,27 @@ def test_rewrite(sushi_context_pre_scheduling, assert_exp_eq): # dialect=context.config.dialect, # ), # ) + + +def test_rewrite_order_by_metric(sushi_context_pre_scheduling): + context = sushi_context_pre_scheduling + graph = ReferenceGraph(context.models.values()) + meta = load_metric_ddl( + d.parse_one( + "METRIC (name sorted_amounts, expression ARRAY_AGG(sushi.orders.amount ORDER BY sushi.orders.amount))", + dialect=context.config.dialect, + ), + dialect=context.config.dialect, + ) + metrics = expand_metrics(UniqueKeyDict("metas", {meta.name: meta})) + + query = rewrite( + "SELECT c.customer_id, METRIC(sorted_amounts) FROM sushi.customers AS c GROUP BY c.customer_id", + graph=graph, + metrics=metrics, + ) + sql = query.sql() + # Stop before the generated NULLS LAST so the assertion doesn't depend on null ordering. + assert "ARRAY_AGG(sushi__orders.amount ORDER BY sushi__orders.amount" in sql + assert "sushi.orders AS sushi__orders" in sql + parse_one(sql, dialect=context.config.dialect)