Pass 1: single-measure filtered aggregates

crates/integration-tests/tests/filtered_aggregates.rs

@@ -328,3 +328,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::export::{NonASAPOpKind, 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::Relational {
+                operator: NonASAPOpKind::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(),
+        &[u64::from(wire.roots[0].0)],
+        &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/logical-optimizer/src/pass1/logical_candidates.rs

@@ -628,11 +628,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,14 +681,16 @@ fn realize(
             input: input.clone(),
             reduction: reduction.clone(),
             grouping: grouping.clone(),
-            filter: None,
+            filter: filter.clone(),
         }))?;
         Ok::<_, LogicalCandidateError>(Rc::new(state.with_coverage(coverage)?))
     };
     let state = match window {
         // Whole-source coverage is declared, not proven: Pass 1 trusts that
         // the state holds every observation of its source that reaches it
-        // (#570).
+        // (#570). A filtered state holds fewer, so this over-states its
+        // population; that only makes disjointness harder to prove, and the
+        // filter is part of the node, so sharing never confuses the two.
         WindowForm::Whole => {
             let coverage = SummaryCoverage {
                 source: single_source(&child)?,

crates/planner/tests/stage_pipeline_selection.rs

@@ -434,6 +434,62 @@ async fn sql_hydra_count_dp_equals_exhaustive() {
     assert_dp_matches_exhaustive(&inventory, &workload, 18);
 }
 
+/// Filtered single-measure aggregates (`FILTER (WHERE …)`) get the same
+/// alternatives as unfiltered ones, each a filtered summary, and the DP
+/// still selects the exhaustive minimum.
+#[tokio::test]
+async fn sql_filtered_aggregates_dp_equals_exhaustive() {
+    let accuracy = AccuracyTarget::Epsilon(0.1);
+    let queries = [
+        "SELECT l_orderkey, COUNT(*) FILTER (WHERE l_extendedprice > 100) FROM lineitem GROUP BY l_orderkey",
+        "SELECT approx_percentile_cont(l_extendedprice, 0.99) FILTER (WHERE l_orderkey = 7) FROM lineitem",
+    ];
+    let workload = PlanningWorkload {
+        query_workload: QueryWorkload {
+            language: QueryLanguage::SQL(SqlDialect::DataFusionSQL),
+            query_batch: Some(queries.iter().map(|q| batch(q, accuracy.clone())).collect()),
+            repeating_queries: None,
+        },
+        data_workload: Some(DataWorkload {
+            arrival: DataArrival::AtRest,
+            ..Default::default()
+        }),
+    };
+    let catalog = SqlCatalog::new().with_table(
+        "lineitem",
+        Schema::new(vec![
+            Field::plain("l_orderkey", DataType::Int64, false),
+            Field::plain("l_extendedprice", DataType::Float64, false),
+        ]),
+    );
+    let mut roots = Vec::new();
+    for (index, query) in queries.iter().enumerate() {
+        let root = lower_sql_dialect(query, &catalog, SqlDialect::DataFusionSQL, accuracy.clone())
+            .await
+            .expect("lowers");
+        roots.push((index, QueryRoot::Operator(root)));
+    }
+    let inventory = stage1_logical_candidates(roots, &Default::default(), &[]).expect("Stage 1");
+    // (pass-through, Count acc, CMS, CountSketch, UnivMon, HydraCms) × (pass-through, KLL, DDSketch).
+    assert_dp_matches_exhaustive(&inventory, &workload, 18);
+    // Every filtered alternative builds through Stages 1 and 2.
+    let exhaustive = select_exhaustive(
+        &inventory,
+        &targets(&workload),
+        &workload.data_workload.clone().unwrap_or_default(),
+        PlanningModels::builtin(),
+        MAX_ENUMERATED_CANDIDATES,
+    )
+    .expect("exhaustive selection");
+    let unbuilt: Vec<_> = exhaustive
+        .selection
+        .rejected
+        .iter()
+        .filter(|r| r.reason.starts_with("Stage "))
+        .collect();
+    assert!(unbuilt.is_empty(), "{unbuilt:#?}");
+}
+
 /// PromQL queries, each with its own ε (δ = 0.001).
 fn promql_with(queries: &[(&str, f64)]) -> PlanningWorkload {
     let mut workload = promql(&[], 1_000);

crates/types/src/ir/operator/asap.rs

@@ -275,29 +275,36 @@ impl ASAPOp {
                 {
                     dtype = DataType::Float64;
                 }
+                // A filtered quantile is NULL for a group without a matching
+                // row, as the filtered aggregate it realizes is.
+                let nullable = matches!(query, SketchStatistic::Quantile { .. })
+                    && built_summary(summary_input).is_some_and(|built| {
+                        matches!(
+                            built.asap(),
+                            Some(SummaryAgg {
+                                filter: Some(_),
+                                ..
+                            })
+                        )
+                    });
                 let mut schema = input.clone();
                 for field in &mut schema.fields {
                     if !field.is_plain() {
-                        *field = Field::plain(name, dtype.clone(), false);
+                        *field = Field::plain(name, dtype.clone(), nullable);
                     }
                 }
                 schema
             }
             FinalizeExactAccumulator { child } => {
                 // A merge's inputs have identical schemas (tumbling panes),
                 // so its first input names the finalized value.
-                let mut built = child;
-                while let Some(SummaryMerge { children }) = built.asap() {
-                    match children.first() {
-                        Some(first) => built = first,
-                        None => break,
-                    }
-                }
+                let built = built_summary(child).unwrap_or(child);
                 let value_result = if let Some(ASAPOp::SummaryAgg {
                     child: source,
                     family: FieldDataType::ExactAggregate(kind, _),
                     input,
                     reduction,
+                    filter,
                     ..
                 }) = built.asap()
                 {
@@ -334,7 +341,9 @@ impl ASAPOp {
                                     reduction: reduction.clone(),
                                     measures: vec![measure],
                                     output_names: vec![],
-                                    filters: vec![],
+                                    // A filtered SUM/MIN/MAX is NULL for a
+                                    // group without a matching row.
+                                    filters: filter.iter().map(|f| Some(f.clone())).collect(),
                                     having: None,
                                 }
                                 .output_schema()
@@ -739,3 +748,14 @@ fn source_kind(node: &OperatorNode) -> OperatorResultKind {
         },
     }
 }
+
+/// The `SummaryAgg` that builds `state`, through any merges. A merge's
+/// inputs have identical schemas (tumbling panes), so its first input
+/// stands for all of them.
+fn built_summary(state: &Rc<OperatorNode>) -> Option<&Rc<OperatorNode>> {
+    let mut built = state;
+    while let Some(ASAPOp::SummaryMerge { children }) = built.asap() {
+        built = children.first()?;
+    }
+    Some(built)
+}