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