1use std::fmt::Display;
16use std::sync::Arc;
17
18use datafusion::arrow::array::{Float64Array, TimestampMillisecondArray};
19use datafusion::arrow::datatypes::TimeUnit;
20use datafusion::common::DataFusionError;
21use datafusion::logical_expr::{ScalarUDF, Volatility};
22use datafusion::physical_plan::ColumnarValue;
23use datafusion_expr::create_udf;
24use datatypes::arrow::array::Array;
25use datatypes::arrow::datatypes::DataType;
26
27use crate::error;
28use crate::functions::extract_array;
29use crate::range_array::RangeArray;
30
31#[derive(Debug)]
34pub struct IDelta<const IS_RATE: bool> {}
35
36impl<const IS_RATE: bool> IDelta<IS_RATE> {
37 pub const fn name() -> &'static str {
38 if IS_RATE { "prom_irate" } else { "prom_idelta" }
39 }
40
41 pub fn scalar_udf() -> ScalarUDF {
42 create_udf(
43 Self::name(),
44 Self::input_type(),
45 Self::return_type(),
46 Volatility::Volatile,
47 Arc::new(Self::calc) as _,
48 )
49 }
50
51 fn input_type() -> Vec<DataType> {
53 vec![
54 RangeArray::convert_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)),
55 RangeArray::convert_data_type(DataType::Float64),
56 ]
57 }
58
59 fn return_type() -> DataType {
60 DataType::Float64
61 }
62
63 fn calc(input: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
64 assert_eq!(input.len(), 2);
66 let ts_array = extract_array(&input[0])?;
67 let value_array = extract_array(&input[1])?;
68
69 let ts_range: RangeArray = RangeArray::try_new(ts_array.to_data().into())?;
70 let value_range: RangeArray = RangeArray::try_new(value_array.to_data().into())?;
71 error::ensure(
72 ts_range.len() == value_range.len(),
73 DataFusionError::Execution(format!(
74 "{}: input arrays should have the same length, found {} and {}",
75 Self::name(),
76 ts_range.len(),
77 value_range.len()
78 )),
79 )?;
80 error::ensure(
81 ts_range.value_type() == DataType::Timestamp(TimeUnit::Millisecond, None),
82 DataFusionError::Execution(format!(
83 "{}: expect TimestampMillisecond as time index array's type, found {}",
84 Self::name(),
85 ts_range.value_type()
86 )),
87 )?;
88 error::ensure(
89 value_range.value_type() == DataType::Float64,
90 DataFusionError::Execution(format!(
91 "{}: expect Float64 as value array's type, found {}",
92 Self::name(),
93 value_range.value_type()
94 )),
95 )?;
96
97 let mut result_array = Vec::with_capacity(ts_range.len());
99
100 for index in 0..ts_range.len() {
101 let timestamps = ts_range.get(index).unwrap();
102 let timestamps = timestamps
103 .as_any()
104 .downcast_ref::<TimestampMillisecondArray>()
105 .unwrap()
106 .values();
107
108 let values = value_range.get(index).unwrap();
109 let values = values
110 .as_any()
111 .downcast_ref::<Float64Array>()
112 .unwrap()
113 .values();
114 error::ensure(
115 timestamps.len() == values.len(),
116 DataFusionError::Execution(format!(
117 "{}: input arrays should have the same length, found {} and {}",
118 Self::name(),
119 timestamps.len(),
120 values.len()
121 )),
122 )?;
123
124 let len = timestamps.len();
125 if len < 2 {
126 result_array.push(None);
127 continue;
128 }
129
130 if !IS_RATE {
132 result_array.push(Some(values[len - 1] - values[len - 2]));
133 continue;
134 }
135
136 let sampled_interval = (timestamps[len - 1] - timestamps[len - 2]) as f64 / 1000.0;
138 let last_value = values[len - 1];
139 let prev_value = values[len - 2];
140 let result_value = if last_value < prev_value {
141 last_value
143 } else {
144 last_value - prev_value
145 };
146
147 result_array.push(Some(result_value / sampled_interval as f64));
148 }
149
150 let result = ColumnarValue::Array(Arc::new(Float64Array::from_iter(result_array)));
151 Ok(result)
152 }
153}
154
155impl<const IS_RATE: bool> Display for IDelta<IS_RATE> {
156 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
157 write!(f, "PromQL Idelta Function (is_rate: {IS_RATE})",)
158 }
159}
160
161#[cfg(test)]
162mod test {
163
164 use super::*;
165 use crate::functions::test_util::simple_range_udf_runner;
166
167 #[test]
168 fn basic_idelta_and_irate() {
169 let ts_array = Arc::new(TimestampMillisecondArray::from_iter(
170 [1000i64, 3000, 5000, 7000, 9000, 11000, 13000, 15000, 17000]
171 .into_iter()
172 .map(Some),
173 ));
174 let ts_ranges = [(0, 2), (0, 5), (1, 1), (3, 3), (8, 1), (9, 0)];
175
176 let values_array = Arc::new(Float64Array::from_iter([
177 1.0, 2.0, 3.0, 5.0, 0.0, 6.0, 7.0, 8.0, 9.0,
178 ]));
179 let values_ranges = [(0, 2), (0, 5), (1, 1), (3, 3), (8, 1), (9, 0)];
180
181 let ts_range_array = RangeArray::from_ranges(ts_array.clone(), ts_ranges).unwrap();
183 let value_range_array =
184 RangeArray::from_ranges(values_array.clone(), values_ranges).unwrap();
185 simple_range_udf_runner(
186 IDelta::<false>::scalar_udf(),
187 ts_range_array,
188 value_range_array,
189 vec![],
190 vec![Some(1.0), Some(-5.0), None, Some(6.0), None, None],
191 );
192
193 let ts_range_array = RangeArray::from_ranges(ts_array, ts_ranges).unwrap();
195 let value_range_array = RangeArray::from_ranges(values_array, values_ranges).unwrap();
196 simple_range_udf_runner(
197 IDelta::<true>::scalar_udf(),
198 ts_range_array,
199 value_range_array,
200 vec![],
201 vec![Some(0.5), Some(0.0), None, Some(3.0), None, None],
203 );
204 }
205}