query/promql/planner/
function_plans.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
15//! Dedicated helper plans for PromQL functions that cannot be expressed as a
16//! plain scalar function over the input plan: the histogram helpers, `vector()`,
17//! `scalar()`, and `absent()`.
18
19use std::sync::Arc;
20
21use common_query::prelude::greptime_value;
22use datafusion::functions_aggregate::expr_fn::first_value;
23use datafusion::logical_expr::expr::ScalarFunction;
24use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder};
25use datafusion::optimizer::simplify_expressions::ExprSimplifier;
26use datafusion::prelude::{Column, Expr as DfExpr};
27use datafusion::scalar::ScalarValue;
28use datafusion_common::DFSchema;
29use datafusion_expr::expr_fn::when;
30use datafusion_expr::simplify::SimplifyContext;
31use datafusion_expr::{ident, lit};
32use datafusion_functions::core::coalesce;
33use datatypes::arrow::datatypes::DataType as ArrowDataType;
34use promql::extension_plan::{
35    Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, ScalarCalculate,
36};
37use promql::functions::{NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile};
38use promql_parser::label::MatchOp;
39use promql_parser::parser::{Expr as PromExpr, FunctionArgs as PromFunctionArgs};
40use snafu::{OptionExt, ResultExt, ensure};
41
42use crate::promql::error::{
43    DataFusionPlanningSnafu, FunctionInvalidArgumentSnafu, MultiFieldsNotSupportedSnafu,
44    PromqlPlanNodeSnafu, Result, TimeIndexNotFoundSnafu, ValueNotFoundSnafu,
45};
46use crate::promql::planner::{
47    LE_COLUMN_NAME, PromPlanner, SCALAR_FUNCTION, SPECIAL_ABSENT_FUNCTION,
48    SPECIAL_HISTOGRAM_FRACTION, SPECIAL_HISTOGRAM_QUANTILE, SPECIAL_TIME_FUNCTION,
49    SPECIAL_VECTOR_FUNCTION,
50};
51use crate::query_engine::QueryEngineState;
52
53impl PromPlanner {
54    /// Create a classic, native, or mixed histogram helper plan.
55    pub(super) async fn create_histogram_plan(
56        &mut self,
57        function_name: &str,
58        args: &PromFunctionArgs,
59        query_engine_state: &QueryEngineState,
60    ) -> Result<LogicalPlan> {
61        let float_literal = |param: &PromExpr| -> Result<f64> {
62            let value = (|| {
63                let expr = Self::get_param_as_literal_expr(
64                    Some(param),
65                    None,
66                    Some(ArrowDataType::Float64),
67                )
68                .ok()?;
69                let simplifier = ExprSimplifier::new(SimplifyContext::default());
70                let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?;
71                let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else {
72                    return None;
73                };
74                let ScalarValue::Float64(Some(value)) =
75                    value.cast_to(&ArrowDataType::Float64).ok()?
76                else {
77                    return None;
78                };
79                Some(value)
80            })()
81            .with_context(|| FunctionInvalidArgumentSnafu {
82                fn_name: function_name.to_string(),
83            })?;
84            Ok(value)
85        };
86        let (function, input) = match (function_name, args.args.as_slice()) {
87            (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => (
88                HistogramFoldOperation::Quantile(float_literal(quantile)?.into()),
89                input.as_ref().clone(),
90            ),
91            (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => (
92                HistogramFoldOperation::Fraction {
93                    lower: float_literal(lower)?.into(),
94                    upper: float_literal(upper)?.into(),
95                },
96                input.as_ref().clone(),
97            ),
98            _ => {
99                return FunctionInvalidArgumentSnafu {
100                    fn_name: function_name.to_string(),
101                }
102                .fail();
103            }
104        };
105
106        let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?;
107        // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a
108        // stable series identifier anymore. HistogramFold must not treat it as a label column.
109        let input_plan = self.strip_tsid_column(input_plan)?;
110        self.ctx.use_tsid = false;
111
112        if let Some((float_field, histogram_field)) =
113            Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
114                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
115        {
116            if self.ctx.has_le_tag() {
117                return self.create_mixed_histogram_plan(
118                    function,
119                    input_plan,
120                    float_field,
121                    histogram_field,
122                );
123            }
124            self.ctx.field_columns = vec![histogram_field];
125        }
126        if self.all_field_columns_are_native_histograms(input_plan.schema()) {
127            return self.create_native_histogram_plan(function, input_plan);
128        }
129
130        if !self.ctx.has_le_tag() {
131            // Return empty result instead of error when 'le' column is not found
132            // This handles the case when histogram metrics don't exist
133            return Ok(LogicalPlan::EmptyRelation(
134                datafusion::logical_expr::EmptyRelation {
135                    produce_one_row: false,
136                    schema: input_plan.schema().clone(),
137                },
138            ));
139        }
140        let time_index_column =
141            self.ctx
142                .time_index_column
143                .clone()
144                .with_context(|| TimeIndexNotFoundSnafu {
145                    table: self.ctx.table_name.clone().unwrap_or_default(),
146                })?;
147        // FIXME(ruihang): support multi fields
148        let field_column = self
149            .ctx
150            .field_columns
151            .first()
152            .with_context(|| FunctionInvalidArgumentSnafu {
153                fn_name: function.function_name().to_string(),
154            })?
155            .clone();
156        // remove le column from tag columns
157        self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME);
158
159        let fold = HistogramFold::new_with_operation(
160            LE_COLUMN_NAME.to_string(),
161            field_column,
162            time_index_column,
163            function,
164            None,
165            input_plan,
166        )
167        .context(DataFusionPlanningSnafu)?;
168        Ok(LogicalPlan::Extension(Extension {
169            node: Arc::new(fold),
170        }))
171    }
172
173    fn create_native_histogram_expr(
174        &self,
175        function: HistogramFoldOperation,
176        field_column: &str,
177    ) -> DfExpr {
178        let field = DfExpr::Column(Column::from_name(field_column));
179        let (func, args) = match function {
180            HistogramFoldOperation::Quantile(quantile) => (
181                Arc::new(NativeHistogramQuantile::scalar_udf_with_collector(
182                    self.promql_annotations.clone(),
183                )),
184                vec![field, lit(f64::from(quantile))],
185            ),
186            HistogramFoldOperation::Fraction { lower, upper } => (
187                Arc::new(NativeHistogramFraction::scalar_udf_with_collector(
188                    self.promql_annotations.clone(),
189                )),
190                vec![field, lit(f64::from(lower)), lit(f64::from(upper))],
191            ),
192        };
193        DfExpr::ScalarFunction(ScalarFunction { func, args })
194    }
195
196    fn create_native_histogram_plan(
197        &mut self,
198        function: HistogramFoldOperation,
199        input_plan: LogicalPlan,
200    ) -> Result<LogicalPlan> {
201        ensure!(
202            self.ctx.field_columns.len() == 1,
203            MultiFieldsNotSupportedSnafu {
204                operator: function.function_name()
205            },
206        );
207
208        let field_column = self.ctx.field_columns[0].clone();
209        let function_expr = self.create_native_histogram_expr(function, &field_column);
210        let display_name = function_expr.schema_name().to_string();
211        self.ctx.field_columns = vec![display_name.clone()];
212
213        let project_exprs = std::iter::once(self.create_time_index_column_expr()?)
214            .chain(std::iter::once(function_expr.alias(display_name)))
215            .chain(self.create_tag_column_exprs()?)
216            .collect::<Vec<_>>();
217
218        LogicalPlanBuilder::from(input_plan)
219            .project(project_exprs)
220            .context(DataFusionPlanningSnafu)?
221            .filter(self.create_empty_values_filter_expr(false)?)
222            .context(DataFusionPlanningSnafu)?
223            .build()
224            .context(DataFusionPlanningSnafu)
225    }
226
227    fn create_mixed_histogram_plan(
228        &mut self,
229        function: HistogramFoldOperation,
230        input_plan: LogicalPlan,
231        float_field: String,
232        histogram_field: String,
233    ) -> Result<LogicalPlan> {
234        let time_index_column =
235            self.ctx
236                .time_index_column
237                .clone()
238                .with_context(|| TimeIndexNotFoundSnafu {
239                    table: self.ctx.table_name.clone().unwrap_or_default(),
240                })?;
241        let tag_columns = self.ctx.tag_columns.clone();
242        let folded = HistogramFold::new_with_operation(
243            LE_COLUMN_NAME.to_string(),
244            float_field.clone(),
245            time_index_column.clone(),
246            function,
247            Some(histogram_field.clone()),
248            input_plan,
249        )
250        .context(DataFusionPlanningSnafu)?;
251        let record_collision = DfExpr::ScalarFunction(ScalarFunction {
252            func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
253                "vector contains a mix of classic and native histograms".to_string(),
254                self.promql_annotations.clone(),
255            )),
256            args: vec![ident(&float_field), ident(&histogram_field)],
257        });
258        let keep = when(
259            ident(&float_field)
260                .is_not_null()
261                .and(ident(&histogram_field).is_not_null()),
262            record_collision,
263        )
264        .otherwise(lit(true))
265        .context(DataFusionPlanningSnafu)?;
266
267        let native_expr = self.create_native_histogram_expr(function, &histogram_field);
268        let output_field = native_expr.schema_name().to_string();
269        let value = DfExpr::ScalarFunction(ScalarFunction {
270            func: coalesce(),
271            args: vec![ident(&float_field), native_expr],
272        });
273        self.ctx.field_columns = vec![output_field.clone()];
274        LogicalPlanBuilder::from(LogicalPlan::Extension(Extension {
275            node: Arc::new(folded),
276        }))
277        .filter(keep)
278        .context(DataFusionPlanningSnafu)?
279        .project(
280            std::iter::once(ident(&time_index_column))
281                .chain(std::iter::once(value.alias(output_field)))
282                .chain(tag_columns.iter().map(ident)),
283        )
284        .context(DataFusionPlanningSnafu)?
285        .build()
286        .context(DataFusionPlanningSnafu)
287    }
288
289    /// Create a [SPECIAL_VECTOR_FUNCTION] plan
290    pub(super) async fn create_vector_plan(
291        &mut self,
292        args: &PromFunctionArgs,
293    ) -> Result<LogicalPlan> {
294        if args.args.len() != 1 {
295            return FunctionInvalidArgumentSnafu {
296                fn_name: SPECIAL_VECTOR_FUNCTION.to_string(),
297            }
298            .fail();
299        }
300        let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?;
301
302        // reuse `SPECIAL_TIME_FUNCTION` as name of time index column
303        self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
304        self.ctx.reset_table_name_and_schema();
305        self.ctx.tag_columns = vec![];
306        self.ctx.aggregation_field_labels.clear();
307        self.ctx.field_columns = vec![greptime_value().to_string()];
308        Ok(LogicalPlan::Extension(Extension {
309            node: Arc::new(
310                EmptyMetric::new(
311                    self.ctx.start,
312                    self.ctx.end,
313                    self.ctx.interval,
314                    SPECIAL_TIME_FUNCTION.to_string(),
315                    greptime_value().to_string(),
316                    Some(lit),
317                )
318                .context(DataFusionPlanningSnafu)?,
319            ),
320        }))
321    }
322
323    /// Create a [SCALAR_FUNCTION] plan
324    pub(super) async fn create_scalar_plan(
325        &mut self,
326        args: &PromFunctionArgs,
327        query_engine_state: &QueryEngineState,
328    ) -> Result<LogicalPlan> {
329        ensure!(
330            args.len() == 1,
331            FunctionInvalidArgumentSnafu {
332                fn_name: SCALAR_FUNCTION
333            }
334        );
335        let input = self
336            .prom_expr_to_plan(&args.args[0], query_engine_state)
337            .await?;
338        let input_schema = input.schema().clone();
339        let alternative_samples =
340            Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns);
341        let histogram_fields = self
342            .ctx
343            .field_columns
344            .iter()
345            .filter(|field| Self::field_column_is_native_histogram(&input_schema, field))
346            .count();
347        ensure!(
348            self.ctx.field_columns.len() == 1 || alternative_samples,
349            MultiFieldsNotSupportedSnafu {
350                operator: SCALAR_FUNCTION
351            },
352        );
353        let scalar_field = self
354            .ctx
355            .field_columns
356            .iter()
357            .find(|field| !Self::field_column_is_native_histogram(&input_schema, field))
358            .or_else(|| self.ctx.field_columns.first())
359            .cloned()
360            .with_context(|| FunctionInvalidArgumentSnafu {
361                fn_name: SCALAR_FUNCTION,
362            })?;
363        let input = if histogram_fields == self.ctx.field_columns.len() {
364            // scalar() ignores histogram samples. An empty input makes ScalarCalculate emit NaN
365            // for every evaluation timestamp without attempting a Struct-to-Float64 cast.
366            LogicalPlanBuilder::from(input)
367                .filter(lit(false))
368                .context(DataFusionPlanningSnafu)?
369                .build()
370                .context(DataFusionPlanningSnafu)?
371        } else if histogram_fields > 0 {
372            // A mixed vector contributes only its float samples to scalar().
373            LogicalPlanBuilder::from(input)
374                .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null())
375                .context(DataFusionPlanningSnafu)?
376                .build()
377                .context(DataFusionPlanningSnafu)?
378        } else {
379            input
380        };
381        let scalar_plan = LogicalPlan::Extension(Extension {
382            node: Arc::new(
383                ScalarCalculate::new(
384                    self.ctx.start,
385                    self.ctx.end,
386                    self.ctx.interval,
387                    input,
388                    self.ctx.time_index_column.as_ref().unwrap(),
389                    &self.ctx.tag_columns,
390                    &scalar_field,
391                    self.ctx.table_name.as_deref(),
392                )
393                .context(PromqlPlanNodeSnafu)?,
394            ),
395        });
396        // scalar plan have no tag columns
397        self.ctx.tag_columns.clear();
398        self.ctx.aggregation_field_labels.clear();
399        self.ctx.field_columns.clear();
400        self.ctx
401            .field_columns
402            .push(scalar_plan.schema().field(1).name().clone());
403        Ok(scalar_plan)
404    }
405
406    /// Create a [SPECIAL_ABSENT_FUNCTION] plan
407    pub(super) async fn create_absent_plan(
408        &mut self,
409        args: &PromFunctionArgs,
410        query_engine_state: &QueryEngineState,
411    ) -> Result<LogicalPlan> {
412        if args.args.len() != 1 {
413            return FunctionInvalidArgumentSnafu {
414                fn_name: SPECIAL_ABSENT_FUNCTION.to_string(),
415            }
416            .fail();
417        }
418        let input = self
419            .prom_expr_to_plan(&args.args[0], query_engine_state)
420            .await?;
421
422        let time_index_expr = self.create_time_index_column_expr()?;
423        let first_field_expr =
424            self.create_field_column_exprs()?
425                .pop()
426                .with_context(|| ValueNotFoundSnafu {
427                    table: self.ctx.table_name.clone().unwrap_or_default(),
428                })?;
429        let first_value_expr = first_value(first_field_expr, vec![]);
430
431        let ordered_aggregated_input = LogicalPlanBuilder::from(input)
432            .aggregate(
433                vec![time_index_expr.clone()],
434                vec![first_value_expr.clone()],
435            )
436            .context(DataFusionPlanningSnafu)?
437            .sort(vec![time_index_expr.sort(true, false)])
438            .context(DataFusionPlanningSnafu)?
439            .build()
440            .context(DataFusionPlanningSnafu)?;
441
442        let fake_labels = self
443            .ctx
444            .selector_matcher
445            .iter()
446            .filter_map(|matcher| match matcher.op {
447                MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())),
448                _ => None,
449            })
450            .collect::<Vec<_>>();
451
452        // Create the absent plan
453        let absent_plan = LogicalPlan::Extension(Extension {
454            node: Arc::new(
455                Absent::try_new(
456                    self.ctx.start,
457                    self.ctx.end,
458                    self.ctx.interval,
459                    self.ctx.time_index_column.as_ref().unwrap().clone(),
460                    self.ctx.field_columns[0].clone(),
461                    fake_labels,
462                    ordered_aggregated_input,
463                )
464                .context(DataFusionPlanningSnafu)?,
465            ),
466        });
467
468        // The absent series carries the equality matchers as labels, not the input's
469        // tags or value fields, so the input's field grouping labels no longer apply.
470        self.ctx.aggregation_field_labels.clear();
471        Ok(absent_plan)
472    }
473}
query/promql/planner/function_plans.rs

query/promql/planner/
function_plans.rs
