Skip to main content

query/promql/
planner.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::collections::{BTreeSet, HashMap, HashSet, VecDeque};
16use std::sync::Arc;
17use std::time::UNIX_EPOCH;
18
19use arrow::datatypes::IntervalDayTime;
20use async_recursion::async_recursion;
21use catalog::table_source::DfTableSourceProvider;
22use common_error::ext::ErrorExt;
23use common_error::status_code::StatusCode;
24use common_function::function::FunctionContext;
25use common_query::prelude::greptime_value;
26use datafusion::common::DFSchemaRef;
27use datafusion::datasource::DefaultTableSource;
28use datafusion::functions_aggregate::average::avg_udaf;
29use datafusion::functions_aggregate::count::count_udaf;
30use datafusion::functions_aggregate::expr_fn::first_value;
31use datafusion::functions_aggregate::min_max::{max_udaf, min_udaf};
32use datafusion::functions_aggregate::stddev::stddev_pop_udaf;
33use datafusion::functions_aggregate::sum::sum_udaf;
34use datafusion::functions_aggregate::variance::var_pop_udaf;
35use datafusion::functions_window::row_number::RowNumber;
36use datafusion::logical_expr::expr::{Alias, ScalarFunction, WindowFunction};
37use datafusion::logical_expr::expr_rewriter::normalize_cols;
38use datafusion::logical_expr::{
39    BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator,
40    ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition,
41};
42use datafusion::prelude as df_prelude;
43use datafusion::prelude::{Column, Expr as DfExpr, JoinType};
44use datafusion::scalar::ScalarValue;
45use datafusion::sql::TableReference;
46use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
47use datafusion_common::{DFSchema, NullEquality};
48use datafusion_expr::expr::WindowFunctionParams;
49use datafusion_expr::utils::conjunction;
50use datafusion_expr::{
51    ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit,
52};
53use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit};
54use datatypes::data_type::ConcreteDataType;
55use itertools::Itertools;
56use once_cell::sync::Lazy;
57use promql::extension_plan::{
58    Absent, EmptyMetric, HistogramFold, InstantManipulate, Millisecond, RangeManipulate,
59    ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, build_special_time_expr,
60};
61use promql::functions::{
62    AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing,
63    IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, PredictLinear, PresentOverTime,
64    QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime,
65    quantile_udaf,
66};
67use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers};
68use promql_parser::parser::token::TokenType;
69use promql_parser::parser::value::ValueType;
70use promql_parser::parser::{
71    AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt, Expr as PromExpr,
72    Function, FunctionArgs as PromFunctionArgs, LabelModifier, MatrixSelector, NumberLiteral,
73    Offset, ParenExpr, StringLiteral, SubqueryExpr, UnaryExpr, VectorMatchCardinality,
74    VectorSelector, token,
75};
76use regex::{self, Regex};
77use snafu::{OptionExt, ResultExt, ensure};
78use store_api::metric_engine_consts::{
79    DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
80    METRIC_ENGINE_NAME, is_metric_engine_internal_column,
81};
82use table::table::adapter::DfTableProviderAdapter;
83
84use crate::parser::{
85    ALIAS_NODE_NAME, ANALYZE_NODE_NAME, ANALYZE_VERBOSE_NODE_NAME, AliasExpr, EXPLAIN_NODE_NAME,
86    EXPLAIN_VERBOSE_NODE_NAME,
87};
88use crate::promql::error::{
89    CatalogSnafu, ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu,
90    ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu,
91    InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu,
92    MultipleMetricMatchersSnafu, MultipleVectorSnafu, NoMetricMatcherSnafu, PromqlPlanNodeSnafu,
93    Result, SameLabelSetSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu,
94    UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu,
95    UnsupportedMatcherOpSnafu, UnsupportedVectorMatchSnafu, ValueNotFoundSnafu,
96    ZeroRangeSelectorSnafu,
97};
98use crate::query_engine::QueryEngineState;
99
100/// `time()` function in PromQL.
101const SPECIAL_TIME_FUNCTION: &str = "time";
102/// `scalar()` function in PromQL.
103const SCALAR_FUNCTION: &str = "scalar";
104/// `absent()` function in PromQL
105const SPECIAL_ABSENT_FUNCTION: &str = "absent";
106/// `histogram_quantile` function in PromQL
107const SPECIAL_HISTOGRAM_QUANTILE: &str = "histogram_quantile";
108/// `vector` function in PromQL
109const SPECIAL_VECTOR_FUNCTION: &str = "vector";
110/// `le` column for conventional histogram.
111const LE_COLUMN_NAME: &str = "le";
112
113/// Static regex for validating label names according to Prometheus specification.
114/// Label names must match the regex: [a-zA-Z_][a-zA-Z0-9_]*
115static LABEL_NAME_REGEX: Lazy<Regex> =
116    Lazy::new(|| Regex::new(r"^[a-zA-Z_][a-zA-Z0-9_]*$").unwrap());
117
118const DEFAULT_TIME_INDEX_COLUMN: &str = "time";
119
120/// default value column name for empty metric
121const DEFAULT_FIELD_COLUMN: &str = "value";
122
123/// Special modifier to project field columns under multi-field mode
124const FIELD_COLUMN_MATCHER: &str = "__field__";
125
126/// Special modifier for cross schema query
127const SCHEMA_COLUMN_MATCHER: &str = "__schema__";
128const DB_COLUMN_MATCHER: &str = "__database__";
129
130/// Prefix for generated binary island leaf aliases.
131const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v";
132
133/// Threshold for scatter scan mode
134const MAX_SCATTER_POINTS: i64 = 400;
135
136/// Interval 1 hour in millisecond
137const INTERVAL_1H: i64 = 60 * 60 * 1000;
138
139#[derive(Default, Debug, Clone)]
140struct PromPlannerContext {
141    // query parameters
142    start: Millisecond,
143    end: Millisecond,
144    interval: Millisecond,
145    lookback_delta: Millisecond,
146
147    // planner states
148    table_name: Option<String>,
149    time_index_column: Option<String>,
150    field_columns: Vec<String>,
151    tag_columns: Vec<String>,
152    /// Use metric engine internal series identifier column (`__tsid`) as series key.
153    ///
154    /// This is enabled only when the underlying scan can provide `__tsid` (`UInt64`). The planner
155    /// uses it internally (e.g. as the series key for [`SeriesDivide`]) and strips it from the
156    /// final output.
157    use_tsid: bool,
158    /// The matcher for field columns `__field__`.
159    field_column_matcher: Option<Vec<Matcher>>,
160    /// The matcher for selectors (normal matchers).
161    selector_matcher: Vec<Matcher>,
162    schema_name: Option<String>,
163    /// The range in millisecond of range selector. None if there is no range selector.
164    range: Option<Millisecond>,
165}
166
167#[derive(Debug, Clone, PartialEq, Eq, Hash)]
168struct VectorLeafKey {
169    metric_name: String,
170    matchers: Vec<(String, String, String)>,
171    or_matchers: Vec<Vec<(String, String, String)>>,
172    offset_ms: i128,
173    at: String,
174}
175
176#[derive(Debug, Clone)]
177struct IslandLeaf {
178    selector: VectorSelector,
179    display_table: String,
180}
181
182#[derive(Debug, Clone)]
183enum IslandExpr {
184    VectorLeaf(usize),
185    Scalar(DfExpr),
186    Unary {
187        input: Box<IslandExpr>,
188    },
189    Binary {
190        op: TokenType,
191        lhs: Box<IslandExpr>,
192        rhs: Box<IslandExpr>,
193    },
194}
195
196impl IslandExpr {
197    fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option<Self> {
198        if let Some(expr) = PromPlanner::try_build_literal_expr(expr) {
199            return Some(Self::Scalar(expr));
200        }
201
202        match expr {
203            PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env),
204            PromExpr::VectorSelector(selector) => {
205                let leaf = env.intern_leaf(selector)?;
206                Some(Self::VectorLeaf(leaf))
207            }
208            PromExpr::Unary(UnaryExpr { expr }) => {
209                let input = Self::try_new(expr, env)?;
210                Some(Self::Unary {
211                    input: Box::new(input),
212                })
213            }
214            PromExpr::Binary(PromBinaryExpr {
215                lhs,
216                rhs,
217                op,
218                modifier,
219            }) if matches!(
220                op.id(),
221                token::T_ADD
222                    | token::T_SUB
223                    | token::T_MUL
224                    | token::T_DIV
225                    | token::T_MOD
226                    | token::T_POW
227                    | token::T_ATAN2
228            ) && modifier.as_ref().is_none_or(|modifier| {
229                !modifier.return_bool
230                    && modifier.matching.is_none()
231                    && matches!(modifier.card, VectorMatchCardinality::OneToOne)
232            }) =>
233            {
234                let lhs = Self::try_new(lhs, env)?;
235                let rhs = Self::try_new(rhs, env)?;
236                Some(Self::Binary {
237                    op: *op,
238                    lhs: Box::new(lhs),
239                    rhs: Box::new(rhs),
240                })
241            }
242            _ => None,
243        }
244    }
245}
246
247#[derive(Debug, Default)]
248struct IslandCollectEnv {
249    leaf_by_key: HashMap<VectorLeafKey, usize>,
250    leaves: Vec<IslandLeaf>,
251    vector_occurrences: usize,
252}
253
254#[derive(Debug)]
255struct PlannedIslandLeaf {
256    plan: LogicalPlan,
257    ctx: PromPlannerContext,
258    alias: TableReference,
259    display_table: String,
260}
261
262#[derive(Debug)]
263struct IslandFieldExprs {
264    exprs: Vec<DfExpr>,
265    names: Vec<String>,
266    scalar: bool,
267}
268
269impl VectorLeafKey {
270    fn from_selector(selector: &VectorSelector) -> Option<Self> {
271        let mut metric_name = selector.name.clone();
272        let mut matchers = Vec::with_capacity(selector.matchers.matchers.len());
273        let matcher_key = |matcher: &Matcher| {
274            (
275                matcher.name.clone(),
276                matcher.op.to_string(),
277                matcher.value.clone(),
278            )
279        };
280
281        for matcher in &selector.matchers.matchers {
282            if matcher.name == METRIC_NAME {
283                if matcher.op != MatchOp::Equal || metric_name.is_some() {
284                    return None;
285                }
286                metric_name = Some(matcher.value.clone());
287            } else {
288                matchers.push(matcher_key(matcher));
289            }
290        }
291        matchers.sort();
292
293        let mut or_matchers = selector
294            .matchers
295            .or_matchers
296            .iter()
297            .map(|group| {
298                let mut group = group.iter().map(matcher_key).collect::<Vec<_>>();
299                group.sort();
300                group
301            })
302            .collect::<Vec<_>>();
303        or_matchers.sort();
304
305        Some(Self {
306            metric_name: metric_name?,
307            matchers,
308            or_matchers,
309            offset_ms: match &selector.offset {
310                Some(Offset::Pos(duration)) => duration.as_millis() as i128,
311                Some(Offset::Neg(duration)) => -(duration.as_millis() as i128),
312                None => 0,
313            },
314            at: format!("{:?}", selector.at),
315        })
316    }
317}
318
319impl IslandCollectEnv {
320    fn intern_leaf(&mut self, selector: &VectorSelector) -> Option<usize> {
321        self.vector_occurrences += 1;
322        let key = VectorLeafKey::from_selector(selector)?;
323        if let Some(id) = self.leaf_by_key.get(&key) {
324            return Some(*id);
325        }
326
327        let id = self.leaves.len();
328        self.leaves.push(IslandLeaf {
329            selector: selector.clone(),
330            display_table: key.metric_name.clone(),
331        });
332        self.leaf_by_key.insert(key, id);
333        Some(id)
334    }
335}
336
337impl PromPlannerContext {
338    fn from_eval_stmt(stmt: &EvalStmt) -> Self {
339        Self {
340            start: stmt.start.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
341            end: stmt.end.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
342            interval: stmt.interval.as_millis() as _,
343            lookback_delta: stmt.lookback_delta.as_millis() as _,
344            ..Default::default()
345        }
346    }
347
348    /// Reset all planner states
349    fn reset(&mut self) {
350        self.table_name = None;
351        self.time_index_column = None;
352        self.field_columns = vec![];
353        self.tag_columns = vec![];
354        self.use_tsid = false;
355        self.field_column_matcher = None;
356        self.selector_matcher.clear();
357        self.schema_name = None;
358        self.range = None;
359    }
360
361    /// Reset table name and schema to empty
362    fn reset_table_name_and_schema(&mut self) {
363        self.table_name = Some(String::new());
364        self.schema_name = None;
365        self.use_tsid = false;
366    }
367
368    /// Check if `le` is present in tag columns
369    fn has_le_tag(&self) -> bool {
370        self.tag_columns.iter().any(|c| c.eq(&LE_COLUMN_NAME))
371    }
372}
373
374pub struct PromPlanner {
375    table_provider: DfTableSourceProvider,
376    ctx: PromPlannerContext,
377}
378
379impl PromPlanner {
380    pub async fn stmt_to_plan(
381        table_provider: DfTableSourceProvider,
382        stmt: &EvalStmt,
383        query_engine_state: &QueryEngineState,
384    ) -> Result<LogicalPlan> {
385        let mut planner = Self {
386            table_provider,
387            ctx: PromPlannerContext::from_eval_stmt(stmt),
388        };
389
390        let plan = planner
391            .prom_expr_to_plan(&stmt.expr, query_engine_state)
392            .await?;
393
394        // Never leak internal series identifier to output.
395        planner.strip_tsid_column(plan)
396    }
397
398    pub async fn prom_expr_to_plan(
399        &mut self,
400        prom_expr: &PromExpr,
401        query_engine_state: &QueryEngineState,
402    ) -> Result<LogicalPlan> {
403        self.prom_expr_to_plan_inner(prom_expr, false, query_engine_state)
404            .await
405    }
406
407    /**
408    Converts a PromQL expression to a logical plan.
409
410    NOTE:
411        The `timestamp_fn` indicates whether the PromQL `timestamp()` function is being evaluated in the current context.
412        If `true`, the planner generates a logical plan that projects the timestamp (time index) column
413        as the value column for each input row, implementing the PromQL `timestamp()` function semantics.
414        If `false`, the planner generates the standard logical plan for the given PromQL expression.
415    */
416    #[async_recursion]
417    async fn prom_expr_to_plan_inner(
418        &mut self,
419        prom_expr: &PromExpr,
420        timestamp_fn: bool,
421        query_engine_state: &QueryEngineState,
422    ) -> Result<LogicalPlan> {
423        let res = match prom_expr {
424            PromExpr::Aggregate(expr) => {
425                self.prom_aggr_expr_to_plan(query_engine_state, expr)
426                    .await?
427            }
428            PromExpr::Unary(expr) => {
429                self.prom_unary_expr_to_plan(query_engine_state, expr)
430                    .await?
431            }
432            PromExpr::Binary(expr) => {
433                self.prom_binary_expr_to_plan(query_engine_state, expr)
434                    .await?
435            }
436            PromExpr::Paren(ParenExpr { expr }) => {
437                self.prom_expr_to_plan_inner(expr, timestamp_fn, query_engine_state)
438                    .await?
439            }
440            PromExpr::Subquery(expr) => {
441                self.prom_subquery_expr_to_plan(query_engine_state, expr)
442                    .await?
443            }
444            PromExpr::NumberLiteral(lit) => self.prom_number_lit_to_plan(lit)?,
445            PromExpr::StringLiteral(lit) => self.prom_string_lit_to_plan(lit)?,
446            PromExpr::VectorSelector(selector) => {
447                self.prom_vector_selector_to_plan(selector, timestamp_fn)
448                    .await?
449            }
450            PromExpr::MatrixSelector(selector) => {
451                self.prom_matrix_selector_to_plan(selector).await?
452            }
453            PromExpr::Call(expr) => {
454                self.prom_call_expr_to_plan(query_engine_state, expr)
455                    .await?
456            }
457            PromExpr::Extension(expr) => {
458                self.prom_ext_expr_to_plan(query_engine_state, expr).await?
459            }
460        };
461
462        Ok(res)
463    }
464
465    async fn prom_subquery_expr_to_plan(
466        &mut self,
467        query_engine_state: &QueryEngineState,
468        subquery_expr: &SubqueryExpr,
469    ) -> Result<LogicalPlan> {
470        let SubqueryExpr {
471            expr, range, step, ..
472        } = subquery_expr;
473
474        let current_interval = self.ctx.interval;
475        if let Some(step) = step {
476            self.ctx.interval = step.as_millis() as _;
477        }
478        let current_start = self.ctx.start;
479        self.ctx.start -= range.as_millis() as i64 - self.ctx.interval;
480        let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
481        self.ctx.interval = current_interval;
482        self.ctx.start = current_start;
483
484        ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
485        let range_ms = range.as_millis() as _;
486        self.ctx.range = Some(range_ms);
487
488        let time_index_column =
489            self.ctx
490                .time_index_column
491                .clone()
492                .with_context(|| TimeIndexNotFoundSnafu {
493                    table: self.ctx.table_name.clone().unwrap_or_default(),
494                })?;
495
496        // `RangeManipulate` assumes each input batch holds exactly one series
497        // (it takes tag column values from row 0 and applies them to every
498        // output row). The inner expression may emit batches that mix series,
499        // so sort by series key + time index and split into per-series batches
500        // with a `SeriesDivide` first.
501        let input_schema = input.schema();
502        let input_has_tsid = input_schema.fields().iter().any(|field| {
503            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
504                && field.data_type() == &ArrowDataType::UInt64
505        });
506        let (series_key_columns, mut sort_exprs) = if input_has_tsid {
507            (
508                vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()],
509                vec![
510                    DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME))
511                        .sort(true, true),
512                ],
513            )
514        } else {
515            // Only use tag columns that survive in the inner plan's schema —
516            // `ctx.tag_columns` can drift from the actual output.
517            let key_columns: Vec<String> = self
518                .ctx
519                .tag_columns
520                .iter()
521                .filter(|name| input_schema.has_column_with_unqualified_name(name))
522                .cloned()
523                .collect();
524            let sort = key_columns
525                .iter()
526                .map(|name| DfExpr::Column(Column::from_name(name)).sort(true, true))
527                .collect::<Vec<_>>();
528            (key_columns, sort)
529        };
530        sort_exprs.push(DfExpr::Column(Column::from_name(&time_index_column)).sort(true, true));
531
532        let sort_plan = LogicalPlanBuilder::from(input)
533            .sort(sort_exprs)
534            .context(DataFusionPlanningSnafu)?
535            .build()
536            .context(DataFusionPlanningSnafu)?;
537        let divide_plan = LogicalPlan::Extension(Extension {
538            node: Arc::new(SeriesDivide::new(
539                series_key_columns,
540                time_index_column.clone(),
541                sort_plan,
542            )),
543        });
544
545        let manipulate = RangeManipulate::new(
546            self.ctx.start,
547            self.ctx.end,
548            self.ctx.interval,
549            range_ms,
550            time_index_column,
551            self.ctx.field_columns.clone(),
552            divide_plan,
553        )
554        .context(DataFusionPlanningSnafu)?;
555
556        Ok(LogicalPlan::Extension(Extension {
557            node: Arc::new(manipulate),
558        }))
559    }
560
561    async fn prom_aggr_expr_to_plan(
562        &mut self,
563        query_engine_state: &QueryEngineState,
564        aggr_expr: &AggregateExpr,
565    ) -> Result<LogicalPlan> {
566        let AggregateExpr {
567            op,
568            expr,
569            modifier,
570            param,
571        } = aggr_expr;
572
573        let mut input = self.prom_expr_to_plan(expr, query_engine_state).await?;
574        let input_has_tsid = input.schema().fields().iter().any(|field| {
575            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
576                && field.data_type() == &ArrowDataType::UInt64
577        });
578
579        // `__tsid` based scan projection may prune tag columns. Ensure tags referenced in
580        // aggregation modifiers (`by`/`without`) are available before planning group keys.
581        let required_group_tags = match modifier {
582            None => BTreeSet::new(),
583            Some(LabelModifier::Include(labels)) => labels
584                .labels
585                .iter()
586                .filter(|label| !is_metric_engine_internal_column(label.as_str()))
587                .cloned()
588                .collect(),
589            Some(LabelModifier::Exclude(labels)) => {
590                let mut all_tags = self.collect_row_key_tag_columns_from_plan(&input)?;
591                for label in &labels.labels {
592                    let _ = all_tags.remove(label);
593                }
594                all_tags
595            }
596        };
597
598        if !required_group_tags.is_empty()
599            && required_group_tags
600                .iter()
601                .any(|tag| Self::find_case_sensitive_column(input.schema(), tag.as_str()).is_none())
602        {
603            input = self.ensure_tag_columns_available(input, &required_group_tags)?;
604            self.refresh_tag_columns_from_schema(input.schema());
605        }
606
607        match (*op).id() {
608            token::T_TOPK | token::T_BOTTOMK => {
609                self.prom_topk_bottomk_to_plan(aggr_expr, input).await
610            }
611            _ => {
612                // When `__tsid` is available, tag columns may have been pruned from the input plan.
613                // For `keep_tsid` decision we should compare against the full row-key label set,
614                // otherwise we may incorrectly treat label-reducing aggregates as preserving labels.
615                let input_tag_columns = if input_has_tsid {
616                    self.collect_row_key_tag_columns_from_plan(&input)?
617                        .into_iter()
618                        .collect::<Vec<_>>()
619                } else {
620                    self.ctx.tag_columns.clone()
621                };
622                // calculate columns to group by
623                // Need to append time index column into group by columns
624                let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
625                // convert op and value columns to aggregate exprs
626                let (mut aggr_exprs, prev_field_exprs) =
627                    self.create_aggregate_exprs(*op, param, &input)?;
628
629                let keep_tsid = op.id() != token::T_COUNT_VALUES
630                    && input_has_tsid
631                    && input_tag_columns.iter().collect::<HashSet<_>>()
632                        == self.ctx.tag_columns.iter().collect::<HashSet<_>>();
633
634                if keep_tsid {
635                    aggr_exprs.push(
636                        first_value(
637                            DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)),
638                            vec![],
639                        )
640                        .alias(DATA_SCHEMA_TSID_COLUMN_NAME),
641                    );
642                }
643                self.ctx.use_tsid = keep_tsid;
644
645                // create plan
646                let builder = LogicalPlanBuilder::from(input);
647                let builder = if op.id() == token::T_COUNT_VALUES {
648                    let label = Self::get_param_value_as_str(*op, param)?;
649                    // `count_values` must be grouped by fields,
650                    // and project the fields to the new label.
651                    group_exprs.extend(prev_field_exprs.clone());
652                    let project_fields = self
653                        .create_field_column_exprs()?
654                        .into_iter()
655                        .chain(self.create_tag_column_exprs()?)
656                        .chain(Some(self.create_time_index_column_expr()?))
657                        .chain(prev_field_exprs.into_iter().map(|expr| expr.alias(label)));
658
659                    builder
660                        .aggregate(group_exprs.clone(), aggr_exprs)
661                        .context(DataFusionPlanningSnafu)?
662                        .project(project_fields)
663                        .context(DataFusionPlanningSnafu)?
664                } else {
665                    builder
666                        .aggregate(group_exprs.clone(), aggr_exprs)
667                        .context(DataFusionPlanningSnafu)?
668                };
669
670                let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false));
671
672                builder
673                    .sort(sort_expr)
674                    .context(DataFusionPlanningSnafu)?
675                    .build()
676                    .context(DataFusionPlanningSnafu)
677            }
678        }
679    }
680
681    /// Create logical plan for PromQL topk and bottomk expr.
682    async fn prom_topk_bottomk_to_plan(
683        &mut self,
684        aggr_expr: &AggregateExpr,
685        input: LogicalPlan,
686    ) -> Result<LogicalPlan> {
687        let AggregateExpr {
688            op,
689            param,
690            modifier,
691            ..
692        } = aggr_expr;
693
694        let input_has_tsid = input.schema().fields().iter().any(|field| {
695            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
696                && field.data_type() == &ArrowDataType::UInt64
697        });
698        self.ctx.use_tsid = input_has_tsid;
699
700        let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?;
701
702        let val = Self::get_param_as_literal_expr(param, Some(*op), Some(ArrowDataType::Float64))?;
703
704        // convert op and value columns to window exprs.
705        let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?;
706
707        let rank_columns: Vec<_> = window_exprs
708            .iter()
709            .map(|expr| expr.schema_name().to_string())
710            .collect();
711
712        // Create ranks filter with `Operator::Or`.
713        // Safety: at least one rank column
714        let filter: DfExpr = rank_columns
715            .iter()
716            .fold(None, |expr, rank| {
717                let predicate = DfExpr::BinaryExpr(BinaryExpr {
718                    left: Box::new(col(rank)),
719                    op: Operator::LtEq,
720                    right: Box::new(val.clone()),
721                });
722
723                match expr {
724                    None => Some(predicate),
725                    Some(expr) => Some(DfExpr::BinaryExpr(BinaryExpr {
726                        left: Box::new(expr),
727                        op: Operator::Or,
728                        right: Box::new(predicate),
729                    })),
730                }
731            })
732            .unwrap();
733
734        let rank_columns: Vec<_> = rank_columns.into_iter().map(col).collect();
735
736        let mut new_group_exprs = group_exprs.clone();
737        // Order by ranks
738        new_group_exprs.extend(rank_columns);
739
740        let group_sort_expr = new_group_exprs
741            .into_iter()
742            .map(|expr| expr.sort(true, false));
743
744        let project_fields = self
745            .create_field_column_exprs()?
746            .into_iter()
747            .chain(self.create_tag_column_exprs()?)
748            .chain(
749                self.ctx
750                    .use_tsid
751                    .then_some(DfExpr::Column(Column::from_name(
752                        DATA_SCHEMA_TSID_COLUMN_NAME,
753                    ))),
754            )
755            .chain(Some(self.create_time_index_column_expr()?));
756
757        LogicalPlanBuilder::from(input)
758            .window(window_exprs)
759            .context(DataFusionPlanningSnafu)?
760            .filter(filter)
761            .context(DataFusionPlanningSnafu)?
762            .sort(group_sort_expr)
763            .context(DataFusionPlanningSnafu)?
764            .project(project_fields)
765            .context(DataFusionPlanningSnafu)?
766            .build()
767            .context(DataFusionPlanningSnafu)
768    }
769
770    async fn prom_unary_expr_to_plan(
771        &mut self,
772        query_engine_state: &QueryEngineState,
773        unary_expr: &UnaryExpr,
774    ) -> Result<LogicalPlan> {
775        let UnaryExpr { expr } = unary_expr;
776        // Unary Expr in PromQL implys the `-` operator
777        let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
778        self.projection_for_each_field_column(input, |col| {
779            Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into()))))
780        })
781    }
782
783    async fn try_plan_binary_island(
784        &mut self,
785        binary_expr: &PromBinaryExpr,
786    ) -> Result<Option<LogicalPlan>> {
787        let original_ctx = self.ctx.clone();
788        let mut collect_env = IslandCollectEnv::default();
789        let Some(island_expr) =
790            IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env)
791        else {
792            return Ok(None);
793        };
794
795        if collect_env.leaves.is_empty()
796            || collect_env.vector_occurrences <= collect_env.leaves.len()
797        {
798            return Ok(None);
799        }
800
801        let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len());
802        for (idx, leaf) in collect_env.leaves.iter().enumerate() {
803            let plan = self
804                .prom_vector_selector_to_plan(&leaf.selector, false)
805                .await?;
806            let ctx = self.ctx.clone();
807            let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}"));
808            let plan = LogicalPlanBuilder::from(plan)
809                .alias(alias.clone())
810                .context(DataFusionPlanningSnafu)?
811                .build()
812                .context(DataFusionPlanningSnafu)?;
813            planned_leaves.push(PlannedIslandLeaf {
814                plan,
815                ctx,
816                alias,
817                display_table: leaf.display_table.clone(),
818            });
819        }
820
821        if !Self::binary_island_join_contexts_supported(&planned_leaves) {
822            self.ctx = original_ctx;
823            return Ok(None);
824        }
825
826        let mut input = planned_leaves[0].plan.clone();
827        for right_idx in 1..planned_leaves.len() {
828            input = self.join_binary_island_leaf(
829                input,
830                &planned_leaves[0],
831                &planned_leaves[right_idx],
832            )?;
833        }
834
835        let field_exprs =
836            Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?;
837        if field_exprs.scalar || field_exprs.exprs.is_empty() {
838            self.ctx = original_ctx;
839            return Ok(None);
840        }
841
842        let plan = self.project_binary_island(
843            input,
844            &planned_leaves[0].alias,
845            &planned_leaves[0].ctx,
846            field_exprs,
847        )?;
848        Ok(Some(plan))
849    }
850
851    fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool {
852        if leaves
853            .iter()
854            .any(|leaf| leaf.ctx.time_index_column.is_none())
855        {
856            return false;
857        }
858
859        if leaves.len() <= 1 {
860            return true;
861        }
862
863        let first_tags = leaves[0].ctx.tag_columns.iter().collect::<BTreeSet<_>>();
864
865        leaves.iter().skip(1).all(|leaf| {
866            (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan))
867                || leaf.ctx.tag_columns.iter().collect::<BTreeSet<_>>() == first_tags
868        })
869    }
870
871    fn join_binary_island_leaf(
872        &self,
873        left: LogicalPlan,
874        first_leaf: &PlannedIslandLeaf,
875        right_leaf: &PlannedIslandLeaf,
876    ) -> Result<LogicalPlan> {
877        let only_join_time_index =
878            first_leaf.ctx.tag_columns.is_empty() || right_leaf.ctx.tag_columns.is_empty();
879        let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns(
880            left.schema(),
881            right_leaf.plan.schema(),
882            &first_leaf.ctx,
883            &right_leaf.ctx,
884            only_join_time_index,
885            &None,
886        )?;
887
888        if let (Some(left_time_index_column), Some(right_time_index_column)) = (
889            first_leaf.ctx.time_index_column.clone(),
890            right_leaf.ctx.time_index_column.clone(),
891        ) {
892            left_keys.insert(left_time_index_column);
893            right_keys.insert(right_time_index_column);
894        }
895
896        LogicalPlanBuilder::from(left)
897            .join_detailed(
898                right_leaf.plan.clone(),
899                JoinType::Inner,
900                (
901                    left_keys
902                        .into_iter()
903                        .map(|name| Column::new(Some(first_leaf.alias.clone()), name))
904                        .collect::<Vec<_>>(),
905                    right_keys
906                        .into_iter()
907                        .map(|name| Column::new(Some(right_leaf.alias.clone()), name))
908                        .collect::<Vec<_>>(),
909                ),
910                force_empty_join.then_some(lit(false)),
911                NullEquality::NullEqualsNull,
912            )
913            .context(DataFusionPlanningSnafu)?
914            .build()
915            .context(DataFusionPlanningSnafu)
916    }
917
918    fn build_binary_island_field_exprs(
919        expr: &IslandExpr,
920        leaves: &[PlannedIslandLeaf],
921        schema: &DFSchemaRef,
922    ) -> Result<IslandFieldExprs> {
923        match expr {
924            IslandExpr::VectorLeaf(id) => {
925                let leaf = &leaves[*id];
926                let exprs = leaf
927                    .ctx
928                    .field_columns
929                    .iter()
930                    .map(|field| {
931                        schema
932                            .qualified_field_with_name(Some(&leaf.alias), field)
933                            .context(DataFusionPlanningSnafu)
934                            .map(|field| DfExpr::Column(field.into()))
935                    })
936                    .collect::<Result<Vec<_>>>()?;
937                let names = leaf
938                    .ctx
939                    .field_columns
940                    .iter()
941                    .map(|field| format!("{}.{}", leaf.display_table, field))
942                    .collect();
943                Ok(IslandFieldExprs {
944                    exprs,
945                    names,
946                    scalar: false,
947                })
948            }
949            IslandExpr::Scalar(expr) => Ok(IslandFieldExprs {
950                exprs: vec![expr.clone()],
951                names: vec![expr.schema_name().to_string()],
952                scalar: true,
953            }),
954            IslandExpr::Unary { input } => {
955                let input = Self::build_binary_island_field_exprs(input, leaves, schema)?;
956                let mut exprs = Vec::with_capacity(input.exprs.len());
957                let mut names = Vec::with_capacity(input.names.len());
958                for (expr, name) in input.exprs.into_iter().zip(input.names) {
959                    exprs.push(DfExpr::Negative(Box::new(expr)));
960                    names.push(format!("-{name}"));
961                }
962                Ok(IslandFieldExprs {
963                    exprs,
964                    names,
965                    scalar: input.scalar,
966                })
967            }
968            IslandExpr::Binary { op, lhs, rhs } => {
969                let same_leaf = match (&**lhs, &**rhs) {
970                    (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right))
971                        if left == right =>
972                    {
973                        Some(*left)
974                    }
975                    _ => None,
976                };
977                let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?;
978                let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?;
979                let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
980                let scalar = lhs.scalar && rhs.scalar;
981                let op = op.to_string();
982
983                let (exprs, names) = match (lhs.scalar, rhs.scalar) {
984                    (true, true) => {
985                        let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?;
986                        let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]);
987                        (vec![expr], vec![name])
988                    }
989                    (true, false) => {
990                        let mut exprs = Vec::with_capacity(rhs.exprs.len());
991                        let mut names = Vec::with_capacity(rhs.names.len());
992                        for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) {
993                            exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?);
994                            names.push(format!("{} {op} {rhs_name}", lhs.names[0]));
995                        }
996                        (exprs, names)
997                    }
998                    (false, true) => {
999                        let mut exprs = Vec::with_capacity(lhs.exprs.len());
1000                        let mut names = Vec::with_capacity(lhs.names.len());
1001                        for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) {
1002                            exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?);
1003                            names.push(format!("{lhs_name} {op} {}", rhs.names[0]));
1004                        }
1005                        (exprs, names)
1006                    }
1007                    (false, false) => {
1008                        let mut exprs = Vec::new();
1009                        let mut names = Vec::new();
1010                        for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs
1011                            .exprs
1012                            .into_iter()
1013                            .zip(rhs.exprs)
1014                            .zip(lhs.names.into_iter().zip(rhs.names))
1015                            .enumerate()
1016                        {
1017                            if let Some(leaf) = same_leaf {
1018                                let field = leaves[leaf]
1019                                    .ctx
1020                                    .field_columns
1021                                    .get(idx)
1022                                    .cloned()
1023                                    .unwrap_or_else(|| lhs_name.clone());
1024                                lhs_name = format!("lhs.{field}");
1025                                rhs_name = format!("rhs.{field}");
1026                            }
1027                            exprs.push(expr_builder(lhs_expr, rhs_expr)?);
1028                            names.push(format!("{lhs_name} {op} {rhs_name}"));
1029                        }
1030                        (exprs, names)
1031                    }
1032                };
1033
1034                Ok(IslandFieldExprs {
1035                    exprs,
1036                    names,
1037                    scalar,
1038                })
1039            }
1040        }
1041    }
1042
1043    fn project_binary_island(
1044        &mut self,
1045        input: LogicalPlan,
1046        base_alias: &TableReference,
1047        base_ctx: &PromPlannerContext,
1048        field_exprs: IslandFieldExprs,
1049    ) -> Result<LogicalPlan> {
1050        self.ctx = base_ctx.clone();
1051
1052        let schema = input.schema();
1053        let non_field_exprs = base_ctx
1054            .tag_columns
1055            .iter()
1056            .chain(base_ctx.time_index_column.iter())
1057            .map(|column| {
1058                schema
1059                    .qualified_field_with_name(Some(base_alias), column)
1060                    .context(DataFusionPlanningSnafu)
1061                    .map(|field| DfExpr::Column(field.into()))
1062            });
1063        let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid)
1064            .into_iter()
1065            .map(Ok);
1066
1067        self.ctx.field_columns = field_exprs.names;
1068        let field_exprs = field_exprs
1069            .exprs
1070            .into_iter()
1071            .zip(self.ctx.field_columns.iter())
1072            .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
1073
1074        let project_exprs = non_field_exprs
1075            .chain(tsid_expr)
1076            .chain(field_exprs)
1077            .collect::<Result<Vec<_>>>()?;
1078
1079        let plan = LogicalPlanBuilder::from(input)
1080            .project(project_exprs)
1081            .context(DataFusionPlanningSnafu)?
1082            .build()
1083            .context(DataFusionPlanningSnafu)?;
1084
1085        self.ctx.table_name = None;
1086        self.ctx.schema_name = None;
1087
1088        Ok(plan)
1089    }
1090
1091    async fn prom_binary_expr_to_plan(
1092        &mut self,
1093        query_engine_state: &QueryEngineState,
1094        binary_expr: &PromBinaryExpr,
1095    ) -> Result<LogicalPlan> {
1096        if let Some(plan) = self.try_plan_binary_island(binary_expr).await? {
1097            return Ok(plan);
1098        }
1099
1100        let PromBinaryExpr {
1101            lhs,
1102            rhs,
1103            op,
1104            modifier,
1105        } = binary_expr;
1106
1107        // if set to true, comparison operator will return 0/1 (for true/false) instead of
1108        // filter on the result column
1109        let should_return_bool = if let Some(m) = modifier {
1110            m.return_bool
1111        } else {
1112            false
1113        };
1114        let is_comparison_op = Self::is_token_a_comparison_op(*op);
1115
1116        // we should build a filter plan here if the op is comparison op and need not
1117        // to return 0/1. Otherwise, we should build a projection plan
1118        match (
1119            Self::try_build_literal_expr(lhs),
1120            Self::try_build_literal_expr(rhs),
1121        ) {
1122            (Some(lhs), Some(rhs)) => {
1123                self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1124                self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1125                self.ctx.reset_table_name_and_schema();
1126                let field_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1127                let mut field_expr = field_expr_builder(lhs, rhs)?;
1128
1129                if is_comparison_op && should_return_bool {
1130                    field_expr = DfExpr::Cast(Cast {
1131                        expr: Box::new(field_expr),
1132                        data_type: ArrowDataType::Float64,
1133                    });
1134                }
1135
1136                Ok(LogicalPlan::Extension(Extension {
1137                    node: Arc::new(
1138                        EmptyMetric::new(
1139                            self.ctx.start,
1140                            self.ctx.end,
1141                            self.ctx.interval,
1142                            SPECIAL_TIME_FUNCTION.to_string(),
1143                            DEFAULT_FIELD_COLUMN.to_string(),
1144                            Some(field_expr),
1145                        )
1146                        .context(DataFusionPlanningSnafu)?,
1147                    ),
1148                }))
1149            }
1150            // lhs is a literal, rhs is a column
1151            (Some(mut expr), None) => {
1152                let input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1153                // check if the literal is a special time expr
1154                if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) {
1155                    expr = time_expr
1156                }
1157                let bin_expr_builder = |col: &String| {
1158                    let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1159                    let mut binary_expr =
1160                        binary_expr_builder(expr.clone(), DfExpr::Column(col.into()))?;
1161
1162                    if is_comparison_op && should_return_bool {
1163                        binary_expr = DfExpr::Cast(Cast {
1164                            expr: Box::new(binary_expr),
1165                            data_type: ArrowDataType::Float64,
1166                        });
1167                    }
1168                    Ok(binary_expr)
1169                };
1170                if is_comparison_op && !should_return_bool {
1171                    self.filter_on_field_column(input, bin_expr_builder)
1172                } else {
1173                    self.projection_for_each_field_column(input, bin_expr_builder)
1174                }
1175            }
1176            // lhs is a column, rhs is a literal
1177            (None, Some(mut expr)) => {
1178                let input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1179                // check if the literal is a special time expr
1180                if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) {
1181                    expr = time_expr
1182                }
1183                let bin_expr_builder = |col: &String| {
1184                    let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1185                    let mut binary_expr =
1186                        binary_expr_builder(DfExpr::Column(col.into()), expr.clone())?;
1187
1188                    if is_comparison_op && should_return_bool {
1189                        binary_expr = DfExpr::Cast(Cast {
1190                            expr: Box::new(binary_expr),
1191                            data_type: ArrowDataType::Float64,
1192                        });
1193                    }
1194                    Ok(binary_expr)
1195                };
1196                if is_comparison_op && !should_return_bool {
1197                    self.filter_on_field_column(input, bin_expr_builder)
1198                } else {
1199                    self.projection_for_each_field_column(input, bin_expr_builder)
1200                }
1201            }
1202            // both are columns. join them on time index
1203            (None, None) => {
1204                let left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1205                let left_field_columns = self.ctx.field_columns.clone();
1206                let left_time_index_column = self.ctx.time_index_column.clone();
1207                let mut left_table_ref = self
1208                    .table_ref()
1209                    .unwrap_or_else(|_| TableReference::bare(""));
1210                let left_context = self.ctx.clone();
1211
1212                let right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1213                let right_field_columns = self.ctx.field_columns.clone();
1214                let right_time_index_column = self.ctx.time_index_column.clone();
1215                let mut right_table_ref = self
1216                    .table_ref()
1217                    .unwrap_or_else(|_| TableReference::bare(""));
1218                let right_context = self.ctx.clone();
1219
1220                // TODO(ruihang): avoid join if left and right are the same table
1221
1222                // set op has "special" join semantics
1223                if Self::is_token_a_set_op(*op) {
1224                    return self.set_op_on_non_field_columns(
1225                        left_input,
1226                        right_input,
1227                        left_context,
1228                        right_context,
1229                        *op,
1230                        modifier,
1231                    );
1232                }
1233
1234                // normal join
1235                if left_table_ref == right_table_ref {
1236                    // rename table references to avoid ambiguity
1237                    left_table_ref = TableReference::bare("lhs");
1238                    right_table_ref = TableReference::bare("rhs");
1239                    // `self.ctx` have ctx in right plan, if right plan have no tag,
1240                    // we use left plan ctx as the ctx for subsequent calculations,
1241                    // to avoid case like `host + scalar(...)`
1242                    // we need preserve tag column on `host` table in subsequent projection,
1243                    // which only show in left plan ctx.
1244                    if self.ctx.tag_columns.is_empty() {
1245                        self.ctx = left_context.clone();
1246                        self.ctx.table_name = Some("lhs".to_string());
1247                    } else {
1248                        self.ctx.table_name = Some("rhs".to_string());
1249                    }
1250                }
1251                let (output_field_columns, field_columns) =
1252                    Self::align_binary_field_columns(&left_field_columns, &right_field_columns);
1253                let left_aligned_field_columns = field_columns
1254                    .iter()
1255                    .map(|(left_col_name, _)| (*left_col_name).clone())
1256                    .collect::<Vec<_>>();
1257                let right_aligned_field_columns = field_columns
1258                    .iter()
1259                    .map(|(_, right_col_name)| (*right_col_name).clone())
1260                    .collect::<Vec<_>>();
1261                // PromQL binary arithmetic only combines the shared prefix of value columns.
1262                // Keep the output field count aligned with that zipped prefix so planning
1263                // remains stable even when the two sides have uneven multi-field schemas.
1264                self.ctx.field_columns = output_field_columns;
1265                let mut field_columns = field_columns.into_iter();
1266
1267                let join_plan = self.join_on_non_field_columns(
1268                    left_input,
1269                    right_input,
1270                    left_table_ref.clone(),
1271                    right_table_ref.clone(),
1272                    left_time_index_column,
1273                    right_time_index_column,
1274                    // if left plan or right plan tag is empty, means case like `scalar(...) + host` or `host + scalar(...)`
1275                    // under this case we only join on time index
1276                    left_context.tag_columns.is_empty() || right_context.tag_columns.is_empty(),
1277                    modifier,
1278                    &left_context,
1279                    &right_context,
1280                )?;
1281                let join_plan_schema = join_plan.schema().clone();
1282
1283                let bin_expr_builder = |_: &String| {
1284                    let (left_col_name, right_col_name) = field_columns.next().unwrap();
1285                    let left_col = join_plan_schema
1286                        .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1287                        .context(DataFusionPlanningSnafu)?
1288                        .into();
1289                    let right_col = join_plan_schema
1290                        .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1291                        .context(DataFusionPlanningSnafu)?
1292                        .into();
1293
1294                    let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1295                    let mut binary_expr =
1296                        binary_expr_builder(DfExpr::Column(left_col), DfExpr::Column(right_col))?;
1297                    if is_comparison_op && should_return_bool {
1298                        binary_expr = DfExpr::Cast(Cast {
1299                            expr: Box::new(binary_expr),
1300                            data_type: ArrowDataType::Float64,
1301                        });
1302                    }
1303                    Ok(binary_expr)
1304                };
1305                if is_comparison_op && !should_return_bool {
1306                    // PromQL comparison operators without `bool` are filters:
1307                    //   - keep the instant-vector side sample values
1308                    //   - drop samples where the comparison is false
1309                    //
1310                    // So we filter on the join result and then project only the side that should
1311                    // be preserved according to PromQL semantics.
1312                    let filtered = self.filter_on_field_column(join_plan, bin_expr_builder)?;
1313                    let (project_table_ref, mut project_context, project_field_columns) =
1314                        match (lhs.value_type(), rhs.value_type()) {
1315                            (ValueType::Scalar, ValueType::Vector) => (
1316                                &right_table_ref,
1317                                right_context.clone(),
1318                                right_aligned_field_columns,
1319                            ),
1320                            _ => (
1321                                &left_table_ref,
1322                                left_context.clone(),
1323                                left_aligned_field_columns,
1324                            ),
1325                        };
1326                    project_context.field_columns = project_field_columns;
1327                    self.project_binary_join_side(filtered, project_table_ref, &project_context)
1328                } else {
1329                    self.projection_for_each_field_column(join_plan, bin_expr_builder)
1330                }
1331            }
1332        }
1333    }
1334
1335    fn project_binary_join_side(
1336        &mut self,
1337        input: LogicalPlan,
1338        table_ref: &TableReference,
1339        context: &PromPlannerContext,
1340    ) -> Result<LogicalPlan> {
1341        let schema = input.schema();
1342
1343        let mut project_exprs =
1344            Vec::with_capacity(context.tag_columns.len() + context.field_columns.len() + 2);
1345
1346        // Project time index from the chosen side.
1347        if let Some(time_index_column) = &context.time_index_column {
1348            let time_index_col = schema
1349                .qualified_field_with_name(Some(table_ref), time_index_column)
1350                .context(DataFusionPlanningSnafu)?
1351                .into();
1352            project_exprs.push(DfExpr::Column(time_index_col));
1353        }
1354
1355        // Project field columns from the chosen side.
1356        for field_column in &context.field_columns {
1357            let field_col = schema
1358                .qualified_field_with_name(Some(table_ref), field_column)
1359                .context(DataFusionPlanningSnafu)?
1360                .into();
1361            project_exprs.push(DfExpr::Column(field_col));
1362        }
1363
1364        // Project tag columns from the chosen side.
1365        for tag_column in &context.tag_columns {
1366            let tag_col = schema
1367                .qualified_field_with_name(Some(table_ref), tag_column)
1368                .context(DataFusionPlanningSnafu)?
1369                .into();
1370            project_exprs.push(DfExpr::Column(tag_col));
1371        }
1372
1373        // Preserve `__tsid` if present, so it can still be used internally downstream. It's
1374        // stripped from the final output anyway.
1375        if let Some(tsid_col) =
1376            Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid)
1377        {
1378            project_exprs.push(tsid_col);
1379        }
1380
1381        let plan = LogicalPlanBuilder::from(input)
1382            .project(project_exprs)
1383            .context(DataFusionPlanningSnafu)?
1384            .build()
1385            .context(DataFusionPlanningSnafu)?;
1386
1387        // Update context to reflect the projected schema. Don't keep a table qualifier since
1388        // the result is a derived expression.
1389        self.ctx = context.clone();
1390        self.ctx.table_name = None;
1391        self.ctx.schema_name = None;
1392
1393        Ok(plan)
1394    }
1395
1396    fn prom_number_lit_to_plan(&mut self, number_literal: &NumberLiteral) -> Result<LogicalPlan> {
1397        let NumberLiteral { val } = number_literal;
1398        self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1399        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1400        self.ctx.reset_table_name_and_schema();
1401        let literal_expr = df_prelude::lit(*val);
1402
1403        let plan = LogicalPlan::Extension(Extension {
1404            node: Arc::new(
1405                EmptyMetric::new(
1406                    self.ctx.start,
1407                    self.ctx.end,
1408                    self.ctx.interval,
1409                    SPECIAL_TIME_FUNCTION.to_string(),
1410                    DEFAULT_FIELD_COLUMN.to_string(),
1411                    Some(literal_expr),
1412                )
1413                .context(DataFusionPlanningSnafu)?,
1414            ),
1415        });
1416        Ok(plan)
1417    }
1418
1419    fn prom_string_lit_to_plan(&mut self, string_literal: &StringLiteral) -> Result<LogicalPlan> {
1420        let StringLiteral { val } = string_literal;
1421        self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1422        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1423        self.ctx.reset_table_name_and_schema();
1424        let literal_expr = df_prelude::lit(val.clone());
1425
1426        let plan = LogicalPlan::Extension(Extension {
1427            node: Arc::new(
1428                EmptyMetric::new(
1429                    self.ctx.start,
1430                    self.ctx.end,
1431                    self.ctx.interval,
1432                    SPECIAL_TIME_FUNCTION.to_string(),
1433                    DEFAULT_FIELD_COLUMN.to_string(),
1434                    Some(literal_expr),
1435                )
1436                .context(DataFusionPlanningSnafu)?,
1437            ),
1438        });
1439        Ok(plan)
1440    }
1441
1442    async fn prom_vector_selector_to_plan(
1443        &mut self,
1444        vector_selector: &VectorSelector,
1445        timestamp_fn: bool,
1446    ) -> Result<LogicalPlan> {
1447        let VectorSelector {
1448            name,
1449            offset,
1450            matchers,
1451            at: _,
1452        } = vector_selector;
1453        let matchers = self.preprocess_label_matchers(matchers, name)?;
1454        if let Some(empty_plan) = self.setup_context().await? {
1455            return Ok(empty_plan);
1456        }
1457        let normalize = self
1458            .selector_to_series_normalize_plan(offset, matchers, false)
1459            .await?;
1460
1461        let normalize = if timestamp_fn {
1462            // If evaluating the PromQL `timestamp()` function, project the time index column as the value column
1463            // before wrapping with [`InstantManipulate`], so the output matches PromQL's `timestamp()` semantics.
1464            self.create_timestamp_func_plan(normalize)?
1465        } else {
1466            normalize
1467        };
1468
1469        let manipulate = InstantManipulate::new(
1470            self.ctx.start,
1471            self.ctx.end,
1472            self.ctx.lookback_delta,
1473            self.ctx.interval,
1474            self.ctx
1475                .time_index_column
1476                .clone()
1477                .expect("time index should be set in `setup_context`"),
1478            if self.ctx.use_tsid {
1479                vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
1480            } else {
1481                self.ctx.tag_columns.clone()
1482            },
1483            self.ctx.field_columns.first().cloned(),
1484            normalize,
1485        );
1486        Ok(LogicalPlan::Extension(Extension {
1487            node: Arc::new(manipulate),
1488        }))
1489    }
1490
1491    /// Builds a projection plan for the PromQL `timestamp()` function.
1492    /// Projects the time index column as the value column for each row.
1493    ///
1494    /// # Arguments
1495    /// * `normalize` - Input [`LogicalPlan`] for the normalized series.
1496    ///
1497    /// # Returns
1498    /// Returns a [`Result<LogicalPlan>`] where the resulting logical plan projects the timestamp
1499    /// column as the value column, along with the original tag and time index columns.
1500    ///
1501    /// # Timestamp vs. Time Function
1502    ///
1503    /// - **Timestamp Function (`timestamp()`)**: In PromQL, the `timestamp()` function returns the
1504    ///   timestamp (time index) of each sample as the value column.
1505    ///
1506    /// - **Time Function (`time()`)**: The `time()` function returns the evaluation time of the query
1507    ///   as a scalar value.
1508    ///
1509    /// # Side Effects
1510    /// Updates the planner context's field columns to the timestamp column name.
1511    ///
1512    fn create_timestamp_func_plan(&mut self, normalize: LogicalPlan) -> Result<LogicalPlan> {
1513        let time_expr = build_special_time_expr(self.ctx.time_index_column.as_ref().unwrap())
1514            .alias(DEFAULT_FIELD_COLUMN);
1515        self.ctx.field_columns = vec![time_expr.schema_name().to_string()];
1516        let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
1517        project_exprs.push(self.create_time_index_column_expr()?);
1518        project_exprs.push(time_expr);
1519        project_exprs.extend(self.create_tag_column_exprs()?);
1520
1521        LogicalPlanBuilder::from(normalize)
1522            .project(project_exprs)
1523            .context(DataFusionPlanningSnafu)?
1524            .build()
1525            .context(DataFusionPlanningSnafu)
1526    }
1527
1528    async fn prom_matrix_selector_to_plan(
1529        &mut self,
1530        matrix_selector: &MatrixSelector,
1531    ) -> Result<LogicalPlan> {
1532        let MatrixSelector { vs, range } = matrix_selector;
1533        let VectorSelector {
1534            name,
1535            offset,
1536            matchers,
1537            ..
1538        } = vs;
1539        let matchers = self.preprocess_label_matchers(matchers, name)?;
1540        ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
1541        let range_ms = range.as_millis() as _;
1542        self.ctx.range = Some(range_ms);
1543
1544        // Some functions like rate may require special fields in the RangeManipulate plan
1545        // so we can't skip RangeManipulate.
1546        let normalize = match self.setup_context().await? {
1547            Some(empty_plan) => empty_plan,
1548            None => {
1549                self.selector_to_series_normalize_plan(offset, matchers, true)
1550                    .await?
1551            }
1552        };
1553        let manipulate = RangeManipulate::new(
1554            self.ctx.start,
1555            self.ctx.end,
1556            self.ctx.interval,
1557            // TODO(ruihang): convert via Timestamp datatypes to support different time units
1558            range_ms,
1559            self.ctx
1560                .time_index_column
1561                .clone()
1562                .expect("time index should be set in `setup_context`"),
1563            self.ctx.field_columns.clone(),
1564            normalize,
1565        )
1566        .context(DataFusionPlanningSnafu)?;
1567
1568        Ok(LogicalPlan::Extension(Extension {
1569            node: Arc::new(manipulate),
1570        }))
1571    }
1572
1573    async fn prom_call_expr_to_plan(
1574        &mut self,
1575        query_engine_state: &QueryEngineState,
1576        call_expr: &Call,
1577    ) -> Result<LogicalPlan> {
1578        let Call { func, args } = call_expr;
1579        // some special functions that are not expression but a plan
1580        match func.name {
1581            SPECIAL_HISTOGRAM_QUANTILE => {
1582                return self.create_histogram_plan(args, query_engine_state).await;
1583            }
1584            SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await,
1585            SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await,
1586            SPECIAL_ABSENT_FUNCTION => {
1587                return self.create_absent_plan(args, query_engine_state).await;
1588            }
1589            _ => {}
1590        }
1591
1592        // transform function arguments
1593        let args = self.create_function_args(&args.args)?;
1594        let input = if let Some(prom_expr) = &args.input {
1595            self.prom_expr_to_plan_inner(prom_expr, func.name == "timestamp", query_engine_state)
1596                .await?
1597        } else {
1598            self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
1599            self.ctx.reset_table_name_and_schema();
1600            self.ctx.tag_columns = vec![];
1601            self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1602            LogicalPlan::Extension(Extension {
1603                node: Arc::new(
1604                    EmptyMetric::new(
1605                        self.ctx.start,
1606                        self.ctx.end,
1607                        self.ctx.interval,
1608                        SPECIAL_TIME_FUNCTION.to_string(),
1609                        DEFAULT_FIELD_COLUMN.to_string(),
1610                        None,
1611                    )
1612                    .context(DataFusionPlanningSnafu)?,
1613                ),
1614            })
1615        };
1616        let (mut func_exprs, new_tags) =
1617            self.create_function_expr(func, args.literals.clone(), query_engine_state)?;
1618        func_exprs.insert(0, self.create_time_index_column_expr()?);
1619        func_exprs.extend_from_slice(&self.create_tag_column_exprs()?);
1620        if let Some(tsid_col) =
1621            Self::optional_tsid_projection(input.schema(), None, self.ctx.use_tsid)
1622        {
1623            func_exprs.push(tsid_col);
1624        }
1625
1626        let builder = LogicalPlanBuilder::from(input)
1627            .project(func_exprs)
1628            .context(DataFusionPlanningSnafu)?
1629            .filter(self.create_empty_values_filter_expr()?)
1630            .context(DataFusionPlanningSnafu)?;
1631
1632        let builder = match func.name {
1633            "sort" => builder
1634                .sort(self.create_field_columns_sort_exprs(true))
1635                .context(DataFusionPlanningSnafu)?,
1636            "sort_desc" => builder
1637                .sort(self.create_field_columns_sort_exprs(false))
1638                .context(DataFusionPlanningSnafu)?,
1639            "sort_by_label" => builder
1640                .sort(Self::create_sort_exprs_by_tags(
1641                    func.name,
1642                    args.literals,
1643                    true,
1644                )?)
1645                .context(DataFusionPlanningSnafu)?,
1646            "sort_by_label_desc" => builder
1647                .sort(Self::create_sort_exprs_by_tags(
1648                    func.name,
1649                    args.literals,
1650                    false,
1651                )?)
1652                .context(DataFusionPlanningSnafu)?,
1653
1654            _ => builder,
1655        };
1656
1657        // Update context tags after building plan
1658        // We can't push them before planning, because they won't exist until projection.
1659        for tag in new_tags {
1660            self.ctx.tag_columns.push(tag);
1661        }
1662
1663        let plan = builder.build().context(DataFusionPlanningSnafu)?;
1664        common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}");
1665
1666        Ok(plan)
1667    }
1668
1669    async fn prom_ext_expr_to_plan(
1670        &mut self,
1671        query_engine_state: &QueryEngineState,
1672        ext_expr: &promql_parser::parser::ast::Extension,
1673    ) -> Result<LogicalPlan> {
1674        // let promql_parser::parser::ast::Extension { expr } = ext_expr;
1675        let expr = &ext_expr.expr;
1676        let children = expr.children();
1677        let plan = self
1678            .prom_expr_to_plan(&children[0], query_engine_state)
1679            .await?;
1680        // Wrapper for the explanation/analyze of the existing plan
1681        // https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/builder/struct.LogicalPlanBuilder.html#method.explain
1682        // if `analyze` is true, runs the actual plan and produces
1683        // information about metrics during run.
1684        // if `verbose` is true, prints out additional details when VERBOSE keyword is specified
1685        match expr.name() {
1686            ANALYZE_NODE_NAME => LogicalPlanBuilder::from(plan)
1687                .explain(false, true)
1688                .unwrap()
1689                .build()
1690                .context(DataFusionPlanningSnafu),
1691            ANALYZE_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
1692                .explain(true, true)
1693                .unwrap()
1694                .build()
1695                .context(DataFusionPlanningSnafu),
1696            EXPLAIN_NODE_NAME => LogicalPlanBuilder::from(plan)
1697                .explain(false, false)
1698                .unwrap()
1699                .build()
1700                .context(DataFusionPlanningSnafu),
1701            EXPLAIN_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
1702                .explain(true, false)
1703                .unwrap()
1704                .build()
1705                .context(DataFusionPlanningSnafu),
1706            ALIAS_NODE_NAME => {
1707                let alias = expr
1708                    .as_any()
1709                    .downcast_ref::<AliasExpr>()
1710                    .context(UnexpectedPlanExprSnafu {
1711                        desc: "Expected AliasExpr",
1712                    })?
1713                    .alias
1714                    .clone();
1715                self.apply_alias(plan, alias)
1716            }
1717            _ => LogicalPlanBuilder::empty(true)
1718                .build()
1719                .context(DataFusionPlanningSnafu),
1720        }
1721    }
1722
1723    /// Extract metric name from `__name__` matcher and set it into [PromPlannerContext].
1724    /// Returns a new [Matchers] that doesn't contain metric name matcher.
1725    ///
1726    /// Each call to this function means new selector is started. Thus, the context will be reset
1727    /// at first.
1728    ///
1729    /// Name rule:
1730    /// - if `name` is some, then the matchers MUST NOT contain `__name__` matcher.
1731    /// - if `name` is none, then the matchers MAY contain NONE OR MULTIPLE `__name__` matchers.
1732    #[allow(clippy::mutable_key_type)]
1733    fn preprocess_label_matchers(
1734        &mut self,
1735        label_matchers: &Matchers,
1736        name: &Option<String>,
1737    ) -> Result<Matchers> {
1738        self.ctx.reset();
1739
1740        let metric_name;
1741        if let Some(name) = name.clone() {
1742            metric_name = Some(name);
1743            ensure!(
1744                label_matchers.find_matchers(METRIC_NAME).is_empty(),
1745                MultipleMetricMatchersSnafu
1746            );
1747        } else {
1748            let mut matches = label_matchers.find_matchers(METRIC_NAME);
1749            ensure!(!matches.is_empty(), NoMetricMatcherSnafu);
1750            ensure!(matches.len() == 1, MultipleMetricMatchersSnafu);
1751            ensure!(
1752                matches[0].op == MatchOp::Equal,
1753                UnsupportedMatcherOpSnafu {
1754                    matcher_op: matches[0].op.to_string(),
1755                    matcher: METRIC_NAME
1756                }
1757            );
1758            metric_name = matches.pop().map(|m| m.value);
1759        }
1760
1761        self.ctx.table_name = metric_name;
1762
1763        let mut matchers = HashSet::new();
1764        for matcher in &label_matchers.matchers {
1765            // TODO(ruihang): support other metric match ops
1766            if matcher.name == FIELD_COLUMN_MATCHER {
1767                self.ctx
1768                    .field_column_matcher
1769                    .get_or_insert_default()
1770                    .push(matcher.clone());
1771            } else if matcher.name == SCHEMA_COLUMN_MATCHER || matcher.name == DB_COLUMN_MATCHER {
1772                ensure!(
1773                    matcher.op == MatchOp::Equal,
1774                    UnsupportedMatcherOpSnafu {
1775                        matcher: matcher.name.clone(),
1776                        matcher_op: matcher.op.to_string(),
1777                    }
1778                );
1779                self.ctx.schema_name = Some(matcher.value.clone());
1780            } else if matcher.name != METRIC_NAME {
1781                self.ctx.selector_matcher.push(matcher.clone());
1782                let _ = matchers.insert(matcher.clone());
1783            }
1784        }
1785
1786        Ok(Matchers::new(matchers.into_iter().collect()))
1787    }
1788
1789    async fn selector_to_series_normalize_plan(
1790        &mut self,
1791        offset: &Option<Offset>,
1792        label_matchers: Matchers,
1793        is_range_selector: bool,
1794    ) -> Result<LogicalPlan> {
1795        // make table scan plan
1796        let table_ref = self.table_ref()?;
1797        let mut table_scan = self.create_table_scan_plan(table_ref.clone()).await?;
1798        let table_schema = table_scan.schema();
1799
1800        // make filter exprs
1801        let offset_duration = match offset {
1802            Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
1803            Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
1804            None => 0,
1805        };
1806        let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?;
1807        if let Some(time_index_filter) = self.build_time_index_filter(offset_duration)? {
1808            scan_filters.push(time_index_filter);
1809        }
1810        table_scan = LogicalPlanBuilder::from(table_scan)
1811            .filter(conjunction(scan_filters).unwrap()) // Safety: `scan_filters` is not empty.
1812            .context(DataFusionPlanningSnafu)?
1813            .build()
1814            .context(DataFusionPlanningSnafu)?;
1815
1816        // make a projection plan if there is any `__field__` matcher
1817        if let Some(field_matchers) = &self.ctx.field_column_matcher {
1818            let col_set = self.ctx.field_columns.iter().collect::<HashSet<_>>();
1819            // opt-in set
1820            let mut result_set = HashSet::new();
1821            // opt-out set
1822            let mut reverse_set = HashSet::new();
1823            for matcher in field_matchers {
1824                match &matcher.op {
1825                    MatchOp::Equal => {
1826                        if col_set.contains(&matcher.value) {
1827                            let _ = result_set.insert(matcher.value.clone());
1828                        } else {
1829                            return Err(ColumnNotFoundSnafu {
1830                                col: matcher.value.clone(),
1831                            }
1832                            .build());
1833                        }
1834                    }
1835                    MatchOp::NotEqual => {
1836                        if col_set.contains(&matcher.value) {
1837                            let _ = reverse_set.insert(matcher.value.clone());
1838                        } else {
1839                            return Err(ColumnNotFoundSnafu {
1840                                col: matcher.value.clone(),
1841                            }
1842                            .build());
1843                        }
1844                    }
1845                    MatchOp::Re(regex) => {
1846                        for col in &self.ctx.field_columns {
1847                            if regex.is_match(col) {
1848                                let _ = result_set.insert(col.clone());
1849                            }
1850                        }
1851                    }
1852                    MatchOp::NotRe(regex) => {
1853                        for col in &self.ctx.field_columns {
1854                            if regex.is_match(col) {
1855                                let _ = reverse_set.insert(col.clone());
1856                            }
1857                        }
1858                    }
1859                }
1860            }
1861            // merge two set
1862            if result_set.is_empty() {
1863                result_set = col_set.into_iter().cloned().collect();
1864            }
1865            for col in reverse_set {
1866                let _ = result_set.remove(&col);
1867            }
1868
1869            // mask the field columns in context using computed result set
1870            self.ctx.field_columns = self
1871                .ctx
1872                .field_columns
1873                .drain(..)
1874                .filter(|col| result_set.contains(col))
1875                .collect();
1876
1877            let exprs = result_set
1878                .into_iter()
1879                .map(|col| DfExpr::Column(Column::new_unqualified(col)))
1880                .chain(self.create_tag_column_exprs()?)
1881                .chain(
1882                    self.ctx
1883                        .use_tsid
1884                        .then_some(DfExpr::Column(Column::new_unqualified(
1885                            DATA_SCHEMA_TSID_COLUMN_NAME,
1886                        ))),
1887                )
1888                .chain(Some(self.create_time_index_column_expr()?))
1889                .collect::<Vec<_>>();
1890
1891            // reuse this variable for simplicity
1892            table_scan = LogicalPlanBuilder::from(table_scan)
1893                .project(exprs)
1894                .context(DataFusionPlanningSnafu)?
1895                .build()
1896                .context(DataFusionPlanningSnafu)?;
1897        }
1898
1899        // make sort plan
1900        let series_key_columns = if self.ctx.use_tsid {
1901            vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
1902        } else {
1903            self.ctx.tag_columns.clone()
1904        };
1905
1906        let sort_exprs = if self.ctx.use_tsid {
1907            vec![
1908                DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)).sort(true, true),
1909                self.create_time_index_column_expr()?.sort(true, true),
1910            ]
1911        } else {
1912            self.create_tag_and_time_index_column_sort_exprs()?
1913        };
1914
1915        let sort_plan = LogicalPlanBuilder::from(table_scan)
1916            .sort(sort_exprs)
1917            .context(DataFusionPlanningSnafu)?
1918            .build()
1919            .context(DataFusionPlanningSnafu)?;
1920
1921        // make divide plan
1922        let time_index_column =
1923            self.ctx
1924                .time_index_column
1925                .clone()
1926                .with_context(|| TimeIndexNotFoundSnafu {
1927                    table: table_ref.to_string(),
1928                })?;
1929        let divide_plan = LogicalPlan::Extension(Extension {
1930            node: Arc::new(SeriesDivide::new(
1931                series_key_columns.clone(),
1932                time_index_column,
1933                sort_plan,
1934            )),
1935        });
1936
1937        // make series_normalize plan
1938        if !is_range_selector && offset_duration == 0 {
1939            return Ok(divide_plan);
1940        }
1941        let series_normalize = SeriesNormalize::new(
1942            offset_duration,
1943            self.ctx
1944                .time_index_column
1945                .clone()
1946                .with_context(|| TimeIndexNotFoundSnafu {
1947                    table: table_ref.to_quoted_string(),
1948                })?,
1949            is_range_selector,
1950            series_key_columns,
1951            divide_plan,
1952        );
1953        let logical_plan = LogicalPlan::Extension(Extension {
1954            node: Arc::new(series_normalize),
1955        });
1956
1957        Ok(logical_plan)
1958    }
1959
1960    /// Convert [LabelModifier] to [Column] exprs for aggregation.
1961    /// Timestamp column and tag columns will be included.
1962    ///
1963    /// # Side effect
1964    ///
1965    /// This method will also change the tag columns in ctx if `update_ctx` is true.
1966    fn agg_modifier_to_col(
1967        &mut self,
1968        input_schema: &DFSchemaRef,
1969        modifier: &Option<LabelModifier>,
1970        update_ctx: bool,
1971    ) -> Result<Vec<DfExpr>> {
1972        match modifier {
1973            None => {
1974                if update_ctx {
1975                    self.ctx.tag_columns.clear();
1976                }
1977                Ok(vec![self.create_time_index_column_expr()?])
1978            }
1979            Some(LabelModifier::Include(labels)) => {
1980                if update_ctx {
1981                    self.ctx.tag_columns.clear();
1982                }
1983                let mut exprs = Vec::with_capacity(labels.labels.len());
1984                for label in &labels.labels {
1985                    if is_metric_engine_internal_column(label) {
1986                        continue;
1987                    }
1988                    // nonexistence label will be ignored
1989                    if let Some(column_name) = Self::find_case_sensitive_column(input_schema, label)
1990                    {
1991                        exprs.push(DfExpr::Column(Column::from_name(column_name.clone())));
1992
1993                        if update_ctx {
1994                            // update the tag columns in context
1995                            self.ctx.tag_columns.push(column_name);
1996                        }
1997                    }
1998                }
1999                // add timestamp column
2000                exprs.push(self.create_time_index_column_expr()?);
2001
2002                Ok(exprs)
2003            }
2004            Some(LabelModifier::Exclude(labels)) => {
2005                let mut all_fields = input_schema
2006                    .fields()
2007                    .iter()
2008                    .map(|f| f.name())
2009                    .collect::<BTreeSet<_>>();
2010
2011                // Exclude metric engine internal columns (not PromQL labels) from the implicit
2012                // "without" label set.
2013                all_fields.retain(|col| !is_metric_engine_internal_column(col.as_str()));
2014
2015                // remove "without"-ed fields
2016                // nonexistence label will be ignored
2017                for label in &labels.labels {
2018                    let _ = all_fields.remove(label);
2019                }
2020
2021                // remove time index and value fields
2022                if let Some(time_index) = &self.ctx.time_index_column {
2023                    let _ = all_fields.remove(time_index);
2024                }
2025                for value in &self.ctx.field_columns {
2026                    let _ = all_fields.remove(value);
2027                }
2028
2029                if update_ctx {
2030                    // change the tag columns in context
2031                    self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect();
2032                }
2033
2034                // collect remaining fields and convert to col expr
2035                let mut exprs = all_fields
2036                    .into_iter()
2037                    .map(|c| DfExpr::Column(Column::from(c)))
2038                    .collect::<Vec<_>>();
2039
2040                // add timestamp column
2041                exprs.push(self.create_time_index_column_expr()?);
2042
2043                Ok(exprs)
2044            }
2045        }
2046    }
2047
2048    // TODO(ruihang): ignore `MetricNameLabel` (`__name__`) matcher
2049    pub fn matchers_to_expr(
2050        label_matchers: Matchers,
2051        table_schema: &DFSchemaRef,
2052    ) -> Result<Vec<DfExpr>> {
2053        let mut exprs = Vec::with_capacity(label_matchers.matchers.len());
2054        for matcher in label_matchers.matchers {
2055            if matcher.name == SCHEMA_COLUMN_MATCHER
2056                || matcher.name == DB_COLUMN_MATCHER
2057                || matcher.name == FIELD_COLUMN_MATCHER
2058            {
2059                continue;
2060            }
2061
2062            let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str());
2063            let col = if let Some(column_name) = column_name {
2064                DfExpr::Column(Column::from_name(column_name))
2065            } else {
2066                DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
2067                    .alias(matcher.name.clone())
2068            };
2069            let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None);
2070            let expr = match matcher.op {
2071                MatchOp::Equal => col.eq(lit),
2072                MatchOp::NotEqual => col.not_eq(lit),
2073                MatchOp::Re(re) => {
2074                    // TODO(ruihang): a more programmatic way to handle this in datafusion
2075
2076                    // This is a hack to handle `.+` and `.*`, and is not strictly correct
2077                    // `.` doesn't match newline (`\n`). Given this is in PromQL context,
2078                    // most of the time it's fine.
2079                    if re.as_str() == "^(?:.*)$" {
2080                        continue;
2081                    }
2082                    if re.as_str() == "^(?:.+)$" {
2083                        col.not_eq(DfExpr::Literal(
2084                            ScalarValue::Utf8(Some(String::new())),
2085                            None,
2086                        ))
2087                    } else {
2088                        DfExpr::BinaryExpr(BinaryExpr {
2089                            left: Box::new(col),
2090                            op: Operator::RegexMatch,
2091                            right: Box::new(DfExpr::Literal(
2092                                ScalarValue::Utf8(Some(re.as_str().to_string())),
2093                                None,
2094                            )),
2095                        })
2096                    }
2097                }
2098                MatchOp::NotRe(re) => {
2099                    if re.as_str() == "^(?:.*)$" {
2100                        DfExpr::Literal(ScalarValue::Boolean(Some(false)), None)
2101                    } else if re.as_str() == "^(?:.+)$" {
2102                        col.eq(DfExpr::Literal(
2103                            ScalarValue::Utf8(Some(String::new())),
2104                            None,
2105                        ))
2106                    } else {
2107                        DfExpr::BinaryExpr(BinaryExpr {
2108                            left: Box::new(col),
2109                            op: Operator::RegexNotMatch,
2110                            right: Box::new(DfExpr::Literal(
2111                                ScalarValue::Utf8(Some(re.as_str().to_string())),
2112                                None,
2113                            )),
2114                        })
2115                    }
2116                }
2117            };
2118            exprs.push(expr);
2119        }
2120
2121        Ok(exprs)
2122    }
2123
2124    fn find_case_sensitive_column(schema: &DFSchemaRef, column: &str) -> Option<String> {
2125        if is_metric_engine_internal_column(column) {
2126            return None;
2127        }
2128        schema
2129            .fields()
2130            .iter()
2131            .find(|field| field.name() == column)
2132            .map(|field| field.name().clone())
2133    }
2134
2135    fn table_from_source(&self, source: &Arc<dyn TableSource>) -> Result<table::TableRef> {
2136        Ok(source
2137            .as_any()
2138            .downcast_ref::<DefaultTableSource>()
2139            .context(UnknownTableSnafu)?
2140            .table_provider
2141            .as_any()
2142            .downcast_ref::<DfTableProviderAdapter>()
2143            .context(UnknownTableSnafu)?
2144            .table())
2145    }
2146
2147    fn table_ref(&self) -> Result<TableReference> {
2148        let table_name = self
2149            .ctx
2150            .table_name
2151            .clone()
2152            .context(TableNameNotFoundSnafu)?;
2153
2154        // set schema name if `__schema__` is given
2155        let table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2156            TableReference::partial(schema_name.as_str(), table_name.as_str())
2157        } else {
2158            TableReference::bare(table_name.as_str())
2159        };
2160
2161        Ok(table_ref)
2162    }
2163
2164    fn build_time_index_filter(&self, offset_duration: i64) -> Result<Option<DfExpr>> {
2165        let start = self.ctx.start;
2166        let end = self.ctx.end;
2167        if end < start {
2168            return InvalidTimeRangeSnafu { start, end }.fail();
2169        }
2170        let lookback_delta = self.ctx.lookback_delta;
2171        let range = self.ctx.range.unwrap_or_default();
2172        let interval = self.ctx.interval;
2173        let time_index_expr = self.create_time_index_column_expr()?;
2174        let num_points = (end - start) / interval;
2175
2176        // Prometheus semantics:
2177        // - Instant selector lookback: (eval_ts - lookback_delta, eval_ts]
2178        // - Range selector:           (eval_ts - range, eval_ts]
2179        //
2180        // So samples positioned exactly at the lower boundary must be excluded. We align the scan
2181        // lower bound with Prometheus by shifting it forward by 1ms (millisecond granularity),
2182        // while still using a `>=` filter.
2183        let selector_window = if range == 0 { lookback_delta } else { range };
2184        let lower_exclusive_adjustment = if selector_window > 0 { 1 } else { 0 };
2185
2186        // Scan a continuous time range
2187        if (end - start) / interval > MAX_SCATTER_POINTS || interval <= INTERVAL_1H {
2188            let single_time_range = time_index_expr
2189                .clone()
2190                .gt_eq(DfExpr::Literal(
2191                    ScalarValue::TimestampMillisecond(
2192                        Some(
2193                            self.ctx.start - offset_duration - selector_window
2194                                + lower_exclusive_adjustment,
2195                        ),
2196                        None,
2197                    ),
2198                    None,
2199                ))
2200                .and(time_index_expr.lt_eq(DfExpr::Literal(
2201                    ScalarValue::TimestampMillisecond(Some(self.ctx.end - offset_duration), None),
2202                    None,
2203                )));
2204            return Ok(Some(single_time_range));
2205        }
2206
2207        // Otherwise scan scatter ranges separately
2208        let mut filters = Vec::with_capacity(num_points as usize + 1);
2209        for timestamp in (start..=end).step_by(interval as usize) {
2210            filters.push(
2211                time_index_expr
2212                    .clone()
2213                    .gt_eq(DfExpr::Literal(
2214                        ScalarValue::TimestampMillisecond(
2215                            Some(
2216                                timestamp - offset_duration - selector_window
2217                                    + lower_exclusive_adjustment,
2218                            ),
2219                            None,
2220                        ),
2221                        None,
2222                    ))
2223                    .and(time_index_expr.clone().lt_eq(DfExpr::Literal(
2224                        ScalarValue::TimestampMillisecond(Some(timestamp - offset_duration), None),
2225                        None,
2226                    ))),
2227            )
2228        }
2229
2230        Ok(filters.into_iter().reduce(DfExpr::or))
2231    }
2232
2233    /// Create a table scan plan and a filter plan with given filter.
2234    ///
2235    /// # Panic
2236    /// If the filter is empty
2237    async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
2238        let provider = self
2239            .table_provider
2240            .resolve_table(table_ref.clone())
2241            .await
2242            .context(CatalogSnafu)?;
2243
2244        let logical_table = self.table_from_source(&provider)?;
2245
2246        // Try to rewrite the table scan to physical table scan if possible.
2247        let mut maybe_phy_table_ref = table_ref.clone();
2248        let mut scan_provider = provider;
2249        let mut table_id_filter: Option<u32> = None;
2250
2251        // If it's a metric engine logical table, scan its physical table directly and filter by
2252        // `__table_id = logical_table_id` to get access to internal columns like `__tsid`.
2253        if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
2254            && let Some(physical_table_name) = logical_table
2255                .table_info()
2256                .meta
2257                .options
2258                .extra_options
2259                .get(LOGICAL_TABLE_METADATA_KEY)
2260        {
2261            let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2262                TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
2263            } else {
2264                TableReference::bare(physical_table_name.as_str())
2265            };
2266
2267            let physical_provider = match self
2268                .table_provider
2269                .resolve_table(physical_table_ref.clone())
2270                .await
2271            {
2272                Ok(provider) => provider,
2273                Err(e) if e.status_code() == StatusCode::TableNotFound => {
2274                    // Fall back to scanning the logical table. It still works, but without
2275                    // `__tsid` optimization.
2276                    scan_provider.clone()
2277                }
2278                Err(e) => return Err(e).context(CatalogSnafu),
2279            };
2280
2281            if !Arc::ptr_eq(&physical_provider, &scan_provider) {
2282                // Only rewrite when internal columns exist in physical schema.
2283                let physical_table = self.table_from_source(&physical_provider)?;
2284
2285                let has_table_id = physical_table
2286                    .schema()
2287                    .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
2288                    .is_some();
2289                let has_tsid = physical_table
2290                    .schema()
2291                    .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2292                    .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2293
2294                if has_table_id && has_tsid {
2295                    scan_provider = physical_provider;
2296                    maybe_phy_table_ref = physical_table_ref;
2297                    table_id_filter = Some(logical_table.table_info().ident.table_id);
2298                }
2299            }
2300        }
2301
2302        let scan_table = self.table_from_source(&scan_provider)?;
2303
2304        let use_tsid = table_id_filter.is_some()
2305            && scan_table
2306                .schema()
2307                .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2308                .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2309        self.ctx.use_tsid = use_tsid;
2310
2311        let all_table_tags = self.ctx.tag_columns.clone();
2312
2313        let scan_tag_columns = if use_tsid {
2314            let mut scan_tags = self.ctx.tag_columns.clone();
2315            for matcher in &self.ctx.selector_matcher {
2316                if is_metric_engine_internal_column(&matcher.name) {
2317                    continue;
2318                }
2319                if all_table_tags.iter().any(|tag| tag == &matcher.name) {
2320                    scan_tags.push(matcher.name.clone());
2321                }
2322            }
2323            scan_tags.sort_unstable();
2324            scan_tags.dedup();
2325            scan_tags
2326        } else {
2327            self.ctx.tag_columns.clone()
2328        };
2329
2330        let is_time_index_ms = scan_table
2331            .schema()
2332            .timestamp_column()
2333            .with_context(|| TimeIndexNotFoundSnafu {
2334                table: maybe_phy_table_ref.to_quoted_string(),
2335            })?
2336            .data_type
2337            == ConcreteDataType::timestamp_millisecond_datatype();
2338
2339        let scan_projection = if table_id_filter.is_some() {
2340            let mut required_columns = HashSet::new();
2341            required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
2342            required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
2343                TimeIndexNotFoundSnafu {
2344                    table: maybe_phy_table_ref.to_quoted_string(),
2345                }
2346            })?);
2347            for col in &scan_tag_columns {
2348                required_columns.insert(col.clone());
2349            }
2350            for col in &self.ctx.field_columns {
2351                required_columns.insert(col.clone());
2352            }
2353            if use_tsid {
2354                required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
2355            }
2356
2357            let arrow_schema = scan_table.schema().arrow_schema().clone();
2358            Some(
2359                arrow_schema
2360                    .fields()
2361                    .iter()
2362                    .enumerate()
2363                    .filter(|(_, field)| required_columns.contains(field.name().as_str()))
2364                    .map(|(idx, _)| idx)
2365                    .collect::<Vec<_>>(),
2366            )
2367        } else {
2368            None
2369        };
2370
2371        let mut scan_plan =
2372            LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
2373                .context(DataFusionPlanningSnafu)?
2374                .build()
2375                .context(DataFusionPlanningSnafu)?;
2376
2377        if let Some(table_id) = table_id_filter {
2378            scan_plan = LogicalPlanBuilder::from(scan_plan)
2379                .filter(
2380                    DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
2381                        .eq(lit(table_id)),
2382                )
2383                .context(DataFusionPlanningSnafu)?
2384                .alias(table_ref.clone()) // rename the relation back to logical table's name after filtering
2385                .context(DataFusionPlanningSnafu)?
2386                .build()
2387                .context(DataFusionPlanningSnafu)?;
2388        }
2389
2390        if !is_time_index_ms {
2391            // cast to ms if time_index not in Millisecond precision
2392            let expr: Vec<_> = self
2393                .create_field_column_exprs()?
2394                .into_iter()
2395                .chain(
2396                    scan_tag_columns
2397                        .iter()
2398                        .map(|tag| DfExpr::Column(Column::from_name(tag))),
2399                )
2400                .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
2401                    Some(table_ref.clone()),
2402                    DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
2403                ))))
2404                .chain(Some(DfExpr::Alias(Alias {
2405                    expr: Box::new(DfExpr::Cast(Cast {
2406                        expr: Box::new(self.create_time_index_column_expr()?),
2407                        data_type: ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
2408                    })),
2409                    relation: Some(table_ref.clone()),
2410                    name: self
2411                        .ctx
2412                        .time_index_column
2413                        .as_ref()
2414                        .with_context(|| TimeIndexNotFoundSnafu {
2415                            table: table_ref.to_quoted_string(),
2416                        })?
2417                        .clone(),
2418                    metadata: None,
2419                })))
2420                .collect::<Vec<_>>();
2421            scan_plan = LogicalPlanBuilder::from(scan_plan)
2422                .project(expr)
2423                .context(DataFusionPlanningSnafu)?
2424                .build()
2425                .context(DataFusionPlanningSnafu)?;
2426        } else if table_id_filter.is_some() {
2427            // Drop the internal `__table_id` column after filtering.
2428            let project_exprs = self
2429                .create_field_column_exprs()?
2430                .into_iter()
2431                .chain(
2432                    scan_tag_columns
2433                        .iter()
2434                        .map(|tag| DfExpr::Column(Column::from_name(tag))),
2435                )
2436                .chain(
2437                    self.ctx
2438                        .use_tsid
2439                        .then_some(DfExpr::Column(Column::from_name(
2440                            DATA_SCHEMA_TSID_COLUMN_NAME,
2441                        ))),
2442                )
2443                .chain(Some(self.create_time_index_column_expr()?))
2444                .collect::<Vec<_>>();
2445
2446            scan_plan = LogicalPlanBuilder::from(scan_plan)
2447                .project(project_exprs)
2448                .context(DataFusionPlanningSnafu)?
2449                .build()
2450                .context(DataFusionPlanningSnafu)?;
2451        }
2452
2453        let result = LogicalPlanBuilder::from(scan_plan)
2454            .build()
2455            .context(DataFusionPlanningSnafu)?;
2456        Ok(result)
2457    }
2458
2459    fn collect_row_key_tag_columns_from_plan(
2460        &self,
2461        plan: &LogicalPlan,
2462    ) -> Result<BTreeSet<String>> {
2463        fn walk(
2464            planner: &PromPlanner,
2465            plan: &LogicalPlan,
2466            out: &mut BTreeSet<String>,
2467        ) -> Result<()> {
2468            // Derived PromQL plans may contain non-Greptime scans without row-key metadata.
2469            if let LogicalPlan::TableScan(scan) = plan
2470                && let Ok(table) = planner.table_from_source(&scan.source)
2471            {
2472                for col in table.table_info().meta.row_key_column_names() {
2473                    if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
2474                        && col != DATA_SCHEMA_TSID_COLUMN_NAME
2475                        && !is_metric_engine_internal_column(col)
2476                    {
2477                        out.insert(col.clone());
2478                    }
2479                }
2480            }
2481
2482            for input in plan.inputs() {
2483                walk(planner, input, out)?;
2484            }
2485            Ok(())
2486        }
2487
2488        let mut out = BTreeSet::new();
2489        walk(self, plan, &mut out)?;
2490        Ok(out)
2491    }
2492
2493    fn ensure_tag_columns_available(
2494        &self,
2495        plan: LogicalPlan,
2496        required_tags: &BTreeSet<String>,
2497    ) -> Result<LogicalPlan> {
2498        if required_tags.is_empty() {
2499            return Ok(plan);
2500        }
2501
2502        struct Rewriter {
2503            required_tags: BTreeSet<String>,
2504        }
2505
2506        impl TreeNodeRewriter for Rewriter {
2507            type Node = LogicalPlan;
2508
2509            fn f_up(
2510                &mut self,
2511                node: Self::Node,
2512            ) -> datafusion_common::Result<Transformed<Self::Node>> {
2513                match node {
2514                    LogicalPlan::TableScan(scan) => {
2515                        let schema = scan.source.schema();
2516                        let mut projection = match scan.projection.clone() {
2517                            Some(p) => p,
2518                            None => {
2519                                // Scanning all columns already covers required tags.
2520                                return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
2521                            }
2522                        };
2523
2524                        let mut changed = false;
2525                        for tag in &self.required_tags {
2526                            if let Some((idx, _)) = schema
2527                                .fields()
2528                                .iter()
2529                                .enumerate()
2530                                .find(|(_, field)| field.name() == tag)
2531                                && !projection.contains(&idx)
2532                            {
2533                                projection.push(idx);
2534                                changed = true;
2535                            }
2536                        }
2537
2538                        if !changed {
2539                            return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
2540                        }
2541
2542                        projection.sort_unstable();
2543                        projection.dedup();
2544
2545                        let new_scan = TableScan::try_new(
2546                            scan.table_name.clone(),
2547                            scan.source.clone(),
2548                            Some(projection),
2549                            scan.filters,
2550                            scan.fetch,
2551                        )?;
2552                        Ok(Transformed::yes(LogicalPlan::TableScan(new_scan)))
2553                    }
2554                    LogicalPlan::Projection(proj) => {
2555                        let input_schema = proj.input.schema();
2556
2557                        let existing = proj
2558                            .schema
2559                            .fields()
2560                            .iter()
2561                            .map(|f| f.name().as_str())
2562                            .collect::<HashSet<_>>();
2563
2564                        let mut expr = proj.expr.clone();
2565                        let mut has_changed = false;
2566                        for tag in &self.required_tags {
2567                            if existing.contains(tag.as_str()) {
2568                                continue;
2569                            }
2570
2571                            if let Some(idx) = input_schema.index_of_column_by_name(None, tag) {
2572                                expr.push(DfExpr::Column(Column::from(
2573                                    input_schema.qualified_field(idx),
2574                                )));
2575                                has_changed = true;
2576                            }
2577                        }
2578
2579                        if !has_changed {
2580                            return Ok(Transformed::no(LogicalPlan::Projection(proj)));
2581                        }
2582
2583                        let new_proj = Projection::try_new(expr, proj.input)?;
2584                        Ok(Transformed::yes(LogicalPlan::Projection(new_proj)))
2585                    }
2586                    other => Ok(Transformed::no(other)),
2587                }
2588            }
2589        }
2590
2591        let mut rewriter = Rewriter {
2592            required_tags: required_tags.clone(),
2593        };
2594        let rewritten = plan
2595            .rewrite(&mut rewriter)
2596            .context(DataFusionPlanningSnafu)?;
2597        Ok(rewritten.data)
2598    }
2599
2600    fn refresh_tag_columns_from_schema(&mut self, schema: &DFSchemaRef) {
2601        let time_index = self.ctx.time_index_column.as_deref();
2602        let field_columns = self.ctx.field_columns.iter().collect::<HashSet<_>>();
2603
2604        let mut tags = schema
2605            .fields()
2606            .iter()
2607            .map(|f| f.name())
2608            .filter(|name| Some(name.as_str()) != time_index)
2609            .filter(|name| !field_columns.contains(name))
2610            .filter(|name| !is_metric_engine_internal_column(name))
2611            .cloned()
2612            .collect::<Vec<_>>();
2613        tags.sort_unstable();
2614        tags.dedup();
2615        self.ctx.tag_columns = tags;
2616    }
2617
2618    /// Setup [PromPlannerContext]'s state fields.
2619    ///
2620    /// Returns a logical plan for an empty metric.
2621    async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
2622        let table_ref = self.table_ref()?;
2623        let source = match self.table_provider.resolve_table(table_ref.clone()).await {
2624            Err(e) if e.status_code() == StatusCode::TableNotFound => {
2625                let plan = self.setup_context_for_empty_metric()?;
2626                return Ok(Some(plan));
2627            }
2628            res => res.context(CatalogSnafu)?,
2629        };
2630        let table = self.table_from_source(&source)?;
2631
2632        // set time index column name
2633        let time_index = table
2634            .schema()
2635            .timestamp_column()
2636            .with_context(|| TimeIndexNotFoundSnafu {
2637                table: table_ref.to_quoted_string(),
2638            })?
2639            .name
2640            .clone();
2641        self.ctx.time_index_column = Some(time_index);
2642
2643        // set values columns
2644        let values = table
2645            .table_info()
2646            .meta
2647            .field_column_names()
2648            .cloned()
2649            .collect();
2650        self.ctx.field_columns = values;
2651
2652        // set primary key (tag) columns
2653        let tags = table
2654            .table_info()
2655            .meta
2656            .row_key_column_names()
2657            .filter(|col| {
2658                // remove metric engine's internal columns
2659                col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
2660            })
2661            .cloned()
2662            .collect();
2663        self.ctx.tag_columns = tags;
2664
2665        self.ctx.use_tsid = false;
2666
2667        Ok(None)
2668    }
2669
2670    /// Setup [PromPlannerContext]'s state fields for a non existent table
2671    /// without any rows.
2672    fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
2673        self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
2674        self.ctx.reset_table_name_and_schema();
2675        self.ctx.tag_columns = vec![];
2676        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2677        self.ctx.use_tsid = false;
2678
2679        // The table doesn't have any data, so we set start to 0 and end to -1.
2680        let plan = LogicalPlan::Extension(Extension {
2681            node: Arc::new(
2682                EmptyMetric::new(
2683                    0,
2684                    -1,
2685                    self.ctx.interval,
2686                    SPECIAL_TIME_FUNCTION.to_string(),
2687                    DEFAULT_FIELD_COLUMN.to_string(),
2688                    Some(lit(0.0f64)),
2689                )
2690                .context(DataFusionPlanningSnafu)?,
2691            ),
2692        });
2693        Ok(plan)
2694    }
2695
2696    // TODO(ruihang): insert column expr
2697    fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
2698        let mut result = FunctionArgs::default();
2699
2700        for arg in args {
2701            // First try to parse as literal expression (including binary expressions like 100.0 + 3.0)
2702            if let Some(expr) = Self::try_build_literal_expr(arg) {
2703                result.literals.push(expr);
2704            } else {
2705                // If not a literal, treat as vector input
2706                match arg.as_ref() {
2707                    PromExpr::Subquery(_)
2708                    | PromExpr::VectorSelector(_)
2709                    | PromExpr::MatrixSelector(_)
2710                    | PromExpr::Extension(_)
2711                    | PromExpr::Aggregate(_)
2712                    | PromExpr::Paren(_)
2713                    | PromExpr::Call(_)
2714                    | PromExpr::Binary(_)
2715                    | PromExpr::Unary(_) => {
2716                        if result.input.replace(*arg.clone()).is_some() {
2717                            MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
2718                        }
2719                    }
2720
2721                    _ => {
2722                        let expr = Self::get_param_as_literal_expr(&Some(arg.clone()), None, None)?;
2723                        result.literals.push(expr);
2724                    }
2725                }
2726            }
2727        }
2728
2729        Ok(result)
2730    }
2731
2732    /// Creates function expressions for projection and returns the expressions and new tags.
2733    ///
2734    /// # Side Effects
2735    ///
2736    /// This method will update [PromPlannerContext]'s fields and tags if needed.
2737    fn create_function_expr(
2738        &mut self,
2739        func: &Function,
2740        other_input_exprs: Vec<DfExpr>,
2741        query_engine_state: &QueryEngineState,
2742    ) -> Result<(Vec<DfExpr>, Vec<String>)> {
2743        // TODO(ruihang): check function args list
2744        let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
2745
2746        // TODO(ruihang): set this according to in-param list
2747        let field_column_pos = 0;
2748        let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
2749        // New labels after executing the function, e.g. `label_replace` etc.
2750        let mut new_tags = vec![];
2751        let scalar_func = match func.name {
2752            "increase" => ScalarFunc::ExtrapolateUdf(
2753                Arc::new(Increase::scalar_udf()),
2754                self.ctx.range.context(ExpectRangeSelectorSnafu)?,
2755            ),
2756            "rate" => ScalarFunc::ExtrapolateUdf(
2757                Arc::new(Rate::scalar_udf()),
2758                self.ctx.range.context(ExpectRangeSelectorSnafu)?,
2759            ),
2760            "delta" => ScalarFunc::ExtrapolateUdf(
2761                Arc::new(Delta::scalar_udf()),
2762                self.ctx.range.context(ExpectRangeSelectorSnafu)?,
2763            ),
2764            "idelta" => ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf())),
2765            "irate" => ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf())),
2766            "resets" => ScalarFunc::Udf(Arc::new(Resets::scalar_udf())),
2767            "changes" => ScalarFunc::Udf(Arc::new(Changes::scalar_udf())),
2768            "deriv" => ScalarFunc::Udf(Arc::new(Deriv::scalar_udf())),
2769            "avg_over_time" => ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf())),
2770            "min_over_time" => ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf())),
2771            "max_over_time" => ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf())),
2772            "sum_over_time" => ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf())),
2773            "count_over_time" => ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf())),
2774            "last_over_time" => ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf())),
2775            "absent_over_time" => ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf())),
2776            "present_over_time" => ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf())),
2777            "stddev_over_time" => ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf())),
2778            "stdvar_over_time" => ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf())),
2779            "quantile_over_time" => ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf())),
2780            "predict_linear" => {
2781                other_input_exprs[0] = DfExpr::Cast(Cast {
2782                    expr: Box::new(other_input_exprs[0].clone()),
2783                    data_type: ArrowDataType::Int64,
2784                });
2785                ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
2786            }
2787            "double_exponential_smoothing" | "holt_winters" => {
2788                ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
2789            }
2790            "time" => {
2791                exprs.push(build_special_time_expr(
2792                    self.ctx.time_index_column.as_ref().unwrap(),
2793                ));
2794                ScalarFunc::GeneratedExpr
2795            }
2796            "minute" => {
2797                // date_part('minute', time_index)
2798                let expr = self.date_part_on_time_index("minute")?;
2799                exprs.push(expr);
2800                ScalarFunc::GeneratedExpr
2801            }
2802            "hour" => {
2803                // date_part('hour', time_index)
2804                let expr = self.date_part_on_time_index("hour")?;
2805                exprs.push(expr);
2806                ScalarFunc::GeneratedExpr
2807            }
2808            "month" => {
2809                // date_part('month', time_index)
2810                let expr = self.date_part_on_time_index("month")?;
2811                exprs.push(expr);
2812                ScalarFunc::GeneratedExpr
2813            }
2814            "year" => {
2815                // date_part('year', time_index)
2816                let expr = self.date_part_on_time_index("year")?;
2817                exprs.push(expr);
2818                ScalarFunc::GeneratedExpr
2819            }
2820            "day_of_month" => {
2821                // date_part('day', time_index)
2822                let expr = self.date_part_on_time_index("day")?;
2823                exprs.push(expr);
2824                ScalarFunc::GeneratedExpr
2825            }
2826            "day_of_week" => {
2827                // date_part('dow', time_index)
2828                let expr = self.date_part_on_time_index("dow")?;
2829                exprs.push(expr);
2830                ScalarFunc::GeneratedExpr
2831            }
2832            "day_of_year" => {
2833                // date_part('doy', time_index)
2834                let expr = self.date_part_on_time_index("doy")?;
2835                exprs.push(expr);
2836                ScalarFunc::GeneratedExpr
2837            }
2838            "days_in_month" => {
2839                // date_part(
2840                //     'days',
2841                //     (date_trunc('month', <TIME INDEX>::date) + interval '1 month - 1 day')
2842                // );
2843                let day_lit_expr = "day".lit();
2844                let month_lit_expr = "month".lit();
2845                let interval_1month_lit_expr =
2846                    DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
2847                let interval_1day_lit_expr = DfExpr::Literal(
2848                    ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
2849                    None,
2850                );
2851                let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
2852                    left: Box::new(interval_1month_lit_expr),
2853                    op: Operator::Minus,
2854                    right: Box::new(interval_1day_lit_expr),
2855                });
2856                let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
2857                    func: datafusion_functions::datetime::date_trunc(),
2858                    args: vec![month_lit_expr, self.create_time_index_column_expr()?],
2859                });
2860                let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
2861                    left: Box::new(date_trunc_expr),
2862                    op: Operator::Plus,
2863                    right: Box::new(the_1month_minus_1day_expr),
2864                });
2865                let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
2866                    func: datafusion_functions::datetime::date_part(),
2867                    args: vec![day_lit_expr, date_trunc_plus_interval_expr],
2868                });
2869
2870                exprs.push(date_part_expr);
2871                ScalarFunc::GeneratedExpr
2872            }
2873
2874            "label_join" => {
2875                self.ctx.use_tsid = false;
2876                let (concat_expr, dst_label) = Self::build_concat_labels_expr(
2877                    &mut other_input_exprs,
2878                    &self.ctx,
2879                    query_engine_state,
2880                )?;
2881
2882                // Reserve the current field columns except the `dst_label`.
2883                for value in &self.ctx.field_columns {
2884                    if *value != dst_label {
2885                        let expr = DfExpr::Column(Column::from_name(value));
2886                        exprs.push(expr);
2887                    }
2888                }
2889
2890                // Remove it from tag columns if exists to avoid duplicated column names
2891                self.ctx.tag_columns.retain(|tag| *tag != dst_label);
2892                new_tags.push(dst_label);
2893                // Add the new label expr to evaluate
2894                exprs.push(concat_expr);
2895
2896                ScalarFunc::GeneratedExpr
2897            }
2898            "label_replace" => {
2899                self.ctx.use_tsid = false;
2900                if let Some((replace_expr, dst_label)) = self
2901                    .build_regexp_replace_label_expr(&mut other_input_exprs, query_engine_state)?
2902                {
2903                    // Reserve the current field columns except the `dst_label`.
2904                    for value in &self.ctx.field_columns {
2905                        if *value != dst_label {
2906                            let expr = DfExpr::Column(Column::from_name(value));
2907                            exprs.push(expr);
2908                        }
2909                    }
2910
2911                    ensure!(
2912                        !self.ctx.tag_columns.contains(&dst_label),
2913                        SameLabelSetSnafu
2914                    );
2915                    new_tags.push(dst_label);
2916                    // Add the new label expr to evaluate
2917                    exprs.push(replace_expr);
2918                } else {
2919                    // Keep the current field columns
2920                    for value in &self.ctx.field_columns {
2921                        let expr = DfExpr::Column(Column::from_name(value));
2922                        exprs.push(expr);
2923                    }
2924                }
2925
2926                ScalarFunc::GeneratedExpr
2927            }
2928            "sort" | "sort_desc" | "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
2929                // These functions are not expression but a part of plan,
2930                // they are processed by `prom_call_expr_to_plan`.
2931                for value in &self.ctx.field_columns {
2932                    let expr = DfExpr::Column(Column::from_name(value));
2933                    exprs.push(expr);
2934                }
2935
2936                ScalarFunc::GeneratedExpr
2937            }
2938            "round" => {
2939                if other_input_exprs.is_empty() {
2940                    other_input_exprs.push_front(0.0f64.lit());
2941                }
2942                ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
2943            }
2944            "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
2945            "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
2946            "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
2947            "pi" => {
2948                // pi functions doesn't accepts any arguments, needs special processing
2949                let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
2950                    func: datafusion::functions::math::pi(),
2951                    args: vec![],
2952                });
2953                exprs.push(fn_expr);
2954
2955                ScalarFunc::GeneratedExpr
2956            }
2957            _ => {
2958                if let Some(f) = query_engine_state
2959                    .session_state()
2960                    .scalar_functions()
2961                    .get(func.name)
2962                {
2963                    ScalarFunc::DataFusionBuiltin(f.clone())
2964                } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
2965                    let func_state = query_engine_state.function_state();
2966                    let query_ctx = self.table_provider.query_ctx();
2967
2968                    ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
2969                        state: func_state,
2970                        query_ctx: query_ctx.clone(),
2971                    })))
2972                } else if let Some(f) = datafusion_functions::math::functions()
2973                    .iter()
2974                    .find(|f| f.name() == func.name)
2975                {
2976                    ScalarFunc::DataFusionUdf(f.clone())
2977                } else {
2978                    return UnsupportedExprSnafu {
2979                        name: func.name.to_string(),
2980                    }
2981                    .fail();
2982                }
2983            }
2984        };
2985
2986        for value in &self.ctx.field_columns {
2987            let col_expr = DfExpr::Column(Column::from_name(value));
2988
2989            match scalar_func.clone() {
2990                ScalarFunc::DataFusionBuiltin(func) => {
2991                    other_input_exprs.insert(field_column_pos, col_expr);
2992                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
2993                        func,
2994                        args: other_input_exprs.clone().into(),
2995                    });
2996                    exprs.push(fn_expr);
2997                    let _ = other_input_exprs.remove(field_column_pos);
2998                }
2999                ScalarFunc::DataFusionUdf(func) => {
3000                    let args = itertools::chain!(
3001                        other_input_exprs.iter().take(field_column_pos).cloned(),
3002                        std::iter::once(col_expr),
3003                        other_input_exprs.iter().skip(field_column_pos).cloned()
3004                    )
3005                    .collect_vec();
3006                    exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
3007                }
3008                ScalarFunc::Udf(func) => {
3009                    let ts_range_expr = DfExpr::Column(Column::from_name(
3010                        RangeManipulate::build_timestamp_range_name(
3011                            self.ctx.time_index_column.as_ref().unwrap(),
3012                        ),
3013                    ));
3014                    other_input_exprs.insert(field_column_pos, ts_range_expr);
3015                    other_input_exprs.insert(field_column_pos + 1, col_expr);
3016                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3017                        func,
3018                        args: other_input_exprs.clone().into(),
3019                    });
3020                    exprs.push(fn_expr);
3021                    let _ = other_input_exprs.remove(field_column_pos + 1);
3022                    let _ = other_input_exprs.remove(field_column_pos);
3023                }
3024                ScalarFunc::ExtrapolateUdf(func, range_length) => {
3025                    let ts_range_expr = DfExpr::Column(Column::from_name(
3026                        RangeManipulate::build_timestamp_range_name(
3027                            self.ctx.time_index_column.as_ref().unwrap(),
3028                        ),
3029                    ));
3030                    other_input_exprs.insert(field_column_pos, ts_range_expr);
3031                    other_input_exprs.insert(field_column_pos + 1, col_expr);
3032                    other_input_exprs
3033                        .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
3034                    other_input_exprs.push_back(lit(range_length));
3035                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3036                        func,
3037                        args: other_input_exprs.clone().into(),
3038                    });
3039                    exprs.push(fn_expr);
3040                    let _ = other_input_exprs.pop_back();
3041                    let _ = other_input_exprs.remove(field_column_pos + 2);
3042                    let _ = other_input_exprs.remove(field_column_pos + 1);
3043                    let _ = other_input_exprs.remove(field_column_pos);
3044                }
3045                ScalarFunc::GeneratedExpr => {}
3046            }
3047        }
3048
3049        // Update value columns' name, and alias them to remove qualifiers
3050        // For label functions such as `label_join`, `label_replace`, etc.,
3051        // we keep the fields unchanged.
3052        if !matches!(func.name, "label_join" | "label_replace") {
3053            let mut new_field_columns = Vec::with_capacity(exprs.len());
3054
3055            exprs = exprs
3056                .into_iter()
3057                .map(|expr| {
3058                    let display_name = expr.schema_name().to_string();
3059                    new_field_columns.push(display_name.clone());
3060                    Ok(expr.alias(display_name))
3061                })
3062                .collect::<std::result::Result<Vec<_>, _>>()
3063                .context(DataFusionPlanningSnafu)?;
3064
3065            self.ctx.field_columns = new_field_columns;
3066        }
3067
3068        Ok((exprs, new_tags))
3069    }
3070
3071    /// Validate label name according to Prometheus specification.
3072    /// Label names must match the regex: [a-zA-Z_][a-zA-Z0-9_]*
3073    /// Additionally, label names starting with double underscores are reserved for internal use.
3074    fn validate_label_name(label_name: &str) -> Result<()> {
3075        // Check if label name starts with double underscores (reserved)
3076        if label_name.starts_with("__") {
3077            return InvalidDestinationLabelNameSnafu { label_name }.fail();
3078        }
3079        // Check if label name matches the required pattern
3080        if !LABEL_NAME_REGEX.is_match(label_name) {
3081            return InvalidDestinationLabelNameSnafu { label_name }.fail();
3082        }
3083
3084        Ok(())
3085    }
3086
3087    /// Build expr for `label_replace` function
3088    fn build_regexp_replace_label_expr(
3089        &self,
3090        other_input_exprs: &mut VecDeque<DfExpr>,
3091        query_engine_state: &QueryEngineState,
3092    ) -> Result<Option<(DfExpr, String)>> {
3093        // label_replace(vector, dst_label, replacement, src_label, regex)
3094        let dst_label = match other_input_exprs.pop_front() {
3095            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
3096            other => UnexpectedPlanExprSnafu {
3097                desc: format!("expected dst_label string literal, but found {:?}", other),
3098            }
3099            .fail()?,
3100        };
3101
3102        // Validate the destination label name
3103        Self::validate_label_name(&dst_label)?;
3104        let replacement = match other_input_exprs.pop_front() {
3105            Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
3106            other => UnexpectedPlanExprSnafu {
3107                desc: format!("expected replacement string literal, but found {:?}", other),
3108            }
3109            .fail()?,
3110        };
3111        let src_label = match other_input_exprs.pop_front() {
3112            Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
3113            other => UnexpectedPlanExprSnafu {
3114                desc: format!("expected src_label string literal, but found {:?}", other),
3115            }
3116            .fail()?,
3117        };
3118
3119        let regex = match other_input_exprs.pop_front() {
3120            Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
3121            other => UnexpectedPlanExprSnafu {
3122                desc: format!("expected regex string literal, but found {:?}", other),
3123            }
3124            .fail()?,
3125        };
3126
3127        // Validate the regex before using it
3128        // doc: https://prometheus.io/docs/prometheus/latest/querying/functions/#label_replace
3129        regex::Regex::new(&regex).map_err(|_| {
3130            InvalidRegularExpressionSnafu {
3131                regex: regex.clone(),
3132            }
3133            .build()
3134        })?;
3135
3136        // If the src_label exists and regex is empty, keep everything unchanged.
3137        if self.ctx.tag_columns.contains(&src_label) && regex.is_empty() {
3138            return Ok(None);
3139        }
3140
3141        // If the src_label doesn't exists, and
3142        if !self.ctx.tag_columns.contains(&src_label) {
3143            if replacement.is_empty() {
3144                // the replacement is empty, keep everything unchanged.
3145                return Ok(None);
3146            } else {
3147                // the replacement is not empty, always adds dst_label with replacement value.
3148                return Ok(Some((
3149                    // alias literal `replacement` as dst_label
3150                    lit(replacement).alias(&dst_label),
3151                    dst_label,
3152                )));
3153            }
3154        }
3155
3156        // Preprocess the regex:
3157        // https://github.com/prometheus/prometheus/blob/d902abc50d6652ba8fe9a81ff8e5cce936114eba/promql/functions.go#L1575C32-L1575C37
3158        let regex = format!("^(?s:{regex})$");
3159
3160        let session_state = query_engine_state.session_state();
3161        let func = session_state
3162            .scalar_functions()
3163            .get("regexp_replace")
3164            .context(UnsupportedExprSnafu {
3165                name: "regexp_replace",
3166            })?;
3167
3168        // regexp_replace(src_label, regex, replacement)
3169        let args = vec![
3170            if src_label.is_empty() {
3171                DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
3172            } else {
3173                DfExpr::Column(Column::from_name(src_label))
3174            },
3175            DfExpr::Literal(ScalarValue::Utf8(Some(regex)), None),
3176            DfExpr::Literal(ScalarValue::Utf8(Some(replacement)), None),
3177        ];
3178
3179        Ok(Some((
3180            DfExpr::ScalarFunction(ScalarFunction {
3181                func: func.clone(),
3182                args,
3183            })
3184            .alias(&dst_label),
3185            dst_label,
3186        )))
3187    }
3188
3189    /// Build expr for `label_join` function
3190    fn build_concat_labels_expr(
3191        other_input_exprs: &mut VecDeque<DfExpr>,
3192        ctx: &PromPlannerContext,
3193        query_engine_state: &QueryEngineState,
3194    ) -> Result<(DfExpr, String)> {
3195        // label_join(vector, dst_label, separator, src_label_1, src_label_2, ...)
3196
3197        let dst_label = match other_input_exprs.pop_front() {
3198            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
3199            other => UnexpectedPlanExprSnafu {
3200                desc: format!("expected dst_label string literal, but found {:?}", other),
3201            }
3202            .fail()?,
3203        };
3204        let separator = match other_input_exprs.pop_front() {
3205            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
3206            other => UnexpectedPlanExprSnafu {
3207                desc: format!("expected separator string literal, but found {:?}", other),
3208            }
3209            .fail()?,
3210        };
3211
3212        // Create a set of available columns (tag columns + field columns + time index column)
3213        let available_columns: HashSet<&str> = ctx
3214            .tag_columns
3215            .iter()
3216            .chain(ctx.field_columns.iter())
3217            .chain(ctx.time_index_column.as_ref())
3218            .map(|s| s.as_str())
3219            .collect();
3220
3221        let src_labels = other_input_exprs
3222            .iter()
3223            .map(|expr| {
3224                // Cast source label into column or null literal
3225                match expr {
3226                    DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
3227                        if label.is_empty() {
3228                            Ok(DfExpr::Literal(ScalarValue::Null, None))
3229                        } else if available_columns.contains(label.as_str()) {
3230                            // Label exists in the table schema
3231                            Ok(DfExpr::Column(Column::from_name(label)))
3232                        } else {
3233                            // Label doesn't exist, treat as empty string (null)
3234                            Ok(DfExpr::Literal(ScalarValue::Null, None))
3235                        }
3236                    }
3237                    other => UnexpectedPlanExprSnafu {
3238                        desc: format!(
3239                            "expected source label string literal, but found {:?}",
3240                            other
3241                        ),
3242                    }
3243                    .fail(),
3244                }
3245            })
3246            .collect::<Result<Vec<_>>>()?;
3247        ensure!(
3248            !src_labels.is_empty(),
3249            FunctionInvalidArgumentSnafu {
3250                fn_name: "label_join"
3251            }
3252        );
3253
3254        let session_state = query_engine_state.session_state();
3255        let func = session_state
3256            .scalar_functions()
3257            .get("concat_ws")
3258            .context(UnsupportedExprSnafu { name: "concat_ws" })?;
3259
3260        // concat_ws(separator, src_label_1, src_label_2, ...) as dst_label
3261        let mut args = Vec::with_capacity(1 + src_labels.len());
3262        args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
3263        args.extend(src_labels);
3264
3265        Ok((
3266            DfExpr::ScalarFunction(ScalarFunction {
3267                func: func.clone(),
3268                args,
3269            })
3270            .alias(&dst_label),
3271            dst_label,
3272        ))
3273    }
3274
3275    fn create_time_index_column_expr(&self) -> Result<DfExpr> {
3276        Ok(DfExpr::Column(Column::from_name(
3277            self.ctx
3278                .time_index_column
3279                .clone()
3280                .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
3281        )))
3282    }
3283
3284    fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
3285        let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
3286        for tag in &self.ctx.tag_columns {
3287            let expr = DfExpr::Column(Column::from_name(tag));
3288            result.push(expr);
3289        }
3290        Ok(result)
3291    }
3292
3293    fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
3294        let mut result = Vec::with_capacity(self.ctx.field_columns.len());
3295        for field in &self.ctx.field_columns {
3296            let expr = DfExpr::Column(Column::from_name(field));
3297            result.push(expr);
3298        }
3299        Ok(result)
3300    }
3301
3302    fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
3303        let mut result = self
3304            .ctx
3305            .tag_columns
3306            .iter()
3307            .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
3308            .collect::<Vec<_>>();
3309        result.push(self.create_time_index_column_expr()?.sort(true, true));
3310        Ok(result)
3311    }
3312
3313    fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
3314        self.ctx
3315            .field_columns
3316            .iter()
3317            .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
3318            .collect::<Vec<_>>()
3319    }
3320
3321    fn create_sort_exprs_by_tags(
3322        func: &str,
3323        tags: Vec<DfExpr>,
3324        asc: bool,
3325    ) -> Result<Vec<SortExpr>> {
3326        ensure!(
3327            !tags.is_empty(),
3328            FunctionInvalidArgumentSnafu { fn_name: func }
3329        );
3330
3331        tags.iter()
3332            .map(|col| match col {
3333                DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
3334                    Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
3335                }
3336                other => UnexpectedPlanExprSnafu {
3337                    desc: format!("expected label string literal, but found {:?}", other),
3338                }
3339                .fail(),
3340            })
3341            .collect::<Result<Vec<_>>>()
3342    }
3343
3344    fn create_empty_values_filter_expr(&self) -> Result<DfExpr> {
3345        let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3346        for value in &self.ctx.field_columns {
3347            let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
3348            exprs.push(expr);
3349        }
3350
3351        // This error context should be computed lazily: the planner may set `ctx.table_name` to
3352        // `None` for derived expressions (e.g. after projecting the LHS of a vector-vector
3353        // comparison filter). Eagerly calling `table_ref()?` here can turn a valid plan into
3354        // a `TableNameNotFound` error even when `conjunction(exprs)` succeeds.
3355        conjunction(exprs).with_context(|| ValueNotFoundSnafu {
3356            table: self
3357                .table_ref()
3358                .map(|t| t.to_quoted_string())
3359                .unwrap_or_else(|_| "unknown".to_string()),
3360        })
3361    }
3362
3363    /// Creates a set of DataFusion `DfExpr::AggregateFunction` expressions for each value column using the specified aggregate function.
3364    ///
3365    /// # Side Effects
3366    ///
3367    /// This method modifies the value columns in the context by replacing them with the new columns
3368    /// created by the aggregate function application.
3369    ///
3370    /// # Returns
3371    ///
3372    /// Returns a tuple of `(aggregate_expressions, previous_field_expressions)` where:
3373    /// - `aggregate_expressions`: Expressions that apply the aggregate function to the original fields
3374    /// - `previous_field_expressions`: Original field expressions before aggregation. This is non-empty
3375    ///   only when the operation is `count_values`, as this operation requires preserving the original
3376    ///   values for grouping.
3377    ///
3378    fn create_aggregate_exprs(
3379        &mut self,
3380        op: TokenType,
3381        param: &Option<Box<PromExpr>>,
3382        input_plan: &LogicalPlan,
3383    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
3384        let mut non_col_args = Vec::new();
3385        let is_group_agg = op.id() == token::T_GROUP;
3386        if is_group_agg {
3387            ensure!(
3388                self.ctx.field_columns.len() == 1,
3389                MultiFieldsNotSupportedSnafu {
3390                    operator: "group()"
3391                }
3392            );
3393        }
3394        let aggr = match op.id() {
3395            token::T_SUM => sum_udaf(),
3396            token::T_QUANTILE => {
3397                let q =
3398                    Self::get_param_as_literal_expr(param, Some(op), Some(ArrowDataType::Float64))?;
3399                non_col_args.push(q);
3400                quantile_udaf()
3401            }
3402            token::T_AVG => avg_udaf(),
3403            token::T_COUNT_VALUES | token::T_COUNT => count_udaf(),
3404            token::T_MIN => min_udaf(),
3405            token::T_MAX => max_udaf(),
3406            // PromQL's `group()` aggregator produces 1 for each group.
3407            // Use `max(1.0)` (per-group) to match semantics and output type (Float64).
3408            token::T_GROUP => max_udaf(),
3409            token::T_STDDEV => stddev_pop_udaf(),
3410            token::T_STDVAR => var_pop_udaf(),
3411            token::T_TOPK | token::T_BOTTOMK => UnsupportedExprSnafu {
3412                name: format!("{op:?}"),
3413            }
3414            .fail()?,
3415            _ => UnexpectedTokenSnafu { token: op }.fail()?,
3416        };
3417
3418        // perform aggregate operation to each value column
3419        let exprs: Vec<DfExpr> = self
3420            .ctx
3421            .field_columns
3422            .iter()
3423            .map(|col| {
3424                if is_group_agg {
3425                    aggr.call(vec![lit(1_f64)])
3426                } else {
3427                    non_col_args.push(DfExpr::Column(Column::from_name(col)));
3428                    let expr = aggr.call(non_col_args.clone());
3429                    non_col_args.pop();
3430                    expr
3431                }
3432            })
3433            .collect::<Vec<_>>();
3434
3435        // if the aggregator is `count_values`, it must be grouped by current fields.
3436        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
3437            let prev_field_exprs: Vec<_> = self
3438                .ctx
3439                .field_columns
3440                .iter()
3441                .map(|col| DfExpr::Column(Column::from_name(col)))
3442                .collect();
3443
3444            ensure!(
3445                self.ctx.field_columns.len() == 1,
3446                UnsupportedExprSnafu {
3447                    name: "count_values on multi-value input"
3448                }
3449            );
3450
3451            prev_field_exprs
3452        } else {
3453            vec![]
3454        };
3455
3456        // update value column name according to the aggregators,
3457        let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
3458
3459        let normalized_exprs =
3460            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
3461        for expr in normalized_exprs {
3462            new_field_columns.push(expr.schema_name().to_string());
3463        }
3464        self.ctx.field_columns = new_field_columns;
3465
3466        Ok((exprs, prev_field_exprs))
3467    }
3468
3469    fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
3470        let param = param
3471            .as_deref()
3472            .with_context(|| FunctionInvalidArgumentSnafu {
3473                fn_name: op.to_string(),
3474            })?;
3475        let PromExpr::StringLiteral(StringLiteral { val }) = param else {
3476            return FunctionInvalidArgumentSnafu {
3477                fn_name: op.to_string(),
3478            }
3479            .fail();
3480        };
3481
3482        Ok(val)
3483    }
3484
3485    fn get_param_as_literal_expr(
3486        param: &Option<Box<PromExpr>>,
3487        op: Option<TokenType>,
3488        expected_type: Option<ArrowDataType>,
3489    ) -> Result<DfExpr> {
3490        let prom_param = param.as_deref().with_context(|| {
3491            if let Some(op) = op {
3492                FunctionInvalidArgumentSnafu {
3493                    fn_name: op.to_string(),
3494                }
3495            } else {
3496                FunctionInvalidArgumentSnafu {
3497                    fn_name: "unknown".to_string(),
3498                }
3499            }
3500        })?;
3501
3502        let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
3503            if let Some(op) = op {
3504                FunctionInvalidArgumentSnafu {
3505                    fn_name: op.to_string(),
3506                }
3507            } else {
3508                FunctionInvalidArgumentSnafu {
3509                    fn_name: "unknown".to_string(),
3510                }
3511            }
3512        })?;
3513
3514        // check if the type is expected
3515        if let Some(expected_type) = expected_type {
3516            // literal should not have reference to column
3517            let expr_type = expr
3518                .get_type(&DFSchema::empty())
3519                .context(DataFusionPlanningSnafu)?;
3520            if expected_type != expr_type {
3521                return FunctionInvalidArgumentSnafu {
3522                    fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
3523                }
3524                .fail();
3525            }
3526        }
3527
3528        Ok(expr)
3529    }
3530
3531    /// Create [DfExpr::WindowFunction] expr for each value column with given window function.
3532    ///
3533    fn create_window_exprs(
3534        &mut self,
3535        op: TokenType,
3536        group_exprs: Vec<DfExpr>,
3537        input_plan: &LogicalPlan,
3538    ) -> Result<Vec<DfExpr>> {
3539        ensure!(
3540            self.ctx.field_columns.len() == 1,
3541            UnsupportedExprSnafu {
3542                name: "topk or bottomk on multi-value input"
3543            }
3544        );
3545
3546        assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
3547
3548        let asc = matches!(op.id(), token::T_BOTTOMK);
3549
3550        let tag_sort_exprs = self
3551            .create_tag_column_exprs()?
3552            .into_iter()
3553            .map(|expr| expr.sort(asc, true));
3554
3555        // perform window operation to each value column
3556        let exprs: Vec<DfExpr> = self
3557            .ctx
3558            .field_columns
3559            .iter()
3560            .map(|col| {
3561                let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1);
3562                // Order by value in the specific order
3563                sort_exprs.push(DfExpr::Column(Column::from(col)).sort(asc, true));
3564                // Then tags if the values are equal,
3565                // Try to ensure the relative stability of the output results.
3566                sort_exprs.extend(tag_sort_exprs.clone());
3567
3568                DfExpr::WindowFunction(Box::new(WindowFunction {
3569                    fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
3570                    params: WindowFunctionParams {
3571                        args: vec![],
3572                        partition_by: group_exprs.clone(),
3573                        order_by: sort_exprs,
3574                        window_frame: WindowFrame::new(Some(true)),
3575                        null_treatment: None,
3576                        distinct: false,
3577                        filter: None,
3578                    },
3579                }))
3580            })
3581            .collect();
3582
3583        let normalized_exprs =
3584            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
3585        Ok(normalized_exprs)
3586    }
3587
3588    /// Try to build a [f64] from [PromExpr].
3589    #[deprecated(
3590        note = "use `Self::get_param_as_literal_expr` instead. This is only for `create_histogram_plan`"
3591    )]
3592    fn try_build_float_literal(expr: &PromExpr) -> Option<f64> {
3593        match expr {
3594            PromExpr::NumberLiteral(NumberLiteral { val }) => Some(*val),
3595            PromExpr::Paren(ParenExpr { expr }) => Self::try_build_float_literal(expr),
3596            PromExpr::Unary(UnaryExpr { expr, .. }) => {
3597                Self::try_build_float_literal(expr).map(|f| -f)
3598            }
3599            PromExpr::StringLiteral(_)
3600            | PromExpr::Binary(_)
3601            | PromExpr::VectorSelector(_)
3602            | PromExpr::MatrixSelector(_)
3603            | PromExpr::Call(_)
3604            | PromExpr::Extension(_)
3605            | PromExpr::Aggregate(_)
3606            | PromExpr::Subquery(_) => None,
3607        }
3608    }
3609
3610    /// Create a [SPECIAL_HISTOGRAM_QUANTILE] plan.
3611    async fn create_histogram_plan(
3612        &mut self,
3613        args: &PromFunctionArgs,
3614        query_engine_state: &QueryEngineState,
3615    ) -> Result<LogicalPlan> {
3616        if args.args.len() != 2 {
3617            return FunctionInvalidArgumentSnafu {
3618                fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(),
3619            }
3620            .fail();
3621        }
3622        #[allow(deprecated)]
3623        let phi = Self::try_build_float_literal(&args.args[0]).with_context(|| {
3624            FunctionInvalidArgumentSnafu {
3625                fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(),
3626            }
3627        })?;
3628
3629        let input = args.args[1].as_ref().clone();
3630        let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?;
3631        // `histogram_quantile` folds buckets across `le`, so `__tsid` (which includes `le`) is not
3632        // a stable series identifier anymore. Also, HistogramFold infers label columns from the
3633        // input schema and must not treat `__tsid` as a label column.
3634        let input_plan = self.strip_tsid_column(input_plan)?;
3635        self.ctx.use_tsid = false;
3636
3637        if !self.ctx.has_le_tag() {
3638            // Return empty result instead of error when 'le' column is not found
3639            // This handles the case when histogram metrics don't exist
3640            return Ok(LogicalPlan::EmptyRelation(
3641                datafusion::logical_expr::EmptyRelation {
3642                    produce_one_row: false,
3643                    schema: Arc::new(DFSchema::empty()),
3644                },
3645            ));
3646        }
3647        let time_index_column =
3648            self.ctx
3649                .time_index_column
3650                .clone()
3651                .with_context(|| TimeIndexNotFoundSnafu {
3652                    table: self.ctx.table_name.clone().unwrap_or_default(),
3653                })?;
3654        // FIXME(ruihang): support multi fields
3655        let field_column = self
3656            .ctx
3657            .field_columns
3658            .first()
3659            .with_context(|| FunctionInvalidArgumentSnafu {
3660                fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(),
3661            })?
3662            .clone();
3663        // remove le column from tag columns
3664        self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME);
3665
3666        Ok(LogicalPlan::Extension(Extension {
3667            node: Arc::new(
3668                HistogramFold::new(
3669                    LE_COLUMN_NAME.to_string(),
3670                    field_column,
3671                    time_index_column,
3672                    phi,
3673                    input_plan,
3674                )
3675                .context(DataFusionPlanningSnafu)?,
3676            ),
3677        }))
3678    }
3679
3680    /// Create a [SPECIAL_VECTOR_FUNCTION] plan
3681    async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result<LogicalPlan> {
3682        if args.args.len() != 1 {
3683            return FunctionInvalidArgumentSnafu {
3684                fn_name: SPECIAL_VECTOR_FUNCTION.to_string(),
3685            }
3686            .fail();
3687        }
3688        let lit = Self::get_param_as_literal_expr(&Some(args.args[0].clone()), None, None)?;
3689
3690        // reuse `SPECIAL_TIME_FUNCTION` as name of time index column
3691        self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3692        self.ctx.reset_table_name_and_schema();
3693        self.ctx.tag_columns = vec![];
3694        self.ctx.field_columns = vec![greptime_value().to_string()];
3695        Ok(LogicalPlan::Extension(Extension {
3696            node: Arc::new(
3697                EmptyMetric::new(
3698                    self.ctx.start,
3699                    self.ctx.end,
3700                    self.ctx.interval,
3701                    SPECIAL_TIME_FUNCTION.to_string(),
3702                    greptime_value().to_string(),
3703                    Some(lit),
3704                )
3705                .context(DataFusionPlanningSnafu)?,
3706            ),
3707        }))
3708    }
3709
3710    /// Create a [SCALAR_FUNCTION] plan
3711    async fn create_scalar_plan(
3712        &mut self,
3713        args: &PromFunctionArgs,
3714        query_engine_state: &QueryEngineState,
3715    ) -> Result<LogicalPlan> {
3716        ensure!(
3717            args.len() == 1,
3718            FunctionInvalidArgumentSnafu {
3719                fn_name: SCALAR_FUNCTION
3720            }
3721        );
3722        let input = self
3723            .prom_expr_to_plan(&args.args[0], query_engine_state)
3724            .await?;
3725        ensure!(
3726            self.ctx.field_columns.len() == 1,
3727            MultiFieldsNotSupportedSnafu {
3728                operator: SCALAR_FUNCTION
3729            },
3730        );
3731        let scalar_plan = LogicalPlan::Extension(Extension {
3732            node: Arc::new(
3733                ScalarCalculate::new(
3734                    self.ctx.start,
3735                    self.ctx.end,
3736                    self.ctx.interval,
3737                    input,
3738                    self.ctx.time_index_column.as_ref().unwrap(),
3739                    &self.ctx.tag_columns,
3740                    &self.ctx.field_columns[0],
3741                    self.ctx.table_name.as_deref(),
3742                )
3743                .context(PromqlPlanNodeSnafu)?,
3744            ),
3745        });
3746        // scalar plan have no tag columns
3747        self.ctx.tag_columns.clear();
3748        self.ctx.field_columns.clear();
3749        self.ctx
3750            .field_columns
3751            .push(scalar_plan.schema().field(1).name().clone());
3752        Ok(scalar_plan)
3753    }
3754
3755    /// Create a [SPECIAL_ABSENT_FUNCTION] plan
3756    async fn create_absent_plan(
3757        &mut self,
3758        args: &PromFunctionArgs,
3759        query_engine_state: &QueryEngineState,
3760    ) -> Result<LogicalPlan> {
3761        if args.args.len() != 1 {
3762            return FunctionInvalidArgumentSnafu {
3763                fn_name: SPECIAL_ABSENT_FUNCTION.to_string(),
3764            }
3765            .fail();
3766        }
3767        let input = self
3768            .prom_expr_to_plan(&args.args[0], query_engine_state)
3769            .await?;
3770
3771        let time_index_expr = self.create_time_index_column_expr()?;
3772        let first_field_expr =
3773            self.create_field_column_exprs()?
3774                .pop()
3775                .with_context(|| ValueNotFoundSnafu {
3776                    table: self.ctx.table_name.clone().unwrap_or_default(),
3777                })?;
3778        let first_value_expr = first_value(first_field_expr, vec![]);
3779
3780        let ordered_aggregated_input = LogicalPlanBuilder::from(input)
3781            .aggregate(
3782                vec![time_index_expr.clone()],
3783                vec![first_value_expr.clone()],
3784            )
3785            .context(DataFusionPlanningSnafu)?
3786            .sort(vec![time_index_expr.sort(true, false)])
3787            .context(DataFusionPlanningSnafu)?
3788            .build()
3789            .context(DataFusionPlanningSnafu)?;
3790
3791        let fake_labels = self
3792            .ctx
3793            .selector_matcher
3794            .iter()
3795            .filter_map(|matcher| match matcher.op {
3796                MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())),
3797                _ => None,
3798            })
3799            .collect::<Vec<_>>();
3800
3801        // Create the absent plan
3802        let absent_plan = LogicalPlan::Extension(Extension {
3803            node: Arc::new(
3804                Absent::try_new(
3805                    self.ctx.start,
3806                    self.ctx.end,
3807                    self.ctx.interval,
3808                    self.ctx.time_index_column.as_ref().unwrap().clone(),
3809                    self.ctx.field_columns[0].clone(),
3810                    fake_labels,
3811                    ordered_aggregated_input,
3812                )
3813                .context(DataFusionPlanningSnafu)?,
3814            ),
3815        });
3816
3817        Ok(absent_plan)
3818    }
3819
3820    /// Try to build a DataFusion Literal Expression from PromQL Expr, return
3821    /// `None` if the input is not a literal expression.
3822    fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
3823        match expr {
3824            PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
3825            PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
3826            PromExpr::VectorSelector(_)
3827            | PromExpr::MatrixSelector(_)
3828            | PromExpr::Extension(_)
3829            | PromExpr::Aggregate(_)
3830            | PromExpr::Subquery(_) => None,
3831            PromExpr::Call(Call { func, .. }) => {
3832                if func.name == SPECIAL_TIME_FUNCTION {
3833                    // For time() function, don't treat it as a literal
3834                    // Let it be handled as a regular function call
3835                    None
3836                } else {
3837                    None
3838                }
3839            }
3840            PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
3841            // TODO(ruihang): support Unary operator
3842            PromExpr::Unary(UnaryExpr { expr, .. }) => Self::try_build_literal_expr(expr),
3843            PromExpr::Binary(PromBinaryExpr {
3844                lhs,
3845                rhs,
3846                op,
3847                modifier,
3848            }) => {
3849                let lhs = Self::try_build_literal_expr(lhs)?;
3850                let rhs = Self::try_build_literal_expr(rhs)?;
3851                let is_comparison_op = Self::is_token_a_comparison_op(*op);
3852                let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
3853                let expr = expr_builder(lhs, rhs).ok()?;
3854
3855                let should_return_bool = if let Some(m) = modifier {
3856                    m.return_bool
3857                } else {
3858                    false
3859                };
3860                if is_comparison_op && should_return_bool {
3861                    Some(DfExpr::Cast(Cast {
3862                        expr: Box::new(expr),
3863                        data_type: ArrowDataType::Float64,
3864                    }))
3865                } else {
3866                    Some(expr)
3867                }
3868            }
3869        }
3870    }
3871
3872    fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
3873        match expr {
3874            PromExpr::Call(Call { func, .. }) => {
3875                if func.name == SPECIAL_TIME_FUNCTION
3876                    && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
3877                {
3878                    Some(build_special_time_expr(time_index_col))
3879                } else {
3880                    None
3881                }
3882            }
3883            _ => None,
3884        }
3885    }
3886
3887    /// Return a lambda to build binary expression from token.
3888    /// Because some binary operator are function in DataFusion like `atan2` or `^`.
3889    #[allow(clippy::type_complexity)]
3890    fn prom_token_to_binary_expr_builder(
3891        token: TokenType,
3892    ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
3893        let cast_float = |expr| {
3894            if matches!(
3895                &expr,
3896                DfExpr::Cast(Cast {
3897                    data_type: ArrowDataType::Float64,
3898                    ..
3899                })
3900            ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
3901            {
3902                expr
3903            } else {
3904                DfExpr::Cast(Cast {
3905                    expr: Box::new(expr),
3906                    data_type: ArrowDataType::Float64,
3907                })
3908            }
3909        };
3910        match token.id() {
3911            token::T_ADD => Ok(Box::new(move |lhs, rhs| {
3912                Ok(cast_float(lhs) + cast_float(rhs))
3913            })),
3914            token::T_SUB => Ok(Box::new(move |lhs, rhs| {
3915                Ok(cast_float(lhs) - cast_float(rhs))
3916            })),
3917            token::T_MUL => Ok(Box::new(move |lhs, rhs| {
3918                Ok(cast_float(lhs) * cast_float(rhs))
3919            })),
3920            token::T_DIV => Ok(Box::new(move |lhs, rhs| {
3921                Ok(cast_float(lhs) / cast_float(rhs))
3922            })),
3923            token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
3924                Ok(cast_float(lhs) % cast_float(rhs))
3925            })),
3926            token::T_EQLC => Ok(Box::new(|lhs, rhs| Ok(lhs.eq(rhs)))),
3927            token::T_NEQ => Ok(Box::new(|lhs, rhs| Ok(lhs.not_eq(rhs)))),
3928            token::T_GTR => Ok(Box::new(|lhs, rhs| Ok(lhs.gt(rhs)))),
3929            token::T_LSS => Ok(Box::new(|lhs, rhs| Ok(lhs.lt(rhs)))),
3930            token::T_GTE => Ok(Box::new(|lhs, rhs| Ok(lhs.gt_eq(rhs)))),
3931            token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))),
3932            token::T_POW => Ok(Box::new(move |lhs, rhs| {
3933                Ok(DfExpr::ScalarFunction(ScalarFunction {
3934                    func: datafusion_functions::math::power(),
3935                    args: vec![cast_float(lhs), cast_float(rhs)],
3936                }))
3937            })),
3938            token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
3939                Ok(DfExpr::ScalarFunction(ScalarFunction {
3940                    func: datafusion_functions::math::atan2(),
3941                    args: vec![cast_float(lhs), cast_float(rhs)],
3942                }))
3943            })),
3944            _ => UnexpectedTokenSnafu { token }.fail(),
3945        }
3946    }
3947
3948    /// Check if the given op is a [comparison operator](https://prometheus.io/docs/prometheus/latest/querying/operators/#comparison-binary-operators).
3949    fn is_token_a_comparison_op(token: TokenType) -> bool {
3950        matches!(
3951            token.id(),
3952            token::T_EQLC
3953                | token::T_NEQ
3954                | token::T_GTR
3955                | token::T_LSS
3956                | token::T_GTE
3957                | token::T_LTE
3958        )
3959    }
3960
3961    /// Check if the given op is a set operator (UNION, INTERSECT and EXCEPT in SQL).
3962    fn is_token_a_set_op(token: TokenType) -> bool {
3963        matches!(
3964            token.id(),
3965            token::T_LAND // INTERSECT
3966                | token::T_LOR // UNION
3967                | token::T_LUNLESS // EXCEPT
3968        )
3969    }
3970
3971    fn align_binary_field_columns<'a>(
3972        left_field_columns: &'a [String],
3973        right_field_columns: &'a [String],
3974    ) -> (Vec<String>, Vec<(&'a String, &'a String)>) {
3975        let field_pairs = left_field_columns
3976            .iter()
3977            .zip(right_field_columns.iter())
3978            .collect::<Vec<_>>();
3979        let output_field_columns = field_pairs
3980            .iter()
3981            .map(|(left_col_name, _)| (*left_col_name).clone())
3982            .collect();
3983        (output_field_columns, field_pairs)
3984    }
3985
3986    fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
3987        plan.schema()
3988            .fields()
3989            .iter()
3990            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
3991    }
3992
3993    fn optional_tsid_projection(
3994        schema: &DFSchemaRef,
3995        table_ref: Option<&TableReference>,
3996        keep_tsid: bool,
3997    ) -> Option<DfExpr> {
3998        keep_tsid.then_some(()).and_then(|_| {
3999            schema
4000                .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
4001                .ok()
4002                .map(|field| DfExpr::Column(field.into()))
4003        })
4004    }
4005
4006    fn binary_join_key_columns(
4007        &self,
4008        left_schema: &DFSchemaRef,
4009        right_schema: &DFSchemaRef,
4010        left_context: &PromPlannerContext,
4011        right_context: &PromPlannerContext,
4012        only_join_time_index: bool,
4013        modifier: &Option<BinModifier>,
4014    ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
4015        let has_tsid = |schema: &DFSchemaRef| {
4016            schema
4017                .fields()
4018                .iter()
4019                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
4020        };
4021        let use_tsid_join = !only_join_time_index
4022            && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
4023            && left_context.use_tsid
4024            && right_context.use_tsid
4025            && has_tsid(left_schema)
4026            && has_tsid(right_schema);
4027
4028        let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
4029            (
4030                BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
4031                BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
4032            )
4033        } else {
4034            if only_join_time_index {
4035                (BTreeSet::new(), BTreeSet::new())
4036            } else {
4037                (
4038                    left_context
4039                        .tag_columns
4040                        .iter()
4041                        .cloned()
4042                        .collect::<BTreeSet<_>>(),
4043                    right_context
4044                        .tag_columns
4045                        .iter()
4046                        .cloned()
4047                        .collect::<BTreeSet<_>>(),
4048                )
4049            }
4050        };
4051
4052        if !use_tsid_join
4053            && let Some(modifier) = modifier
4054            && let Some(matching) = &modifier.matching
4055        {
4056            match matching {
4057                LabelModifier::Include(on) => {
4058                    let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
4059                    left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
4060                    right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
4061                }
4062                LabelModifier::Exclude(ignoring) => {
4063                    for label in &ignoring.labels {
4064                        let _ = left_tag_columns.remove(label);
4065                        let _ = right_tag_columns.remove(label);
4066                    }
4067                }
4068            }
4069        }
4070
4071        let force_empty_join =
4072            !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
4073        if force_empty_join {
4074            let common_tag_columns = left_tag_columns
4075                .intersection(&right_tag_columns)
4076                .cloned()
4077                .collect::<BTreeSet<_>>();
4078            left_tag_columns = common_tag_columns.clone();
4079            right_tag_columns = common_tag_columns;
4080        }
4081
4082        Ok((left_tag_columns, right_tag_columns, force_empty_join))
4083    }
4084
4085    fn binary_modifier_preserves_tsid_join_key(
4086        &self,
4087        left_context: &PromPlannerContext,
4088        right_context: &PromPlannerContext,
4089        modifier: &Option<BinModifier>,
4090    ) -> bool {
4091        let Some(modifier) = modifier else {
4092            return true;
4093        };
4094
4095        if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
4096            return false;
4097        }
4098
4099        match &modifier.matching {
4100            None => true,
4101            Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
4102                !left_context.tag_columns.contains(label)
4103                    && !right_context.tag_columns.contains(label)
4104            }),
4105            Some(LabelModifier::Include(on)) => {
4106                let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
4107                let left_labels = left_context
4108                    .tag_columns
4109                    .iter()
4110                    .cloned()
4111                    .collect::<BTreeSet<_>>();
4112                let right_labels = right_context
4113                    .tag_columns
4114                    .iter()
4115                    .cloned()
4116                    .collect::<BTreeSet<_>>();
4117
4118                on_labels == left_labels && on_labels == right_labels
4119            }
4120        }
4121    }
4122
4123    /// Build a inner join on time index column and tag columns to concat two logical plans.
4124    /// When `only_join_time_index == true` we only join on the time index, because these two plan may not have the same tag columns
4125    #[allow(clippy::too_many_arguments)]
4126    fn join_on_non_field_columns(
4127        &self,
4128        left: LogicalPlan,
4129        right: LogicalPlan,
4130        left_table_ref: TableReference,
4131        right_table_ref: TableReference,
4132        left_time_index_column: Option<String>,
4133        right_time_index_column: Option<String>,
4134        only_join_time_index: bool,
4135        modifier: &Option<BinModifier>,
4136        left_context: &PromPlannerContext,
4137        right_context: &PromPlannerContext,
4138    ) -> Result<LogicalPlan> {
4139        let (mut left_tag_columns, mut right_tag_columns, force_empty_join) = self
4140            .binary_join_key_columns(
4141                left.schema(),
4142                right.schema(),
4143                left_context,
4144                right_context,
4145                only_join_time_index,
4146                modifier,
4147            )?;
4148
4149        // push time index column if it exists
4150        if let (Some(left_time_index_column), Some(right_time_index_column)) =
4151            (left_time_index_column, right_time_index_column)
4152        {
4153            left_tag_columns.insert(left_time_index_column);
4154            right_tag_columns.insert(right_time_index_column);
4155        }
4156
4157        let right = LogicalPlanBuilder::from(right)
4158            .alias(right_table_ref)
4159            .context(DataFusionPlanningSnafu)?
4160            .build()
4161            .context(DataFusionPlanningSnafu)?;
4162
4163        // Inner Join on time index column to concat two operator
4164        LogicalPlanBuilder::from(left)
4165            .alias(left_table_ref)
4166            .context(DataFusionPlanningSnafu)?
4167            .join_detailed(
4168                right,
4169                JoinType::Inner,
4170                (
4171                    left_tag_columns
4172                        .into_iter()
4173                        .map(Column::from_name)
4174                        .collect::<Vec<_>>(),
4175                    right_tag_columns
4176                        .into_iter()
4177                        .map(Column::from_name)
4178                        .collect::<Vec<_>>(),
4179                ),
4180                force_empty_join.then_some(lit(false)),
4181                NullEquality::NullEqualsNull,
4182            )
4183            .context(DataFusionPlanningSnafu)?
4184            .build()
4185            .context(DataFusionPlanningSnafu)
4186    }
4187
4188    /// Build a set operator (AND/OR/UNLESS)
4189    fn set_op_on_non_field_columns(
4190        &mut self,
4191        left: LogicalPlan,
4192        mut right: LogicalPlan,
4193        left_context: PromPlannerContext,
4194        right_context: PromPlannerContext,
4195        op: TokenType,
4196        modifier: &Option<BinModifier>,
4197    ) -> Result<LogicalPlan> {
4198        let mut left_tag_col_set = left_context
4199            .tag_columns
4200            .iter()
4201            .cloned()
4202            .collect::<HashSet<_>>();
4203        let mut right_tag_col_set = right_context
4204            .tag_columns
4205            .iter()
4206            .cloned()
4207            .collect::<HashSet<_>>();
4208
4209        if matches!(op.id(), token::T_LOR) {
4210            return self.or_operator(
4211                left,
4212                right,
4213                left_tag_col_set,
4214                right_tag_col_set,
4215                left_context,
4216                right_context,
4217                modifier,
4218            );
4219        }
4220
4221        // apply modifier
4222        if let Some(modifier) = modifier {
4223            // one-to-many and many-to-one are not supported
4224            ensure!(
4225                matches!(
4226                    modifier.card,
4227                    VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany
4228                ),
4229                UnsupportedVectorMatchSnafu {
4230                    name: modifier.card.clone(),
4231                },
4232            );
4233            // apply label modifier
4234            if let Some(matching) = &modifier.matching {
4235                match matching {
4236                    // keeps columns mentioned in `on`
4237                    LabelModifier::Include(on) => {
4238                        let mask = on.labels.iter().cloned().collect::<HashSet<_>>();
4239                        left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect();
4240                        right_tag_col_set =
4241                            right_tag_col_set.intersection(&mask).cloned().collect();
4242                    }
4243                    // removes columns memtioned in `ignoring`
4244                    LabelModifier::Exclude(ignoring) => {
4245                        // doesn't check existence of label
4246                        for label in &ignoring.labels {
4247                            let _ = left_tag_col_set.remove(label);
4248                            let _ = right_tag_col_set.remove(label);
4249                        }
4250                    }
4251                }
4252            }
4253        }
4254        // ensure two sides have the same tag columns
4255        if !matches!(op.id(), token::T_LOR) {
4256            ensure!(
4257                left_tag_col_set == right_tag_col_set,
4258                CombineTableColumnMismatchSnafu {
4259                    left: left_tag_col_set.into_iter().collect::<Vec<_>>(),
4260                    right: right_tag_col_set.into_iter().collect::<Vec<_>>(),
4261                }
4262            )
4263        };
4264        let left_time_index = left_context.time_index_column.clone().unwrap();
4265        let right_time_index = right_context.time_index_column.clone().unwrap();
4266        let join_keys = left_tag_col_set
4267            .iter()
4268            .cloned()
4269            .chain([left_time_index.clone()])
4270            .collect::<Vec<_>>();
4271        self.ctx.time_index_column = Some(left_time_index.clone());
4272        self.ctx.use_tsid = left_context.use_tsid;
4273
4274        // alias right time index column if necessary
4275        if left_context.time_index_column != right_context.time_index_column {
4276            let right_project_exprs = right
4277                .schema()
4278                .fields()
4279                .iter()
4280                .map(|field| {
4281                    if field.name() == &right_time_index {
4282                        DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index)
4283                    } else {
4284                        DfExpr::Column(Column::from_name(field.name()))
4285                    }
4286                })
4287                .collect::<Vec<_>>();
4288
4289            right = LogicalPlanBuilder::from(right)
4290                .project(right_project_exprs)
4291                .context(DataFusionPlanningSnafu)?
4292                .build()
4293                .context(DataFusionPlanningSnafu)?;
4294        }
4295
4296        ensure!(
4297            left_context.field_columns.len() == 1,
4298            MultiFieldsNotSupportedSnafu {
4299                operator: "AND operator"
4300            }
4301        );
4302        // Update the field column in context.
4303        // The AND/UNLESS operator only keep the field column in left input.
4304        let left_field_col = left_context.field_columns.first().unwrap();
4305        self.ctx.field_columns = vec![left_field_col.clone()];
4306
4307        // Generate join plan.
4308        // All set operations in PromQL are "distinct"
4309        match op.id() {
4310            token::T_LAND => LogicalPlanBuilder::from(left)
4311                .distinct()
4312                .context(DataFusionPlanningSnafu)?
4313                .join_detailed(
4314                    right,
4315                    JoinType::LeftSemi,
4316                    (join_keys.clone(), join_keys),
4317                    None,
4318                    NullEquality::NullEqualsNull,
4319                )
4320                .context(DataFusionPlanningSnafu)?
4321                .build()
4322                .context(DataFusionPlanningSnafu),
4323            token::T_LUNLESS => LogicalPlanBuilder::from(left)
4324                .distinct()
4325                .context(DataFusionPlanningSnafu)?
4326                .join_detailed(
4327                    right,
4328                    JoinType::LeftAnti,
4329                    (join_keys.clone(), join_keys),
4330                    None,
4331                    NullEquality::NullEqualsNull,
4332                )
4333                .context(DataFusionPlanningSnafu)?
4334                .build()
4335                .context(DataFusionPlanningSnafu),
4336            token::T_LOR => {
4337                // OR is handled at the beginning of this function, as it cannot
4338                // be expressed using JOIN like AND and UNLESS.
4339                unreachable!()
4340            }
4341            _ => UnexpectedTokenSnafu { token: op }.fail(),
4342        }
4343    }
4344
4345    // TODO(ruihang): change function name
4346    #[allow(clippy::too_many_arguments)]
4347    fn or_operator(
4348        &mut self,
4349        left: LogicalPlan,
4350        right: LogicalPlan,
4351        left_tag_cols_set: HashSet<String>,
4352        right_tag_cols_set: HashSet<String>,
4353        left_context: PromPlannerContext,
4354        right_context: PromPlannerContext,
4355        modifier: &Option<BinModifier>,
4356    ) -> Result<LogicalPlan> {
4357        // checks
4358        ensure!(
4359            left_context.field_columns.len() == right_context.field_columns.len(),
4360            CombineTableColumnMismatchSnafu {
4361                left: left_context.field_columns.clone(),
4362                right: right_context.field_columns.clone()
4363            }
4364        );
4365        ensure!(
4366            left_context.field_columns.len() == 1,
4367            MultiFieldsNotSupportedSnafu {
4368                operator: "OR operator"
4369            }
4370        );
4371
4372        // prepare hash sets
4373        let all_tags = left_tag_cols_set
4374            .union(&right_tag_cols_set)
4375            .cloned()
4376            .collect::<HashSet<_>>();
4377        let tags_not_in_left = all_tags
4378            .difference(&left_tag_cols_set)
4379            .cloned()
4380            .collect::<Vec<_>>();
4381        let tags_not_in_right = all_tags
4382            .difference(&right_tag_cols_set)
4383            .cloned()
4384            .collect::<Vec<_>>();
4385        let left_qualifier = left.schema().qualified_field(0).0.cloned();
4386        let right_qualifier = right.schema().qualified_field(0).0.cloned();
4387        let left_qualifier_string = left_qualifier
4388            .as_ref()
4389            .map(|l| l.to_string())
4390            .unwrap_or_default();
4391        let right_qualifier_string = right_qualifier
4392            .as_ref()
4393            .map(|r| r.to_string())
4394            .unwrap_or_default();
4395        let left_time_index_column =
4396            left_context
4397                .time_index_column
4398                .clone()
4399                .with_context(|| TimeIndexNotFoundSnafu {
4400                    table: left_qualifier_string.clone(),
4401                })?;
4402        let right_time_index_column =
4403            right_context
4404                .time_index_column
4405                .clone()
4406                .with_context(|| TimeIndexNotFoundSnafu {
4407                    table: right_qualifier_string.clone(),
4408                })?;
4409        // Take the name of first field column. The length is checked above.
4410        let left_field_col = left_context.field_columns.first().unwrap();
4411        let right_field_col = right_context.field_columns.first().unwrap();
4412        let left_has_tsid = left
4413            .schema()
4414            .fields()
4415            .iter()
4416            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
4417        let right_has_tsid = right
4418            .schema()
4419            .fields()
4420            .iter()
4421            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
4422
4423        // step 0: fill all columns in output schema
4424        let mut all_columns_set = left
4425            .schema()
4426            .fields()
4427            .iter()
4428            .chain(right.schema().fields().iter())
4429            .map(|field| field.name().clone())
4430            .collect::<HashSet<_>>();
4431        // Keep `__tsid` only when both sides contain it, otherwise it may break schema alignment
4432        // (e.g. `unknown_metric or some_metric`).
4433        if !(left_has_tsid && right_has_tsid) {
4434            all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME);
4435        }
4436        // remove time index column
4437        all_columns_set.remove(&left_time_index_column);
4438        all_columns_set.remove(&right_time_index_column);
4439        // remove field column in the right
4440        if left_field_col != right_field_col {
4441            all_columns_set.remove(right_field_col);
4442        }
4443        let mut all_columns = all_columns_set.into_iter().collect::<Vec<_>>();
4444        // sort to ensure the generated schema is not volatile
4445        all_columns.sort_unstable();
4446        // use left time index column name as the result time index column name
4447        all_columns.insert(0, left_time_index_column.clone());
4448
4449        // step 1: align schema using project, fill non-exist columns with null
4450        let left_proj_exprs = all_columns.iter().map(|col| {
4451            if tags_not_in_left.contains(col) {
4452                DfExpr::Literal(ScalarValue::Utf8(None), None).alias(col.clone())
4453            } else {
4454                DfExpr::Column(Column::new(None::<String>, col))
4455            }
4456        });
4457        let right_time_index_expr = DfExpr::Column(Column::new(
4458            right_qualifier.clone(),
4459            right_time_index_column,
4460        ))
4461        .alias(left_time_index_column.clone());
4462        // The field column in right side may not have qualifier (it may be removed by join operation),
4463        // so we need to find it from the schema.
4464        let right_qualifier_for_field = right
4465            .schema()
4466            .iter()
4467            .find(|(_, f)| f.name() == right_field_col)
4468            .map(|(q, _)| q)
4469            .with_context(|| ColumnNotFoundSnafu {
4470                col: right_field_col.clone(),
4471            })?
4472            .cloned();
4473
4474        // `skip(1)` to skip the time index column
4475        let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| {
4476            // expr
4477            if col == left_field_col && left_field_col != right_field_col {
4478                // qualify field in right side if necessary to handle different field name
4479                DfExpr::Column(Column::new(
4480                    right_qualifier_for_field.clone(),
4481                    right_field_col,
4482                ))
4483            } else if tags_not_in_right.contains(col) {
4484                DfExpr::Literal(ScalarValue::Utf8(None), None).alias(col.clone())
4485            } else {
4486                DfExpr::Column(Column::new(None::<String>, col))
4487            }
4488        });
4489        let right_proj_exprs = [right_time_index_expr]
4490            .into_iter()
4491            .chain(right_proj_exprs_without_time_index);
4492
4493        let left_projected = LogicalPlanBuilder::from(left)
4494            .project(left_proj_exprs)
4495            .context(DataFusionPlanningSnafu)?
4496            .alias(left_qualifier_string.clone())
4497            .context(DataFusionPlanningSnafu)?
4498            .build()
4499            .context(DataFusionPlanningSnafu)?;
4500        let right_projected = LogicalPlanBuilder::from(right)
4501            .project(right_proj_exprs)
4502            .context(DataFusionPlanningSnafu)?
4503            .alias(right_qualifier_string.clone())
4504            .context(DataFusionPlanningSnafu)?
4505            .build()
4506            .context(DataFusionPlanningSnafu)?;
4507
4508        // step 2: compute match columns
4509        let mut match_columns = if let Some(modifier) = modifier
4510            && let Some(matching) = &modifier.matching
4511        {
4512            match matching {
4513                // keeps columns mentioned in `on`
4514                LabelModifier::Include(on) => on.labels.clone(),
4515                // removes columns memtioned in `ignoring`
4516                LabelModifier::Exclude(ignoring) => {
4517                    let ignoring = ignoring.labels.iter().cloned().collect::<HashSet<_>>();
4518                    all_tags.difference(&ignoring).cloned().collect()
4519                }
4520            }
4521        } else {
4522            all_tags.iter().cloned().collect()
4523        };
4524        // sort to ensure the generated plan is not volatile
4525        match_columns.sort_unstable();
4526        // step 3: build `UnionDistinctOn` plan
4527        let schema = left_projected.schema().clone();
4528        let union_distinct_on = UnionDistinctOn::new(
4529            left_projected,
4530            right_projected,
4531            match_columns,
4532            left_time_index_column.clone(),
4533            schema,
4534        );
4535        let result = LogicalPlan::Extension(Extension {
4536            node: Arc::new(union_distinct_on),
4537        });
4538
4539        // step 4: update context
4540        self.ctx.time_index_column = Some(left_time_index_column);
4541        self.ctx.tag_columns = all_tags.into_iter().collect();
4542        self.ctx.field_columns = vec![left_field_col.clone()];
4543        self.ctx.use_tsid = left_has_tsid && right_has_tsid;
4544
4545        Ok(result)
4546    }
4547
4548    /// Build a projection that project and perform operation expr for every value columns.
4549    /// Non-value columns (tag and timestamp) will be preserved in the projection.
4550    ///
4551    /// # Side effect
4552    ///
4553    /// This function will update the value columns in the context. Those new column names
4554    /// don't contains qualifier.
4555    fn projection_for_each_field_column<F>(
4556        &mut self,
4557        input: LogicalPlan,
4558        name_to_expr: F,
4559    ) -> Result<LogicalPlan>
4560    where
4561        F: FnMut(&String) -> Result<DfExpr>,
4562    {
4563        let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
4564        let non_field_columns_iter = self
4565            .ctx
4566            .tag_columns
4567            .iter()
4568            .chain(self.ctx.time_index_column.iter())
4569            .map(|col| Ok(DfExpr::Column(Column::new(table_ref.clone(), col))));
4570        let tsid_iter =
4571            Self::optional_tsid_projection(input.schema(), table_ref.as_ref(), self.ctx.use_tsid)
4572                .into_iter()
4573                .map(Ok);
4574
4575        // build computation exprs
4576        let result_field_columns = self
4577            .ctx
4578            .field_columns
4579            .iter()
4580            .map(name_to_expr)
4581            .collect::<Result<Vec<_>>>()?;
4582
4583        // alias the computation exprs to remove qualifier
4584        self.ctx.field_columns = result_field_columns
4585            .iter()
4586            .map(|expr| expr.schema_name().to_string())
4587            .collect();
4588        let field_columns_iter = result_field_columns
4589            .into_iter()
4590            .zip(self.ctx.field_columns.iter())
4591            .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
4592
4593        // chain non-field columns (unchanged) and field columns (applied computation then alias)
4594        let project_fields = non_field_columns_iter
4595            .chain(tsid_iter)
4596            .chain(field_columns_iter)
4597            .collect::<Result<Vec<_>>>()?;
4598
4599        LogicalPlanBuilder::from(input)
4600            .project(project_fields)
4601            .context(DataFusionPlanningSnafu)?
4602            .build()
4603            .context(DataFusionPlanningSnafu)
4604    }
4605
4606    /// Build a filter plan that filter on value column. Notice that only one value column
4607    /// is expected.
4608    fn filter_on_field_column<F>(
4609        &self,
4610        input: LogicalPlan,
4611        mut name_to_expr: F,
4612    ) -> Result<LogicalPlan>
4613    where
4614        F: FnMut(&String) -> Result<DfExpr>,
4615    {
4616        ensure!(
4617            self.ctx.field_columns.len() == 1,
4618            UnsupportedExprSnafu {
4619                name: "filter on multi-value input"
4620            }
4621        );
4622
4623        let field_column_filter = name_to_expr(&self.ctx.field_columns[0])?;
4624
4625        LogicalPlanBuilder::from(input)
4626            .filter(field_column_filter)
4627            .context(DataFusionPlanningSnafu)?
4628            .build()
4629            .context(DataFusionPlanningSnafu)
4630    }
4631
4632    /// Generate an expr like `date_part("hour", <TIME_INDEX>)`. Caller should ensure the
4633    /// time index column in context is set
4634    fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
4635        let input_expr = datafusion::logical_expr::col(
4636            self.ctx
4637                .time_index_column
4638                .as_ref()
4639                // table name doesn't matters here
4640                .with_context(|| TimeIndexNotFoundSnafu {
4641                    table: "<doesn't matter>",
4642                })?,
4643        );
4644        let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4645            func: datafusion_functions::datetime::date_part(),
4646            args: vec![date_part.lit(), input_expr],
4647        });
4648        Ok(fn_expr)
4649    }
4650
4651    fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
4652        let schema = plan.schema();
4653        if !schema
4654            .fields()
4655            .iter()
4656            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
4657        {
4658            return Ok(plan);
4659        }
4660
4661        // Preserve column qualifiers so downstream plan nodes can keep referencing
4662        // the columns by their original qualified names.
4663        let project_exprs = schema
4664            .iter()
4665            .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
4666            .map(|(qualifier, field)| {
4667                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
4668            })
4669            .collect::<Vec<_>>();
4670
4671        LogicalPlanBuilder::from(plan)
4672            .project(project_exprs)
4673            .context(DataFusionPlanningSnafu)?
4674            .build()
4675            .context(DataFusionPlanningSnafu)
4676    }
4677
4678    /// Apply an alias to the query result by adding a projection with the alias name
4679    fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
4680        let fields_expr = self.create_field_column_exprs()?;
4681
4682        // TODO(dennis): how to support multi-value aliasing?
4683        ensure!(
4684            fields_expr.len() == 1,
4685            UnsupportedExprSnafu {
4686                name: "alias on multi-value result"
4687            }
4688        );
4689
4690        let project_fields = fields_expr
4691            .into_iter()
4692            .map(|expr| expr.alias(&alias_name))
4693            .chain(self.create_tag_column_exprs()?)
4694            .chain(Some(self.create_time_index_column_expr()?));
4695
4696        LogicalPlanBuilder::from(plan)
4697            .project(project_fields)
4698            .context(DataFusionPlanningSnafu)?
4699            .build()
4700            .context(DataFusionPlanningSnafu)
4701    }
4702}
4703
4704#[derive(Default, Debug)]
4705struct FunctionArgs {
4706    input: Option<PromExpr>,
4707    literals: Vec<DfExpr>,
4708}
4709
4710/// Represents different types of scalar functions supported in PromQL expressions.
4711/// Each variant defines how the function should be processed and what arguments it expects.
4712#[derive(Debug, Clone)]
4713enum ScalarFunc {
4714    /// DataFusion's registered(including built-in) scalar functions (e.g., abs, sqrt, round, clamp).
4715    /// These are passed through directly to DataFusion's execution engine.
4716    /// Processing: Simple argument insertion at the specified position.
4717    DataFusionBuiltin(Arc<ScalarUdfDef>),
4718    /// User-defined functions registered in DataFusion's function registry.
4719    /// Similar to DataFusionBuiltin but for custom functions not built into DataFusion.
4720    /// Processing: Direct pass-through with argument positioning.
4721    DataFusionUdf(Arc<ScalarUdfDef>),
4722    /// PromQL-specific functions that operate on time series data with temporal context.
4723    /// These functions require both timestamp ranges and values to perform calculations.
4724    /// Processing: Automatically injects timestamp_range and value columns as first arguments.
4725    /// Examples: idelta, irate, resets, changes, deriv, *_over_time function
4726    Udf(Arc<ScalarUdfDef>),
4727    /// PromQL functions requiring extrapolation calculations with explicit range information.
4728    /// These functions need to know the time range length to perform rate calculations.
4729    /// The second field contains the range length in milliseconds.
4730    /// Processing: Injects timestamp_range, value, time_index columns and appends range_length.
4731    /// Examples: increase, rate, delta
4732    // TODO(ruihang): maybe merge with Udf later
4733    ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
4734    /// Functions that generate expressions directly without external UDF calls.
4735    /// The expression is constructed during function matching and requires no additional processing.
4736    /// Examples: time(), minute(), hour(), month(), year() and other date/time extractors
4737    GeneratedExpr,
4738}
4739
4740#[cfg(test)]
4741mod test {
4742    use std::time::{Duration, UNIX_EPOCH};
4743
4744    use catalog::RegisterTableRequest;
4745    use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager};
4746    use common_base::Plugins;
4747    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
4748    use common_query::prelude::greptime_timestamp;
4749    use common_query::test_util::DummyDecoder;
4750    use datafusion::arrow::datatypes::Schema as ArrowSchema;
4751    use datafusion::datasource::memory::MemorySourceConfig;
4752    use datafusion::datasource::source::DataSourceExec;
4753    use datafusion::logical_expr::Extension;
4754    use datatypes::prelude::ConcreteDataType;
4755    use datatypes::schema::{ColumnSchema, Schema};
4756    use promql_parser::label::Labels;
4757    use promql_parser::parser;
4758    use session::context::QueryContext;
4759    use table::metadata::{TableInfoBuilder, TableMetaBuilder};
4760    use table::test_util::EmptyTable;
4761
4762    use super::*;
4763    use crate::QueryEngineContext;
4764    use crate::options::QueryOptions;
4765    use crate::parser::QueryLanguageParser;
4766
4767    fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> {
4768        if let LogicalPlan::Extension(Extension { node }) = plan
4769            && let Some(instant_manipulate) = node.as_any().downcast_ref::<InstantManipulate>()
4770        {
4771            return Some(instant_manipulate);
4772        }
4773
4774        plan.inputs().into_iter().find_map(find_instant_manipulate)
4775    }
4776
4777    fn build_query_engine_state() -> QueryEngineState {
4778        QueryEngineState::new(
4779            new_memory_catalog_manager().unwrap(),
4780            None,
4781            None,
4782            None,
4783            None,
4784            None,
4785            false,
4786            Plugins::default(),
4787            QueryOptions::default(),
4788        )
4789    }
4790
4791    async fn build_optimized_promql_plan(
4792        table_provider: DfTableSourceProvider,
4793        eval_stmt: &EvalStmt,
4794    ) -> LogicalPlan {
4795        let state = build_query_engine_state();
4796        let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state)
4797            .await
4798            .unwrap();
4799        let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
4800        state
4801            .optimize_by_extension_rules(raw_plan, &context)
4802            .unwrap()
4803    }
4804
4805    async fn build_optimized_tsid_plan(
4806        query: &str,
4807        num_tag: usize,
4808        num_field: usize,
4809        end_secs: u64,
4810        lookback_secs: u64,
4811    ) -> String {
4812        let eval_stmt = EvalStmt {
4813            expr: parser::parse(query).unwrap(),
4814            start: UNIX_EPOCH,
4815            end: UNIX_EPOCH
4816                .checked_add(Duration::from_secs(end_secs))
4817                .unwrap(),
4818            interval: Duration::from_secs(5),
4819            lookback_delta: Duration::from_secs(lookback_secs),
4820        };
4821        let table_provider = build_test_table_provider_with_tsid(
4822            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
4823            num_tag,
4824            num_field,
4825        )
4826        .await;
4827
4828        build_optimized_promql_plan(table_provider, &eval_stmt)
4829            .await
4830            .display_indent_schema()
4831            .to_string()
4832    }
4833
4834    async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) {
4835        let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
4836
4837        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
4838        assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0"));
4839        assert!(plan_str.contains("Distinct:"));
4840        assert!(plan_str.contains(expected_outer_agg), "{plan_str}");
4841        assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
4842    }
4843
4844    async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) {
4845        let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await;
4846        assert!(!plan_str.contains("Distinct:"), "{plan_str}");
4847    }
4848
4849    fn build_eval_stmt(expr: &str) -> EvalStmt {
4850        EvalStmt {
4851            expr: parser::parse(expr).unwrap(),
4852            start: UNIX_EPOCH,
4853            end: UNIX_EPOCH
4854                .checked_add(Duration::from_secs(100_000))
4855                .unwrap(),
4856            interval: Duration::from_secs(5),
4857            lookback_delta: Duration::from_secs(1),
4858        }
4859    }
4860
4861    async fn build_test_table_provider(
4862        table_name_tuples: &[(String, String)],
4863        num_tag: usize,
4864        num_field: usize,
4865    ) -> DfTableSourceProvider {
4866        let catalog_list = MemoryCatalogManager::with_default_setup();
4867        for (schema_name, table_name) in table_name_tuples {
4868            let mut columns = vec![];
4869            for i in 0..num_tag {
4870                columns.push(ColumnSchema::new(
4871                    format!("tag_{i}"),
4872                    ConcreteDataType::string_datatype(),
4873                    false,
4874                ));
4875            }
4876            columns.push(
4877                ColumnSchema::new(
4878                    "timestamp".to_string(),
4879                    ConcreteDataType::timestamp_millisecond_datatype(),
4880                    false,
4881                )
4882                .with_time_index(true),
4883            );
4884            for i in 0..num_field {
4885                columns.push(ColumnSchema::new(
4886                    format!("field_{i}"),
4887                    ConcreteDataType::float64_datatype(),
4888                    true,
4889                ));
4890            }
4891            let schema = Arc::new(Schema::new(columns));
4892            let table_meta = TableMetaBuilder::empty()
4893                .schema(schema)
4894                .primary_key_indices((0..num_tag).collect())
4895                .value_indices((num_tag + 1..num_tag + 1 + num_field).collect())
4896                .next_column_id(1024)
4897                .build()
4898                .unwrap();
4899            let table_info = TableInfoBuilder::default()
4900                .name(table_name.clone())
4901                .meta(table_meta)
4902                .build()
4903                .unwrap();
4904            let table = EmptyTable::from_table_info(&table_info);
4905
4906            assert!(
4907                catalog_list
4908                    .register_table_sync(RegisterTableRequest {
4909                        catalog: DEFAULT_CATALOG_NAME.to_string(),
4910                        schema: schema_name.clone(),
4911                        table_name: table_name.clone(),
4912                        table_id: 1024,
4913                        table,
4914                    })
4915                    .is_ok()
4916            );
4917        }
4918
4919        DfTableSourceProvider::new(
4920            catalog_list,
4921            false,
4922            QueryContext::arc(),
4923            DummyDecoder::arc(),
4924            false,
4925        )
4926    }
4927
4928    async fn build_test_table_provider_with_tsid(
4929        table_name_tuples: &[(String, String)],
4930        num_tag: usize,
4931        num_field: usize,
4932    ) -> DfTableSourceProvider {
4933        let table_specs = table_name_tuples
4934            .iter()
4935            .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field))
4936            .collect::<Vec<_>>();
4937        build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await
4938    }
4939
4940    async fn build_test_table_provider_with_tsid_fields(
4941        table_specs: &[((String, String), usize)],
4942        num_tag: usize,
4943    ) -> DfTableSourceProvider {
4944        let table_specs = table_specs
4945            .iter()
4946            .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field))
4947            .collect::<Vec<_>>();
4948        build_test_table_provider_with_tsid_tag_fields(&table_specs).await
4949    }
4950
4951    async fn build_test_table_provider_with_tsid_tag_fields(
4952        table_specs: &[((String, String), usize, usize)],
4953    ) -> DfTableSourceProvider {
4954        let catalog_list = MemoryCatalogManager::with_default_setup();
4955
4956        let physical_table_name = "phy";
4957        let physical_table_id = 999u32;
4958        let physical_num_tag = table_specs
4959            .iter()
4960            .map(|(_, num_tag, _)| *num_tag)
4961            .max()
4962            .unwrap_or(0);
4963        let physical_num_field = table_specs
4964            .iter()
4965            .map(|(_, _, num_field)| *num_field)
4966            .max()
4967            .unwrap_or(0);
4968
4969        // Register a metric engine physical table with internal columns.
4970        {
4971            let mut columns = vec![
4972                ColumnSchema::new(
4973                    DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
4974                    ConcreteDataType::uint32_datatype(),
4975                    false,
4976                ),
4977                ColumnSchema::new(
4978                    DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
4979                    ConcreteDataType::uint64_datatype(),
4980                    false,
4981                ),
4982            ];
4983            for i in 0..physical_num_tag {
4984                columns.push(ColumnSchema::new(
4985                    format!("tag_{i}"),
4986                    ConcreteDataType::string_datatype(),
4987                    false,
4988                ));
4989            }
4990            columns.push(
4991                ColumnSchema::new(
4992                    "timestamp".to_string(),
4993                    ConcreteDataType::timestamp_millisecond_datatype(),
4994                    false,
4995                )
4996                .with_time_index(true),
4997            );
4998            for i in 0..physical_num_field {
4999                columns.push(ColumnSchema::new(
5000                    format!("field_{i}"),
5001                    ConcreteDataType::float64_datatype(),
5002                    true,
5003                ));
5004            }
5005
5006            let schema = Arc::new(Schema::new(columns));
5007            let primary_key_indices = (0..(2 + physical_num_tag)).collect::<Vec<_>>();
5008            let table_meta = TableMetaBuilder::empty()
5009                .schema(schema)
5010                .primary_key_indices(primary_key_indices)
5011                .value_indices(
5012                    (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(),
5013                )
5014                .engine(METRIC_ENGINE_NAME.to_string())
5015                .next_column_id(1024)
5016                .build()
5017                .unwrap();
5018            let table_info = TableInfoBuilder::default()
5019                .table_id(physical_table_id)
5020                .name(physical_table_name)
5021                .meta(table_meta)
5022                .build()
5023                .unwrap();
5024            let table = EmptyTable::from_table_info(&table_info);
5025
5026            assert!(
5027                catalog_list
5028                    .register_table_sync(RegisterTableRequest {
5029                        catalog: DEFAULT_CATALOG_NAME.to_string(),
5030                        schema: DEFAULT_SCHEMA_NAME.to_string(),
5031                        table_name: physical_table_name.to_string(),
5032                        table_id: physical_table_id,
5033                        table,
5034                    })
5035                    .is_ok()
5036            );
5037        }
5038
5039        // Register metric engine logical tables without `__tsid`, referencing the physical table.
5040        for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate()
5041        {
5042            let mut columns = vec![];
5043            for i in 0..*num_tag {
5044                columns.push(ColumnSchema::new(
5045                    format!("tag_{i}"),
5046                    ConcreteDataType::string_datatype(),
5047                    false,
5048                ));
5049            }
5050            columns.push(
5051                ColumnSchema::new(
5052                    "timestamp".to_string(),
5053                    ConcreteDataType::timestamp_millisecond_datatype(),
5054                    false,
5055                )
5056                .with_time_index(true),
5057            );
5058            for i in 0..*num_field {
5059                columns.push(ColumnSchema::new(
5060                    format!("field_{i}"),
5061                    ConcreteDataType::float64_datatype(),
5062                    true,
5063                ));
5064            }
5065
5066            let schema = Arc::new(Schema::new(columns));
5067            let mut options = table::requests::TableOptions::default();
5068            options.extra_options.insert(
5069                LOGICAL_TABLE_METADATA_KEY.to_string(),
5070                physical_table_name.to_string(),
5071            );
5072            let table_id = 1024u32 + idx as u32;
5073            let table_meta = TableMetaBuilder::empty()
5074                .schema(schema)
5075                .primary_key_indices((0..*num_tag).collect())
5076                .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect())
5077                .engine(METRIC_ENGINE_NAME.to_string())
5078                .options(options)
5079                .next_column_id(1024)
5080                .build()
5081                .unwrap();
5082            let table_info = TableInfoBuilder::default()
5083                .table_id(table_id)
5084                .name(table_name.clone())
5085                .meta(table_meta)
5086                .build()
5087                .unwrap();
5088            let table = EmptyTable::from_table_info(&table_info);
5089
5090            assert!(
5091                catalog_list
5092                    .register_table_sync(RegisterTableRequest {
5093                        catalog: DEFAULT_CATALOG_NAME.to_string(),
5094                        schema: schema_name.clone(),
5095                        table_name: table_name.clone(),
5096                        table_id,
5097                        table,
5098                    })
5099                    .is_ok()
5100            );
5101        }
5102
5103        DfTableSourceProvider::new(
5104            catalog_list,
5105            false,
5106            QueryContext::arc(),
5107            DummyDecoder::arc(),
5108            false,
5109        )
5110    }
5111
5112    async fn build_test_table_provider_with_fields(
5113        table_name_tuples: &[(String, String)],
5114        tags: &[&str],
5115    ) -> DfTableSourceProvider {
5116        let catalog_list = MemoryCatalogManager::with_default_setup();
5117        for (schema_name, table_name) in table_name_tuples {
5118            let mut columns = vec![];
5119            let num_tag = tags.len();
5120            for tag in tags {
5121                columns.push(ColumnSchema::new(
5122                    tag.to_string(),
5123                    ConcreteDataType::string_datatype(),
5124                    false,
5125                ));
5126            }
5127            columns.push(
5128                ColumnSchema::new(
5129                    greptime_timestamp().to_string(),
5130                    ConcreteDataType::timestamp_millisecond_datatype(),
5131                    false,
5132                )
5133                .with_time_index(true),
5134            );
5135            columns.push(ColumnSchema::new(
5136                greptime_value().to_string(),
5137                ConcreteDataType::float64_datatype(),
5138                true,
5139            ));
5140            let schema = Arc::new(Schema::new(columns));
5141            let table_meta = TableMetaBuilder::empty()
5142                .schema(schema)
5143                .primary_key_indices((0..num_tag).collect())
5144                .next_column_id(1024)
5145                .build()
5146                .unwrap();
5147            let table_info = TableInfoBuilder::default()
5148                .name(table_name.clone())
5149                .meta(table_meta)
5150                .build()
5151                .unwrap();
5152            let table = EmptyTable::from_table_info(&table_info);
5153
5154            assert!(
5155                catalog_list
5156                    .register_table_sync(RegisterTableRequest {
5157                        catalog: DEFAULT_CATALOG_NAME.to_string(),
5158                        schema: schema_name.clone(),
5159                        table_name: table_name.clone(),
5160                        table_id: 1024,
5161                        table,
5162                    })
5163                    .is_ok()
5164            );
5165        }
5166
5167        DfTableSourceProvider::new(
5168            catalog_list,
5169            false,
5170            QueryContext::arc(),
5171            DummyDecoder::arc(),
5172            false,
5173        )
5174    }
5175
5176    // {
5177    //     input: `abs(some_metric{foo!="bar"})`,
5178    //     expected: &Call{
5179    //         Func: MustGetFunction("abs"),
5180    //         Args: Expressions{
5181    //             &VectorSelector{
5182    //                 Name: "some_metric",
5183    //                 LabelMatchers: []*labels.Matcher{
5184    //                     MustLabelMatcher(labels.MatchNotEqual, "foo", "bar"),
5185    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
5186    //                 },
5187    //             },
5188    //         },
5189    //     },
5190    // },
5191    async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) {
5192        let prom_expr =
5193            parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap();
5194        let eval_stmt = EvalStmt {
5195            expr: prom_expr,
5196            start: UNIX_EPOCH,
5197            end: UNIX_EPOCH
5198                .checked_add(Duration::from_secs(100_000))
5199                .unwrap(),
5200            interval: Duration::from_secs(5),
5201            lookback_delta: Duration::from_secs(1),
5202        };
5203
5204        let table_provider = build_test_table_provider(
5205            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5206            1,
5207            1,
5208        )
5209        .await;
5210        let plan =
5211            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5212                .await
5213                .unwrap();
5214
5215        let expected = String::from(
5216            "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\
5217            \n  Projection: some_metric.timestamp, TEMPLATE(some_metric.field_0) AS TEMPLATE(field_0), some_metric.tag_0 [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\
5218            \n    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
5219            \n      PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
5220            \n        Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
5221	            \n          Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
5222            \n            TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"
5223        ).replace("TEMPLATE", plan_name);
5224
5225        assert_eq!(plan.display_indent_schema().to_string(), expected);
5226    }
5227
5228    #[tokio::test]
5229    async fn single_abs() {
5230        do_single_instant_function_call("abs", "abs").await;
5231    }
5232
5233    #[tokio::test]
5234    #[should_panic]
5235    async fn single_absent() {
5236        do_single_instant_function_call("absent", "").await;
5237    }
5238
5239    #[tokio::test]
5240    async fn single_ceil() {
5241        do_single_instant_function_call("ceil", "ceil").await;
5242    }
5243
5244    #[tokio::test]
5245    async fn single_exp() {
5246        do_single_instant_function_call("exp", "exp").await;
5247    }
5248
5249    #[tokio::test]
5250    async fn single_ln() {
5251        do_single_instant_function_call("ln", "ln").await;
5252    }
5253
5254    #[tokio::test]
5255    async fn single_log2() {
5256        do_single_instant_function_call("log2", "log2").await;
5257    }
5258
5259    #[tokio::test]
5260    async fn single_log10() {
5261        do_single_instant_function_call("log10", "log10").await;
5262    }
5263
5264    #[tokio::test]
5265    #[should_panic]
5266    async fn single_scalar() {
5267        do_single_instant_function_call("scalar", "").await;
5268    }
5269
5270    #[tokio::test]
5271    #[should_panic]
5272    async fn single_sgn() {
5273        do_single_instant_function_call("sgn", "").await;
5274    }
5275
5276    #[tokio::test]
5277    #[should_panic]
5278    async fn single_sort() {
5279        do_single_instant_function_call("sort", "").await;
5280    }
5281
5282    #[tokio::test]
5283    #[should_panic]
5284    async fn single_sort_desc() {
5285        do_single_instant_function_call("sort_desc", "").await;
5286    }
5287
5288    #[tokio::test]
5289    async fn single_sqrt() {
5290        do_single_instant_function_call("sqrt", "sqrt").await;
5291    }
5292
5293    #[tokio::test]
5294    #[should_panic]
5295    async fn single_timestamp() {
5296        do_single_instant_function_call("timestamp", "").await;
5297    }
5298
5299    #[tokio::test]
5300    async fn single_acos() {
5301        do_single_instant_function_call("acos", "acos").await;
5302    }
5303
5304    #[tokio::test]
5305    #[should_panic]
5306    async fn single_acosh() {
5307        do_single_instant_function_call("acosh", "").await;
5308    }
5309
5310    #[tokio::test]
5311    async fn single_asin() {
5312        do_single_instant_function_call("asin", "asin").await;
5313    }
5314
5315    #[tokio::test]
5316    #[should_panic]
5317    async fn single_asinh() {
5318        do_single_instant_function_call("asinh", "").await;
5319    }
5320
5321    #[tokio::test]
5322    async fn single_atan() {
5323        do_single_instant_function_call("atan", "atan").await;
5324    }
5325
5326    #[tokio::test]
5327    #[should_panic]
5328    async fn single_atanh() {
5329        do_single_instant_function_call("atanh", "").await;
5330    }
5331
5332    #[tokio::test]
5333    async fn single_cos() {
5334        do_single_instant_function_call("cos", "cos").await;
5335    }
5336
5337    #[tokio::test]
5338    #[should_panic]
5339    async fn single_cosh() {
5340        do_single_instant_function_call("cosh", "").await;
5341    }
5342
5343    #[tokio::test]
5344    async fn single_sin() {
5345        do_single_instant_function_call("sin", "sin").await;
5346    }
5347
5348    #[tokio::test]
5349    #[should_panic]
5350    async fn single_sinh() {
5351        do_single_instant_function_call("sinh", "").await;
5352    }
5353
5354    #[tokio::test]
5355    async fn single_tan() {
5356        do_single_instant_function_call("tan", "tan").await;
5357    }
5358
5359    #[tokio::test]
5360    #[should_panic]
5361    async fn single_tanh() {
5362        do_single_instant_function_call("tanh", "").await;
5363    }
5364
5365    #[tokio::test]
5366    #[should_panic]
5367    async fn single_deg() {
5368        do_single_instant_function_call("deg", "").await;
5369    }
5370
5371    #[tokio::test]
5372    #[should_panic]
5373    async fn single_rad() {
5374        do_single_instant_function_call("rad", "").await;
5375    }
5376
5377    // {
5378    //     input: "avg by (foo)(some_metric)",
5379    //     expected: &AggregateExpr{
5380    //         Op: AVG,
5381    //         Expr: &VectorSelector{
5382    //             Name: "some_metric",
5383    //             LabelMatchers: []*labels.Matcher{
5384    //                 MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
5385    //             },
5386    //             PosRange: PositionRange{
5387    //                 Start: 13,
5388    //                 End:   24,
5389    //             },
5390    //         },
5391    //         Grouping: []string{"foo"},
5392    //         PosRange: PositionRange{
5393    //             Start: 0,
5394    //             End:   25,
5395    //         },
5396    //     },
5397    // },
5398    async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) {
5399        let prom_expr = parser::parse(&format!(
5400            "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})",
5401        ))
5402        .unwrap();
5403        let mut eval_stmt = EvalStmt {
5404            expr: prom_expr,
5405            start: UNIX_EPOCH,
5406            end: UNIX_EPOCH
5407                .checked_add(Duration::from_secs(100_000))
5408                .unwrap(),
5409            interval: Duration::from_secs(5),
5410            lookback_delta: Duration::from_secs(1),
5411        };
5412
5413        // test group by
5414        let table_provider = build_test_table_provider(
5415            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5416            2,
5417            2,
5418        )
5419        .await;
5420        let plan =
5421            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5422                .await
5423                .unwrap();
5424        let expected_no_without = String::from(
5425            "Sort: some_metric.tag_1 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\
5426            \n  Aggregate: groupBy=[[some_metric.tag_1, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\
5427            \n    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
5428            \n      PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
5429            \n        Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
5430            \n          Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
5431            \n            TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
5432        ).replace("TEMPLATE", plan_name);
5433        assert_eq!(
5434            plan.display_indent_schema().to_string(),
5435            expected_no_without
5436        );
5437
5438        // test group without
5439        if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr {
5440            *modifier = Some(LabelModifier::Exclude(Labels {
5441                labels: vec![String::from("tag_1")].into_iter().collect(),
5442            }));
5443        }
5444        let table_provider = build_test_table_provider(
5445            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5446            2,
5447            2,
5448        )
5449        .await;
5450        let plan =
5451            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5452                .await
5453                .unwrap();
5454        let expected_without = String::from(
5455            "Sort: some_metric.tag_0 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\
5456            \n  Aggregate: groupBy=[[some_metric.tag_0, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\
5457            \n    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
5458            \n      PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
5459            \n        Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
5460            \n          Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
5461            \n            TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
5462        ).replace("TEMPLATE", plan_name);
5463        assert_eq!(plan.display_indent_schema().to_string(), expected_without);
5464    }
5465
5466    #[tokio::test]
5467    async fn aggregate_sum() {
5468        do_aggregate_expr_plan("sum", "sum").await;
5469    }
5470
5471    #[tokio::test]
5472    async fn tsid_is_used_for_series_divide_when_available() {
5473        let prom_expr = parser::parse("some_metric").unwrap();
5474        let eval_stmt = EvalStmt {
5475            expr: prom_expr,
5476            start: UNIX_EPOCH,
5477            end: UNIX_EPOCH
5478                .checked_add(Duration::from_secs(100_000))
5479                .unwrap(),
5480            interval: Duration::from_secs(5),
5481            lookback_delta: Duration::from_secs(1),
5482        };
5483
5484        let table_provider = build_test_table_provider_with_tsid(
5485            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5486            1,
5487            1,
5488        )
5489        .await;
5490        let plan =
5491            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5492                .await
5493                .unwrap();
5494
5495        let plan_str = plan.display_indent_schema().to_string();
5496        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
5497        assert!(plan_str.contains("__tsid ASC NULLS FIRST"));
5498        assert!(
5499            !plan
5500                .schema()
5501                .fields()
5502                .iter()
5503                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5504        );
5505
5506        let manipulate = find_instant_manipulate(&plan).unwrap();
5507        let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new(
5508            MemorySourceConfig::try_new(&[], Arc::new(ArrowSchema::empty()), None).unwrap(),
5509        ))));
5510        assert!(format!("{exec:?}").contains("reuse_tsid_column: true"));
5511    }
5512
5513    #[tokio::test]
5514    async fn default_binary_join_uses_tsid_when_available() {
5515        let eval_stmt = build_eval_stmt("some_metric / some_alt_metric");
5516
5517        let table_provider = build_test_table_provider_with_tsid(
5518            &[
5519                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5520                (
5521                    DEFAULT_SCHEMA_NAME.to_string(),
5522                    "some_alt_metric".to_string(),
5523                ),
5524            ],
5525            1,
5526            1,
5527        )
5528        .await;
5529        let plan =
5530            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5531                .await
5532                .unwrap();
5533
5534        let plan_str = plan.display_indent_schema().to_string();
5535        assert!(
5536            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
5537            "{plan_str}"
5538        );
5539        assert!(
5540            !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
5541            "{plan_str}"
5542        );
5543    }
5544
5545    #[tokio::test]
5546    async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() {
5547        for query in [
5548            "timestamp(some_metric) / some_metric",
5549            "some_metric / timestamp(some_metric)",
5550        ] {
5551            let eval_stmt = build_eval_stmt(query);
5552
5553            let table_provider = build_test_table_provider_with_tsid(
5554                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5555                1,
5556                1,
5557            )
5558            .await;
5559            let plan =
5560                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5561                    .await
5562                    .unwrap();
5563
5564            let plan_str = plan.display_indent_schema().to_string();
5565            assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}");
5566            assert!(
5567                plan_str.contains("lhs.tag_0 = rhs.tag_0"),
5568                "{query}: {plan_str}"
5569            );
5570            assert!(
5571                !plan
5572                    .schema()
5573                    .fields()
5574                    .iter()
5575                    .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
5576                "{query}: {plan_str}"
5577            );
5578        }
5579    }
5580
5581    #[tokio::test]
5582    async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() {
5583        let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job");
5584
5585        let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
5586            (
5587                (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()),
5588                2,
5589                1,
5590            ),
5591            (
5592                (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()),
5593                1,
5594                1,
5595            ),
5596        ])
5597        .await;
5598        let plan =
5599            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5600                .await
5601                .unwrap();
5602        let plan_str = plan.display_indent_schema().to_string();
5603
5604        assert!(
5605            plan_str.contains("Boolean(false)") || plan_str.contains("false"),
5606            "{plan_str}"
5607        );
5608    }
5609
5610    #[tokio::test]
5611    async fn tsid_is_preserved_for_nested_default_binary_joins() {
5612        let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric");
5613
5614        let table_provider = build_test_table_provider_with_tsid(
5615            &[
5616                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5617                (
5618                    DEFAULT_SCHEMA_NAME.to_string(),
5619                    "some_alt_metric".to_string(),
5620                ),
5621                (
5622                    DEFAULT_SCHEMA_NAME.to_string(),
5623                    "some_third_metric".to_string(),
5624                ),
5625            ],
5626            1,
5627            1,
5628        )
5629        .await;
5630        let plan =
5631            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5632                .await
5633                .unwrap();
5634
5635        let plan_str = plan.display_indent_schema().to_string();
5636        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
5637        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
5638    }
5639
5640    #[tokio::test]
5641    async fn repeated_tsid_binary_operand_reuses_leaf_plan() {
5642        let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100");
5643
5644        let table_provider = build_test_table_provider_with_tsid(
5645            &[
5646                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5647                (
5648                    DEFAULT_SCHEMA_NAME.to_string(),
5649                    "some_alt_metric".to_string(),
5650                ),
5651            ],
5652            1,
5653            1,
5654        )
5655        .await;
5656        let plan =
5657            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5658                .await
5659                .unwrap();
5660
5661        let plan_str = plan.display_indent_schema().to_string();
5662        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
5663        assert_eq!(
5664            plan_str
5665                .matches("Filter: phy.__table_id = UInt32(1024)")
5666                .count(),
5667            1,
5668            "{plan_str}"
5669        );
5670        assert_eq!(
5671            plan_str.matches("PromInstantManipulate").count(),
5672            2,
5673            "{plan_str}"
5674        );
5675        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
5676    }
5677
5678    #[tokio::test]
5679    async fn repeated_tsid_binary_operand_reuses_shorter_field_side() {
5680        let eval_stmt =
5681            build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100");
5682
5683        let table_provider = build_test_table_provider_with_tsid_fields(
5684            &[
5685                (
5686                    (
5687                        DEFAULT_SCHEMA_NAME.to_string(),
5688                        "two_field_metric".to_string(),
5689                    ),
5690                    2,
5691                ),
5692                (
5693                    (
5694                        DEFAULT_SCHEMA_NAME.to_string(),
5695                        "one_field_metric".to_string(),
5696                    ),
5697                    1,
5698                ),
5699            ],
5700            1,
5701        )
5702        .await;
5703        let plan =
5704            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5705                .await
5706                .unwrap();
5707
5708        let field_names = plan
5709            .schema()
5710            .fields()
5711            .iter()
5712            .map(|field| field.name().clone())
5713            .collect::<Vec<_>>();
5714        let value_columns = field_names
5715            .iter()
5716            .filter(|name| {
5717                *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
5718            })
5719            .count();
5720        assert_eq!(value_columns, 1, "{field_names:?}");
5721        let plan_str = plan.display_indent_schema().to_string();
5722        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
5723        assert_eq!(
5724            plan_str
5725                .matches("Filter: phy.__table_id = UInt32(1025)")
5726                .count(),
5727            1,
5728            "{plan_str}"
5729        );
5730        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
5731    }
5732
5733    #[tokio::test]
5734    async fn binary_island_reuses_self_operand_without_join() {
5735        let eval_stmt = build_eval_stmt("some_metric / some_metric");
5736
5737        let table_provider = build_test_table_provider_with_tsid(
5738            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
5739            1,
5740            1,
5741        )
5742        .await;
5743        let plan =
5744            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5745                .await
5746                .unwrap();
5747
5748        let plan_str = plan.display_indent_schema().to_string();
5749        assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}");
5750        assert_eq!(
5751            plan_str
5752                .matches("Filter: phy.__table_id = UInt32(1024)")
5753                .count(),
5754            1,
5755            "{plan_str}"
5756        );
5757        assert_eq!(
5758            plan_str.matches("PromInstantManipulate").count(),
5759            1,
5760            "{plan_str}"
5761        );
5762    }
5763
5764    #[tokio::test]
5765    async fn binary_island_reuses_leaf_across_two_branches() {
5766        let eval_stmt =
5767            build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)");
5768
5769        let table_provider = build_test_table_provider_with_tsid(
5770            &[
5771                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5772                (
5773                    DEFAULT_SCHEMA_NAME.to_string(),
5774                    "some_alt_metric".to_string(),
5775                ),
5776                (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()),
5777            ],
5778            1,
5779            1,
5780        )
5781        .await;
5782        let plan =
5783            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5784                .await
5785                .unwrap();
5786
5787        let plan_str = plan.display_indent_schema().to_string();
5788        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
5789        assert_eq!(
5790            plan_str
5791                .matches("Filter: phy.__table_id = UInt32(1024)")
5792                .count(),
5793            1,
5794            "{plan_str}"
5795        );
5796        assert_eq!(
5797            plan_str.matches("PromInstantManipulate").count(),
5798            3,
5799            "{plan_str}"
5800        );
5801    }
5802
5803    #[tokio::test]
5804    async fn binary_island_generated_alias_avoids_user_column_names() {
5805        let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric");
5806
5807        let table_provider = build_test_table_provider_with_fields(
5808            &[
5809                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5810                (
5811                    DEFAULT_SCHEMA_NAME.to_string(),
5812                    "some_alt_metric".to_string(),
5813                ),
5814            ],
5815            &["prom_v0", "__prom_v0"],
5816        )
5817        .await;
5818        let plan =
5819            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5820                .await
5821                .unwrap();
5822
5823        let field_names = plan.schema().field_names();
5824        assert!(field_names.iter().any(|name| name.ends_with(".prom_v0")));
5825        assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0")));
5826
5827        let plan_str = plan.display_indent_schema().to_string();
5828        assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}");
5829        assert_eq!(
5830            plan_str.matches("PromInstantManipulate").count(),
5831            2,
5832            "{plan_str}"
5833        );
5834    }
5835
5836    #[tokio::test]
5837    async fn binary_island_clears_qualifier_for_nested_unary_projection() {
5838        let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)");
5839
5840        let table_provider = build_test_table_provider_with_tsid(
5841            &[
5842                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5843                (
5844                    DEFAULT_SCHEMA_NAME.to_string(),
5845                    "some_alt_metric".to_string(),
5846                ),
5847            ],
5848            1,
5849            1,
5850        )
5851        .await;
5852        let plan =
5853            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5854                .await
5855                .unwrap();
5856
5857        let plan_str = plan.display_indent_schema().to_string();
5858        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
5859        assert_eq!(
5860            plan_str.matches("PromInstantManipulate").count(),
5861            2,
5862            "{plan_str}"
5863        );
5864    }
5865
5866    #[tokio::test]
5867    async fn binary_island_keeps_distinct_matcher_leaves() {
5868        let eval_stmt = build_eval_stmt(
5869            "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}",
5870        );
5871
5872        let table_provider = build_test_table_provider_with_tsid(
5873            &[
5874                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5875                (
5876                    DEFAULT_SCHEMA_NAME.to_string(),
5877                    "some_alt_metric".to_string(),
5878                ),
5879            ],
5880            1,
5881            1,
5882        )
5883        .await;
5884        let plan =
5885            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5886                .await
5887                .unwrap();
5888
5889        let plan_str = plan.display_indent_schema().to_string();
5890        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
5891        assert_eq!(
5892            plan_str.matches("PromInstantManipulate").count(),
5893            3,
5894            "{plan_str}"
5895        );
5896    }
5897
5898    #[tokio::test]
5899    async fn binary_island_keeps_offset_leaves_distinct() {
5900        let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric");
5901
5902        let table_provider = build_test_table_provider_with_tsid(
5903            &[
5904                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5905                (
5906                    DEFAULT_SCHEMA_NAME.to_string(),
5907                    "some_alt_metric".to_string(),
5908                ),
5909            ],
5910            1,
5911            1,
5912        )
5913        .await;
5914        let plan =
5915            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5916                .await
5917                .unwrap();
5918
5919        let plan_str = plan.display_indent_schema().to_string();
5920        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
5921        assert_eq!(
5922            plan_str.matches("PromInstantManipulate").count(),
5923            3,
5924            "{plan_str}"
5925        );
5926    }
5927
5928    #[tokio::test]
5929    async fn binary_island_falls_back_for_group_modifier() {
5930        let eval_stmt = build_eval_stmt(
5931            "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric",
5932        );
5933
5934        let table_provider = build_test_table_provider_with_tsid(
5935            &[
5936                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5937                (
5938                    DEFAULT_SCHEMA_NAME.to_string(),
5939                    "some_alt_metric".to_string(),
5940                ),
5941            ],
5942            1,
5943            1,
5944        )
5945        .await;
5946        let plan =
5947            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5948                .await
5949                .unwrap();
5950
5951        let plan_str = plan.display_indent_schema().to_string();
5952        assert_eq!(
5953            plan_str.matches("PromInstantManipulate").count(),
5954            3,
5955            "{plan_str}"
5956        );
5957    }
5958
5959    #[tokio::test]
5960    async fn binary_island_falls_back_for_comparison_filter() {
5961        let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric");
5962
5963        let table_provider = build_test_table_provider_with_tsid(
5964            &[
5965                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
5966                (
5967                    DEFAULT_SCHEMA_NAME.to_string(),
5968                    "some_alt_metric".to_string(),
5969                ),
5970            ],
5971            1,
5972            1,
5973        )
5974        .await;
5975        let plan =
5976            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
5977                .await
5978                .unwrap();
5979
5980        let plan_str = plan.display_indent_schema().to_string();
5981        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
5982        assert_eq!(
5983            plan_str.matches("PromInstantManipulate").count(),
5984            3,
5985            "{plan_str}"
5986        );
5987    }
5988
5989    #[tokio::test]
5990    async fn tsid_binary_join_uses_shorter_field_side() {
5991        let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric");
5992
5993        let table_provider = build_test_table_provider_with_tsid_fields(
5994            &[
5995                (
5996                    (
5997                        DEFAULT_SCHEMA_NAME.to_string(),
5998                        "one_field_metric".to_string(),
5999                    ),
6000                    1,
6001                ),
6002                (
6003                    (
6004                        DEFAULT_SCHEMA_NAME.to_string(),
6005                        "two_field_metric".to_string(),
6006                    ),
6007                    2,
6008                ),
6009            ],
6010            1,
6011        )
6012        .await;
6013        let plan =
6014            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6015                .await
6016                .unwrap();
6017
6018        let field_names = plan
6019            .schema()
6020            .fields()
6021            .iter()
6022            .map(|field| field.name().clone())
6023            .collect::<Vec<_>>();
6024        let value_columns = field_names
6025            .iter()
6026            .filter(|name| {
6027                *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
6028            })
6029            .count();
6030        assert_eq!(value_columns, 1, "{field_names:?}");
6031    }
6032
6033    #[tokio::test]
6034    async fn comparison_binary_join_uses_shorter_field_side() {
6035        let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric");
6036
6037        let table_provider = build_test_table_provider_with_tsid_fields(
6038            &[
6039                (
6040                    (
6041                        DEFAULT_SCHEMA_NAME.to_string(),
6042                        "two_field_metric".to_string(),
6043                    ),
6044                    2,
6045                ),
6046                (
6047                    (
6048                        DEFAULT_SCHEMA_NAME.to_string(),
6049                        "one_field_metric".to_string(),
6050                    ),
6051                    1,
6052                ),
6053            ],
6054            1,
6055        )
6056        .await;
6057        let plan =
6058            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6059                .await
6060                .unwrap();
6061
6062        let field_names = plan
6063            .schema()
6064            .fields()
6065            .iter()
6066            .map(|field| field.name().clone())
6067            .collect::<Vec<_>>();
6068        assert!(
6069            field_names.iter().any(|name| name == "field_0"),
6070            "{field_names:?}"
6071        );
6072        assert!(
6073            !field_names.iter().any(|name| name == "field_1"),
6074            "{field_names:?}"
6075        );
6076    }
6077
6078    #[tokio::test]
6079    async fn label_matching_modifier_disables_tsid_binary_join() {
6080        let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric");
6081
6082        let table_provider = build_test_table_provider_with_tsid(
6083            &[
6084                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6085                (
6086                    DEFAULT_SCHEMA_NAME.to_string(),
6087                    "some_alt_metric".to_string(),
6088                ),
6089            ],
6090            2,
6091            1,
6092        )
6093        .await;
6094        let plan =
6095            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6096                .await
6097                .unwrap();
6098
6099        let plan_str = plan.display_indent_schema().to_string();
6100        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
6101        assert!(
6102            plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"),
6103            "{plan_str}"
6104        );
6105    }
6106
6107    #[tokio::test]
6108    async fn ignoring_absent_label_keeps_tsid_binary_join() {
6109        let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric");
6110
6111        let table_provider = build_test_table_provider_with_tsid(
6112            &[
6113                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6114                (
6115                    DEFAULT_SCHEMA_NAME.to_string(),
6116                    "some_alt_metric".to_string(),
6117                ),
6118            ],
6119            2,
6120            1,
6121        )
6122        .await;
6123        let plan =
6124            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6125                .await
6126                .unwrap();
6127
6128        let plan_str = plan.display_indent_schema().to_string();
6129        assert!(
6130            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
6131            "{plan_str}"
6132        );
6133        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
6134        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6135    }
6136
6137    #[tokio::test]
6138    async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() {
6139        let eval_stmt =
6140            build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric");
6141
6142        let table_provider = build_test_table_provider_with_tsid(
6143            &[
6144                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6145                (
6146                    DEFAULT_SCHEMA_NAME.to_string(),
6147                    "some_alt_metric".to_string(),
6148                ),
6149            ],
6150            2,
6151            1,
6152        )
6153        .await;
6154        let plan =
6155            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6156                .await
6157                .unwrap();
6158
6159        let plan_str = plan.display_indent_schema().to_string();
6160        assert!(
6161            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
6162            "{plan_str}"
6163        );
6164        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
6165        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6166    }
6167
6168    #[tokio::test]
6169    async fn on_full_label_set_keeps_tsid_binary_join() {
6170        let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric");
6171
6172        let table_provider = build_test_table_provider_with_tsid(
6173            &[
6174                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6175                (
6176                    DEFAULT_SCHEMA_NAME.to_string(),
6177                    "some_alt_metric".to_string(),
6178                ),
6179            ],
6180            2,
6181            1,
6182        )
6183        .await;
6184        let plan =
6185            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6186                .await
6187                .unwrap();
6188
6189        let plan_str = plan.display_indent_schema().to_string();
6190        assert!(
6191            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
6192            "{plan_str}"
6193        );
6194        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
6195        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6196    }
6197
6198    #[tokio::test]
6199    async fn on_partial_label_set_disables_tsid_binary_join() {
6200        let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
6201
6202        let table_provider = build_test_table_provider_with_tsid(
6203            &[
6204                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6205                (
6206                    DEFAULT_SCHEMA_NAME.to_string(),
6207                    "some_alt_metric".to_string(),
6208                ),
6209            ],
6210            2,
6211            1,
6212        )
6213        .await;
6214        let plan =
6215            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6216                .await
6217                .unwrap();
6218
6219        let plan_str = plan.display_indent_schema().to_string();
6220        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
6221        assert!(
6222            plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
6223            "{plan_str}"
6224        );
6225        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6226    }
6227
6228    #[tokio::test]
6229    async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() {
6230        let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
6231
6232        let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
6233            (
6234                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6235                2,
6236                1,
6237            ),
6238            (
6239                (
6240                    DEFAULT_SCHEMA_NAME.to_string(),
6241                    "some_alt_metric".to_string(),
6242                ),
6243                1,
6244                1,
6245            ),
6246        ])
6247        .await;
6248        let plan =
6249            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6250                .await
6251                .unwrap();
6252
6253        let plan_str = plan.display_indent_schema().to_string();
6254        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
6255        assert!(
6256            plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
6257            "{plan_str}"
6258        );
6259        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6260    }
6261
6262    #[tokio::test]
6263    async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() {
6264        let eval_stmt = build_eval_stmt("some_metric > some_alt_metric");
6265
6266        let table_provider = build_test_table_provider_with_tsid(
6267            &[
6268                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6269                (
6270                    DEFAULT_SCHEMA_NAME.to_string(),
6271                    "some_alt_metric".to_string(),
6272                ),
6273            ],
6274            2,
6275            1,
6276        )
6277        .await;
6278        let mut planner = PromPlanner {
6279            table_provider,
6280            ctx: PromPlannerContext::from_eval_stmt(&eval_stmt),
6281        };
6282        let plan = planner
6283            .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state())
6284            .await
6285            .unwrap();
6286
6287        let plan_str = plan.display_indent_schema().to_string();
6288        assert!(
6289            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
6290            "{plan_str}"
6291        );
6292        assert!(
6293            plan.schema()
6294                .fields()
6295                .iter()
6296                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
6297            "{plan_str}"
6298        );
6299        assert!(planner.ctx.use_tsid, "{plan_str}");
6300    }
6301
6302    #[tokio::test]
6303    async fn comparison_bool_binary_join_uses_tsid_when_available() {
6304        let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric");
6305
6306        let table_provider = build_test_table_provider_with_tsid(
6307            &[
6308                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
6309                (
6310                    DEFAULT_SCHEMA_NAME.to_string(),
6311                    "some_alt_metric".to_string(),
6312                ),
6313            ],
6314            2,
6315            1,
6316        )
6317        .await;
6318        let plan =
6319            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6320                .await
6321                .unwrap();
6322
6323        let plan_str = plan.display_indent_schema().to_string();
6324        assert!(
6325            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
6326            "{plan_str}"
6327        );
6328        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
6329        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
6330    }
6331
6332    #[tokio::test]
6333    async fn scalar_count_count_range_keeps_full_window() {
6334        let plan_str = build_optimized_tsid_plan(
6335            "scalar(count(count(some_metric) by (tag_0)))",
6336            1,
6337            1,
6338            100_000,
6339            1,
6340        )
6341        .await;
6342        assert!(plan_str.contains("ScalarCalculate: tags=[]"));
6343        assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]"));
6344        assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]"));
6345    }
6346
6347    #[tokio::test]
6348    async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() {
6349        let plan_str = build_optimized_tsid_plan(
6350            "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))",
6351            2,
6352            1,
6353            10,
6354            300,
6355        )
6356        .await;
6357        assert!(plan_str.contains("Distinct:"), "{plan_str}");
6358    }
6359
6360    #[tokio::test]
6361    async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() {
6362        assert_nested_count_rewrite_applies(
6363            "count(count(some_metric) by (tag_0))",
6364            "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]"
6365        )
6366        .await;
6367    }
6368
6369    #[tokio::test]
6370    async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() {
6371        assert_nested_count_rewrite_applies(
6372            "count(sum(some_metric) by (tag_0))",
6373            "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]"
6374        )
6375        .await;
6376    }
6377
6378    #[tokio::test]
6379    async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() {
6380        for (query, expected_outer_agg) in [
6381            (
6382                "count(avg(some_metric) by (tag_0))",
6383                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]",
6384            ),
6385            (
6386                "count(min(some_metric) by (tag_0))",
6387                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]",
6388            ),
6389            (
6390                "count(max(some_metric) by (tag_0))",
6391                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]",
6392            ),
6393            (
6394                "count(stddev(some_metric) by (tag_0))",
6395                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]",
6396            ),
6397            (
6398                "count(stdvar(some_metric) by (tag_0))",
6399                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]",
6400            ),
6401        ] {
6402            assert_nested_count_rewrite_applies(query, expected_outer_agg).await;
6403        }
6404    }
6405
6406    #[tokio::test]
6407    async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() {
6408        let count_plan =
6409            build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1)
6410                .await;
6411        assert!(
6412            !count_plan.contains("some_metric.field_0 IS NOT NULL"),
6413            "{count_plan}"
6414        );
6415
6416        for query in [
6417            "count(sum(some_metric) by (tag_0))",
6418            "count(avg(some_metric) by (tag_0))",
6419            "count(min(some_metric) by (tag_0))",
6420            "count(max(some_metric) by (tag_0))",
6421            "count(stddev(some_metric) by (tag_0))",
6422            "count(stdvar(some_metric) by (tag_0))",
6423        ] {
6424            let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
6425            assert!(
6426                plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"),
6427                "{query}: {plan_str}"
6428            );
6429        }
6430    }
6431
6432    #[tokio::test]
6433    async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() {
6434        assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await;
6435        assert_nested_count_rewrite_missing(
6436            "count(sum(irate(some_metric[1h])) by (tag_0))",
6437            2,
6438            300,
6439        )
6440        .await;
6441    }
6442
6443    #[tokio::test]
6444    async fn physical_table_name_is_not_leaked_in_plan() {
6445        let prom_expr = parser::parse("some_metric").unwrap();
6446        let eval_stmt = EvalStmt {
6447            expr: prom_expr,
6448            start: UNIX_EPOCH,
6449            end: UNIX_EPOCH
6450                .checked_add(Duration::from_secs(100_000))
6451                .unwrap(),
6452            interval: Duration::from_secs(5),
6453            lookback_delta: Duration::from_secs(1),
6454        };
6455
6456        let table_provider = build_test_table_provider_with_tsid(
6457            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6458            1,
6459            1,
6460        )
6461        .await;
6462        let plan =
6463            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6464                .await
6465                .unwrap();
6466
6467        let plan_str = plan.display_indent_schema().to_string();
6468        assert!(plan_str.contains("TableScan: phy"), "{plan}");
6469        assert!(plan_str.contains("SubqueryAlias: some_metric"));
6470        assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)"));
6471        assert!(!plan_str.contains("TableScan: some_metric"));
6472    }
6473
6474    #[tokio::test]
6475    async fn sum_without_does_not_group_by_tsid() {
6476        let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap();
6477        let eval_stmt = EvalStmt {
6478            expr: prom_expr,
6479            start: UNIX_EPOCH,
6480            end: UNIX_EPOCH
6481                .checked_add(Duration::from_secs(100_000))
6482                .unwrap(),
6483            interval: Duration::from_secs(5),
6484            lookback_delta: Duration::from_secs(1),
6485        };
6486
6487        let table_provider = build_test_table_provider_with_tsid(
6488            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6489            1,
6490            1,
6491        )
6492        .await;
6493        let plan =
6494            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6495                .await
6496                .unwrap();
6497
6498        let plan_str = plan.display_indent_schema().to_string();
6499        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
6500
6501        let aggr_line = plan_str
6502            .lines()
6503            .find(|line| line.contains("Aggregate: groupBy="))
6504            .unwrap();
6505        assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
6506    }
6507
6508    #[tokio::test]
6509    async fn topk_without_does_not_partition_by_tsid() {
6510        let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap();
6511        let eval_stmt = EvalStmt {
6512            expr: prom_expr,
6513            start: UNIX_EPOCH,
6514            end: UNIX_EPOCH
6515                .checked_add(Duration::from_secs(100_000))
6516                .unwrap(),
6517            interval: Duration::from_secs(5),
6518            lookback_delta: Duration::from_secs(1),
6519        };
6520
6521        let table_provider = build_test_table_provider_with_tsid(
6522            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6523            1,
6524            1,
6525        )
6526        .await;
6527        let plan =
6528            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6529                .await
6530                .unwrap();
6531
6532        let plan_str = plan.display_indent_schema().to_string();
6533        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
6534
6535        let window_line = plan_str
6536            .lines()
6537            .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
6538            .unwrap();
6539        let partition_by = window_line
6540            .split("PARTITION BY [")
6541            .nth(1)
6542            .and_then(|s| s.split("] ORDER BY").next())
6543            .unwrap();
6544        assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
6545    }
6546
6547    #[tokio::test]
6548    async fn sum_by_does_not_group_by_tsid() {
6549        let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap();
6550        let eval_stmt = EvalStmt {
6551            expr: prom_expr,
6552            start: UNIX_EPOCH,
6553            end: UNIX_EPOCH
6554                .checked_add(Duration::from_secs(100_000))
6555                .unwrap(),
6556            interval: Duration::from_secs(5),
6557            lookback_delta: Duration::from_secs(1),
6558        };
6559
6560        let table_provider = build_test_table_provider_with_tsid(
6561            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6562            1,
6563            1,
6564        )
6565        .await;
6566        let plan =
6567            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6568                .await
6569                .unwrap();
6570
6571        let plan_str = plan.display_indent_schema().to_string();
6572        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
6573
6574        let aggr_line = plan_str
6575            .lines()
6576            .find(|line| line.contains("Aggregate: groupBy="))
6577            .unwrap();
6578        assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
6579    }
6580
6581    #[tokio::test]
6582    async fn aggregate_over_binary_time_function_expr() {
6583        for op in ["sum", "min", "max", "avg"] {
6584            let prom_expr = parser::parse(&format!(
6585                "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)"
6586            ))
6587            .unwrap();
6588            let eval_stmt = EvalStmt {
6589                expr: prom_expr,
6590                start: UNIX_EPOCH,
6591                end: UNIX_EPOCH
6592                    .checked_add(Duration::from_secs(100_000))
6593                    .unwrap(),
6594                interval: Duration::from_secs(5),
6595                lookback_delta: Duration::from_secs(1),
6596            };
6597
6598            let table_provider = build_test_table_provider_with_tsid(
6599                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6600                3,
6601                1,
6602            )
6603            .await;
6604            let plan =
6605                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6606                    .await
6607                    .unwrap();
6608
6609            let plan_str = plan.display_indent_schema().to_string();
6610            let aggr_line = plan_str
6611                .lines()
6612                .find(|line| line.contains("Aggregate: groupBy="))
6613                .unwrap();
6614            assert!(aggr_line.contains(op), "{plan_str}");
6615            assert!(aggr_line.contains("first_value"), "{plan_str}");
6616            assert!(
6617                !plan
6618                    .schema()
6619                    .fields()
6620                    .iter()
6621                    .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME })
6622            );
6623        }
6624    }
6625
6626    #[tokio::test]
6627    async fn topk_by_does_not_partition_by_tsid() {
6628        let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap();
6629        let eval_stmt = EvalStmt {
6630            expr: prom_expr,
6631            start: UNIX_EPOCH,
6632            end: UNIX_EPOCH
6633                .checked_add(Duration::from_secs(100_000))
6634                .unwrap(),
6635            interval: Duration::from_secs(5),
6636            lookback_delta: Duration::from_secs(1),
6637        };
6638
6639        let table_provider = build_test_table_provider_with_tsid(
6640            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6641            1,
6642            1,
6643        )
6644        .await;
6645        let plan =
6646            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6647                .await
6648                .unwrap();
6649
6650        let plan_str = plan.display_indent_schema().to_string();
6651        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
6652
6653        let window_line = plan_str
6654            .lines()
6655            .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
6656            .unwrap();
6657        let partition_by = window_line
6658            .split("PARTITION BY [")
6659            .nth(1)
6660            .and_then(|s| s.split("] ORDER BY").next())
6661            .unwrap();
6662        assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
6663    }
6664
6665    #[tokio::test]
6666    async fn selector_matcher_on_tsid_does_not_use_internal_column() {
6667        let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap();
6668        let eval_stmt = EvalStmt {
6669            expr: prom_expr,
6670            start: UNIX_EPOCH,
6671            end: UNIX_EPOCH
6672                .checked_add(Duration::from_secs(100_000))
6673                .unwrap(),
6674            interval: Duration::from_secs(5),
6675            lookback_delta: Duration::from_secs(1),
6676        };
6677
6678        let table_provider = build_test_table_provider_with_tsid(
6679            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6680            1,
6681            1,
6682        )
6683        .await;
6684        let plan =
6685            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6686                .await
6687                .unwrap();
6688
6689        fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet<Column>) {
6690            if let LogicalPlan::Filter(filter) = plan {
6691                datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap();
6692            }
6693            for input in plan.inputs() {
6694                collect_filter_cols(input, out);
6695            }
6696        }
6697
6698        let mut filter_cols = HashSet::new();
6699        collect_filter_cols(&plan, &mut filter_cols);
6700        assert!(
6701            !filter_cols
6702                .iter()
6703                .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME)
6704        );
6705    }
6706
6707    #[tokio::test]
6708    async fn tsid_is_not_used_when_physical_table_is_missing() {
6709        let prom_expr = parser::parse("some_metric").unwrap();
6710        let eval_stmt = EvalStmt {
6711            expr: prom_expr,
6712            start: UNIX_EPOCH,
6713            end: UNIX_EPOCH
6714                .checked_add(Duration::from_secs(100_000))
6715                .unwrap(),
6716            interval: Duration::from_secs(5),
6717            lookback_delta: Duration::from_secs(1),
6718        };
6719
6720        let catalog_list = MemoryCatalogManager::with_default_setup();
6721
6722        // Register a metric engine logical table referencing a missing physical table.
6723        let mut columns = vec![ColumnSchema::new(
6724            "tag_0".to_string(),
6725            ConcreteDataType::string_datatype(),
6726            false,
6727        )];
6728        columns.push(
6729            ColumnSchema::new(
6730                "timestamp".to_string(),
6731                ConcreteDataType::timestamp_millisecond_datatype(),
6732                false,
6733            )
6734            .with_time_index(true),
6735        );
6736        columns.push(ColumnSchema::new(
6737            "field_0".to_string(),
6738            ConcreteDataType::float64_datatype(),
6739            true,
6740        ));
6741        let schema = Arc::new(Schema::new(columns));
6742        let mut options = table::requests::TableOptions::default();
6743        options
6744            .extra_options
6745            .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string());
6746        let table_meta = TableMetaBuilder::empty()
6747            .schema(schema)
6748            .primary_key_indices(vec![0])
6749            .value_indices(vec![2])
6750            .engine(METRIC_ENGINE_NAME.to_string())
6751            .options(options)
6752            .next_column_id(1024)
6753            .build()
6754            .unwrap();
6755        let table_info = TableInfoBuilder::default()
6756            .table_id(1024)
6757            .name("some_metric")
6758            .meta(table_meta)
6759            .build()
6760            .unwrap();
6761        let table = EmptyTable::from_table_info(&table_info);
6762        catalog_list
6763            .register_table_sync(RegisterTableRequest {
6764                catalog: DEFAULT_CATALOG_NAME.to_string(),
6765                schema: DEFAULT_SCHEMA_NAME.to_string(),
6766                table_name: "some_metric".to_string(),
6767                table_id: 1024,
6768                table,
6769            })
6770            .unwrap();
6771
6772        let table_provider = DfTableSourceProvider::new(
6773            catalog_list,
6774            false,
6775            QueryContext::arc(),
6776            DummyDecoder::arc(),
6777            false,
6778        );
6779
6780        let plan =
6781            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6782                .await
6783                .unwrap();
6784
6785        let plan_str = plan.display_indent_schema().to_string();
6786        assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
6787        assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
6788    }
6789
6790    #[tokio::test]
6791    async fn tsid_is_carried_only_when_aggregate_preserves_label_set() {
6792        let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap();
6793        let eval_stmt = EvalStmt {
6794            expr: prom_expr,
6795            start: UNIX_EPOCH,
6796            end: UNIX_EPOCH
6797                .checked_add(Duration::from_secs(100_000))
6798                .unwrap(),
6799            interval: Duration::from_secs(5),
6800            lookback_delta: Duration::from_secs(1),
6801        };
6802
6803        let table_provider = build_test_table_provider_with_tsid(
6804            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6805            1,
6806            1,
6807        )
6808        .await;
6809        let plan =
6810            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6811                .await
6812                .unwrap();
6813
6814        let plan_str = plan.display_indent_schema().to_string();
6815        assert!(plan_str.contains("first_value") && plan_str.contains("__tsid"));
6816        assert!(
6817            !plan
6818                .schema()
6819                .fields()
6820                .iter()
6821                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
6822        );
6823
6824        // Merging aggregate: label set is reduced, tsid should not be carried.
6825        let prom_expr = parser::parse("sum(some_metric)").unwrap();
6826        let eval_stmt = EvalStmt {
6827            expr: prom_expr,
6828            start: UNIX_EPOCH,
6829            end: UNIX_EPOCH
6830                .checked_add(Duration::from_secs(100_000))
6831                .unwrap(),
6832            interval: Duration::from_secs(5),
6833            lookback_delta: Duration::from_secs(1),
6834        };
6835        let table_provider = build_test_table_provider_with_tsid(
6836            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6837            1,
6838            1,
6839        )
6840        .await;
6841        let plan =
6842            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6843                .await
6844                .unwrap();
6845        let plan_str = plan.display_indent_schema().to_string();
6846        assert!(!plan_str.contains("first_value"));
6847    }
6848
6849    #[tokio::test]
6850    async fn or_operator_with_unknown_metric_does_not_require_tsid() {
6851        let prom_expr = parser::parse("unknown_metric or some_metric").unwrap();
6852        let eval_stmt = EvalStmt {
6853            expr: prom_expr,
6854            start: UNIX_EPOCH,
6855            end: UNIX_EPOCH
6856                .checked_add(Duration::from_secs(100_000))
6857                .unwrap(),
6858            interval: Duration::from_secs(5),
6859            lookback_delta: Duration::from_secs(1),
6860        };
6861
6862        let table_provider = build_test_table_provider_with_tsid(
6863            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6864            1,
6865            1,
6866        )
6867        .await;
6868
6869        let plan =
6870            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6871                .await
6872                .unwrap();
6873
6874        assert!(
6875            !plan
6876                .schema()
6877                .fields()
6878                .iter()
6879                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
6880        );
6881    }
6882
6883    #[tokio::test]
6884    async fn aggregate_avg() {
6885        do_aggregate_expr_plan("avg", "avg").await;
6886    }
6887
6888    #[tokio::test]
6889    #[should_panic] // output type doesn't match
6890    async fn aggregate_count() {
6891        do_aggregate_expr_plan("count", "count").await;
6892    }
6893
6894    #[tokio::test]
6895    async fn aggregate_min() {
6896        do_aggregate_expr_plan("min", "min").await;
6897    }
6898
6899    #[tokio::test]
6900    async fn aggregate_max() {
6901        do_aggregate_expr_plan("max", "max").await;
6902    }
6903
6904    #[tokio::test]
6905    async fn aggregate_group() {
6906        // Regression test for `group()` aggregator.
6907        // PromQL: sum(group by (cluster)(kubernetes_build_info{service="kubernetes",job="apiserver"}))
6908        // should be plannable, and `group()` should produce constant 1 for each group.
6909        let prom_expr = parser::parse(
6910            "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))",
6911        )
6912        .unwrap();
6913        let eval_stmt = EvalStmt {
6914            expr: prom_expr,
6915            start: UNIX_EPOCH,
6916            end: UNIX_EPOCH
6917                .checked_add(Duration::from_secs(100_000))
6918                .unwrap(),
6919            interval: Duration::from_secs(5),
6920            lookback_delta: Duration::from_secs(1),
6921        };
6922
6923        let table_provider = build_test_table_provider_with_fields(
6924            &[(
6925                DEFAULT_SCHEMA_NAME.to_string(),
6926                "kubernetes_build_info".to_string(),
6927            )],
6928            &["cluster", "service", "job"],
6929        )
6930        .await;
6931        let plan =
6932            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6933                .await
6934                .unwrap();
6935
6936        let plan_str = plan.display_indent_schema().to_string();
6937        assert!(plan_str.contains("max(Float64(1"));
6938    }
6939
6940    #[tokio::test]
6941    async fn aggregate_stddev() {
6942        do_aggregate_expr_plan("stddev", "stddev_pop").await;
6943    }
6944
6945    #[tokio::test]
6946    async fn aggregate_stdvar() {
6947        do_aggregate_expr_plan("stdvar", "var_pop").await;
6948    }
6949
6950    // TODO(ruihang): add range fn tests once exprs are ready.
6951
6952    // {
6953    //     input: "some_metric{tag_0="foo"} + some_metric{tag_0="bar"}",
6954    //     expected: &BinaryExpr{
6955    //         Op: ADD,
6956    //         LHS: &VectorSelector{
6957    //             Name: "a",
6958    //             LabelMatchers: []*labels.Matcher{
6959    //                     MustLabelMatcher(labels.MatchEqual, "tag_0", "foo"),
6960    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
6961    //             },
6962    //         },
6963    //         RHS: &VectorSelector{
6964    //             Name: "sum",
6965    //             LabelMatchers: []*labels.Matcher{
6966    //                     MustLabelMatcher(labels.MatchxEqual, "tag_0", "bar"),
6967    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
6968    //             },
6969    //         },
6970    //         VectorMatching: &VectorMatching{},
6971    //     },
6972    // },
6973    #[tokio::test]
6974    async fn binary_op_column_column() {
6975        let prom_expr =
6976            parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap();
6977        let eval_stmt = EvalStmt {
6978            expr: prom_expr,
6979            start: UNIX_EPOCH,
6980            end: UNIX_EPOCH
6981                .checked_add(Duration::from_secs(100_000))
6982                .unwrap(),
6983            interval: Duration::from_secs(5),
6984            lookback_delta: Duration::from_secs(1),
6985        };
6986
6987        let table_provider = build_test_table_provider(
6988            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
6989            1,
6990            1,
6991        )
6992        .await;
6993        let plan =
6994            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
6995                .await
6996                .unwrap();
6997
6998        let expected = String::from(
6999            "Projection: rhs.tag_0, rhs.timestamp, CAST(lhs.field_0 AS Float64) + CAST(rhs.field_0 AS Float64) AS lhs.field_0 + rhs.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), lhs.field_0 + rhs.field_0:Float64;N]\
7000            \n  Inner Join: lhs.tag_0 = rhs.tag_0, lhs.timestamp = rhs.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7001            \n    SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7002            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7003            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7004            \n          Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7005            \n            Filter: some_metric.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7006            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7007            \n    SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7008            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7009            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7010            \n          Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7011            \n            Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7012            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7013        );
7014
7015        assert_eq!(plan.display_indent_schema().to_string(), expected);
7016    }
7017
7018    async fn indie_query_plan_compare<T: AsRef<str>>(query: &str, expected: T) {
7019        let prom_expr = parser::parse(query).unwrap();
7020        let eval_stmt = EvalStmt {
7021            expr: prom_expr,
7022            start: UNIX_EPOCH,
7023            end: UNIX_EPOCH
7024                .checked_add(Duration::from_secs(100_000))
7025                .unwrap(),
7026            interval: Duration::from_secs(5),
7027            lookback_delta: Duration::from_secs(1),
7028        };
7029
7030        let table_provider = build_test_table_provider(
7031            &[
7032                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
7033                (
7034                    "greptime_private".to_string(),
7035                    "some_alt_metric".to_string(),
7036                ),
7037            ],
7038            1,
7039            1,
7040        )
7041        .await;
7042        let plan =
7043            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7044                .await
7045                .unwrap();
7046
7047        assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref());
7048    }
7049
7050    #[tokio::test]
7051    async fn binary_op_literal_column() {
7052        let query = r#"1 + some_metric{tag_0="bar"}"#;
7053        let expected = String::from(
7054            "Projection: some_metric.tag_0, some_metric.timestamp, Float64(1) + CAST(some_metric.field_0 AS Float64) AS Float64(1) + field_0 [tag_0:Utf8, timestamp:Timestamp(ms), Float64(1) + field_0:Float64;N]\
7055            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7056            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7057            \n      Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7058            \n        Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7059            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7060        );
7061
7062        indie_query_plan_compare(query, expected).await;
7063    }
7064
7065    #[tokio::test]
7066    async fn binary_op_literal_literal() {
7067        let query = r#"1 + 1"#;
7068        let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N]
7069  TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#;
7070        indie_query_plan_compare(query, expected).await;
7071    }
7072
7073    #[tokio::test]
7074    async fn simple_bool_grammar() {
7075        let query = "some_metric != bool 1.2345";
7076        let expected = String::from(
7077            "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 != Float64(1.2345) AS Float64) AS field_0 != Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 != Float64(1.2345):Float64;N]\
7078            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7079            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7080            \n      Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7081            \n        Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7082            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7083        );
7084
7085        indie_query_plan_compare(query, expected).await;
7086    }
7087
7088    #[tokio::test]
7089    async fn bool_with_additional_arithmetic() {
7090        let query = "some_metric + (1 == bool 2)";
7091        let expected = String::from(
7092            "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 AS Float64) + CAST(Float64(1) = Float64(2) AS Float64) AS field_0 + Float64(1) = Float64(2) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 + Float64(1) = Float64(2):Float64;N]\
7093            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7094            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7095            \n      Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7096            \n        Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7097            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7098        );
7099
7100        indie_query_plan_compare(query, expected).await;
7101    }
7102
7103    #[tokio::test]
7104    async fn simple_unary() {
7105        let query = "-some_metric";
7106        let expected = String::from(
7107            "Projection: some_metric.tag_0, some_metric.timestamp, (- some_metric.field_0) AS (- field_0) [tag_0:Utf8, timestamp:Timestamp(ms), (- field_0):Float64;N]\
7108            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7109            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7110            \n      Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7111            \n        Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7112            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7113        );
7114
7115        indie_query_plan_compare(query, expected).await;
7116    }
7117
7118    #[tokio::test]
7119    async fn increase_aggr() {
7120        let query = "increase(some_metric[5m])";
7121        let expected = String::from(
7122            "Filter: prom_increase(timestamp_range,field_0,timestamp,Int64(300000)) IS NOT NULL [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\
7123            \n  Projection: some_metric.timestamp, prom_increase(timestamp_range, field_0, some_metric.timestamp, Int64(300000)) AS prom_increase(timestamp_range,field_0,timestamp,Int64(300000)), some_metric.tag_0 [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\
7124            \n    PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\
7125            \n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7126            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7127            \n          Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7128            \n            Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7129            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7130        );
7131
7132        indie_query_plan_compare(query, expected).await;
7133    }
7134
7135    #[tokio::test]
7136    async fn less_filter_on_value() {
7137        let query = "some_metric < 1.2345";
7138        let expected = String::from(
7139            "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7140            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7141            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7142            \n      Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7143            \n        Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7144            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7145        );
7146
7147        indie_query_plan_compare(query, expected).await;
7148    }
7149
7150    #[tokio::test]
7151    async fn count_over_time() {
7152        let query = "count_over_time(some_metric[5m])";
7153        let expected = String::from(
7154            "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\
7155            \n  Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\
7156            \n    PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\
7157            \n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7158            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7159            \n          Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7160            \n            Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7161            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7162        );
7163
7164        indie_query_plan_compare(query, expected).await;
7165    }
7166
7167    /// The outer `PromRangeManipulate` from a subquery must be preceded by
7168    /// `Sort` + `PromSeriesDivide`.
7169    #[tokio::test]
7170    async fn count_over_time_subquery() {
7171        let query = "count_over_time(some_metric[10m:1m])";
7172        let expected = String::from(
7173            "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\
7174            \n  Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\
7175            \n    PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\
7176            \n      PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7177            \n        Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7178            \n          PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7179            \n            PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7180            \n              Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7181            \n                Filter: some_metric.timestamp >= TimestampMillisecond(-540999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7182            \n                  TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7183        );
7184        indie_query_plan_compare(query, expected).await;
7185    }
7186
7187    #[tokio::test]
7188    async fn test_hash_join() {
7189        let mut eval_stmt = EvalStmt {
7190            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7191            start: UNIX_EPOCH,
7192            end: UNIX_EPOCH
7193                .checked_add(Duration::from_secs(100_000))
7194                .unwrap(),
7195            interval: Duration::from_secs(5),
7196            lookback_delta: Duration::from_secs(1),
7197        };
7198
7199        let case = r#"http_server_requests_seconds_sum{uri="/accounts/login"} / ignoring(kubernetes_pod_name,kubernetes_namespace) http_server_requests_seconds_count{uri="/accounts/login"}"#;
7200
7201        let prom_expr = parser::parse(case).unwrap();
7202        eval_stmt.expr = prom_expr;
7203        let table_provider = build_test_table_provider_with_fields(
7204            &[
7205                (
7206                    DEFAULT_SCHEMA_NAME.to_string(),
7207                    "http_server_requests_seconds_sum".to_string(),
7208                ),
7209                (
7210                    DEFAULT_SCHEMA_NAME.to_string(),
7211                    "http_server_requests_seconds_count".to_string(),
7212                ),
7213            ],
7214            &["uri", "kubernetes_namespace", "kubernetes_pod_name"],
7215        )
7216        .await;
7217        // Should be ok
7218        let plan =
7219            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7220                .await
7221                .unwrap();
7222        let expected = "Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\
7223            \n  Inner Join: http_server_requests_seconds_sum.greptime_timestamp = http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_sum.uri = http_server_requests_seconds_count.uri\
7224            \n    SubqueryAlias: http_server_requests_seconds_sum\
7225            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
7226            \n        PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
7227            \n          Sort: http_server_requests_seconds_sum.uri ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_sum.greptime_timestamp ASC NULLS FIRST\
7228            \n            Filter: http_server_requests_seconds_sum.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_sum.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_sum.greptime_timestamp <= TimestampMillisecond(100000000, None)\
7229            \n              TableScan: http_server_requests_seconds_sum\
7230            \n    SubqueryAlias: http_server_requests_seconds_count\
7231            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
7232            \n        PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
7233            \n          Sort: http_server_requests_seconds_count.uri ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_count.greptime_timestamp ASC NULLS FIRST\
7234            \n            Filter: http_server_requests_seconds_count.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_count.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_count.greptime_timestamp <= TimestampMillisecond(100000000, None)\
7235            \n              TableScan: http_server_requests_seconds_count";
7236        assert_eq!(plan.to_string(), expected);
7237    }
7238
7239    #[tokio::test]
7240    async fn test_nested_histogram_quantile() {
7241        let mut eval_stmt = EvalStmt {
7242            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7243            start: UNIX_EPOCH,
7244            end: UNIX_EPOCH
7245                .checked_add(Duration::from_secs(100_000))
7246                .unwrap(),
7247            interval: Duration::from_secs(5),
7248            lookback_delta: Duration::from_secs(1),
7249        };
7250
7251        let case = r#"label_replace(histogram_quantile(0.99, sum by(pod, le, path, code) (rate(greptime_servers_grpc_requests_elapsed_bucket{container="frontend"}[1m0s]))), "pod_new", "$1", "pod", "greptimedb-frontend-[0-9a-z]*-(.*)")"#;
7252
7253        let prom_expr = parser::parse(case).unwrap();
7254        eval_stmt.expr = prom_expr;
7255        let table_provider = build_test_table_provider_with_fields(
7256            &[(
7257                DEFAULT_SCHEMA_NAME.to_string(),
7258                "greptime_servers_grpc_requests_elapsed_bucket".to_string(),
7259            )],
7260            &["pod", "le", "path", "code", "container"],
7261        )
7262        .await;
7263        // Should be ok
7264        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7265            .await
7266            .unwrap();
7267    }
7268
7269    #[tokio::test]
7270    async fn test_histogram_quantile_binary_op() {
7271        let mut eval_stmt = EvalStmt {
7272            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7273            start: UNIX_EPOCH,
7274            end: UNIX_EPOCH
7275                .checked_add(Duration::from_secs(100_000))
7276                .unwrap(),
7277            interval: Duration::from_secs(5),
7278            lookback_delta: Duration::from_secs(1),
7279        };
7280
7281        // Arithmetic applied to a histogram_quantile() result. Regression for #8144:
7282        // HistogramFold used to drop the input column qualifiers, so the binary-op
7283        // projection failed to resolve the qualified tag column.
7284        let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#;
7285
7286        let prom_expr = parser::parse(case).unwrap();
7287        eval_stmt.expr = prom_expr;
7288        let table_provider = build_test_table_provider_with_fields(
7289            &[(
7290                DEFAULT_SCHEMA_NAME.to_string(),
7291                "http_request_duration_seconds_bucket".to_string(),
7292            )],
7293            &["pod", "le"],
7294        )
7295        .await;
7296        // Should plan without a "No field named ..." error.
7297        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7298            .await
7299            .unwrap();
7300    }
7301
7302    #[tokio::test]
7303    async fn test_parse_and_operator() {
7304        let mut eval_stmt = EvalStmt {
7305            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7306            start: UNIX_EPOCH,
7307            end: UNIX_EPOCH
7308                .checked_add(Duration::from_secs(100_000))
7309                .unwrap(),
7310            interval: Duration::from_secs(5),
7311            lookback_delta: Duration::from_secs(1),
7312        };
7313
7314        let cases = [
7315            r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) and (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#,
7316            r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) unless (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#,
7317        ];
7318
7319        for case in cases {
7320            let prom_expr = parser::parse(case).unwrap();
7321            eval_stmt.expr = prom_expr;
7322            let table_provider = build_test_table_provider_with_fields(
7323                &[
7324                    (
7325                        DEFAULT_SCHEMA_NAME.to_string(),
7326                        "kubelet_volume_stats_used_bytes".to_string(),
7327                    ),
7328                    (
7329                        DEFAULT_SCHEMA_NAME.to_string(),
7330                        "kubelet_volume_stats_capacity_bytes".to_string(),
7331                    ),
7332                ],
7333                &["namespace", "persistentvolumeclaim"],
7334            )
7335            .await;
7336            // Should be ok
7337            let _ =
7338                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7339                    .await
7340                    .unwrap();
7341        }
7342    }
7343
7344    #[tokio::test]
7345    async fn test_nested_binary_op() {
7346        let mut eval_stmt = EvalStmt {
7347            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7348            start: UNIX_EPOCH,
7349            end: UNIX_EPOCH
7350                .checked_add(Duration::from_secs(100_000))
7351                .unwrap(),
7352            interval: Duration::from_secs(5),
7353            lookback_delta: Duration::from_secs(1),
7354        };
7355
7356        let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) -
7357        (
7358            sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m]))
7359            or
7360            vector(0)
7361        )"#;
7362
7363        let prom_expr = parser::parse(case).unwrap();
7364        eval_stmt.expr = prom_expr;
7365        let table_provider = build_test_table_provider_with_fields(
7366            &[(
7367                DEFAULT_SCHEMA_NAME.to_string(),
7368                "nginx_ingress_controller_requests".to_string(),
7369            )],
7370            &["namespace", "job"],
7371        )
7372        .await;
7373        // Should be ok
7374        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7375            .await
7376            .unwrap();
7377    }
7378
7379    #[tokio::test]
7380    async fn test_parse_or_operator() {
7381        let mut eval_stmt = EvalStmt {
7382            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7383            start: UNIX_EPOCH,
7384            end: UNIX_EPOCH
7385                .checked_add(Duration::from_secs(100_000))
7386                .unwrap(),
7387            interval: Duration::from_secs(5),
7388            lookback_delta: Duration::from_secs(1),
7389        };
7390
7391        let case = r#"
7392        sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) /
7393        (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100)
7394            or
7395        200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) /
7396        sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#;
7397
7398        let table_provider = build_test_table_provider_with_fields(
7399            &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
7400            &["tenant_name", "cluster_name"],
7401        )
7402        .await;
7403        eval_stmt.expr = parser::parse(case).unwrap();
7404        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7405            .await
7406            .unwrap();
7407
7408        let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
7409            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) +
7410            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
7411            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
7412            or
7413            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
7414            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
7415            or
7416            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
7417            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#;
7418        let table_provider = build_test_table_provider_with_fields(
7419            &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
7420            &["tenant_name", "cluster_name"],
7421        )
7422        .await;
7423        eval_stmt.expr = parser::parse(case).unwrap();
7424        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7425            .await
7426            .unwrap();
7427
7428        let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) +
7429            sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
7430            (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
7431            (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#;
7432        let table_provider = build_test_table_provider_with_fields(
7433            &[
7434                (
7435                    DEFAULT_SCHEMA_NAME.to_string(),
7436                    "background_waitevent_cnt".to_string(),
7437                ),
7438                (
7439                    DEFAULT_SCHEMA_NAME.to_string(),
7440                    "foreground_waitevent_cnt".to_string(),
7441                ),
7442            ],
7443            &["tenant_name", "cluster_name"],
7444        )
7445        .await;
7446        eval_stmt.expr = parser::parse(case).unwrap();
7447        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7448            .await
7449            .unwrap();
7450
7451        let case = r#"avg(node_load1{cluster_name=~"cluster1"}) by (cluster_name,host_name) or max(container_cpu_load_average_10s{cluster_name=~"cluster1"}) by (cluster_name,host_name) * 100 / max(container_spec_cpu_quota{cluster_name=~"cluster1"}) by (cluster_name,host_name)"#;
7452        let table_provider = build_test_table_provider_with_fields(
7453            &[
7454                (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()),
7455                (
7456                    DEFAULT_SCHEMA_NAME.to_string(),
7457                    "container_cpu_load_average_10s".to_string(),
7458                ),
7459                (
7460                    DEFAULT_SCHEMA_NAME.to_string(),
7461                    "container_spec_cpu_quota".to_string(),
7462                ),
7463            ],
7464            &["cluster_name", "host_name"],
7465        )
7466        .await;
7467        eval_stmt.expr = parser::parse(case).unwrap();
7468        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7469            .await
7470            .unwrap();
7471    }
7472
7473    #[tokio::test]
7474    async fn value_matcher() {
7475        // template
7476        let mut eval_stmt = EvalStmt {
7477            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7478            start: UNIX_EPOCH,
7479            end: UNIX_EPOCH
7480                .checked_add(Duration::from_secs(100_000))
7481                .unwrap(),
7482            interval: Duration::from_secs(5),
7483            lookback_delta: Duration::from_secs(1),
7484        };
7485
7486        let cases = [
7487            // single equal matcher
7488            (
7489                r#"some_metric{__field__="field_1"}"#,
7490                vec![
7491                    "some_metric.field_1",
7492                    "some_metric.tag_0",
7493                    "some_metric.tag_1",
7494                    "some_metric.tag_2",
7495                    "some_metric.timestamp",
7496                ],
7497            ),
7498            // two equal matchers
7499            (
7500                r#"some_metric{__field__="field_1", __field__="field_0"}"#,
7501                vec![
7502                    "some_metric.field_0",
7503                    "some_metric.field_1",
7504                    "some_metric.tag_0",
7505                    "some_metric.tag_1",
7506                    "some_metric.tag_2",
7507                    "some_metric.timestamp",
7508                ],
7509            ),
7510            // single not_eq matcher
7511            (
7512                r#"some_metric{__field__!="field_1"}"#,
7513                vec![
7514                    "some_metric.field_0",
7515                    "some_metric.field_2",
7516                    "some_metric.tag_0",
7517                    "some_metric.tag_1",
7518                    "some_metric.tag_2",
7519                    "some_metric.timestamp",
7520                ],
7521            ),
7522            // two not_eq matchers
7523            (
7524                r#"some_metric{__field__!="field_1", __field__!="field_2"}"#,
7525                vec![
7526                    "some_metric.field_0",
7527                    "some_metric.tag_0",
7528                    "some_metric.tag_1",
7529                    "some_metric.tag_2",
7530                    "some_metric.timestamp",
7531                ],
7532            ),
7533            // equal and not_eq matchers (no conflict)
7534            (
7535                r#"some_metric{__field__="field_1", __field__!="field_0"}"#,
7536                vec![
7537                    "some_metric.field_1",
7538                    "some_metric.tag_0",
7539                    "some_metric.tag_1",
7540                    "some_metric.tag_2",
7541                    "some_metric.timestamp",
7542                ],
7543            ),
7544            // equal and not_eq matchers (conflict)
7545            (
7546                r#"some_metric{__field__="field_2", __field__!="field_2"}"#,
7547                vec![
7548                    "some_metric.tag_0",
7549                    "some_metric.tag_1",
7550                    "some_metric.tag_2",
7551                    "some_metric.timestamp",
7552                ],
7553            ),
7554            // single regex eq matcher
7555            (
7556                r#"some_metric{__field__=~"field_1|field_2"}"#,
7557                vec![
7558                    "some_metric.field_1",
7559                    "some_metric.field_2",
7560                    "some_metric.tag_0",
7561                    "some_metric.tag_1",
7562                    "some_metric.tag_2",
7563                    "some_metric.timestamp",
7564                ],
7565            ),
7566            // single regex not_eq matcher
7567            (
7568                r#"some_metric{__field__!~"field_1|field_2"}"#,
7569                vec![
7570                    "some_metric.field_0",
7571                    "some_metric.tag_0",
7572                    "some_metric.tag_1",
7573                    "some_metric.tag_2",
7574                    "some_metric.timestamp",
7575                ],
7576            ),
7577        ];
7578
7579        for case in cases {
7580            let prom_expr = parser::parse(case.0).unwrap();
7581            eval_stmt.expr = prom_expr;
7582            let table_provider = build_test_table_provider(
7583                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7584                3,
7585                3,
7586            )
7587            .await;
7588            let plan =
7589                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7590                    .await
7591                    .unwrap();
7592            let mut fields = plan.schema().field_names();
7593            let mut expected = case.1.into_iter().map(String::from).collect::<Vec<_>>();
7594            fields.sort();
7595            expected.sort();
7596            assert_eq!(fields, expected, "case: {:?}", case.0);
7597        }
7598
7599        let bad_cases = [
7600            r#"some_metric{__field__="nonexistent"}"#,
7601            r#"some_metric{__field__!="nonexistent"}"#,
7602        ];
7603
7604        for case in bad_cases {
7605            let prom_expr = parser::parse(case).unwrap();
7606            eval_stmt.expr = prom_expr;
7607            let table_provider = build_test_table_provider(
7608                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7609                3,
7610                3,
7611            )
7612            .await;
7613            let plan =
7614                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7615                    .await;
7616            assert!(plan.is_err(), "case: {:?}", case);
7617        }
7618    }
7619
7620    #[tokio::test]
7621    async fn custom_schema() {
7622        let query = "some_alt_metric{__schema__=\"greptime_private\"}";
7623        let expected = String::from(
7624            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7625            \n  PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7626            \n    Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7627            \n      Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7628            \n        TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7629        );
7630
7631        indie_query_plan_compare(query, expected).await;
7632
7633        let query = "some_alt_metric{__database__=\"greptime_private\"}";
7634        let expected = String::from(
7635            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7636            \n  PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7637            \n    Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7638            \n      Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7639            \n        TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7640        );
7641
7642        indie_query_plan_compare(query, expected).await;
7643
7644        let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric";
7645        let expected = String::from(
7646            "Projection: some_metric.tag_0, some_metric.timestamp, CAST(greptime_private.some_alt_metric.field_0 AS Float64) / CAST(some_metric.field_0 AS Float64) AS greptime_private.some_alt_metric.field_0 / some_metric.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), greptime_private.some_alt_metric.field_0 / some_metric.field_0:Float64;N]\
7647            \n  Inner Join: greptime_private.some_alt_metric.tag_0 = some_metric.tag_0, greptime_private.some_alt_metric.timestamp = some_metric.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7648            \n    SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7649            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7650            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7651            \n          Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7652            \n            Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7653            \n              TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7654            \n    SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7655            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7656            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7657            \n          Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7658            \n            Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
7659            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
7660        );
7661
7662        indie_query_plan_compare(query, expected).await;
7663    }
7664
7665    #[tokio::test]
7666    async fn only_equals_is_supported_for_special_matcher() {
7667        let queries = &[
7668            "some_alt_metric{__schema__!=\"greptime_private\"}",
7669            "some_alt_metric{__schema__=~\"lalala\"}",
7670            "some_alt_metric{__database__!=\"greptime_private\"}",
7671            "some_alt_metric{__database__=~\"lalala\"}",
7672        ];
7673
7674        for query in queries {
7675            let prom_expr = parser::parse(query).unwrap();
7676            let eval_stmt = EvalStmt {
7677                expr: prom_expr,
7678                start: UNIX_EPOCH,
7679                end: UNIX_EPOCH
7680                    .checked_add(Duration::from_secs(100_000))
7681                    .unwrap(),
7682                interval: Duration::from_secs(5),
7683                lookback_delta: Duration::from_secs(1),
7684            };
7685
7686            let table_provider = build_test_table_provider(
7687                &[
7688                    (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
7689                    (
7690                        "greptime_private".to_string(),
7691                        "some_alt_metric".to_string(),
7692                    ),
7693                ],
7694                1,
7695                1,
7696            )
7697            .await;
7698
7699            let plan =
7700                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7701                    .await;
7702            assert!(plan.is_err(), "query: {:?}", query);
7703        }
7704    }
7705
7706    #[tokio::test]
7707    async fn test_non_ms_precision() {
7708        let catalog_list = MemoryCatalogManager::with_default_setup();
7709        let columns = vec![
7710            ColumnSchema::new(
7711                "tag".to_string(),
7712                ConcreteDataType::string_datatype(),
7713                false,
7714            ),
7715            ColumnSchema::new(
7716                "timestamp".to_string(),
7717                ConcreteDataType::timestamp_nanosecond_datatype(),
7718                false,
7719            )
7720            .with_time_index(true),
7721            ColumnSchema::new(
7722                "field".to_string(),
7723                ConcreteDataType::float64_datatype(),
7724                true,
7725            ),
7726        ];
7727        let schema = Arc::new(Schema::new(columns));
7728        let table_meta = TableMetaBuilder::empty()
7729            .schema(schema)
7730            .primary_key_indices(vec![0])
7731            .value_indices(vec![2])
7732            .next_column_id(1024)
7733            .build()
7734            .unwrap();
7735        let table_info = TableInfoBuilder::default()
7736            .name("metrics".to_string())
7737            .meta(table_meta)
7738            .build()
7739            .unwrap();
7740        let table = EmptyTable::from_table_info(&table_info);
7741        assert!(
7742            catalog_list
7743                .register_table_sync(RegisterTableRequest {
7744                    catalog: DEFAULT_CATALOG_NAME.to_string(),
7745                    schema: DEFAULT_SCHEMA_NAME.to_string(),
7746                    table_name: "metrics".to_string(),
7747                    table_id: 1024,
7748                    table,
7749                })
7750                .is_ok()
7751        );
7752
7753        let plan = PromPlanner::stmt_to_plan(
7754            DfTableSourceProvider::new(
7755                catalog_list.clone(),
7756                false,
7757                QueryContext::arc(),
7758                DummyDecoder::arc(),
7759                true,
7760            ),
7761            &EvalStmt {
7762                expr: parser::parse("metrics{tag = \"1\"}").unwrap(),
7763                start: UNIX_EPOCH,
7764                end: UNIX_EPOCH
7765                    .checked_add(Duration::from_secs(100_000))
7766                    .unwrap(),
7767                interval: Duration::from_secs(5),
7768                lookback_delta: Duration::from_secs(1),
7769            },
7770            &build_query_engine_state(),
7771        )
7772        .await
7773        .unwrap();
7774        assert_eq!(
7775            plan.display_indent_schema().to_string(),
7776            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7777            \n  PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7778            \n    Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7779            \n      Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp >= TimestampMillisecond(-999, None) AND metrics.timestamp <= TimestampMillisecond(100000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7780            \n        Projection: metrics.field, metrics.tag, CAST(metrics.timestamp AS Timestamp(ms)) AS timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7781            \n          TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
7782        );
7783        let plan = PromPlanner::stmt_to_plan(
7784            DfTableSourceProvider::new(
7785                catalog_list.clone(),
7786                false,
7787                QueryContext::arc(),
7788                DummyDecoder::arc(),
7789                true,
7790            ),
7791            &EvalStmt {
7792                expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(),
7793                start: UNIX_EPOCH,
7794                end: UNIX_EPOCH
7795                    .checked_add(Duration::from_secs(100_000))
7796                    .unwrap(),
7797                interval: Duration::from_secs(5),
7798                lookback_delta: Duration::from_secs(1),
7799            },
7800            &build_query_engine_state(),
7801        )
7802        .await
7803        .unwrap();
7804        assert_eq!(
7805            plan.display_indent_schema().to_string(),
7806            "Filter: prom_avg_over_time(timestamp_range,field) IS NOT NULL [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\
7807            \n  Projection: metrics.timestamp, prom_avg_over_time(timestamp_range, field) AS prom_avg_over_time(timestamp_range,field), metrics.tag [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\
7808            \n    PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[5000], time index=[timestamp], values=[\"field\"] [field:Dictionary(Int64, Float64);N, tag:Utf8, timestamp:Timestamp(ms), timestamp_range:Dictionary(Int64, Timestamp(ms))]\
7809            \n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7810            \n        PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7811            \n          Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7812            \n            Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp >= TimestampMillisecond(-4999, None) AND metrics.timestamp <= TimestampMillisecond(100000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7813            \n              Projection: metrics.field, metrics.tag, CAST(metrics.timestamp AS Timestamp(ms)) AS timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
7814            \n                TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
7815        );
7816    }
7817
7818    #[tokio::test]
7819    async fn test_nonexistent_label() {
7820        // template
7821        let mut eval_stmt = EvalStmt {
7822            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7823            start: UNIX_EPOCH,
7824            end: UNIX_EPOCH
7825                .checked_add(Duration::from_secs(100_000))
7826                .unwrap(),
7827            interval: Duration::from_secs(5),
7828            lookback_delta: Duration::from_secs(1),
7829        };
7830
7831        let case = r#"some_metric{nonexistent="hi"}"#;
7832        let prom_expr = parser::parse(case).unwrap();
7833        eval_stmt.expr = prom_expr;
7834        let table_provider = build_test_table_provider(
7835            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7836            3,
7837            3,
7838        )
7839        .await;
7840        // Should be ok
7841        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7842            .await
7843            .unwrap();
7844    }
7845
7846    #[tokio::test]
7847    async fn test_label_join() {
7848        let prom_expr = parser::parse(
7849            "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')",
7850        )
7851        .unwrap();
7852        let eval_stmt = EvalStmt {
7853            expr: prom_expr,
7854            start: UNIX_EPOCH,
7855            end: UNIX_EPOCH
7856                .checked_add(Duration::from_secs(100_000))
7857                .unwrap(),
7858            interval: Duration::from_secs(5),
7859            lookback_delta: Duration::from_secs(1),
7860        };
7861
7862        let table_provider =
7863            build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1)
7864                .await;
7865        let plan =
7866            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7867                .await
7868                .unwrap();
7869
7870        let expected = r#"
7871Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8]
7872  Projection: up.timestamp, up.field_0, concat_ws(Utf8(","), up.tag_1, up.tag_2, up.tag_3) AS foo, up.tag_0, up.tag_1, up.tag_2, up.tag_3 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8]
7873    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7874      PromSeriesDivide: tags=["tag_0", "tag_1", "tag_2", "tag_3"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7875        Sort: up.tag_0 ASC NULLS FIRST, up.tag_1 ASC NULLS FIRST, up.tag_2 ASC NULLS FIRST, up.tag_3 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7876          Filter: up.tag_0 = Utf8("api-server") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7877            TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
7878
7879        let ret = plan.display_indent_schema().to_string();
7880        assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
7881    }
7882
7883    #[tokio::test]
7884    async fn test_label_replace() {
7885        let prom_expr = parser::parse(
7886            "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")",
7887        )
7888        .unwrap();
7889        let eval_stmt = EvalStmt {
7890            expr: prom_expr,
7891            start: UNIX_EPOCH,
7892            end: UNIX_EPOCH
7893                .checked_add(Duration::from_secs(100_000))
7894                .unwrap(),
7895            interval: Duration::from_secs(5),
7896            lookback_delta: Duration::from_secs(1),
7897        };
7898
7899        let table_provider =
7900            build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1)
7901                .await;
7902        let plan =
7903            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7904                .await
7905                .unwrap();
7906
7907        let expected = r#"
7908Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8]
7909  Projection: up.timestamp, up.field_0, regexp_replace(up.tag_0, Utf8("^(?s:(.*):.*)$"), Utf8("$1")) AS foo, up.tag_0 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8]
7910    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7911      PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7912        Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7913          Filter: up.tag_0 = Utf8("a:c") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
7914            TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
7915
7916        let ret = plan.display_indent_schema().to_string();
7917        assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
7918    }
7919
7920    #[tokio::test]
7921    async fn test_matchers_to_expr() {
7922        let mut eval_stmt = EvalStmt {
7923            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7924            start: UNIX_EPOCH,
7925            end: UNIX_EPOCH
7926                .checked_add(Duration::from_secs(100_000))
7927                .unwrap(),
7928            interval: Duration::from_secs(5),
7929            lookback_delta: Duration::from_secs(1),
7930        };
7931        let case =
7932            r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#;
7933
7934        let prom_expr = parser::parse(case).unwrap();
7935        eval_stmt.expr = prom_expr;
7936        let table_provider = build_test_table_provider(
7937            &[(
7938                DEFAULT_SCHEMA_NAME.to_string(),
7939                "prometheus_tsdb_head_series".to_string(),
7940            )],
7941            3,
7942            3,
7943        )
7944        .await;
7945        let plan =
7946            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7947                .await
7948                .unwrap();
7949        let expected = "Sort: prometheus_tsdb_head_series.timestamp ASC NULLS LAST [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\
7950        \n  Aggregate: groupBy=[[prometheus_tsdb_head_series.timestamp]], aggr=[[sum(prometheus_tsdb_head_series.field_0), sum(prometheus_tsdb_head_series.field_1), sum(prometheus_tsdb_head_series.field_2)]] [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\
7951        \n    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\
7952        \n      PromSeriesDivide: tags=[\"tag_0\", \"tag_1\", \"tag_2\"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\
7953        \n        Sort: prometheus_tsdb_head_series.tag_0 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_1 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_2 ASC NULLS FIRST, prometheus_tsdb_head_series.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\
7954        \n          Filter: prometheus_tsdb_head_series.tag_1 ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\
7955        \n            TableScan: prometheus_tsdb_head_series [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]";
7956        assert_eq!(plan.display_indent_schema().to_string(), expected);
7957    }
7958
7959    #[tokio::test]
7960    async fn test_topk_expr() {
7961        let mut eval_stmt = EvalStmt {
7962            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
7963            start: UNIX_EPOCH,
7964            end: UNIX_EPOCH
7965                .checked_add(Duration::from_secs(100_000))
7966                .unwrap(),
7967            interval: Duration::from_secs(5),
7968            lookback_delta: Duration::from_secs(1),
7969        };
7970        let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
7971
7972        let prom_expr = parser::parse(case).unwrap();
7973        eval_stmt.expr = prom_expr;
7974        let table_provider = build_test_table_provider_with_fields(
7975            &[
7976                (
7977                    DEFAULT_SCHEMA_NAME.to_string(),
7978                    "prometheus_tsdb_head_series".to_string(),
7979                ),
7980                (
7981                    DEFAULT_SCHEMA_NAME.to_string(),
7982                    "http_server_requests_seconds_count".to_string(),
7983                ),
7984            ],
7985            &["ip"],
7986        )
7987        .await;
7988
7989        let plan =
7990            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
7991                .await
7992                .unwrap();
7993        let expected = "Projection: sum(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [sum(prometheus_tsdb_head_series.greptime_value):Float64;N, ip:Utf8, greptime_timestamp:Timestamp(ms)]\
7994        \n  Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\
7995        \n    Filter: row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= Float64(10) [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\
7996        \n      WindowAggr: windowExpr=[[row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\
7997        \n        Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\
7998        \n          Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\
7999        \n            PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8000        \n              PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8001        \n                Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8002        \n                  Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8003        \n                    TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
8004
8005        assert_eq!(plan.display_indent_schema().to_string(), expected);
8006    }
8007
8008    #[tokio::test]
8009    async fn test_count_values_expr() {
8010        let mut eval_stmt = EvalStmt {
8011            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
8012            start: UNIX_EPOCH,
8013            end: UNIX_EPOCH
8014                .checked_add(Duration::from_secs(100_000))
8015                .unwrap(),
8016            interval: Duration::from_secs(5),
8017            lookback_delta: Duration::from_secs(1),
8018        };
8019        let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
8020
8021        let prom_expr = parser::parse(case).unwrap();
8022        eval_stmt.expr = prom_expr;
8023        let table_provider = build_test_table_provider_with_fields(
8024            &[
8025                (
8026                    DEFAULT_SCHEMA_NAME.to_string(),
8027                    "prometheus_tsdb_head_series".to_string(),
8028                ),
8029                (
8030                    DEFAULT_SCHEMA_NAME.to_string(),
8031                    "http_server_requests_seconds_count".to_string(),
8032                ),
8033            ],
8034            &["ip"],
8035        )
8036        .await;
8037
8038        let plan =
8039            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8040                .await
8041                .unwrap();
8042        let expected = "Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\
8043        \n  Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.greptime_value ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]\
8044        \n    Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series, prometheus_tsdb_head_series.greptime_value [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]\
8045        \n      Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\
8046        \n        PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8047        \n          PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8048        \n            Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8049        \n              Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8050        \n                TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
8051
8052        assert_eq!(plan.display_indent_schema().to_string(), expected);
8053    }
8054
8055    #[tokio::test]
8056    async fn test_value_alias() {
8057        let mut eval_stmt = EvalStmt {
8058            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
8059            start: UNIX_EPOCH,
8060            end: UNIX_EPOCH
8061                .checked_add(Duration::from_secs(100_000))
8062                .unwrap(),
8063            interval: Duration::from_secs(5),
8064            lookback_delta: Duration::from_secs(1),
8065        };
8066        let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
8067
8068        let prom_expr = parser::parse(case).unwrap();
8069        eval_stmt.expr = prom_expr;
8070        eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series");
8071        let table_provider = build_test_table_provider_with_fields(
8072            &[
8073                (
8074                    DEFAULT_SCHEMA_NAME.to_string(),
8075                    "prometheus_tsdb_head_series".to_string(),
8076                ),
8077                (
8078                    DEFAULT_SCHEMA_NAME.to_string(),
8079                    "http_server_requests_seconds_count".to_string(),
8080                ),
8081            ],
8082            &["ip"],
8083        )
8084        .await;
8085
8086        let plan =
8087            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8088                .await
8089                .unwrap();
8090        let expected = r#"
8091Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, greptime_timestamp:Timestamp(ms)]
8092  Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]
8093    Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.greptime_value ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]
8094      Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series, prometheus_tsdb_head_series.greptime_value [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]
8095        Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]
8096          PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8097            PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8098              Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8099                Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8100                  TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#;
8101        assert_eq!(format!("\n{}", plan.display_indent_schema()), expected);
8102    }
8103
8104    #[tokio::test]
8105    async fn test_quantile_expr() {
8106        let mut eval_stmt = EvalStmt {
8107            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
8108            start: UNIX_EPOCH,
8109            end: UNIX_EPOCH
8110                .checked_add(Duration::from_secs(100_000))
8111                .unwrap(),
8112            interval: Duration::from_secs(5),
8113            lookback_delta: Duration::from_secs(1),
8114        };
8115        let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
8116
8117        let prom_expr = parser::parse(case).unwrap();
8118        eval_stmt.expr = prom_expr;
8119        let table_provider = build_test_table_provider_with_fields(
8120            &[
8121                (
8122                    DEFAULT_SCHEMA_NAME.to_string(),
8123                    "prometheus_tsdb_head_series".to_string(),
8124                ),
8125                (
8126                    DEFAULT_SCHEMA_NAME.to_string(),
8127                    "http_server_requests_seconds_count".to_string(),
8128                ),
8129            ],
8130            &["ip"],
8131        )
8132        .await;
8133
8134        let plan =
8135            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8136                .await
8137                .unwrap();
8138        let expected = "Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\
8139        \n  Aggregate: groupBy=[[prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[quantile(Float64(0.3), sum(prometheus_tsdb_head_series.greptime_value))]] [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\
8140        \n    Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\
8141        \n      Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\
8142        \n        PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8143        \n          PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8144        \n            Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8145        \n              Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
8146        \n                TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
8147
8148        assert_eq!(plan.display_indent_schema().to_string(), expected);
8149    }
8150
8151    #[tokio::test]
8152    async fn test_or_not_exists_table_label() {
8153        let mut eval_stmt = EvalStmt {
8154            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
8155            start: UNIX_EPOCH,
8156            end: UNIX_EPOCH
8157                .checked_add(Duration::from_secs(100_000))
8158                .unwrap(),
8159            interval: Duration::from_secs(5),
8160            lookback_delta: Duration::from_secs(1),
8161        };
8162        let case = r#"sum by (job, tag0, tag2) (metric_exists) or sum by (job, tag0, tag2) (metric_not_exists)"#;
8163
8164        let prom_expr = parser::parse(case).unwrap();
8165        eval_stmt.expr = prom_expr;
8166        let table_provider = build_test_table_provider_with_fields(
8167            &[(DEFAULT_SCHEMA_NAME.to_string(), "metric_exists".to_string())],
8168            &["job"],
8169        )
8170        .await;
8171
8172        let plan =
8173            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8174                .await
8175                .unwrap();
8176        let expected = r#"UnionDistinctOn: on col=[["job"]], ts_col=[greptime_timestamp] [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N]
8177  SubqueryAlias: metric_exists [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N]
8178    Projection: metric_exists.greptime_timestamp, metric_exists.job, sum(metric_exists.greptime_value) [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N]
8179      Sort: metric_exists.job ASC NULLS LAST, metric_exists.greptime_timestamp ASC NULLS LAST [job:Utf8, greptime_timestamp:Timestamp(ms), sum(metric_exists.greptime_value):Float64;N]
8180        Aggregate: groupBy=[[metric_exists.job, metric_exists.greptime_timestamp]], aggr=[[sum(metric_exists.greptime_value)]] [job:Utf8, greptime_timestamp:Timestamp(ms), sum(metric_exists.greptime_value):Float64;N]
8181          PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8182            PromSeriesDivide: tags=["job"] [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8183              Sort: metric_exists.job ASC NULLS FIRST, metric_exists.greptime_timestamp ASC NULLS FIRST [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8184                Filter: metric_exists.greptime_timestamp >= TimestampMillisecond(-999, None) AND metric_exists.greptime_timestamp <= TimestampMillisecond(100000000, None) [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8185                  TableScan: metric_exists [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
8186  SubqueryAlias:  [greptime_timestamp:Timestamp(ms), job:Utf8;N, sum(.value):Float64;N]
8187    Projection: .time AS greptime_timestamp, Utf8(NULL) AS job, sum(.value) [greptime_timestamp:Timestamp(ms), job:Utf8;N, sum(.value):Float64;N]
8188      Sort: .time ASC NULLS LAST [time:Timestamp(ms), sum(.value):Float64;N]
8189        Aggregate: groupBy=[[.time]], aggr=[[sum(.value)]] [time:Timestamp(ms), sum(.value):Float64;N]
8190          EmptyMetric: range=[0..-1], interval=[5000] [time:Timestamp(ms), value:Float64;N]
8191            TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#;
8192
8193        assert_eq!(plan.display_indent_schema().to_string(), expected);
8194    }
8195
8196    #[tokio::test]
8197    async fn test_histogram_quantile_missing_le_column() {
8198        let mut eval_stmt = EvalStmt {
8199            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
8200            start: UNIX_EPOCH,
8201            end: UNIX_EPOCH
8202                .checked_add(Duration::from_secs(100_000))
8203                .unwrap(),
8204            interval: Duration::from_secs(5),
8205            lookback_delta: Duration::from_secs(1),
8206        };
8207
8208        // Test case: histogram_quantile with a table that doesn't have 'le' column
8209        let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#;
8210
8211        let prom_expr = parser::parse(case).unwrap();
8212        eval_stmt.expr = prom_expr;
8213
8214        // Create a table provider with a table that doesn't have 'le' column
8215        let table_provider = build_test_table_provider_with_fields(
8216            &[(
8217                DEFAULT_SCHEMA_NAME.to_string(),
8218                "non_existent_histogram_bucket".to_string(),
8219            )],
8220            &["pod", "instance"], // Note: no 'le' column
8221        )
8222        .await;
8223
8224        // Should return empty result instead of error
8225        let result =
8226            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8227                .await;
8228
8229        // This should succeed now (returning empty result) instead of failing with "Cannot find column le"
8230        assert!(
8231            result.is_ok(),
8232            "Expected successful plan creation with empty result, but got error: {:?}",
8233            result.err()
8234        );
8235
8236        // Verify that the result is an EmptyRelation
8237        let plan = result.unwrap();
8238        match plan {
8239            LogicalPlan::EmptyRelation(_) => {
8240                // This is what we expect
8241            }
8242            _ => panic!("Expected EmptyRelation, but got: {:?}", plan),
8243        }
8244    }
8245}