Skip to main content

chronix_query/
downsample.rs

1//! Time-bucket downsampling engine.
2//!
3//! Groups data points into fixed-width time buckets and applies an
4//! aggregation function per bucket. This enables efficient data reduction
5//! for dashboards, alerting queries, and long-term storage.
6//!
7//! # Two entry points
8//!
9//! [`downsample`] takes a whole `RecordBatch` and returns every bucket in it.
10//! It is correct only when the batch contains *all* the rows for the buckets
11//! it covers.
12//!
13//! [`StreamingDownsampler`] is the one to use over a chunked scan. Bucket
14//! boundaries do not line up with batch boundaries, so downsampling each
15//! batch independently emits the straddling interval twice — once per side —
16//! which a caller then sums, averages or plots as two points. The
17//! streaming form keeps one open bucket across batches and closes it only
18//! when a row from a later bucket arrives, so a bucket is emitted exactly
19//! once no matter how the input was chunked. Peak memory is one bucket's
20//! accumulator, not one batch's worth of rows.
21
22use std::sync::Arc;
23
24use fnv::FnvHashMap;
25
26use arrow::array::{Array, ArrayRef, Float64Array, Int64Array};
27use arrow::datatypes::{DataType, Field, Schema};
28use arrow::record_batch::RecordBatch;
29
30use crate::aggregate::AggFn;
31use crate::error::{QueryError, Result};
32
33/// Extract the value at `index` from a numeric column as `f64`.
34///
35/// Returns `None` if the value is null.
36#[allow(clippy::cast_precision_loss)]
37fn extract_f64(col: &dyn Array, index: usize) -> Option<f64> {
38    crate::extract_f64(col, index)
39}
40
41/// Downsample a `RecordBatch` into time buckets.
42///
43/// For each bucket `[bucket.start_of(ts), bucket.next(start))`, the
44/// aggregation function is applied to the values in the `value_column`.
45/// The bucket may be a fixed span or a calendar one — it is the same type
46/// SQL's `time_bucket()` and a rollup tier use, so the three agree.
47///
48/// Returns a new `RecordBatch` with two columns: `_time` (bucket start)
49/// and the aggregated value column.
50///
51/// # Arguments
52///
53/// * `batch` – input data (must have a `_time` Int64 column)
54/// * `value_column` – name of the column to aggregate
55/// * `bucket` – the bucket width and the calendar it is read against
56/// * `function` – aggregation function to apply per bucket
57///
58/// # Errors
59///
60/// Returns an error if columns are missing or have unexpected types.
61#[allow(clippy::cast_precision_loss)]
62pub fn downsample(
63    batch: &RecordBatch,
64    value_column: &str,
65    bucket: &chronix_core::TimeBucket,
66    function: &AggFn,
67    memory_tracker: Option<&crate::memory::MemoryTracker>,
68) -> Result<RecordBatch> {
69    let schema = batch.schema();
70
71    let time_idx = schema
72        .index_of(chronix_core::TIME_COLUMN)
73        .map_err(|_| QueryError::Validation("'timestamp' column not found".into()))?;
74    let val_idx = schema
75        .index_of(value_column)
76        .map_err(|_| QueryError::FieldNotFound {
77            measurement: String::new(),
78            field: value_column.to_string(),
79        })?;
80
81    let timestamps = batch
82        .column(time_idx)
83        .as_any()
84        .downcast_ref::<Int64Array>()
85        .ok_or_else(|| QueryError::Validation("'timestamp' column is not Int64".into()))?;
86
87    let value_col = batch.column(val_idx).as_ref();
88
89    // Validate the value column is a supported numeric type
90    match value_col.data_type() {
91        DataType::Float64 | DataType::Int64 | DataType::UInt64 | DataType::Decimal128(_, _) => {}
92        dt => {
93            return Err(QueryError::Validation(format!(
94                "'{value_column}' column has unsupported type {dt} \
95                 (expected Float64, Int64, UInt64, or Decimal128)"
96            )));
97        }
98    }
99
100    // Group indices by bucket – use a hash map for O(1) inserts,
101    // then sort the keys once at the end to restore time order.
102    let mut buckets: FnvHashMap<i64, Vec<usize>> = FnvHashMap::default();
103    for i in 0..timestamps.len() {
104        if timestamps.is_null(i) {
105            continue;
106        }
107        let ts = timestamps.value(i);
108        // The bucket knows its own calendar: a fixed width aligns on the
109        // epoch, a `1d` in a zone on local midnight, a `1mo` on the first of
110        // the month. Doing the arithmetic here is what made this surface
111        // disagree with SQL's `time_bucket()`.
112        let bucket_start = bucket.start_of(ts);
113        buckets.entry(bucket_start).or_default().push(i);
114    }
115
116    // Sort bucket keys so the output is ordered by time.
117    let mut sorted_keys: Vec<i64> = buckets.keys().copied().collect();
118    sorted_keys.sort_unstable();
119
120    // Compute aggregate per bucket
121    let mut out_times = Vec::with_capacity(buckets.len());
122    let mut out_values = Vec::with_capacity(buckets.len());
123
124    for bucket_start in &sorted_keys {
125        let indices = &buckets[bucket_start];
126        // Collect (timestamp, value) pairs for this bucket.
127        // Timestamps are needed for First/Last semantics.
128        let raw_pairs: Vec<(i64, f64)> = indices
129            .iter()
130            .filter_map(|&i| {
131                let val = extract_f64(value_col, i)?;
132                Some((timestamps.value(i), val))
133            })
134            .collect();
135
136        if raw_pairs.is_empty() {
137            continue;
138        }
139
140        let agg_value = match function {
141            // Count counts all non-null, non-NaN rows for consistency
142            // with the other aggregation functions.
143            AggFn::Count => raw_pairs.iter().filter(|(_, v)| !v.is_nan()).count() as f64,
144            _ => {
145                // For numeric aggregations, filter out NaN
146                let clean: Vec<(i64, f64)> =
147                    raw_pairs.into_iter().filter(|(_, v)| !v.is_nan()).collect();
148                if clean.is_empty() {
149                    continue;
150                }
151                match function {
152                    AggFn::Sum => clean.iter().map(|(_, v)| v).sum(),
153                    AggFn::Min => clean.iter().map(|(_, v)| *v).fold(f64::INFINITY, f64::min),
154                    AggFn::Max => clean
155                        .iter()
156                        .map(|(_, v)| *v)
157                        .fold(f64::NEG_INFINITY, f64::max),
158                    AggFn::Avg => {
159                        let sum: f64 = clean.iter().map(|(_, v)| v).sum();
160                        sum / clean.len() as f64
161                    }
162                    // First: value with the smallest timestamp in the bucket
163                    AggFn::First => clean
164                        .iter()
165                        .min_by_key(|(ts, _)| *ts)
166                        .map(|(_, v)| *v)
167                        .expect("clean is non-empty"),
168                    // Last: value with the largest timestamp in the bucket
169                    AggFn::Last => clean
170                        .iter()
171                        .max_by_key(|(ts, _)| *ts)
172                        .map(|(_, v)| *v)
173                        .expect("clean is non-empty"),
174                    AggFn::Count => {
175                        // Count is handled in the outer branch; defensive error.
176                        return Err(QueryError::Validation(
177                            "Count aggregation reached inner branch unexpectedly".into(),
178                        ));
179                    }
180                }
181            }
182        };
183
184        out_times.push(*bucket_start);
185        out_values.push(agg_value);
186    }
187
188    let out_schema = Arc::new(Schema::new(vec![
189        Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
190        Field::new(value_column, DataType::Float64, false),
191    ]));
192
193    let result = RecordBatch::try_new(
194        out_schema,
195        vec![
196            Arc::new(Int64Array::from(out_times)) as ArrayRef,
197            Arc::new(Float64Array::from(out_values)) as ArrayRef,
198        ],
199    )?;
200
201    if let Some(tracker) = memory_tracker {
202        tracker.try_allocate(result.get_array_memory_size())?;
203    }
204
205    Ok(result)
206}
207
208/// Boundary-safe downsampling over a sequence of time-ordered batches.
209///
210/// Rows must arrive ascending by timestamp — which is what the query engine's
211/// merge guarantees — so a bucket is complete as soon as a row belonging to a
212/// later bucket is seen. Buckets are emitted in order and exactly once.
213///
214/// # Example
215///
216/// ```no_run
217/// # use chronix_query::downsample::StreamingDownsampler;
218/// # use chronix_query::aggregate::AggFn;
219/// # use arrow::record_batch::RecordBatch;
220/// # use chronix_core::TimeBucket;
221/// # use std::time::Duration;
222/// # fn batches() -> Vec<RecordBatch> { vec![] }
223/// // A fixed minute. `TimeBucket::parse("1d", Some("Europe/Berlin"))` for a
224/// // calendar bucket — the same type SQL's `time_bucket()` uses.
225/// let bucket = TimeBucket::fixed(Duration::from_secs(60));
226/// let mut ds = StreamingDownsampler::new("value", bucket, AggFn::Avg)?;
227/// let mut out = Vec::new();
228/// for batch in batches() {
229///     out.extend(ds.push(&batch)?);
230/// }
231/// out.extend(ds.finish()?);
232/// # Ok::<(), chronix_query::QueryError>(())
233/// ```
234pub struct StreamingDownsampler {
235    value_column: String,
236    bucket: chronix_core::TimeBucket,
237    function: AggFn,
238    /// Start timestamp of the bucket currently being accumulated.
239    open_bucket: Option<i64>,
240    /// Accumulated state for the open bucket.
241    acc: BucketAcc,
242    /// Completed buckets not yet drained into an output batch.
243    pending: Vec<(i64, f64)>,
244    /// Rows per emitted batch.
245    chunk_rows: usize,
246}
247
248/// Running state for one time bucket.
249#[derive(Default)]
250struct BucketAcc {
251    count: u64,
252    sum: f64,
253    min: f64,
254    max: f64,
255    first_ts: i64,
256    first_val: f64,
257    last_ts: i64,
258    last_val: f64,
259}
260
261impl BucketAcc {
262    fn reset(&mut self) {
263        self.count = 0;
264        self.sum = 0.0;
265        self.min = f64::INFINITY;
266        self.max = f64::NEG_INFINITY;
267        self.first_ts = i64::MAX;
268        self.first_val = 0.0;
269        self.last_ts = i64::MIN;
270        self.last_val = 0.0;
271    }
272
273    fn update(&mut self, ts: i64, v: f64) {
274        // NaN is excluded, matching `downsample`'s per-bucket filtering.
275        if v.is_nan() {
276            return;
277        }
278        self.count += 1;
279        self.sum += v;
280        if v < self.min {
281            self.min = v;
282        }
283        if v > self.max {
284            self.max = v;
285        }
286        if ts < self.first_ts {
287            self.first_ts = ts;
288            self.first_val = v;
289        }
290        // `>=` so the last-encountered value wins on ties.
291        if ts >= self.last_ts {
292            self.last_ts = ts;
293            self.last_val = v;
294        }
295    }
296
297    #[allow(clippy::cast_precision_loss)]
298    fn finalize(&self, function: AggFn) -> Option<f64> {
299        if self.count == 0 {
300            // `Count` still reports a bucket that had only NaN rows as absent,
301            // matching `downsample`, which skips such buckets entirely.
302            return None;
303        }
304        Some(match function {
305            AggFn::Count => self.count as f64,
306            AggFn::Sum => self.sum,
307            AggFn::Min => self.min,
308            AggFn::Max => self.max,
309            AggFn::Avg => self.sum / self.count as f64,
310            AggFn::First => self.first_val,
311            AggFn::Last => self.last_val,
312        })
313    }
314}
315
316impl StreamingDownsampler {
317    /// Rows per emitted batch.
318    const DEFAULT_CHUNK_ROWS: usize = 65_536;
319
320    /// Create a downsampler for `value_column` over `bucket`-wide buckets.
321    ///
322    /// # Errors
323    ///
324    /// Never fails today; the signature is fallible because `push` is, and
325    /// callers already handle a `Result` here.
326    pub fn new(
327        value_column: &str,
328        bucket: chronix_core::TimeBucket,
329        function: AggFn,
330    ) -> Result<Self> {
331        let mut acc = BucketAcc::default();
332        acc.reset();
333        Ok(Self {
334            value_column: value_column.to_string(),
335            bucket,
336            function,
337            open_bucket: None,
338            acc,
339            pending: Vec::new(),
340            chunk_rows: Self::DEFAULT_CHUNK_ROWS,
341        })
342    }
343
344    /// Fold one batch in, returning any batches of completed buckets.
345    ///
346    /// # Errors
347    ///
348    /// Returns an error if the timestamp or value column is missing or has an
349    /// unsupported type, or if a row arrives out of timestamp order.
350    pub fn push(&mut self, batch: &RecordBatch) -> Result<Vec<RecordBatch>> {
351        if batch.num_rows() == 0 {
352            return Ok(Vec::new());
353        }
354        let schema = batch.schema();
355        let time_idx = schema
356            .index_of(chronix_core::TIME_COLUMN)
357            .map_err(|_| QueryError::Validation("'timestamp' column not found".into()))?;
358        let val_idx =
359            schema
360                .index_of(&self.value_column)
361                .map_err(|_| QueryError::FieldNotFound {
362                    measurement: String::new(),
363                    field: self.value_column.clone(),
364                })?;
365
366        let timestamps = batch
367            .column(time_idx)
368            .as_any()
369            .downcast_ref::<Int64Array>()
370            .ok_or_else(|| QueryError::Validation("'timestamp' column is not Int64".into()))?;
371        let value_col = batch.column(val_idx).as_ref();
372        match value_col.data_type() {
373            DataType::Float64 | DataType::Int64 | DataType::UInt64 | DataType::Decimal128(_, _) => {
374            }
375            dt => {
376                return Err(QueryError::Validation(format!(
377                    "'{}' column has unsupported type {dt} \
378                     (expected Float64, Int64, UInt64, or Decimal128)",
379                    self.value_column
380                )));
381            }
382        }
383
384        for i in 0..timestamps.len() {
385            if timestamps.is_null(i) {
386                continue;
387            }
388            let ts = timestamps.value(i);
389            // See `downsample` above: the bucket owns its own calendar.
390            let bucket = self.bucket.start_of(ts);
391
392            match self.open_bucket {
393                Some(open) if bucket == open => {}
394                Some(open) if bucket < open => {
395                    // Out-of-order input would silently corrupt a closed
396                    // bucket, so say so instead of producing a wrong number.
397                    return Err(QueryError::Validation(format!(
398                        "downsample input is not ascending by timestamp: bucket {bucket} \
399                         arrived after {open} was closed"
400                    )));
401                }
402                Some(_) => self.close_open_bucket(),
403                None => {}
404            }
405            self.open_bucket = Some(bucket);
406
407            if let Some(v) = extract_f64(value_col, i) {
408                self.acc.update(ts, v);
409            }
410        }
411
412        self.drain(false)
413    }
414
415    /// Close the final bucket and return any remaining output.
416    ///
417    /// # Errors
418    ///
419    /// Returns an error if the output batch cannot be built.
420    pub fn finish(mut self) -> Result<Vec<RecordBatch>> {
421        if self.open_bucket.is_some() {
422            self.close_open_bucket();
423        }
424        self.drain(true)
425    }
426
427    fn close_open_bucket(&mut self) {
428        if let (Some(start), Some(v)) = (self.open_bucket, self.acc.finalize(self.function)) {
429            self.pending.push((start, v));
430        }
431        self.acc.reset();
432        self.open_bucket = None;
433    }
434
435    /// Turn completed buckets into output batches. Unless `flush`, a partial
436    /// chunk is kept back so emitted batches stay a uniform size.
437    fn drain(&mut self, flush: bool) -> Result<Vec<RecordBatch>> {
438        let mut out = Vec::new();
439        while self.pending.len() >= self.chunk_rows || (flush && !self.pending.is_empty()) {
440            let take = self.chunk_rows.min(self.pending.len());
441            let rows: Vec<(i64, f64)> = self.pending.drain(..take).collect();
442            // Propagate rather than skip: dropping a failed batch here would
443            // silently lose completed buckets.
444            out.push(build_batch(&self.value_column, &rows)?);
445        }
446        Ok(out)
447    }
448}
449
450/// Build the two-column output batch for a run of completed buckets.
451fn build_batch(value_column: &str, rows: &[(i64, f64)]) -> Result<RecordBatch> {
452    let schema = Arc::new(Schema::new(vec![
453        Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
454        Field::new(value_column, DataType::Float64, false),
455    ]));
456    let times: Vec<i64> = rows.iter().map(|(t, _)| *t).collect();
457    let values: Vec<f64> = rows.iter().map(|(_, v)| *v).collect();
458    RecordBatch::try_new(
459        schema,
460        vec![
461            Arc::new(Int64Array::from(times)) as ArrayRef,
462            Arc::new(Float64Array::from(values)) as ArrayRef,
463        ],
464    )
465    .map_err(Into::into)
466}
467
468#[cfg(test)]
469mod tests {
470    use super::*;
471    use arrow::array::UInt64Array;
472    use chronix_core::TimeBucket;
473
474    fn make_batch(times: Vec<i64>, values: Vec<f64>) -> RecordBatch {
475        let schema = Arc::new(Schema::new(vec![
476            Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
477            Field::new("value", DataType::Float64, false),
478        ]));
479
480        RecordBatch::try_new(
481            schema,
482            vec![
483                Arc::new(Int64Array::from(times)),
484                Arc::new(Float64Array::from(values)),
485            ],
486        )
487        .unwrap()
488    }
489
490    #[test]
491    fn downsample_avg_10s_buckets() {
492        let batch = make_batch(
493            vec![0, 5, 10, 15, 20, 25],
494            vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
495        );
496
497        let result = downsample(
498            &batch,
499            "value",
500            &TimeBucket::fixed_ns(10),
501            &AggFn::Avg,
502            None,
503        )
504        .unwrap();
505        assert_eq!(result.num_rows(), 3);
506
507        let times = result
508            .column(0)
509            .as_any()
510            .downcast_ref::<Int64Array>()
511            .unwrap();
512        let values = result
513            .column(1)
514            .as_any()
515            .downcast_ref::<Float64Array>()
516            .unwrap();
517
518        // Bucket [0, 10): ts=0, 5 → avg 1.5
519        assert_eq!(times.value(0), 0);
520        assert!((values.value(0) - 1.5).abs() < f64::EPSILON);
521
522        // Bucket [10, 20): ts=10, 15 → avg 3.5
523        assert_eq!(times.value(1), 10);
524        assert!((values.value(1) - 3.5).abs() < f64::EPSILON);
525
526        // Bucket [20, 30): ts=20, 25 → avg 5.5
527        assert_eq!(times.value(2), 20);
528        assert!((values.value(2) - 5.5).abs() < f64::EPSILON);
529    }
530
531    #[test]
532    fn downsample_sum() {
533        let batch = make_batch(vec![0, 5, 10, 15], vec![10.0, 20.0, 30.0, 40.0]);
534
535        let result = downsample(
536            &batch,
537            "value",
538            &TimeBucket::fixed_ns(10),
539            &AggFn::Sum,
540            None,
541        )
542        .unwrap();
543        assert_eq!(result.num_rows(), 2);
544
545        let values = result
546            .column(1)
547            .as_any()
548            .downcast_ref::<Float64Array>()
549            .unwrap();
550
551        assert!((values.value(0) - 30.0).abs() < f64::EPSILON);
552        assert!((values.value(1) - 70.0).abs() < f64::EPSILON);
553    }
554
555    #[test]
556    fn downsample_count() {
557        let batch = make_batch(vec![0, 1, 2, 10, 20, 21, 22, 23], vec![1.0; 8]);
558
559        let result = downsample(
560            &batch,
561            "value",
562            &TimeBucket::fixed_ns(10),
563            &AggFn::Count,
564            None,
565        )
566        .unwrap();
567        assert_eq!(result.num_rows(), 3);
568
569        let values = result
570            .column(1)
571            .as_any()
572            .downcast_ref::<Float64Array>()
573            .unwrap();
574
575        assert!((values.value(0) - 3.0).abs() < f64::EPSILON); // [0,10): 3 points
576        assert!((values.value(1) - 1.0).abs() < f64::EPSILON); // [10,20): 1 point
577        assert!((values.value(2) - 4.0).abs() < f64::EPSILON); // [20,30): 4 points
578    }
579
580    #[test]
581    fn downsample_min_max() {
582        let batch = make_batch(vec![0, 5, 10, 15], vec![3.0, 1.0, 4.0, 2.0]);
583
584        let min_result = downsample(
585            &batch,
586            "value",
587            &TimeBucket::fixed_ns(10),
588            &AggFn::Min,
589            None,
590        )
591        .unwrap();
592        let max_result = downsample(
593            &batch,
594            "value",
595            &TimeBucket::fixed_ns(10),
596            &AggFn::Max,
597            None,
598        )
599        .unwrap();
600
601        let min_vals = min_result
602            .column(1)
603            .as_any()
604            .downcast_ref::<Float64Array>()
605            .unwrap();
606        let max_vals = max_result
607            .column(1)
608            .as_any()
609            .downcast_ref::<Float64Array>()
610            .unwrap();
611
612        assert!((min_vals.value(0) - 1.0).abs() < f64::EPSILON); // min(3, 1)
613        assert!((min_vals.value(1) - 2.0).abs() < f64::EPSILON); // min(4, 2)
614        assert!((max_vals.value(0) - 3.0).abs() < f64::EPSILON); // max(3, 1)
615        assert!((max_vals.value(1) - 4.0).abs() < f64::EPSILON); // max(4, 2)
616    }
617
618    #[test]
619    fn downsample_first_last() {
620        let batch = make_batch(vec![0, 5, 10, 15], vec![10.0, 20.0, 30.0, 40.0]);
621
622        let first = downsample(
623            &batch,
624            "value",
625            &TimeBucket::fixed_ns(10),
626            &AggFn::First,
627            None,
628        )
629        .unwrap();
630        let last = downsample(
631            &batch,
632            "value",
633            &TimeBucket::fixed_ns(10),
634            &AggFn::Last,
635            None,
636        )
637        .unwrap();
638
639        let first_vals = first
640            .column(1)
641            .as_any()
642            .downcast_ref::<Float64Array>()
643            .unwrap();
644        let last_vals = last
645            .column(1)
646            .as_any()
647            .downcast_ref::<Float64Array>()
648            .unwrap();
649
650        assert!((first_vals.value(0) - 10.0).abs() < f64::EPSILON);
651        assert!((first_vals.value(1) - 30.0).abs() < f64::EPSILON);
652        assert!((last_vals.value(0) - 20.0).abs() < f64::EPSILON);
653        assert!((last_vals.value(1) - 40.0).abs() < f64::EPSILON);
654    }
655
656    #[test]
657    fn empty_batch() {
658        let batch = make_batch(vec![], vec![]);
659        let result = downsample(
660            &batch,
661            "value",
662            &TimeBucket::fixed_ns(10),
663            &AggFn::Avg,
664            None,
665        )
666        .unwrap();
667        assert_eq!(result.num_rows(), 0);
668    }
669
670    #[test]
671    fn downsample_i64_column() {
672        let schema = Arc::new(Schema::new(vec![
673            Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
674            Field::new("value", DataType::Int64, false),
675        ]));
676        let batch = RecordBatch::try_new(
677            schema,
678            vec![
679                Arc::new(Int64Array::from(vec![0, 5, 10, 15])) as ArrayRef,
680                Arc::new(Int64Array::from(vec![10, 20, 30, 40])) as ArrayRef,
681            ],
682        )
683        .unwrap();
684
685        let result = downsample(
686            &batch,
687            "value",
688            &TimeBucket::fixed_ns(10),
689            &AggFn::Sum,
690            None,
691        )
692        .unwrap();
693        assert_eq!(result.num_rows(), 2);
694
695        let values = result
696            .column(1)
697            .as_any()
698            .downcast_ref::<Float64Array>()
699            .unwrap();
700        assert!((values.value(0) - 30.0).abs() < f64::EPSILON);
701        assert!((values.value(1) - 70.0).abs() < f64::EPSILON);
702    }
703
704    #[test]
705    fn downsample_u64_column() {
706        let schema = Arc::new(Schema::new(vec![
707            Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
708            Field::new("value", DataType::UInt64, false),
709        ]));
710        let batch = RecordBatch::try_new(
711            schema,
712            vec![
713                Arc::new(Int64Array::from(vec![0, 5, 10, 15])) as ArrayRef,
714                Arc::new(UInt64Array::from(vec![100u64, 200, 300, 400])) as ArrayRef,
715            ],
716        )
717        .unwrap();
718
719        let result = downsample(
720            &batch,
721            "value",
722            &TimeBucket::fixed_ns(10),
723            &AggFn::Avg,
724            None,
725        )
726        .unwrap();
727        assert_eq!(result.num_rows(), 2);
728
729        let values = result
730            .column(1)
731            .as_any()
732            .downcast_ref::<Float64Array>()
733            .unwrap();
734        assert!((values.value(0) - 150.0).abs() < f64::EPSILON);
735        assert!((values.value(1) - 350.0).abs() < f64::EPSILON);
736    }
737
738    #[test]
739    fn downsample_unsupported_type_fails() {
740        let schema = Arc::new(Schema::new(vec![
741            Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
742            Field::new("value", DataType::Boolean, false),
743        ]));
744        let batch = RecordBatch::try_new(
745            schema,
746            vec![
747                Arc::new(Int64Array::from(vec![0i64])) as ArrayRef,
748                Arc::new(arrow::array::BooleanArray::from(vec![true])) as ArrayRef,
749            ],
750        )
751        .unwrap();
752
753        let err = downsample(
754            &batch,
755            "value",
756            &TimeBucket::fixed_ns(10),
757            &AggFn::Avg,
758            None,
759        )
760        .unwrap_err();
761        let msg = format!("{err}");
762        assert!(
763            msg.contains("unsupported"),
764            "expected unsupported type error, got: {msg}"
765        );
766    }
767
768    #[test]
769    fn downsample_respects_memory_tracker() {
770        let batch = make_batch(vec![0, 5, 10, 15], vec![1.0, 2.0, 3.0, 4.0]);
771        let tracker = crate::memory::MemoryTracker::new(1_000_000);
772        let result = downsample(
773            &batch,
774            "value",
775            &TimeBucket::fixed_ns(10),
776            &AggFn::Avg,
777            Some(&tracker),
778        )
779        .unwrap();
780        assert!(tracker.allocated() > 0);
781        assert_eq!(result.num_rows(), 2);
782    }
783
784    #[test]
785    fn downsample_exceeds_memory_budget() {
786        let batch = make_batch(vec![0, 5, 10, 15], vec![1.0, 2.0, 3.0, 4.0]);
787        let tracker = crate::memory::MemoryTracker::new(1);
788        let err = downsample(
789            &batch,
790            "value",
791            &TimeBucket::fixed_ns(10),
792            &AggFn::Avg,
793            Some(&tracker),
794        )
795        .unwrap_err();
796        assert!(matches!(
797            err,
798            crate::error::QueryError::QueryMemoryExceeded { .. }
799        ));
800    }
801}