chronix_analytics/preprocess/
arrow_adapters.rs1use std::sync::Arc;
9
10use arrow::array::{ArrayRef, Float64Array, Int64Array};
11
12use crate::preprocess::{Interpolator, PreprocessConfig, PreprocessPipeline, Smoother};
13
14#[derive(Debug)]
16pub struct ArrowPreprocessResult {
17 pub values: Float64Array,
19 pub timestamps: Int64Array,
21 pub gaps_filled: usize,
23}
24
25impl ArrowPreprocessResult {
26 pub fn values_array_ref(&self) -> ArrayRef {
28 Arc::new(self.values.clone())
29 }
30
31 pub fn timestamps_array_ref(&self) -> ArrayRef {
33 Arc::new(self.timestamps.clone())
34 }
35}
36
37pub 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 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
58pub fn arrow_smooth(smoother: &Smoother, values: &Float64Array) -> Float64Array {
60 let slice = values.values().as_ref();
61 Float64Array::from(smoother.smooth(slice))
62}
63
64pub 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 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 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}