micromegas_datafusion_extensions/histogram/
accumulator.rs

1use std::sync::Arc;
2
3use datafusion::{
4    arrow::{
5        array::{
6            Array, ArrayBuilder, ArrayRef, Float64Array, Int64Array, ListBuilder, PrimitiveBuilder,
7            StructBuilder, UInt64Builder,
8        },
9        datatypes::{DataType, Field, Float64Type, UInt64Type},
10    },
11    common::utils::proxy::VecAllocExt,
12    error::DataFusionError,
13    logical_expr::Accumulator,
14    scalar::ScalarValue,
15};
16
17use super::histogram_udaf::HistogramArray;
18
19/// An accumulator for computing histograms.
20#[derive(Debug)]
21pub struct HistogramAccumulator {
22    start: Option<f64>,
23    end: Option<f64>,
24    min: f64,
25    max: f64,
26    sum: f64,
27    sum_sq: f64,
28    count: u64,
29    bins: Vec<u64>,
30}
31
32impl HistogramAccumulator {
33    pub fn new(start: f64, end: f64, nb_bins: usize) -> Self {
34        let bins = vec![0; nb_bins];
35        Self {
36            start: Some(start),
37            end: Some(end),
38            bins,
39            min: f64::MAX,
40            max: f64::MIN,
41            sum: 0.0,
42            sum_sq: 0.0,
43            count: 0,
44        }
45    }
46
47    pub fn new_non_configured() -> Self {
48        Self {
49            start: None,
50            end: None,
51            min: f64::MAX,
52            max: f64::MIN,
53            sum: 0.0,
54            sum_sq: 0.0,
55            count: 0,
56            bins: Vec::new(),
57        }
58    }
59
60    /// Validates and sets histogram bounds from scalar parameters.
61    pub fn configure_from_params(
62        &mut self,
63        start: f64,
64        end: f64,
65        nb_bins: i64,
66    ) -> datafusion::error::Result<()> {
67        if nb_bins < 1 {
68            return Err(DataFusionError::Execution(format!(
69                "make_histogram: nb_bins must be >= 1, got {nb_bins}"
70            )));
71        }
72        if !start.is_finite() {
73            return Err(DataFusionError::Execution(format!(
74                "make_histogram: start must be finite, got {start}"
75            )));
76        }
77        if !end.is_finite() {
78            return Err(DataFusionError::Execution(format!(
79                "make_histogram: end must be finite, got {end}"
80            )));
81        }
82        if start > end {
83            return Err(DataFusionError::Execution(format!(
84                "make_histogram: start ({start}) must be <= end ({end})"
85            )));
86        }
87        self.start = Some(start);
88        self.end = Some(end);
89        self.bins.resize(nb_bins as usize, 0);
90        Ok(())
91    }
92
93    /// If not configured, scans for the first non-null row and uses it as a template.
94    /// If already configured or if all rows are null, does nothing.
95    pub fn configure(&mut self, histo_array: &HistogramArray) -> datafusion::error::Result<()> {
96        if self.start.is_some() {
97            return Ok(());
98        }
99        let Some(idx) = (0..histo_array.len()).find(|&i| !histo_array.is_null_at(i)) else {
100            return Ok(());
101        };
102        self.start = Some(histo_array.get_start(idx)?);
103        self.end = Some(histo_array.get_end(idx)?);
104        self.bins.resize(histo_array.get_bins(idx)?.len(), 0);
105        Ok(())
106    }
107
108    pub fn update_batch_scalars(
109        &mut self,
110        scalars: &Float64Array,
111    ) -> datafusion::error::Result<()> {
112        if self.start.is_none() || self.end.is_none() {
113            return Err(DataFusionError::Execution(
114                "can't record scalar in a non-configured histogram".into(),
115            ));
116        }
117        let start = self.start.unwrap();
118        let range = self.end.unwrap() - start;
119        let bin_width = range / (self.bins.len() as f64);
120        for i in 0..scalars.len() {
121            if !scalars.is_null(i) {
122                let v = scalars.value(i);
123                self.min = self.min.min(v);
124                self.max = self.max.max(v);
125                self.sum += v;
126                self.sum_sq += v * v;
127                self.count += 1;
128                let bin_index = (((v - start) / bin_width).floor()) as usize;
129                let bin_index = bin_index.clamp(0, self.bins.len() - 1);
130                self.bins[bin_index] += 1;
131            }
132        }
133        Ok(())
134    }
135
136    pub fn merge_histograms(
137        &mut self,
138        histo_array: &HistogramArray,
139    ) -> datafusion::error::Result<()> {
140        self.configure(histo_array)?;
141        for index_histo in 0..histo_array.len() {
142            if histo_array.is_null_at(index_histo) {
143                continue;
144            }
145            let start = histo_array.get_start(index_histo)?;
146            if self.start.unwrap() != start {
147                return Err(DataFusionError::Execution(
148                    "Error merging incompatible histograms".into(),
149                ));
150            }
151            let end = histo_array.get_end(index_histo)?;
152            if self.end.unwrap() != end {
153                return Err(DataFusionError::Execution(
154                    "Error merging incompatible histograms".into(),
155                ));
156            }
157
158            let min = histo_array.get_min(index_histo)?;
159            let max = histo_array.get_max(index_histo)?;
160            let sum = histo_array.get_sum(index_histo)?;
161            let sum_sq = histo_array.get_sum_sq(index_histo)?;
162            let count = histo_array.get_count(index_histo)?;
163            let bins = histo_array.get_bins(index_histo)?;
164            if bins.len() != self.bins.len() {
165                return Err(DataFusionError::Execution(
166                    "Error merging incompatible histograms".into(),
167                ));
168            }
169            self.min = self.min.min(min);
170            self.max = self.max.max(max);
171            self.sum += sum;
172            self.sum_sq += sum_sq;
173            self.count += count;
174
175            // optim opportunity: use arrow compute
176            for i in 0..self.bins.len() {
177                self.bins[i] += bins.value(i);
178            }
179        }
180        Ok(())
181    }
182}
183
184impl Accumulator for HistogramAccumulator {
185    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion::error::Result<()> {
186        // we support two signatures
187        // scalar case: [starts, ends, bin_counts, scalars_to_reduce]
188        // merge case: [histograms]
189
190        match values.len() {
191            4 => {
192                if values[0].is_empty() {
193                    return Ok(());
194                }
195                for (i, name) in [(0usize, "start"), (1, "end"), (2, "nb_bins")] {
196                    if values[i].is_null(0) {
197                        return Err(DataFusionError::Execution(format!(
198                            "make_histogram: {name} argument is null"
199                        )));
200                    }
201                }
202                let batch_start = values[0]
203                    .as_any()
204                    .downcast_ref::<Float64Array>()
205                    .ok_or_else(|| {
206                        DataFusionError::Internal("values[0] should be a Float64Array".into())
207                    })?
208                    .value(0);
209                let batch_end = values[1]
210                    .as_any()
211                    .downcast_ref::<Float64Array>()
212                    .ok_or_else(|| {
213                        DataFusionError::Internal("values[1] should be a Float64Array".into())
214                    })?
215                    .value(0);
216                let batch_nb_bins = values[2]
217                    .as_any()
218                    .downcast_ref::<Int64Array>()
219                    .ok_or_else(|| {
220                        DataFusionError::Internal("values[2] should be an Int64Array".into())
221                    })?
222                    .value(0);
223                if let Some(configured_start) = self.start {
224                    let configured_end = self.end.expect("end is set whenever start is set");
225                    if configured_start != batch_start
226                        || configured_end != batch_end
227                        || self.bins.len() != batch_nb_bins as usize
228                    {
229                        return Err(DataFusionError::Execution(
230                            "make_histogram: bounds/bins changed between batches".into(),
231                        ));
232                    }
233                } else {
234                    self.configure_from_params(batch_start, batch_end, batch_nb_bins)?;
235                }
236                let scalars = values[3]
237                    .as_any()
238                    .downcast_ref::<Float64Array>()
239                    .ok_or_else(|| {
240                        DataFusionError::Internal("values[3] should be a Float64Array".into())
241                    })?;
242                self.update_batch_scalars(scalars)
243            }
244            1 => {
245                let histo_array: HistogramArray = values[0].as_ref().try_into()?;
246                self.merge_histograms(&histo_array)
247            }
248
249            other => Err(DataFusionError::Execution(format!(
250                "invalid arguments to HistogramAccumulator::update_batch, nb_values={other}"
251            ))),
252        }
253    }
254
255    fn evaluate(&mut self) -> datafusion::error::Result<datafusion::scalar::ScalarValue> {
256        let fields = state_arrow_fields();
257        let mut struct_builder = StructBuilder::from_fields(fields, 1);
258        let start_builder = struct_builder
259            .field_builder::<PrimitiveBuilder<Float64Type>>(0)
260            .ok_or_else(|| DataFusionError::Execution("Error accessing to start builder".into()))?;
261        start_builder.append_value(self.start.unwrap_or(0.0));
262
263        let end_builder = struct_builder
264            .field_builder::<PrimitiveBuilder<Float64Type>>(1)
265            .ok_or_else(|| DataFusionError::Execution("Error accessing to end builder".into()))?;
266        end_builder.append_value(self.end.unwrap_or(0.0));
267
268        let min_builder = struct_builder
269            .field_builder::<PrimitiveBuilder<Float64Type>>(2)
270            .ok_or_else(|| DataFusionError::Execution("Error accessing to min builder".into()))?;
271        min_builder.append_value(self.min);
272
273        let max_builder = struct_builder
274            .field_builder::<PrimitiveBuilder<Float64Type>>(3)
275            .ok_or_else(|| DataFusionError::Execution("Error accessing to max builder".into()))?;
276        max_builder.append_value(self.max);
277
278        let sum_builder = struct_builder
279            .field_builder::<PrimitiveBuilder<Float64Type>>(4)
280            .ok_or_else(|| DataFusionError::Execution("Error accessing to sum builder".into()))?;
281        sum_builder.append_value(self.sum);
282
283        let sum_sq_builder = struct_builder
284            .field_builder::<PrimitiveBuilder<Float64Type>>(5)
285            .ok_or_else(|| {
286                DataFusionError::Execution("Error accessing to sum_sq builder".into())
287            })?;
288        sum_sq_builder.append_value(self.sum_sq);
289
290        let count_builder = struct_builder
291            .field_builder::<PrimitiveBuilder<UInt64Type>>(6)
292            .ok_or_else(|| DataFusionError::Execution("Error accessing to count builder".into()))?;
293        count_builder.append_value(self.count);
294
295        let bins_builder = struct_builder
296            .field_builder::<ListBuilder<Box<dyn ArrayBuilder>>>(7)
297            .ok_or_else(|| DataFusionError::Execution("Error accessing to bins builder".into()))?;
298        let bin_array_builder = bins_builder
299            .values()
300            .as_any_mut()
301            .downcast_mut::<UInt64Builder>()
302            .ok_or_else(|| {
303                DataFusionError::Execution("Error accessing to bins array builder".into())
304            })?;
305        bin_array_builder.append_slice(&self.bins);
306        bins_builder.append(true);
307        struct_builder.append(self.start.is_some());
308        Ok(ScalarValue::Struct(Arc::new(struct_builder.finish())))
309    }
310
311    fn size(&self) -> usize {
312        size_of_val(self) + self.bins.allocated_size()
313    }
314
315    fn state(&mut self) -> datafusion::error::Result<Vec<datafusion::scalar::ScalarValue>> {
316        Ok(vec![self.evaluate()?])
317    }
318
319    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion::error::Result<()> {
320        for state in states {
321            let histo_array: HistogramArray = state.try_into()?;
322            self.merge_histograms(&histo_array)?;
323        }
324        Ok(())
325    }
326}
327
328/// Returns the Arrow fields for the histogram state.
329pub fn state_arrow_fields() -> Vec<Field> {
330    vec![
331        Field::new("start", DataType::Float64, false),
332        Field::new("end", DataType::Float64, false),
333        Field::new("min", DataType::Float64, false),
334        Field::new("max", DataType::Float64, false),
335        Field::new("sum", DataType::Float64, false),
336        Field::new("sum_sq", DataType::Float64, false),
337        Field::new("count", DataType::UInt64, false),
338        Field::new(
339            "bins",
340            DataType::List(Arc::new(Field::new("bin", DataType::UInt64, false))),
341            false,
342        ),
343    ]
344}
micromegas_datafusion_extensions/histogram/accumulator.rs

micromegas_datafusion_extensions/histogram/
accumulator.rs
