refactor(stage3): baseline selection MILP over x, y, z, t with (A)-(I), realizer and plan export with θ and coverage (#1207 PR 6)

crates/cost/src/state_bytes.rs

@@ -1,6 +1,7 @@
 //! The bytes one summary or exact state holds and its update operations
 //! per input row, from its parameters and the rows it folds.
 
+use crate::analytical::AnalyticalCostError;
 use asap_types::ir::operator::AggIntent;
 use asap_types::ir::schema::sketch::KLL_MAX_K;
 use asap_types::ir::schema::{
@@ -58,15 +59,21 @@ pub fn keeps_every_value(measure: &AggIntent) -> bool {
 /// Count-Min and Count-Sketch, plus one heap update for their top-k variants;
 /// for UnivMon, the same per layer it updates (`univmon_expected_layers`);
 /// 1 for every other sketch.
-pub fn summary_shape(family: &FieldDataType, input_rows: StateInputRows) -> (f64, u64) {
-    match family {
-        FieldDataType::Sketch(..) => {
-            let _ = input_rows;
-            todo!("#1207 PR 6: the state shape of θ from the candidate's LogicalProperties::parameters")
-        }
-        FieldDataType::ExactAggregate(..) => (1.0, EXACT_ACCUMULATOR_STATE_BYTES),
+///
+/// `parameters` is the build's θ (`PhysicalASAPDAGNode::parameters`). Err
+/// [`AnalyticalCostError::UnsupportedCandidate`] for a sketch without θ: a
+/// planner bug, since every sketch build carries its θ.
+pub fn summary_shape(
+    family: &FieldDataType,
+    parameters: Option<&SketchParams>,
+    input_rows: StateInputRows,
+) -> Result<(f64, u64), AnalyticalCostError> {
+    Ok(match (family, parameters) {
+        (FieldDataType::Sketch(..), Some(parameters)) => sketch_shape(parameters, input_rows),
+        (FieldDataType::Sketch(..), None) => return Err(AnalyticalCostError::UnsupportedCandidate),
+        (FieldDataType::ExactAggregate(..), _) => (1.0, EXACT_ACCUMULATOR_STATE_BYTES),
         _ => (1.0, 8),
-    }
+    })
 }
 
 /// What one summary state's bytes depend on besides its parameters (only
@@ -290,9 +297,10 @@ mod tests {
             assert_eq!(
                 summary_shape(
                     &FieldDataType::ExactAggregate(kind, params),
+                    None,
                     StateInputRows::UNKNOWN
                 ),
-                (1.0, 56)
+                Ok((1.0, 56))
             );
         }
         assert_eq!(asap_types::ir::schema::EXACT_ACCUMULATOR_STATE_BYTES, 56);
@@ -357,6 +365,147 @@ mod tests {
         assert_eq!(priced(Some(10_000), false), priced(None, false));
     }
 
+    /// The state shape of a sketch of `algorithm` with θ `params`, its
+    /// samples unknown.
+    fn shape(algorithm: SketchAlgorithm, params: SketchParams) -> (f64, u64) {
+        summary_shape(
+            &FieldDataType::Sketch(algorithm, GroupingStrategy::default()),
+            Some(&params),
+            StateInputRows::UNKNOWN,
+        )
+        .unwrap()
+    }
+
+    /// A sketch build without θ is a planner bug: it cannot be sized.
+    #[test]
+    fn a_sketch_without_parameters_is_not_sized() {
+        assert_eq!(
+            summary_shape(
+                &FieldDataType::Sketch(SketchAlgorithm::Kll, GroupingStrategy::default()),
+                None,
+                StateInputRows::UNKNOWN,
+            ),
+            Err(AnalyticalCostError::UnsupportedCandidate)
+        );
+    }
+
+    /// DDSketch state grows as its relative accuracy α tightens.
+    #[test]
+    fn ddsketch_state_grows_as_alpha_shrinks() {
+        let bytes = |alpha| shape(SketchAlgorithm::DDSketch, SketchParams::DDSketch { alpha }).1;
+        assert!(
+            bytes(0.1) < bytes(0.01),
+            "{} vs {}",
+            bytes(0.1),
+            bytes(0.01)
+        );
+        assert!(
+            bytes(0.01) < bytes(0.001),
+            "{} vs {}",
+            bytes(0.01),
+            bytes(0.001)
+        );
+    }
+
+    /// Every sketch family's state is sized from its parameters, as
+    /// `asap_sketchlib` allocates it: an HLL's registers, a KMV's `k`
+    /// hashes once it has seen `k` distinct items, and a UnivMon's layers
+    /// once keys of the assumed size have churned through every heap.
+    #[test]
+    fn summary_state_is_sized_from_sketch_params() {
+        use asap_sketchlib::DataInput;
+        assert_eq!(
+            shape(
+                SketchAlgorithm::Hll,
+                SketchParams::Hll {
+                    precision: 14,
+                    relative_error: None
+                }
+            )
+            .1,
+            asap_sketchlib::HyperLogLog::<asap_sketchlib::Classic>::new().memory_bytes() as u64
+        );
+        let mut kmv: asap_sketchlib::Kmv = asap_sketchlib::Kmv::new(1_000);
+        for i in 0..2_000u64 {
+            kmv.update(&DataInput::U64(i));
+        }
+        assert_eq!(
+            shape(
+                SketchAlgorithm::Kmv,
+                SketchParams::Kmv {
+                    k: 1_000,
+                    relative_error: None
+                }
+            )
+            .1,
+            kmv.memory_bytes() as u64
+        );
+        let mut univmon = asap_sketchlib::UnivMon::new(8, 3, 64, 2);
+        for i in 0..10_000u64 {
+            let mut key = [0; ASSUMED_VALUE_KEY_BYTES];
+            key[..8].copy_from_slice(&i.to_le_bytes());
+            univmon.update_weighted(&DataInput::Bytes(&key), 1);
+        }
+        assert_eq!(
+            shape(
+                SketchAlgorithm::UnivMon,
+                SketchParams::UnivMon {
+                    heap_size: 8,
+                    sketch_rows: 3,
+                    sketch_cols: 64,
+                    layers: 2,
+                    distinct_failure_probability: None,
+                }
+            )
+            .1,
+            univmon.memory_bytes() as u64
+        );
+    }
+
+    /// KLL state whose samples are unknown is charged its fixed 744 bytes
+    /// and the item array at the capacity `asap_sketchlib` grows it to:
+    /// `compute_max_capacity(k, 8)` 8-byte items, with `k` capped at 26 602.
+    #[test]
+    fn kll_state_is_sized_as_asap_sketchlib_allocates_it() {
+        let bytes = |k| {
+            shape(
+                SketchAlgorithm::Kll,
+                SketchParams::Kll {
+                    k,
+                    failure_probability: None,
+                },
+            )
+            .1
+        };
+        assert_eq!(bytes(269), 744 + 8 * 1_206);
+        assert_eq!(bytes(26_602), 744 + 8 * 80_118);
+        assert_eq!(bytes(30_000), 744 + 8 * 80_118);
+    }
+
+    /// A UnivMon row updates layers `0..=b`, where its bottom layer `b` is
+    /// geometric (`asap_sketchlib`'s `bottom_layer_for_hash`), so
+    /// `2 - 2^(1 - layers)` layers in expectation; each updates one counter
+    /// per sketch row and its heap.
+    #[test]
+    fn univmon_update_cost_is_expected_layers_times_rows_and_heap() {
+        let ops = |layers| {
+            shape(
+                SketchAlgorithm::UnivMon,
+                SketchParams::UnivMon {
+                    heap_size: 256,
+                    sketch_rows: 5,
+                    sketch_cols: 1_024,
+                    layers,
+                    distinct_failure_probability: None,
+                },
+            )
+            .0
+        };
+        assert_eq!(ops(1), 6.0);
+        assert_eq!(ops(2), 1.5 * 6.0);
+        assert_eq!(ops(16), (2.0 - 2f64.powi(-15)) * 6.0);
+    }
+
     /// A KLL state is charged what `asap_sketchlib` allocates for the
     /// samples it sees, not its full capacity: its fixed struct and level
     /// array, plus the item array the library starts empty and doubles from

crates/cost/src/usage.rs

@@ -458,15 +458,15 @@ fn each_node_usage<'a>(
     })
 }
 
-/// The state family of a node with `payload` whose first producer has
-/// `producer_family`: a build's own (one sketch instance per group; a
-/// shared Hydra structure is not that sketch), inherited by the merges,
+/// The state family of `node` whose first producer has `producer_family`:
+/// a build's own, its θ (`node.parameters`; one sketch instance per group;
+/// a shared Hydra structure is not that sketch), inherited by the merges,
 /// estimates and finalizations over it, none otherwise.
 fn state_family<'a>(
-    payload: &'a Payload,
+    node: &'a asap_types::ir::physical_export::PhysicalASAPDAGNode,
     producer_family: Option<StateFamily<'a>>,
 ) -> StateFamily<'a> {
