feat(optimizer): select atomic costs by observed data shape

asap-planner-rs/src/bin/candidate_gen_dump.rs

@@ -6,8 +6,9 @@ use std::path::PathBuf;
 
 use asap_planner::{
     optimizer::{
-        enumerate_candidates, extract_aqes, load_optional_selected_atomic_cost_table,
-        resolve_atomic_costs, AtomicCostTable, AtomicCosts, CandidateConfig, RQE,
+        enumerate_candidates, extract_aqes, load_nearest_atomic_cost_table,
+        load_optional_selected_atomic_cost_table, resolve_atomic_costs, AtomicCostTable,
+        AtomicCosts, CandidateConfig, DataShape, ShapeMatchPolicy, RQE,
     },
     ControllerConfig,
 };
@@ -38,16 +39,43 @@ struct Args {
     /// JSON `profiles[].workload` value selecting exactly one measured profile.
     #[arg(long = "atomic-cost-workload", requires = "atomic_costs")]
     atomic_cost_workload: Option<PathBuf>,
+    #[arg(
+        long = "atomic-cost-observed-shape",
+        requires = "atomic_costs",
+        conflicts_with = "atomic_cost_workload"
+    )]
+    atomic_cost_observed_shape: Option<PathBuf>,
+    #[arg(long, default_value_t = 100_000)]
+    minimum_benchmark_events: u64,
+    #[arg(long, default_value_t = 1.0)]
+    max_log2_cardinality_distance: f64,
+    #[arg(long, default_value_t = 0.2)]
+    max_zipf_distance: f64,
 }
 
 fn main() -> anyhow::Result<()> {
     let args = Args::parse();
 
-    let atomic_cost_table = load_optional_selected_atomic_cost_table(
-        args.atomic_costs.as_deref(),
-        args.atomic_cost_workload.as_deref(),
-    )?
-    .unwrap_or_default();
+    let atomic_cost_table = if let Some(shape_path) = args.atomic_cost_observed_shape.as_deref() {
+        let observed: DataShape = serde_json::from_str(&std::fs::read_to_string(shape_path)?)?;
+        load_nearest_atomic_cost_table(
+            args.atomic_costs
+                .as_deref()
+                .expect("clap requires --atomic-costs"),
+            observed,
+            ShapeMatchPolicy {
+                minimum_benchmark_events: args.minimum_benchmark_events,
+                max_log2_cardinality_distance: args.max_log2_cardinality_distance,
+                max_zipf_distance: args.max_zipf_distance,
+            },
+        )?
+    } else {
+        load_optional_selected_atomic_cost_table(
+            args.atomic_costs.as_deref(),
+            args.atomic_cost_workload.as_deref(),
+        )?
+        .unwrap_or_default()
+    };
 
     let yaml_str = std::fs::read_to_string(&args.input_config)?;
     let config: ControllerConfig = serde_yaml::from_str(&yaml_str)?;

asap-planner-rs/src/bin/optimizer_cli.rs

@@ -7,7 +7,8 @@
 use std::path::PathBuf;
 
 use asap_planner::optimizer::{
-    load_optional_selected_atomic_cost_table, run_greedy_pipeline, AtomicCostTable, SeriesDataset,
+    load_nearest_atomic_cost_table, load_optional_selected_atomic_cost_table, run_greedy_pipeline,
+    AtomicCostTable, DataShape, SeriesDataset, ShapeMatchPolicy,
 };
 use asap_planner::ControllerConfig;
 use clap::Parser;
@@ -50,6 +51,21 @@ struct Args {
     #[arg(long = "atomic-cost-workload", requires = "atomic_costs")]
     atomic_cost_workload: Option<PathBuf>,
 
+    /// JSON DataShape observed by ASAPQuery-backend. Selects the nearest safe
+    /// benchmark profile instead of requiring descriptor equality.
+    #[arg(
+        long = "atomic-cost-observed-shape",
+        requires = "atomic_costs",
+        conflicts_with = "atomic_cost_workload"
+    )]
+    atomic_cost_observed_shape: Option<PathBuf>,
+    #[arg(long, default_value_t = 100_000)]
+    minimum_benchmark_events: u64,
+    #[arg(long, default_value_t = 1.0)]
+    max_log2_cardinality_distance: f64,
+    #[arg(long, default_value_t = 0.2)]
+    max_zipf_distance: f64,
+
     #[arg(short, long, action = clap::ArgAction::Count)]
     verbose: u8,
 }
