promql/functions/
edge_count.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use datafusion::arrow::array::{Float64Array, TimestampMillisecondArray};
18use datafusion::common::DataFusionError;
19use datafusion::logical_expr::{ScalarUDF, Volatility};
20use datafusion::physical_plan::ColumnarValue;
21use datafusion_expr::create_udf;
22use datatypes::arrow::array::Array;
23use datatypes::arrow::datatypes::DataType;
24
25use crate::functions::extract_range_array;
26use crate::range_array::RangeArray;
27
28#[derive(Clone, Copy)]
29pub(super) enum EdgeKind {
30    Changes,
31    Resets,
32}
33
34pub(super) fn scalar_udf(name: &'static str, kind: EdgeKind) -> ScalarUDF {
35    create_udf(
36        name,
37        input_type(),
38        DataType::Float64,
39        Volatility::Volatile,
40        Arc::new(
41            move |input: &[ColumnarValue]| -> Result<ColumnarValue, DataFusionError> {
42                calc(input, name, kind)
43            },
44        ) as _,
45    )
46}
47
48fn input_type() -> Vec<DataType> {
49    vec![
50        RangeArray::convert_data_type(TimestampMillisecondArray::new_null(0).data_type().clone()),
51        RangeArray::convert_data_type(Float64Array::new_null(0).data_type().clone()),
52    ]
53}
54
55fn calc(
56    input: &[ColumnarValue],
57    name: &str,
58    kind: EdgeKind,
59) -> Result<ColumnarValue, DataFusionError> {
60    assert_eq!(input.len(), 2);
61
62    let timestamp_ranges = extract_range_array(&input[0])?;
63    let value_ranges = extract_range_array(&input[1])?;
64    if timestamp_ranges.len() != value_ranges.len() {
65        return Err(DataFusionError::Execution(format!(
66            "RangeArray have different lengths in PromQL function {name}: array1={}, array2={}",
67            timestamp_ranges.len(),
68            value_ranges.len()
69        )));
70    }
71
72    timestamp_ranges
73        .values()
74        .as_any()
75        .downcast_ref::<TimestampMillisecondArray>()
76        .unwrap();
77    let values = value_ranges
78        .values()
79        .as_any()
80        .downcast_ref::<Float64Array>()
81        .unwrap();
82    let requested_edges = validate_windows(&timestamp_ranges, &value_ranges, name)?;
83    let raw_values = values.values();
84    // A NULL field value means the series has no sample at that timestamp. The prefix sums
85    // encode edges between physically adjacent slots, which no longer holds once nulls are
86    // skipped, so a null-bearing input falls back to scanning each window.
87    let has_nulls = values.null_count() > 0;
88    let direct = has_nulls || should_scan_direct(requested_edges, raw_values.len());
89    let prefix = (!direct).then(|| build_prefix(raw_values.as_ref(), kind));
90
91    let mut result = Vec::with_capacity(value_ranges.len());
92    for index in 0..value_ranges.len() {
93        let (offset, len) = range_at(&value_ranges, index, name)?;
94        let end = checked_end(offset, len, index, name)?;
95        let count = match len {
96            0 => None,
97            _ if has_nulls => count_edges_skipping_nulls(values, offset, end, kind),
98            1 => Some(0),
99            _ if direct => Some(count_edges(raw_values.as_ref(), offset, end, kind)),
100            _ => {
101                let prefix = prefix.as_ref().unwrap();
102                Some(prefix[end - 1] - prefix[offset])
103            }
104        };
105        result.push(count.map(|count| count as f64));
106    }
107
108    Ok(ColumnarValue::Array(Arc::new(Float64Array::from_iter(
109        result,
110    ))))
111}
112
113fn validate_windows(
114    timestamps: &RangeArray,
115    values: &RangeArray,
116    name: &str,
117) -> Result<usize, DataFusionError> {
118    let mut requested_edges = 0usize;
119    for index in 0..values.len() {
120        let (timestamp_offset, timestamp_len) = range_at(timestamps, index, name)?;
121        let (value_offset, value_len) = range_at(values, index, name)?;
122        if timestamp_len != value_len {
123            return Err(DataFusionError::Execution(format!(
124                "RangeArray's element {index} have different lengths in PromQL function {name}: array1={timestamp_len}, array2={value_len}"
125            )));
126        }
127        checked_end(timestamp_offset, timestamp_len, index, name)?;
128        checked_end(value_offset, value_len, index, name)?;
129        requested_edges = requested_edges.saturating_add(value_len.saturating_sub(1));
130    }
131    Ok(requested_edges)
132}
133
134fn range_at(
135    ranges: &RangeArray,
136    index: usize,
137    name: &str,
138) -> Result<(usize, usize), DataFusionError> {
139    ranges.get_offset_length(index).ok_or_else(|| {
140        DataFusionError::Execution(format!(
141            "RangeArray's element {index} is unavailable in PromQL function {name}"
142        ))
143    })
144}
145
146fn checked_end(
147    offset: usize,
148    len: usize,
149    index: usize,
150    name: &str,
151) -> Result<usize, DataFusionError> {
152    offset.checked_add(len).ok_or_else(|| {
153        DataFusionError::Execution(format!(
154            "RangeArray's element {index} has an invalid range in PromQL function {name}"
155        ))
156    })
157}
158
159fn should_scan_direct(requested_edges: usize, backing_len: usize) -> bool {
160    requested_edges <= backing_len.saturating_sub(1)
161}
162
163fn build_prefix(values: &[f64], kind: EdgeKind) -> Vec<u64> {
164    let mut prefix = Vec::with_capacity(values.len());
165    prefix.push(0);
166    for index in 1..values.len() {
167        prefix.push(prefix[index - 1] + u64::from(is_edge(values[index - 1], values[index], kind)));
168    }
169    prefix
170}
171
172fn count_edges(values: &[f64], offset: usize, end: usize, kind: EdgeKind) -> u64 {
173    let mut count = 0;
174    for index in offset + 1..end {
175        count += u64::from(is_edge(values[index - 1], values[index], kind));
176    }
177    count
178}
179
180/// Counts edges between consecutive samples in `[offset, end)`, treating null slots as
181/// absent. Returns `None` when the window holds no sample.
182fn count_edges_skipping_nulls(
183    values: &Float64Array,
184    offset: usize,
185    end: usize,
186    kind: EdgeKind,
187) -> Option<u64> {
188    let raw_values = values.values();
189    let mut count = 0;
190    let mut previous = None;
191    for index in offset..end {
192        if values.is_null(index) {
193            continue;
194        }
195        let current = raw_values[index];
196        if let Some(previous) = previous.replace(current) {
197            count += u64::from(is_edge(previous, current, kind));
198        }
199    }
200    previous.is_some().then_some(count)
201}
202
203fn is_edge(previous: f64, current: f64, kind: EdgeKind) -> bool {
204    match kind {
205        EdgeKind::Changes => previous != current && !(previous.is_nan() && current.is_nan()),
206        EdgeKind::Resets => current < previous,
207    }
208}
209
210#[cfg(test)]
211mod test {
212    use super::*;
213
214    #[test]
215    fn adaptive_gate_uses_direct_scan_at_threshold() {
216        assert!(should_scan_direct(4, 5));
217        assert!(!should_scan_direct(5, 5));
218        assert!(should_scan_direct(0, 0));
219    }
220}
promql/functions/edge_count.rs

promql/functions/
edge_count.rs
