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    let direct = should_scan_direct(requested_edges, raw_values.len());
85    let prefix = (!direct).then(|| build_prefix(raw_values.as_ref(), kind));
86
87    let mut result = Vec::with_capacity(value_ranges.len());
88    for index in 0..value_ranges.len() {
89        let (offset, len) = range_at(&value_ranges, index, name)?;
90        let end = checked_end(offset, len, index, name)?;
91        let count = match len {
92            0 => None,
93            1 => Some(0),
94            _ if direct => Some(count_edges(raw_values.as_ref(), offset, end, kind)),
95            _ => {
96                let prefix = prefix.as_ref().unwrap();
97                Some(prefix[end - 1] - prefix[offset])
98            }
99        };
100        result.push(count.map(|count| count as f64));
101    }
102
103    Ok(ColumnarValue::Array(Arc::new(Float64Array::from_iter(
104        result,
105    ))))
106}
107
108fn validate_windows(
109    timestamps: &RangeArray,
110    values: &RangeArray,
111    name: &str,
112) -> Result<usize, DataFusionError> {
113    let mut requested_edges = 0usize;
114    for index in 0..values.len() {
115        let (timestamp_offset, timestamp_len) = range_at(timestamps, index, name)?;
116        let (value_offset, value_len) = range_at(values, index, name)?;
117        if timestamp_len != value_len {
118            return Err(DataFusionError::Execution(format!(
119                "RangeArray's element {index} have different lengths in PromQL function {name}: array1={timestamp_len}, array2={value_len}"
120            )));
121        }
122        checked_end(timestamp_offset, timestamp_len, index, name)?;
123        checked_end(value_offset, value_len, index, name)?;
124        requested_edges = requested_edges.saturating_add(value_len.saturating_sub(1));
125    }
126    Ok(requested_edges)
127}
128
129fn range_at(
130    ranges: &RangeArray,
131    index: usize,
132    name: &str,
133) -> Result<(usize, usize), DataFusionError> {
134    ranges.get_offset_length(index).ok_or_else(|| {
135        DataFusionError::Execution(format!(
136            "RangeArray's element {index} is unavailable in PromQL function {name}"
137        ))
138    })
139}
140
141fn checked_end(
142    offset: usize,
143    len: usize,
144    index: usize,
145    name: &str,
146) -> Result<usize, DataFusionError> {
147    offset.checked_add(len).ok_or_else(|| {
148        DataFusionError::Execution(format!(
149            "RangeArray's element {index} has an invalid range in PromQL function {name}"
150        ))
151    })
152}
153
154fn should_scan_direct(requested_edges: usize, backing_len: usize) -> bool {
155    requested_edges <= backing_len.saturating_sub(1)
156}
157
158fn build_prefix(values: &[f64], kind: EdgeKind) -> Vec<u64> {
159    let mut prefix = Vec::with_capacity(values.len());
160    prefix.push(0);
161    for index in 1..values.len() {
162        prefix.push(prefix[index - 1] + u64::from(is_edge(values[index - 1], values[index], kind)));
163    }
164    prefix
165}
166
167fn count_edges(values: &[f64], offset: usize, end: usize, kind: EdgeKind) -> u64 {
168    let mut count = 0;
169    for index in offset + 1..end {
170        count += u64::from(is_edge(values[index - 1], values[index], kind));
171    }
172    count
173}
174
175fn is_edge(previous: f64, current: f64, kind: EdgeKind) -> bool {
176    match kind {
177        EdgeKind::Changes => previous != current && !(previous.is_nan() && current.is_nan()),
178        EdgeKind::Resets => current < previous,
179    }
180}
181
182#[cfg(test)]
183mod test {
184    use super::*;
185
186    #[test]
187    fn adaptive_gate_uses_direct_scan_at_threshold() {
188        assert!(should_scan_direct(4, 5));
189        assert!(!should_scan_direct(5, 5));
190        assert!(should_scan_direct(0, 0));
191    }
192}
promql/functions/edge_count.rs

promql/functions/
edge_count.rs
