Skip to main content

dag_ml_data_core/
aggregation.rs

1//! ADR-07 canonical prediction-aggregation reducer contract.
2//!
3//! `dag-ml-data` owns the reducer *names, parameters and validation* so the
4//! conformance pack is complete and every binding agrees on the surface. The
5//! numerical execution (mean/median/robust-mean, the Hotelling T² outlier
6//! boundary, vote tallies) belongs to the host bridge — there is no reducer
7//! math in this layer. `skipna = true` is fixed ADR-07 contract behaviour and is
8//! not configurable in v1.
9//!
10//! The policy is a flat struct with `deny_unknown_fields` (not an internally
11//! tagged enum) so the JSON wire shape stays flat — `{"reducer": "robust_mean",
12//! "trim_fraction": 0.1}` — yet unknown keys such as `skipna` are rejected at the
13//! boundary, keeping the Rust deserializer and the C ABI validator
14//! (`dagmldata_aggregation_policy_validate_json`) in agreement. Type-level
15//! looseness (a parameter that belongs to another reducer) is caught
16//! semantically by [`AggregationPolicy::validate`]. No shared JSON schema is
17//! published yet — the conformance pack pins only the validator symbol (the
18//! dag-ml coordinator-level aggregation policy stays a distinct contract; see the
19//! ADR-07 reconciliation follow-up).
20
21use serde::{Deserialize, Serialize};
22
23use crate::error::{DataError, Result};
24
25/// Default `trim_fraction` for `robust_mean` when unspecified (ADR-07).
26pub const DEFAULT_TRIM_FRACTION: f64 = 0.1;
27/// Default `threshold` for `exclude_outliers` when unspecified (ADR-07).
28pub const DEFAULT_THRESHOLD: f64 = 0.95;
29
30/// A reducer label (weight column, custom reducer id) must be non-empty, carry no
31/// leading or trailing whitespace, and contain no control characters.
32fn is_clean_label(value: &str) -> bool {
33    !value.is_empty() && value == value.trim() && !value.chars().any(char::is_control)
34}
35
36/// The canonical reducer names. The `custom` escape hatch is intentionally not
37/// listed; it is addressed by `custom_reducer_id` instead. This array is the
38/// source of truth for the canonical set (the conformance pack pins only the C
39/// ABI validator symbol, not a shared reducer list).
40pub const CANONICAL_AGGREGATION_REDUCERS: [&str; 6] = [
41    "mean",
42    "weighted_mean",
43    "median",
44    "vote",
45    "robust_mean",
46    "exclude_outliers",
47];
48
49/// Canonical per-sample aggregation reducer name (ADR-07).
50#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
51#[serde(rename_all = "snake_case")]
52pub enum AggregationReducerName {
53    /// Arithmetic mean of valid values.
54    Mean,
55    /// Mean weighted by a provider-supplied per-row weight column.
56    WeightedMean,
57    /// Sample median.
58    Median,
59    /// Majority vote (classification only); refused on a regression task.
60    Vote,
61    /// Trimmed mean: drop the top/bottom `trim_fraction` before averaging.
62    RobustMean,
63    /// Drop rows outside a Hotelling T² confidence boundary at `threshold`, then mean.
64    ExcludeOutliers,
65    /// Host-controller reducer addressed by a stable `custom_reducer_id`.
66    Custom,
67}
68
69impl AggregationReducerName {
70    /// The snake_case wire name for diagnostics.
71    pub fn as_str(self) -> &'static str {
72        match self {
73            Self::Mean => "mean",
74            Self::WeightedMean => "weighted_mean",
75            Self::Median => "median",
76            Self::Vote => "vote",
77            Self::RobustMean => "robust_mean",
78            Self::ExcludeOutliers => "exclude_outliers",
79            Self::Custom => "custom",
80        }
81    }
82}
83
84/// Prediction task kind needed to validate task-sensitive reducers. This is a
85/// deliberately tiny enum; the bridge maps its richer task type into it. It does
86/// not make `dag-ml-data` the home of DAG-ML task semantics.
87#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
88#[serde(rename_all = "snake_case")]
89pub enum PredictionTaskKind {
90    Regression,
91    Classification,
92}
93
94/// Aggregation policy: which reducer collapses repeated predictions to one value
95/// per sample, plus the single parameter that reducer accepts. Parameters that
96/// do not belong to the chosen reducer must be absent (enforced by
97/// [`validate`](Self::validate)). `None` numeric parameters take the ADR-07
98/// defaults ([`DEFAULT_TRIM_FRACTION`], [`DEFAULT_THRESHOLD`]).
99#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
100#[serde(deny_unknown_fields, rename_all = "snake_case")]
101pub struct AggregationPolicy {
102    pub reducer: AggregationReducerName,
103    /// `robust_mean` trim fraction in `[0.0, 0.5)`.
104    #[serde(default, skip_serializing_if = "Option::is_none")]
105    pub trim_fraction: Option<f64>,
106    /// `exclude_outliers` confidence threshold in `(0.0, 1.0)`.
107    #[serde(default, skip_serializing_if = "Option::is_none")]
108    pub threshold: Option<f64>,
109    /// `weighted_mean` provider-supplied per-row weight column.
110    #[serde(default, skip_serializing_if = "Option::is_none")]
111    pub weight_column: Option<String>,
112    /// `custom` host-controller reducer id.
113    #[serde(default, skip_serializing_if = "Option::is_none")]
114    pub custom_reducer_id: Option<String>,
115}
116
117impl AggregationPolicy {
118    /// Validate the reducer and its parameters: the reducer's own parameter is
119    /// range-checked, and any parameter belonging to a different reducer must be
120    /// absent. No task context (see [`validate_for_task`](Self::validate_for_task)).
121    pub fn validate(&self) -> Result<()> {
122        use AggregationReducerName as R;
123
124        // Each reducer permits exactly one optional parameter (or none).
125        let (allow_trim, allow_threshold, allow_weight, allow_custom) = match self.reducer {
126            R::Mean | R::Median | R::Vote => (false, false, false, false),
127            R::RobustMean => (true, false, false, false),
128            R::ExcludeOutliers => (false, true, false, false),
129            R::WeightedMean => (false, false, true, false),
130            R::Custom => (false, false, false, true),
131        };
132        self.reject_unexpected("trim_fraction", self.trim_fraction.is_some(), allow_trim)?;
133        self.reject_unexpected("threshold", self.threshold.is_some(), allow_threshold)?;
134        self.reject_unexpected("weight_column", self.weight_column.is_some(), allow_weight)?;
135        self.reject_unexpected(
136            "custom_reducer_id",
137            self.custom_reducer_id.is_some(),
138            allow_custom,
139        )?;
140
141        match self.reducer {
142            R::Mean | R::Median | R::Vote => {}
143            R::RobustMean => {
144                if let Some(trim) = self.trim_fraction {
145                    if !trim.is_finite() || !(0.0..0.5).contains(&trim) {
146                        return Err(DataError::Validation(format!(
147                            "aggregation reducer `robust_mean` trim_fraction must be in [0.0, 0.5), got {trim}"
148                        )));
149                    }
150                }
151            }
152            R::ExcludeOutliers => {
153                if let Some(threshold) = self.threshold {
154                    if !threshold.is_finite() || threshold <= 0.0 || threshold >= 1.0 {
155                        return Err(DataError::Validation(format!(
156                            "aggregation reducer `exclude_outliers` threshold must be in (0.0, 1.0), got {threshold}"
157                        )));
158                    }
159                }
160            }
161            R::WeightedMean => {
162                let column = self.weight_column.as_deref().unwrap_or_default();
163                if !is_clean_label(column) {
164                    return Err(DataError::Validation(
165                        "aggregation reducer `weighted_mean` requires a non-empty weight_column with no surrounding whitespace or control characters"
166                            .to_string(),
167                    ));
168                }
169            }
170            R::Custom => {
171                let id = self.custom_reducer_id.as_deref().unwrap_or_default();
172                if !is_clean_label(id) {
173                    return Err(DataError::Validation(
174                        "aggregation reducer `custom` requires a non-empty custom_reducer_id with no surrounding whitespace or control characters"
175                            .to_string(),
176                    ));
177                }
178                if CANONICAL_AGGREGATION_REDUCERS.contains(&id) || id == "custom" {
179                    return Err(DataError::Validation(format!(
180                        "custom_reducer_id `{id}` collides with a reserved reducer name"
181                    )));
182                }
183            }
184        }
185        Ok(())
186    }
187
188    /// Validate parameters and refuse task-incompatible reducers (`vote` on a
189    /// regression task), raising [`DataError::IncompatibleReducer`].
190    pub fn validate_for_task(&self, task: PredictionTaskKind) -> Result<()> {
191        self.validate()?;
192        if matches!(
193            (self.reducer, task),
194            (AggregationReducerName::Vote, PredictionTaskKind::Regression)
195        ) {
196            return Err(DataError::IncompatibleReducer {
197                reducer: "vote",
198                task: "regression",
199            });
200        }
201        Ok(())
202    }
203
204    fn reject_unexpected(&self, param: &str, present: bool, allowed: bool) -> Result<()> {
205        if present && !allowed {
206            return Err(DataError::Validation(format!(
207                "aggregation reducer `{}` does not accept parameter `{param}`",
208                self.reducer.as_str()
209            )));
210        }
211        Ok(())
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218
219    fn policy(reducer: AggregationReducerName) -> AggregationPolicy {
220        AggregationPolicy {
221            reducer,
222            trim_fraction: None,
223            threshold: None,
224            weight_column: None,
225            custom_reducer_id: None,
226        }
227    }
228
229    #[test]
230    fn flat_wire_shape_and_strict_unknown_fields() {
231        let text = serde_json::to_string(&policy(AggregationReducerName::Mean)).unwrap();
232        assert_eq!(text, r#"{"reducer":"mean"}"#);
233        let mut robust = policy(AggregationReducerName::RobustMean);
234        robust.trim_fraction = Some(0.2);
235        let text = serde_json::to_string(&robust).unwrap();
236        assert_eq!(text, r#"{"reducer":"robust_mean","trim_fraction":0.2}"#);
237
238        // unknown keys (e.g. skipna) are rejected at the boundary
239        assert!(
240            serde_json::from_str::<AggregationPolicy>(r#"{"reducer":"mean","skipna":false}"#)
241                .is_err()
242        );
243        // a missing reducer is rejected
244        assert!(serde_json::from_str::<AggregationPolicy>(r#"{}"#).is_err());
245        // a bare string is not the policy shape
246        assert!(serde_json::from_str::<AggregationPolicy>(r#""mean""#).is_err());
247    }
248
249    #[test]
250    fn cross_parameter_contamination_is_rejected() {
251        let mut mean = policy(AggregationReducerName::Mean);
252        mean.trim_fraction = Some(0.1);
253        assert!(mean.validate().is_err());
254        let mut robust = policy(AggregationReducerName::RobustMean);
255        robust.threshold = Some(0.9);
256        assert!(robust.validate().is_err());
257        let mut custom = policy(AggregationReducerName::Custom);
258        custom.custom_reducer_id = Some("ok.id".to_string());
259        custom.weight_column = Some("w".to_string());
260        assert!(custom.validate().is_err());
261    }
262
263    #[test]
264    fn no_param_reducers_validate_clean() {
265        for reducer in [
266            AggregationReducerName::Mean,
267            AggregationReducerName::Median,
268            AggregationReducerName::Vote,
269        ] {
270            assert!(policy(reducer).validate().is_ok());
271        }
272    }
273
274    #[test]
275    fn robust_mean_trim_fraction_bounds() {
276        let mut p = policy(AggregationReducerName::RobustMean);
277        assert!(p.validate().is_ok()); // None -> default applies downstream
278        for good in [0.0, 0.25, 0.49] {
279            p.trim_fraction = Some(good);
280            assert!(p.validate().is_ok(), "trim {good} should pass");
281        }
282        for bad in [-0.1, 0.5, 0.9, f64::NAN, f64::INFINITY] {
283            p.trim_fraction = Some(bad);
284            assert!(p.validate().is_err(), "trim {bad} should fail");
285        }
286    }
287
288    #[test]
289    fn exclude_outliers_threshold_bounds() {
290        let mut p = policy(AggregationReducerName::ExcludeOutliers);
291        assert!(p.validate().is_ok());
292        p.threshold = Some(0.95);
293        assert!(p.validate().is_ok());
294        for bad in [0.0, 1.0, -0.5, 1.5, f64::NAN] {
295            p.threshold = Some(bad);
296            assert!(p.validate().is_err(), "threshold {bad} should fail");
297        }
298    }
299
300    #[test]
301    fn weighted_mean_requires_weight_column() {
302        let mut p = policy(AggregationReducerName::WeightedMean);
303        assert!(p.validate().is_err()); // missing
304        p.weight_column = Some("   ".to_string());
305        assert!(p.validate().is_err()); // blank
306        p.weight_column = Some("weight".to_string());
307        assert!(p.validate().is_ok());
308        // surrounding whitespace / control characters are rejected
309        p.weight_column = Some(" weight ".to_string());
310        assert!(p.validate().is_err());
311        p.weight_column = Some("we\tight".to_string());
312        assert!(p.validate().is_err());
313    }
314
315    #[test]
316    fn custom_reducer_id_non_empty_and_not_reserved() {
317        let mut p = policy(AggregationReducerName::Custom);
318        assert!(p.validate().is_err()); // missing
319        p.custom_reducer_id = Some("site.special".to_string());
320        assert!(p.validate().is_ok());
321        for reserved in ["mean", "robust_mean", "custom"] {
322            p.custom_reducer_id = Some(reserved.to_string());
323            assert!(p.validate().is_err(), "reserved `{reserved}` should fail");
324        }
325        // surrounding whitespace / control characters are rejected
326        for dirty in [" site.special", "site.special ", "site\tspecial"] {
327            p.custom_reducer_id = Some(dirty.to_string());
328            assert!(p.validate().is_err(), "dirty `{dirty:?}` should fail");
329        }
330    }
331
332    #[test]
333    fn vote_is_classification_only() {
334        let vote = policy(AggregationReducerName::Vote);
335        assert!(vote
336            .validate_for_task(PredictionTaskKind::Classification)
337            .is_ok());
338        let error = vote
339            .validate_for_task(PredictionTaskKind::Regression)
340            .unwrap_err();
341        assert_eq!(error.code(), "incompatible_reducer");
342        assert_eq!(error.error_code(), 0x0002_0003);
343        assert!(policy(AggregationReducerName::Mean)
344            .validate_for_task(PredictionTaskKind::Regression)
345            .is_ok());
346    }
347
348    #[test]
349    fn canonical_names_round_trip() {
350        for name in CANONICAL_AGGREGATION_REDUCERS {
351            let json = format!(r#"{{"reducer":"{name}"}}"#);
352            let parsed: AggregationPolicy = serde_json::from_str(&json).unwrap();
353            assert_eq!(parsed.reducer.as_str(), name);
354        }
355        assert_eq!(CANONICAL_AGGREGATION_REDUCERS.len(), 6);
356    }
357}