Pass 1: single-measure filtered aggregates

crates/integration-tests/tests/filtered_aggregates.rs

@@ -324,3 +324,198 @@ async fn hand_filtered_kll_compiles_and_executes() {
         );
     }
 }
+
+/// Whether `root` binds in the executor against an in-memory `events`.
+fn binds(root: &Rc<OperatorNode>) -> Result<(), String> {
+    use asap_executor::sources::{DataSources, MemorySource};
+    use asap_executor::values::Batch;
+    use asap_types::ir::physical_export::PhysicalASAPOperatorPayload;
+    let wire = physical_common::compile_physical_asap_dag(root).map_err(|e| e.to_string())?;
+    let (source, schema) = wire
+        .nodes
+        .iter()
+        .find_map(|node| match &node.payload {
+            PhysicalASAPOperatorPayload::NonASAP(NonASAPOp::Scan { source, .. }) => {
+                Some((source.clone(), node.output_schema.clone()))
+            }
+            _ => None,
+        })
+        .unwrap();
+    let schema = std::sync::Arc::new(schema);
+    let batch = Batch::try_new(schema.clone(), vec![]).unwrap();
+    let mut sources = DataSources::default();
+    sources
+        .register(
+            source,
+            std::sync::Arc::new(MemorySource::new(schema, vec![batch]).unwrap()),
+        )
+        .unwrap();
+    asap_executor::physical_planner::bind_with_data_sources(
+        &wire,
+        BTreeMap::new(),
+        &[wire.roots[0] as u64],
+        &sources,
+    )
+    .map(|_| ())
+    .map_err(|e| e.to_string())
+}
+
+/// Every Pass 1 alternative for `sql` at `target`, composed, with the
+/// families of its `SummaryAgg`s and whether they all carry a filter.
+async fn alternatives(sql: &str, target: AccuracyTarget) -> Vec<(String, Rc<OperatorNode>)> {
+    let root = lower_sql(sql, &catalog(), target).await.unwrap();
+    let inventory =
+        enumerate_local_logical_candidates(vec![(0, QueryRoot::Operator(root))], &BTreeMap::new())
+            .unwrap();
+    enumerate_choices(&inventory, usize::MAX)
+        .into_iter()
+        .map(|choice| {
+            let roots = compose_logical_candidate(&inventory, &choice)
+                .unwrap_or_else(|e| panic!("{choice:?} composes: {e}"));
+            let QueryRoot::Operator(root) = &roots[0].1 else {
+                panic!("operator root")
+            };
+            let label = summary_builds(root)
+                .iter()
+                .map(|build| {
+                    let Some(ASAPOp::SummaryAgg {
+                        family,
+                        grouping,
+                        filter,
+                        ..
+                    }) = build.asap()
+                    else {
+                        unreachable!()
+                    };
+                    assert!(filter.is_some(), "{choice:?} keeps the filter");
+                    let family = match family {
+                        asap_types::ir::schema::FieldDataType::Sketch(kind, _) => {
+                            format!("{:?}", kind.algorithm())
+                        }
+                        other => format!("{other:?}"),
+                    };
+                    match grouping == &Default::default() {
+                        true => family,
+                        false => format!("Hydra{family}"),
+                    }
+                })
+                .collect::<Vec<_>>()
+                .join("+");
+            (label, root.clone())
+        })
+        .collect()
+}
+
+/// Pass 1 offers a filtered single-measure aggregate the same alternatives
+/// as the unfiltered one, each with `SummaryAgg.filter` set. All compose;
+/// an alternative binds in the executor exactly when its unfiltered
+/// counterpart does. The exact `Count` accumulator and HydraCms execute to
+/// the exact plan's counts, `b` included.
+#[tokio::test]
+async fn pass1_offers_filtered_count_alternatives() {
+    // ε = 0.1 keeps the Hydra grid inside the default memory limit.
+    let target = AccuracyTarget::EpsilonDelta {
+        epsilon: 0.1,
+        delta: 0.01,
+    };
+    let filtered = alternatives(
+        "SELECT g, COUNT(*) FILTER (WHERE x > 0) AS c FROM events GROUP BY g",
+        target.clone(),
+    )
+    .await;
+    let plain = lower_sql(
+        "SELECT g, COUNT(*) AS c FROM events GROUP BY g",
+        &catalog(),
+        target.clone(),
+    )
+    .await
+    .unwrap();
+    let inventory =
+        enumerate_local_logical_candidates(vec![(0, QueryRoot::Operator(plain))], &BTreeMap::new())
+            .unwrap();
+    let plain: Vec<_> = enumerate_choices(&inventory, usize::MAX)
+        .into_iter()
+        .map(|choice| {
+            let roots = compose_logical_candidate(&inventory, &choice).unwrap();
+            let QueryRoot::Operator(root) = &roots[0].1 else {
+                panic!("operator root")
+            };
+            root.clone()
+        })
+        .collect();
+    let labels: Vec<_> = filtered.iter().map(|(label, _)| label.as_str()).collect();
+    assert_eq!(
+        labels,
+        [
+            "",
+            "ExactAggregate(Count, Count)",
+            "Cms",
+            "CountSketch",
+            "UnivMon",
+            "HydraCms"
+        ]
+    );
+    let expected = printed([
+        vec![s("a"), Value::Int64(2)],
+        vec![s("b"), Value::Int64(0)],
+        vec![s("c"), Value::Int64(1)],
+    ]);
+    for ((label, root), plain) in filtered.iter().zip(&plain) {
+        assert_eq!(binds(root).is_ok(), binds(plain).is_ok(), "{label}");
+        if matches!(
+            label.as_str(),
+            "" | "ExactAggregate(Count, Count)" | "HydraCms"
+        ) {
+            // Few groups in a wide grid: Hydra's estimate is exact here.
+            assert_eq!(sorted(root), expected, "{label}");
+        }
+    }
+}
+
+/// A filtered SUM and a filtered percentile: every alternative composes
+/// with the filter, and the exact `Sum` accumulator, KLL and DDSketch
+/// execute to the exact answer, reading NULL for `b`, which has no `x > 0`
+/// row.
+#[tokio::test]
+async fn pass1_filtered_sum_and_quantile_alternatives_execute() {
+    let target = AccuracyTarget::EpsilonDelta {
+        epsilon: 0.01,
+        delta: 0.01,
+    };
+    // SQL's exact percentile has no native implementation.
+    for (sql, expected, executable) in [
+        (
+            "SELECT g, SUM(x) FILTER (WHERE x > 0) AS v FROM events GROUP BY g",
+            ["a 3.0", "b NULL", "c 5.0"],
+            vec!["", "ExactAggregate(Sum, Sum)"],
+        ),
+        (
+            "SELECT g, approx_percentile_cont(x, 0.5) FILTER (WHERE x > 0) AS v FROM events GROUP BY g",
+            ["a 1.0", "b NULL", "c 5.0"],
+            vec!["Kll", "DDSketch"],
+        ),
+    ] {
+        let alternatives = alternatives(sql, target.clone()).await;
+        let mut executed = vec![];
+        for (label, root) in &alternatives {
+            if binds(root).is_ok() {
+                let actual = sorted(root);
+                // Within 2%: DDSketch's relative-error guarantee at ε = 0.01.
+                let close = |a: &str, e: &str| {
+                    a == e
+                        || matches!((a.parse::<f64>(), e.parse::<f64>()),
+                            (Ok(a), Ok(e)) if (a - e).abs() <= 0.02 * e.abs())
+                };
+                assert!(
+                    actual.len() == expected.len()
+                        && actual.iter().zip(expected).all(|(a, e)| {
+                            a.split(' ').zip(e.split(' ')).all(|(a, e)| close(a, e))
+                        }),
+                    "{sql}: {label}: {actual:?}"
+                );
+                executed.push(label.as_str());
+            }
+        }
+        assert_eq!(executed, executable, "{sql}");
+    }
+}