-    match payload {
+    match &node.payload {
         Payload::ASAP(
             ASAPOp::SummaryAgg {
                 family, grouping, ..
@@ -478,7 +478,9 @@ fn state_family<'a>(
             FieldDataType::Sketch(..)
                 if *grouping == GroupingStrategy::PerSubpopulationInstance =>
             {
-                todo!("#1207 PR 6: the state family of θ from the candidate's LogicalProperties::parameters")
+                node.parameters
+                    .as_ref()
+                    .map_or(StateFamily::None, StateFamily::Sketch)
             }
             FieldDataType::ExactAggregate(kind, _) => StateFamily::Exact(kind),
             _ => StateFamily::None,
@@ -594,7 +596,7 @@ pub fn node_usage<'a>(
         missing: "an arrival rate (the arrival's rows_per_second, or the \
                   source's series_count with sample_interval)",
     };
-    let family = state_family(&node.payload, producers.first().map(|p| p.family));
+    let family = state_family(node, producers.first().map(|p| p.family));
     let unchanged_columns = unchanged_source_columns(
         node,
         &producers
@@ -916,7 +918,7 @@ pub fn node_usage<'a>(
                 rows_per_state: (!ingestion).then(|| input.rows.div_ceil(group_count.max(1))),
                 non_negative: raw_data_samples_non_negative(),
             };
-            let (depth, state_bytes) = summary_shape(family, input_rows);
+            let (depth, state_bytes) = summary_shape(family, node.parameters.as_ref(), input_rows)?;
             let state_bytes = measured_state_bytes.unwrap_or(state_bytes);
             let out = EdgeStatistics {
                 rows: group_count,
@@ -972,7 +974,7 @@ pub fn node_usage<'a>(
                 rows_per_state: Some(rows_per_tumbling_window.div_ceil(group_count.max(1))),
                 non_negative: raw_data_samples_non_negative(),
             };
-            let (depth, state_bytes) = summary_shape(family, input_rows);
+            let (depth, state_bytes) = summary_shape(family, node.parameters.as_ref(), input_rows)?;
             let state_bytes = measured_state_bytes.unwrap_or(state_bytes);
             let out = EdgeStatistics {
                 rows: states,

crates/devtools/src/bin/asap_planner.rs

@@ -63,7 +63,7 @@
 //     `P1`, the selected plan as Stage 2 built it;
 //   - stage3_plan_selection (from `stage3`): Stage 3's `Selection`:
 //     `selected_dag` (`candidate_of_node`, `edges`,
-//     `node_of_output_reference`), `nodes_with_proven_non_negative_weights`,
+//     `node_of_output_reference`),
 //     the plan's `cost` (its `usage`, `cost_per_second` and per-node
 //     charges), `accepted_accuracy_per_query`, `pruned_candidates` and
 //     `search_report`; `error` when no plan was selected.
@@ -734,8 +734,6 @@ fn selection_json(selection: &Result<Selection, SelectionError>) -> Value {
                 .collect::<Vec<_>>(),
             "node_of_output_reference": nodes(&dag.node_of_output_reference),
         },
-        "nodes_with_proven_non_negative_weights":
-            nodes(&selection.nodes_with_proven_non_negative_weights),
         "cost": cost_json(&selection.cost),
         "accepted_accuracy_per_query": accepted_json(&selection.accepted_accuracy_per_query),
         "pruned_candidates": selection

crates/devtools/src/planner_fuzz/mod.rs

@@ -34,7 +34,7 @@ pub mod relations;
 pub mod report;
 pub mod shrink;
 
-use std::collections::HashMap;
+use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::rc::Rc;
 
 use serde::{Deserialize, Serialize};
@@ -451,25 +451,86 @@ impl Planner {
         })
     }
 
-    /// Stage 2 and Stage 3 over the logical plan `roots` of `workload`'s
-    /// queries (one root per query, in order): every materialization
-    /// assignment implemented and priced, the cheapest selected, with the
-    /// same models as [`plan`](Self::plan). Guarantees are composed again
-    /// first, so a derived node carries the guarantee of what it reads, and
-    /// summaries are sized for their consumers as Pass 1 sizes them.
-    /// With `last_root_is_key_filter`, the last root is a `Filter` on
-    /// group-key columns only and keeps its input's guarantee (whole groups
-    /// selected by an exact key keep each remaining row's bound), which
-    /// Stage 1's composition, giving no guarantee to any filter over
-    /// approximate rows, does not know.
+    /// Stages 0–3 over the witness plan `roots` of `workload`'s queries
+    /// (one root per query, in order), with the same models as
+    /// [`plan`](Self::plan): the witness's own operators are hash-consed
+    /// into the memo and Stage 3 selects the cheapest plan of them, each
+    /// summary at the θ its consumers' parameter floors allow. The cost is
+    /// that plan's, and `derived_root` the nodes of the realized plan only
+    /// the last root reads. `last_root_is_key_filter` is accepted for the
+    /// relations' interface; accuracy is not composed again here (constraint
+    /// (K), #1207 PR 8).
     pub fn price(
         &self,
         workload: &FuzzWorkload,
         roots: &[QueryRoot],
         last_root_is_key_filter: bool,
     ) -> Result<Priced, String> {
-        let _ = (workload, roots, last_root_is_key_filter);
-        todo!("#1207 PR 6: price a witness plan through the new Stage 2 and Stage 3")
+        let _ = last_root_is_key_filter;
+        let planning = workload.planning_workload();
+        let entries: Vec<_> = planning.query_workload.entries().collect();
+        let cost_inputs = self.cost_inputs();
+        let deployment = workload.deployment();
+        let models = PlanningModels::builtin()
+            .with_cost_inputs(&cost_inputs)
+            .with_deployment_capabilities(&deployment)
+            .with_ignore_accuracy_guarantees(!self.settings.enforce_accuracy_guarantees);
+        let run = run_search_over(
+            roots.to_vec(),
+            &entries,
+            &workload.data,
+            models,
+            PlanSearchSettings::default(),
+        )
+        .map_err(|e| format!("witness: {e}"))?;
+        let selection = run
+            .run
+            .selection
+            .as_ref()
+            .map_err(|e| format!("witness selection: {e}"))?;
+        let cost = &selection.cost;
+        let cost_per_second = cost
+            .cost_per_second
+            .ok_or("no cost per second (cost model 3)")?;
+        let dag = &run
+            .run
+            .selected_physical_plan
+            .as_ref()
+            .ok_or("a selection has a realized plan")?
+            .plan
+            .dag;
+        let mut inputs: BTreeMap<_, Vec<_>> = BTreeMap::new();
+        for edge in &dag.edges {
+            inputs.entry(edge.consumer).or_default().push(edge.producer);
+        }
+        let reached = |root| {
+            let mut seen = BTreeSet::new();
+            let mut stack = vec![root];
+            while let Some(node) = stack.pop() {
+                if seen.insert(node) {
+                    stack.extend(inputs.get(&node).into_iter().flatten().copied());
+                }
+            }
+            seen
+        };
+        let reached_by_root: Vec<BTreeSet<_>> = dag.roots.iter().map(|r| reached(*r)).collect();
+        let (last, others) = reached_by_root.split_last().ok_or("no roots")?;
+        let private: Vec<_> = last
+            .iter()
+            .filter(|node| others.iter().all(|set| !set.contains(node)))
+            .filter_map(|node| cost.per_node.get(node))
+            .collect();
+        Ok(Priced {
+            selected: "selected".into(),
+            cost_per_second,
+            derived_root: PrivateNodes {
+                cost_per_second: private.iter().filter_map(|n| n.cost_per_second).sum(),
+                operators: private
+                    .iter()
+                    .map(|n| format!("{:?}", n.operator))
+                    .collect(),
+            },
+        })
     }
 
     /// [`price`](Self::price) of `workload`'s selected logical plan