1use std::sync::Arc;
24
25use arrow::array::{Array, ArrayRef, Float64Array, Int64Array};
26use arrow::datatypes::{DataType, Field, Schema};
27use arrow::record_batch::RecordBatch;
28
29use crate::error::{QueryError, Result};
30
31#[derive(Debug, Clone, PartialEq)]
33pub enum WindowFn {
34 RowNumber,
36 Rank {
42 order_column: Option<String>,
44 },
45 DenseRank {
49 order_column: Option<String>,
51 },
52 Lead {
54 offset: usize,
56 default: f64,
58 },
59 Lag {
61 offset: usize,
63 default: f64,
65 },
66 Delta,
68 Rate,
78 IRate,
83 MovingAverage {
85 window_size: usize,
87 },
88 CumulativeSum,
90 Increase,
92 RangeMovingAverage {
98 range_ns: i64,
100 },
101 RangeSum {
103 range_ns: i64,
105 },
106 RangeCount {
108 range_ns: i64,
110 },
111}
112
113#[must_use]
115pub fn window_fn_name(func: &WindowFn) -> String {
116 match func {
117 WindowFn::RowNumber => "row_number".to_string(),
118 WindowFn::Rank { .. } => "rank".to_string(),
119 WindowFn::DenseRank { .. } => "dense_rank".to_string(),
120 WindowFn::Lead { offset, .. } => format!("lead_{offset}"),
121 WindowFn::Lag { offset, .. } => format!("lag_{offset}"),
122 WindowFn::Delta => "delta".to_string(),
123 WindowFn::Rate => "rate".to_string(),
124 WindowFn::IRate => "irate".to_string(),
125 WindowFn::MovingAverage { window_size } => format!("moving_avg_{window_size}"),
126 WindowFn::CumulativeSum => "cumulative_sum".to_string(),
127 WindowFn::Increase => "increase".to_string(),
128 WindowFn::RangeMovingAverage { range_ns } => format!("range_avg_{range_ns}ns"),
129 WindowFn::RangeSum { range_ns } => format!("range_sum_{range_ns}ns"),
130 WindowFn::RangeCount { range_ns } => format!("range_count_{range_ns}ns"),
131 }
132}
133
134#[allow(clippy::cast_precision_loss)]
136fn extract_f64(col: &dyn Array, index: usize) -> Option<f64> {
137 crate::extract_f64(col, index)
138}
139
140fn f64_opt_eq(a: Option<f64>, b: Option<f64>) -> bool {
142 match (a, b) {
143 (Some(x), Some(y)) => x.to_bits() == y.to_bits(),
144 (None, None) => true,
145 _ => false,
146 }
147}
148
149pub fn apply_window(
163 batch: &RecordBatch,
164 functions: &[WindowFn],
165 partition_by: &[String],
166 value_column: &str,
167 memory_tracker: Option<&crate::memory::MemoryTracker>,
168) -> Result<RecordBatch> {
169 if functions.is_empty() {
170 return Ok(batch.clone());
171 }
172
173 let schema = batch.schema();
174 let n = batch.num_rows();
175
176 if let Some(tracker) = memory_tracker {
179 let estimated_bytes = n * functions.len() * 8 + n * 16; tracker.try_allocate(estimated_bytes)?;
181 }
182
183 let val_idx = schema
185 .index_of(value_column)
186 .map_err(|_| QueryError::FieldNotFound {
187 measurement: String::new(),
188 field: value_column.to_string(),
189 })?;
190 let val_col = batch.column(val_idx).as_ref();
191
192 let ts_col = batch
194 .column_by_name(chronix_core::TIME_COLUMN)
195 .and_then(|c| c.as_any().downcast_ref::<Int64Array>().cloned());
196
197 let needs_ts = functions.iter().any(|f| {
207 matches!(
208 f,
209 WindowFn::Rate
210 | WindowFn::IRate
211 | WindowFn::RangeMovingAverage { .. }
212 | WindowFn::RangeSum { .. }
213 | WindowFn::RangeCount { .. }
214 )
215 });
216 if needs_ts
217 && let Some(ref ts) = ts_col
218 && !ts.is_empty()
219 {
220 let max_ts = (0..ts.len()).map(|i| ts.value(i)).max().unwrap_or(0);
222 if (1_000_000_000..100_000_000_000).contains(&max_ts) {
227 tracing::warn!(
228 max_ts,
229 "timestamps may be in seconds instead of nanoseconds; \
230 Rate/IRate assume nanosecond timestamps"
231 );
232 }
233 }
234
235 let partitions = build_partitions(batch, partition_by)?;
237
238 let mut new_fields: Vec<Field> = Vec::with_capacity(functions.len());
240 let mut new_columns: Vec<ArrayRef> = Vec::with_capacity(functions.len());
241
242 for func in functions {
243 let col_name = format!("{value_column}_{}", window_fn_name(func));
244 new_fields.push(Field::new(&col_name, DataType::Float64, true));
245
246 let mut result = vec![None; n];
247
248 for partition in &partitions {
249 compute_window_fn(
250 func,
251 partition,
252 val_col,
253 ts_col.as_ref(),
254 batch,
255 &mut result,
256 )?;
257 }
258
259 let arr: Float64Array = result.into_iter().collect();
260 new_columns.push(Arc::new(arr));
261 }
262
263 let mut out_fields: Vec<Field> = schema.fields().iter().map(|f| f.as_ref().clone()).collect();
265 out_fields.extend(new_fields);
266
267 let mut out_columns: Vec<ArrayRef> = (0..batch.num_columns())
268 .map(|i| batch.column(i).clone())
269 .collect();
270 out_columns.extend(new_columns);
271
272 let out_schema = Arc::new(Schema::new(out_fields));
273 let out_batch = RecordBatch::try_new(out_schema, out_columns)?;
274 if let Some(tracker) = memory_tracker {
275 tracker.try_allocate(out_batch.get_array_memory_size())?;
276 }
277 Ok(out_batch)
278}
279
280fn build_partitions(batch: &RecordBatch, partition_by: &[String]) -> Result<Vec<Vec<usize>>> {
284 let n = batch.num_rows();
285
286 if partition_by.is_empty() {
287 return Ok(vec![(0..n).collect()]);
288 }
289
290 use arrow::array::StringArray;
291 use std::collections::HashMap;
292
293 let part_cols: Vec<&StringArray> = partition_by
294 .iter()
295 .map(|name| {
296 batch
297 .column_by_name(name)
298 .and_then(|c| c.as_any().downcast_ref::<StringArray>())
299 .ok_or_else(|| {
300 QueryError::Validation(format!(
301 "partition-by column '{name}' not found or not string"
302 ))
303 })
304 })
305 .collect::<Result<Vec<_>>>()?;
306
307 let mut groups: Vec<(Vec<Option<&str>>, Vec<usize>)> = Vec::new();
309 let mut index: HashMap<u64, Vec<usize>> = HashMap::new();
310
311 for row in 0..n {
312 let mut hasher = std::hash::DefaultHasher::new();
313 for col in &part_cols {
314 if col.is_null(row) {
315 std::hash::Hash::hash(&0u8, &mut hasher);
316 } else {
317 std::hash::Hash::hash(&1u8, &mut hasher);
318 std::hash::Hash::hash(col.value(row), &mut hasher);
319 }
320 }
321 let hash = std::hash::Hasher::finish(&hasher);
322
323 let key: Vec<Option<&str>> = part_cols
324 .iter()
325 .map(|col| {
326 if col.is_null(row) {
327 None
328 } else {
329 Some(col.value(row))
330 }
331 })
332 .collect();
333
334 let bucket = index.entry(hash).or_default();
335 let mut found = false;
336 for &gidx in bucket.iter() {
337 if groups[gidx].0 == key {
338 groups[gidx].1.push(row);
339 found = true;
340 break;
341 }
342 }
343 if !found {
344 let gidx = groups.len();
345 groups.push((key, vec![row]));
346 bucket.push(gidx);
347 }
348 }
349
350 Ok(groups.into_iter().map(|(_, rows)| rows).collect())
351}
352
353#[allow(clippy::cast_precision_loss)]
363fn compute_window_fn(
364 func: &WindowFn,
365 partition: &[usize],
366 val_col: &dyn Array,
367 ts_col: Option<&Int64Array>,
368 batch: &RecordBatch,
369 output: &mut [Option<f64>],
370) -> std::result::Result<(), QueryError> {
371 let len = partition.len();
372 if len == 0 {
373 return Ok(());
374 }
375
376 let needs_values = !matches!(
380 func,
381 WindowFn::RowNumber | WindowFn::Rank { .. } | WindowFn::DenseRank { .. }
382 );
383 let values: Vec<Option<f64>> = if needs_values {
384 partition.iter().map(|&r| extract_f64(val_col, r)).collect()
385 } else {
386 Vec::new()
387 };
388
389 let timestamps: Vec<i64> = if matches!(
391 func,
392 WindowFn::Rate
393 | WindowFn::IRate
394 | WindowFn::RangeMovingAverage { .. }
395 | WindowFn::RangeSum { .. }
396 | WindowFn::RangeCount { .. }
397 ) {
398 if let Some(ts) = ts_col {
399 partition.iter().map(|&r| ts.value(r)).collect()
400 } else {
401 Vec::new()
402 }
403 } else {
404 Vec::new()
405 };
406
407 match func {
408 WindowFn::RowNumber => {
409 for (i, &row) in partition.iter().enumerate() {
410 output[row] = Some((i + 1) as f64);
411 }
412 }
413
414 WindowFn::Rank { order_column } => {
415 let order_col: &dyn Array = order_column
418 .as_ref()
419 .and_then(|name| batch.column_by_name(name))
420 .map(std::convert::AsRef::as_ref)
421 .unwrap_or(val_col);
422 let order_values: Vec<Option<f64>> = partition
423 .iter()
424 .map(|&r| extract_f64(order_col, r))
425 .collect();
426 let mut rank = 1usize;
427 output[partition[0]] = Some(1.0);
428 for i in 1..len {
429 if !f64_opt_eq(order_values[i], order_values[i - 1]) {
430 rank = i + 1;
431 }
432 output[partition[i]] = Some(rank as f64);
433 }
434 }
435
436 WindowFn::DenseRank { order_column } => {
437 let order_col: &dyn Array = order_column
440 .as_ref()
441 .and_then(|name| batch.column_by_name(name))
442 .map(std::convert::AsRef::as_ref)
443 .unwrap_or(val_col);
444 let order_values: Vec<Option<f64>> = partition
445 .iter()
446 .map(|&r| extract_f64(order_col, r))
447 .collect();
448 let mut rank = 1usize;
449 output[partition[0]] = Some(1.0);
450 for i in 1..len {
451 if !f64_opt_eq(order_values[i], order_values[i - 1]) {
452 rank += 1;
453 }
454 output[partition[i]] = Some(rank as f64);
455 }
456 }
457
458 WindowFn::Lead { offset, default } => {
459 for (i, &row) in partition.iter().enumerate() {
460 let target = i + offset;
461 output[row] = if target < len {
462 values[target].or(Some(*default))
463 } else {
464 Some(*default)
465 };
466 }
467 }
468
469 WindowFn::Lag { offset, default } => {
470 for (i, &row) in partition.iter().enumerate() {
471 output[row] = if i >= *offset {
472 values[i - offset].or(Some(*default))
473 } else {
474 Some(*default)
475 };
476 }
477 }
478
479 WindowFn::Delta => {
480 output[partition[0]] = None; for i in 1..len {
482 output[partition[i]] = match (values[i], values[i - 1]) {
483 (Some(c), Some(p)) => Some(c - p),
484 _ => None,
485 };
486 }
487 }
488
489 WindowFn::Rate => {
490 output[partition[0]] = None;
493 if !timestamps.is_empty() {
494 for i in 1..len {
495 let dt = (timestamps[i] - timestamps[i - 1]) as f64 / 1_000_000_000.0;
496 output[partition[i]] = match (values[i], values[i - 1]) {
497 (Some(c), Some(p)) if dt > 0.0 => Some((c - p) / dt),
498 _ => None,
499 };
500 }
501 }
502 }
503
504 WindowFn::IRate => {
505 let irate = if len >= 2 && !timestamps.is_empty() {
507 let dt = (timestamps[len - 1] - timestamps[len - 2]) as f64 / 1_000_000_000.0;
508 match (values[len - 1], values[len - 2]) {
509 (Some(vl), Some(vp)) if dt > 0.0 => Some((vl - vp) / dt),
510 _ => None,
511 }
512 } else {
513 None
514 };
515 for &row in partition {
516 output[row] = irate;
517 }
518 }
519
520 WindowFn::MovingAverage { window_size } => {
521 let ws = *window_size;
522 let mut sum = 0.0;
524 let mut count = 0usize;
525 for i in 0..len {
526 if let Some(val) = values[i] {
528 sum += val;
529 count += 1;
530 }
531 if i >= ws
533 && let Some(val) = values[i - ws]
534 {
535 sum -= val;
536 count -= 1;
537 }
538 if i + 1 < ws {
539 output[partition[i]] = None;
541 } else {
542 output[partition[i]] = if count > 0 {
543 Some(sum / count as f64)
544 } else {
545 None
546 };
547 }
548 }
549 }
550
551 WindowFn::CumulativeSum => {
552 let mut running = 0.0;
553 for (i, &row) in partition.iter().enumerate() {
554 if let Some(v) = values[i] {
555 running += v;
556 output[row] = Some(running);
557 } else {
558 output[row] = None;
559 }
560 }
561 }
562
563 WindowFn::Increase => {
564 let mut total = 0.0;
567 output[partition[0]] = Some(0.0);
568 for i in 1..len {
569 match (values[i], values[i - 1]) {
570 (Some(c), Some(p)) => {
571 if c >= p {
572 total += c - p;
573 }
574 output[partition[i]] = Some(total);
576 }
577 _ => {
578 output[partition[i]] = None;
579 }
580 }
581 }
582 }
583
584 WindowFn::RangeMovingAverage { range_ns } => {
591 if timestamps.is_empty() {
592 return Err(QueryError::Validation(
593 "timestamp column required for RangeMovingAverage".into(),
594 ));
595 }
596 let range = *range_ns;
597
598 let mut start = 0usize;
599 let mut sum = 0.0;
600 let mut count = 0usize;
601
602 for i in 0..len {
603 if let Some(v) = values[i] {
604 sum += v;
605 count += 1;
606 }
607 while start < i && timestamps[i] - timestamps[start] > range {
608 if let Some(v) = values[start] {
609 sum -= v;
610 count -= 1;
611 }
612 start += 1;
613 }
614 output[partition[i]] = if count > 0 {
615 Some(sum / count as f64)
616 } else {
617 None
618 };
619 }
620 }
621
622 WindowFn::RangeSum { range_ns } => {
623 if timestamps.is_empty() {
624 return Err(QueryError::Validation(
625 "timestamp column required for RangeSum".into(),
626 ));
627 }
628 let range = *range_ns;
629
630 let mut start = 0usize;
631 let mut sum = 0.0;
632 let mut count = 0usize;
633
634 for i in 0..len {
635 if let Some(v) = values[i] {
636 sum += v;
637 count += 1;
638 }
639 while start < i && timestamps[i] - timestamps[start] > range {
640 if let Some(v) = values[start] {
641 sum -= v;
642 count -= 1;
643 }
644 start += 1;
645 }
646 output[partition[i]] = if count > 0 { Some(sum) } else { None };
647 }
648 }
649
650 WindowFn::RangeCount { range_ns } => {
651 if timestamps.is_empty() {
652 return Err(QueryError::Validation(
653 "timestamp column required for RangeCount".into(),
654 ));
655 }
656 let range = *range_ns;
657
658 let mut start = 0usize;
659 let mut count = 0usize;
660
661 for i in 0..len {
662 if values[i].is_some() {
663 count += 1;
664 }
665 while start < i && timestamps[i] - timestamps[start] > range {
666 if values[start].is_some() {
667 count -= 1;
668 }
669 start += 1;
670 }
671 output[partition[i]] = if count > 0 { Some(count as f64) } else { None };
672 }
673 }
674 }
675
676 Ok(())
677}
678
679#[cfg(test)]
680mod tests {
681 use super::*;
682 use arrow::array::StringArray;
683 use arrow::array::UInt64Array;
684
685 fn test_batch() -> RecordBatch {
686 let schema = Arc::new(Schema::new(vec![
687 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
688 Field::new("host", DataType::Utf8, false),
689 Field::new("value", DataType::Float64, false),
690 ]));
691
692 let timestamps = Arc::new(Int64Array::from(vec![
693 1_000_000_000,
694 2_000_000_000,
695 3_000_000_000,
696 4_000_000_000,
697 5_000_000_000,
698 ])); let hosts = Arc::new(StringArray::from(vec!["a", "a", "a", "a", "a"]));
700 let values = Arc::new(Float64Array::from(vec![10.0, 20.0, 15.0, 25.0, 30.0]));
701
702 RecordBatch::try_new(schema, vec![timestamps, hosts, values]).unwrap()
703 }
704
705 fn partitioned_batch() -> RecordBatch {
706 let schema = Arc::new(Schema::new(vec![
707 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
708 Field::new("host", DataType::Utf8, false),
709 Field::new("value", DataType::Float64, false),
710 ]));
711
712 let timestamps = Arc::new(Int64Array::from(vec![
713 1_000_000_000,
714 1_000_000_000,
715 2_000_000_000,
716 2_000_000_000,
717 3_000_000_000,
718 3_000_000_000,
719 ]));
720 let hosts = Arc::new(StringArray::from(vec!["a", "b", "a", "b", "a", "b"]));
721 let values = Arc::new(Float64Array::from(vec![
722 10.0, 100.0, 20.0, 200.0, 30.0, 300.0,
723 ]));
724
725 RecordBatch::try_new(schema, vec![timestamps, hosts, values]).unwrap()
726 }
727
728 #[test]
729 fn row_number() {
730 let batch = test_batch();
731 let result = apply_window(&batch, &[WindowFn::RowNumber], &[], "value", None).unwrap();
732
733 assert_eq!(result.num_columns(), 4); assert_eq!(result.num_rows(), 5);
736
737 let rn = result
738 .column_by_name("value_row_number")
739 .unwrap()
740 .as_any()
741 .downcast_ref::<Float64Array>()
742 .unwrap();
743 for i in 0..5 {
744 assert!((rn.value(i) - (i + 1) as f64).abs() < f64::EPSILON);
745 }
746 }
747
748 #[test]
749 fn rank_and_dense_rank() {
750 let schema = Arc::new(Schema::new(vec![
752 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
753 Field::new("value", DataType::Float64, false),
754 ]));
755 let ts = Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5]));
756 let vals = Arc::new(Float64Array::from(vec![10.0, 20.0, 20.0, 25.0, 30.0]));
757 let batch = RecordBatch::try_new(schema, vec![ts, vals]).unwrap();
758
759 let result = apply_window(
760 &batch,
761 &[
762 WindowFn::Rank { order_column: None },
763 WindowFn::DenseRank { order_column: None },
764 ],
765 &[],
766 "value",
767 None,
768 )
769 .unwrap();
770
771 let rank = result
772 .column_by_name("value_rank")
773 .unwrap()
774 .as_any()
775 .downcast_ref::<Float64Array>()
776 .unwrap();
777 assert_eq!(rank.value(0), 1.0);
778 assert_eq!(rank.value(1), 2.0);
779 assert_eq!(rank.value(2), 2.0);
780 assert_eq!(rank.value(3), 4.0);
781 assert_eq!(rank.value(4), 5.0);
782
783 let drank = result
784 .column_by_name("value_dense_rank")
785 .unwrap()
786 .as_any()
787 .downcast_ref::<Float64Array>()
788 .unwrap();
789 assert_eq!(drank.value(0), 1.0);
790 assert_eq!(drank.value(1), 2.0);
791 assert_eq!(drank.value(2), 2.0);
792 assert_eq!(drank.value(3), 3.0);
793 assert_eq!(drank.value(4), 4.0);
794 }
795
796 #[test]
797 fn lead_and_lag() {
798 let batch = test_batch(); let result = apply_window(
800 &batch,
801 &[
802 WindowFn::Lead {
803 offset: 1,
804 default: -1.0,
805 },
806 WindowFn::Lag {
807 offset: 1,
808 default: -1.0,
809 },
810 ],
811 &[],
812 "value",
813 None,
814 )
815 .unwrap();
816
817 let lead = result
818 .column_by_name("value_lead_1")
819 .unwrap()
820 .as_any()
821 .downcast_ref::<Float64Array>()
822 .unwrap();
823 assert_eq!(lead.value(0), 20.0);
824 assert_eq!(lead.value(3), 30.0);
825 assert_eq!(lead.value(4), -1.0); let lag = result
828 .column_by_name("value_lag_1")
829 .unwrap()
830 .as_any()
831 .downcast_ref::<Float64Array>()
832 .unwrap();
833 assert_eq!(lag.value(0), -1.0); assert_eq!(lag.value(1), 10.0);
835 assert_eq!(lag.value(4), 25.0);
836 }
837
838 #[test]
839 fn lead_offset_2() {
840 let batch = test_batch(); let result = apply_window(
842 &batch,
843 &[WindowFn::Lead {
844 offset: 2,
845 default: 0.0,
846 }],
847 &[],
848 "value",
849 None,
850 )
851 .unwrap();
852
853 let lead = result
854 .column_by_name("value_lead_2")
855 .unwrap()
856 .as_any()
857 .downcast_ref::<Float64Array>()
858 .unwrap();
859 assert_eq!(lead.value(0), 15.0);
860 assert_eq!(lead.value(2), 30.0);
861 assert_eq!(lead.value(3), 0.0); assert_eq!(lead.value(4), 0.0); }
864
865 #[test]
866 fn delta() {
867 let batch = test_batch(); let result = apply_window(&batch, &[WindowFn::Delta], &[], "value", None).unwrap();
869
870 let delta = result
871 .column_by_name("value_delta")
872 .unwrap()
873 .as_any()
874 .downcast_ref::<Float64Array>()
875 .unwrap();
876 assert!(delta.is_null(0)); assert_eq!(delta.value(1), 10.0); assert_eq!(delta.value(2), -5.0); assert_eq!(delta.value(3), 10.0); assert_eq!(delta.value(4), 5.0); }
882
883 #[test]
884 fn rate() {
885 let batch = test_batch(); let result = apply_window(&batch, &[WindowFn::Rate], &[], "value", None).unwrap();
887
888 let rate = result
889 .column_by_name("value_rate")
890 .unwrap()
891 .as_any()
892 .downcast_ref::<Float64Array>()
893 .unwrap();
894 assert!(rate.is_null(0));
895 assert!((rate.value(1) - 10.0).abs() < f64::EPSILON);
897 assert!((rate.value(2) - (-5.0)).abs() < f64::EPSILON);
899 }
900
901 #[test]
902 fn irate() {
903 let batch = test_batch(); let result = apply_window(&batch, &[WindowFn::IRate], &[], "value", None).unwrap();
905
906 let irate = result
907 .column_by_name("value_irate")
908 .unwrap()
909 .as_any()
910 .downcast_ref::<Float64Array>()
911 .unwrap();
912 for i in 0..5 {
914 assert!((irate.value(i) - 5.0).abs() < f64::EPSILON);
915 }
916 }
917
918 #[test]
919 fn moving_average() {
920 let batch = test_batch(); let result = apply_window(
922 &batch,
923 &[WindowFn::MovingAverage { window_size: 3 }],
924 &[],
925 "value",
926 None,
927 )
928 .unwrap();
929
930 let ma = result
931 .column_by_name("value_moving_avg_3")
932 .unwrap()
933 .as_any()
934 .downcast_ref::<Float64Array>()
935 .unwrap();
936 assert!(ma.is_null(0)); assert!(ma.is_null(1));
938 assert!((ma.value(2) - 15.0).abs() < f64::EPSILON);
940 assert!((ma.value(3) - 20.0).abs() < f64::EPSILON);
942 assert!((ma.value(4) - 70.0 / 3.0).abs() < 1e-10);
944 }
945
946 #[test]
947 fn cumulative_sum() {
948 let batch = test_batch(); let result = apply_window(&batch, &[WindowFn::CumulativeSum], &[], "value", None).unwrap();
950
951 let cs = result
952 .column_by_name("value_cumulative_sum")
953 .unwrap()
954 .as_any()
955 .downcast_ref::<Float64Array>()
956 .unwrap();
957 assert_eq!(cs.value(0), 10.0);
958 assert_eq!(cs.value(1), 30.0);
959 assert_eq!(cs.value(2), 45.0);
960 assert_eq!(cs.value(3), 70.0);
961 assert_eq!(cs.value(4), 100.0);
962 }
963
964 #[test]
965 fn increase() {
966 let batch = test_batch();
968 let result = apply_window(&batch, &[WindowFn::Increase], &[], "value", None).unwrap();
969
970 let inc = result
971 .column_by_name("value_increase")
972 .unwrap()
973 .as_any()
974 .downcast_ref::<Float64Array>()
975 .unwrap();
976 assert_eq!(inc.value(0), 0.0); assert_eq!(inc.value(1), 10.0); assert_eq!(inc.value(2), 10.0); assert_eq!(inc.value(3), 20.0); assert_eq!(inc.value(4), 25.0); }
982
983 #[test]
984 fn partitioned_row_number() {
985 let batch = partitioned_batch();
986 let result = apply_window(
987 &batch,
988 &[WindowFn::RowNumber],
989 &["host".to_string()],
990 "value",
991 None,
992 )
993 .unwrap();
994
995 let rn = result
996 .column_by_name("value_row_number")
997 .unwrap()
998 .as_any()
999 .downcast_ref::<Float64Array>()
1000 .unwrap();
1001 assert_eq!(rn.value(0), 1.0);
1003 assert_eq!(rn.value(2), 2.0);
1004 assert_eq!(rn.value(4), 3.0);
1005 assert_eq!(rn.value(1), 1.0);
1007 assert_eq!(rn.value(3), 2.0);
1008 assert_eq!(rn.value(5), 3.0);
1009 }
1010
1011 #[test]
1012 fn partitioned_delta() {
1013 let batch = partitioned_batch();
1014 let result = apply_window(
1015 &batch,
1016 &[WindowFn::Delta],
1017 &["host".to_string()],
1018 "value",
1019 None,
1020 )
1021 .unwrap();
1022
1023 let delta = result
1024 .column_by_name("value_delta")
1025 .unwrap()
1026 .as_any()
1027 .downcast_ref::<Float64Array>()
1028 .unwrap();
1029 assert!(delta.is_null(0));
1031 assert_eq!(delta.value(2), 10.0);
1032 assert_eq!(delta.value(4), 10.0);
1033 assert!(delta.is_null(1));
1035 assert_eq!(delta.value(3), 100.0);
1036 assert_eq!(delta.value(5), 100.0);
1037 }
1038
1039 #[test]
1040 fn empty_functions_returns_clone() {
1041 let batch = test_batch();
1042 let result = apply_window(&batch, &[], &[], "value", None).unwrap();
1043 assert_eq!(result.num_columns(), 3);
1044 assert_eq!(result.num_rows(), 5);
1045 }
1046
1047 #[test]
1048 fn multiple_window_functions() {
1049 let batch = test_batch();
1050 let result = apply_window(
1051 &batch,
1052 &[
1053 WindowFn::RowNumber,
1054 WindowFn::Delta,
1055 WindowFn::CumulativeSum,
1056 ],
1057 &[],
1058 "value",
1059 None,
1060 )
1061 .unwrap();
1062
1063 assert_eq!(result.num_columns(), 6); assert!(result.column_by_name("value_row_number").is_some());
1065 assert!(result.column_by_name("value_delta").is_some());
1066 assert!(result.column_by_name("value_cumulative_sum").is_some());
1067 }
1068
1069 #[test]
1070 fn missing_value_column_errors() {
1071 let batch = test_batch();
1072 let err = apply_window(&batch, &[WindowFn::Delta], &[], "nonexistent", None).unwrap_err();
1073 assert!(matches!(err, QueryError::FieldNotFound { .. }));
1074 }
1075
1076 #[test]
1077 fn i64_and_u64_columns() {
1078 let schema = Arc::new(Schema::new(vec![
1079 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
1080 Field::new("icount", DataType::Int64, false),
1081 Field::new("ucount", DataType::UInt64, false),
1082 ]));
1083 let ts = Arc::new(Int64Array::from(vec![
1084 1_000_000_000,
1085 2_000_000_000,
1086 3_000_000_000,
1087 ]));
1088 let ival = Arc::new(Int64Array::from(vec![10, 20, 30]));
1089 let uval = Arc::new(UInt64Array::from(vec![100u64, 200, 300]));
1090 let batch = RecordBatch::try_new(schema, vec![ts, ival, uval]).unwrap();
1091
1092 let result = apply_window(&batch, &[WindowFn::CumulativeSum], &[], "icount", None).unwrap();
1093 let cs = result
1094 .column_by_name("icount_cumulative_sum")
1095 .unwrap()
1096 .as_any()
1097 .downcast_ref::<Float64Array>()
1098 .unwrap();
1099 assert_eq!(cs.value(0), 10.0);
1100 assert_eq!(cs.value(1), 30.0);
1101 assert_eq!(cs.value(2), 60.0);
1102
1103 let result2 = apply_window(&batch, &[WindowFn::Delta], &[], "ucount", None).unwrap();
1104 let delta = result2
1105 .column_by_name("ucount_delta")
1106 .unwrap()
1107 .as_any()
1108 .downcast_ref::<Float64Array>()
1109 .unwrap();
1110 assert!(delta.is_null(0));
1111 assert_eq!(delta.value(1), 100.0);
1112 assert_eq!(delta.value(2), 100.0);
1113 }
1114
1115 #[test]
1116 fn window_respects_memory_tracker() {
1117 let batch = test_batch();
1118 let tracker = crate::memory::MemoryTracker::new(1_000_000);
1119 let result = apply_window(
1120 &batch,
1121 &[WindowFn::CumulativeSum],
1122 &[],
1123 "value",
1124 Some(&tracker),
1125 )
1126 .unwrap();
1127 assert!(tracker.allocated() > 0);
1128 assert_eq!(result.num_rows(), 5);
1129 }
1130
1131 #[test]
1132 fn window_exceeds_memory_budget() {
1133 let batch = test_batch();
1134 let tracker = crate::memory::MemoryTracker::new(1);
1135 let err = apply_window(
1136 &batch,
1137 &[WindowFn::CumulativeSum],
1138 &[],
1139 "value",
1140 Some(&tracker),
1141 )
1142 .unwrap_err();
1143 assert!(matches!(
1144 err,
1145 crate::error::QueryError::QueryMemoryExceeded { .. }
1146 ));
1147 }
1148
1149 #[test]
1152 fn range_moving_average() {
1153 let batch = test_batch();
1154 let result = apply_window(
1158 &batch,
1159 &[WindowFn::RangeMovingAverage {
1160 range_ns: 2_000_000_000,
1161 }],
1162 &[],
1163 "value",
1164 None,
1165 )
1166 .unwrap();
1167
1168 let avg = result
1169 .column_by_name("value_range_avg_2000000000ns")
1170 .unwrap()
1171 .as_any()
1172 .downcast_ref::<Float64Array>()
1173 .unwrap();
1174
1175 assert!((avg.value(0) - 10.0).abs() < f64::EPSILON);
1177 assert!((avg.value(1) - 15.0).abs() < f64::EPSILON);
1179 assert!((avg.value(2) - 15.0).abs() < f64::EPSILON);
1181 assert!((avg.value(3) - 20.0).abs() < f64::EPSILON);
1183 assert!((avg.value(4) - 70.0 / 3.0).abs() < 0.01);
1185 }
1186
1187 #[test]
1188 fn range_sum() {
1189 let batch = test_batch();
1190 let result = apply_window(
1192 &batch,
1193 &[WindowFn::RangeSum {
1194 range_ns: 1_000_000_000,
1195 }],
1196 &[],
1197 "value",
1198 None,
1199 )
1200 .unwrap();
1201
1202 let sum = result
1203 .column_by_name("value_range_sum_1000000000ns")
1204 .unwrap()
1205 .as_any()
1206 .downcast_ref::<Float64Array>()
1207 .unwrap();
1208
1209 assert!((sum.value(0) - 10.0).abs() < f64::EPSILON);
1212 assert!((sum.value(1) - 30.0).abs() < f64::EPSILON);
1214 assert!((sum.value(2) - 35.0).abs() < f64::EPSILON);
1216 assert!((sum.value(3) - 40.0).abs() < f64::EPSILON);
1218 assert!((sum.value(4) - 55.0).abs() < f64::EPSILON);
1220 }
1221
1222 #[test]
1223 fn range_count() {
1224 let batch = test_batch();
1225 let result = apply_window(
1226 &batch,
1227 &[WindowFn::RangeCount {
1228 range_ns: 1_500_000_000,
1229 }],
1230 &[],
1231 "value",
1232 None,
1233 )
1234 .unwrap();
1235
1236 let cnt = result
1237 .column_by_name("value_range_count_1500000000ns")
1238 .unwrap()
1239 .as_any()
1240 .downcast_ref::<Float64Array>()
1241 .unwrap();
1242
1243 assert!((cnt.value(0) - 1.0).abs() < f64::EPSILON);
1246 assert!((cnt.value(1) - 2.0).abs() < f64::EPSILON);
1248 assert!((cnt.value(2) - 2.0).abs() < f64::EPSILON);
1250 assert!((cnt.value(3) - 2.0).abs() < f64::EPSILON);
1252 }
1253
1254 #[test]
1255 fn range_moving_average_partitioned() {
1256 let batch = test_batch();
1257 let result = apply_window(
1260 &batch,
1261 &[WindowFn::RangeMovingAverage {
1262 range_ns: 5_000_000_000,
1263 }],
1264 &["host".to_string()],
1265 "value",
1266 None,
1267 )
1268 .unwrap();
1269
1270 let avg = result
1271 .column_by_name("value_range_avg_5000000000ns")
1272 .unwrap()
1273 .as_any()
1274 .downcast_ref::<Float64Array>()
1275 .unwrap();
1276
1277 assert!((avg.value(4) - 20.0).abs() < f64::EPSILON);
1280 }
1281
1282 #[test]
1283 fn range_zero_range_is_self_only() {
1284 let batch = test_batch();
1286 let result = apply_window(
1287 &batch,
1288 &[WindowFn::RangeMovingAverage { range_ns: 0 }],
1289 &[],
1290 "value",
1291 None,
1292 )
1293 .unwrap();
1294
1295 let avg = result
1296 .column_by_name("value_range_avg_0ns")
1297 .unwrap()
1298 .as_any()
1299 .downcast_ref::<Float64Array>()
1300 .unwrap();
1301
1302 assert!((avg.value(0) - 10.0).abs() < f64::EPSILON);
1304 assert!((avg.value(1) - 20.0).abs() < f64::EPSILON);
1305 assert!((avg.value(2) - 15.0).abs() < f64::EPSILON);
1306 assert!((avg.value(3) - 25.0).abs() < f64::EPSILON);
1307 assert!((avg.value(4) - 30.0).abs() < f64::EPSILON);
1308 }
1309
1310 #[test]
1311 fn range_with_null_values() {
1312 let schema = Arc::new(Schema::new(vec![
1313 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
1314 Field::new("value", DataType::Float64, true),
1315 ]));
1316 let ts = Arc::new(Int64Array::from(vec![
1317 1_000_000_000,
1318 2_000_000_000,
1319 3_000_000_000,
1320 4_000_000_000,
1321 ]));
1322 let values = Arc::new(Float64Array::from(vec![Some(10.0), None, None, Some(40.0)]));
1324 let batch = RecordBatch::try_new(schema, vec![ts, values]).unwrap();
1325
1326 let result = apply_window(
1327 &batch,
1328 &[
1329 WindowFn::RangeMovingAverage {
1330 range_ns: 2_000_000_000,
1331 },
1332 WindowFn::RangeSum {
1333 range_ns: 2_000_000_000,
1334 },
1335 WindowFn::RangeCount {
1336 range_ns: 2_000_000_000,
1337 },
1338 ],
1339 &[],
1340 "value",
1341 None,
1342 )
1343 .unwrap();
1344
1345 let avg = result
1346 .column_by_name("value_range_avg_2000000000ns")
1347 .unwrap()
1348 .as_any()
1349 .downcast_ref::<Float64Array>()
1350 .unwrap();
1351 let sum = result
1352 .column_by_name("value_range_sum_2000000000ns")
1353 .unwrap()
1354 .as_any()
1355 .downcast_ref::<Float64Array>()
1356 .unwrap();
1357 let cnt = result
1358 .column_by_name("value_range_count_2000000000ns")
1359 .unwrap()
1360 .as_any()
1361 .downcast_ref::<Float64Array>()
1362 .unwrap();
1363
1364 assert!((avg.value(0) - 10.0).abs() < f64::EPSILON);
1366 assert!((sum.value(0) - 10.0).abs() < f64::EPSILON);
1367 assert!((cnt.value(0) - 1.0).abs() < f64::EPSILON);
1368
1369 assert!((avg.value(1) - 10.0).abs() < f64::EPSILON);
1371 assert!((sum.value(1) - 10.0).abs() < f64::EPSILON);
1372 assert!((cnt.value(1) - 1.0).abs() < f64::EPSILON);
1373
1374 assert!((avg.value(2) - 10.0).abs() < f64::EPSILON);
1377
1378 assert!((avg.value(3) - 40.0).abs() < f64::EPSILON);
1380 assert!((sum.value(3) - 40.0).abs() < f64::EPSILON);
1381 assert!((cnt.value(3) - 1.0).abs() < f64::EPSILON);
1382 }
1383
1384 #[test]
1385 fn range_all_null_returns_none() {
1386 let schema = Arc::new(Schema::new(vec![
1387 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
1388 Field::new("value", DataType::Float64, true),
1389 ]));
1390 let ts = Arc::new(Int64Array::from(vec![1_000_000_000, 2_000_000_000]));
1391 let values = Arc::new(Float64Array::from(vec![None, None]));
1392 let batch = RecordBatch::try_new(schema, vec![ts, values]).unwrap();
1393
1394 let result = apply_window(
1395 &batch,
1396 &[
1397 WindowFn::RangeMovingAverage {
1398 range_ns: 5_000_000_000,
1399 },
1400 WindowFn::RangeSum {
1401 range_ns: 5_000_000_000,
1402 },
1403 WindowFn::RangeCount {
1404 range_ns: 5_000_000_000,
1405 },
1406 ],
1407 &[],
1408 "value",
1409 None,
1410 )
1411 .unwrap();
1412
1413 let avg = result
1414 .column_by_name("value_range_avg_5000000000ns")
1415 .unwrap()
1416 .as_any()
1417 .downcast_ref::<Float64Array>()
1418 .unwrap();
1419 let sum = result
1420 .column_by_name("value_range_sum_5000000000ns")
1421 .unwrap()
1422 .as_any()
1423 .downcast_ref::<Float64Array>()
1424 .unwrap();
1425 let cnt = result
1426 .column_by_name("value_range_count_5000000000ns")
1427 .unwrap()
1428 .as_any()
1429 .downcast_ref::<Float64Array>()
1430 .unwrap();
1431
1432 assert!(avg.is_null(0));
1434 assert!(avg.is_null(1));
1435 assert!(sum.is_null(0));
1436 assert!(sum.is_null(1));
1437 assert!(cnt.is_null(0));
1438 assert!(cnt.is_null(1));
1439 }
1440
1441 #[test]
1442 fn range_missing_timestamp_errors() {
1443 let schema = Arc::new(Schema::new(vec![Field::new(
1444 "value",
1445 DataType::Float64,
1446 false,
1447 )]));
1448 let values = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0]));
1449 let batch = RecordBatch::try_new(schema, vec![values]).unwrap();
1450
1451 let err = apply_window(
1452 &batch,
1453 &[WindowFn::RangeMovingAverage { range_ns: 1000 }],
1454 &[],
1455 "value",
1456 None,
1457 )
1458 .unwrap_err();
1459 assert!(matches!(err, QueryError::Validation(_)));
1460
1461 let err = apply_window(
1462 &batch,
1463 &[WindowFn::RangeSum { range_ns: 1000 }],
1464 &[],
1465 "value",
1466 None,
1467 )
1468 .unwrap_err();
1469 assert!(matches!(err, QueryError::Validation(_)));
1470
1471 let err = apply_window(
1472 &batch,
1473 &[WindowFn::RangeCount { range_ns: 1000 }],
1474 &[],
1475 "value",
1476 None,
1477 )
1478 .unwrap_err();
1479 assert!(matches!(err, QueryError::Validation(_)));
1480 }
1481
1482 #[test]
1483 fn range_duplicate_timestamps() {
1484 let schema = Arc::new(Schema::new(vec![
1485 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
1486 Field::new("value", DataType::Float64, false),
1487 ]));
1488 let ts = Arc::new(Int64Array::from(vec![
1490 1_000_000_000,
1491 1_000_000_000,
1492 2_000_000_000,
1493 ]));
1494 let values = Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0]));
1495 let batch = RecordBatch::try_new(schema, vec![ts, values]).unwrap();
1496
1497 let result = apply_window(
1498 &batch,
1499 &[WindowFn::RangeMovingAverage {
1500 range_ns: 500_000_000,
1501 }],
1502 &[],
1503 "value",
1504 None,
1505 )
1506 .unwrap();
1507
1508 let avg = result
1509 .column_by_name("value_range_avg_500000000ns")
1510 .unwrap()
1511 .as_any()
1512 .downcast_ref::<Float64Array>()
1513 .unwrap();
1514
1515 assert!((avg.value(0) - 10.0).abs() < f64::EPSILON);
1517 assert!((avg.value(1) - 15.0).abs() < f64::EPSILON);
1519 assert!((avg.value(2) - 30.0).abs() < f64::EPSILON);
1521 }
1522
1523 #[test]
1524 fn range_single_row() {
1525 let schema = Arc::new(Schema::new(vec![
1526 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
1527 Field::new("value", DataType::Float64, false),
1528 ]));
1529 let ts = Arc::new(Int64Array::from(vec![1_000_000_000]));
1530 let values = Arc::new(Float64Array::from(vec![42.0]));
1531 let batch = RecordBatch::try_new(schema, vec![ts, values]).unwrap();
1532
1533 let result = apply_window(
1534 &batch,
1535 &[
1536 WindowFn::RangeMovingAverage {
1537 range_ns: 1_000_000_000,
1538 },
1539 WindowFn::RangeSum {
1540 range_ns: 1_000_000_000,
1541 },
1542 WindowFn::RangeCount {
1543 range_ns: 1_000_000_000,
1544 },
1545 ],
1546 &[],
1547 "value",
1548 None,
1549 )
1550 .unwrap();
1551
1552 let avg = result
1553 .column_by_name("value_range_avg_1000000000ns")
1554 .unwrap()
1555 .as_any()
1556 .downcast_ref::<Float64Array>()
1557 .unwrap();
1558 let sum = result
1559 .column_by_name("value_range_sum_1000000000ns")
1560 .unwrap()
1561 .as_any()
1562 .downcast_ref::<Float64Array>()
1563 .unwrap();
1564 let cnt = result
1565 .column_by_name("value_range_count_1000000000ns")
1566 .unwrap()
1567 .as_any()
1568 .downcast_ref::<Float64Array>()
1569 .unwrap();
1570
1571 assert!((avg.value(0) - 42.0).abs() < f64::EPSILON);
1572 assert!((sum.value(0) - 42.0).abs() < f64::EPSILON);
1573 assert!((cnt.value(0) - 1.0).abs() < f64::EPSILON);
1574 }
1575}