crates/integration-tests/tests/summary_coverage_e2e.rs

@@ -1,20 +1,22 @@
 //! Query string → planned summary state → `coverage()`, for the examples of
-//! the ASAP primitive schema design doc (#573 §4.2.2). Each query is lowered,
-//! searched and selected with the default cost model, and the quantile's KLL
-//! `SummaryAgg` in the selected DAG is checked.
+//! the ASAP primitive schema design doc (#573 §4.2.2). Each query is lowered
+//! and its Stage 1 logical candidates are enumerated; the quantile's KLL
+//! `SummaryAgg` is checked. Candidates, not the selected plan: the plan the
+//! cost models select for these small queries may build no summary.
 
+use std::collections::BTreeMap;
 use std::ops::Bound;
 use std::rc::Rc;
 
 use asap_frontend_sql::{lower_sql, SqlCatalog};
 use asap_integration_tests::fixtures::lower_promql;
-use asap_logical_optimizer::search_workload;
-use asap_plan_selection::candidate_selection::global_selection;
-use asap_plan_selection::DefaultCostModel;
+use asap_logical_optimizer::pass1::logical_candidates::{
+    compose_logical_candidate, enumerate_choices, enumerate_local_logical_candidates,
+};
 use asap_types::ir::properties::summary_coverage::{ColumnIdentity, Constraint, SelectionBox};
 use asap_types::ir::scalar::ScalarValue;
 use asap_types::ir::schema::{DataType, Field, FieldDataType, Schema, SketchAlgorithm};
-use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode, Predicate, ScalarExpr};
+use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode, Predicate, QueryRoot, ScalarExpr};
 use asap_types::types::AccuracyTarget;
 
 fn accuracy() -> AccuracyTarget {
@@ -24,25 +26,25 @@ fn accuracy() -> AccuracyTarget {
     }
 }
 
