feat(asap-physical-operators): execute HydraCms grouped counts and filtered aggregates

crates/asap-physical-operators/src/capability.rs

@@ -22,7 +22,7 @@ pub fn validate_summary_kernel(
     grouping: &GroupingStrategy,
 ) -> Result<(), String> {
     if grouping != &GroupingStrategy::PerSubpopulationInstance {
-        return Err("shared summary grouping has no registered kernel".into());
+        return validate_hydra_cms(family, input, grouping);
     }
     let keyed = match family {
         SummaryFamilyType::ExactAggregate(kind, params) => {
@@ -116,6 +116,94 @@ pub fn validate_summary_kernel(
     Ok(())
 }
 
+/// The only shared grouping with a kernel: Hydra over Count-Min. Each update
+/// adds one item column's value to its group, with weight 1 (a row count) or
+/// a weight column proven non-negative (a Count-Min point count only
+/// overestimates when no weight is negative).
+fn validate_hydra_cms(
+    family: &SummaryFamilyType,
+    input: &SummaryUpdate,
+    grouping: &GroupingStrategy,
+) -> Result<(), String> {
+    use planner_types::post_asap::{NonNegativeWeightProof, SummaryInputExpr, WeightDomain};
+    if !matches!(family, SummaryFamilyType::Sketch(_, layout) if layout == grouping) {
+        return Err("Planner family and operator grouping disagree".into());
+    }
+    hydra_cms_shape(family)?;
+    if !matches!(input.item, Some(SummaryInputExpr::Column(_))) {
+        return Err("HydraCms requires one item column".into());
+    }
+    let WeightDomain::NonNegative { proof } = &input.weight_domain else {
+        return Err("HydraCms weights must be non-negative; the weight's sign is unknown".into());
+    };
+    match (&input.weight, proof) {
+        (SummaryInputExpr::Constant(w), NonNegativeWeightProof::UnitCount) if *w == 1.0 => Ok(()),
+        (SummaryInputExpr::Column(_), _) => Ok(()),
+        _ => Err("HydraCms weight must be the unit count or a non-negative weight column".into()),
+    }
+}
+
+/// `(shared_rows, shared_columns, width, depth)` of a HydraCms family. The
+/// family is the per-group Count-Min it emulates, with the Hydra's own width
+/// and depth.
+pub(crate) fn hydra_cms_shape(
+    family: &SummaryFamilyType,
+) -> Result<(usize, usize, usize, usize), String> {
+    use planner_types::post_asap::{HydraKind, HydraParams};
+    let SummaryFamilyType::Sketch(
+        kind,
+        GroupingStrategy::SharedMultiSubpopulation {
+            kind: HydraKind::HydraCms,
+            params:
+                HydraParams::HydraCms {
+                    width,
+                    depth,
+                    shared_rows,
+                    shared_columns,
+                },
+        },
+    ) = family
+    else {
+        return Err("shared summary grouping has no registered kernel".into());
+    };
+    if kind.algorithm() != &SketchAlgorithm::Cms
+        || kind.params()
+            != &(SketchParams::Cms {
+                width: *width,
+                depth: *depth,
+            })
+    {
+        return Err("HydraCms family must be Count-Min with the Hydra width and depth".into());
+    }
+    if !valid_matrix(*width, *depth) || !valid_matrix(*shared_columns, *shared_rows) {
+        return Err("invalid HydraCms dimensions".into());
+    }
+    Ok((
+        *shared_rows as usize,
+        *shared_columns as usize,
+        *width as usize,
+        *depth as usize,
+    ))
+}
+
+/// A filtered `SummaryAgg` never runs at ingestion time. An ingestion-time
+/// build reads raw samples (`$population`, `$timestamp`, `value`), which carry
+/// none of the query's columns for the filter predicate to read, so only a
+/// query-time build can apply it. The ingestion-time compiler and the
+/// binding of any ingestion-timed `SummaryAgg` apply this check. Planner-side
+/// timing validation (materialization timings, Stage 2 eligibility) does not
+/// apply it yet: until it gains a check or a `PruneReason`, a plan placing a
+/// filtered `SummaryAgg` at ingestion is accepted by the planner and rejected
+/// here at compile time.
+pub fn validate_ingestion_time_summary<C>(
+    filter: Option<&planner_types::ir::Predicate<C>>,
+) -> Result<(), String> {
+    match filter {
+        Some(_) => Err("a filtered summary update runs only at query time".into()),
+        None => Ok(()),
+    }
+}
+
 fn valid_matrix(width: u32, depth: u32) -> bool {
     // Construction uses the kernel's native row hashing, so no encoded-size
     // limit applies here.
@@ -147,9 +235,13 @@ pub fn validate_native_family(family: &SummaryFamilyType) -> Result<(), Error> {
     use planner_types::post_asap::SketchAlgorithm as A;
     if let SummaryFamilyType::Sketch(kind, grouping) = family {
         // Plain Count-Min is native as stored state only: it merges and reads
-        // its bare count, but the DAG does not build it from rows.
+        // its bare count, but the DAG does not build it from rows. HydraCms
+        // is built by the shared summary operator.
         if let (A::Cms, SketchParams::Cms { width, depth }) = (kind.algorithm(), kind.params()) {
-            return if valid_matrix(*width, *depth) && grouping == &Default::default() {
+            if grouping != &GroupingStrategy::PerSubpopulationInstance {
+                return hydra_cms_shape(family).map(|_| ()).map_err(Error::Invalid);
+            }
+            return if valid_matrix(*width, *depth) {
                 Ok(())
             } else {
                 Err(Error::Invalid(
@@ -212,6 +304,21 @@ pub fn validate_sketch_evaluation(
             (A::DDSketch, _) => bare_count,
             (A::Hll, SketchStatistic::Cardinality) => true,
             (A::Hll, _) => bare_count,
+            // HydraCms answers a group's item frequency. A group's total
+            // needs a group-mass query asap_sketchlib's Hydra lacks.
+            (A::Cms, SketchStatistic::PointCount { value, .. })
+                if matches!(family, SummaryFamilyType::Sketch(_, grouping)
+                    if grouping != &GroupingStrategy::PerSubpopulationInstance) =>
+            {
+                if value.is_none() {
+                    return Err(Error::Invalid(
+                        "HydraCms does not answer a group total: asap_sketchlib's Hydra has no \
+                         group-mass query; only an item's point count within a group is answered"
+                            .into(),
+                    ));
+                }
+                true
+            }
             // Only count intents read a Count-Min bare count, and their
             // updates have unit weight; the evaluation is typed Int64 on that basis.
             (A::Cms, _) => bare_count,
@@ -259,3 +366,18 @@ pub fn validate_exact_evaluation(
     }
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Only a query-time build applies a summary filter.
+    #[test]
+    fn filtered_summary_runs_only_at_query_time() {
+        let filter = planner_types::ir::Predicate::<usize>(planner_types::ir::ScalarExpr::Literal(
+            planner_types::pre_asap::ScalarValue::Boolean(true),
+        ));
+        assert!(validate_ingestion_time_summary(Some(&filter)).is_err());
+        assert!(validate_ingestion_time_summary::<usize>(None).is_ok());
+    }
+}

crates/asap-physical-operators/src/operators/aggregate/mod.rs

@@ -48,11 +48,56 @@ impl Operator {
             kind: Kind::Aggregate {
                 groups,
                 measures: measures.into_iter().map(|(_, r)| r).collect(),
+                filters: vec![],
             },
             inputs: vec![input],
             output: schema(fields),
         })
     }
+
+    /// Give each measure an optional row filter (SQL `FILTER (WHERE …)`).
+    /// A group is formed from all its rows, so a group with no matching row
+    /// still appears, with the measure's empty-input value: 0 for counts,
+    /// NULL for SQL SUM/AVG/MIN/MAX. No filter at all keeps the plain form.
+    pub fn with_measure_filters(mut self, filters: Vec<Option<Expression>>) -> Result<Self, Error> {
+        let Kind::Aggregate {
+            filters: slot,
+            measures,
+            groups,
+        } = &mut self.kind
+        else {
+            return Err(invalid("measure filters need an aggregate"));
+        };
+        if filters.iter().all(Option::is_none) {
+            return Ok(self);
+        }
+        if filters.len() != measures.len() {
+            return Err(invalid("one filter slot per aggregate measure required"));
+        }
+        // PromQL emits no per-measure filters (label matchers are scan
+        // predicates), and its outputs are non-nullable.
+        if self.inputs[0].has_promql_series_identity() {
+            return Err(invalid(
+                "measure filters are not supported over PromQL series identity",
+            ));
+        }
+        let mut output = (*self.output).clone();
+        for (index, filter) in filters.iter().enumerate() {
+            let Some(filter) = filter else { continue };
+            if filter.dtype(&self.inputs[0])?.0 != DataType::Bool {
+                return Err(invalid("measure filter must be boolean"));
+            }
+            if matches!(
+                measures[index],
+                Reduction::Sum(_) | Reduction::Avg(_) | Reduction::Min(_) | Reduction::Max(_)
+            ) {
+                output.fields[groups.len() + index].nullable = true;
+            }
+        }
+        *slot = filters;
+        self.output = Arc::new(output);
+        Ok(self)
+    }
     pub fn window(
         input: SchemaRef,
         intent: planner_types::pre_asap::AggIntent<ColumnRef>,
@@ -163,8 +208,20 @@ pub(super) fn execute<'a>(
                 )
                 .await?
             }
-            Kind::Aggregate { groups, measures } => {
-                reduce(rows, groups, measures, &operator.inputs[0], &context).await?
+            Kind::Aggregate {
+                groups,
+                measures,
+                filters,
+            } => {
+                reduce(
+                    rows,
+                    groups,
+                    measures,
+                    filters,
+                    &operator.inputs[0],
+                    &context,
+                )
+                .await?
             }
             _ => unreachable!(),
         };
@@ -178,6 +235,7 @@ async fn reduce(
     rows: Vec<Vec<Value>>,
     groups: &[usize],
     measures: &[Reduction],
+    filters: &[Option<Expression>],
     input: &SchemaRef,
     context: &RunContext,
 ) -> Result<Vec<Vec<Value>>, Error> {
@@ -204,8 +262,21 @@ async fn reduce(
             .iter()
             .map(|&i| rows[0][i].clone())
             .collect::<Vec<_>>();
-        for measure in measures {
-            result.push(reduce_one(&rows, measure, input, &mut work).await?);
+        for (index, measure) in measures.iter().enumerate() {
+            let mut selected = Vec::with_capacity(rows.len());
+            for row in &rows {
+                // SQL FILTER keeps a row only when the predicate is true, not NULL.
+                if filters
+                    .get(index)
+                    .and_then(Option::as_ref)
+                    .map_or(Ok(true), |filter| {
+                        filter.evaluate(row).map(|v| matches!(v, Value::Bool(true)))
+                    })?
+                {
+                    selected.push(row);
+                }
+            }
+            result.push(reduce_one(&selected, measure, input, &mut work).await?);
         }
         workspace.grow(row_bytes(&result))?;
         output.push(result);
@@ -239,7 +310,7 @@ pub(super) fn quantile(q: f64, mut values: Vec<f64>) -> f64 {
 }
 
 async fn reduce_one(
-    rows: &[Vec<Value>],
+    rows: &[&Vec<Value>],
     measure: &Reduction,
     input: &SchemaRef,
     work: &mut Cooperative,

crates/asap-physical-operators/src/operators/mod.rs

@@ -123,6 +123,9 @@ enum Kind {
     Aggregate {
         groups: Vec<usize>,
         measures: Vec<Reduction>,
+        /// Per-measure row filters (SQL `FILTER (WHERE …)`); empty when none.
+        #[serde(default, skip_serializing_if = "Vec::is_empty")]
+        filters: Vec<Option<Expression>>,
     },
     SemiJoin {
         keys: Vec<(usize, usize)>,
@@ -137,12 +140,30 @@ enum Kind {
         value: usize,
         time: Option<usize>,
         groups: Vec<usize>,
+        /// Rows for which this is not true update no state; their group is kept.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        filter: Option<Box<Expression>>,
     },
     KeyedSummaryBuild {
         family: SummaryFamilyType,
         value: usize,
         items: Vec<usize>,
         groups: Vec<usize>,
+        /// Rows for which this is not true update no state; their group is kept.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        filter: Option<Box<Expression>>,
+    },
+    /// One HydraCms state shared by all groups, adding each row's item weight.
+    SharedSummaryBuild {
+        family: SummaryFamilyType,
+        item: usize,
+        /// The weight column; each row adds 1 without one.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        weight: Option<usize>,
+        groups: Vec<usize>,
+        /// Rows for which this is not true update no state; their group is kept.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        filter: Option<Box<Expression>>,
     },
     KeyedEvaluation {
         state: usize,
@@ -287,6 +308,7 @@ impl PhysicalOperator<Batch, SchemaRef> for Operator {
                 | Kind::SemiJoin { .. }
                 | Kind::SummaryBuild { .. }
                 | Kind::KeyedSummaryBuild { .. }
+                | Kind::SharedSummaryBuild { .. }
                 | Kind::SummaryMerge { .. }
                 | Kind::VectorToScalar { .. }
         )
@@ -337,7 +359,9 @@ impl PhysicalOperator<Batch, SchemaRef> for Operator {
             Kind::Window { .. } => "WindowAggregate",
             Kind::SemiJoin { .. } => "SemiJoin",
             Kind::Join { .. } => "RelationalJoin",
-            Kind::SummaryBuild { .. } | Kind::KeyedSummaryBuild { .. } => "SummaryAgg",
+            Kind::SummaryBuild { .. }
+            | Kind::KeyedSummaryBuild { .. }
+            | Kind::SharedSummaryBuild { .. } => "SummaryAgg",
             Kind::KeyedEvaluation { .. } => "SummaryEstimate",
             Kind::SummaryMerge { .. } => "SummaryMerge",
             Kind::Evaluation { .. } => "SummaryEvaluation",
@@ -392,6 +416,7 @@ impl PhysicalOperator<Batch, SchemaRef> for Operator {
             Kind::SummaryBuild { .. }
             | Kind::Evaluation { .. }
             | Kind::KeyedSummaryBuild { .. }
+            | Kind::SharedSummaryBuild { .. }
             | Kind::KeyedEvaluation { .. } => summary::execute(self, inputs, context),
         }
     }