promql/
extension_plan.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
15mod absent;
16mod empty_metric;
17mod histogram_fold;
18mod instant_manipulate;
19mod normalize;
20mod planner;
21mod range_manipulate;
22mod scalar_calculate;
23mod series_divide;
24#[cfg(test)]
25mod test_util;
26mod union_distinct_on;
27
28pub use absent::{Absent, AbsentExec, AbsentStream};
29use common_query::native_histogram::{SUM_FIELD, native_histogram_value_type};
30use common_query::prometheus::is_prometheus_stale_nan;
31use datafusion::arrow::array::{Array, Float64Array, StructArray};
32use datafusion::arrow::datatypes::{ArrowPrimitiveType, TimestampMillisecondType};
33use datafusion::common::DFSchemaRef;
34use datafusion::error::{DataFusionError, Result as DataFusionResult};
35use datatypes::data_type::DataType as _;
36pub use empty_metric::{EmptyMetric, EmptyMetricExec, EmptyMetricStream, build_special_time_expr};
37pub use histogram_fold::{
38    HistogramFold, HistogramFoldExec, HistogramFoldOperation, HistogramFoldStream,
39};
40pub use instant_manipulate::{InstantManipulate, InstantManipulateExec, InstantManipulateStream};
41pub use normalize::{SeriesNormalize, SeriesNormalizeExec, SeriesNormalizeStream};
42pub use planner::PromExtensionPlanner;
43pub use range_manipulate::{RangeManipulate, RangeManipulateExec, RangeManipulateStream};
44pub use scalar_calculate::ScalarCalculate;
45pub use series_divide::{SeriesDivide, SeriesDivideExec, SeriesDivideStream};
46pub use union_distinct_on::{UnionDistinctOn, UnionDistinctOnExec, UnionDistinctOnStream};
47
48pub type Millisecond = <TimestampMillisecondType as ArrowPrimitiveType>::Native;
49
50const METRIC_NUM_SERIES: &str = "num_series";
51
52fn prometheus_stale_sample_column(column: &dyn Array) -> Option<(&dyn Array, &Float64Array)> {
53    let values = if let Some(values) = column.as_any().downcast_ref::<Float64Array>() {
54        values
55    } else {
56        let histograms = column.as_any().downcast_ref::<StructArray>()?;
57        if histograms.data_type() != &native_histogram_value_type().as_arrow_type() {
58            return None;
59        }
60        histograms
61            .column_by_name(SUM_FIELD)?
62            .as_any()
63            .downcast_ref::<Float64Array>()?
64    };
65    Some((column, values))
66}
67
68fn is_prometheus_stale_sample((column, values): (&dyn Array, &Float64Array), row: usize) -> bool {
69    column.is_valid(row) && values.is_valid(row) && is_prometheus_stale_nan(values.value(row))
70}
71
72/// Utilities for handling unfix logic in extension plans
73/// Convert column name to index for serialization
74pub fn serialize_column_index(schema: &DFSchemaRef, column_name: &str) -> u64 {
75    schema
76        .index_of_column_by_name(None, column_name)
77        .map(|idx| idx as u64)
78        .unwrap_or(u64::MAX) // make sure if not found, it will report error in deserialization
79}
80
81/// Convert index back to column name for deserialization
82pub fn resolve_column_name(
83    index: u64,
84    schema: &DFSchemaRef,
85    context: &str,
86    column_type: &str,
87) -> DataFusionResult<String> {
88    let columns = schema.columns();
89    columns
90        .get(index as usize)
91        .ok_or_else(|| {
92            DataFusionError::Internal(format!(
93                "Failed to get {} column at idx {} during unfixing {} with columns:{:?}",
94                column_type, index, context, columns
95            ))
96        })
97        .map(|field| field.name().to_string())
98}
99
100/// Batch process multiple column indices
101pub fn resolve_column_names(
102    indices: &[u64],
103    schema: &DFSchemaRef,
104    context: &str,
105    column_type: &str,
106) -> DataFusionResult<Vec<String>> {
107    indices
108        .iter()
109        .map(|idx| resolve_column_name(*idx, schema, context, column_type))
110        .collect()
111}
promql/extension_plan.rs

promql/
extension_plan.rs
