Skip to main content

chronix_analytics/preprocess/
arrow_adapters.rs

1//! Arrow-native wrappers for preprocessing operations.
2//!
3//! These adapters provide zero-copy interop between Arrow `Float64Array` /
4//! `Int64Array` and the slice-based preprocessing functions.  Input arrays
5//! are accessed via `.values()` (O(1), no copy), and output arrays are
6//! constructed from the resulting `Vec`s.
7
8use std::sync::Arc;
9
10use arrow::array::{ArrayRef, Float64Array, Int64Array};
11
12use crate::preprocess::{Interpolator, PreprocessConfig, PreprocessPipeline, Smoother};
13
14/// Result of an Arrow-native preprocessing operation.
15#[derive(Debug)]
16pub struct ArrowPreprocessResult {
17    /// Processed values as Arrow Float64Array.
18    pub values: Float64Array,
19    /// Aligned timestamps as Arrow Int64Array.
20    pub timestamps: Int64Array,
21    /// Number of gaps filled during interpolation.
22    pub gaps_filled: usize,
23}
24
25impl ArrowPreprocessResult {
26    /// Returns the values column as an `ArrayRef` (for DataFusion interop).
27    pub fn values_array_ref(&self) -> ArrayRef {
28        Arc::new(self.values.clone())
29    }
30
31    /// Returns the timestamps column as an `ArrayRef`.
32    pub fn timestamps_array_ref(&self) -> ArrayRef {
33        Arc::new(self.timestamps.clone())
34    }
35}
36
37/// Interpolate gap-fills on Arrow arrays — zero-copy input, new array output.
38pub fn arrow_interpolate(
39    interpolator: &Interpolator,
40    timestamps: &Int64Array,
41    values: &Float64Array,
42    expected_interval_ns: i64,
43) -> Result<ArrowPreprocessResult, crate::preprocess::interpolation::InterpolationError> {
44    // Zero-copy slice access
45    let ts_slice = timestamps.values().as_ref();
46    let val_slice = values.values().as_ref();
47
48    let (out_ts, out_vals, gaps_filled, _gaps_skipped) =
49        interpolator.fill(ts_slice, val_slice, expected_interval_ns)?;
50
51    Ok(ArrowPreprocessResult {
52        values: Float64Array::from(out_vals),
53        timestamps: Int64Array::from(out_ts),
54        gaps_filled,
55    })
56}
57
58/// Smooth an Arrow Float64Array — zero-copy input.
59pub fn arrow_smooth(smoother: &Smoother, values: &Float64Array) -> Float64Array {
60    let slice = values.values().as_ref();
61    Float64Array::from(smoother.smooth(slice))
62}
63
64/// Run the full preprocessing pipeline on Arrow arrays.
65pub fn arrow_preprocess(
66    config: &PreprocessConfig,
67    timestamps: &Int64Array,
68    values: &Float64Array,
69) -> Result<ArrowPreprocessResult, crate::preprocess::PreprocessError> {
70    let ts_slice = timestamps.values().as_ref();
71    let val_slice = values.values().as_ref();
72
73    let result = PreprocessPipeline::run(config, ts_slice, val_slice)?;
74
75    Ok(ArrowPreprocessResult {
76        values: Float64Array::from(result.values),
77        timestamps: Int64Array::from(result.timestamps),
78        gaps_filled: result.gaps_filled,
79    })
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    #[test]
87    fn arrow_interpolate_zero_copy() {
88        let ts = Int64Array::from(vec![0, 1_000_000_000, 3_000_000_000, 4_000_000_000]);
89        let vals = Float64Array::from(vec![1.0, 2.0, 4.0, 5.0]);
90
91        let result = arrow_interpolate(&Interpolator::Linear, &ts, &vals, 1_000_000_000).unwrap();
92
93        // Gap at t=2s should be filled
94        assert!(result.gaps_filled > 0);
95        assert!(result.values.len() >= 4);
96        assert!(result.timestamps.len() == result.values.len());
97    }
98
99    #[test]
100    fn arrow_smooth_produces_float64array() {
101        let vals = Float64Array::from(vec![1.0, 3.0, 5.0, 7.0, 9.0, 11.0]);
102        let smoother = Smoother::MovingAverage { window: 3 };
103        let smoothed = arrow_smooth(&smoother, &vals);
104        assert_eq!(smoothed.len(), vals.len());
105        // Centered moving average at index 2: mean of values[1..4] = (3+5+7)/3 = 5.0
106        assert!((smoothed.value(2) - 5.0).abs() < 1e-10);
107    }
108
109    #[test]
110    fn arrow_preprocess_pipeline() {
111        let ts = Int64Array::from(vec![0i64, 1_000_000_000, 2_000_000_000, 3_000_000_000]);
112        let vals = Float64Array::from(vec![10.0, 20.0, 30.0, 40.0]);
113
114        let config = PreprocessConfig {
115            smoothing: Some(Smoother::Exponential { alpha: 0.3 }),
116            ..PreprocessConfig::default()
117        };
118
119        let result = arrow_preprocess(&config, &ts, &vals).unwrap();
120        assert_eq!(result.values.len(), 4);
121    }
122
123    #[test]
124    fn arrow_result_array_refs() {
125        let result = ArrowPreprocessResult {
126            values: Float64Array::from(vec![1.0, 2.0]),
127            timestamps: Int64Array::from(vec![100, 200]),
128            gaps_filled: 0,
129        };
130        let val_ref = result.values_array_ref();
131        assert_eq!(val_ref.len(), 2);
132        let ts_ref = result.timestamps_array_ref();
133        assert_eq!(ts_ref.len(), 2);
134    }
135}