shape_runtime/
columnar_aggregations.rs

1//! SIMD-optimized aggregations for columnar object series
2//!
3//! Uses the `wide` crate for portable SIMD that works on stable Rust
4//! and automatically selects AVX2/SSE2/NEON based on target platform.
5
6use shape_value::aligned_vec::AlignedVec;
7use wide::f64x4;
8
9/// Sum a numeric column using portable SIMD
10///
11/// Performance: ~4x faster than scalar sum for large arrays (>1000 elements)
12pub fn sum_f64_column(data: &AlignedVec<f64>) -> f64 {
13    let slice = data.as_slice();
14    let len = slice.len();
15
16    if len == 0 {
17        return 0.0;
18    }
19
20    // Process 4 f64 values at a time using wide's portable SIMD
21    let mut sum = f64x4::ZERO;
22    let chunks = len / 4;
23
24    for i in 0..chunks {
25        let idx = i * 4;
26        // Load 4 f64 values - wide automatically uses best SIMD for platform
27        let chunk = f64x4::new([slice[idx], slice[idx + 1], slice[idx + 2], slice[idx + 3]]);
28        sum += chunk;
29    }
30
31    // Horizontal sum of the SIMD vector
32    let mut result = sum.reduce_add();
33
34    // Handle remaining elements (0-3 values)
35    for i in (chunks * 4)..len {
36        result += slice[i];
37    }
38
39    result
40}
41
42/// Calculate mean of a numeric column using SIMD sum
43pub fn mean_f64_column(data: &AlignedVec<f64>) -> f64 {
44    if data.is_empty() {
45        return f64::NAN;
46    }
47
48    sum_f64_column(data) / data.len() as f64
49}
50
51/// Count non-NaN values in a numeric column
52pub fn count_valid_f64(data: &AlignedVec<f64>) -> usize {
53    let slice = data.as_slice();
54    slice.iter().filter(|&&v| !v.is_nan()).count()
55}
56
57/// Find minimum value in a numeric column using SIMD
58pub fn min_f64_column(data: &AlignedVec<f64>) -> f64 {
59    let slice = data.as_slice();
60    let len = slice.len();
61
62    if len == 0 {
63        return f64::NAN;
64    }
65
66    // Process 4 values at a time
67    let mut min_vec = f64x4::splat(f64::INFINITY);
68    let chunks = len / 4;
69
70    for i in 0..chunks {
71        let idx = i * 4;
72        let chunk = f64x4::new([slice[idx], slice[idx + 1], slice[idx + 2], slice[idx + 3]]);
73        min_vec = min_vec.min(chunk);
74    }
75
76    // Find minimum across the SIMD vector
77    let arr: [f64; 4] = min_vec.to_array();
78    let mut result = arr.iter().copied().fold(f64::INFINITY, f64::min);
79
80    // Handle remaining elements
81    for i in (chunks * 4)..len {
82        result = result.min(slice[i]);
83    }
84
85    result
86}
87
88/// Find maximum value in a numeric column using SIMD
89pub fn max_f64_column(data: &AlignedVec<f64>) -> f64 {
90    let slice = data.as_slice();
91    let len = slice.len();
92
93    if len == 0 {
94        return f64::NAN;
95    }
96
97    // Process 4 values at a time
98    let mut max_vec = f64x4::splat(f64::NEG_INFINITY);
99    let chunks = len / 4;
100
101    for i in 0..chunks {
102        let idx = i * 4;
103        let chunk = f64x4::new([slice[idx], slice[idx + 1], slice[idx + 2], slice[idx + 3]]);
104        max_vec = max_vec.max(chunk);
105    }
106
107    // Find maximum across the SIMD vector
108    let arr: [f64; 4] = max_vec.to_array();
109    let mut result = arr.iter().copied().fold(f64::NEG_INFINITY, f64::max);
110
111    // Handle remaining elements
112    for i in (chunks * 4)..len {
113        result = result.max(slice[i]);
114    }
115
116    result
117}
118
119// ============================================================================
120// Slice-based wrappers for Arrow Float64Array::values() -> &[f64]
121// ============================================================================
122
123/// Sum a raw f64 slice using portable SIMD.
124pub fn sum_f64_slice(data: &[f64]) -> f64 {
125    let len = data.len();
126    if len == 0 {
127        return 0.0;
128    }
129
130    let mut sum = f64x4::ZERO;
131    let chunks = len / 4;
132
133    for i in 0..chunks {
134        let idx = i * 4;
135        let chunk = f64x4::new([data[idx], data[idx + 1], data[idx + 2], data[idx + 3]]);
136        sum += chunk;
137    }
138
139    let mut result = sum.reduce_add();
140    for i in (chunks * 4)..len {
141        result += data[i];
142    }
143    result
144}
145
146/// Mean of a raw f64 slice using SIMD sum.
147pub fn mean_f64_slice(data: &[f64]) -> f64 {
148    if data.is_empty() {
149        return f64::NAN;
150    }
151    sum_f64_slice(data) / data.len() as f64
152}
153
154/// Minimum of a raw f64 slice using SIMD.
155pub fn min_f64_slice(data: &[f64]) -> f64 {
156    let len = data.len();
157    if len == 0 {
158        return f64::NAN;
159    }
160
161    let mut min_vec = f64x4::splat(f64::INFINITY);
162    let chunks = len / 4;
163
164    for i in 0..chunks {
165        let idx = i * 4;
166        let chunk = f64x4::new([data[idx], data[idx + 1], data[idx + 2], data[idx + 3]]);
167        min_vec = min_vec.min(chunk);
168    }
169
170    let arr: [f64; 4] = min_vec.to_array();
171    let mut result = arr.iter().copied().fold(f64::INFINITY, f64::min);
172    for i in (chunks * 4)..len {
173        result = result.min(data[i]);
174    }
175    result
176}
177
178/// Maximum of a raw f64 slice using SIMD.
179pub fn max_f64_slice(data: &[f64]) -> f64 {
180    let len = data.len();
181    if len == 0 {
182        return f64::NAN;
183    }
184
185    let mut max_vec = f64x4::splat(f64::NEG_INFINITY);
186    let chunks = len / 4;
187
188    for i in 0..chunks {
189        let idx = i * 4;
190        let chunk = f64x4::new([data[idx], data[idx + 1], data[idx + 2], data[idx + 3]]);
191        max_vec = max_vec.max(chunk);
192    }
193
194    let arr: [f64; 4] = max_vec.to_array();
195    let mut result = arr.iter().copied().fold(f64::NEG_INFINITY, f64::max);
196    for i in (chunks * 4)..len {
197        result = result.max(data[i]);
198    }
199    result
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[test]
207    fn test_simd_sum_matches_scalar() {
208        let data = AlignedVec::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]);
209        let simd_result = sum_f64_column(&data);
210        let scalar_result: f64 = data.iter().sum();
211        assert!((simd_result - scalar_result).abs() < 1e-10);
212        assert_eq!(simd_result, 55.0);
213    }
214
215    #[test]
216    fn test_simd_mean() {
217        let data = AlignedVec::from_vec(vec![10.0, 20.0, 30.0, 40.0]);
218        let result = mean_f64_column(&data);
219        assert_eq!(result, 25.0);
220    }
221
222    #[test]
223    fn test_simd_min_max() {
224        let data = AlignedVec::from_vec(vec![5.0, 2.0, 9.0, 1.0, 7.0, 3.0]);
225        assert_eq!(min_f64_column(&data), 1.0);
226        assert_eq!(max_f64_column(&data), 9.0);
227    }
228
229    #[test]
230    fn test_sum_with_odd_length() {
231        // Test with length not divisible by 4
232        let data = AlignedVec::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
233        let result = sum_f64_column(&data);
234        assert_eq!(result, 15.0);
235    }
236
237    #[test]
238    fn test_slice_sum() {
239        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
240        assert!((sum_f64_slice(&data) - 55.0).abs() < 1e-10);
241        assert_eq!(sum_f64_slice(&[]), 0.0);
242    }
243
244    #[test]
245    fn test_slice_mean() {
246        assert_eq!(mean_f64_slice(&[10.0, 20.0, 30.0, 40.0]), 25.0);
247        assert!(mean_f64_slice(&[]).is_nan());
248    }
249
250    #[test]
251    fn test_slice_min_max() {
252        let data = vec![5.0, 2.0, 9.0, 1.0, 7.0, 3.0];
253        assert_eq!(min_f64_slice(&data), 1.0);
254        assert_eq!(max_f64_slice(&data), 9.0);
255    }
256}
shape_runtime/columnar_aggregations.rs

shape_runtime/
columnar_aggregations.rs
