From d02b5968fc75f5e49220f878d70251322c8967ef Mon Sep 17 00:00:00 2001 From: Steven Kearnes Date: Mon, 17 Aug 2026 21:52:46 -0400 Subject: [PATCH] Let a query order by a value under a repeated level An ordering key and an aggregate's argument both had to be scalar, which left "the ten highest-yielding reactions" unwritable: a yield lives under outcomes, products, and measurements, so the path resolves to a list rather than a number. Both a cheap model and an expensive one reached for it independently while translating questions, and neither could have succeeded. A Reduction names how to reduce that list to the one value the reaction is judged by, and resolve() already returns a list expression for a repeated path, so it compiles to one DuckDB list aggregate around what the resolver produces. It is refused over a scalar path, which needs no reduction, and inside an aggregated query's order_by, where there is no reaction left to reduce over. Co-Authored-By: Claude Opus 5 (1M context) --- ord_schema/search/README.md | 11 ++- ord_schema/search/query.py | 72 +++++++++++++++- ord_schema/search/query_test.py | 148 ++++++++++++++++++++++++++++++++ 3 files changed, 226 insertions(+), 5 deletions(-) diff --git a/ord_schema/search/README.md b/ord_schema/search/README.md index 3c23e66a..d7d957ab 100644 --- a/ord_schema/search/README.md +++ b/ord_schema/search/README.md @@ -55,9 +55,11 @@ Value = { literal: } | { compound: } Aggregate = { group_by: [Path], measures: [{ fn: "count"|"count_distinct"|"sum"|"avg"|"min"|"max", - path?: Path, name: string }] } + path?: Path | Reduction, name: string }] } -Order = { key: string, descending?: bool } +Reduction = { reduce: "min"|"max"|"avg"|"sum"|"count", path: Path } + +Order = { key: string | Reduction, descending?: bool } ``` A `Path` is a dotted column path such as `conditions.temperature.setpoint_kelvin`. Inside an @@ -76,6 +78,10 @@ is a compile error rather than a wrong answer: - Operators must suit the leaf type — `contains` is text-only, ordering is numeric-only. - `group_by` paths must be scalar, so the number of groups is bounded by the values a column holds rather than by an explosion over a repeated level. +- A `Reduction` is the one place a repeated path is read without a quantifier: it reduces + that reaction's own elements to a single value, so `max` over `outcomes.products.measurements.percentage.value` + is the reaction's best yield rather than the corpus's. Its path must cross a repeated level; + a scalar one is refused, since it would give the same query two spellings. - A `{"compound": ...}` value is resolved through [`ord_schema.resolvers`](../resolvers.py) and **bound as a parameter**, so the model names compounds and never spells structures. - A `substructure`/`similarity` path must name a compound's `smiles`, inside a @@ -131,6 +137,7 @@ in what they read. Worked examples, with the route each clause takes: | solvent-free: no component is a solvent | `forall inputs.components` | pivot — the index shows which elements match, never that all of them do | | pyridine **and** a boronic acid in one component | `exists inputs.components` | pivot — two structure predicates is one more than an occurrence row can carry | | a desired product with a yield above 50% | `exists outcomes.products`, nested `exists measurements` | both levels' pivots, joined on the ordinal prefix | +| the ten highest-yielding reactions | `order_by` a `reduce` over `outcomes.products.measurements` | no quantifier: a list aggregate over the projection | Every pivot row above falls to the elements when no pivot is available — a level the budget refused, or one with neither an artifact nor room to build. The answer does not diff --git a/ord_schema/search/query.py b/ord_schema/search/query.py index a37731da..458fd90c 100644 --- a/ord_schema/search/query.py +++ b/ord_schema/search/query.py @@ -458,11 +458,40 @@ def _check(self) -> "Similarity": Quantifier.model_rebuild() +# DuckDB's list aggregates, which ignore the nulls a list may hold. `count` filters +# rather than taking len(), which would count them. +_REDUCERS = { + "min": "list_min({expression})", + "max": "list_max({expression})", + "avg": "list_avg({expression})", + "sum": "list_sum({expression})", + "count": "len(list_filter({expression}, value -> value IS NOT NULL))", +} + + +class Reduction(BaseModel): + """One value per reaction, reduced from a path that crosses a repeated level. + + An ordering key and an aggregate's argument both have to be scalar, which leaves + "the highest-yielding reactions" unwritable: a yield lives under outcomes, products, + and measurements, so the path resolves to a list rather than a number. This reduces + that list to the one value the reaction is judged by, leaving the aggregate to + combine those across reactions. + + Attributes: + reduce: How to reduce the list. ``count`` counts the values that are present. + path: A dotted path crossing at least one repeated level. + """ + + reduce: Literal["min", "max", "avg", "sum", "count"] + path: str + + class Measure(BaseModel): """One aggregate over the matching rows.""" fn: Literal["count", "count_distinct", "sum", "avg", "min", "max"] - path: str | None = None + path: str | Reduction | None = None name: str @model_validator(mode="after") @@ -488,7 +517,7 @@ class Aggregate(BaseModel): class Order(BaseModel): """How to sort the result.""" - key: str + key: str | Reduction descending: bool = False @@ -988,6 +1017,30 @@ def _scalar(path: str, schema: pa.Schema, what: str) -> str: return resolved.expression +def _reduced(reduction: Reduction, schema: pa.Schema) -> str: + """Returns the expression reducing a repeated path to one value per reaction. + + Args: + reduction: What to reduce, and how. + schema: Schema the path resolves against. + + Returns: + A DuckDB expression yielding one scalar per reaction, NULL where the reaction + holds no elements under that path at all. + + Raises: + QueryError: If the path is already scalar, which needs no reduction; accepting + one would give the same query two spellings. + """ + resolved = resolve(reduction.path, schema=schema) + if not resolved.repeated: + raise QueryError( + f"{reduction.path}: {reduction.reduce} reduces a repeated level, and this " + f"path is already scalar; order by the path itself" + ) + return _REDUCERS[reduction.reduce].format(expression=resolved.expression) + + def compile_query( query: Query, *, @@ -1033,6 +1086,10 @@ def compile_query( for measure in query.aggregate.measures: if measure.path is None: argument = "*" + elif isinstance(measure.path, Reduction): + argument = _reduced(measure.path, schema) + if measure.fn == "count_distinct": + argument = f"DISTINCT {argument}" else: argument = _scalar(measure.path, schema, measure.fn) if measure.fn == "count_distinct": @@ -1071,7 +1128,16 @@ def compile_query( if query.order_by: keys = [] for order in query.order_by: - if orderable is None: + if isinstance(order.key, Reduction): + if orderable is not None: + # After grouping there is no reaction left to reduce over: the + # reduction is one input to a measure, not a key beside it. + raise QueryError( + "an aggregated query orders by a measure name or a group_by " + "path; reduce inside a measure instead" + ) + key = _reduced(order.key, schema) + elif orderable is None: key = _scalar(order.key, schema, "order_by") elif order.key in orderable: key = ( diff --git a/ord_schema/search/query_test.py b/ord_schema/search/query_test.py index ca118a5d..ac3e404f 100644 --- a/ord_schema/search/query_test.py +++ b/ord_schema/search/query_test.py @@ -393,6 +393,154 @@ def test_rejected_before_compilation(payload): query.Query.model_validate(payload) +# Reductions over a repeated level + + +def test_ordering_by_a_reduction_over_a_repeated_path(): + # A yield lives under outcomes, products, and measurements, so ordering by it needs + # the list reduced to the one number the reaction is judged by. + compiled = _compile( + { + "order_by": [ + { + "key": { + "reduce": "max", + "path": "outcomes.products.measurements.percentage.value", + }, + "descending": True, + } + ], + "limit": 10, + } + ) + assert "list_max(" in compiled.sql + assert compiled.sql.endswith("DESC LIMIT 10") + + +def test_a_reduction_reaches_the_same_rows_the_elements_do(): + # The reduction is over the reaction's own elements, so it agrees with the list the + # projection holds rather than with a corpus-wide aggregate. + compiled = _compile( + { + "order_by": [ + { + "key": { + "reduce": "max", + "path": "outcomes.products.measurements.percentage.value", + } + } + ] + } + ) + resolved = query.resolve("outcomes.products.measurements.percentage.value") + assert resolved.expression in compiled.sql + + +def test_a_measure_may_reduce_a_repeated_path(): + compiled = _compile( + { + "aggregate": { + "group_by": ["conditions.temperature.setpoint_kelvin"], + "measures": [ + { + "fn": "avg", + "path": { + "reduce": "max", + "path": "outcomes.products.measurements.percentage.value", + }, + "name": "best_yield", + } + ], + } + } + ) + assert "avg(list_max(" in compiled.sql + + +@pytest.mark.parametrize( + ("reducer", "expected"), + [ + ("min", "list_min("), + ("max", "list_max("), + ("avg", "list_avg("), + ("sum", "list_sum("), + # Counting what is there means filtering the nulls a list may hold, since len() + # would count them. + ("count", "len(list_filter("), + ], +) +def test_each_reducer_compiles_to_its_list_aggregate(reducer, expected): + compiled = _compile( + { + "order_by": [ + { + "key": { + "reduce": reducer, + "path": "outcomes.products.measurements.percentage.value", + } + } + ] + } + ) + assert expected in compiled.sql + + +def test_a_reduction_over_a_scalar_path_is_refused(): + # A scalar needs no reducing, and accepting one would give the same query two + # spellings, one of which wraps a value in a single-element list. + with pytest.raises(query.QueryError, match="already scalar"): + _compile( + { + "order_by": [ + { + "key": { + "reduce": "max", + "path": "conditions.temperature.setpoint_kelvin", + } + } + ] + } + ) + + +def test_an_aggregated_query_cannot_order_by_a_reduction(): + # After grouping there is no reaction left to reduce over; the reduction belongs + # inside a measure, where it is one input to the aggregate. + with pytest.raises(query.QueryError, match="reduce inside a measure"): + _compile( + { + "aggregate": {"measures": [{"fn": "count", "name": "n"}]}, + "order_by": [ + { + "key": { + "reduce": "max", + "path": "outcomes.products.measurements.percentage.value", + } + } + ], + } + ) + + +def test_a_reduction_runs(): + # The expression has to be DuckDB the planner accepts, not merely a string. + compiled = _compile( + { + "order_by": [ + { + "key": { + "reduce": "max", + "path": "outcomes.products.measurements.percentage.value", + }, + "descending": True, + } + ], + "limit": 5, + } + ) + assert _run(compiled) == [] + + # The whole point