@@ -77,10 +93,26 @@ fn main() -> anyhow::Result<()> {
     let config: ControllerConfig = serde_yaml::from_str(&yaml_str)?;
     let dataset = SeriesDataset::from_path(&args.dataset)?;
 
-    let atomic_cost_table = match load_optional_selected_atomic_cost_table(
-        args.atomic_costs.as_deref(),
-        args.atomic_cost_workload.as_deref(),
-    )? {
+    let selected = if let Some(shape_path) = args.atomic_cost_observed_shape.as_deref() {
+        let observed: DataShape = serde_json::from_str(&std::fs::read_to_string(shape_path)?)?;
+        Some(load_nearest_atomic_cost_table(
+            args.atomic_costs
+                .as_deref()
+                .expect("clap requires --atomic-costs"),
+            observed,
+            ShapeMatchPolicy {
+                minimum_benchmark_events: args.minimum_benchmark_events,
+                max_log2_cardinality_distance: args.max_log2_cardinality_distance,
+                max_zipf_distance: args.max_zipf_distance,
+            },
+        )?)
+    } else {
+        load_optional_selected_atomic_cost_table(
+            args.atomic_costs.as_deref(),
+            args.atomic_cost_workload.as_deref(),
+        )?
+    };
+    let atomic_cost_table = match selected {
         Some(table) => table,
         None => {
             tracing::warn!(

asap-planner-rs/src/optimizer/atomic_costs.rs

@@ -39,9 +39,26 @@ pub struct AtomicCostDocument {
 #[serde(deny_unknown_fields)]
 pub struct AtomicCostProfile {
     pub workload: WorkloadDescription,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub shape: Option<DataShape>,
     pub entries: Vec<AtomicCostEntry>,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct DataShape {
+    pub cardinality: u64,
+    pub zipf_exponent: Option<f64>,
+    pub benchmark_events: u64,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct ShapeMatchPolicy {
+    pub minimum_benchmark_events: u64,
+    pub max_log2_cardinality_distance: f64,
+    pub max_zipf_distance: f64,
+}
+
 /// The provenance of the input data on which atomic costs were measured.
 /// Synthetic descriptions remain opaque because their generator schema evolves
 /// independently; external profiles are represented explicitly for selection.
@@ -145,6 +162,72 @@ pub fn load_atomic_cost_table(
     }
 }
 
+/// Load the nearest compatible benchmark profile. Distribution families never
+/// interpolate; benchmark event count is a sufficiency gate rather than a
+/// distance axis once enough data has been measured.
+pub fn load_nearest_atomic_cost_table(
+    path: &Path,
+    observed: DataShape,
+    policy: ShapeMatchPolicy,
+) -> anyhow::Result<AtomicCostTable> {
+    if observed.cardinality == 0
+        || policy.minimum_benchmark_events == 0
+        || !policy.max_log2_cardinality_distance.is_finite()
+        || policy.max_log2_cardinality_distance <= 0.0
+        || !policy.max_zipf_distance.is_finite()
+        || policy.max_zipf_distance <= 0.0
+    {
+        anyhow::bail!("invalid shape matching request");
+    }
+    let raw = std::fs::read_to_string(path).map_err(|error| {
+        anyhow::anyhow!("reading atomic-cost table {}: {error}", path.display())
+    })?;
+    let document: AtomicCostDocument = serde_json::from_str(&raw).map_err(|error| {
+        anyhow::anyhow!("parsing atomic-cost document {}: {error}", path.display())
+    })?;
+    if document.schema_version != ATOMIC_COST_SCHEMA_VERSION {
+        anyhow::bail!(
+            "unsupported atomic-cost schema_version {}",
+            document.schema_version
+        );
+    }
+    document
+        .profiles
+        .iter()
+        .filter_map(|profile| Some((profile, shape_distance(observed, profile.shape?, policy)?)))
+        .filter(|(_, distance)| *distance <= 1.0)
+        .min_by(|(_, left), (_, right)| left.total_cmp(right))
+        .map(|(profile, _)| profile.entries.clone())
+        .ok_or_else(|| anyhow::anyhow!("no benchmark profile is within the observed-shape bounds"))
+}
+
+fn shape_distance(
+    observed: DataShape,
+    candidate: DataShape,
+    policy: ShapeMatchPolicy,
+) -> Option<f64> {
+    if candidate.cardinality == 0
+        || candidate.benchmark_events < policy.minimum_benchmark_events
+        || candidate.zipf_exponent.is_some() != observed.zipf_exponent.is_some()
+    {
+        return None;
+    }
+    let cardinality =
+        ((candidate.cardinality as f64).log2() - (observed.cardinality as f64).log2()).abs()
+            / policy.max_log2_cardinality_distance;
+    let zipf = match (candidate.zipf_exponent, observed.zipf_exponent) {
+        (None, None) => 0.0,
+        (Some(candidate), Some(observed)) => {
+            if !candidate.is_finite() || !observed.is_finite() {
+                return None;
+            }
+            (candidate - observed).abs() / policy.max_zipf_distance
+        }
+        _ => return None,
+    };
+    Some(cardinality.max(zipf))
+}
+
 /// Load the selector artifact and return the corresponding empirical table.
 /// Offline callers use this single interface so selector validation cannot
 /// drift between planner tools.
@@ -754,4 +837,59 @@ mod tests {
                 .is_some()
         );
     }
+
+    #[test]
+    fn nearest_profile_uses_shape_not_cheapest_distant_scenario() {
+        let document = serde_json::json!({
+            "schema_version": ATOMIC_COST_SCHEMA_VERSION,
+            "profiles": [
+                {"workload":{"synthetic":{"description":{}}},
+                 "shape":{"cardinality":1000,"zipf_exponent":1.2,"benchmark_events":100000},
+                 "entries":[cms_entry(3, 512)]},
+                {"workload":{"synthetic":{"description":{}}},
+                 "shape":{"cardinality":1000000,"zipf_exponent":1.2,"benchmark_events":100000},
+                 "entries":[cms_entry(3, 256)]}
+            ]
+        });
+        let path = std::env::temp_dir().join(format!("atomic-shape-{}.json", std::process::id()));
+        std::fs::write(&path, serde_json::to_vec(&document).unwrap()).unwrap();
+        let selected = load_nearest_atomic_cost_table(
+            &path,
+            DataShape {
+                cardinality: 1200,
+                zipf_exponent: Some(1.1),
+                benchmark_events: 10_000,
+            },
+            ShapeMatchPolicy {
+                minimum_benchmark_events: 10_000,
+                max_log2_cardinality_distance: 4.0,
+                max_zipf_distance: 0.5,
+            },
+        )
+        .unwrap();
+        std::fs::remove_file(path).unwrap();
+        assert_eq!(selected[0].sketch_config["params"]["cols"], 512);
+    }
+
+    #[test]
+    fn nearest_profile_rejects_uniform_zipf_mismatch() {
+        assert!(shape_distance(
+            DataShape {
+                cardinality: 1000,
+                zipf_exponent: None,
+                benchmark_events: 1
+            },
+            DataShape {
+                cardinality: 1000,
+                zipf_exponent: Some(1.0),
+                benchmark_events: 10000
+            },
+            ShapeMatchPolicy {
+                minimum_benchmark_events: 1000,
+                max_log2_cardinality_distance: 1.0,
+                max_zipf_distance: 0.5
+            },
+        )
+        .is_none());
+    }
 }

asap-planner-rs/src/optimizer/mod.rs

@@ -12,9 +12,10 @@ pub mod translator;
 
 pub use aqe_extractor::{extract_aqes, RQE};
 pub use atomic_costs::{
-    load_atomic_cost_table, load_optional_selected_atomic_cost_table,
-    load_selected_atomic_cost_table, resolve_atomic_costs, AtomicCostEntry, AtomicCostTable,
-    ExternalWorkload, WorkloadDescription,
+    load_atomic_cost_table, load_nearest_atomic_cost_table,
+    load_optional_selected_atomic_cost_table, load_selected_atomic_cost_table,
+    resolve_atomic_costs, AtomicCostEntry, AtomicCostTable, DataShape, ExternalWorkload,
+    ShapeMatchPolicy, WorkloadDescription,
 };
 pub use candidate_gen::{
     enumerate_candidates, enumerate_candidates_with_label_group_count, CandidateConfig,