1use serde::{Deserialize, Serialize};
22
23use crate::error::{DataError, Result};
24
25pub const DEFAULT_TRIM_FRACTION: f64 = 0.1;
27pub const DEFAULT_THRESHOLD: f64 = 0.95;
29
30fn is_clean_label(value: &str) -> bool {
33 !value.is_empty() && value == value.trim() && !value.chars().any(char::is_control)
34}
35
36pub const CANONICAL_AGGREGATION_REDUCERS: [&str; 6] = [
41 "mean",
42 "weighted_mean",
43 "median",
44 "vote",
45 "robust_mean",
46 "exclude_outliers",
47];
48
49#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
51#[serde(rename_all = "snake_case")]
52pub enum AggregationReducerName {
53 Mean,
55 WeightedMean,
57 Median,
59 Vote,
61 RobustMean,
63 ExcludeOutliers,
65 Custom,
67}
68
69impl AggregationReducerName {
70 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#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
88#[serde(rename_all = "snake_case")]
89pub enum PredictionTaskKind {
90 Regression,
91 Classification,
92}
93
94#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
100#[serde(deny_unknown_fields, rename_all = "snake_case")]
101pub struct AggregationPolicy {
102 pub reducer: AggregationReducerName,
103 #[serde(default, skip_serializing_if = "Option::is_none")]
105 pub trim_fraction: Option<f64>,
106 #[serde(default, skip_serializing_if = "Option::is_none")]
108 pub threshold: Option<f64>,
109 #[serde(default, skip_serializing_if = "Option::is_none")]
111 pub weight_column: Option<String>,
112 #[serde(default, skip_serializing_if = "Option::is_none")]
114 pub custom_reducer_id: Option<String>,
115}
116
117impl AggregationPolicy {
118 pub fn validate(&self) -> Result<()> {
122 use AggregationReducerName as R;
123
124 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 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 assert!(
240 serde_json::from_str::<AggregationPolicy>(r#"{"reducer":"mean","skipna":false}"#)
241 .is_err()
242 );
243 assert!(serde_json::from_str::<AggregationPolicy>(r#"{}"#).is_err());
245 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()); 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()); p.weight_column = Some(" ".to_string());
305 assert!(p.validate().is_err()); p.weight_column = Some("weight".to_string());
307 assert!(p.validate().is_ok());
308 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()); 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 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}