query/promql/
label_values.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::time::{SystemTime, UNIX_EPOCH};
16
17use common_time::Timestamp;
18use common_time::timestamp::TimeUnit;
19use datafusion_common::{Column, ScalarValue};
20use datafusion_expr::utils::conjunction;
21use datafusion_expr::{Expr, LogicalPlan, LogicalPlanBuilder, col};
22use snafu::{OptionExt, ResultExt};
23use table::TableRef;
24
25use crate::promql::error::{
26    DataFusionPlanningSnafu, Result, SystemTimeOutOfRangeSnafu, TimeIndexNotFoundSnafu,
27    TimestampOutOfRangeSnafu,
28};
29
30/// Converts a [`SystemTime`] to a millisecond [`Timestamp`].
31///
32/// `duration_since` reports an instant before the epoch as an error rather than
33/// a negative duration, and an RFC3339 query parameter can name one, so the sign
34/// is recovered here instead of unwrapping.
35fn millis_since_epoch(time: SystemTime) -> Result<Timestamp> {
36    Ok(Timestamp::new_millisecond(signed_millis_since_epoch(time)?))
37}
38
39/// Converts a [`SystemTime`] to signed milliseconds since the Unix epoch, rejecting instants
40/// outside the `i64` millisecond range.
41///
42/// Sub-millisecond parts round down, also before the epoch, like Prometheus'
43/// `timestamp.FromTime`.
44pub(crate) fn signed_millis_since_epoch(time: SystemTime) -> Result<i64> {
45    let (millis, before_epoch) = match time.duration_since(UNIX_EPOCH) {
46        Ok(duration) => (duration.as_millis(), false),
47        Err(earlier) => (earlier.duration().as_nanos().div_ceil(1_000_000), true),
48    };
49    signed_millis(millis, before_epoch).with_context(|| SystemTimeOutOfRangeSnafu { time })
50}
51
52/// Converts a millisecond magnitude and epoch direction without platform time limits.
53fn signed_millis(millis: u128, before_epoch: bool) -> Option<i64> {
54    let millis = i64::try_from(millis).ok()?;
55    Some(if before_epoch { -millis } else { millis })
56}
57
58fn build_time_filter(time_index_expr: Expr, start: Timestamp, end: Timestamp) -> Expr {
59    time_index_expr
60        .clone()
61        .gt_eq(Expr::Literal(timestamp_to_scalar_value(start), None))
62        .and(time_index_expr.lt_eq(Expr::Literal(timestamp_to_scalar_value(end), None)))
63}
64
65fn timestamp_to_scalar_value(timestamp: Timestamp) -> ScalarValue {
66    let value = timestamp.value();
67    match timestamp.unit() {
68        TimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None),
69        TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None),
70        TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None),
71        TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None),
72    }
73}
74
75/// Rewrite label values query to DataFusion logical plan.
76pub fn rewrite_label_values_query(
77    table: TableRef,
78    scan_plan: LogicalPlan,
79    mut conditions: Vec<Expr>,
80    label_name: String,
81    start: SystemTime,
82    end: SystemTime,
83) -> Result<LogicalPlan> {
84    let schema = table.schema();
85    let ts_column = schema
86        .timestamp_column()
87        .with_context(|| TimeIndexNotFoundSnafu {
88            table: table.table_info().full_table_name(),
89        })?;
90    let unit = ts_column
91        .data_type
92        .as_timestamp()
93        .map(|data_type| data_type.unit())
94        .with_context(|| TimeIndexNotFoundSnafu {
95            table: table.table_info().full_table_name(),
96        })?;
97
98    // We only support millisecond precision at most.
99    let start = millis_since_epoch(start)?;
100    let start = start.convert_to(unit).context(TimestampOutOfRangeSnafu {
101        timestamp: start.value(),
102        unit,
103    })?;
104    let end = millis_since_epoch(end)?;
105    let end = end.convert_to(unit).context(TimestampOutOfRangeSnafu {
106        timestamp: end.value(),
107        unit,
108    })?;
109    let time_index_expr = col(Column::from_name(ts_column.name.clone()));
110
111    conditions.push(build_time_filter(time_index_expr, start, end));
112    // Safety: `conditions` is not empty.
113    let filter = conjunction(conditions).unwrap();
114
115    // Builds time filter
116    let logical_plan = LogicalPlanBuilder::from(scan_plan)
117        .filter(filter)
118        .context(DataFusionPlanningSnafu)?
119        .project(vec![col(Column::from_name(label_name))])
120        .context(DataFusionPlanningSnafu)?
121        .distinct()
122        .context(DataFusionPlanningSnafu)?
123        .build()
124        .context(DataFusionPlanningSnafu)?;
125
126    Ok(logical_plan)
127}
128
129#[cfg(test)]
130mod tests {
131    use std::time::Duration;
132
133    use super::*;
134
135    #[test]
136    fn millis_before_the_epoch_are_negative() {
137        // `SystemTime::duration_since` reports these as an error; unwrapping it
138        // panicked on any request whose RFC3339 start named a pre-epoch instant.
139        let time = UNIX_EPOCH - Duration::from_millis(1);
140        assert_eq!(millis_since_epoch(time).unwrap().value(), -1);
141        assert_eq!(millis_since_epoch(UNIX_EPOCH).unwrap().value(), 0);
142
143        let time = UNIX_EPOCH + Duration::from_millis(1);
144        assert_eq!(millis_since_epoch(time).unwrap().value(), 1);
145    }
146
147    #[test]
148    fn millis_beyond_i64_are_rejected() {
149        let limit = i64::MAX as u128;
150        assert_eq!(signed_millis(limit, false), Some(i64::MAX));
151        assert_eq!(signed_millis(limit, true), Some(-i64::MAX));
152        assert_eq!(signed_millis(limit + 1, false), None);
153        assert_eq!(signed_millis(limit + 1, true), None);
154
155        // Windows cannot represent this SystemTime. The conversion boundary is
156        // tested above on every platform; exercise the wrapper where possible.
157        if let Some(time) = UNIX_EPOCH.checked_add(Duration::from_millis(i64::MAX as u64 + 1)) {
158            assert!(millis_since_epoch(time).is_err());
159        }
160    }
161}
query/promql/label_values.rs

query/promql/
label_values.rs
