fix: read sketch capabilities and state bytes from asap_sketchlib (#953, #946, #947)

Cargo.toml

@@ -23,8 +23,10 @@ resolver = "2"
 arrow = { version = "58.4", default-features = false }
 # The hasher Arrow already uses, for the executor's hot interning maps.
 ahash = { version = "0.8", default-features = false, features = ["std", "runtime-rng"] }
-# `planner-pin/0.4.0-plus`: release 0.4.0 (eb04e7f) plus the open sketchlib
-# PRs #245 (checked f64 updates, heap rescoring), #246 (KLL lazy-buffer
-# follow-ups), #247 (Count Sketch stack medians) and #248
-# (`DdSketch::quantile_interpolated`). Moves to a release once they merge.
-asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "63e565506426020c505ef0b84e0c91f450d09e7c" }
+# main after release 0.4.0 (eb04e7f): #245 (checked f64 updates, heap
+# rescoring), #246 (KLL lazy-buffer follow-ups), #247 (Count Sketch stack
+# medians), #248 (`DdSketch::quantile_interpolated`), #249 (`memory_bytes` and
+# `estimated_memory_bytes`; KMV without `experimental`), #251 (the public
+# capability table) and #254 (Planner names in that table). Moves to 0.4.1
+# once that is released.
+asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "21b63578c2e612e153f5aaa4d3848f1035a13c60" }

crates/executor/Cargo.toml

@@ -9,7 +9,7 @@ ahash.workspace = true
 chrono = { version = "0.4.39", default-features = false, features = ["std"] }
 futures = "0.3"
 planner-types = { package = "asap-types", path = "../types" }
-asap_sketchlib = { workspace = true, features = ["experimental"] }
+asap_sketchlib = { workspace = true }
 serde = { version = "1", features = ["derive", "rc"] }
 serde_json = "1"
 thiserror = "2"

crates/executor/src/capability.rs

@@ -325,37 +325,41 @@ pub fn validate_native_family(family: &SummaryFamilyType) -> Result<(), Error> {
     .map_err(Error::Invalid)
 }
 