-fn plan(pre: Rc<OperatorNode>) -> Rc<OperatorNode> {
-    let space = search_workload(vec![("query", pre)]);
-    global_selection(&space, &DefaultCostModel)
-        .assemble_selected_dag(&space.roots[0].1)
-        .expect("materialization failed")
-        .expect("root must be discovered")
-}
-
-/// The single KLL `SummaryAgg` of the selected DAG.
-fn kll_state(root: &Rc<OperatorNode>) -> Rc<OperatorNode> {
-    let states: Vec<_> = OperatorNode::reachable(root)
-        .into_iter()
-        .filter(|node| {
+/// The KLL `SummaryAgg` of the first logical candidate that builds one.
+fn kll_state(root: Rc<OperatorNode>) -> Rc<OperatorNode> {
+    let inventory =
+        enumerate_local_logical_candidates(vec![(0, QueryRoot::Operator(root))], &BTreeMap::new())
+            .unwrap();
+    for choice in enumerate_choices(&inventory, usize::MAX) {
+        let roots = compose_logical_candidate(&inventory, &choice).unwrap();
+        let QueryRoot::Operator(root) = &roots[0].1 else {
+            panic!("operator root")
+        };
+        let state = OperatorNode::reachable(root).into_iter().find(|node| {
             matches!(node.asap(), Some(ASAPOp::SummaryAgg { family: FieldDataType::Sketch(kind, _), .. })
                 if kind.algorithm() == &SketchAlgorithm::Kll)
-        })
-        .collect();
-    assert_eq!(states.len(), 1, "expected one KLL state in {root:#?}");
-    states[0].clone()
+        });
+        if let Some(state) = state {
+            return state;
+        }
+    }
+    panic!("no candidate builds a KLL state");
 }
 
 async fn sql_state(sql: &str) -> Rc<OperatorNode> {
@@ -57,13 +59,11 @@ async fn sql_state(sql: &str) -> Rc<OperatorNode> {
     let pre = lower_sql(sql, &catalog, accuracy())
         .await
         .unwrap_or_else(|e| panic!("lower failed for {sql:?}: {e}"));
-    kll_state(&plan(pre))
+    kll_state(pre)
 }
 
 fn promql_state(query: &str) -> Rc<OperatorNode> {
-    kll_state(&plan(
-        lower_promql(query, accuracy()).expect("lowering failed"),
-    ))
+    kll_state(lower_promql(query, accuracy()).expect("lowering failed"))
 }
 
 fn column(table: Option<&str>, name: &str) -> ColumnIdentity {
@@ -157,6 +157,41 @@ async fn sql_renamed_column_moves_and_expression_stays() {
     );
 }
 
+/// The worked example of #573 §4.2.2: the `FILTER` range and the renamed
+/// value set move to the selection, the expression condition stays.
+#[tokio::test]
+async fn sql_worked_example() {
+    let state = sql_state(
+        "SELECT job, approx_percentile_cont(latency, 0.99) FILTER (WHERE latency < 100) \
+         FROM (SELECT job, region AS r, latency FROM t \
+               WHERE region = 'us' AND latency * 2 > 10) \
+         GROUP BY job",
+    )
+    .await;
+    assert_eq!(
+        state.coverage().unwrap().selection,
+        vec![SelectionBox {
+            columns: [
+                (
+                    column(None, "latency"),
+                    Constraint::Interval {
+                        lower: Bound::Unbounded,
+                        upper: Bound::Excluded(ScalarValue::Float64(100.0)),
+                    },
+                ),
+                (column(None, "r"), Constraint::In(vec![utf8("us")])),
+            ]
+            .into(),
+            relative_time: None,
+        }]
+    );
+    let definition = &state.coverage().unwrap().definition;
+    assert!(matches!(
+        definition.asap(),
+        Some(ASAPOp::SummaryAgg { filter: None, .. })
+    ));
+}
+
 #[test]
 fn promql_offset_window_is_relative_time() {
     let state = promql_state("quantile_over_time(0.99, m[1m] offset 2m)");

crates/logical-optimizer/src/pass1/logical_candidates.rs

@@ -627,11 +627,13 @@ fn realize(
             "multi-measure aggregate",
         ));
     };
-    if !filters.is_empty() || having.is_some() {
-        return Err(LogicalCandidateError::Unsupported(
-            "filtered or HAVING aggregate",
-        ));
+    if having.is_some() {
+        return Err(LogicalCandidateError::Unsupported("HAVING aggregate"));
     }
+    // A single measure's row filter (SQL `FILTER (WHERE …)`) becomes the
+    // summary's filter over the same input rows; every group is kept. A
+    // whole-expression target is never filtered.
+    let filter = filters.first().cloned().flatten();
     let whole =
         match absorbs {
             true => Some(whole_expression_input(target).ok_or(
@@ -679,7 +681,7 @@ fn realize(
                 input: input.clone(),
                 reduction: reduction.clone(),
                 grouping: grouping.clone(),
-                filter: None,
+                filter: filter.clone(),
             },
         ))?)
     };