rill_ml/preprocessing/
forward_fill.rs

1//! Forward-fill imputer for missing data.
2//!
3//! Replaces `NaN` values with the last observed non-NaN value for
4//! each feature. If no value has been seen yet, `NaN` is replaced
5//! with `0.0`.
6
7use crate::error::{RillError, checked_increment, ensure_finite};
8use crate::traits::Transformer;
9
10/// Replaces `NaN` values with the last observed non-NaN value.
11///
12/// This transformer accepts `NaN` in its input. When no value has
13/// been seen for a feature, `0.0` is used as the fill value.
14#[derive(Debug, Clone)]
15#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
16pub struct ForwardFill {
17    feature_count: usize,
18    last_values: Vec<Option<f64>>,
19    samples_seen: u64,
20}
21
22impl ForwardFill {
23    /// Create a new forward-fill imputer for `feature_count` features.
24    ///
25    /// # Errors
26    /// Returns [`RillError::EmptyFeatures`] if `feature_count` is `0`.
27    pub fn new(feature_count: usize) -> Result<Self, RillError> {
28        if feature_count == 0 {
29            return Err(RillError::EmptyFeatures);
30        }
31        Ok(Self {
32            feature_count,
33            last_values: vec![None; feature_count],
34            samples_seen: 0,
35        })
36    }
37
38    /// The per-feature last observed non-NaN values (`None` if unseen).
39    pub fn last_values(&self) -> &[Option<f64>] {
40        &self.last_values
41    }
42
43    /// Validate only the dimension, allowing NaN values.
44    fn check_dimension(&self, features: &[f64]) -> Result<(), RillError> {
45        if features.is_empty() {
46            return Err(RillError::EmptyFeatures);
47        }
48        if features.len() != self.feature_count {
49            return Err(RillError::DimensionMismatch {
50                expected: self.feature_count,
51                actual: features.len(),
52            });
53        }
54        Ok(())
55    }
56}
57
58impl Transformer for ForwardFill {
59    fn input_dim(&self) -> usize {
60        self.feature_count
61    }
62
63    fn output_dim(&self) -> usize {
64        self.feature_count
65    }
66
67    fn transform(&self, features: &[f64]) -> Result<Vec<f64>, RillError> {
68        self.check_dimension(features)?;
69        let mut out = Vec::with_capacity(features.len());
70        for (i, &x) in features.iter().enumerate() {
71            if x.is_nan() {
72                out.push(self.last_values[i].unwrap_or(0.0));
73            } else {
74                ensure_finite("feature", x)?;
75                out.push(x);
76            }
77        }
78        Ok(out)
79    }
80
81    fn update(&mut self, features: &[f64]) -> Result<(), RillError> {
82        self.check_dimension(features)?;
83        for (i, &x) in features.iter().enumerate() {
84            if x.is_nan() {
85                continue;
86            }
87            ensure_finite("feature", x)?;
88            self.last_values[i] = Some(x);
89        }
90        self.samples_seen = checked_increment(self.samples_seen, "samples_seen")?;
91        Ok(())
92    }
93
94    fn samples_seen(&self) -> u64 {
95        self.samples_seen
96    }
97
98    fn reset(&mut self) {
99        for v in &mut self.last_values {
100            *v = None;
101        }
102        self.samples_seen = 0;
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn nan_replaced_with_last_value() {
112        let mut ff = ForwardFill::new(2).unwrap();
113        ff.update(&[1.0, 2.0]).unwrap();
114        ff.update(&[3.0, 4.0]).unwrap();
115        let out = ff.transform(&[f64::NAN, f64::NAN]).unwrap();
116        assert_eq!(out, vec![3.0, 4.0]);
117    }
118
119    #[test]
120    fn nan_replaced_with_zero_when_no_data() {
121        let ff = ForwardFill::new(2).unwrap();
122        let out = ff.transform(&[f64::NAN, f64::NAN]).unwrap();
123        assert_eq!(out, vec![0.0, 0.0]);
124    }
125
126    #[test]
127    fn non_nan_passed_through() {
128        let mut ff = ForwardFill::new(2).unwrap();
129        ff.update(&[5.0, 6.0]).unwrap();
130        let out = ff.transform(&[1.5, -2.0]).unwrap();
131        assert_eq!(out, vec![1.5, -2.0]);
132    }
133
134    #[test]
135    fn last_value_updates_correctly() {
136        let mut ff = ForwardFill::new(2).unwrap();
137        ff.update(&[1.0, 10.0]).unwrap();
138        assert_eq!(ff.last_values()[0], Some(1.0));
139        assert_eq!(ff.last_values()[1], Some(10.0));
140        ff.update(&[2.0, f64::NAN]).unwrap();
141        assert_eq!(ff.last_values()[0], Some(2.0));
142        assert_eq!(ff.last_values()[1], Some(10.0));
143    }
144
145    #[test]
146    fn nan_skipped_in_update() {
147        let mut ff = ForwardFill::new(2).unwrap();
148        ff.update(&[f64::NAN, 5.0]).unwrap();
149        assert_eq!(ff.last_values()[0], None);
150        assert_eq!(ff.last_values()[1], Some(5.0));
151        ff.update(&[7.0, f64::NAN]).unwrap();
152        assert_eq!(ff.last_values()[0], Some(7.0));
153        assert_eq!(ff.last_values()[1], Some(5.0));
154    }
155
156    #[test]
157    fn dimension_mismatch_rejected() {
158        let ff = ForwardFill::new(3).unwrap();
159        assert!(matches!(
160            ff.transform(&[1.0, 2.0]),
161            Err(RillError::DimensionMismatch { .. })
162        ));
163        let mut ff = ff;
164        assert!(matches!(
165            ff.update(&[1.0, 2.0, 3.0, 4.0]),
166            Err(RillError::DimensionMismatch { .. })
167        ));
168    }
169
170    #[test]
171    fn reset_clears_state() {
172        let mut ff = ForwardFill::new(2).unwrap();
173        ff.update(&[1.0, 2.0]).unwrap();
174        ff.update(&[3.0, 4.0]).unwrap();
175        assert_eq!(ff.samples_seen(), 2);
176        ff.reset();
177        assert_eq!(ff.samples_seen(), 0);
178        assert_eq!(ff.last_values()[0], None);
179        assert_eq!(ff.last_values()[1], None);
180    }
181
182    #[test]
183    #[cfg(feature = "serde")]
184    fn serde_roundtrip() {
185        let mut ff = ForwardFill::new(2).unwrap();
186        ff.update(&[1.0, f64::NAN]).unwrap();
187        ff.update(&[3.0, 5.0]).unwrap();
188        let json = serde_json::to_string(&ff).unwrap();
189        let restored: ForwardFill = serde_json::from_str(&json).unwrap();
190        assert_eq!(restored.input_dim(), ff.input_dim());
191        assert_eq!(restored.output_dim(), ff.output_dim());
192        assert_eq!(restored.samples_seen(), ff.samples_seen());
193        assert_eq!(restored.last_values(), ff.last_values());
194    }
195}
rill_ml/preprocessing/forward_fill.rs

rill_ml/preprocessing/
forward_fill.rs