-/// The families the executor builds per instance and reads, in
-/// `SketchAlgorithm` order. A family it builds but reads nothing from
-/// (plain Count Sketch) and one it does not build (Theta) are absent.
-pub(crate) const EVALUATED_FAMILIES: [SketchAlgorithm; 8] = [
-    SketchAlgorithm::UnivMon,
-    SketchAlgorithm::Kll,
-    SketchAlgorithm::Cms,
-    SketchAlgorithm::Hll,
-    SketchAlgorithm::DDSketch,
-    SketchAlgorithm::CmsWithHeap,
-    SketchAlgorithm::Kmv,
-    SketchAlgorithm::CountSketchWithHeap,
+/// What the executor evaluates from a per-instance state of each family it
+/// builds and reads: its supported-summary declaration. Every readout here is
+/// in the family's row of the family-capability table; a readout the row
+/// lists and this omits is not evaluated yet (Count-Min's and the heap
+/// families' item frequency, which the executor answers only from a Hydra
+/// grid; KLL's rank; UnivMon's G-sum). Plain Count Sketch is built but read
+/// from nothing.
+pub(crate) const EVALUATED_READOUTS: &[(SketchAlgorithm, &[StatisticKind])] = &[
+    (SketchAlgorithm::Cms, &[StatisticKind::Total]),
+    (SketchAlgorithm::CmsWithHeap, &[StatisticKind::TopK]),
+    (SketchAlgorithm::CountSketchWithHeap, &[StatisticKind::TopK]),
+    (SketchAlgorithm::Hll, &[StatisticKind::Distinct]),
+    (SketchAlgorithm::Kmv, &[StatisticKind::Distinct]),
+    (SketchAlgorithm::Kll, &[StatisticKind::Quantile]),
+    (SketchAlgorithm::DDSketch, &[StatisticKind::Quantile]),
+    (
+        SketchAlgorithm::UnivMon,
+        &[
+            StatisticKind::Total,
+            StatisticKind::FrequencyL2,
+            StatisticKind::Distinct,
+            StatisticKind::FrequencyEntropy,
+        ],
+    ),
 ];
 
-/// The readouts the family-capability table lists for an evaluated family
-/// that the executor does not evaluate yet: a per-instance Count-Min's item
-/// frequency (the executor answers an item's frequency only from a Hydra
-/// grid, which Stage 2 does not produce yet; `SupportedSummary` gains a
-/// grouping dimension when it does).
-pub(crate) const UNEVALUATED_READOUTS: [(SketchAlgorithm, StatisticKind); 1] =
-    [(SketchAlgorithm::Cms, StatisticKind::ItemFrequency)];
-
 /// Whether the executor reads `readout` from a state of `family` whose
 /// family-capability table row answers it.
 fn evaluates(family: &SummaryFamilyType, readout: StatisticKind) -> bool {
     use planner_types::ir::schema::HydraKind;
     match family {
         SummaryFamilyType::Sketch(kind, GroupingStrategy::PerSubpopulationInstance) => {
-            let algorithm = kind.algorithm();
-            EVALUATED_FAMILIES.contains(algorithm)
-                && !UNEVALUATED_READOUTS.contains(&(algorithm.clone(), readout))
+            EVALUATED_READOUTS
+                .iter()
+                .any(|(family, readouts)| family == kind.algorithm() && readouts.contains(&readout))
         }
         SummaryFamilyType::Sketch(_, GroupingStrategy::SharedMultiSubpopulation { kind, .. }) => {
             matches!(kind, HydraKind::HydraCms | HydraKind::HydraKll)

crates/executor/src/summary_kernels/datasketches_kll.rs

@@ -54,11 +54,13 @@ impl AggregateCore for DatasketchesKLLAccumulator {
             .as_any()
             .downcast_ref::<Self>()
             .ok_or("KLL merges only with KLL")?;
-        if self.inner.k() != other.inner.k() {
+        // sketchlib's merge panics on a different k or m.
+        let (k, m) = (self.inner.k(), self.inner.m());
+        if (k, m) != (other.inner.k(), other.inner.m()) {
             return Err(format!(
-                "KLL merges only with the same k ({} vs {})",
-                self.inner.k(),
-                other.inner.k()
+                "KLL merges only with the same k and m ({k}, {m} vs {}, {})",
+                other.inner.k(),
+                other.inner.m()
             )
             .into());
         }
@@ -136,13 +138,18 @@ mod tests {
         assert!(kll.approx_memory_bytes() > empty);
     }
 
-    // States of different k do not merge.
+    // States of different k or m do not merge: an error, not sketchlib's panic.
     #[test]
-    fn merge_rejects_a_different_k() {
+    fn merge_rejects_a_different_k_or_m() {
         let kll = DatasketchesKLLAccumulator::new(200);
         assert!(kll
             .merge_with(&DatasketchesKLLAccumulator::new(100))
             .is_err());
+        assert!(kll
+            .merge_with(&DatasketchesKLLAccumulator {
+                inner: Kll::new(200, 4),
+            })
+            .is_err());
         assert!(kll
             .merge_with(&DatasketchesKLLAccumulator::new(200))
             .is_ok());

crates/executor/src/summary_kernels/dd_sketch.rs

@@ -42,6 +42,12 @@ impl AggregateCore for DDSketchAccumulator {
         self.inner.get_count() == 0
     }
 
+    /// sketchlib's count of what the sketch holds now (its struct
+    /// included); the bucket stores grow with the values' range.
+    fn approx_memory_bytes(&self) -> usize {
+        self.inner.memory_bytes()
+    }
+
     /// Quantiles.
     fn estimate(&self, query: &SketchStatistic) -> Result<f64, KernelError> {
         match query {

crates/executor/src/summary_kernels/factory.rs

@@ -263,11 +263,7 @@ impl AccumulatorUpdater for DDSketchAccumulatorUpdater {
     }
 
     fn memory_usage_bytes(&self) -> usize {
-        // The struct and the 8-byte counts both bucket stores allocate,
-        // which grow with the values' range.
-        let sketch = &self.acc.inner;
-        let bins = sketch.positive_store_counts().len() + sketch.negative_store_counts().len();
-        std::mem::size_of::<DDSketchAccumulator>() + bins * std::mem::size_of::<u64>()
+        self.acc.approx_memory_bytes()
     }
 }
 
@@ -716,10 +712,8 @@ mod planner_parameter_regression {
                 },
             ),
         ] {
-            let family = SummaryFamilyType::Sketch(
-                SketchKind::new(algorithm.clone(), params),
-                Default::default(),
-            );
+            let family =
+                SummaryFamilyType::Sketch(SketchKind::new(algorithm, params), Default::default());
             let update = SummaryUpdate {
                 item: Some(SummaryInputExpr::Column(
                     planner_types::ir::scalar::ColumnRef::Named("host".into()),
@@ -872,9 +866,9 @@ mod ddsketch_domain {
 mod tests {
     use super::*;
 
-    /// A DDSketch updater reports the bins its stores allocate: none
-    /// before a sample, one 128-bin chunk of the positive store after a
-    /// few positive samples.
+    /// A DDSketch updater reports sketchlib's count of the bins its stores
+    /// allocate: none before a sample, one 128-bin chunk of the positive
+    /// store after a few positive samples.
     #[test]
     fn ddsketch_memory_is_its_allocated_bins() {
         let mut updater = DDSketchAccumulatorUpdater::new(0.05);
@@ -884,6 +878,10 @@ mod tests {
             updater.update_single(value, 0).unwrap();
         }
         assert_eq!(updater.memory_usage_bytes(), empty + 128 * 8);
+        assert_eq!(
+            updater.memory_usage_bytes(),
+            updater.acc.inner.memory_bytes()
+        );
     }
 
     /// A numeric sketch takes an Int64 as its exact Float64, and refuses

crates/executor/src/summary_kernels/univmon.rs

@@ -160,29 +160,30 @@ impl UnivMonAccumulator {
 }
 
 impl AggregateCore for UnivMonAccumulator {
-    /// Counters and heap slots of every layer, plus an upper bound on the
-    /// variable-length keys the heaps retain: the smaller of every slot
-    /// holding the longest key seen and every layer holding a copy of every
-    /// key inserted. Constant time, since a build charges memory after
-    /// every row; the library exposes no eviction hook for an exact running
-    /// total, so a stream of long keys is still overcounted.
+    /// This accumulator's own fields, sketchlib's bytes for the sketch's
+    /// shape with keys that own no bytes ([`UnivMon::estimated_memory_bytes`]:
+    /// at least what it holds), and an upper bound on the variable-length
+    /// keys the heaps retain: the smaller of every slot holding the longest
+    /// key seen and every layer holding a copy of every key inserted.
+    /// Constant time, since a build charges memory after every row and
+    /// [`UnivMon::memory_bytes`] walks the heaps; the library exposes no
+    /// eviction hook for an exact running total, so a stream of long keys is
+    /// still overcounted.
     fn approx_memory_bytes(&self) -> usize {
-        let layers = self.inner.layers;
+        let (heap_size, rows, cols, layers) = self.dimensions();
         let key_bytes = layers
-            .saturating_mul(self.inner.heap_size)
+            .saturating_mul(heap_size)
             .saturating_mul(self.max_key_bytes)
             .min(layers.saturating_mul(self.inserted_key_bytes));
-        key_bytes
-            .saturating_add(std::mem::size_of::<Self>())
-            .saturating_add(
-                self.inner.layers.saturating_mul(
-                    self.inner
-                        .rows
-                        .saturating_mul(self.inner.cols)
-                        .saturating_mul(16)
-                        .saturating_add(self.inner.heap_size.saturating_mul(256)),
-                ),
-            )
+        (std::mem::size_of::<Self>() - std::mem::size_of::<UnivMon>())
+            .saturating_add(UnivMon::estimated_memory_bytes(
+                heap_size,
+                rows,
+                cols,
+                layers,
+                &DataInput::U64(0),
+            ))
+            .saturating_add(key_bytes)
     }
     fn clone_boxed_core(&self) -> Box<dyn AggregateCore> {
         Box::new(self.clone())

crates/executor/src/supported_summaries.rs

@@ -3,24 +3,18 @@
 //! `DeploymentCapabilities::supported_summaries`, so the Planner offers no
 //! summary the executor would reject when it binds the plan.
 
-use crate::capability::{EVALUATED_FAMILIES, UNEVALUATED_READOUTS};
-use planner_types::ir::schema::{FamilyCapability, SupportedSummary};
+use crate::capability::EVALUATED_READOUTS;
+use planner_types::ir::schema::SupportedSummary;
 
 /// Every sketch family this executor builds from rows and the statistics it
-/// answers from one state of it: the family-capability table's row of each
-/// family it evaluates, less the readouts it does not evaluate yet. Exact
-/// accumulators are always supported and are not listed.
+/// answers from one state of it (`EVALUATED_READOUTS`). Exact accumulators
+/// are always supported and are not listed.
 pub fn supported_summary_families() -> Vec<SupportedSummary> {
-    EVALUATED_FAMILIES
-        .into_iter()
-        .map(|family| SupportedSummary {
-            readouts: FamilyCapability::of_per_instance_sketch(&family)
-                .readouts
-                .iter()
-                .map(|&(readout, _)| readout)
-                .filter(|readout| !UNEVALUATED_READOUTS.contains(&(family.clone(), *readout)))
-                .collect(),
+    EVALUATED_READOUTS
+        .iter()
+        .map(|&(family, readouts)| SupportedSummary {
             family,
+            readouts: readouts.to_vec(),
         })
         .collect()
 }
@@ -34,20 +28,20 @@ mod tests {
     use planner_types::ir::operator::operator_properties::{Reduction, Source};
     use planner_types::ir::scalar::ColumnRef;
     use planner_types::ir::schema::{
-        DataType, Field, FieldDataType, GroupingStrategy, Schema, SketchAlgorithm, SketchKind,
-        SketchParams, SketchStatistic, StatisticKind, SummaryInputExpr, SummaryUpdate,
+        DataType, FamilyCapability, Field, FieldDataType, GroupingStrategy, Schema,
+        SketchAlgorithm, SketchKind, SketchParams, SketchStatistic, StatisticKind,
+        SummaryInputExpr, SummaryUpdate,
     };
     use planner_types::ir::{ASAPOp, NonASAPOp, Operator, OperatorNode};
 
-    const ALGORITHMS: [SketchAlgorithm; 10] = [
+    const ALGORITHMS: [SketchAlgorithm; 9] = [
         SketchAlgorithm::UnivMon,
         SketchAlgorithm::Kll,
         SketchAlgorithm::Cms,
         SketchAlgorithm::Hll,
         SketchAlgorithm::DDSketch,
         SketchAlgorithm::CmsWithHeap,
         SketchAlgorithm::Kmv,
-        SketchAlgorithm::Theta,
         SketchAlgorithm::CountSketch,
         SketchAlgorithm::CountSketchWithHeap,
     ];
@@ -95,7 +89,6 @@ mod tests {
                 k: 1024,
                 relative_error: None,
             },
-            A::Theta => SketchParams::Theta { k: 1024 },
             A::CountSketch => SketchParams::CountSketch {
                 width: 64,
                 depth: 3,
@@ -105,12 +98,13 @@ mod tests {
                 depth: 3,
                 heap_size: 10,
             },
+            other => unreachable!("{other:?} is not probed"),
         }
     }
 
     fn per_instance(algorithm: &SketchAlgorithm) -> FieldDataType {
         FieldDataType::Sketch(
-            SketchKind::new(algorithm.clone(), params(algorithm)),
+            SketchKind::new(*algorithm, params(algorithm)),
             GroupingStrategy::PerSubpopulationInstance,
         )
     }
@@ -133,6 +127,7 @@ mod tests {
             StatisticKind::FrequencyL2 => SketchStatistic::FrequencyL2,
             StatisticKind::FrequencyEntropy => SketchStatistic::FrequencyEntropy,
             StatisticKind::TopK => SketchStatistic::TopK { k: 5 },
+            other => unreachable!("{other:?} is not probed"),
         }
     }
 
@@ -230,9 +225,8 @@ mod tests {
     /// The list holds exactly the (family, readout) pairs the factory builds
     /// per instance and an evaluation answers from that state, and the
     /// IR's `SummaryEstimate` validation accepts exactly what the
-    /// family-capability table answers. A listed family's readouts are its
-    /// table row less the ones named not evaluated, each of which the table
-    /// answers and the executor does not.
+    /// family-capability table answers. Every declared readout is in the
+    /// family's table row.
     #[test]
     fn supported_summary_families_match_the_factory_and_evaluations() {
         let listed = supported_summary_families();
@@ -252,11 +246,10 @@ mod tests {
                     answered,
                     "IR: {algorithm:?} {readout:?}"
                 );
-                if UNEVALUATED_READOUTS.contains(&(algorithm.clone(), readout)) {
-                    assert!(answered && !executed, "{algorithm:?} {readout:?}");
-                } else if EVALUATED_FAMILIES.contains(algorithm) {
-                    assert_eq!(answered, executed, "table: {algorithm:?} {readout:?}");
-                }
+                assert!(
+                    !declared || answered,
+                    "declared but not in the table: {algorithm:?} {readout:?}"
+                );
             }
         }
         assert!(listed.iter().all(|s| !s.readouts.is_empty()));

crates/executor/tests/deployment.rs

@@ -311,7 +311,7 @@ fn factory_and_dag_heap_states_are_interchangeable() {
             FrequencyAlgorithm::CountSketch,
         ),
     ] {
-        let family = heap_sketch(algorithm.clone(), 1_024, 5, 4);
+        let family = heap_sketch(algorithm, 1_024, 5, 4);
         let factory = build_weighted_frequency(&family, [(0, 1.0), (1, 2.0)]);
         let mut dag = WeightedFrequency::new(frequency, 1_024, 5, 4).unwrap();
         dag.update(&[Value::Utf8("host-0".into())], 3.0).unwrap();