1use 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#[allow(clippy::cast_precision_loss)]
37fn extract_f64(col: &dyn Array, index: usize) -> Option<f64> {
38 crate::extract_f64(col, index)
39}
40
41#[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 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 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 let bucket_start = bucket.start_of(ts);
113 buckets.entry(bucket_start).or_default().push(i);
114 }
115
116 let mut sorted_keys: Vec<i64> = buckets.keys().copied().collect();
118 sorted_keys.sort_unstable();
119
120 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 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 AggFn::Count => raw_pairs.iter().filter(|(_, v)| !v.is_nan()).count() as f64,
144 _ => {
145 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 AggFn::First => clean
164 .iter()
165 .min_by_key(|(ts, _)| *ts)
166 .map(|(_, v)| *v)
167 .expect("clean is non-empty"),
168 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 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
208pub struct StreamingDownsampler {
235 value_column: String,
236 bucket: chronix_core::TimeBucket,
237 function: AggFn,
238 open_bucket: Option<i64>,
240 acc: BucketAcc,
242 pending: Vec<(i64, f64)>,
244 chunk_rows: usize,
246}
247
248#[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 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 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 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 const DEFAULT_CHUNK_ROWS: usize = 65_536;
319
320 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 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 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 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 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 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 out.push(build_batch(&self.value_column, &rows)?);
445 }
446 Ok(out)
447 }
448}
449
450fn 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 assert_eq!(times.value(0), 0);
520 assert!((values.value(0) - 1.5).abs() < f64::EPSILON);
521
522 assert_eq!(times.value(1), 10);
524 assert!((values.value(1) - 3.5).abs() < f64::EPSILON);
525
526 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); assert!((values.value(1) - 1.0).abs() < f64::EPSILON); assert!((values.value(2) - 4.0).abs() < f64::EPSILON); }
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); assert!((min_vals.value(1) - 2.0).abs() < f64::EPSILON); assert!((max_vals.value(0) - 3.0).abs() < f64::EPSILON); assert!((max_vals.value(1) - 4.0).abs() < f64::EPSILON); }
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}