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::native_histogram::native_histogram_value_type;
26use common_query::prelude::{
27    GREPTIME_TEMPORALITY_DELTA, OTLP_AGGREGATION_TEMPORALITY_LABEL, greptime_native_histogram,
28    greptime_value,
29};
30use common_query::promql_annotations::PromqlAnnotationCollector;
31use datafusion::common::DFSchemaRef;
32use datafusion::datasource::DefaultTableSource;
33use datafusion::functions_aggregate::average::avg_udaf;
34use datafusion::functions_aggregate::count::count_udaf;
35use datafusion::functions_aggregate::expr_fn::first_value;
36use datafusion::functions_aggregate::min_max::{max_udaf, min_udaf};
37use datafusion::functions_aggregate::stddev::stddev_pop_udaf;
38use datafusion::functions_aggregate::sum::sum_udaf;
39use datafusion::functions_aggregate::variance::var_pop_udaf;
40use datafusion::functions_window::row_number::RowNumber;
41use datafusion::logical_expr::expr::{Alias, ScalarFunction, WindowFunction};
42use datafusion::logical_expr::expr_rewriter::normalize_cols;
43use datafusion::logical_expr::{
44    BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator,
45    ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition,
46};
47use datafusion::optimizer::simplify_expressions::ExprSimplifier;
48use datafusion::prelude as df_prelude;
49use datafusion::prelude::{Column, Expr as DfExpr, JoinType};
50use datafusion::scalar::ScalarValue;
51use datafusion::sql::TableReference;
52use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
53use datafusion_common::{DFSchema, NullEquality};
54use datafusion_expr::expr::WindowFunctionParams;
55use datafusion_expr::expr_fn::when;
56use datafusion_expr::simplify::SimplifyContext;
57use datafusion_expr::utils::{conjunction, disjunction};
58use datafusion_expr::{
59    ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit,
60};
61use datafusion_functions::core::coalesce;
62use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit};
63use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType};
64use itertools::Itertools;
65use once_cell::sync::Lazy;
66use promql::extension_plan::{
67    Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, InstantManipulate, Millisecond,
68    RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn,
69    build_special_time_expr,
70};
71use promql::functions::{
72    AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing,
73    IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange,
74    NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg,
75    NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges,
76    NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta,
77    NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction,
78    NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease,
79    NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg,
80    NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile,
81    NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev,
82    NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime,
83    NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime,
84    Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf,
85};
86use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers};
87use promql_parser::parser::token::TokenType;
88use promql_parser::parser::value::ValueType;
89use promql_parser::parser::{
90    AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt, Expr as PromExpr,
91    Function, FunctionArgs as PromFunctionArgs, LabelModifier, MatrixSelector, NumberLiteral,
92    Offset, ParenExpr, StringLiteral, SubqueryExpr, UnaryExpr, VectorMatchCardinality,
93    VectorSelector, token,
94};
95use regex::{self, Regex};
96use snafu::{OptionExt, ResultExt, ensure};
97use store_api::metric_engine_consts::{
98    DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
99    METRIC_ENGINE_NAME, is_metric_engine_internal_column,
100};
101use table::table::adapter::DfTableProviderAdapter;
102
103use crate::parser::{
104    ALIAS_NODE_NAME, ANALYZE_NODE_NAME, ANALYZE_VERBOSE_NODE_NAME, AliasExpr, EXPLAIN_NODE_NAME,
105    EXPLAIN_VERBOSE_NODE_NAME,
106};
107use crate::promql::error::{
108    CatalogSnafu, ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu,
109    ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu,
110    InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu,
111    MultipleMetricMatchersSnafu, MultipleVectorSnafu, NoMetricMatcherSnafu, PromqlPlanNodeSnafu,
112    Result, SameLabelSetSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu,
113    UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu,
114    UnsupportedMatcherOpSnafu, UnsupportedVectorMatchSnafu, ValueNotFoundSnafu,
115    ZeroRangeSelectorSnafu,
116};
117use crate::query_engine::QueryEngineState;
118
119/// `time()` function in PromQL.
120const SPECIAL_TIME_FUNCTION: &str = "time";
121/// `scalar()` function in PromQL.
122const SCALAR_FUNCTION: &str = "scalar";
123/// `absent()` function in PromQL
124const SPECIAL_ABSENT_FUNCTION: &str = "absent";
125/// `histogram_quantile` function in PromQL
126const SPECIAL_HISTOGRAM_QUANTILE: &str = "histogram_quantile";
127/// `histogram_fraction` function in PromQL
128const SPECIAL_HISTOGRAM_FRACTION: &str = "histogram_fraction";
129/// `vector` function in PromQL
130const SPECIAL_VECTOR_FUNCTION: &str = "vector";
131/// `le` column for conventional histogram.
132const LE_COLUMN_NAME: &str = "le";
133
134/// Static regex for validating label names according to Prometheus specification.
135/// Label names must match the regex: [a-zA-Z_][a-zA-Z0-9_]*
136static LABEL_NAME_REGEX: Lazy<Regex> =
137    Lazy::new(|| Regex::new(r"^[a-zA-Z_][a-zA-Z0-9_]*$").unwrap());
138
139const DEFAULT_TIME_INDEX_COLUMN: &str = "time";
140
141/// default value column name for empty metric
142const DEFAULT_FIELD_COLUMN: &str = "value";
143
144/// Special modifier to project field columns under multi-field mode
145const FIELD_COLUMN_MATCHER: &str = "__field__";
146
147/// Special modifier for cross schema query
148const SCHEMA_COLUMN_MATCHER: &str = "__schema__";
149const DB_COLUMN_MATCHER: &str = "__database__";
150
151/// Prefix for generated binary island leaf aliases.
152const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v";
153const OR_FLOAT_FIELD_PREFIX: &str = "__promql_or_float_";
154const OR_HISTOGRAM_FIELD_PREFIX: &str = "__promql_or_histogram_";
155const TIMESTAMP_VALUE_PREFIX: &str = "__promql_timestamp_value_";
156
157/// Threshold for scatter scan mode
158const MAX_SCATTER_POINTS: i64 = 400;
159
160/// Interval 1 hour in millisecond
161const INTERVAL_1H: i64 = 60 * 60 * 1000;
162
163#[derive(Default, Debug, Clone)]
164struct PromPlannerContext {
165    // query parameters
166    start: Millisecond,
167    end: Millisecond,
168    interval: Millisecond,
169    lookback_delta: Millisecond,
170
171    // planner states
172    table_name: Option<String>,
173    time_index_column: Option<String>,
174    field_columns: Vec<String>,
175    tag_columns: Vec<String>,
176    /// Use metric engine internal series identifier column (`__tsid`) as series key.
177    ///
178    /// This is enabled only when the underlying scan can provide `__tsid` (`UInt64`). The planner
179    /// uses it internally (e.g. as the series key for [`SeriesDivide`]) and strips it from the
180    /// final output.
181    use_tsid: bool,
182    /// The matcher for field columns `__field__`.
183    field_column_matcher: Option<Vec<Matcher>>,
184    /// The matcher for selectors (normal matchers).
185    selector_matcher: Vec<Matcher>,
186    schema_name: Option<String>,
187    /// The range in millisecond of range selector. None if there is no range selector.
188    range: Option<Millisecond>,
189}
190
191#[derive(Debug, Clone, PartialEq, Eq, Hash)]
192struct VectorLeafKey {
193    metric_name: String,
194    matchers: Vec<(String, String, String)>,
195    or_matchers: Vec<Vec<(String, String, String)>>,
196    offset_ms: i128,
197    at: String,
198}
199
200#[derive(Debug, Clone)]
201struct IslandLeaf {
202    selector: VectorSelector,
203    display_table: String,
204}
205
206#[derive(Debug, Clone)]
207enum IslandExpr {
208    VectorLeaf(usize),
209    Scalar(DfExpr),
210    Unary {
211        input: Box<IslandExpr>,
212    },
213    Binary {
214        op: TokenType,
215        lhs: Box<IslandExpr>,
216        rhs: Box<IslandExpr>,
217    },
218}
219
220impl IslandExpr {
221    fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option<Self> {
222        if let Some(expr) = PromPlanner::try_build_literal_expr(expr) {
223            return Some(Self::Scalar(expr));
224        }
225
226        match expr {
227            PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env),
228            PromExpr::VectorSelector(selector) => {
229                let leaf = env.intern_leaf(selector)?;
230                Some(Self::VectorLeaf(leaf))
231            }
232            PromExpr::Unary(UnaryExpr { expr }) => {
233                let input = Self::try_new(expr, env)?;
234                Some(Self::Unary {
235                    input: Box::new(input),
236                })
237            }
238            PromExpr::Binary(PromBinaryExpr {
239                lhs,
240                rhs,
241                op,
242                modifier,
243            }) if matches!(
244                op.id(),
245                token::T_ADD
246                    | token::T_SUB
247                    | token::T_MUL
248                    | token::T_DIV
249                    | token::T_MOD
250                    | token::T_POW
251                    | token::T_ATAN2
252            ) && modifier.as_ref().is_none_or(|modifier| {
253                !modifier.return_bool
254                    && modifier.matching.is_none()
255                    && matches!(modifier.card, VectorMatchCardinality::OneToOne)
256                    && modifier.fill_values.lhs.is_none()
257                    && modifier.fill_values.rhs.is_none()
258            }) =>
259            {
260                let lhs = Self::try_new(lhs, env)?;
261                let rhs = Self::try_new(rhs, env)?;
262                Some(Self::Binary {
263                    op: *op,
264                    lhs: Box::new(lhs),
265                    rhs: Box::new(rhs),
266                })
267            }
268            _ => None,
269        }
270    }
271}
272
273#[derive(Debug, Default)]
274struct IslandCollectEnv {
275    leaf_by_key: HashMap<VectorLeafKey, usize>,
276    leaves: Vec<IslandLeaf>,
277    vector_occurrences: usize,
278}
279
280#[derive(Debug)]
281struct PlannedIslandLeaf {
282    plan: LogicalPlan,
283    ctx: PromPlannerContext,
284    alias: TableReference,
285    display_table: String,
286}
287
288#[derive(Debug)]
289struct IslandFieldExprs {
290    exprs: Vec<DfExpr>,
291    names: Vec<String>,
292    scalar: bool,
293}
294
295impl VectorLeafKey {
296    fn from_selector(selector: &VectorSelector) -> Option<Self> {
297        let mut metric_name = selector.name.clone();
298        let mut matchers = Vec::with_capacity(selector.matchers.matchers.len());
299        let matcher_key = |matcher: &Matcher| {
300            (
301                matcher.name.clone(),
302                matcher.op.to_string(),
303                matcher.value.clone(),
304            )
305        };
306
307        for matcher in &selector.matchers.matchers {
308            if matcher.name == METRIC_NAME {
309                if matcher.op != MatchOp::Equal || metric_name.is_some() {
310                    return None;
311                }
312                metric_name = Some(matcher.value.clone());
313            } else {
314                matchers.push(matcher_key(matcher));
315            }
316        }
317        matchers.sort();
318
319        let mut or_matchers = selector
320            .matchers
321            .or_matchers
322            .iter()
323            .map(|group| {
324                let mut group = group.iter().map(matcher_key).collect::<Vec<_>>();
325                group.sort();
326                group
327            })
328            .collect::<Vec<_>>();
329        or_matchers.sort();
330
331        Some(Self {
332            metric_name: metric_name?,
333            matchers,
334            or_matchers,
335            offset_ms: match &selector.offset {
336                Some(Offset::Pos(duration)) => duration.as_millis() as i128,
337                Some(Offset::Neg(duration)) => -(duration.as_millis() as i128),
338                None => 0,
339            },
340            at: format!("{:?}", selector.at),
341        })
342    }
343}
344
345impl IslandCollectEnv {
346    fn intern_leaf(&mut self, selector: &VectorSelector) -> Option<usize> {
347        self.vector_occurrences += 1;
348        let key = VectorLeafKey::from_selector(selector)?;
349        if let Some(id) = self.leaf_by_key.get(&key) {
350            return Some(*id);
351        }
352
353        let id = self.leaves.len();
354        self.leaves.push(IslandLeaf {
355            selector: selector.clone(),
356            display_table: key.metric_name.clone(),
357        });
358        self.leaf_by_key.insert(key, id);
359        Some(id)
360    }
361}
362
363impl PromPlannerContext {
364    fn from_eval_stmt(stmt: &EvalStmt) -> Self {
365        Self {
366            start: stmt.start.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
367            end: stmt.end.duration_since(UNIX_EPOCH).unwrap().as_millis() as _,
368            interval: stmt.interval.as_millis() as _,
369            lookback_delta: stmt.lookback_delta.as_millis() as _,
370            ..Default::default()
371        }
372    }
373
374    /// Reset all planner states
375    fn reset(&mut self) {
376        self.table_name = None;
377        self.time_index_column = None;
378        self.field_columns = vec![];
379        self.tag_columns = vec![];
380        self.use_tsid = false;
381        self.field_column_matcher = None;
382        self.selector_matcher.clear();
383        self.schema_name = None;
384        self.range = None;
385    }
386
387    /// Reset table name and schema to empty
388    fn reset_table_name_and_schema(&mut self) {
389        self.table_name = Some(String::new());
390        self.schema_name = None;
391        self.use_tsid = false;
392    }
393
394    /// Check if `le` is present in tag columns
395    fn has_le_tag(&self) -> bool {
396        self.tag_columns.iter().any(|c| c.eq(&LE_COLUMN_NAME))
397    }
398}
399
400pub struct PromPlanner {
401    table_provider: DfTableSourceProvider,
402    ctx: PromPlannerContext,
403    /// Optional collector passed to native histogram UDFs.
404    promql_annotations: Option<PromqlAnnotationCollector>,
405}
406
407type BinaryFieldPair<'a> = (&'a String, &'a String);
408
409impl PromPlanner {
410    pub async fn stmt_to_plan(
411        table_provider: DfTableSourceProvider,
412        stmt: &EvalStmt,
413        query_engine_state: &QueryEngineState,
414    ) -> Result<LogicalPlan> {
415        Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None).await
416    }
417
418    /// Plans a PromQL statement and passes the optional collector to histogram UDFs.
419    pub async fn stmt_to_plan_with_annotations(
420        table_provider: DfTableSourceProvider,
421        stmt: &EvalStmt,
422        query_engine_state: &QueryEngineState,
423        promql_annotations: Option<PromqlAnnotationCollector>,
424    ) -> Result<LogicalPlan> {
425        let mut planner = Self {
426            table_provider,
427            ctx: PromPlannerContext::from_eval_stmt(stmt),
428            promql_annotations,
429        };
430
431        let plan = planner
432            .prom_expr_to_plan(&stmt.expr, query_engine_state)
433            .await?;
434
435        // Never leak internal series identifier to output.
436        planner.strip_tsid_column(plan)
437    }
438
439    pub async fn prom_expr_to_plan(
440        &mut self,
441        prom_expr: &PromExpr,
442        query_engine_state: &QueryEngineState,
443    ) -> Result<LogicalPlan> {
444        self.prom_expr_to_plan_inner(prom_expr, false, query_engine_state)
445            .await
446    }
447
448    /**
449    Converts a PromQL expression to a logical plan.
450
451    NOTE:
452        The `timestamp_fn` indicates whether the PromQL `timestamp()` function is being evaluated in the current context.
453        If `true`, the planner generates a logical plan that projects the timestamp (time index) column
454        as the value column for each input row, implementing the PromQL `timestamp()` function semantics.
455        If `false`, the planner generates the standard logical plan for the given PromQL expression.
456    */
457    #[async_recursion]
458    async fn prom_expr_to_plan_inner(
459        &mut self,
460        prom_expr: &PromExpr,
461        timestamp_fn: bool,
462        query_engine_state: &QueryEngineState,
463    ) -> Result<LogicalPlan> {
464        let res = match prom_expr {
465            PromExpr::Aggregate(expr) => {
466                self.prom_aggr_expr_to_plan(query_engine_state, expr)
467                    .await?
468            }
469            PromExpr::Unary(expr) => {
470                self.prom_unary_expr_to_plan(query_engine_state, expr)
471                    .await?
472            }
473            PromExpr::Binary(expr) => {
474                self.prom_binary_expr_to_plan(query_engine_state, expr)
475                    .await?
476            }
477            PromExpr::Paren(ParenExpr { expr }) => {
478                self.prom_expr_to_plan_inner(expr, timestamp_fn, query_engine_state)
479                    .await?
480            }
481            PromExpr::Subquery(expr) => {
482                self.prom_subquery_expr_to_plan(query_engine_state, expr)
483                    .await?
484            }
485            PromExpr::NumberLiteral(lit) => self.prom_number_lit_to_plan(lit)?,
486            PromExpr::StringLiteral(lit) => self.prom_string_lit_to_plan(lit)?,
487            PromExpr::VectorSelector(selector) => {
488                self.prom_vector_selector_to_plan(selector, timestamp_fn)
489                    .await?
490            }
491            PromExpr::MatrixSelector(selector) => {
492                self.prom_matrix_selector_to_plan(selector).await?
493            }
494            PromExpr::Call(expr) => {
495                self.prom_call_expr_to_plan(query_engine_state, expr)
496                    .await?
497            }
498            PromExpr::Extension(expr) => {
499                self.prom_ext_expr_to_plan(query_engine_state, expr).await?
500            }
501        };
502
503        Ok(res)
504    }
505
506    async fn prom_subquery_expr_to_plan(
507        &mut self,
508        query_engine_state: &QueryEngineState,
509        subquery_expr: &SubqueryExpr,
510    ) -> Result<LogicalPlan> {
511        let SubqueryExpr {
512            expr, range, step, ..
513        } = subquery_expr;
514
515        let current_interval = self.ctx.interval;
516        if let Some(step) = step {
517            self.ctx.interval = step.as_millis() as _;
518        }
519        let current_start = self.ctx.start;
520        self.ctx.start -= range.as_millis() as i64 - self.ctx.interval;
521        let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
522        self.ctx.interval = current_interval;
523        self.ctx.start = current_start;
524
525        ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
526        let range_ms = range.as_millis() as _;
527        self.ctx.range = Some(range_ms);
528
529        let time_index_column =
530            self.ctx
531                .time_index_column
532                .clone()
533                .with_context(|| TimeIndexNotFoundSnafu {
534                    table: self.ctx.table_name.clone().unwrap_or_default(),
535                })?;
536
537        // `RangeManipulate` assumes each input batch holds exactly one series
538        // (it takes tag column values from row 0 and applies them to every
539        // output row). The inner expression may emit batches that mix series,
540        // so sort by series key + time index and split into per-series batches
541        // with a `SeriesDivide` first.
542        let input_schema = input.schema();
543        let input_has_tsid = input_schema.fields().iter().any(|field| {
544            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
545                && field.data_type() == &ArrowDataType::UInt64
546        });
547        let (series_key_columns, mut sort_exprs) = if input_has_tsid {
548            (
549                vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()],
550                vec![
551                    DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME))
552                        .sort(true, true),
553                ],
554            )
555        } else {
556            // Only use tag columns that survive in the inner plan's schema —
557            // `ctx.tag_columns` can drift from the actual output.
558            let key_columns: Vec<String> = self
559                .ctx
560                .tag_columns
561                .iter()
562                .filter(|name| input_schema.has_column_with_unqualified_name(name))
563                .cloned()
564                .collect();
565            let sort = key_columns
566                .iter()
567                .map(|name| DfExpr::Column(Column::from_name(name)).sort(true, true))
568                .collect::<Vec<_>>();
569            (key_columns, sort)
570        };
571        sort_exprs.push(DfExpr::Column(Column::from_name(&time_index_column)).sort(true, true));
572
573        let sort_plan = LogicalPlanBuilder::from(input)
574            .sort(sort_exprs)
575            .context(DataFusionPlanningSnafu)?
576            .build()
577            .context(DataFusionPlanningSnafu)?;
578        let divide_plan = LogicalPlan::Extension(Extension {
579            node: Arc::new(SeriesDivide::new(
580                series_key_columns,
581                time_index_column.clone(),
582                sort_plan,
583            )),
584        });
585
586        let manipulate = RangeManipulate::new(
587            self.ctx.start,
588            self.ctx.end,
589            self.ctx.interval,
590            range_ms,
591            time_index_column,
592            self.ctx.field_columns.clone(),
593            divide_plan,
594        )
595        .context(DataFusionPlanningSnafu)?;
596
597        Ok(LogicalPlan::Extension(Extension {
598            node: Arc::new(manipulate),
599        }))
600    }
601
602    async fn prom_aggr_expr_to_plan(
603        &mut self,
604        query_engine_state: &QueryEngineState,
605        aggr_expr: &AggregateExpr,
606    ) -> Result<LogicalPlan> {
607        let AggregateExpr {
608            op,
609            expr,
610            modifier,
611            param,
612        } = aggr_expr;
613
614        let mut input = self.prom_expr_to_plan(expr, query_engine_state).await?;
615        let input_has_tsid = input.schema().fields().iter().any(|field| {
616            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
617                && field.data_type() == &ArrowDataType::UInt64
618        });
619
620        // `__tsid` based scan projection may prune tag columns. Ensure tags referenced in
621        // aggregation modifiers (`by`/`without`) are available before planning group keys.
622        let required_group_tags = match modifier {
623            None => BTreeSet::new(),
624            Some(LabelModifier::Include(labels)) => labels
625                .labels
626                .iter()
627                .filter(|label| !is_metric_engine_internal_column(label.as_str()))
628                .cloned()
629                .collect(),
630            Some(LabelModifier::Exclude(labels)) => {
631                let mut all_tags = self.collect_row_key_tag_columns_from_plan(&input)?;
632                for label in &labels.labels {
633                    let _ = all_tags.remove(label);
634                }
635                all_tags
636            }
637        };
638
639        if !required_group_tags.is_empty()
640            && required_group_tags
641                .iter()
642                .any(|tag| Self::find_case_sensitive_column(input.schema(), tag.as_str()).is_none())
643        {
644            input = self.ensure_tag_columns_available(input, &required_group_tags)?;
645            self.refresh_tag_columns_from_schema(input.schema());
646        }
647
648        match (*op).id() {
649            token::T_TOPK | token::T_BOTTOMK => {
650                self.prom_topk_bottomk_to_plan(aggr_expr, input).await
651            }
652            _ => {
653                // When `__tsid` is available, tag columns may have been pruned from the input plan.
654                // For `keep_tsid` decision we should compare against the full row-key label set,
655                // otherwise we may incorrectly treat label-reducing aggregates as preserving labels.
656                let input_tag_columns = if input_has_tsid {
657                    self.collect_row_key_tag_columns_from_plan(&input)?
658                        .into_iter()
659                        .collect::<Vec<_>>()
660                } else {
661                    self.ctx.tag_columns.clone()
662                };
663                // calculate columns to group by
664                // Need to append time index column into group by columns
665                let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
666                let mixed_sample_columns =
667                    Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
668                        .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
669                // Aggregates over native histogram inputs may drop every sample in a group
670                // (e.g. `min` over histogram-only samples, or `sum` over histograms with
671                // incompatible schemas) and leave a NULL-valued group row behind. Compute this
672                // before `create_aggregate_exprs` mutates `ctx.field_columns`.
673                let preserve_any_value = mixed_sample_columns.is_some();
674                let has_native_histogram = preserve_any_value
675                    || self.all_field_columns_are_native_histograms(input.schema());
676                // convert op and value columns to aggregate exprs
677                let (mut aggr_exprs, prev_field_exprs) =
678                    self.create_aggregate_exprs(*op, param, &input)?;
679                let prev_field_exprs =
680                    normalize_cols(prev_field_exprs, &input).context(DataFusionPlanningSnafu)?;
681
682                let keep_tsid = op.id() != token::T_COUNT_VALUES
683                    && input_has_tsid
684                    && input_tag_columns.iter().collect::<HashSet<_>>()
685                        == self.ctx.tag_columns.iter().collect::<HashSet<_>>();
686
687                if keep_tsid {
688                    aggr_exprs.push(
689                        first_value(
690                            DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)),
691                            vec![],
692                        )
693                        .alias(DATA_SCHEMA_TSID_COLUMN_NAME),
694                    );
695                }
696                self.ctx.use_tsid = keep_tsid;
697
698                // create plan
699                let builder = LogicalPlanBuilder::from(input);
700                let builder = if op.id() == token::T_COUNT_VALUES {
701                    let label = Self::get_param_value_as_str(*op, param)?;
702                    // `count_values` must be grouped by fields,
703                    // and project the fields to the new label.
704                    let count_value_exprs = prev_field_exprs.iter().map(|expr| {
705                        match expr {
706                            DfExpr::Column(column) => DfExpr::Column(column.clone()),
707                            _ => DfExpr::Column(Column::from_name(expr.schema_name().to_string())),
708                        }
709                        .alias(label)
710                    });
711                    let aggregate_group_exprs = group_exprs
712                        .iter()
713                        .cloned()
714                        .chain(prev_field_exprs.clone())
715                        .collect::<Vec<_>>();
716                    group_exprs.push(col(label));
717                    let project_fields = self
718                        .create_field_column_exprs()?
719                        .into_iter()
720                        .chain(self.create_tag_column_exprs()?)
721                        .chain(Some(self.create_time_index_column_expr()?))
722                        .chain(count_value_exprs);
723
724                    builder
725                        .aggregate(aggregate_group_exprs, aggr_exprs)
726                        .context(DataFusionPlanningSnafu)?
727                        .project(project_fields)
728                        .context(DataFusionPlanningSnafu)?
729                } else {
730                    builder
731                        .aggregate(group_exprs.clone(), aggr_exprs)
732                        .context(DataFusionPlanningSnafu)?
733                };
734
735                let builder = if let Some((float, histogram)) = mixed_sample_columns {
736                    let builder = match op.id() {
737                        token::T_SUM | token::T_AVG => builder
738                            .filter(self.mixed_aggregate_filter_expr(*op, &float, &histogram)?)
739                            .context(DataFusionPlanningSnafu)?,
740                        token::T_MIN
741                        | token::T_MAX
742                        | token::T_STDDEV
743                        | token::T_STDVAR
744                        | token::T_QUANTILE => builder
745                            .filter(self.mixed_ignored_histogram_filter_expr(*op, &histogram)?)
746                            .context(DataFusionPlanningSnafu)?,
747                        _ => builder,
748                    };
749
750                    match op.id() {
751                        token::T_SUM
752                        | token::T_AVG
753                        | token::T_MIN
754                        | token::T_MAX
755                        | token::T_STDDEV
756                        | token::T_STDVAR
757                        | token::T_QUANTILE => {
758                            let project_fields = self
759                                .create_field_column_exprs()?
760                                .into_iter()
761                                .chain(self.create_tag_column_exprs()?)
762                                .chain(self.ctx.use_tsid.then_some(DfExpr::Column(
763                                    Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME),
764                                )))
765                                .chain(Some(self.create_time_index_column_expr()?));
766                            builder
767                                .project(project_fields)
768                                .context(DataFusionPlanningSnafu)?
769                        }
770                        _ => builder,
771                    }
772                } else {
773                    builder
774                };
775
776                // Drop group rows whose every aggregated sample was discarded (NULL), so that
777                // e.g. `group(min(native_histogram))` doesn't resurrect groups Prometheus
778                // considers unseen. For alternative float/histogram fields keep the row if any
779                // field survived.
780                let builder = if has_native_histogram {
781                    builder
782                        .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
783                        .context(DataFusionPlanningSnafu)?
784                } else {
785                    builder
786                };
787
788                let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false));
789
790                builder
791                    .sort(sort_expr)
792                    .context(DataFusionPlanningSnafu)?
793                    .build()
794                    .context(DataFusionPlanningSnafu)
795            }
796        }
797    }
798
799    /// Create logical plan for PromQL topk and bottomk expr.
800    async fn prom_topk_bottomk_to_plan(
801        &mut self,
802        aggr_expr: &AggregateExpr,
803        input: LogicalPlan,
804    ) -> Result<LogicalPlan> {
805        let AggregateExpr {
806            op,
807            param,
808            modifier,
809            ..
810        } = aggr_expr;
811
812        let input_has_tsid = input.schema().fields().iter().any(|field| {
813            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
814                && field.data_type() == &ArrowDataType::UInt64
815        });
816        self.ctx.use_tsid = input_has_tsid;
817
818        let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?;
819
820        let mut input = input;
821        if let Some((float_column, histogram_column)) =
822            Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
823                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
824        {
825            let drop_histogram = DfExpr::ScalarFunction(ScalarFunction {
826                func: Arc::new(NativeHistogramDrop::bool_false_udf(
827                    format!(
828                        "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
829                        op
830                    ),
831                    self.promql_annotations.clone(),
832                )),
833                args: vec![col(&histogram_column)],
834            });
835            let keep_float = when(col(&histogram_column).is_not_null(), drop_histogram)
836                .otherwise(col(&float_column).is_not_null())
837                .context(DataFusionPlanningSnafu)?;
838            input = LogicalPlanBuilder::from(input)
839                .filter(keep_float)
840                .context(DataFusionPlanningSnafu)?
841                .build()
842                .context(DataFusionPlanningSnafu)?;
843            self.ctx.field_columns = vec![float_column];
844        }
845
846        if self.all_field_columns_are_native_histograms(input.schema()) {
847            let promql_annotations = self.promql_annotations.clone();
848            let input = self.projection_for_each_field_column(input, |col| {
849                Ok(DfExpr::ScalarFunction(ScalarFunction {
850                    func: Arc::new(NativeHistogramDrop::float_null_udf(
851                        format!(
852                            "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
853                            op
854                        ),
855                        promql_annotations.clone(),
856                    )),
857                    args: vec![DfExpr::Column(Column::from_name(col))],
858                }))
859            })?;
860            return LogicalPlanBuilder::from(input)
861                .filter(self.create_empty_values_filter_expr(false)?)
862                .context(DataFusionPlanningSnafu)?
863                .build()
864                .context(DataFusionPlanningSnafu);
865        }
866
867        let val = Self::get_param_as_literal_expr(
868            param.as_deref(),
869            Some(*op),
870            Some(ArrowDataType::Float64),
871        )?;
872
873        // convert op and value columns to window exprs.
874        let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?;
875
876        let rank_columns: Vec<_> = window_exprs
877            .iter()
878            .map(|expr| expr.schema_name().to_string())
879            .collect();
880
881        // Create ranks filter with `Operator::Or`.
882        // Safety: at least one rank column
883        let filter: DfExpr = rank_columns
884            .iter()
885            .fold(None, |expr, rank| {
886                let predicate = DfExpr::BinaryExpr(BinaryExpr {
887                    left: Box::new(col(rank)),
888                    op: Operator::LtEq,
889                    right: Box::new(val.clone()),
890                });
891
892                match expr {
893                    None => Some(predicate),
894                    Some(expr) => Some(DfExpr::BinaryExpr(BinaryExpr {
895                        left: Box::new(expr),
896                        op: Operator::Or,
897                        right: Box::new(predicate),
898                    })),
899                }
900            })
901            .unwrap();
902
903        let rank_columns: Vec<_> = rank_columns.into_iter().map(col).collect();
904
905        let mut new_group_exprs = group_exprs.clone();
906        // Order by ranks
907        new_group_exprs.extend(rank_columns);
908
909        let group_sort_expr = new_group_exprs
910            .into_iter()
911            .map(|expr| expr.sort(true, false));
912
913        let project_fields = self
914            .create_field_column_exprs()?
915            .into_iter()
916            .chain(self.create_tag_column_exprs()?)
917            .chain(
918                self.ctx
919                    .use_tsid
920                    .then_some(DfExpr::Column(Column::from_name(
921                        DATA_SCHEMA_TSID_COLUMN_NAME,
922                    ))),
923            )
924            .chain(Some(self.create_time_index_column_expr()?));
925
926        LogicalPlanBuilder::from(input)
927            .window(window_exprs)
928            .context(DataFusionPlanningSnafu)?
929            .filter(filter)
930            .context(DataFusionPlanningSnafu)?
931            .sort(group_sort_expr)
932            .context(DataFusionPlanningSnafu)?
933            .project(project_fields)
934            .context(DataFusionPlanningSnafu)?
935            .build()
936            .context(DataFusionPlanningSnafu)
937    }
938
939    async fn prom_unary_expr_to_plan(
940        &mut self,
941        query_engine_state: &QueryEngineState,
942        unary_expr: &UnaryExpr,
943    ) -> Result<LogicalPlan> {
944        let UnaryExpr { expr } = unary_expr;
945        // Unary Expr in PromQL implys the `-` operator
946        let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
947        self.negate_field_columns(input)
948    }
949
950    fn negate_field_columns(&mut self, input: LogicalPlan) -> Result<LogicalPlan> {
951        let input_schema = input.schema().clone();
952        self.projection_for_each_field_column(input, |col| {
953            if Self::field_column_is_native_histogram(&input_schema, col) {
954                Ok(DfExpr::ScalarFunction(ScalarFunction {
955                    func: Arc::new(NativeHistogramNeg::scalar_udf()),
956                    args: vec![DfExpr::Column(col.into())],
957                }))
958            } else {
959                Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into()))))
960            }
961        })
962    }
963
964    async fn try_plan_binary_island(
965        &mut self,
966        binary_expr: &PromBinaryExpr,
967    ) -> Result<Option<LogicalPlan>> {
968        let original_ctx = self.ctx.clone();
969        let mut collect_env = IslandCollectEnv::default();
970        let Some(island_expr) =
971            IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env)
972        else {
973            return Ok(None);
974        };
975
976        if collect_env.leaves.is_empty()
977            || collect_env.vector_occurrences <= collect_env.leaves.len()
978        {
979            return Ok(None);
980        }
981
982        let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len());
983        for (idx, leaf) in collect_env.leaves.iter().enumerate() {
984            let plan = self
985                .prom_vector_selector_to_plan(&leaf.selector, false)
986                .await?;
987            let ctx = self.ctx.clone();
988            let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}"));
989            let plan = LogicalPlanBuilder::from(plan)
990                .alias(alias.clone())
991                .context(DataFusionPlanningSnafu)?
992                .build()
993                .context(DataFusionPlanningSnafu)?;
994            planned_leaves.push(PlannedIslandLeaf {
995                plan,
996                ctx,
997                alias,
998                display_table: leaf.display_table.clone(),
999            });
1000        }
1001
1002        if planned_leaves.iter().any(|leaf| {
1003            Self::field_columns_contain_native_histogram(
1004                leaf.plan.schema(),
1005                &leaf.ctx.field_columns,
1006            )
1007        }) {
1008            self.ctx = original_ctx;
1009            return Ok(None);
1010        }
1011
1012        if !Self::binary_island_join_contexts_supported(&planned_leaves) {
1013            self.ctx = original_ctx;
1014            return Ok(None);
1015        }
1016
1017        let mut input = planned_leaves[0].plan.clone();
1018        for right_idx in 1..planned_leaves.len() {
1019            input = self.join_binary_island_leaf(
1020                input,
1021                &planned_leaves[0],
1022                &planned_leaves[right_idx],
1023            )?;
1024        }
1025
1026        let field_exprs =
1027            Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?;
1028        if field_exprs.scalar || field_exprs.exprs.is_empty() {
1029            self.ctx = original_ctx;
1030            return Ok(None);
1031        }
1032
1033        let plan = self.project_binary_island(
1034            input,
1035            &planned_leaves[0].alias,
1036            &planned_leaves[0].ctx,
1037            field_exprs,
1038        )?;
1039        Ok(Some(plan))
1040    }
1041
1042    fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool {
1043        if leaves
1044            .iter()
1045            .any(|leaf| leaf.ctx.time_index_column.is_none())
1046        {
1047            return false;
1048        }
1049
1050        if leaves.len() <= 1 {
1051            return true;
1052        }
1053
1054        let first_tags = leaves[0].ctx.tag_columns.iter().collect::<BTreeSet<_>>();
1055
1056        leaves.iter().skip(1).all(|leaf| {
1057            (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan))
1058                || leaf.ctx.tag_columns.iter().collect::<BTreeSet<_>>() == first_tags
1059        })
1060    }
1061
1062    fn join_binary_island_leaf(
1063        &self,
1064        left: LogicalPlan,
1065        first_leaf: &PlannedIslandLeaf,
1066        right_leaf: &PlannedIslandLeaf,
1067    ) -> Result<LogicalPlan> {
1068        let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty()
1069            || right_leaf.ctx.tag_columns.is_empty())
1070            && !first_leaf
1071                .ctx
1072                .tag_columns
1073                .iter()
1074                .chain(&right_leaf.ctx.tag_columns)
1075                .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL);
1076        let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns(
1077            left.schema(),
1078            right_leaf.plan.schema(),
1079            &first_leaf.ctx,
1080            &right_leaf.ctx,
1081            only_join_time_index,
1082            &None,
1083        )?;
1084
1085        if let (Some(left_time_index_column), Some(right_time_index_column)) = (
1086            first_leaf.ctx.time_index_column.clone(),
1087            right_leaf.ctx.time_index_column.clone(),
1088        ) {
1089            left_keys.insert(left_time_index_column);
1090            right_keys.insert(right_time_index_column);
1091        }
1092
1093        LogicalPlanBuilder::from(left)
1094            .join_detailed(
1095                right_leaf.plan.clone(),
1096                JoinType::Inner,
1097                (
1098                    left_keys
1099                        .into_iter()
1100                        .map(|name| Column::new(Some(first_leaf.alias.clone()), name))
1101                        .collect::<Vec<_>>(),
1102                    right_keys
1103                        .into_iter()
1104                        .map(|name| Column::new(Some(right_leaf.alias.clone()), name))
1105                        .collect::<Vec<_>>(),
1106                ),
1107                force_empty_join.then_some(lit(false)),
1108                NullEquality::NullEqualsNull,
1109            )
1110            .context(DataFusionPlanningSnafu)?
1111            .build()
1112            .context(DataFusionPlanningSnafu)
1113    }
1114
1115    fn build_binary_island_field_exprs(
1116        expr: &IslandExpr,
1117        leaves: &[PlannedIslandLeaf],
1118        schema: &DFSchemaRef,
1119    ) -> Result<IslandFieldExprs> {
1120        match expr {
1121            IslandExpr::VectorLeaf(id) => {
1122                let leaf = &leaves[*id];
1123                let exprs = leaf
1124                    .ctx
1125                    .field_columns
1126                    .iter()
1127                    .map(|field| {
1128                        schema
1129                            .qualified_field_with_name(Some(&leaf.alias), field)
1130                            .context(DataFusionPlanningSnafu)
1131                            .map(|field| DfExpr::Column(field.into()))
1132                    })
1133                    .collect::<Result<Vec<_>>>()?;
1134                let names = leaf
1135                    .ctx
1136                    .field_columns
1137                    .iter()
1138                    .map(|field| format!("{}.{}", leaf.display_table, field))
1139                    .collect();
1140                Ok(IslandFieldExprs {
1141                    exprs,
1142                    names,
1143                    scalar: false,
1144                })
1145            }
1146            IslandExpr::Scalar(expr) => Ok(IslandFieldExprs {
1147                exprs: vec![expr.clone()],
1148                names: vec![expr.schema_name().to_string()],
1149                scalar: true,
1150            }),
1151            IslandExpr::Unary { input } => {
1152                let input = Self::build_binary_island_field_exprs(input, leaves, schema)?;
1153                let mut exprs = Vec::with_capacity(input.exprs.len());
1154                let mut names = Vec::with_capacity(input.names.len());
1155                for (expr, name) in input.exprs.into_iter().zip(input.names) {
1156                    exprs.push(DfExpr::Negative(Box::new(expr)));
1157                    names.push(format!("-{name}"));
1158                }
1159                Ok(IslandFieldExprs {
1160                    exprs,
1161                    names,
1162                    scalar: input.scalar,
1163                })
1164            }
1165            IslandExpr::Binary { op, lhs, rhs } => {
1166                let same_leaf = match (&**lhs, &**rhs) {
1167                    (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right))
1168                        if left == right =>
1169                    {
1170                        Some(*left)
1171                    }
1172                    _ => None,
1173                };
1174                let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?;
1175                let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?;
1176                let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1177                let scalar = lhs.scalar && rhs.scalar;
1178                let op = op.to_string();
1179
1180                let (exprs, names) = match (lhs.scalar, rhs.scalar) {
1181                    (true, true) => {
1182                        let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?;
1183                        let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]);
1184                        (vec![expr], vec![name])
1185                    }
1186                    (true, false) => {
1187                        let mut exprs = Vec::with_capacity(rhs.exprs.len());
1188                        let mut names = Vec::with_capacity(rhs.names.len());
1189                        for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) {
1190                            exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?);
1191                            names.push(format!("{} {op} {rhs_name}", lhs.names[0]));
1192                        }
1193                        (exprs, names)
1194                    }
1195                    (false, true) => {
1196                        let mut exprs = Vec::with_capacity(lhs.exprs.len());
1197                        let mut names = Vec::with_capacity(lhs.names.len());
1198                        for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) {
1199                            exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?);
1200                            names.push(format!("{lhs_name} {op} {}", rhs.names[0]));
1201                        }
1202                        (exprs, names)
1203                    }
1204                    (false, false) => {
1205                        let mut exprs = Vec::new();
1206                        let mut names = Vec::new();
1207                        for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs
1208                            .exprs
1209                            .into_iter()
1210                            .zip(rhs.exprs)
1211                            .zip(lhs.names.into_iter().zip(rhs.names))
1212                            .enumerate()
1213                        {
1214                            if let Some(leaf) = same_leaf {
1215                                let field = leaves[leaf]
1216                                    .ctx
1217                                    .field_columns
1218                                    .get(idx)
1219                                    .cloned()
1220                                    .unwrap_or_else(|| lhs_name.clone());
1221                                lhs_name = format!("lhs.{field}");
1222                                rhs_name = format!("rhs.{field}");
1223                            }
1224                            exprs.push(expr_builder(lhs_expr, rhs_expr)?);
1225                            names.push(format!("{lhs_name} {op} {rhs_name}"));
1226                        }
1227                        (exprs, names)
1228                    }
1229                };
1230
1231                Ok(IslandFieldExprs {
1232                    exprs,
1233                    names,
1234                    scalar,
1235                })
1236            }
1237        }
1238    }
1239
1240    fn project_binary_island(
1241        &mut self,
1242        input: LogicalPlan,
1243        base_alias: &TableReference,
1244        base_ctx: &PromPlannerContext,
1245        field_exprs: IslandFieldExprs,
1246    ) -> Result<LogicalPlan> {
1247        self.ctx = base_ctx.clone();
1248
1249        let schema = input.schema();
1250        let non_field_exprs = base_ctx
1251            .tag_columns
1252            .iter()
1253            .chain(base_ctx.time_index_column.iter())
1254            .map(|column| {
1255                schema
1256                    .qualified_field_with_name(Some(base_alias), column)
1257                    .context(DataFusionPlanningSnafu)
1258                    .map(|field| DfExpr::Column(field.into()))
1259            });
1260        let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid)
1261            .into_iter()
1262            .map(Ok);
1263
1264        self.ctx.field_columns = field_exprs.names;
1265        let field_exprs = field_exprs
1266            .exprs
1267            .into_iter()
1268            .zip(self.ctx.field_columns.iter())
1269            .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
1270
1271        let project_exprs = non_field_exprs
1272            .chain(tsid_expr)
1273            .chain(field_exprs)
1274            .collect::<Result<Vec<_>>>()?;
1275
1276        let plan = LogicalPlanBuilder::from(input)
1277            .project(project_exprs)
1278            .context(DataFusionPlanningSnafu)?
1279            .build()
1280            .context(DataFusionPlanningSnafu)?;
1281
1282        self.ctx.table_name = None;
1283        self.ctx.schema_name = None;
1284
1285        Ok(plan)
1286    }
1287
1288    async fn prom_binary_expr_to_plan(
1289        &mut self,
1290        query_engine_state: &QueryEngineState,
1291        binary_expr: &PromBinaryExpr,
1292    ) -> Result<LogicalPlan> {
1293        // promql-parser accepts fill modifiers, but Greptime does not implement the
1294        // required outer joins and missing-value substitution. Reject them before the
1295        // binary-island fast path so they cannot silently behave like normal inner joins.
1296        if let Some(modifier) = &binary_expr.modifier {
1297            ensure!(
1298                modifier.fill_values.lhs.is_none() && modifier.fill_values.rhs.is_none(),
1299                UnsupportedExprSnafu {
1300                    name: "PromQL fill modifiers"
1301                }
1302            );
1303        }
1304
1305        if let Some(plan) = self.try_plan_binary_island(binary_expr).await? {
1306            return Ok(plan);
1307        }
1308
1309        let PromBinaryExpr {
1310            lhs,
1311            rhs,
1312            op,
1313            modifier,
1314        } = binary_expr;
1315
1316        // if set to true, comparison operator will return 0/1 (for true/false) instead of
1317        // filter on the result column
1318        let should_return_bool = if let Some(m) = modifier {
1319            m.return_bool
1320        } else {
1321            false
1322        };
1323        let is_comparison_op = Self::is_token_a_comparison_op(*op);
1324
1325        // we should build a filter plan here if the op is comparison op and need not
1326        // to return 0/1. Otherwise, we should build a projection plan
1327        match (
1328            Self::try_build_literal_expr(lhs),
1329            Self::try_build_literal_expr(rhs),
1330        ) {
1331            (Some(lhs), Some(rhs)) => {
1332                self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1333                self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1334                self.ctx.reset_table_name_and_schema();
1335                let field_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1336                let mut field_expr = field_expr_builder(lhs, rhs)?;
1337
1338                if is_comparison_op && should_return_bool {
1339                    field_expr = DfExpr::Cast(Cast {
1340                        expr: Box::new(field_expr),
1341                        data_type: ArrowDataType::Float64,
1342                    });
1343                }
1344
1345                Ok(LogicalPlan::Extension(Extension {
1346                    node: Arc::new(
1347                        EmptyMetric::new(
1348                            self.ctx.start,
1349                            self.ctx.end,
1350                            self.ctx.interval,
1351                            SPECIAL_TIME_FUNCTION.to_string(),
1352                            DEFAULT_FIELD_COLUMN.to_string(),
1353                            Some(field_expr),
1354                        )
1355                        .context(DataFusionPlanningSnafu)?,
1356                    ),
1357                }))
1358            }
1359            // lhs is a literal, rhs is a column
1360            (Some(mut expr), None) => {
1361                let input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1362                // check if the literal is a special time expr
1363                if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) {
1364                    expr = time_expr
1365                }
1366                let input_schema = input.schema().clone();
1367                let preserve_any_value = Self::field_columns_are_alternative_samples(
1368                    &input_schema,
1369                    &self.ctx.field_columns,
1370                );
1371                let has_native_histogram = Self::field_columns_contain_native_histogram(
1372                    &input_schema,
1373                    &self.ctx.field_columns,
1374                );
1375                let retain_field_columns = self
1376                    .ctx
1377                    .field_columns
1378                    .iter()
1379                    .map(|col| {
1380                        Self::binary_result_is_histogram(
1381                            *op,
1382                            false,
1383                            Self::field_column_is_native_histogram(&input_schema, col),
1384                        )
1385                        .is_some()
1386                    })
1387                    .collect();
1388                let promql_annotations = self.promql_annotations.clone();
1389                let bin_expr_builder = |col: &String| {
1390                    let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1391                    let rhs_is_histogram =
1392                        Self::field_column_is_native_histogram(&input_schema, col);
1393                    let rhs = DfExpr::Column(col.into());
1394                    let mut binary_expr = match Self::native_histogram_binary_expr(
1395                        *op,
1396                        expr.clone(),
1397                        false,
1398                        rhs.clone(),
1399                        rhs_is_histogram,
1400                        is_comparison_op && !should_return_bool,
1401                        promql_annotations.clone(),
1402                    )? {
1403                        Some(expr) => expr,
1404                        None => binary_expr_builder(expr.clone(), rhs)?,
1405                    };
1406
1407                    if is_comparison_op && should_return_bool {
1408                        binary_expr = DfExpr::Cast(Cast {
1409                            expr: Box::new(binary_expr),
1410                            data_type: ArrowDataType::Float64,
1411                        });
1412                    }
1413                    Ok(binary_expr)
1414                };
1415                if is_comparison_op && !should_return_bool {
1416                    self.filter_on_field_column(input, bin_expr_builder)
1417                } else {
1418                    let projected =
1419                        self.projection_for_each_field_column(input, bin_expr_builder)?;
1420                    self.filter_binary_projection(
1421                        projected,
1422                        has_native_histogram,
1423                        preserve_any_value,
1424                        retain_field_columns,
1425                    )
1426                }
1427            }
1428            // lhs is a column, rhs is a literal
1429            (None, Some(mut expr)) => {
1430                let input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1431                // check if the literal is a special time expr
1432                if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) {
1433                    expr = time_expr
1434                }
1435                let input_schema = input.schema().clone();
1436                let preserve_any_value = Self::field_columns_are_alternative_samples(
1437                    &input_schema,
1438                    &self.ctx.field_columns,
1439                );
1440                let has_native_histogram = Self::field_columns_contain_native_histogram(
1441                    &input_schema,
1442                    &self.ctx.field_columns,
1443                );
1444                let retain_field_columns = self
1445                    .ctx
1446                    .field_columns
1447                    .iter()
1448                    .map(|col| {
1449                        Self::binary_result_is_histogram(
1450                            *op,
1451                            Self::field_column_is_native_histogram(&input_schema, col),
1452                            false,
1453                        )
1454                        .is_some()
1455                    })
1456                    .collect();
1457                let promql_annotations = self.promql_annotations.clone();
1458                let bin_expr_builder = |col: &String| {
1459                    let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1460                    let lhs_is_histogram =
1461                        Self::field_column_is_native_histogram(&input_schema, col);
1462                    let lhs = DfExpr::Column(col.into());
1463                    let mut binary_expr = match Self::native_histogram_binary_expr(
1464                        *op,
1465                        lhs.clone(),
1466                        lhs_is_histogram,
1467                        expr.clone(),
1468                        false,
1469                        is_comparison_op && !should_return_bool,
1470                        promql_annotations.clone(),
1471                    )? {
1472                        Some(expr) => expr,
1473                        None => binary_expr_builder(lhs, expr.clone())?,
1474                    };
1475
1476                    if is_comparison_op && should_return_bool {
1477                        binary_expr = DfExpr::Cast(Cast {
1478                            expr: Box::new(binary_expr),
1479                            data_type: ArrowDataType::Float64,
1480                        });
1481                    }
1482                    Ok(binary_expr)
1483                };
1484                if is_comparison_op && !should_return_bool {
1485                    self.filter_on_field_column(input, bin_expr_builder)
1486                } else {
1487                    let projected =
1488                        self.projection_for_each_field_column(input, bin_expr_builder)?;
1489                    self.filter_binary_projection(
1490                        projected,
1491                        has_native_histogram,
1492                        preserve_any_value,
1493                        retain_field_columns,
1494                    )
1495                }
1496            }
1497            // both are columns. join them on time index
1498            (None, None) => {
1499                let left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1500                let left_field_columns = self.ctx.field_columns.clone();
1501                let left_time_index_column = self.ctx.time_index_column.clone();
1502                let mut left_table_ref = self
1503                    .table_ref()
1504                    .unwrap_or_else(|_| TableReference::bare(""));
1505                let left_context = self.ctx.clone();
1506
1507                let right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1508                let right_field_columns = self.ctx.field_columns.clone();
1509                let right_time_index_column = self.ctx.time_index_column.clone();
1510                let mut right_table_ref = self
1511                    .table_ref()
1512                    .unwrap_or_else(|_| TableReference::bare(""));
1513                let right_context = self.ctx.clone();
1514                let left_is_empty_metric = Self::is_empty_metric(&left_input);
1515                let right_is_empty_metric = Self::is_empty_metric(&right_input);
1516
1517                // TODO(ruihang): avoid join if left and right are the same table
1518
1519                // set op has "special" join semantics
1520                if Self::is_token_a_set_op(*op) {
1521                    return self.set_op_on_non_field_columns(
1522                        left_input,
1523                        right_input,
1524                        left_context,
1525                        right_context,
1526                        *op,
1527                        modifier,
1528                    );
1529                }
1530
1531                let has_native_histogram = Self::field_columns_contain_native_histogram(
1532                    left_input.schema(),
1533                    &left_field_columns,
1534                ) || Self::field_columns_contain_native_histogram(
1535                    right_input.schema(),
1536                    &right_field_columns,
1537                );
1538
1539                // normal join
1540                if left_table_ref == right_table_ref {
1541                    // rename table references to avoid ambiguity
1542                    left_table_ref = TableReference::bare("lhs");
1543                    right_table_ref = TableReference::bare("rhs");
1544                    // `self.ctx` have ctx in right plan, if right plan have no tag,
1545                    // we use left plan ctx as the ctx for subsequent calculations,
1546                    // to avoid case like `host + scalar(...)`
1547                    // we need preserve tag column on `host` table in subsequent projection,
1548                    // which only show in left plan ctx.
1549                    if self.ctx.tag_columns.is_empty() {
1550                        self.ctx = left_context.clone();
1551                        self.ctx.table_name = Some("lhs".to_string());
1552                    } else {
1553                        self.ctx.table_name = Some("rhs".to_string());
1554                    }
1555                } else if right_is_empty_metric && !left_is_empty_metric {
1556                    self.ctx = left_context.clone();
1557                }
1558                // Computed scalars reach this join path instead of the literal projection paths.
1559                // Broadcast them for arithmetic in the same way as literal scalars.
1560                let broadcast_scalar = !is_comparison_op;
1561                let (field_groups, invalid_field_pairs) = Self::align_binary_field_columns(
1562                    left_input.schema(),
1563                    right_input.schema(),
1564                    &left_field_columns,
1565                    &right_field_columns,
1566                    *op,
1567                    broadcast_scalar && lhs.value_type() == ValueType::Scalar,
1568                    broadcast_scalar && rhs.value_type() == ValueType::Scalar,
1569                );
1570                let left_aligned_field_columns = field_groups
1571                    .iter()
1572                    .flat_map(|(_, pairs)| {
1573                        pairs
1574                            .iter()
1575                            .map(|(left_col_name, _)| (*left_col_name).clone())
1576                    })
1577                    .collect::<Vec<_>>();
1578                let right_aligned_field_columns = field_groups
1579                    .iter()
1580                    .flat_map(|(_, pairs)| {
1581                        pairs
1582                            .iter()
1583                            .map(|(_, right_col_name)| (*right_col_name).clone())
1584                    })
1585                    .collect::<Vec<_>>();
1586                // Regular multi-field vectors combine their shared prefix. Alternative
1587                // float/histogram lanes instead align by valid PromQL sample combinations.
1588                self.ctx.field_columns = field_groups
1589                    .iter()
1590                    .map(|(output, _)| output.clone())
1591                    .collect();
1592                let mut field_groups = field_groups.into_iter();
1593                // `vector()` uses EmptyMetric and keeps GreptimeDB's timestamp broadcast.
1594                let has_empty_metric_operand = left_is_empty_metric || right_is_empty_metric;
1595
1596                let join_plan = self.join_on_non_field_columns(
1597                    left_input,
1598                    right_input,
1599                    left_table_ref.clone(),
1600                    right_table_ref.clone(),
1601                    left_time_index_column,
1602                    right_time_index_column,
1603                    lhs.value_type() == ValueType::Scalar
1604                        || rhs.value_type() == ValueType::Scalar
1605                        || has_empty_metric_operand
1606                        || ((left_context.tag_columns.is_empty()
1607                            || right_context.tag_columns.is_empty())
1608                            && !left_context
1609                                .tag_columns
1610                                .iter()
1611                                .chain(&right_context.tag_columns)
1612                                .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)),
1613                    modifier,
1614                    &left_context,
1615                    &right_context,
1616                )?;
1617                let join_plan_schema = join_plan.schema().clone();
1618                let promql_annotations = self.promql_annotations.clone();
1619                // These predicates always pass; they only evaluate otherwise-discarded pairs
1620                // while collecting annotations.
1621                let invalid_pair_predicates = invalid_field_pairs
1622                    .into_iter()
1623                    .filter(|_| promql_annotations.is_some())
1624                    .map(|(left_col_name, right_col_name)| {
1625                        let left_field = join_plan_schema
1626                            .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1627                            .context(DataFusionPlanningSnafu)?;
1628                        let right_field = join_plan_schema
1629                            .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1630                            .context(DataFusionPlanningSnafu)?;
1631                        let left_is_histogram =
1632                            left_field.1.data_type() == &Self::native_histogram_arrow_type();
1633                        let right_is_histogram =
1634                            right_field.1.data_type() == &Self::native_histogram_arrow_type();
1635                        let drop_expr = Self::native_histogram_binary_expr(
1636                            *op,
1637                            DfExpr::Column(left_field.into()),
1638                            left_is_histogram,
1639                            DfExpr::Column(right_field.into()),
1640                            right_is_histogram,
1641                            true,
1642                            promql_annotations.clone(),
1643                        )?
1644                        .with_context(|| UnexpectedPlanExprSnafu {
1645                            desc: "invalid native histogram pair produced no drop expression",
1646                        })?;
1647                        Ok(DfExpr::Not(Box::new(drop_expr)))
1648                    })
1649                    .collect::<Result<Vec<_>>>()?;
1650                let join_plan = if let Some(predicate) = conjunction(invalid_pair_predicates) {
1651                    LogicalPlanBuilder::from(join_plan)
1652                        .filter(predicate)
1653                        .context(DataFusionPlanningSnafu)?
1654                        .build()
1655                        .context(DataFusionPlanningSnafu)?
1656                } else {
1657                    join_plan
1658                };
1659
1660                let bin_expr_builder = |_: &String| {
1661                    let (_, field_pairs) =
1662                        field_groups
1663                            .next()
1664                            .with_context(|| UnexpectedPlanExprSnafu {
1665                                desc: "missing binary field group",
1666                            })?;
1667                    let binary_exprs = field_pairs
1668                        .into_iter()
1669                        .map(|(left_col_name, right_col_name)| {
1670                            let left_field = join_plan_schema
1671                                .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1672                                .context(DataFusionPlanningSnafu)?;
1673                            let right_field = join_plan_schema
1674                                .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1675                                .context(DataFusionPlanningSnafu)?;
1676                            let left_is_histogram =
1677                                left_field.1.data_type() == &Self::native_histogram_arrow_type();
1678                            let right_is_histogram =
1679                                right_field.1.data_type() == &Self::native_histogram_arrow_type();
1680                            let left_col = left_field.into();
1681                            let right_col = right_field.into();
1682
1683                            let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1684                            let lhs = DfExpr::Column(left_col);
1685                            let rhs = DfExpr::Column(right_col);
1686                            let mut binary_expr = match Self::native_histogram_binary_expr(
1687                                *op,
1688                                lhs.clone(),
1689                                left_is_histogram,
1690                                rhs.clone(),
1691                                right_is_histogram,
1692                                is_comparison_op && !should_return_bool,
1693                                promql_annotations.clone(),
1694                            )? {
1695                                Some(expr) => expr,
1696                                None => binary_expr_builder(lhs, rhs)?,
1697                            };
1698                            if is_comparison_op && should_return_bool {
1699                                binary_expr = DfExpr::Cast(Cast {
1700                                    expr: Box::new(binary_expr),
1701                                    data_type: ArrowDataType::Float64,
1702                                });
1703                            }
1704                            Ok(binary_expr)
1705                        })
1706                        .collect::<Result<Vec<_>>>()?;
1707                    if let [binary_expr] = binary_exprs.as_slice() {
1708                        Ok(binary_expr.clone())
1709                    } else {
1710                        Ok(DfExpr::ScalarFunction(ScalarFunction {
1711                            func: coalesce(),
1712                            args: binary_exprs,
1713                        }))
1714                    }
1715                };
1716                if is_comparison_op && !should_return_bool {
1717                    // PromQL comparison operators without `bool` are filters:
1718                    //   - keep the instant-vector side sample values
1719                    //   - drop samples where the comparison is false
1720                    //
1721                    // So we filter on the join result and then project only the side that should
1722                    // be preserved according to PromQL semantics.
1723                    let filtered = self.filter_on_field_column(join_plan, bin_expr_builder)?;
1724                    let (project_table_ref, mut project_context, project_field_columns) =
1725                        match (lhs.value_type(), rhs.value_type()) {
1726                            (ValueType::Scalar, ValueType::Vector) => (
1727                                &right_table_ref,
1728                                right_context.clone(),
1729                                right_aligned_field_columns,
1730                            ),
1731                            _ => (
1732                                &left_table_ref,
1733                                left_context.clone(),
1734                                left_aligned_field_columns,
1735                            ),
1736                        };
1737                    project_context.field_columns = project_field_columns;
1738                    self.project_binary_join_side(filtered, project_table_ref, &project_context)
1739                } else {
1740                    let projected =
1741                        self.projection_for_each_field_column(join_plan, bin_expr_builder)?;
1742                    let preserve_any_value = Self::field_columns_are_alternative_samples(
1743                        projected.schema(),
1744                        &self.ctx.field_columns,
1745                    );
1746                    let retain_field_columns = vec![true; self.ctx.field_columns.len()];
1747                    self.filter_binary_projection(
1748                        projected,
1749                        has_native_histogram,
1750                        preserve_any_value,
1751                        retain_field_columns,
1752                    )
1753                }
1754            }
1755        }
1756    }
1757
1758    fn filter_binary_projection(
1759        &mut self,
1760        input: LogicalPlan,
1761        has_native_histogram: bool,
1762        preserve_any_value: bool,
1763        retain_field_columns: Vec<bool>,
1764    ) -> Result<LogicalPlan> {
1765        if !has_native_histogram {
1766            return Ok(input);
1767        }
1768
1769        ensure!(
1770            retain_field_columns.len() == self.ctx.field_columns.len(),
1771            UnexpectedPlanExprSnafu {
1772                desc: "binary output field count changed unexpectedly",
1773            }
1774        );
1775
1776        let filtered = LogicalPlanBuilder::from(input)
1777            .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
1778            .context(DataFusionPlanningSnafu)?
1779            .build()
1780            .context(DataFusionPlanningSnafu)?;
1781        if retain_field_columns.iter().all(|retain| *retain) {
1782            return Ok(filtered);
1783        }
1784
1785        let retained = self
1786            .ctx
1787            .field_columns
1788            .iter()
1789            .zip(retain_field_columns)
1790            .filter(|(_, retain)| *retain)
1791            .map(|(field, _)| field.clone())
1792            .collect::<Vec<_>>();
1793        if retained.is_empty() {
1794            return Ok(filtered);
1795        }
1796        self.ctx.field_columns = retained;
1797
1798        let mut output_columns = self
1799            .ctx
1800            .field_columns
1801            .iter()
1802            .chain(&self.ctx.tag_columns)
1803            .cloned()
1804            .collect::<HashSet<_>>();
1805        output_columns.extend(self.ctx.time_index_column.iter().cloned());
1806        if self.ctx.use_tsid {
1807            output_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
1808        }
1809        let project_exprs = filtered
1810            .schema()
1811            .iter()
1812            .filter(|(_, field)| output_columns.contains(field.name()))
1813            .map(|(qualifier, field)| {
1814                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1815            })
1816            .collect::<Vec<_>>();
1817        LogicalPlanBuilder::from(filtered)
1818            .project(project_exprs)
1819            .context(DataFusionPlanningSnafu)?
1820            .build()
1821            .context(DataFusionPlanningSnafu)
1822    }
1823
1824    fn project_binary_join_side(
1825        &mut self,
1826        input: LogicalPlan,
1827        table_ref: &TableReference,
1828        context: &PromPlannerContext,
1829    ) -> Result<LogicalPlan> {
1830        let schema = input.schema();
1831
1832        let mut project_exprs =
1833            Vec::with_capacity(context.tag_columns.len() + context.field_columns.len() + 2);
1834
1835        // Project time index from the chosen side.
1836        if let Some(time_index_column) = &context.time_index_column {
1837            let time_index_col = schema
1838                .qualified_field_with_name(Some(table_ref), time_index_column)
1839                .context(DataFusionPlanningSnafu)?
1840                .into();
1841            project_exprs.push(DfExpr::Column(time_index_col));
1842        }
1843
1844        // Project field columns from the chosen side.
1845        for field_column in &context.field_columns {
1846            let field_col = schema
1847                .qualified_field_with_name(Some(table_ref), field_column)
1848                .context(DataFusionPlanningSnafu)?
1849                .into();
1850            project_exprs.push(DfExpr::Column(field_col));
1851        }
1852
1853        // Project tag columns from the chosen side.
1854        for tag_column in &context.tag_columns {
1855            let tag_col = schema
1856                .qualified_field_with_name(Some(table_ref), tag_column)
1857                .context(DataFusionPlanningSnafu)?
1858                .into();
1859            project_exprs.push(DfExpr::Column(tag_col));
1860        }
1861
1862        // Preserve `__tsid` if present, so it can still be used internally downstream. It's
1863        // stripped from the final output anyway.
1864        if let Some(tsid_col) =
1865            Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid)
1866        {
1867            project_exprs.push(tsid_col);
1868        }
1869
1870        let plan = LogicalPlanBuilder::from(input)
1871            .project(project_exprs)
1872            .context(DataFusionPlanningSnafu)?
1873            .build()
1874            .context(DataFusionPlanningSnafu)?;
1875
1876        // Update context to reflect the projected schema. Don't keep a table qualifier since
1877        // the result is a derived expression.
1878        self.ctx = context.clone();
1879        self.ctx.table_name = None;
1880        self.ctx.schema_name = None;
1881
1882        Ok(plan)
1883    }
1884
1885    fn prom_number_lit_to_plan(&mut self, number_literal: &NumberLiteral) -> Result<LogicalPlan> {
1886        let NumberLiteral { val } = number_literal;
1887        self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1888        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1889        self.ctx.reset_table_name_and_schema();
1890        let literal_expr = df_prelude::lit(*val);
1891
1892        let plan = LogicalPlan::Extension(Extension {
1893            node: Arc::new(
1894                EmptyMetric::new(
1895                    self.ctx.start,
1896                    self.ctx.end,
1897                    self.ctx.interval,
1898                    SPECIAL_TIME_FUNCTION.to_string(),
1899                    DEFAULT_FIELD_COLUMN.to_string(),
1900                    Some(literal_expr),
1901                )
1902                .context(DataFusionPlanningSnafu)?,
1903            ),
1904        });
1905        Ok(plan)
1906    }
1907
1908    fn prom_string_lit_to_plan(&mut self, string_literal: &StringLiteral) -> Result<LogicalPlan> {
1909        let StringLiteral { val } = string_literal;
1910        self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1911        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1912        self.ctx.reset_table_name_and_schema();
1913        let literal_expr = df_prelude::lit(val.clone());
1914
1915        let plan = LogicalPlan::Extension(Extension {
1916            node: Arc::new(
1917                EmptyMetric::new(
1918                    self.ctx.start,
1919                    self.ctx.end,
1920                    self.ctx.interval,
1921                    SPECIAL_TIME_FUNCTION.to_string(),
1922                    DEFAULT_FIELD_COLUMN.to_string(),
1923                    Some(literal_expr),
1924                )
1925                .context(DataFusionPlanningSnafu)?,
1926            ),
1927        });
1928        Ok(plan)
1929    }
1930
1931    async fn prom_vector_selector_to_plan(
1932        &mut self,
1933        vector_selector: &VectorSelector,
1934        timestamp_fn: bool,
1935    ) -> Result<LogicalPlan> {
1936        let VectorSelector {
1937            name,
1938            offset,
1939            matchers,
1940            at: _,
1941        } = vector_selector;
1942        let matchers = self.preprocess_label_matchers(matchers, name)?;
1943        if let Some(empty_plan) = self.setup_context().await? {
1944            return Ok(empty_plan);
1945        }
1946        let normalize = self
1947            .selector_to_series_normalize_plan(offset, matchers, false)
1948            .await?;
1949        let time_index_column =
1950            self.ctx
1951                .time_index_column
1952                .clone()
1953                .with_context(|| TimeIndexNotFoundSnafu {
1954                    table: self.ctx.table_name.clone().unwrap_or_default(),
1955                })?;
1956
1957        let (normalize, timestamp_value_column) = if timestamp_fn {
1958            // Keep the original sample for stale-marker detection while carrying
1959            // its timestamp through InstantManipulate in a private value column.
1960            let occupied = normalize
1961                .schema()
1962                .fields()
1963                .iter()
1964                .map(|field| field.name().as_str())
1965                .collect::<HashSet<_>>();
1966            let mut timestamp_value_column = TIMESTAMP_VALUE_PREFIX.to_string();
1967            while occupied.contains(timestamp_value_column.as_str()) {
1968                timestamp_value_column.push('_');
1969            }
1970            let mut project_exprs = normalize
1971                .schema()
1972                .iter()
1973                .map(|(qualifier, field)| {
1974                    DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1975                })
1976                .collect::<Vec<_>>();
1977            project_exprs
1978                .push(build_special_time_expr(&time_index_column).alias(&timestamp_value_column));
1979            let normalize = LogicalPlanBuilder::from(normalize)
1980                .project(project_exprs)
1981                .context(DataFusionPlanningSnafu)?
1982                .build()
1983                .context(DataFusionPlanningSnafu)?;
1984            (normalize, Some(timestamp_value_column))
1985        } else {
1986            (normalize, None)
1987        };
1988
1989        let field_column = self.ctx.field_columns.first().cloned();
1990        let manipulate = InstantManipulate::new(
1991            self.ctx.start,
1992            self.ctx.end,
1993            self.ctx.lookback_delta,
1994            self.ctx.interval,
1995            time_index_column,
1996            if self.ctx.use_tsid {
1997                vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
1998            } else {
1999                self.ctx.tag_columns.clone()
2000            },
2001            field_column,
2002            normalize,
2003        );
2004        let manipulate = LogicalPlan::Extension(Extension {
2005            node: Arc::new(manipulate),
2006        });
2007        if let Some(timestamp_value_column) = timestamp_value_column {
2008            self.create_timestamp_func_plan(manipulate, &timestamp_value_column)
2009        } else {
2010            Ok(manipulate)
2011        }
2012    }
2013
2014    /// Builds a projection plan for the PromQL `timestamp()` function.
2015    /// Projects the time index column as the value column for each row.
2016    ///
2017    /// # Arguments
2018    /// * `input` - Input [`LogicalPlan`] after instant-vector selection.
2019    /// * `timestamp_value_column` - Private column containing each selected sample's timestamp.
2020    ///
2021    /// # Returns
2022    /// Returns a [`Result<LogicalPlan>`] where the resulting logical plan projects the timestamp
2023    /// column as the value column, along with the original tag and time index columns.
2024    ///
2025    /// # Timestamp vs. Time Function
2026    ///
2027    /// - **Timestamp Function (`timestamp()`)**: In PromQL, the `timestamp()` function returns the
2028    ///   timestamp (time index) of each sample as the value column.
2029    ///
2030    /// - **Time Function (`time()`)**: The `time()` function returns the evaluation time of the query
2031    ///   as a scalar value.
2032    ///
2033    /// # Side Effects
2034    /// Updates the planner context's field columns to the timestamp column name.
2035    ///
2036    fn create_timestamp_func_plan(
2037        &mut self,
2038        input: LogicalPlan,
2039        timestamp_value_column: &str,
2040    ) -> Result<LogicalPlan> {
2041        let time_expr = col(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN);
2042        self.ctx.field_columns = vec![time_expr.schema_name().to_string()];
2043        let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
2044        project_exprs.push(self.create_time_index_column_expr()?);
2045        project_exprs.push(time_expr);
2046        project_exprs.extend(self.create_tag_column_exprs()?);
2047
2048        LogicalPlanBuilder::from(input)
2049            .project(project_exprs)
2050            .context(DataFusionPlanningSnafu)?
2051            .build()
2052            .context(DataFusionPlanningSnafu)
2053    }
2054
2055    async fn prom_matrix_selector_to_plan(
2056        &mut self,
2057        matrix_selector: &MatrixSelector,
2058    ) -> Result<LogicalPlan> {
2059        let MatrixSelector { vs, range } = matrix_selector;
2060        let VectorSelector {
2061            name,
2062            offset,
2063            matchers,
2064            ..
2065        } = vs;
2066        let matchers = self.preprocess_label_matchers(matchers, name)?;
2067        ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
2068        let range_ms = range.as_millis() as _;
2069        self.ctx.range = Some(range_ms);
2070
2071        // Some functions like rate may require special fields in the RangeManipulate plan
2072        // so we can't skip RangeManipulate.
2073        let normalize = match self.setup_context().await? {
2074            Some(empty_plan) => empty_plan,
2075            None => {
2076                self.selector_to_series_normalize_plan(offset, matchers, true)
2077                    .await?
2078            }
2079        };
2080        let manipulate = RangeManipulate::new(
2081            self.ctx.start,
2082            self.ctx.end,
2083            self.ctx.interval,
2084            // TODO(ruihang): convert via Timestamp datatypes to support different time units
2085            range_ms,
2086            self.ctx
2087                .time_index_column
2088                .clone()
2089                .expect("time index should be set in `setup_context`"),
2090            self.ctx.field_columns.clone(),
2091            normalize,
2092        )
2093        .context(DataFusionPlanningSnafu)?;
2094
2095        Ok(LogicalPlan::Extension(Extension {
2096            node: Arc::new(manipulate),
2097        }))
2098    }
2099
2100    async fn prom_call_expr_to_plan(
2101        &mut self,
2102        query_engine_state: &QueryEngineState,
2103        call_expr: &Call,
2104    ) -> Result<LogicalPlan> {
2105        let Call { func, args } = call_expr;
2106        // some special functions that are not expression but a plan
2107        match func.name {
2108            SPECIAL_HISTOGRAM_QUANTILE | SPECIAL_HISTOGRAM_FRACTION => {
2109                return self
2110                    .create_histogram_plan(func.name, args, query_engine_state)
2111                    .await;
2112            }
2113            SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await,
2114            SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await,
2115            SPECIAL_ABSENT_FUNCTION => {
2116                return self.create_absent_plan(args, query_engine_state).await;
2117            }
2118            _ => {}
2119        }
2120
2121        // transform function arguments
2122        let args = self.create_function_args(&args.args)?;
2123        let input = if let Some(prom_expr) = &args.input {
2124            self.prom_expr_to_plan_inner(prom_expr, func.name == "timestamp", query_engine_state)
2125                .await?
2126        } else {
2127            self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
2128            self.ctx.reset_table_name_and_schema();
2129            self.ctx.tag_columns = vec![];
2130            self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2131            LogicalPlan::Extension(Extension {
2132                node: Arc::new(
2133                    EmptyMetric::new(
2134                        self.ctx.start,
2135                        self.ctx.end,
2136                        self.ctx.interval,
2137                        SPECIAL_TIME_FUNCTION.to_string(),
2138                        DEFAULT_FIELD_COLUMN.to_string(),
2139                        None,
2140                    )
2141                    .context(DataFusionPlanningSnafu)?,
2142                ),
2143            })
2144        };
2145        let preserve_any_value =
2146            Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
2147        let (mut func_exprs, new_tags) = self.create_function_expr(
2148            func,
2149            args.literals.clone(),
2150            input.schema(),
2151            query_engine_state,
2152        )?;
2153        func_exprs.insert(0, self.create_time_index_column_expr()?);
2154        func_exprs.extend_from_slice(&self.create_tag_column_exprs()?);
2155        if let Some(tsid_col) =
2156            Self::optional_tsid_projection(input.schema(), None, self.ctx.use_tsid)
2157        {
2158            func_exprs.push(tsid_col);
2159        }
2160
2161        let builder = LogicalPlanBuilder::from(input)
2162            .project(func_exprs)
2163            .context(DataFusionPlanningSnafu)?
2164            .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
2165            .context(DataFusionPlanningSnafu)?;
2166
2167        let builder = match func.name {
2168            "sort" => builder
2169                .sort(self.create_field_columns_sort_exprs(true))
2170                .context(DataFusionPlanningSnafu)?,
2171            "sort_desc" => builder
2172                .sort(self.create_field_columns_sort_exprs(false))
2173                .context(DataFusionPlanningSnafu)?,
2174            "sort_by_label" => builder
2175                .sort(Self::create_sort_exprs_by_tags(
2176                    func.name,
2177                    args.literals,
2178                    true,
2179                )?)
2180                .context(DataFusionPlanningSnafu)?,
2181            "sort_by_label_desc" => builder
2182                .sort(Self::create_sort_exprs_by_tags(
2183                    func.name,
2184                    args.literals,
2185                    false,
2186                )?)
2187                .context(DataFusionPlanningSnafu)?,
2188
2189            _ => builder,
2190        };
2191
2192        // Update context tags after building plan
2193        // We can't push them before planning, because they won't exist until projection.
2194        for tag in new_tags {
2195            self.ctx.tag_columns.push(tag);
2196        }
2197
2198        let plan = builder.build().context(DataFusionPlanningSnafu)?;
2199        common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}");
2200
2201        Ok(plan)
2202    }
2203
2204    async fn prom_ext_expr_to_plan(
2205        &mut self,
2206        query_engine_state: &QueryEngineState,
2207        ext_expr: &promql_parser::parser::ast::Extension,
2208    ) -> Result<LogicalPlan> {
2209        // let promql_parser::parser::ast::Extension { expr } = ext_expr;
2210        let expr = &ext_expr.expr;
2211        let children = expr.children();
2212        let plan = self
2213            .prom_expr_to_plan(&children[0], query_engine_state)
2214            .await?;
2215        // Wrapper for the explanation/analyze of the existing plan
2216        // https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/builder/struct.LogicalPlanBuilder.html#method.explain
2217        // if `analyze` is true, runs the actual plan and produces
2218        // information about metrics during run.
2219        // if `verbose` is true, prints out additional details when VERBOSE keyword is specified
2220        match expr.name() {
2221            ANALYZE_NODE_NAME => LogicalPlanBuilder::from(plan)
2222                .explain(false, true)
2223                .unwrap()
2224                .build()
2225                .context(DataFusionPlanningSnafu),
2226            ANALYZE_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2227                .explain(true, true)
2228                .unwrap()
2229                .build()
2230                .context(DataFusionPlanningSnafu),
2231            EXPLAIN_NODE_NAME => LogicalPlanBuilder::from(plan)
2232                .explain(false, false)
2233                .unwrap()
2234                .build()
2235                .context(DataFusionPlanningSnafu),
2236            EXPLAIN_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2237                .explain(true, false)
2238                .unwrap()
2239                .build()
2240                .context(DataFusionPlanningSnafu),
2241            ALIAS_NODE_NAME => {
2242                let alias = expr
2243                    .as_any()
2244                    .downcast_ref::<AliasExpr>()
2245                    .context(UnexpectedPlanExprSnafu {
2246                        desc: "Expected AliasExpr",
2247                    })?
2248                    .alias
2249                    .clone();
2250                self.apply_alias(plan, alias)
2251            }
2252            _ => LogicalPlanBuilder::empty(true)
2253                .build()
2254                .context(DataFusionPlanningSnafu),
2255        }
2256    }
2257
2258    /// Extract metric name from `__name__` matcher and set it into [PromPlannerContext].
2259    /// Returns a new [Matchers] that doesn't contain metric name matcher.
2260    ///
2261    /// Each call to this function means new selector is started. Thus, the context will be reset
2262    /// at first.
2263    ///
2264    /// Name rule:
2265    /// - if `name` is some, then the matchers MUST NOT contain `__name__` matcher.
2266    /// - if `name` is none, then the matchers MAY contain NONE OR MULTIPLE `__name__` matchers.
2267    #[allow(clippy::mutable_key_type)]
2268    fn preprocess_label_matchers(
2269        &mut self,
2270        label_matchers: &Matchers,
2271        name: &Option<String>,
2272    ) -> Result<Matchers> {
2273        self.ctx.reset();
2274
2275        let metric_name;
2276        if let Some(name) = name.clone() {
2277            metric_name = Some(name);
2278            ensure!(
2279                label_matchers.find_matchers(METRIC_NAME).is_empty(),
2280                MultipleMetricMatchersSnafu
2281            );
2282        } else {
2283            let mut matches = label_matchers.find_matchers(METRIC_NAME);
2284            ensure!(!matches.is_empty(), NoMetricMatcherSnafu);
2285            ensure!(matches.len() == 1, MultipleMetricMatchersSnafu);
2286            ensure!(
2287                matches[0].op == MatchOp::Equal,
2288                UnsupportedMatcherOpSnafu {
2289                    matcher_op: matches[0].op.to_string(),
2290                    matcher: METRIC_NAME
2291                }
2292            );
2293            metric_name = matches.pop().map(|m| m.value);
2294        }
2295
2296        self.ctx.table_name = metric_name;
2297
2298        let mut matchers = HashSet::new();
2299        for matcher in &label_matchers.matchers {
2300            // TODO(ruihang): support other metric match ops
2301            if matcher.name == FIELD_COLUMN_MATCHER {
2302                self.ctx
2303                    .field_column_matcher
2304                    .get_or_insert_default()
2305                    .push(matcher.clone());
2306            } else if matcher.name == SCHEMA_COLUMN_MATCHER || matcher.name == DB_COLUMN_MATCHER {
2307                ensure!(
2308                    matcher.op == MatchOp::Equal,
2309                    UnsupportedMatcherOpSnafu {
2310                        matcher: matcher.name.clone(),
2311                        matcher_op: matcher.op.to_string(),
2312                    }
2313                );
2314                self.ctx.schema_name = Some(matcher.value.clone());
2315            } else if matcher.name != METRIC_NAME {
2316                self.ctx.selector_matcher.push(matcher.clone());
2317                let _ = matchers.insert(matcher.clone());
2318            }
2319        }
2320
2321        Ok(Matchers::new(matchers.into_iter().collect()))
2322    }
2323
2324    async fn selector_to_series_normalize_plan(
2325        &mut self,
2326        offset: &Option<Offset>,
2327        label_matchers: Matchers,
2328        is_range_selector: bool,
2329    ) -> Result<LogicalPlan> {
2330        // make table scan plan
2331        let table_ref = self.table_ref()?;
2332        let mut table_scan = self.create_table_scan_plan(table_ref.clone()).await?;
2333        let table_schema = table_scan.schema();
2334
2335        // make filter exprs
2336        let offset_duration = match offset {
2337            Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2338            Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2339            None => 0,
2340        };
2341        let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?;
2342        if let Some(time_index_filter) = self.build_time_index_filter(offset_duration)? {
2343            scan_filters.push(time_index_filter);
2344        }
2345        table_scan = LogicalPlanBuilder::from(table_scan)
2346            .filter(conjunction(scan_filters).unwrap()) // Safety: `scan_filters` is not empty.
2347            .context(DataFusionPlanningSnafu)?
2348            .build()
2349            .context(DataFusionPlanningSnafu)?;
2350
2351        // make a projection plan if there is any `__field__` matcher
2352        if let Some(field_matchers) = &self.ctx.field_column_matcher {
2353            let col_set = self.ctx.field_columns.iter().collect::<HashSet<_>>();
2354            // opt-in set
2355            let mut result_set = HashSet::new();
2356            // opt-out set
2357            let mut reverse_set = HashSet::new();
2358            for matcher in field_matchers {
2359                match &matcher.op {
2360                    MatchOp::Equal => {
2361                        if col_set.contains(&matcher.value) {
2362                            let _ = result_set.insert(matcher.value.clone());
2363                        } else {
2364                            return Err(ColumnNotFoundSnafu {
2365                                col: matcher.value.clone(),
2366                            }
2367                            .build());
2368                        }
2369                    }
2370                    MatchOp::NotEqual => {
2371                        if col_set.contains(&matcher.value) {
2372                            let _ = reverse_set.insert(matcher.value.clone());
2373                        } else {
2374                            return Err(ColumnNotFoundSnafu {
2375                                col: matcher.value.clone(),
2376                            }
2377                            .build());
2378                        }
2379                    }
2380                    MatchOp::Re(regex) => {
2381                        for col in &self.ctx.field_columns {
2382                            if regex.is_match(col) {
2383                                let _ = result_set.insert(col.clone());
2384                            }
2385                        }
2386                    }
2387                    MatchOp::NotRe(regex) => {
2388                        for col in &self.ctx.field_columns {
2389                            if regex.is_match(col) {
2390                                let _ = reverse_set.insert(col.clone());
2391                            }
2392                        }
2393                    }
2394                }
2395            }
2396            // merge two set
2397            if result_set.is_empty() {
2398                result_set = col_set.into_iter().cloned().collect();
2399            }
2400            for col in reverse_set {
2401                let _ = result_set.remove(&col);
2402            }
2403
2404            // mask the field columns in context using computed result set
2405            self.ctx.field_columns = self
2406                .ctx
2407                .field_columns
2408                .drain(..)
2409                .filter(|col| result_set.contains(col))
2410                .collect();
2411
2412            let exprs = result_set
2413                .into_iter()
2414                .map(|col| DfExpr::Column(Column::new_unqualified(col)))
2415                .chain(self.create_tag_column_exprs()?)
2416                .chain(
2417                    self.ctx
2418                        .use_tsid
2419                        .then_some(DfExpr::Column(Column::new_unqualified(
2420                            DATA_SCHEMA_TSID_COLUMN_NAME,
2421                        ))),
2422                )
2423                .chain(Some(self.create_time_index_column_expr()?))
2424                .collect::<Vec<_>>();
2425
2426            // reuse this variable for simplicity
2427            table_scan = LogicalPlanBuilder::from(table_scan)
2428                .project(exprs)
2429                .context(DataFusionPlanningSnafu)?
2430                .build()
2431                .context(DataFusionPlanningSnafu)?;
2432        }
2433
2434        // make sort plan
2435        let series_key_columns = if self.ctx.use_tsid {
2436            vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
2437        } else {
2438            self.ctx.tag_columns.clone()
2439        };
2440
2441        let sort_exprs = if self.ctx.use_tsid {
2442            vec![
2443                DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)).sort(true, true),
2444                self.create_time_index_column_expr()?.sort(true, true),
2445            ]
2446        } else {
2447            self.create_tag_and_time_index_column_sort_exprs()?
2448        };
2449
2450        let sort_plan = LogicalPlanBuilder::from(table_scan)
2451            .sort(sort_exprs)
2452            .context(DataFusionPlanningSnafu)?
2453            .build()
2454            .context(DataFusionPlanningSnafu)?;
2455
2456        // make divide plan
2457        let time_index_column =
2458            self.ctx
2459                .time_index_column
2460                .clone()
2461                .with_context(|| TimeIndexNotFoundSnafu {
2462                    table: table_ref.to_string(),
2463                })?;
2464        let divide_plan = LogicalPlan::Extension(Extension {
2465            node: Arc::new(SeriesDivide::new(
2466                series_key_columns.clone(),
2467                time_index_column,
2468                sort_plan,
2469            )),
2470        });
2471
2472        // make series_normalize plan
2473        if !is_range_selector && offset_duration == 0 {
2474            return Ok(divide_plan);
2475        }
2476        let series_normalize = SeriesNormalize::new(
2477            offset_duration,
2478            self.ctx
2479                .time_index_column
2480                .clone()
2481                .with_context(|| TimeIndexNotFoundSnafu {
2482                    table: table_ref.to_quoted_string(),
2483                })?,
2484            is_range_selector,
2485            series_key_columns,
2486            divide_plan,
2487        );
2488        let logical_plan = LogicalPlan::Extension(Extension {
2489            node: Arc::new(series_normalize),
2490        });
2491
2492        Ok(logical_plan)
2493    }
2494
2495    /// Convert [LabelModifier] to [Column] exprs for aggregation.
2496    /// Timestamp column and tag columns will be included.
2497    ///
2498    /// # Side effect
2499    ///
2500    /// This method will also change the tag columns in ctx if `update_ctx` is true.
2501    fn agg_modifier_to_col(
2502        &mut self,
2503        input_schema: &DFSchemaRef,
2504        modifier: &Option<LabelModifier>,
2505        update_ctx: bool,
2506    ) -> Result<Vec<DfExpr>> {
2507        match modifier {
2508            None => {
2509                if update_ctx {
2510                    self.ctx.tag_columns.clear();
2511                }
2512                Ok(vec![self.create_time_index_column_expr()?])
2513            }
2514            Some(LabelModifier::Include(labels)) => {
2515                if update_ctx {
2516                    self.ctx.tag_columns.clear();
2517                }
2518                let mut exprs = Vec::with_capacity(labels.labels.len());
2519                for label in &labels.labels {
2520                    if is_metric_engine_internal_column(label) {
2521                        continue;
2522                    }
2523                    // nonexistence label will be ignored
2524                    if let Some(column_name) = Self::find_case_sensitive_column(input_schema, label)
2525                    {
2526                        exprs.push(DfExpr::Column(Column::from_name(column_name.clone())));
2527
2528                        if update_ctx {
2529                            // update the tag columns in context
2530                            self.ctx.tag_columns.push(column_name);
2531                        }
2532                    }
2533                }
2534                // add timestamp column
2535                exprs.push(self.create_time_index_column_expr()?);
2536
2537                Ok(exprs)
2538            }
2539            Some(LabelModifier::Exclude(labels)) => {
2540                let mut all_fields = input_schema
2541                    .fields()
2542                    .iter()
2543                    .map(|f| f.name())
2544                    .collect::<BTreeSet<_>>();
2545
2546                // Exclude metric engine internal columns (not PromQL labels) from the implicit
2547                // "without" label set.
2548                all_fields.retain(|col| !is_metric_engine_internal_column(col.as_str()));
2549
2550                // remove "without"-ed fields
2551                // nonexistence label will be ignored
2552                for label in &labels.labels {
2553                    let _ = all_fields.remove(label);
2554                }
2555
2556                // remove time index and value fields
2557                if let Some(time_index) = &self.ctx.time_index_column {
2558                    let _ = all_fields.remove(time_index);
2559                }
2560                for value in &self.ctx.field_columns {
2561                    let _ = all_fields.remove(value);
2562                }
2563
2564                if update_ctx {
2565                    // change the tag columns in context
2566                    self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect();
2567                }
2568
2569                // collect remaining fields and convert to col expr
2570                let mut exprs = all_fields
2571                    .into_iter()
2572                    .map(|c| DfExpr::Column(Column::from(c)))
2573                    .collect::<Vec<_>>();
2574
2575                // add timestamp column
2576                exprs.push(self.create_time_index_column_expr()?);
2577
2578                Ok(exprs)
2579            }
2580        }
2581    }
2582
2583    // TODO(ruihang): ignore `MetricNameLabel` (`__name__`) matcher
2584    pub fn matchers_to_expr(
2585        label_matchers: Matchers,
2586        table_schema: &DFSchemaRef,
2587    ) -> Result<Vec<DfExpr>> {
2588        let mut exprs = Vec::with_capacity(label_matchers.matchers.len());
2589        for matcher in label_matchers.matchers {
2590            if matcher.name == SCHEMA_COLUMN_MATCHER
2591                || matcher.name == DB_COLUMN_MATCHER
2592                || matcher.name == FIELD_COLUMN_MATCHER
2593            {
2594                continue;
2595            }
2596
2597            let accepts_empty = matcher.is_match("");
2598            let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str());
2599            let col = if let Some(column_name) = column_name {
2600                let column = DfExpr::Column(Column::from_name(&column_name));
2601                let field = table_schema
2602                    .index_of_column_by_name(None, &column_name)
2603                    .map(|index| table_schema.field(index));
2604                if accepts_empty
2605                    && column_name == OTLP_AGGREGATION_TEMPORALITY_LABEL
2606                    && let Some(data_type) = field
2607                        .filter(|field| {
2608                            field.is_nullable()
2609                                && Self::string_value_data_type(field.data_type()).is_some()
2610                        })
2611                        .map(|field| field.data_type())
2612                {
2613                    let empty = Self::string_scalar_value(data_type, Some(String::new()))
2614                        .expect("nullable label has a string type");
2615                    DfExpr::ScalarFunction(ScalarFunction {
2616                        func: coalesce(),
2617                        args: vec![column, DfExpr::Literal(empty, None)],
2618                    })
2619                } else {
2620                    column
2621                }
2622            } else {
2623                DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
2624                    .alias(matcher.name.clone())
2625            };
2626            let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None);
2627            let expr = match matcher.op {
2628                MatchOp::Equal => col.eq(lit),
2629                MatchOp::NotEqual => col.not_eq(lit),
2630                MatchOp::Re(re) => {
2631                    // TODO(ruihang): a more programmatic way to handle this in datafusion
2632
2633                    // This is a hack to handle `.+` and `.*`, and is not strictly correct
2634                    // `.` doesn't match newline (`\n`). Given this is in PromQL context,
2635                    // most of the time it's fine.
2636                    if re.as_str() == "^(?:.*)$" {
2637                        continue;
2638                    }
2639                    if re.as_str() == "^(?:.+)$" {
2640                        col.not_eq(DfExpr::Literal(
2641                            ScalarValue::Utf8(Some(String::new())),
2642                            None,
2643                        ))
2644                    } else {
2645                        DfExpr::BinaryExpr(BinaryExpr {
2646                            left: Box::new(col),
2647                            op: Operator::RegexMatch,
2648                            right: Box::new(DfExpr::Literal(
2649                                ScalarValue::Utf8(Some(re.as_str().to_string())),
2650                                None,
2651                            )),
2652                        })
2653                    }
2654                }
2655                MatchOp::NotRe(re) => {
2656                    if re.as_str() == "^(?:.*)$" {
2657                        DfExpr::Literal(ScalarValue::Boolean(Some(false)), None)
2658                    } else if re.as_str() == "^(?:.+)$" {
2659                        col.eq(DfExpr::Literal(
2660                            ScalarValue::Utf8(Some(String::new())),
2661                            None,
2662                        ))
2663                    } else {
2664                        DfExpr::BinaryExpr(BinaryExpr {
2665                            left: Box::new(col),
2666                            op: Operator::RegexNotMatch,
2667                            right: Box::new(DfExpr::Literal(
2668                                ScalarValue::Utf8(Some(re.as_str().to_string())),
2669                                None,
2670                            )),
2671                        })
2672                    }
2673                }
2674            };
2675            exprs.push(expr);
2676        }
2677
2678        Ok(exprs)
2679    }
2680
2681    fn find_case_sensitive_column(schema: &DFSchemaRef, column: &str) -> Option<String> {
2682        if is_metric_engine_internal_column(column) {
2683            return None;
2684        }
2685        schema
2686            .fields()
2687            .iter()
2688            .find(|field| field.name() == column)
2689            .map(|field| field.name().clone())
2690    }
2691
2692    fn table_from_source(&self, source: &Arc<dyn TableSource>) -> Result<table::TableRef> {
2693        Ok(source
2694            .as_any()
2695            .downcast_ref::<DefaultTableSource>()
2696            .context(UnknownTableSnafu)?
2697            .table_provider
2698            .as_any()
2699            .downcast_ref::<DfTableProviderAdapter>()
2700            .context(UnknownTableSnafu)?
2701            .table())
2702    }
2703
2704    fn table_ref(&self) -> Result<TableReference> {
2705        let table_name = self
2706            .ctx
2707            .table_name
2708            .clone()
2709            .context(TableNameNotFoundSnafu)?;
2710
2711        // set schema name if `__schema__` is given
2712        let table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2713            TableReference::partial(schema_name.as_str(), table_name.as_str())
2714        } else {
2715            TableReference::bare(table_name.as_str())
2716        };
2717
2718        Ok(table_ref)
2719    }
2720
2721    fn build_time_index_filter(&self, offset_duration: i64) -> Result<Option<DfExpr>> {
2722        let start = self.ctx.start;
2723        let end = self.ctx.end;
2724        if end < start {
2725            return InvalidTimeRangeSnafu { start, end }.fail();
2726        }
2727        let lookback_delta = self.ctx.lookback_delta;
2728        let range = self.ctx.range.unwrap_or_default();
2729        let interval = self.ctx.interval;
2730        let time_index_expr = self.create_time_index_column_expr()?;
2731        let num_points = (end - start) / interval;
2732
2733        // Prometheus semantics:
2734        // - Instant selector lookback: (eval_ts - lookback_delta, eval_ts]
2735        // - Range selector:           (eval_ts - range, eval_ts]
2736        //
2737        // So samples positioned exactly at the lower boundary must be excluded. We align the scan
2738        // lower bound with Prometheus by shifting it forward by 1ms (millisecond granularity),
2739        // while still using a `>=` filter.
2740        let selector_window = if range == 0 { lookback_delta } else { range };
2741        let lower_exclusive_adjustment = if selector_window > 0 { 1 } else { 0 };
2742
2743        // Scan a continuous time range
2744        if (end - start) / interval > MAX_SCATTER_POINTS || interval <= INTERVAL_1H {
2745            let single_time_range = time_index_expr
2746                .clone()
2747                .gt_eq(DfExpr::Literal(
2748                    ScalarValue::TimestampMillisecond(
2749                        Some(
2750                            self.ctx.start - offset_duration - selector_window
2751                                + lower_exclusive_adjustment,
2752                        ),
2753                        None,
2754                    ),
2755                    None,
2756                ))
2757                .and(time_index_expr.lt_eq(DfExpr::Literal(
2758                    ScalarValue::TimestampMillisecond(Some(self.ctx.end - offset_duration), None),
2759                    None,
2760                )));
2761            return Ok(Some(single_time_range));
2762        }
2763
2764        // Otherwise scan scatter ranges separately
2765        let mut filters = Vec::with_capacity(num_points as usize + 1);
2766        for timestamp in (start..=end).step_by(interval as usize) {
2767            filters.push(
2768                time_index_expr
2769                    .clone()
2770                    .gt_eq(DfExpr::Literal(
2771                        ScalarValue::TimestampMillisecond(
2772                            Some(
2773                                timestamp - offset_duration - selector_window
2774                                    + lower_exclusive_adjustment,
2775                            ),
2776                            None,
2777                        ),
2778                        None,
2779                    ))
2780                    .and(time_index_expr.clone().lt_eq(DfExpr::Literal(
2781                        ScalarValue::TimestampMillisecond(Some(timestamp - offset_duration), None),
2782                        None,
2783                    ))),
2784            )
2785        }
2786
2787        Ok(filters.into_iter().reduce(DfExpr::or))
2788    }
2789
2790    /// Create a table scan plan and a filter plan with given filter.
2791    ///
2792    /// # Panic
2793    /// If the filter is empty
2794    async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
2795        let provider = self
2796            .table_provider
2797            .resolve_table(table_ref.clone())
2798            .await
2799            .context(CatalogSnafu)?;
2800
2801        let logical_table = self.table_from_source(&provider)?;
2802
2803        // Try to rewrite the table scan to physical table scan if possible.
2804        let mut maybe_phy_table_ref = table_ref.clone();
2805        let mut scan_provider = provider;
2806        let mut table_id_filter: Option<u32> = None;
2807
2808        // If it's a metric engine logical table, scan its physical table directly and filter by
2809        // `__table_id = logical_table_id` to get access to internal columns like `__tsid`.
2810        if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
2811            && let Some(physical_table_name) = logical_table
2812                .table_info()
2813                .meta
2814                .options
2815                .extra_options
2816                .get(LOGICAL_TABLE_METADATA_KEY)
2817        {
2818            let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2819                TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
2820            } else {
2821                TableReference::bare(physical_table_name.as_str())
2822            };
2823
2824            let physical_provider = match self
2825                .table_provider
2826                .resolve_table(physical_table_ref.clone())
2827                .await
2828            {
2829                Ok(provider) => provider,
2830                Err(e) if e.status_code() == StatusCode::TableNotFound => {
2831                    // Fall back to scanning the logical table. It still works, but without
2832                    // `__tsid` optimization.
2833                    scan_provider.clone()
2834                }
2835                Err(e) => return Err(e).context(CatalogSnafu),
2836            };
2837
2838            if !Arc::ptr_eq(&physical_provider, &scan_provider) {
2839                // Only rewrite when internal columns exist in physical schema.
2840                let physical_table = self.table_from_source(&physical_provider)?;
2841
2842                let has_table_id = physical_table
2843                    .schema()
2844                    .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
2845                    .is_some();
2846                let has_tsid = physical_table
2847                    .schema()
2848                    .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2849                    .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2850
2851                if has_table_id && has_tsid {
2852                    scan_provider = physical_provider;
2853                    maybe_phy_table_ref = physical_table_ref;
2854                    table_id_filter = Some(logical_table.table_info().ident.table_id);
2855                }
2856            }
2857        }
2858
2859        let scan_table = self.table_from_source(&scan_provider)?;
2860
2861        let use_tsid = table_id_filter.is_some()
2862            && scan_table
2863                .schema()
2864                .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2865                .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2866        self.ctx.use_tsid = use_tsid;
2867
2868        let all_table_tags = self.ctx.tag_columns.clone();
2869
2870        let scan_tag_columns = if use_tsid {
2871            let mut scan_tags = self.ctx.tag_columns.clone();
2872            for matcher in &self.ctx.selector_matcher {
2873                if is_metric_engine_internal_column(&matcher.name) {
2874                    continue;
2875                }
2876                if all_table_tags.iter().any(|tag| tag == &matcher.name) {
2877                    scan_tags.push(matcher.name.clone());
2878                }
2879            }
2880            scan_tags.sort_unstable();
2881            scan_tags.dedup();
2882            scan_tags
2883        } else {
2884            self.ctx.tag_columns.clone()
2885        };
2886
2887        let is_time_index_ms = scan_table
2888            .schema()
2889            .timestamp_column()
2890            .with_context(|| TimeIndexNotFoundSnafu {
2891                table: maybe_phy_table_ref.to_quoted_string(),
2892            })?
2893            .data_type
2894            == ConcreteDataType::timestamp_millisecond_datatype();
2895
2896        let scan_projection = if table_id_filter.is_some() {
2897            let mut required_columns = HashSet::new();
2898            required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
2899            required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
2900                TimeIndexNotFoundSnafu {
2901                    table: maybe_phy_table_ref.to_quoted_string(),
2902                }
2903            })?);
2904            for col in &scan_tag_columns {
2905                required_columns.insert(col.clone());
2906            }
2907            for col in &self.ctx.field_columns {
2908                required_columns.insert(col.clone());
2909            }
2910            if use_tsid {
2911                required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
2912            }
2913
2914            let arrow_schema = scan_table.schema().arrow_schema().clone();
2915            Some(
2916                arrow_schema
2917                    .fields()
2918                    .iter()
2919                    .enumerate()
2920                    .filter(|(_, field)| required_columns.contains(field.name().as_str()))
2921                    .map(|(idx, _)| idx)
2922                    .collect::<Vec<_>>(),
2923            )
2924        } else {
2925            None
2926        };
2927
2928        let mut scan_plan =
2929            LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
2930                .context(DataFusionPlanningSnafu)?
2931                .build()
2932                .context(DataFusionPlanningSnafu)?;
2933
2934        if let Some(table_id) = table_id_filter {
2935            scan_plan = LogicalPlanBuilder::from(scan_plan)
2936                .filter(
2937                    DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
2938                        .eq(lit(table_id)),
2939                )
2940                .context(DataFusionPlanningSnafu)?
2941                .alias(table_ref.clone()) // rename the relation back to logical table's name after filtering
2942                .context(DataFusionPlanningSnafu)?
2943                .build()
2944                .context(DataFusionPlanningSnafu)?;
2945        }
2946
2947        if !is_time_index_ms {
2948            // cast to ms if time_index not in Millisecond precision
2949            let expr: Vec<_> = self
2950                .create_field_column_exprs()?
2951                .into_iter()
2952                .chain(
2953                    scan_tag_columns
2954                        .iter()
2955                        .map(|tag| DfExpr::Column(Column::from_name(tag))),
2956                )
2957                .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
2958                    Some(table_ref.clone()),
2959                    DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
2960                ))))
2961                .chain(Some(DfExpr::Alias(Alias {
2962                    expr: Box::new(DfExpr::Cast(Cast {
2963                        expr: Box::new(self.create_time_index_column_expr()?),
2964                        data_type: ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
2965                    })),
2966                    relation: Some(table_ref.clone()),
2967                    name: self
2968                        .ctx
2969                        .time_index_column
2970                        .as_ref()
2971                        .with_context(|| TimeIndexNotFoundSnafu {
2972                            table: table_ref.to_quoted_string(),
2973                        })?
2974                        .clone(),
2975                    metadata: None,
2976                })))
2977                .collect::<Vec<_>>();
2978            scan_plan = LogicalPlanBuilder::from(scan_plan)
2979                .project(expr)
2980                .context(DataFusionPlanningSnafu)?
2981                .build()
2982                .context(DataFusionPlanningSnafu)?;
2983        } else if table_id_filter.is_some() {
2984            // Drop the internal `__table_id` column after filtering.
2985            let project_exprs = self
2986                .create_field_column_exprs()?
2987                .into_iter()
2988                .chain(
2989                    scan_tag_columns
2990                        .iter()
2991                        .map(|tag| DfExpr::Column(Column::from_name(tag))),
2992                )
2993                .chain(
2994                    self.ctx
2995                        .use_tsid
2996                        .then_some(DfExpr::Column(Column::from_name(
2997                            DATA_SCHEMA_TSID_COLUMN_NAME,
2998                        ))),
2999                )
3000                .chain(Some(self.create_time_index_column_expr()?))
3001                .collect::<Vec<_>>();
3002
3003            scan_plan = LogicalPlanBuilder::from(scan_plan)
3004                .project(project_exprs)
3005                .context(DataFusionPlanningSnafu)?
3006                .build()
3007                .context(DataFusionPlanningSnafu)?;
3008        }
3009
3010        let result = LogicalPlanBuilder::from(scan_plan)
3011            .build()
3012            .context(DataFusionPlanningSnafu)?;
3013        Ok(result)
3014    }
3015
3016    fn collect_row_key_tag_columns_from_plan(
3017        &self,
3018        plan: &LogicalPlan,
3019    ) -> Result<BTreeSet<String>> {
3020        fn walk(
3021            planner: &PromPlanner,
3022            plan: &LogicalPlan,
3023            out: &mut BTreeSet<String>,
3024        ) -> Result<()> {
3025            // Derived PromQL plans may contain non-Greptime scans without row-key metadata.
3026            if let LogicalPlan::TableScan(scan) = plan
3027                && let Ok(table) = planner.table_from_source(&scan.source)
3028            {
3029                for col in table.table_info().meta.row_key_column_names() {
3030                    if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
3031                        && col != DATA_SCHEMA_TSID_COLUMN_NAME
3032                        && !is_metric_engine_internal_column(col)
3033                    {
3034                        out.insert(col.clone());
3035                    }
3036                }
3037            }
3038
3039            for input in plan.inputs() {
3040                walk(planner, input, out)?;
3041            }
3042            Ok(())
3043        }
3044
3045        let mut out = BTreeSet::new();
3046        walk(self, plan, &mut out)?;
3047        Ok(out)
3048    }
3049
3050    fn ensure_tag_columns_available(
3051        &self,
3052        plan: LogicalPlan,
3053        required_tags: &BTreeSet<String>,
3054    ) -> Result<LogicalPlan> {
3055        if required_tags.is_empty() {
3056            return Ok(plan);
3057        }
3058
3059        struct Rewriter {
3060            required_tags: BTreeSet<String>,
3061        }
3062
3063        impl TreeNodeRewriter for Rewriter {
3064            type Node = LogicalPlan;
3065
3066            fn f_up(
3067                &mut self,
3068                node: Self::Node,
3069            ) -> datafusion_common::Result<Transformed<Self::Node>> {
3070                match node {
3071                    LogicalPlan::TableScan(scan) => {
3072                        let schema = scan.source.schema();
3073                        let mut projection = match scan.projection.clone() {
3074                            Some(p) => p,
3075                            None => {
3076                                // Scanning all columns already covers required tags.
3077                                return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
3078                            }
3079                        };
3080
3081                        let mut changed = false;
3082                        for tag in &self.required_tags {
3083                            if let Some((idx, _)) = schema
3084                                .fields()
3085                                .iter()
3086                                .enumerate()
3087                                .find(|(_, field)| field.name() == tag)
3088                                && !projection.contains(&idx)
3089                            {
3090                                projection.push(idx);
3091                                changed = true;
3092                            }
3093                        }
3094
3095                        if !changed {
3096                            return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
3097                        }
3098
3099                        projection.sort_unstable();
3100                        projection.dedup();
3101
3102                        let new_scan = TableScan::try_new(
3103                            scan.table_name.clone(),
3104                            scan.source.clone(),
3105                            Some(projection),
3106                            scan.filters,
3107                            scan.fetch,
3108                        )?;
3109                        Ok(Transformed::yes(LogicalPlan::TableScan(new_scan)))
3110                    }
3111                    LogicalPlan::Projection(proj) => {
3112                        let input_schema = proj.input.schema();
3113
3114                        let existing = proj
3115                            .schema
3116                            .fields()
3117                            .iter()
3118                            .map(|f| f.name().as_str())
3119                            .collect::<HashSet<_>>();
3120
3121                        let mut expr = proj.expr.clone();
3122                        let mut has_changed = false;
3123                        for tag in &self.required_tags {
3124                            if existing.contains(tag.as_str()) {
3125                                continue;
3126                            }
3127
3128                            if let Some(idx) = input_schema.index_of_column_by_name(None, tag) {
3129                                expr.push(DfExpr::Column(Column::from(
3130                                    input_schema.qualified_field(idx),
3131                                )));
3132                                has_changed = true;
3133                            }
3134                        }
3135
3136                        if !has_changed {
3137                            return Ok(Transformed::no(LogicalPlan::Projection(proj)));
3138                        }
3139
3140                        let new_proj = Projection::try_new(expr, proj.input)?;
3141                        Ok(Transformed::yes(LogicalPlan::Projection(new_proj)))
3142                    }
3143                    other => Ok(Transformed::no(other)),
3144                }
3145            }
3146        }
3147
3148        let mut rewriter = Rewriter {
3149            required_tags: required_tags.clone(),
3150        };
3151        let rewritten = plan
3152            .rewrite(&mut rewriter)
3153            .context(DataFusionPlanningSnafu)?;
3154        Ok(rewritten.data)
3155    }
3156
3157    fn refresh_tag_columns_from_schema(&mut self, schema: &DFSchemaRef) {
3158        let time_index = self.ctx.time_index_column.as_deref();
3159        let field_columns = self.ctx.field_columns.iter().collect::<HashSet<_>>();
3160
3161        let mut tags = schema
3162            .fields()
3163            .iter()
3164            .map(|f| f.name())
3165            .filter(|name| Some(name.as_str()) != time_index)
3166            .filter(|name| !field_columns.contains(name))
3167            .filter(|name| !is_metric_engine_internal_column(name))
3168            .cloned()
3169            .collect::<Vec<_>>();
3170        tags.sort_unstable();
3171        tags.dedup();
3172        self.ctx.tag_columns = tags;
3173    }
3174
3175    /// Setup [PromPlannerContext]'s state fields.
3176    ///
3177    /// Returns a logical plan for an empty metric.
3178    async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
3179        let table_ref = self.table_ref()?;
3180        let source = match self.table_provider.resolve_table(table_ref.clone()).await {
3181            Err(e) if e.status_code() == StatusCode::TableNotFound => {
3182                let plan = self.setup_context_for_empty_metric()?;
3183                return Ok(Some(plan));
3184            }
3185            res => res.context(CatalogSnafu)?,
3186        };
3187        let table = self.table_from_source(&source)?;
3188
3189        // set time index column name
3190        let time_index = table
3191            .schema()
3192            .timestamp_column()
3193            .with_context(|| TimeIndexNotFoundSnafu {
3194                table: table_ref.to_quoted_string(),
3195            })?
3196            .name
3197            .clone();
3198        self.ctx.time_index_column = Some(time_index);
3199
3200        // set values columns
3201        let values = table
3202            .table_info()
3203            .meta
3204            .field_column_names()
3205            .cloned()
3206            .collect();
3207        self.ctx.field_columns = values;
3208
3209        // set primary key (tag) columns
3210        let tags = table
3211            .table_info()
3212            .meta
3213            .row_key_column_names()
3214            .filter(|col| {
3215                // remove metric engine's internal columns
3216                col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
3217            })
3218            .cloned()
3219            .collect();
3220        self.ctx.tag_columns = tags;
3221
3222        self.ctx.use_tsid = false;
3223
3224        Ok(None)
3225    }
3226
3227    /// Setup [PromPlannerContext]'s state fields for a non existent table
3228    /// without any rows.
3229    fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
3230        self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3231        self.ctx.reset_table_name_and_schema();
3232        self.ctx.tag_columns = vec![];
3233        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
3234        self.ctx.use_tsid = false;
3235
3236        // The table doesn't have any data, so we set start to 0 and end to -1.
3237        let plan = LogicalPlan::Extension(Extension {
3238            node: Arc::new(
3239                EmptyMetric::new(
3240                    0,
3241                    -1,
3242                    self.ctx.interval,
3243                    SPECIAL_TIME_FUNCTION.to_string(),
3244                    DEFAULT_FIELD_COLUMN.to_string(),
3245                    Some(lit(0.0f64)),
3246                )
3247                .context(DataFusionPlanningSnafu)?,
3248            ),
3249        });
3250        Ok(plan)
3251    }
3252
3253    // TODO(ruihang): insert column expr
3254    fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
3255        let mut result = FunctionArgs::default();
3256
3257        for arg in args {
3258            // First try to parse as literal expression (including binary expressions like 100.0 + 3.0)
3259            if let Some(expr) = Self::try_build_literal_expr(arg) {
3260                result.literals.push(expr);
3261            } else {
3262                // If not a literal, treat as vector input
3263                match arg.as_ref() {
3264                    PromExpr::Subquery(_)
3265                    | PromExpr::VectorSelector(_)
3266                    | PromExpr::MatrixSelector(_)
3267                    | PromExpr::Extension(_)
3268                    | PromExpr::Aggregate(_)
3269                    | PromExpr::Paren(_)
3270                    | PromExpr::Call(_)
3271                    | PromExpr::Binary(_)
3272                    | PromExpr::Unary(_) => {
3273                        if result.input.replace(*arg.clone()).is_some() {
3274                            MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
3275                        }
3276                    }
3277
3278                    _ => {
3279                        let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?;
3280                        result.literals.push(expr);
3281                    }
3282                }
3283            }
3284        }
3285
3286        Ok(result)
3287    }
3288
3289    fn create_mixed_range_function_exprs(
3290        &mut self,
3291        func: &Function,
3292        mut other_input_exprs: VecDeque<DfExpr>,
3293        float_field: &str,
3294        histogram_field: &str,
3295        input_schema: &DFSchemaRef,
3296    ) -> Result<Option<Vec<DfExpr>>> {
3297        let returns_histogram = matches!(
3298            func.name,
3299            "rate"
3300                | "increase"
3301                | "delta"
3302                | "idelta"
3303                | "irate"
3304                | "avg_over_time"
3305                | "sum_over_time"
3306                | "last_over_time"
3307        );
3308        if !returns_histogram
3309            && !matches!(
3310                func.name,
3311                "changes"
3312                    | "resets"
3313                    | "deriv"
3314                    | "min_over_time"
3315                    | "max_over_time"
3316                    | "count_over_time"
3317                    | "absent_over_time"
3318                    | "present_over_time"
3319                    | "stddev_over_time"
3320                    | "stdvar_over_time"
3321                    | "quantile_over_time"
3322                    | "predict_linear"
3323                    | "double_exponential_smoothing"
3324                    | "holt_winters"
3325            )
3326        {
3327            return Ok(None);
3328        }
3329
3330        if func.name == "predict_linear" {
3331            other_input_exprs[0] = DfExpr::Cast(Cast {
3332                expr: Box::new(other_input_exprs[0].clone()),
3333                data_type: ArrowDataType::Int64,
3334            });
3335        }
3336
3337        let timestamp_range = DfExpr::Column(Column::from_name(
3338            RangeManipulate::build_timestamp_range_name(
3339                self.ctx.time_index_column.as_ref().unwrap(),
3340            ),
3341        ));
3342        let float_range = DfExpr::Column(Column::from_name(float_field));
3343        let histogram_range = DfExpr::Column(Column::from_name(histogram_field));
3344        let mut args = Vec::with_capacity(other_input_exprs.len() + 6);
3345        args.push(lit(func.name));
3346        args.push(timestamp_range.clone());
3347        args.push(float_range.clone());
3348        args.push(histogram_range.clone());
3349        args.extend(other_input_exprs);
3350        if matches!(func.name, "rate" | "increase" | "delta") {
3351            args.push(self.create_time_index_column_expr()?);
3352            args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?));
3353        }
3354
3355        let mut float_expr = DfExpr::ScalarFunction(ScalarFunction {
3356            func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3357            args: args.clone(),
3358        });
3359        if matches!(func.name, "rate" | "increase") {
3360            let raw_delta_function = if func.name == "rate" {
3361                "raw_delta_rate"
3362            } else {
3363                "raw_delta_increase"
3364            };
3365            let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3366                func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3367                args: vec![
3368                    lit(raw_delta_function),
3369                    timestamp_range,
3370                    float_range,
3371                    histogram_range,
3372                ],
3373            });
3374            float_expr = self.select_delta_range_math(
3375                func.name,
3376                input_schema,
3377                self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3378                delta_sum,
3379                float_expr,
3380            )?;
3381        }
3382        let exprs = if returns_histogram {
3383            self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()];
3384            vec![
3385                float_expr.alias(float_field),
3386                DfExpr::ScalarFunction(ScalarFunction {
3387                    func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())),
3388                    args,
3389                })
3390                .alias(histogram_field),
3391            ]
3392        } else {
3393            let display_name = float_expr.schema_name().to_string();
3394            self.ctx.field_columns = vec![display_name.clone()];
3395            vec![float_expr.alias(display_name)]
3396        };
3397        Ok(Some(exprs))
3398    }
3399
3400    /// Creates function expressions for projection and returns the expressions and new tags.
3401    ///
3402    /// # Side Effects
3403    ///
3404    /// This method will update [PromPlannerContext]'s fields and tags if needed.
3405    fn create_function_expr(
3406        &mut self,
3407        func: &Function,
3408        other_input_exprs: Vec<DfExpr>,
3409        input_schema: &DFSchemaRef,
3410        query_engine_state: &QueryEngineState,
3411    ) -> Result<(Vec<DfExpr>, Vec<String>)> {
3412        // TODO(ruihang): check function args list
3413        let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
3414        if let Some((float_field, histogram_field)) =
3415            Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns)
3416                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
3417            && let Some(exprs) = self.create_mixed_range_function_exprs(
3418                func,
3419                other_input_exprs.clone(),
3420                &float_field,
3421                &histogram_field,
3422                input_schema,
3423            )?
3424        {
3425            return Ok((exprs, vec![]));
3426        }
3427        let alternative_samples =
3428            Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns);
3429        let all_field_columns_are_native_histogram_ranges =
3430            self.all_field_columns_are_native_histogram_ranges(input_schema);
3431
3432        // TODO(ruihang): set this according to in-param list
3433        let field_column_pos = 0;
3434        let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3435        // New labels after executing the function, e.g. `label_replace` etc.
3436        let mut new_tags = vec![];
3437        let promql_annotations = self.promql_annotations.clone();
3438        let native_histogram_drop_udf = |name: &str| {
3439            Arc::new(NativeHistogramDrop::float_null_udf(
3440                format!(
3441                    "{name}: dropped native histogram samples because this function is not supported for native histograms"
3442                ),
3443                promql_annotations.clone(),
3444            ))
3445        };
3446        let scalar_func = match func.name {
3447            "increase" => {
3448                if all_field_columns_are_native_histogram_ranges {
3449                    ScalarFunc::ExtrapolateUdf(
3450                        Arc::new(NativeHistogramIncrease::scalar_udf_with_collector(
3451                            self.promql_annotations.clone(),
3452                        )),
3453                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3454                    )
3455                } else {
3456                    ScalarFunc::ExtrapolateUdf(
3457                        Arc::new(Increase::scalar_udf()),
3458                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3459                    )
3460                }
3461            }
3462            "rate" => {
3463                if all_field_columns_are_native_histogram_ranges {
3464                    ScalarFunc::ExtrapolateUdf(
3465                        Arc::new(NativeHistogramRate::scalar_udf_with_collector(
3466                            self.promql_annotations.clone(),
3467                        )),
3468                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3469                    )
3470                } else {
3471                    ScalarFunc::ExtrapolateUdf(
3472                        Arc::new(Rate::scalar_udf()),
3473                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3474                    )
3475                }
3476            }
3477            "delta" => {
3478                if all_field_columns_are_native_histogram_ranges {
3479                    ScalarFunc::ExtrapolateUdf(
3480                        Arc::new(NativeHistogramDelta::scalar_udf_with_collector(
3481                            self.promql_annotations.clone(),
3482                        )),
3483                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3484                    )
3485                } else {
3486                    ScalarFunc::ExtrapolateUdf(
3487                        Arc::new(Delta::scalar_udf()),
3488                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3489                    )
3490                }
3491            }
3492            "idelta" => {
3493                if all_field_columns_are_native_histogram_ranges {
3494                    ScalarFunc::Udf(Arc::new(NativeHistogramIDelta::scalar_udf_with_collector(
3495                        self.promql_annotations.clone(),
3496                    )))
3497                } else {
3498                    ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf()))
3499                }
3500            }
3501            "irate" => {
3502                if all_field_columns_are_native_histogram_ranges {
3503                    ScalarFunc::Udf(Arc::new(NativeHistogramIRate::scalar_udf_with_collector(
3504                        self.promql_annotations.clone(),
3505                    )))
3506                } else {
3507                    ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf()))
3508                }
3509            }
3510            "resets" => {
3511                if all_field_columns_are_native_histogram_ranges {
3512                    ScalarFunc::Udf(Arc::new(NativeHistogramResets::scalar_udf()))
3513                } else {
3514                    ScalarFunc::Udf(Arc::new(Resets::scalar_udf()))
3515                }
3516            }
3517            "changes" => {
3518                if all_field_columns_are_native_histogram_ranges {
3519                    ScalarFunc::Udf(Arc::new(NativeHistogramChanges::scalar_udf()))
3520                } else {
3521                    ScalarFunc::Udf(Arc::new(Changes::scalar_udf()))
3522                }
3523            }
3524            "deriv" => {
3525                if all_field_columns_are_native_histogram_ranges {
3526                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3527                } else {
3528                    ScalarFunc::Udf(Arc::new(Deriv::scalar_udf()))
3529                }
3530            }
3531            "avg_over_time" => {
3532                if all_field_columns_are_native_histogram_ranges {
3533                    ScalarFunc::Udf(Arc::new(
3534                        NativeHistogramAvgOverTime::scalar_udf_with_collector(
3535                            self.promql_annotations.clone(),
3536                        ),
3537                    ))
3538                } else {
3539                    ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf()))
3540                }
3541            }
3542            "min_over_time" => {
3543                if all_field_columns_are_native_histogram_ranges {
3544                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3545                } else {
3546                    ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf()))
3547                }
3548            }
3549            "max_over_time" => {
3550                if all_field_columns_are_native_histogram_ranges {
3551                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3552                } else {
3553                    ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf()))
3554                }
3555            }
3556            "sum_over_time" => {
3557                if all_field_columns_are_native_histogram_ranges {
3558                    ScalarFunc::Udf(Arc::new(
3559                        NativeHistogramSumOverTime::scalar_udf_with_collector(
3560                            self.promql_annotations.clone(),
3561                        ),
3562                    ))
3563                } else {
3564                    ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf()))
3565                }
3566            }
3567            "count_over_time" => {
3568                if all_field_columns_are_native_histogram_ranges {
3569                    ScalarFunc::Udf(Arc::new(NativeHistogramCountOverTime::scalar_udf()))
3570                } else {
3571                    ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf()))
3572                }
3573            }
3574            "last_over_time" => {
3575                if all_field_columns_are_native_histogram_ranges {
3576                    ScalarFunc::Udf(Arc::new(NativeHistogramLastOverTime::scalar_udf()))
3577                } else {
3578                    ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf()))
3579                }
3580            }
3581            "absent_over_time" => {
3582                if all_field_columns_are_native_histogram_ranges {
3583                    ScalarFunc::Udf(Arc::new(NativeHistogramAbsentOverTime::scalar_udf()))
3584                } else {
3585                    ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf()))
3586                }
3587            }
3588            "present_over_time" => {
3589                if all_field_columns_are_native_histogram_ranges {
3590                    ScalarFunc::Udf(Arc::new(NativeHistogramPresentOverTime::scalar_udf()))
3591                } else {
3592                    ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf()))
3593                }
3594            }
3595            "stddev_over_time" => {
3596                if all_field_columns_are_native_histogram_ranges {
3597                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3598                } else {
3599                    ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf()))
3600                }
3601            }
3602            "stdvar_over_time" => {
3603                if all_field_columns_are_native_histogram_ranges {
3604                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3605                } else {
3606                    ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf()))
3607                }
3608            }
3609            "quantile_over_time" => {
3610                if all_field_columns_are_native_histogram_ranges {
3611                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3612                } else {
3613                    ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf()))
3614                }
3615            }
3616            "predict_linear" => {
3617                if all_field_columns_are_native_histogram_ranges {
3618                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3619                } else {
3620                    other_input_exprs[0] = DfExpr::Cast(Cast {
3621                        expr: Box::new(other_input_exprs[0].clone()),
3622                        data_type: ArrowDataType::Int64,
3623                    });
3624                    ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
3625                }
3626            }
3627            "double_exponential_smoothing" | "holt_winters" => {
3628                if all_field_columns_are_native_histogram_ranges {
3629                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3630                } else {
3631                    ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
3632                }
3633            }
3634            "histogram_count" => {
3635                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramCount::scalar_udf()))
3636            }
3637            "histogram_sum" => {
3638                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramSum::scalar_udf()))
3639            }
3640            "histogram_avg" => {
3641                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramAvg::scalar_udf()))
3642            }
3643            "histogram_stddev" => {
3644                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStddev::scalar_udf()))
3645            }
3646            "histogram_stdvar" => {
3647                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf()))
3648            }
3649            "time" => {
3650                exprs.push(build_special_time_expr(
3651                    self.ctx.time_index_column.as_ref().unwrap(),
3652                ));
3653                ScalarFunc::GeneratedExpr
3654            }
3655            "minute" => {
3656                // date_part('minute', time_index)
3657                let expr = self.date_part_on_time_index("minute")?;
3658                exprs.push(expr);
3659                ScalarFunc::GeneratedExpr
3660            }
3661            "hour" => {
3662                // date_part('hour', time_index)
3663                let expr = self.date_part_on_time_index("hour")?;
3664                exprs.push(expr);
3665                ScalarFunc::GeneratedExpr
3666            }
3667            "month" => {
3668                // date_part('month', time_index)
3669                let expr = self.date_part_on_time_index("month")?;
3670                exprs.push(expr);
3671                ScalarFunc::GeneratedExpr
3672            }
3673            "year" => {
3674                // date_part('year', time_index)
3675                let expr = self.date_part_on_time_index("year")?;
3676                exprs.push(expr);
3677                ScalarFunc::GeneratedExpr
3678            }
3679            "day_of_month" => {
3680                // date_part('day', time_index)
3681                let expr = self.date_part_on_time_index("day")?;
3682                exprs.push(expr);
3683                ScalarFunc::GeneratedExpr
3684            }
3685            "day_of_week" => {
3686                // date_part('dow', time_index)
3687                let expr = self.date_part_on_time_index("dow")?;
3688                exprs.push(expr);
3689                ScalarFunc::GeneratedExpr
3690            }
3691            "day_of_year" => {
3692                // date_part('doy', time_index)
3693                let expr = self.date_part_on_time_index("doy")?;
3694                exprs.push(expr);
3695                ScalarFunc::GeneratedExpr
3696            }
3697            "days_in_month" => {
3698                // date_part(
3699                //     'days',
3700                //     (date_trunc('month', <TIME INDEX>::date) + interval '1 month - 1 day')
3701                // );
3702                let day_lit_expr = "day".lit();
3703                let month_lit_expr = "month".lit();
3704                let interval_1month_lit_expr =
3705                    DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
3706                let interval_1day_lit_expr = DfExpr::Literal(
3707                    ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
3708                    None,
3709                );
3710                let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
3711                    left: Box::new(interval_1month_lit_expr),
3712                    op: Operator::Minus,
3713                    right: Box::new(interval_1day_lit_expr),
3714                });
3715                let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
3716                    func: datafusion_functions::datetime::date_trunc(),
3717                    args: vec![month_lit_expr, self.create_time_index_column_expr()?],
3718                });
3719                let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
3720                    left: Box::new(date_trunc_expr),
3721                    op: Operator::Plus,
3722                    right: Box::new(the_1month_minus_1day_expr),
3723                });
3724                let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
3725                    func: datafusion_functions::datetime::date_part(),
3726                    args: vec![day_lit_expr, date_trunc_plus_interval_expr],
3727                });
3728
3729                exprs.push(date_part_expr);
3730                ScalarFunc::GeneratedExpr
3731            }
3732
3733            "label_join" => {
3734                self.ctx.use_tsid = false;
3735                let (concat_expr, dst_label) = Self::build_concat_labels_expr(
3736                    &mut other_input_exprs,
3737                    &self.ctx,
3738                    query_engine_state,
3739                )?;
3740
3741                // Reserve the current field columns except the `dst_label`.
3742                for value in &self.ctx.field_columns {
3743                    if *value != dst_label {
3744                        let expr = DfExpr::Column(Column::from_name(value));
3745                        exprs.push(expr);
3746                    }
3747                }
3748
3749                // Remove it from tag columns if exists to avoid duplicated column names
3750                self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3751                new_tags.push(dst_label);
3752                // Add the new label expr to evaluate
3753                exprs.push(concat_expr);
3754
3755                ScalarFunc::GeneratedExpr
3756            }
3757            "label_replace" => {
3758                self.ctx.use_tsid = false;
3759                if let Some((replace_expr, dst_label)) = self
3760                    .build_regexp_replace_label_expr(&mut other_input_exprs, query_engine_state)?
3761                {
3762                    // Reserve the current field columns except the `dst_label`.
3763                    for value in &self.ctx.field_columns {
3764                        if *value != dst_label {
3765                            let expr = DfExpr::Column(Column::from_name(value));
3766                            exprs.push(expr);
3767                        }
3768                    }
3769
3770                    ensure!(
3771                        !self.ctx.tag_columns.contains(&dst_label),
3772                        SameLabelSetSnafu
3773                    );
3774                    new_tags.push(dst_label);
3775                    // Add the new label expr to evaluate
3776                    exprs.push(replace_expr);
3777                } else {
3778                    // Keep the current field columns
3779                    for value in &self.ctx.field_columns {
3780                        let expr = DfExpr::Column(Column::from_name(value));
3781                        exprs.push(expr);
3782                    }
3783                }
3784
3785                ScalarFunc::GeneratedExpr
3786            }
3787            "sort" | "sort_desc" => {
3788                // Value sorting silently ignores native histogram samples.
3789                for value in &self.ctx.field_columns {
3790                    if !Self::field_column_is_native_histogram(input_schema, value) {
3791                        exprs.push(DfExpr::Column(Column::from_name(value)));
3792                    }
3793                }
3794                // Keep a nullable float field so the normal empty-value filter produces an
3795                // empty vector when the input contains only histograms.
3796                if exprs.is_empty() {
3797                    exprs.push(DfExpr::Literal(ScalarValue::Float64(None), None));
3798                }
3799
3800                ScalarFunc::GeneratedExpr
3801            }
3802            "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
3803                // These functions are not expression but a part of plan,
3804                // they are processed by `prom_call_expr_to_plan`.
3805                for value in &self.ctx.field_columns {
3806                    let expr = DfExpr::Column(Column::from_name(value));
3807                    exprs.push(expr);
3808                }
3809
3810                ScalarFunc::GeneratedExpr
3811            }
3812            "round" if self.all_field_columns_are_native_histograms(input_schema) => {
3813                if other_input_exprs.is_empty() {
3814                    other_input_exprs.push_front(0.0f64.lit());
3815                }
3816                ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3817            }
3818            "round" => {
3819                if other_input_exprs.is_empty() {
3820                    other_input_exprs.push_front(0.0f64.lit());
3821                }
3822                ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
3823            }
3824            "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => {
3825                ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3826            }
3827            "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
3828            "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
3829            "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
3830            "pi" => {
3831                // pi functions doesn't accepts any arguments, needs special processing
3832                let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3833                    func: datafusion::functions::math::pi(),
3834                    args: vec![],
3835                });
3836                exprs.push(fn_expr);
3837
3838                ScalarFunc::GeneratedExpr
3839            }
3840            _ => {
3841                if let Some(f) = query_engine_state
3842                    .session_state()
3843                    .scalar_functions()
3844                    .get(func.name)
3845                {
3846                    if self.all_field_columns_are_native_histograms(input_schema) {
3847                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3848                    } else {
3849                        ScalarFunc::DataFusionBuiltin(f.clone())
3850                    }
3851                } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
3852                    if self.all_field_columns_are_native_histograms(input_schema) {
3853                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3854                    } else {
3855                        let func_state = query_engine_state.function_state();
3856                        let query_ctx = self.table_provider.query_ctx();
3857
3858                        ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
3859                            state: func_state,
3860                            query_ctx: query_ctx.clone(),
3861                        })))
3862                    }
3863                } else if let Some(f) = datafusion_functions::math::functions()
3864                    .iter()
3865                    .find(|f| f.name() == func.name)
3866                {
3867                    if self.all_field_columns_are_native_histograms(input_schema) {
3868                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3869                    } else {
3870                        ScalarFunc::DataFusionUdf(f.clone())
3871                    }
3872                } else {
3873                    return UnsupportedExprSnafu {
3874                        name: func.name.to_string(),
3875                    }
3876                    .fail();
3877                }
3878            }
3879        };
3880
3881        for value in &self.ctx.field_columns {
3882            let col_expr = DfExpr::Column(Column::from_name(value));
3883            let value_is_histogram = Self::field_column_is_native_histogram(input_schema, value);
3884
3885            match scalar_func.clone() {
3886                ScalarFunc::DataFusionBuiltin(func) => {
3887                    if alternative_samples && value_is_histogram {
3888                        continue;
3889                    }
3890                    other_input_exprs.insert(field_column_pos, col_expr);
3891                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3892                        func,
3893                        args: other_input_exprs.clone().into(),
3894                    });
3895                    exprs.push(fn_expr);
3896                    let _ = other_input_exprs.remove(field_column_pos);
3897                }
3898                ScalarFunc::DataFusionUdf(func) => {
3899                    if alternative_samples && value_is_histogram {
3900                        continue;
3901                    }
3902                    let args = itertools::chain!(
3903                        other_input_exprs.iter().take(field_column_pos).cloned(),
3904                        std::iter::once(col_expr),
3905                        other_input_exprs.iter().skip(field_column_pos).cloned()
3906                    )
3907                    .collect_vec();
3908                    exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
3909                }
3910                ScalarFunc::NativeHistogramUdf(func) => {
3911                    if value_is_histogram {
3912                        let args = itertools::chain!(
3913                            other_input_exprs.iter().take(field_column_pos).cloned(),
3914                            std::iter::once(col_expr),
3915                            other_input_exprs.iter().skip(field_column_pos).cloned()
3916                        )
3917                        .collect_vec();
3918                        exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }));
3919                    } else if !alternative_samples {
3920                        exprs.push(
3921                            DfExpr::Literal(ScalarValue::Float64(None), None).alias(format!(
3922                                "{}_{}",
3923                                func.name(),
3924                                value
3925                            )),
3926                        );
3927                    }
3928                }
3929                ScalarFunc::Udf(func) => {
3930                    let ts_range_expr = DfExpr::Column(Column::from_name(
3931                        RangeManipulate::build_timestamp_range_name(
3932                            self.ctx.time_index_column.as_ref().unwrap(),
3933                        ),
3934                    ));
3935                    other_input_exprs.insert(field_column_pos, ts_range_expr);
3936                    other_input_exprs.insert(field_column_pos + 1, col_expr);
3937                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3938                        func,
3939                        args: other_input_exprs.clone().into(),
3940                    });
3941                    exprs.push(fn_expr);
3942                    let _ = other_input_exprs.remove(field_column_pos + 1);
3943                    let _ = other_input_exprs.remove(field_column_pos);
3944                }
3945                ScalarFunc::ExtrapolateUdf(udf, range_length) => {
3946                    let ts_range_expr = DfExpr::Column(Column::from_name(
3947                        RangeManipulate::build_timestamp_range_name(
3948                            self.ctx.time_index_column.as_ref().unwrap(),
3949                        ),
3950                    ));
3951                    other_input_exprs.insert(field_column_pos, ts_range_expr.clone());
3952                    other_input_exprs.insert(field_column_pos + 1, col_expr.clone());
3953                    other_input_exprs
3954                        .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
3955                    other_input_exprs.push_back(lit(range_length));
3956                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3957                        func: udf,
3958                        args: other_input_exprs.clone().into(),
3959                    });
3960                    let fn_expr = if matches!(func.name, "rate" | "increase")
3961                        && !all_field_columns_are_native_histogram_ranges
3962                    {
3963                        let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3964                            func: Arc::new(SumOverTime::scalar_udf()),
3965                            args: vec![ts_range_expr, col_expr],
3966                        });
3967                        self.select_delta_range_math(
3968                            func.name,
3969                            input_schema,
3970                            range_length,
3971                            delta_sum,
3972                            fn_expr,
3973                        )?
3974                    } else {
3975                        fn_expr
3976                    };
3977                    exprs.push(fn_expr);
3978                    let _ = other_input_exprs.pop_back();
3979                    let _ = other_input_exprs.remove(field_column_pos + 2);
3980                    let _ = other_input_exprs.remove(field_column_pos + 1);
3981                    let _ = other_input_exprs.remove(field_column_pos);
3982                }
3983                ScalarFunc::GeneratedExpr => {}
3984            }
3985        }
3986
3987        // Update value columns' name, and alias them to remove qualifiers
3988        // For label functions such as `label_join`, `label_replace`, etc.,
3989        // we keep the fields unchanged.
3990        if !matches!(func.name, "label_join" | "label_replace") {
3991            let mut new_field_columns = Vec::with_capacity(exprs.len());
3992
3993            exprs = exprs
3994                .into_iter()
3995                .map(|expr| {
3996                    let display_name = expr.schema_name().to_string();
3997                    new_field_columns.push(display_name.clone());
3998                    Ok(expr.alias(display_name))
3999                })
4000                .collect::<std::result::Result<Vec<_>, _>>()
4001                .context(DataFusionPlanningSnafu)?;
4002
4003            self.ctx.field_columns = new_field_columns;
4004        }
4005
4006        Ok((exprs, new_tags))
4007    }
4008
4009    fn select_delta_range_math(
4010        &self,
4011        function: &str,
4012        input_schema: &DFSchemaRef,
4013        range_length: Millisecond,
4014        delta_sum: DfExpr,
4015        cumulative: DfExpr,
4016    ) -> Result<DfExpr> {
4017        let marker_is_delta = if self
4018            .ctx
4019            .tag_columns
4020            .iter()
4021            .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)
4022        {
4023            Self::field_column_type(input_schema, OTLP_AGGREGATION_TEMPORALITY_LABEL)
4024                .filter(|data_type| Self::string_value_data_type(data_type).is_some())
4025                .map(|_| {
4026                    DfExpr::Column(Column::from_name(OTLP_AGGREGATION_TEMPORALITY_LABEL))
4027                        .eq(lit(GREPTIME_TEMPORALITY_DELTA))
4028                })
4029        } else {
4030            None
4031        };
4032        let Some(marker_is_delta) = marker_is_delta else {
4033            return Ok(cumulative);
4034        };
4035
4036        let delta = if function == "rate" {
4037            DfExpr::BinaryExpr(BinaryExpr {
4038                left: Box::new(delta_sum),
4039                op: Operator::Divide,
4040                right: Box::new(lit(range_length as f64 / 1000.0)),
4041            })
4042        } else {
4043            delta_sum
4044        };
4045        let display_name = cumulative.schema_name().to_string();
4046        when(marker_is_delta, delta)
4047            .otherwise(cumulative)
4048            .context(DataFusionPlanningSnafu)
4049            .map(|expr| expr.alias(display_name))
4050    }
4051
4052    /// Validate label name according to Prometheus specification.
4053    /// Label names must match the regex: [a-zA-Z_][a-zA-Z0-9_]*
4054    /// Additionally, label names starting with double underscores are reserved for internal use.
4055    fn validate_label_name(label_name: &str) -> Result<()> {
4056        // Check if label name starts with double underscores (reserved)
4057        if label_name.starts_with("__") {
4058            return InvalidDestinationLabelNameSnafu { label_name }.fail();
4059        }
4060        // Check if label name matches the required pattern
4061        if !LABEL_NAME_REGEX.is_match(label_name) {
4062            return InvalidDestinationLabelNameSnafu { label_name }.fail();
4063        }
4064
4065        Ok(())
4066    }
4067
4068    /// Build expr for `label_replace` function
4069    fn build_regexp_replace_label_expr(
4070        &self,
4071        other_input_exprs: &mut VecDeque<DfExpr>,
4072        query_engine_state: &QueryEngineState,
4073    ) -> Result<Option<(DfExpr, String)>> {
4074        // label_replace(vector, dst_label, replacement, src_label, regex)
4075        let dst_label = match other_input_exprs.pop_front() {
4076            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4077            other => UnexpectedPlanExprSnafu {
4078                desc: format!("expected dst_label string literal, but found {:?}", other),
4079            }
4080            .fail()?,
4081        };
4082
4083        // Validate the destination label name
4084        Self::validate_label_name(&dst_label)?;
4085        let replacement = match other_input_exprs.pop_front() {
4086            Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
4087            other => UnexpectedPlanExprSnafu {
4088                desc: format!("expected replacement string literal, but found {:?}", other),
4089            }
4090            .fail()?,
4091        };
4092        let src_label = match other_input_exprs.pop_front() {
4093            Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
4094            other => UnexpectedPlanExprSnafu {
4095                desc: format!("expected src_label string literal, but found {:?}", other),
4096            }
4097            .fail()?,
4098        };
4099
4100        let regex = match other_input_exprs.pop_front() {
4101            Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
4102            other => UnexpectedPlanExprSnafu {
4103                desc: format!("expected regex string literal, but found {:?}", other),
4104            }
4105            .fail()?,
4106        };
4107
4108        // Validate the regex before using it
4109        // doc: https://prometheus.io/docs/prometheus/latest/querying/functions/#label_replace
4110        regex::Regex::new(&regex).map_err(|_| {
4111            InvalidRegularExpressionSnafu {
4112                regex: regex.clone(),
4113            }
4114            .build()
4115        })?;
4116
4117        // If the src_label exists and regex is empty, keep everything unchanged.
4118        if self.ctx.tag_columns.contains(&src_label) && regex.is_empty() {
4119            return Ok(None);
4120        }
4121
4122        // If the src_label doesn't exists, and
4123        if !self.ctx.tag_columns.contains(&src_label) {
4124            if replacement.is_empty() {
4125                // the replacement is empty, keep everything unchanged.
4126                return Ok(None);
4127            } else {
4128                // the replacement is not empty, always adds dst_label with replacement value.
4129                return Ok(Some((
4130                    // alias literal `replacement` as dst_label
4131                    lit(replacement).alias(&dst_label),
4132                    dst_label,
4133                )));
4134            }
4135        }
4136
4137        // Preprocess the regex:
4138        // https://github.com/prometheus/prometheus/blob/d902abc50d6652ba8fe9a81ff8e5cce936114eba/promql/functions.go#L1575C32-L1575C37
4139        let regex = format!("^(?s:{regex})$");
4140
4141        let session_state = query_engine_state.session_state();
4142        let func = session_state
4143            .scalar_functions()
4144            .get("regexp_replace")
4145            .context(UnsupportedExprSnafu {
4146                name: "regexp_replace",
4147            })?;
4148
4149        // regexp_replace(src_label, regex, replacement)
4150        let args = vec![
4151            if src_label.is_empty() {
4152                DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
4153            } else {
4154                DfExpr::Column(Column::from_name(src_label))
4155            },
4156            DfExpr::Literal(ScalarValue::Utf8(Some(regex)), None),
4157            DfExpr::Literal(ScalarValue::Utf8(Some(replacement)), None),
4158        ];
4159
4160        Ok(Some((
4161            DfExpr::ScalarFunction(ScalarFunction {
4162                func: func.clone(),
4163                args,
4164            })
4165            .alias(&dst_label),
4166            dst_label,
4167        )))
4168    }
4169
4170    /// Build expr for `label_join` function
4171    fn build_concat_labels_expr(
4172        other_input_exprs: &mut VecDeque<DfExpr>,
4173        ctx: &PromPlannerContext,
4174        query_engine_state: &QueryEngineState,
4175    ) -> Result<(DfExpr, String)> {
4176        // label_join(vector, dst_label, separator, src_label_1, src_label_2, ...)
4177
4178        let dst_label = match other_input_exprs.pop_front() {
4179            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4180            other => UnexpectedPlanExprSnafu {
4181                desc: format!("expected dst_label string literal, but found {:?}", other),
4182            }
4183            .fail()?,
4184        };
4185        let separator = match other_input_exprs.pop_front() {
4186            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4187            other => UnexpectedPlanExprSnafu {
4188                desc: format!("expected separator string literal, but found {:?}", other),
4189            }
4190            .fail()?,
4191        };
4192
4193        // Create a set of available columns (tag columns + field columns + time index column)
4194        let available_columns: HashSet<&str> = ctx
4195            .tag_columns
4196            .iter()
4197            .chain(ctx.field_columns.iter())
4198            .chain(ctx.time_index_column.as_ref())
4199            .map(|s| s.as_str())
4200            .collect();
4201
4202        let src_labels = other_input_exprs
4203            .iter()
4204            .map(|expr| {
4205                // Cast source label into column or null literal
4206                match expr {
4207                    DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
4208                        if label.is_empty() {
4209                            Ok(DfExpr::Literal(ScalarValue::Null, None))
4210                        } else if available_columns.contains(label.as_str()) {
4211                            // Label exists in the table schema
4212                            Ok(DfExpr::Column(Column::from_name(label)))
4213                        } else {
4214                            // Label doesn't exist, treat as empty string (null)
4215                            Ok(DfExpr::Literal(ScalarValue::Null, None))
4216                        }
4217                    }
4218                    other => UnexpectedPlanExprSnafu {
4219                        desc: format!(
4220                            "expected source label string literal, but found {:?}",
4221                            other
4222                        ),
4223                    }
4224                    .fail(),
4225                }
4226            })
4227            .collect::<Result<Vec<_>>>()?;
4228        ensure!(
4229            !src_labels.is_empty(),
4230            FunctionInvalidArgumentSnafu {
4231                fn_name: "label_join"
4232            }
4233        );
4234
4235        let session_state = query_engine_state.session_state();
4236        let func = session_state
4237            .scalar_functions()
4238            .get("concat_ws")
4239            .context(UnsupportedExprSnafu { name: "concat_ws" })?;
4240
4241        // concat_ws(separator, src_label_1, src_label_2, ...) as dst_label
4242        let mut args = Vec::with_capacity(1 + src_labels.len());
4243        args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
4244        args.extend(src_labels);
4245
4246        Ok((
4247            DfExpr::ScalarFunction(ScalarFunction {
4248                func: func.clone(),
4249                args,
4250            })
4251            .alias(&dst_label),
4252            dst_label,
4253        ))
4254    }
4255
4256    fn create_time_index_column_expr(&self) -> Result<DfExpr> {
4257        Ok(DfExpr::Column(Column::from_name(
4258            self.ctx
4259                .time_index_column
4260                .clone()
4261                .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
4262        )))
4263    }
4264
4265    fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
4266        let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
4267        for tag in &self.ctx.tag_columns {
4268            let expr = DfExpr::Column(Column::from_name(tag));
4269            result.push(expr);
4270        }
4271        Ok(result)
4272    }
4273
4274    fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
4275        let mut result = Vec::with_capacity(self.ctx.field_columns.len());
4276        for field in &self.ctx.field_columns {
4277            let expr = DfExpr::Column(Column::from_name(field));
4278            result.push(expr);
4279        }
4280        Ok(result)
4281    }
4282
4283    fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
4284        let mut result = self
4285            .ctx
4286            .tag_columns
4287            .iter()
4288            .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
4289            .collect::<Vec<_>>();
4290        result.push(self.create_time_index_column_expr()?.sort(true, true));
4291        Ok(result)
4292    }
4293
4294    fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
4295        self.ctx
4296            .field_columns
4297            .iter()
4298            .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
4299            .collect::<Vec<_>>()
4300    }
4301
4302    fn create_sort_exprs_by_tags(
4303        func: &str,
4304        tags: Vec<DfExpr>,
4305        asc: bool,
4306    ) -> Result<Vec<SortExpr>> {
4307        ensure!(
4308            !tags.is_empty(),
4309            FunctionInvalidArgumentSnafu { fn_name: func }
4310        );
4311
4312        tags.iter()
4313            .map(|col| match col {
4314                DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
4315                    Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
4316                }
4317                other => UnexpectedPlanExprSnafu {
4318                    desc: format!("expected label string literal, but found {:?}", other),
4319                }
4320                .fail(),
4321            })
4322            .collect::<Result<Vec<_>>>()
4323    }
4324
4325    fn create_empty_values_filter_expr(&self, preserve_any_value: bool) -> Result<DfExpr> {
4326        let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
4327        for value in &self.ctx.field_columns {
4328            let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
4329            exprs.push(expr);
4330        }
4331
4332        // This error context should be computed lazily: the planner may set `ctx.table_name` to
4333        // `None` for derived expressions (e.g. after projecting the LHS of a vector-vector
4334        // comparison filter). Eagerly calling `table_ref()?` here can turn a valid plan into
4335        // a `TableNameNotFound` error even when predicate construction succeeds.
4336        let predicate = if preserve_any_value {
4337            disjunction(exprs)
4338        } else {
4339            conjunction(exprs)
4340        };
4341        predicate.with_context(|| ValueNotFoundSnafu {
4342            table: self
4343                .table_ref()
4344                .map(|t| t.to_quoted_string())
4345                .unwrap_or_else(|_| "unknown".to_string()),
4346        })
4347    }
4348
4349    /// Creates a set of DataFusion `DfExpr::AggregateFunction` expressions for each value column using the specified aggregate function.
4350    ///
4351    /// # Side Effects
4352    ///
4353    /// This method modifies the value columns in the context by replacing them with the new columns
4354    /// created by the aggregate function application.
4355    ///
4356    /// # Returns
4357    ///
4358    /// Returns a tuple of `(aggregate_expressions, previous_field_expressions)` where:
4359    /// - `aggregate_expressions`: Expressions that apply the aggregate function to the original fields
4360    /// - `previous_field_expressions`: Original field expressions before aggregation. This is non-empty
4361    ///   only when the operation is `count_values`, as this operation requires preserving the original
4362    ///   values for grouping.
4363    ///
4364    fn create_aggregate_exprs(
4365        &mut self,
4366        op: TokenType,
4367        param: &Option<Box<PromExpr>>,
4368        input_plan: &LogicalPlan,
4369    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4370        let mixed_sample_columns =
4371            Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4372                .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
4373        let is_group_agg = op.id() == token::T_GROUP;
4374        if is_group_agg && mixed_sample_columns.is_none() {
4375            ensure!(
4376                self.ctx.field_columns.len() == 1,
4377                MultiFieldsNotSupportedSnafu {
4378                    operator: "group()"
4379                }
4380            );
4381        }
4382
4383        if let Some((float, histogram)) = mixed_sample_columns {
4384            return self.create_mixed_aggregate_exprs(op, param, &float, &histogram);
4385        }
4386
4387        if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4388            return self.create_native_histogram_aggregate_exprs(op, input_plan);
4389        }
4390
4391        // perform aggregate operation to each value column
4392        let exprs = self
4393            .ctx
4394            .field_columns
4395            .iter()
4396            .map(|col| {
4397                Self::create_numeric_aggregate_expr(
4398                    op,
4399                    param,
4400                    DfExpr::Column(Column::from_name(col)),
4401                )
4402            })
4403            .collect::<Result<Vec<_>>>()?;
4404
4405        // if the aggregator is `count_values`, it must be grouped by current fields.
4406        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4407            let prev_field_exprs: Vec<_> = self
4408                .ctx
4409                .field_columns
4410                .iter()
4411                .map(|col| DfExpr::Column(Column::from_name(col)))
4412                .collect();
4413
4414            ensure!(
4415                self.ctx.field_columns.len() == 1,
4416                UnsupportedExprSnafu {
4417                    name: "count_values on multi-value input"
4418                }
4419            );
4420
4421            prev_field_exprs
4422        } else {
4423            vec![]
4424        };
4425
4426        // update value column name according to the aggregators,
4427        let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
4428
4429        let normalized_exprs =
4430            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4431        for expr in normalized_exprs {
4432            new_field_columns.push(expr.schema_name().to_string());
4433        }
4434        self.ctx.field_columns = new_field_columns;
4435
4436        Ok((exprs, prev_field_exprs))
4437    }
4438
4439    fn create_numeric_aggregate_expr(
4440        op: TokenType,
4441        param: &Option<Box<PromExpr>>,
4442        input: DfExpr,
4443    ) -> Result<DfExpr> {
4444        let expr = match op.id() {
4445            token::T_SUM => sum_udaf().call(vec![input]),
4446            token::T_QUANTILE => {
4447                let q = Self::get_param_as_literal_expr(
4448                    param.as_deref(),
4449                    Some(op),
4450                    Some(ArrowDataType::Float64),
4451                )?;
4452                quantile_udaf().call(vec![q, input])
4453            }
4454            token::T_AVG => avg_udaf().call(vec![input]),
4455            token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]),
4456            token::T_MIN => min_udaf().call(vec![input]),
4457            token::T_MAX => max_udaf().call(vec![input]),
4458            // PromQL's `group()` aggregator produces 1 for each group.
4459            // Use `max(1.0)` (per-group) to match semantics and output type (Float64).
4460            token::T_GROUP => max_udaf().call(vec![lit(1_f64)]),
4461            token::T_STDDEV => stddev_pop_udaf().call(vec![input]),
4462            token::T_STDVAR => var_pop_udaf().call(vec![input]),
4463            token::T_TOPK | token::T_BOTTOMK => {
4464                return UnsupportedExprSnafu {
4465                    name: format!("{op:?}"),
4466                }
4467                .fail();
4468            }
4469            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4470        };
4471        Ok(expr)
4472    }
4473
4474    fn create_mixed_aggregate_exprs(
4475        &mut self,
4476        op: TokenType,
4477        param: &Option<Box<PromExpr>>,
4478        float_column: &str,
4479        histogram_column: &str,
4480    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4481        let float_input = DfExpr::Column(Column::from_name(float_column));
4482        let histogram_input = DfExpr::Column(Column::from_name(histogram_column));
4483        let float_count = count_udaf().call(vec![float_input.clone()]);
4484        let histogram_count = count_udaf().call(vec![histogram_input.clone()]);
4485        let mixed_sample_value = || {
4486            DfExpr::ScalarFunction(ScalarFunction {
4487                func: coalesce(),
4488                args: vec![
4489                    DfExpr::ScalarFunction(ScalarFunction {
4490                        func: Arc::new(PromqlFloatToString::scalar_udf()),
4491                        args: vec![float_input.clone()],
4492                    }),
4493                    DfExpr::ScalarFunction(ScalarFunction {
4494                        func: Arc::new(NativeHistogramToString::scalar_udf()),
4495                        args: vec![histogram_input.clone()],
4496                    }),
4497                ],
4498            })
4499        };
4500
4501        let (exprs, prev_field_exprs, field_columns) = match op.id() {
4502            token::T_SUM | token::T_AVG => (
4503                vec![
4504                    Self::create_numeric_aggregate_expr(op, param, float_input)?
4505                        .alias(float_column),
4506                    self.create_native_histogram_aggregate_expr(op, histogram_column)?,
4507                    float_count.alias(Self::mixed_sample_count_name(float_column)),
4508                    histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4509                ],
4510                vec![],
4511                vec![float_column.to_string(), histogram_column.to_string()],
4512            ),
4513            token::T_COUNT => {
4514                let present = when(
4515                    float_input
4516                        .clone()
4517                        .is_not_null()
4518                        .or(histogram_input.clone().is_not_null()),
4519                    lit(1_i64),
4520                )
4521                .otherwise(lit(ScalarValue::Int64(None)))
4522                .context(DataFusionPlanningSnafu)?;
4523                (
4524                    vec![count_udaf().call(vec![present]).alias(float_column)],
4525                    vec![],
4526                    vec![float_column.to_string()],
4527                )
4528            }
4529            token::T_GROUP => (
4530                vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)],
4531                vec![],
4532                vec![float_column.to_string()],
4533            ),
4534            token::T_COUNT_VALUES => {
4535                let value = mixed_sample_value();
4536                (
4537                    vec![count_udaf().call(vec![value.clone()]).alias(float_column)],
4538                    vec![value],
4539                    vec![float_column.to_string()],
4540                )
4541            }
4542            token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => {
4543                (
4544                    vec![
4545                        Self::create_numeric_aggregate_expr(op, param, float_input)?
4546                            .alias(float_column),
4547                        histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4548                    ],
4549                    vec![],
4550                    vec![float_column.to_string()],
4551                )
4552            }
4553            token::T_TOPK | token::T_BOTTOMK => {
4554                return UnsupportedExprSnafu {
4555                    name: format!("{op:?}"),
4556                }
4557                .fail();
4558            }
4559            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4560        };
4561
4562        self.ctx.field_columns = field_columns;
4563        Ok((exprs, prev_field_exprs))
4564    }
4565
4566    fn mixed_sample_count_column(column: &str) -> DfExpr {
4567        DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column)))
4568    }
4569
4570    fn mixed_sample_count_name(column: &str) -> String {
4571        format!("__promql_sample_count({column})")
4572    }
4573
4574    fn mixed_aggregate_filter_expr(
4575        &self,
4576        op: TokenType,
4577        float_column: &str,
4578        histogram_column: &str,
4579    ) -> Result<DfExpr> {
4580        let float_count = Self::mixed_sample_count_column(float_column);
4581        let histogram_count = Self::mixed_sample_count_column(histogram_column);
4582        let mixed = float_count
4583            .clone()
4584            .gt(lit(0_i64))
4585            .and(histogram_count.clone().gt(lit(0_i64)));
4586        let drop_mixed = DfExpr::ScalarFunction(ScalarFunction {
4587            func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
4588                format!(
4589                    "{op}: dropped aggregation result containing both float and native histogram samples"
4590                ),
4591                self.promql_annotations.clone(),
4592            )),
4593            args: vec![float_count, histogram_count],
4594        });
4595
4596        when(mixed, drop_mixed)
4597            .otherwise(lit(true))
4598            .context(DataFusionPlanningSnafu)
4599    }
4600
4601    fn mixed_ignored_histogram_filter_expr(
4602        &self,
4603        op: TokenType,
4604        histogram_column: &str,
4605    ) -> Result<DfExpr> {
4606        let histogram_count = Self::mixed_sample_count_column(histogram_column);
4607        let has_histograms = histogram_count.clone().gt(lit(0_i64));
4608        let record_info = DfExpr::ScalarFunction(ScalarFunction {
4609            func: Arc::new(NativeHistogramDrop::bool_true_udf(
4610                format!(
4611                    "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4612                ),
4613                self.promql_annotations.clone(),
4614            )),
4615            args: vec![histogram_count],
4616        });
4617
4618        when(has_histograms, record_info)
4619            .otherwise(lit(true))
4620            .context(DataFusionPlanningSnafu)
4621    }
4622
4623    fn create_native_histogram_aggregate_expr(
4624        &self,
4625        op: TokenType,
4626        column: &str,
4627    ) -> Result<DfExpr> {
4628        let input = DfExpr::Column(Column::from_name(column));
4629        let expr = match op.id() {
4630            token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector(
4631                self.promql_annotations.clone(),
4632            ))
4633            .call(vec![input])
4634            .alias(column),
4635            token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector(
4636                self.promql_annotations.clone(),
4637            ))
4638            .call(vec![input])
4639            .alias(column),
4640            token::T_COUNT_VALUES | token::T_COUNT => {
4641                count_udaf().call(vec![input]).alias(column)
4642            }
4643            token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column),
4644            token::T_MIN
4645            | token::T_MAX
4646            | token::T_STDDEV
4647            | token::T_STDVAR
4648            | token::T_QUANTILE
4649            | token::T_TOPK
4650            | token::T_BOTTOMK => sum_udaf()
4651                .call(vec![DfExpr::ScalarFunction(ScalarFunction {
4652                    func: Arc::new(NativeHistogramDrop::float_null_udf(
4653                        format!(
4654                            "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4655                        ),
4656                        self.promql_annotations.clone(),
4657                    )),
4658                    args: vec![input],
4659                })])
4660                .alias(column),
4661            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4662        };
4663        Ok(expr)
4664    }
4665
4666    fn create_native_histogram_aggregate_exprs(
4667        &mut self,
4668        op: TokenType,
4669        input_plan: &LogicalPlan,
4670    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4671        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4672            ensure!(
4673                self.ctx.field_columns.len() == 1,
4674                UnsupportedExprSnafu {
4675                    name: "count_values on multi-value input"
4676                }
4677            );
4678            self.ctx
4679                .field_columns
4680                .iter()
4681                .map(|col| {
4682                    DfExpr::ScalarFunction(ScalarFunction {
4683                        func: Arc::new(NativeHistogramToString::scalar_udf()),
4684                        args: vec![DfExpr::Column(Column::from_name(col))],
4685                    })
4686                })
4687                .collect::<Vec<_>>()
4688        } else {
4689            vec![]
4690        };
4691
4692        let exprs = self
4693            .ctx
4694            .field_columns
4695            .iter()
4696            .map(|col| self.create_native_histogram_aggregate_expr(op, col))
4697            .collect::<Result<Vec<_>>>()?;
4698
4699        let normalized_exprs =
4700            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4701        self.ctx.field_columns = normalized_exprs
4702            .into_iter()
4703            .map(|expr| expr.schema_name().to_string())
4704            .collect();
4705
4706        Ok((exprs, prev_field_exprs))
4707    }
4708
4709    fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
4710        let param = param
4711            .as_deref()
4712            .with_context(|| FunctionInvalidArgumentSnafu {
4713                fn_name: op.to_string(),
4714            })?;
4715        let PromExpr::StringLiteral(StringLiteral { val }) = param else {
4716            return FunctionInvalidArgumentSnafu {
4717                fn_name: op.to_string(),
4718            }
4719            .fail();
4720        };
4721
4722        Ok(val)
4723    }
4724
4725    fn get_param_as_literal_expr(
4726        param: Option<&PromExpr>,
4727        op: Option<TokenType>,
4728        expected_type: Option<ArrowDataType>,
4729    ) -> Result<DfExpr> {
4730        let prom_param = param.with_context(|| {
4731            if let Some(op) = op {
4732                FunctionInvalidArgumentSnafu {
4733                    fn_name: op.to_string(),
4734                }
4735            } else {
4736                FunctionInvalidArgumentSnafu {
4737                    fn_name: "unknown".to_string(),
4738                }
4739            }
4740        })?;
4741
4742        let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
4743            if let Some(op) = op {
4744                FunctionInvalidArgumentSnafu {
4745                    fn_name: op.to_string(),
4746                }
4747            } else {
4748                FunctionInvalidArgumentSnafu {
4749                    fn_name: "unknown".to_string(),
4750                }
4751            }
4752        })?;
4753
4754        // check if the type is expected
4755        if let Some(expected_type) = expected_type {
4756            // literal should not have reference to column
4757            let expr_type = expr
4758                .get_type(&DFSchema::empty())
4759                .context(DataFusionPlanningSnafu)?;
4760            if expected_type != expr_type {
4761                return FunctionInvalidArgumentSnafu {
4762                    fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
4763                }
4764                .fail();
4765            }
4766        }
4767
4768        Ok(expr)
4769    }
4770
4771    /// Create [DfExpr::WindowFunction] expr for each value column with given window function.
4772    ///
4773    fn create_window_exprs(
4774        &mut self,
4775        op: TokenType,
4776        group_exprs: Vec<DfExpr>,
4777        input_plan: &LogicalPlan,
4778    ) -> Result<Vec<DfExpr>> {
4779        ensure!(
4780            self.ctx.field_columns.len() == 1,
4781            UnsupportedExprSnafu {
4782                name: "topk or bottomk on multi-value input"
4783            }
4784        );
4785
4786        assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
4787
4788        let asc = matches!(op.id(), token::T_BOTTOMK);
4789
4790        let tag_sort_exprs = self
4791            .create_tag_column_exprs()?
4792            .into_iter()
4793            .map(|expr| expr.sort(asc, true));
4794
4795        // perform window operation to each value column
4796        let exprs: Vec<DfExpr> = self
4797            .ctx
4798            .field_columns
4799            .iter()
4800            .map(|col| {
4801                let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1);
4802                // Order by value in the specific order
4803                sort_exprs.push(DfExpr::Column(Column::from(col)).sort(asc, true));
4804                // Then tags if the values are equal,
4805                // Try to ensure the relative stability of the output results.
4806                sort_exprs.extend(tag_sort_exprs.clone());
4807
4808                DfExpr::WindowFunction(Box::new(WindowFunction {
4809                    fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
4810                    params: WindowFunctionParams {
4811                        args: vec![],
4812                        partition_by: group_exprs.clone(),
4813                        order_by: sort_exprs,
4814                        window_frame: WindowFrame::new(Some(true)),
4815                        null_treatment: None,
4816                        distinct: false,
4817                        filter: None,
4818                    },
4819                }))
4820            })
4821            .collect();
4822
4823        let normalized_exprs =
4824            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4825        Ok(normalized_exprs)
4826    }
4827
4828    /// Create a classic, native, or mixed histogram helper plan.
4829    async fn create_histogram_plan(
4830        &mut self,
4831        function_name: &str,
4832        args: &PromFunctionArgs,
4833        query_engine_state: &QueryEngineState,
4834    ) -> Result<LogicalPlan> {
4835        let float_literal = |param: &PromExpr| -> Result<f64> {
4836            let value = (|| {
4837                let expr = Self::get_param_as_literal_expr(
4838                    Some(param),
4839                    None,
4840                    Some(ArrowDataType::Float64),
4841                )
4842                .ok()?;
4843                let simplifier = ExprSimplifier::new(SimplifyContext::default());
4844                let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?;
4845                let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else {
4846                    return None;
4847                };
4848                let ScalarValue::Float64(Some(value)) =
4849                    value.cast_to(&ArrowDataType::Float64).ok()?
4850                else {
4851                    return None;
4852                };
4853                Some(value)
4854            })()
4855            .with_context(|| FunctionInvalidArgumentSnafu {
4856                fn_name: function_name.to_string(),
4857            })?;
4858            Ok(value)
4859        };
4860        let (function, input) = match (function_name, args.args.as_slice()) {
4861            (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => (
4862                HistogramFoldOperation::Quantile(float_literal(quantile)?.into()),
4863                input.as_ref().clone(),
4864            ),
4865            (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => (
4866                HistogramFoldOperation::Fraction {
4867                    lower: float_literal(lower)?.into(),
4868                    upper: float_literal(upper)?.into(),
4869                },
4870                input.as_ref().clone(),
4871            ),
4872            _ => {
4873                return FunctionInvalidArgumentSnafu {
4874                    fn_name: function_name.to_string(),
4875                }
4876                .fail();
4877            }
4878        };
4879
4880        let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?;
4881        // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a
4882        // stable series identifier anymore. HistogramFold must not treat it as a label column.
4883        let input_plan = self.strip_tsid_column(input_plan)?;
4884        self.ctx.use_tsid = false;
4885
4886        if let Some((float_field, histogram_field)) =
4887            Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4888                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
4889        {
4890            if self.ctx.has_le_tag() {
4891                return self.create_mixed_histogram_plan(
4892                    function,
4893                    input_plan,
4894                    float_field,
4895                    histogram_field,
4896                );
4897            }
4898            self.ctx.field_columns = vec![histogram_field];
4899        }
4900        if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4901            return self.create_native_histogram_plan(function, input_plan);
4902        }
4903
4904        if !self.ctx.has_le_tag() {
4905            // Return empty result instead of error when 'le' column is not found
4906            // This handles the case when histogram metrics don't exist
4907            return Ok(LogicalPlan::EmptyRelation(
4908                datafusion::logical_expr::EmptyRelation {
4909                    produce_one_row: false,
4910                    schema: input_plan.schema().clone(),
4911                },
4912            ));
4913        }
4914        let time_index_column =
4915            self.ctx
4916                .time_index_column
4917                .clone()
4918                .with_context(|| TimeIndexNotFoundSnafu {
4919                    table: self.ctx.table_name.clone().unwrap_or_default(),
4920                })?;
4921        // FIXME(ruihang): support multi fields
4922        let field_column = self
4923            .ctx
4924            .field_columns
4925            .first()
4926            .with_context(|| FunctionInvalidArgumentSnafu {
4927                fn_name: function.function_name().to_string(),
4928            })?
4929            .clone();
4930        // remove le column from tag columns
4931        self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME);
4932
4933        let fold = HistogramFold::new_with_operation(
4934            LE_COLUMN_NAME.to_string(),
4935            field_column,
4936            time_index_column,
4937            function,
4938            None,
4939            input_plan,
4940        )
4941        .context(DataFusionPlanningSnafu)?;
4942        Ok(LogicalPlan::Extension(Extension {
4943            node: Arc::new(fold),
4944        }))
4945    }
4946
4947    fn create_native_histogram_expr(
4948        &self,
4949        function: HistogramFoldOperation,
4950        field_column: &str,
4951    ) -> DfExpr {
4952        let field = DfExpr::Column(Column::from_name(field_column));
4953        let (func, args) = match function {
4954            HistogramFoldOperation::Quantile(quantile) => (
4955                Arc::new(NativeHistogramQuantile::scalar_udf_with_collector(
4956                    self.promql_annotations.clone(),
4957                )),
4958                vec![field, lit(f64::from(quantile))],
4959            ),
4960            HistogramFoldOperation::Fraction { lower, upper } => (
4961                Arc::new(NativeHistogramFraction::scalar_udf_with_collector(
4962                    self.promql_annotations.clone(),
4963                )),
4964                vec![field, lit(f64::from(lower)), lit(f64::from(upper))],
4965            ),
4966        };
4967        DfExpr::ScalarFunction(ScalarFunction { func, args })
4968    }
4969
4970    fn create_native_histogram_plan(
4971        &mut self,
4972        function: HistogramFoldOperation,
4973        input_plan: LogicalPlan,
4974    ) -> Result<LogicalPlan> {
4975        ensure!(
4976            self.ctx.field_columns.len() == 1,
4977            MultiFieldsNotSupportedSnafu {
4978                operator: function.function_name()
4979            },
4980        );
4981
4982        let field_column = self.ctx.field_columns[0].clone();
4983        let function_expr = self.create_native_histogram_expr(function, &field_column);
4984        let display_name = function_expr.schema_name().to_string();
4985        self.ctx.field_columns = vec![display_name.clone()];
4986
4987        let project_exprs = std::iter::once(self.create_time_index_column_expr()?)
4988            .chain(std::iter::once(function_expr.alias(display_name)))
4989            .chain(self.create_tag_column_exprs()?)
4990            .collect::<Vec<_>>();
4991
4992        LogicalPlanBuilder::from(input_plan)
4993            .project(project_exprs)
4994            .context(DataFusionPlanningSnafu)?
4995            .filter(self.create_empty_values_filter_expr(false)?)
4996            .context(DataFusionPlanningSnafu)?
4997            .build()
4998            .context(DataFusionPlanningSnafu)
4999    }
5000
5001    fn create_mixed_histogram_plan(
5002        &mut self,
5003        function: HistogramFoldOperation,
5004        input_plan: LogicalPlan,
5005        float_field: String,
5006        histogram_field: String,
5007    ) -> Result<LogicalPlan> {
5008        let time_index_column =
5009            self.ctx
5010                .time_index_column
5011                .clone()
5012                .with_context(|| TimeIndexNotFoundSnafu {
5013                    table: self.ctx.table_name.clone().unwrap_or_default(),
5014                })?;
5015        let tag_columns = self.ctx.tag_columns.clone();
5016        let folded = HistogramFold::new_with_operation(
5017            LE_COLUMN_NAME.to_string(),
5018            float_field.clone(),
5019            time_index_column.clone(),
5020            function,
5021            Some(histogram_field.clone()),
5022            input_plan,
5023        )
5024        .context(DataFusionPlanningSnafu)?;
5025        let record_collision = DfExpr::ScalarFunction(ScalarFunction {
5026            func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
5027                "vector contains a mix of classic and native histograms".to_string(),
5028                self.promql_annotations.clone(),
5029            )),
5030            args: vec![col(&float_field), col(&histogram_field)],
5031        });
5032        let keep = when(
5033            col(&float_field)
5034                .is_not_null()
5035                .and(col(&histogram_field).is_not_null()),
5036            record_collision,
5037        )
5038        .otherwise(lit(true))
5039        .context(DataFusionPlanningSnafu)?;
5040
5041        let native_expr = self.create_native_histogram_expr(function, &histogram_field);
5042        let output_field = native_expr.schema_name().to_string();
5043        let value = DfExpr::ScalarFunction(ScalarFunction {
5044            func: coalesce(),
5045            args: vec![col(&float_field), native_expr],
5046        });
5047        self.ctx.field_columns = vec![output_field.clone()];
5048        LogicalPlanBuilder::from(LogicalPlan::Extension(Extension {
5049            node: Arc::new(folded),
5050        }))
5051        .filter(keep)
5052        .context(DataFusionPlanningSnafu)?
5053        .project(
5054            std::iter::once(col(&time_index_column))
5055                .chain(std::iter::once(value.alias(output_field)))
5056                .chain(tag_columns.iter().map(col)),
5057        )
5058        .context(DataFusionPlanningSnafu)?
5059        .build()
5060        .context(DataFusionPlanningSnafu)
5061    }
5062
5063    /// Create a [SPECIAL_VECTOR_FUNCTION] plan
5064    async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result<LogicalPlan> {
5065        if args.args.len() != 1 {
5066            return FunctionInvalidArgumentSnafu {
5067                fn_name: SPECIAL_VECTOR_FUNCTION.to_string(),
5068            }
5069            .fail();
5070        }
5071        let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?;
5072
5073        // reuse `SPECIAL_TIME_FUNCTION` as name of time index column
5074        self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
5075        self.ctx.reset_table_name_and_schema();
5076        self.ctx.tag_columns = vec![];
5077        self.ctx.field_columns = vec![greptime_value().to_string()];
5078        Ok(LogicalPlan::Extension(Extension {
5079            node: Arc::new(
5080                EmptyMetric::new(
5081                    self.ctx.start,
5082                    self.ctx.end,
5083                    self.ctx.interval,
5084                    SPECIAL_TIME_FUNCTION.to_string(),
5085                    greptime_value().to_string(),
5086                    Some(lit),
5087                )
5088                .context(DataFusionPlanningSnafu)?,
5089            ),
5090        }))
5091    }
5092
5093    /// Create a [SCALAR_FUNCTION] plan
5094    async fn create_scalar_plan(
5095        &mut self,
5096        args: &PromFunctionArgs,
5097        query_engine_state: &QueryEngineState,
5098    ) -> Result<LogicalPlan> {
5099        ensure!(
5100            args.len() == 1,
5101            FunctionInvalidArgumentSnafu {
5102                fn_name: SCALAR_FUNCTION
5103            }
5104        );
5105        let input = self
5106            .prom_expr_to_plan(&args.args[0], query_engine_state)
5107            .await?;
5108        let input_schema = input.schema().clone();
5109        let alternative_samples =
5110            Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns);
5111        let histogram_fields = self
5112            .ctx
5113            .field_columns
5114            .iter()
5115            .filter(|field| Self::field_column_is_native_histogram(&input_schema, field))
5116            .count();
5117        ensure!(
5118            self.ctx.field_columns.len() == 1 || alternative_samples,
5119            MultiFieldsNotSupportedSnafu {
5120                operator: SCALAR_FUNCTION
5121            },
5122        );
5123        let scalar_field = self
5124            .ctx
5125            .field_columns
5126            .iter()
5127            .find(|field| !Self::field_column_is_native_histogram(&input_schema, field))
5128            .or_else(|| self.ctx.field_columns.first())
5129            .cloned()
5130            .with_context(|| FunctionInvalidArgumentSnafu {
5131                fn_name: SCALAR_FUNCTION,
5132            })?;
5133        let input = if histogram_fields == self.ctx.field_columns.len() {
5134            // scalar() ignores histogram samples. An empty input makes ScalarCalculate emit NaN
5135            // for every evaluation timestamp without attempting a Struct-to-Float64 cast.
5136            LogicalPlanBuilder::from(input)
5137                .filter(lit(false))
5138                .context(DataFusionPlanningSnafu)?
5139                .build()
5140                .context(DataFusionPlanningSnafu)?
5141        } else if histogram_fields > 0 {
5142            // A mixed vector contributes only its float samples to scalar().
5143            LogicalPlanBuilder::from(input)
5144                .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null())
5145                .context(DataFusionPlanningSnafu)?
5146                .build()
5147                .context(DataFusionPlanningSnafu)?
5148        } else {
5149            input
5150        };
5151        let scalar_plan = LogicalPlan::Extension(Extension {
5152            node: Arc::new(
5153                ScalarCalculate::new(
5154                    self.ctx.start,
5155                    self.ctx.end,
5156                    self.ctx.interval,
5157                    input,
5158                    self.ctx.time_index_column.as_ref().unwrap(),
5159                    &self.ctx.tag_columns,
5160                    &scalar_field,
5161                    self.ctx.table_name.as_deref(),
5162                )
5163                .context(PromqlPlanNodeSnafu)?,
5164            ),
5165        });
5166        // scalar plan have no tag columns
5167        self.ctx.tag_columns.clear();
5168        self.ctx.field_columns.clear();
5169        self.ctx
5170            .field_columns
5171            .push(scalar_plan.schema().field(1).name().clone());
5172        Ok(scalar_plan)
5173    }
5174
5175    /// Create a [SPECIAL_ABSENT_FUNCTION] plan
5176    async fn create_absent_plan(
5177        &mut self,
5178        args: &PromFunctionArgs,
5179        query_engine_state: &QueryEngineState,
5180    ) -> Result<LogicalPlan> {
5181        if args.args.len() != 1 {
5182            return FunctionInvalidArgumentSnafu {
5183                fn_name: SPECIAL_ABSENT_FUNCTION.to_string(),
5184            }
5185            .fail();
5186        }
5187        let input = self
5188            .prom_expr_to_plan(&args.args[0], query_engine_state)
5189            .await?;
5190
5191        let time_index_expr = self.create_time_index_column_expr()?;
5192        let first_field_expr =
5193            self.create_field_column_exprs()?
5194                .pop()
5195                .with_context(|| ValueNotFoundSnafu {
5196                    table: self.ctx.table_name.clone().unwrap_or_default(),
5197                })?;
5198        let first_value_expr = first_value(first_field_expr, vec![]);
5199
5200        let ordered_aggregated_input = LogicalPlanBuilder::from(input)
5201            .aggregate(
5202                vec![time_index_expr.clone()],
5203                vec![first_value_expr.clone()],
5204            )
5205            .context(DataFusionPlanningSnafu)?
5206            .sort(vec![time_index_expr.sort(true, false)])
5207            .context(DataFusionPlanningSnafu)?
5208            .build()
5209            .context(DataFusionPlanningSnafu)?;
5210
5211        let fake_labels = self
5212            .ctx
5213            .selector_matcher
5214            .iter()
5215            .filter_map(|matcher| match matcher.op {
5216                MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())),
5217                _ => None,
5218            })
5219            .collect::<Vec<_>>();
5220
5221        // Create the absent plan
5222        let absent_plan = LogicalPlan::Extension(Extension {
5223            node: Arc::new(
5224                Absent::try_new(
5225                    self.ctx.start,
5226                    self.ctx.end,
5227                    self.ctx.interval,
5228                    self.ctx.time_index_column.as_ref().unwrap().clone(),
5229                    self.ctx.field_columns[0].clone(),
5230                    fake_labels,
5231                    ordered_aggregated_input,
5232                )
5233                .context(DataFusionPlanningSnafu)?,
5234            ),
5235        });
5236
5237        Ok(absent_plan)
5238    }
5239
5240    /// Try to build a DataFusion Literal Expression from PromQL Expr, return
5241    /// `None` if the input is not a literal expression.
5242    fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
5243        match expr {
5244            PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
5245            PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
5246            PromExpr::VectorSelector(_)
5247            | PromExpr::MatrixSelector(_)
5248            | PromExpr::Extension(_)
5249            | PromExpr::Aggregate(_)
5250            | PromExpr::Subquery(_) => None,
5251            PromExpr::Call(Call { func, .. }) => {
5252                if func.name == SPECIAL_TIME_FUNCTION {
5253                    // For time() function, don't treat it as a literal
5254                    // Let it be handled as a regular function call
5255                    None
5256                } else {
5257                    None
5258                }
5259            }
5260            PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
5261            PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new(
5262                Self::try_build_literal_expr(expr)?,
5263            ))),
5264            PromExpr::Binary(PromBinaryExpr {
5265                lhs,
5266                rhs,
5267                op,
5268                modifier,
5269            }) => {
5270                let lhs = Self::try_build_literal_expr(lhs)?;
5271                let rhs = Self::try_build_literal_expr(rhs)?;
5272                let is_comparison_op = Self::is_token_a_comparison_op(*op);
5273                let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
5274                let expr = expr_builder(lhs, rhs).ok()?;
5275
5276                let should_return_bool = if let Some(m) = modifier {
5277                    m.return_bool
5278                } else {
5279                    false
5280                };
5281                if is_comparison_op && should_return_bool {
5282                    Some(DfExpr::Cast(Cast {
5283                        expr: Box::new(expr),
5284                        data_type: ArrowDataType::Float64,
5285                    }))
5286                } else {
5287                    Some(expr)
5288                }
5289            }
5290        }
5291    }
5292
5293    fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
5294        match expr {
5295            PromExpr::Call(Call { func, .. }) => {
5296                if func.name == SPECIAL_TIME_FUNCTION
5297                    && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
5298                {
5299                    Some(build_special_time_expr(time_index_col))
5300                } else {
5301                    None
5302                }
5303            }
5304            _ => None,
5305        }
5306    }
5307
5308    fn native_histogram_binary_expr(
5309        token: TokenType,
5310        lhs: DfExpr,
5311        lhs_is_histogram: bool,
5312        rhs: DfExpr,
5313        rhs_is_histogram: bool,
5314        filter_context: bool,
5315        promql_annotations: Option<PromqlAnnotationCollector>,
5316    ) -> Result<Option<DfExpr>> {
5317        if !lhs_is_histogram && !rhs_is_histogram {
5318            return Ok(None);
5319        }
5320
5321        let scalar_fn = |func: ScalarUdfDef, args| {
5322            DfExpr::ScalarFunction(ScalarFunction {
5323                func: Arc::new(func),
5324                args,
5325            })
5326        };
5327        let invalid_expr = || {
5328            let message = format!(
5329                "{}: dropped native histogram samples because this binary operation is not supported for native histograms",
5330                token
5331            );
5332            let func = if filter_context {
5333                NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone())
5334            } else {
5335                NativeHistogramDrop::float_null_udf(message, promql_annotations.clone())
5336            };
5337            let args = vec![lhs.clone(), rhs.clone()];
5338            scalar_fn(func, args)
5339        };
5340
5341        let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5342            (token::T_ADD, true, true) => scalar_fn(
5343                NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()),
5344                vec![lhs, rhs],
5345            ),
5346            (token::T_SUB, true, true) => scalar_fn(
5347                NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()),
5348                vec![lhs, rhs],
5349            ),
5350            (token::T_MUL, true, false) => {
5351                scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs])
5352            }
5353            (token::T_MUL, false, true) => {
5354                scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs])
5355            }
5356            (token::T_DIV, true, false) => {
5357                scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs])
5358            }
5359            (token::T_EQLC, true, true) => {
5360                scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs])
5361            }
5362            (token::T_NEQ, true, true) => {
5363                scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs])
5364            }
5365            _ => invalid_expr(),
5366        };
5367
5368        Ok(Some(expr))
5369    }
5370
5371    /// Return a lambda to build binary expression from token.
5372    /// Because some binary operator are function in DataFusion like `atan2` or `^`.
5373    #[allow(clippy::type_complexity)]
5374    fn prom_token_to_binary_expr_builder(
5375        token: TokenType,
5376    ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
5377        let cast_float = |expr| {
5378            if matches!(
5379                &expr,
5380                DfExpr::Cast(Cast {
5381                    data_type: ArrowDataType::Float64,
5382                    ..
5383                })
5384            ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
5385            {
5386                expr
5387            } else {
5388                DfExpr::Cast(Cast {
5389                    expr: Box::new(expr),
5390                    data_type: ArrowDataType::Float64,
5391                })
5392            }
5393        };
5394        match token.id() {
5395            token::T_ADD => Ok(Box::new(move |lhs, rhs| {
5396                Ok(cast_float(lhs) + cast_float(rhs))
5397            })),
5398            token::T_SUB => Ok(Box::new(move |lhs, rhs| {
5399                Ok(cast_float(lhs) - cast_float(rhs))
5400            })),
5401            token::T_MUL => Ok(Box::new(move |lhs, rhs| {
5402                Ok(cast_float(lhs) * cast_float(rhs))
5403            })),
5404            token::T_DIV => Ok(Box::new(move |lhs, rhs| {
5405                Ok(cast_float(lhs) / cast_float(rhs))
5406            })),
5407            token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
5408                Ok(cast_float(lhs) % cast_float(rhs))
5409            })),
5410            token::T_EQLC => Ok(Box::new(|lhs, rhs| Ok(lhs.eq(rhs)))),
5411            token::T_NEQ => Ok(Box::new(|lhs, rhs| Ok(lhs.not_eq(rhs)))),
5412            token::T_GTR => Ok(Box::new(|lhs, rhs| Ok(lhs.gt(rhs)))),
5413            token::T_LSS => Ok(Box::new(|lhs, rhs| Ok(lhs.lt(rhs)))),
5414            token::T_GTE => Ok(Box::new(|lhs, rhs| Ok(lhs.gt_eq(rhs)))),
5415            token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))),
5416            token::T_POW => Ok(Box::new(move |lhs, rhs| {
5417                Ok(DfExpr::ScalarFunction(ScalarFunction {
5418                    func: datafusion_functions::math::power(),
5419                    args: vec![cast_float(lhs), cast_float(rhs)],
5420                }))
5421            })),
5422            token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
5423                Ok(DfExpr::ScalarFunction(ScalarFunction {
5424                    func: datafusion_functions::math::atan2(),
5425                    args: vec![cast_float(lhs), cast_float(rhs)],
5426                }))
5427            })),
5428            _ => UnexpectedTokenSnafu { token }.fail(),
5429        }
5430    }
5431
5432    /// Check if the given op is a [comparison operator](https://prometheus.io/docs/prometheus/latest/querying/operators/#comparison-binary-operators).
5433    fn is_token_a_comparison_op(token: TokenType) -> bool {
5434        matches!(
5435            token.id(),
5436            token::T_EQLC
5437                | token::T_NEQ
5438                | token::T_GTR
5439                | token::T_LSS
5440                | token::T_GTE
5441                | token::T_LTE
5442        )
5443    }
5444
5445    /// Check if the given op is a set operator (UNION, INTERSECT and EXCEPT in SQL).
5446    fn is_token_a_set_op(token: TokenType) -> bool {
5447        matches!(
5448            token.id(),
5449            token::T_LAND // INTERSECT
5450                | token::T_LOR // UNION
5451                | token::T_LUNLESS // EXCEPT
5452        )
5453    }
5454
5455    fn align_binary_field_columns<'a>(
5456        left_schema: &DFSchemaRef,
5457        right_schema: &DFSchemaRef,
5458        left_field_columns: &'a [String],
5459        right_field_columns: &'a [String],
5460        op: TokenType,
5461        left_is_scalar: bool,
5462        right_is_scalar: bool,
5463    ) -> (
5464        Vec<(String, Vec<BinaryFieldPair<'a>>)>,
5465        Vec<BinaryFieldPair<'a>>,
5466    ) {
5467        // Mixed vectors store mutually exclusive float and histogram samples in two columns.
5468        // Retain each valid sample combination and group expressions by their output lane.
5469        let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns);
5470        let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns);
5471        let alternative_alignment = match (left_alternative, right_alternative) {
5472            (Some(output_names), Some(_)) => Some((
5473                output_names,
5474                left_field_columns
5475                    .iter()
5476                    .flat_map(|left| right_field_columns.iter().map(move |right| (left, right)))
5477                    .collect::<Vec<_>>(),
5478            )),
5479            (Some(output_names), None) if right_field_columns.len() == 1 => Some((
5480                output_names,
5481                left_field_columns
5482                    .iter()
5483                    .map(|left| (left, &right_field_columns[0]))
5484                    .collect::<Vec<_>>(),
5485            )),
5486            (None, Some(output_names)) if left_field_columns.len() == 1 => Some((
5487                output_names,
5488                right_field_columns
5489                    .iter()
5490                    .map(|right| (&left_field_columns[0], right))
5491                    .collect::<Vec<_>>(),
5492            )),
5493            _ => None,
5494        };
5495        let mut invalid_pairs = Vec::new();
5496        if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment {
5497            let mut float_pairs = Vec::new();
5498            let mut histogram_pairs = Vec::new();
5499            for (left, right) in field_pairs {
5500                let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left);
5501                let right_is_histogram =
5502                    Self::field_column_is_native_histogram(right_schema, right);
5503                match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) {
5504                    Some(false) => float_pairs.push((left, right)),
5505                    Some(true) => histogram_pairs.push((left, right)),
5506                    None => invalid_pairs.push((left, right)),
5507                }
5508            }
5509            if !float_pairs.is_empty() || !histogram_pairs.is_empty() {
5510                return (
5511                    [
5512                        (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)),
5513                        (!histogram_pairs.is_empty())
5514                            .then(|| (histogram_output.to_string(), histogram_pairs)),
5515                    ]
5516                    .into_iter()
5517                    .flatten()
5518                    .collect(),
5519                    invalid_pairs,
5520                );
5521            }
5522        }
5523
5524        if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 {
5525            return (
5526                right_field_columns
5527                    .iter()
5528                    .map(|right| (right.clone(), vec![(&left_field_columns[0], right)]))
5529                    .collect(),
5530                invalid_pairs,
5531            );
5532        }
5533        if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 {
5534            return (
5535                left_field_columns
5536                    .iter()
5537                    .map(|left| (left.clone(), vec![(left, &right_field_columns[0])]))
5538                    .collect(),
5539                invalid_pairs,
5540            );
5541        }
5542
5543        (
5544            left_field_columns
5545                .iter()
5546                .zip(right_field_columns.iter())
5547                .map(|(left, right)| (left.clone(), vec![(left, right)]))
5548                .collect(),
5549            invalid_pairs,
5550        )
5551    }
5552
5553    fn binary_result_is_histogram(
5554        token: TokenType,
5555        lhs_is_histogram: bool,
5556        rhs_is_histogram: bool,
5557    ) -> Option<bool> {
5558        match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5559            (_, false, false) => Some(false),
5560            (token::T_ADD | token::T_SUB, true, true)
5561            | (token::T_MUL, true, false)
5562            | (token::T_MUL, false, true)
5563            | (token::T_DIV, true, false) => Some(true),
5564            (token::T_EQLC | token::T_NEQ, true, true) => Some(false),
5565            _ => None,
5566        }
5567    }
5568
5569    fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
5570        plan.schema()
5571            .fields()
5572            .iter()
5573            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5574    }
5575
5576    fn is_empty_metric(plan: &LogicalPlan) -> bool {
5577        matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<EmptyMetric>())
5578    }
5579
5580    fn native_histogram_arrow_type() -> ArrowDataType {
5581        native_histogram_value_type().as_arrow_type()
5582    }
5583
5584    fn field_column_type<'a>(
5585        schema: &'a DFSchemaRef,
5586        field_column: &str,
5587    ) -> Option<&'a ArrowDataType> {
5588        schema
5589            .index_of_column_by_name(None, field_column)
5590            .map(|idx| schema.field(idx).data_type())
5591    }
5592
5593    fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool {
5594        Self::field_column_type(schema, field_column)
5595            .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type())
5596    }
5597
5598    fn field_columns_contain_native_histogram(
5599        schema: &DFSchemaRef,
5600        field_columns: &[String],
5601    ) -> bool {
5602        field_columns
5603            .iter()
5604            .any(|field| Self::field_column_is_native_histogram(schema, field))
5605    }
5606
5607    fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5608        Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5609            matches!(
5610                data_type,
5611                ArrowDataType::Dictionary(key_type, value_type)
5612                    if key_type.as_ref() == &ArrowDataType::Int64
5613                        && value_type.as_ref() == &ArrowDataType::Float64
5614            )
5615        })
5616    }
5617
5618    fn field_columns_are_alternative_samples(
5619        schema: &DFSchemaRef,
5620        field_columns: &[String],
5621    ) -> bool {
5622        Self::alternative_sample_columns(schema, field_columns).is_some()
5623    }
5624
5625    fn alternative_sample_columns<'a>(
5626        schema: &DFSchemaRef,
5627        field_columns: &'a [String],
5628    ) -> Option<(&'a str, &'a str)> {
5629        if field_columns.len() != 2 {
5630            return None;
5631        }
5632
5633        let canonical_float = field_columns.iter().find(|field| {
5634            field.as_str() == greptime_value()
5635                && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5636                    || Self::field_column_is_float_range(schema, field))
5637        });
5638        let canonical_histogram = field_columns.iter().find(|field| {
5639            field.as_str() == greptime_native_histogram()
5640                && (Self::field_column_is_native_histogram(schema, field)
5641                    || Self::field_column_is_native_histogram_range(schema, field))
5642        });
5643        if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) {
5644            return Some((float, histogram));
5645        }
5646
5647        let float = field_columns.iter().find(|field| {
5648            field.starts_with(OR_FLOAT_FIELD_PREFIX)
5649                && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5650                    || Self::field_column_is_float_range(schema, field))
5651        })?;
5652        let histogram = field_columns.iter().find(|field| {
5653            field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)
5654                && (Self::field_column_is_native_histogram(schema, field)
5655                    || Self::field_column_is_native_histogram_range(schema, field))
5656        })?;
5657        Some((float, histogram))
5658    }
5659
5660    fn alternative_sample_range_columns<'a>(
5661        schema: &DFSchemaRef,
5662        field_columns: &'a [String],
5663    ) -> Option<(&'a str, &'a str)> {
5664        Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| {
5665            Self::field_column_is_float_range(schema, float)
5666                && Self::field_column_is_native_histogram_range(schema, histogram)
5667        })
5668    }
5669
5670    fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5671        Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5672            matches!(
5673                data_type,
5674                ArrowDataType::Dictionary(key_type, value_type)
5675                    if key_type.as_ref() == &ArrowDataType::Int64
5676                        && value_type.as_ref() == &Self::native_histogram_arrow_type()
5677            )
5678        })
5679    }
5680
5681    fn all_field_columns_are_native_histograms(&self, schema: &DFSchemaRef) -> bool {
5682        !self.ctx.field_columns.is_empty()
5683            && self
5684                .ctx
5685                .field_columns
5686                .iter()
5687                .all(|field| Self::field_column_is_native_histogram(schema, field))
5688    }
5689
5690    fn all_field_columns_are_native_histogram_ranges(&self, schema: &DFSchemaRef) -> bool {
5691        !self.ctx.field_columns.is_empty()
5692            && self
5693                .ctx
5694                .field_columns
5695                .iter()
5696                .all(|field| Self::field_column_is_native_histogram_range(schema, field))
5697    }
5698
5699    fn optional_tsid_projection(
5700        schema: &DFSchemaRef,
5701        table_ref: Option<&TableReference>,
5702        keep_tsid: bool,
5703    ) -> Option<DfExpr> {
5704        keep_tsid.then_some(()).and_then(|_| {
5705            schema
5706                .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
5707                .ok()
5708                .map(|field| DfExpr::Column(field.into()))
5709        })
5710    }
5711
5712    fn binary_join_key_columns(
5713        &self,
5714        left_schema: &DFSchemaRef,
5715        right_schema: &DFSchemaRef,
5716        left_context: &PromPlannerContext,
5717        right_context: &PromPlannerContext,
5718        only_join_time_index: bool,
5719        modifier: &Option<BinModifier>,
5720    ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
5721        let has_tsid = |schema: &DFSchemaRef| {
5722            schema
5723                .fields()
5724                .iter()
5725                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5726        };
5727        let use_tsid_join = !only_join_time_index
5728            && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
5729            && left_context.use_tsid
5730            && right_context.use_tsid
5731            && has_tsid(left_schema)
5732            && has_tsid(right_schema);
5733
5734        let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
5735            (
5736                BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5737                BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5738            )
5739        } else {
5740            if only_join_time_index {
5741                (BTreeSet::new(), BTreeSet::new())
5742            } else {
5743                (
5744                    left_context
5745                        .tag_columns
5746                        .iter()
5747                        .cloned()
5748                        .collect::<BTreeSet<_>>(),
5749                    right_context
5750                        .tag_columns
5751                        .iter()
5752                        .cloned()
5753                        .collect::<BTreeSet<_>>(),
5754                )
5755            }
5756        };
5757
5758        if !use_tsid_join
5759            && let Some(modifier) = modifier
5760            && let Some(matching) = &modifier.matching
5761        {
5762            match matching {
5763                LabelModifier::Include(on) => {
5764                    let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5765                    left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
5766                    right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
5767                }
5768                LabelModifier::Exclude(ignoring) => {
5769                    for label in &ignoring.labels {
5770                        let _ = left_tag_columns.remove(label);
5771                        let _ = right_tag_columns.remove(label);
5772                    }
5773                }
5774            }
5775        }
5776
5777        let force_empty_join =
5778            !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
5779        if force_empty_join {
5780            let common_tag_columns = left_tag_columns
5781                .intersection(&right_tag_columns)
5782                .cloned()
5783                .collect::<BTreeSet<_>>();
5784            left_tag_columns = common_tag_columns.clone();
5785            right_tag_columns = common_tag_columns;
5786        }
5787
5788        Ok((left_tag_columns, right_tag_columns, force_empty_join))
5789    }
5790
5791    fn binary_modifier_preserves_tsid_join_key(
5792        &self,
5793        left_context: &PromPlannerContext,
5794        right_context: &PromPlannerContext,
5795        modifier: &Option<BinModifier>,
5796    ) -> bool {
5797        let Some(modifier) = modifier else {
5798            return true;
5799        };
5800
5801        if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
5802            return false;
5803        }
5804
5805        match &modifier.matching {
5806            None => true,
5807            Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
5808                !left_context.tag_columns.contains(label)
5809                    && !right_context.tag_columns.contains(label)
5810            }),
5811            Some(LabelModifier::Include(on)) => {
5812                let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5813                let left_labels = left_context
5814                    .tag_columns
5815                    .iter()
5816                    .cloned()
5817                    .collect::<BTreeSet<_>>();
5818                let right_labels = right_context
5819                    .tag_columns
5820                    .iter()
5821                    .cloned()
5822                    .collect::<BTreeSet<_>>();
5823
5824                on_labels == left_labels && on_labels == right_labels
5825            }
5826        }
5827    }
5828
5829    /// Build a inner join on time index column and tag columns to concat two logical plans.
5830    /// When `only_join_time_index == true` we only join on the time index, because these two plan may not have the same tag columns
5831    #[allow(clippy::too_many_arguments)]
5832    fn join_on_non_field_columns(
5833        &self,
5834        left: LogicalPlan,
5835        right: LogicalPlan,
5836        left_table_ref: TableReference,
5837        right_table_ref: TableReference,
5838        left_time_index_column: Option<String>,
5839        right_time_index_column: Option<String>,
5840        only_join_time_index: bool,
5841        modifier: &Option<BinModifier>,
5842        left_context: &PromPlannerContext,
5843        right_context: &PromPlannerContext,
5844    ) -> Result<LogicalPlan> {
5845        let (mut left_tag_columns, mut right_tag_columns, mut force_empty_join) = self
5846            .binary_join_key_columns(
5847                left.schema(),
5848                right.schema(),
5849                left_context,
5850                right_context,
5851                only_join_time_index,
5852                modifier,
5853            )?;
5854        let use_tsid_join = !only_join_time_index
5855            && !force_empty_join
5856            && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()])
5857            && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]);
5858        let (left, right) = if !only_join_time_index
5859            && !use_tsid_join
5860            && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier)
5861        {
5862            let mut aligned_left_context = left_context.clone();
5863            let mut aligned_right_context = right_context.clone();
5864            let (left, right, _) = Self::align_temporality_match_column(
5865                left,
5866                right,
5867                &mut aligned_left_context,
5868                &mut aligned_right_context,
5869            )?;
5870            (left_tag_columns, right_tag_columns, force_empty_join) = self
5871                .binary_join_key_columns(
5872                    left.schema(),
5873                    right.schema(),
5874                    &aligned_left_context,
5875                    &aligned_right_context,
5876                    false,
5877                    modifier,
5878                )?;
5879            (left, right)
5880        } else {
5881            (left, right)
5882        };
5883
5884        // push time index column if it exists
5885        if let (Some(left_time_index_column), Some(right_time_index_column)) =
5886            (left_time_index_column, right_time_index_column)
5887        {
5888            left_tag_columns.insert(left_time_index_column);
5889            right_tag_columns.insert(right_time_index_column);
5890        }
5891
5892        let right = LogicalPlanBuilder::from(right)
5893            .alias(right_table_ref)
5894            .context(DataFusionPlanningSnafu)?
5895            .build()
5896            .context(DataFusionPlanningSnafu)?;
5897
5898        // Inner Join on time index column to concat two operator
5899        LogicalPlanBuilder::from(left)
5900            .alias(left_table_ref)
5901            .context(DataFusionPlanningSnafu)?
5902            .join_detailed(
5903                right,
5904                JoinType::Inner,
5905                (
5906                    left_tag_columns
5907                        .into_iter()
5908                        .map(Column::from_name)
5909                        .collect::<Vec<_>>(),
5910                    right_tag_columns
5911                        .into_iter()
5912                        .map(Column::from_name)
5913                        .collect::<Vec<_>>(),
5914                ),
5915                force_empty_join.then_some(lit(false)),
5916                NullEquality::NullEqualsNull,
5917            )
5918            .context(DataFusionPlanningSnafu)?
5919            .build()
5920            .context(DataFusionPlanningSnafu)
5921    }
5922
5923    fn selected_binary_match_labels(
5924        left_context: &PromPlannerContext,
5925        right_context: &PromPlannerContext,
5926        modifier: &Option<BinModifier>,
5927    ) -> BTreeSet<String> {
5928        let mut labels = left_context
5929            .tag_columns
5930            .iter()
5931            .chain(&right_context.tag_columns)
5932            .cloned()
5933            .collect::<BTreeSet<_>>();
5934        if let Some(matching) = modifier
5935            .as_ref()
5936            .and_then(|modifier| modifier.matching.as_ref())
5937        {
5938            match matching {
5939                LabelModifier::Include(on) => {
5940                    labels = on
5941                        .labels
5942                        .iter()
5943                        .filter(|label| {
5944                            left_context.tag_columns.contains(label)
5945                                || right_context.tag_columns.contains(label)
5946                        })
5947                        .cloned()
5948                        .collect();
5949                }
5950                LabelModifier::Exclude(ignoring) => {
5951                    for label in &ignoring.labels {
5952                        labels.remove(label);
5953                    }
5954                }
5955            }
5956        }
5957        labels
5958    }
5959
5960    fn only_temporality_match_label_mismatches(
5961        left_context: &PromPlannerContext,
5962        right_context: &PromPlannerContext,
5963        modifier: &Option<BinModifier>,
5964    ) -> bool {
5965        let mut mismatches =
5966            Self::selected_binary_match_labels(left_context, right_context, modifier)
5967                .into_iter()
5968                .filter(|label| {
5969                    left_context.tag_columns.contains(label)
5970                        != right_context.tag_columns.contains(label)
5971                });
5972        matches!(
5973            (mismatches.next(), mismatches.next()),
5974            (Some(label), None) if label == OTLP_AGGREGATION_TEMPORALITY_LABEL
5975        )
5976    }
5977
5978    fn align_temporality_match_column(
5979        mut left: LogicalPlan,
5980        mut right: LogicalPlan,
5981        left_context: &mut PromPlannerContext,
5982        right_context: &mut PromPlannerContext,
5983    ) -> Result<(LogicalPlan, LogicalPlan, bool)> {
5984        let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
5985        let left_has_marker = left_context.tag_columns.iter().any(|tag| tag == marker);
5986        let (present, add_to_left) = if left_has_marker {
5987            (&left, false)
5988        } else {
5989            (&right, true)
5990        };
5991        let data_type = present
5992            .schema()
5993            .fields()
5994            .iter()
5995            .find(|field| field.name() == marker)
5996            .map(|field| field.data_type().clone())
5997            .with_context(|| ColumnNotFoundSnafu {
5998                col: marker.to_string(),
5999            })?;
6000        let null = Self::string_scalar_value(&data_type, None).with_context(|| {
6001            UnexpectedPlanExprSnafu {
6002                desc: format!("temporality match label {marker} must be a string"),
6003            }
6004        })?;
6005        let add_marker = |plan: LogicalPlan| {
6006            let visible = plan
6007                .schema()
6008                .iter()
6009                .map(|(qualifier, field)| {
6010                    DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6011                })
6012                .collect::<Vec<_>>();
6013            LogicalPlanBuilder::from(plan)
6014                .project(
6015                    visible
6016                        .into_iter()
6017                        .chain([DfExpr::Literal(null, None).alias(marker)]),
6018                )
6019                .context(DataFusionPlanningSnafu)?
6020                .build()
6021                .context(DataFusionPlanningSnafu)
6022        };
6023
6024        if add_to_left {
6025            left = add_marker(left)?;
6026            left_context.tag_columns.push(marker.to_string());
6027        } else {
6028            right = add_marker(right)?;
6029            right_context.tag_columns.push(marker.to_string());
6030        }
6031        Ok((left, right, add_to_left))
6032    }
6033
6034    fn normalized_match_key_expr(
6035        label: &str,
6036        field: Option<(Option<TableReference>, ArrowDataType)>,
6037        value_type: &ArrowDataType,
6038        internal_name: &str,
6039    ) -> DfExpr {
6040        let empty = Self::string_scalar_value(value_type, Some(String::new()))
6041            .expect("match label value type is a string");
6042        let expr = if let Some((qualifier, data_type)) = field {
6043            let column = DfExpr::Column(Column::new(qualifier, label));
6044            let column = if &data_type == value_type {
6045                column
6046            } else {
6047                DfExpr::Cast(Cast {
6048                    expr: Box::new(column),
6049                    data_type: value_type.clone(),
6050                })
6051            };
6052            DfExpr::ScalarFunction(ScalarFunction {
6053                func: coalesce(),
6054                args: vec![column, DfExpr::Literal(empty, None)],
6055            })
6056        } else {
6057            DfExpr::Literal(empty, None)
6058        };
6059        expr.alias(internal_name)
6060    }
6061
6062    fn is_zero_row_empty_relation(plan: &LogicalPlan) -> bool {
6063        // `produce_one_row` is used for input-free plans that still emit one row;
6064        // only the false case is a statically proven empty vector.
6065        matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row)
6066    }
6067
6068    /// Build a set operator (AND/OR/UNLESS)
6069    fn set_op_on_non_field_columns(
6070        &mut self,
6071        mut left: LogicalPlan,
6072        mut right: LogicalPlan,
6073        left_context: PromPlannerContext,
6074        right_context: PromPlannerContext,
6075        op: TokenType,
6076        modifier: &Option<BinModifier>,
6077    ) -> Result<LogicalPlan> {
6078        let left_tag_col_set = left_context
6079            .tag_columns
6080            .iter()
6081            .cloned()
6082            .collect::<HashSet<_>>();
6083        let right_tag_col_set = right_context
6084            .tag_columns
6085            .iter()
6086            .cloned()
6087            .collect::<HashSet<_>>();
6088
6089        if matches!(op.id(), token::T_LOR) {
6090            return self.or_operator(
6091                left,
6092                right,
6093                left_tag_col_set,
6094                right_tag_col_set,
6095                left_context,
6096                right_context,
6097                modifier,
6098            );
6099        }
6100
6101        if let Some(modifier) = modifier {
6102            ensure!(
6103                matches!(
6104                    modifier.card,
6105                    VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany
6106                ),
6107                UnsupportedVectorMatchSnafu {
6108                    name: modifier.card.clone(),
6109                },
6110            );
6111        }
6112
6113        let output_context = left_context.clone();
6114        let visible_left_schema = left.schema().clone();
6115        let mut left_context = left_context;
6116        let mut right_context = right_context;
6117        let added_marker_to_left = if Self::only_temporality_match_label_mismatches(
6118            &left_context,
6119            &right_context,
6120            modifier,
6121        ) {
6122            let aligned = Self::align_temporality_match_column(
6123                left,
6124                right,
6125                &mut left_context,
6126                &mut right_context,
6127            )?;
6128            left = aligned.0;
6129            right = aligned.1;
6130            aligned.2
6131        } else {
6132            false
6133        };
6134
6135        let mut left_tag_col_set = left_context
6136            .tag_columns
6137            .iter()
6138            .cloned()
6139            .collect::<BTreeSet<_>>();
6140        let mut right_tag_col_set = right_context
6141            .tag_columns
6142            .iter()
6143            .cloned()
6144            .collect::<BTreeSet<_>>();
6145        if let Some(matching) = modifier
6146            .as_ref()
6147            .and_then(|modifier| modifier.matching.as_ref())
6148        {
6149            match matching {
6150                LabelModifier::Include(on) => {
6151                    let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
6152                    left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect();
6153                    right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect();
6154                }
6155                LabelModifier::Exclude(ignoring) => {
6156                    for label in &ignoring.labels {
6157                        let _ = left_tag_col_set.remove(label);
6158                        let _ = right_tag_col_set.remove(label);
6159                    }
6160                }
6161            }
6162        }
6163        ensure!(
6164            left_tag_col_set == right_tag_col_set,
6165            CombineTableColumnMismatchSnafu {
6166                left: left_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6167                right: right_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6168            }
6169        );
6170
6171        let left_time_index = left_context.time_index_column.clone().unwrap();
6172        let right_time_index = right_context.time_index_column.clone().unwrap();
6173
6174        // alias right time index column if necessary
6175        if left_context.time_index_column != right_context.time_index_column {
6176            let right_project_exprs = right
6177                .schema()
6178                .fields()
6179                .iter()
6180                .map(|field| {
6181                    if field.name() == &right_time_index {
6182                        DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index)
6183                    } else {
6184                        DfExpr::Column(Column::from_name(field.name()))
6185                    }
6186                })
6187                .collect::<Vec<_>>();
6188
6189            right = LogicalPlanBuilder::from(right)
6190                .project(right_project_exprs)
6191                .context(DataFusionPlanningSnafu)?
6192                .build()
6193                .context(DataFusionPlanningSnafu)?;
6194        }
6195
6196        let join_keys = left_tag_col_set
6197            .into_iter()
6198            .chain([left_time_index])
6199            .collect::<Vec<_>>();
6200
6201        ensure!(
6202            left_context.field_columns.len() == 1
6203                || Self::field_columns_are_alternative_samples(
6204                    left.schema(),
6205                    &left_context.field_columns,
6206                ),
6207            MultiFieldsNotSupportedSnafu {
6208                operator: "AND/UNLESS operator"
6209            }
6210        );
6211        // Generate join plan.
6212        // All set operations in PromQL are "distinct"
6213        let result = match op.id() {
6214            token::T_LAND => LogicalPlanBuilder::from(left)
6215                .distinct()
6216                .context(DataFusionPlanningSnafu)?
6217                .join_detailed(
6218                    right,
6219                    JoinType::LeftSemi,
6220                    (join_keys.clone(), join_keys),
6221                    None,
6222                    NullEquality::NullEqualsNull,
6223                )
6224                .context(DataFusionPlanningSnafu)?
6225                .build()
6226                .context(DataFusionPlanningSnafu),
6227            token::T_LUNLESS => LogicalPlanBuilder::from(left)
6228                .distinct()
6229                .context(DataFusionPlanningSnafu)?
6230                .join_detailed(
6231                    right,
6232                    JoinType::LeftAnti,
6233                    (join_keys.clone(), join_keys),
6234                    None,
6235                    NullEquality::NullEqualsNull,
6236                )
6237                .context(DataFusionPlanningSnafu)?
6238                .build()
6239                .context(DataFusionPlanningSnafu),
6240            token::T_LOR => {
6241                // OR is handled at the beginning of this function, as it cannot
6242                // be expressed using JOIN like AND and UNLESS.
6243                unreachable!()
6244            }
6245            _ => UnexpectedTokenSnafu { token: op }.fail(),
6246        }?;
6247        let result = if added_marker_to_left {
6248            LogicalPlanBuilder::from(result)
6249                .project(visible_left_schema.iter().map(|(qualifier, field)| {
6250                    DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6251                }))
6252                .context(DataFusionPlanningSnafu)?
6253                .build()
6254                .context(DataFusionPlanningSnafu)?
6255        } else {
6256            result
6257        };
6258
6259        // AND/UNLESS preserve the complete left operand schema and metadata.
6260        self.ctx = output_context;
6261        Ok(result)
6262    }
6263
6264    fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> {
6265        match data_type {
6266            data_type if data_type.is_string() => Some(data_type),
6267            ArrowDataType::Dictionary(_, value_type) if value_type.is_string() => Some(value_type),
6268            _ => None,
6269        }
6270    }
6271
6272    fn string_scalar_value(
6273        data_type: &ArrowDataType,
6274        value: Option<String>,
6275    ) -> Option<ScalarValue> {
6276        match data_type {
6277            ArrowDataType::Utf8 => Some(ScalarValue::Utf8(value)),
6278            ArrowDataType::LargeUtf8 => Some(ScalarValue::LargeUtf8(value)),
6279            ArrowDataType::Utf8View => Some(ScalarValue::Utf8View(value)),
6280            ArrowDataType::Dictionary(key_type, value_type) => Some(ScalarValue::Dictionary(
6281                key_type.clone(),
6282                Box::new(Self::string_scalar_value(value_type, value)?),
6283            )),
6284            _ => None,
6285        }
6286    }
6287
6288    fn common_label_data_type(
6289        left: Option<&ArrowDataType>,
6290        right: Option<&ArrowDataType>,
6291    ) -> Option<ArrowDataType> {
6292        match (left, right) {
6293            (Some(left), Some(right)) if left == right => {
6294                Self::string_value_data_type(left).map(|_| left.clone())
6295            }
6296            (Some(left), Some(right)) => {
6297                let left_value_type = Self::string_value_data_type(left)?;
6298                let right_value_type = Self::string_value_data_type(right)?;
6299                // DataFusion projections can decode dictionaries, but do not encode plain strings
6300                // as dictionaries. Preserve the encoding only when both inputs already share it.
6301                match (left_value_type, right_value_type) {
6302                    (left, right) if left == right => Some(left.clone()),
6303                    (ArrowDataType::LargeUtf8, _) | (_, ArrowDataType::LargeUtf8) => {
6304                        Some(ArrowDataType::LargeUtf8)
6305                    }
6306                    (ArrowDataType::Utf8View, ArrowDataType::Utf8View) => {
6307                        Some(ArrowDataType::Utf8View)
6308                    }
6309                    _ => Some(ArrowDataType::Utf8),
6310                }
6311            }
6312            (Some(data_type), None) | (None, Some(data_type)) => {
6313                Self::string_value_data_type(data_type).cloned()
6314            }
6315            (None, None) => Some(ArrowDataType::Utf8),
6316        }
6317    }
6318
6319    // TODO(ruihang): change function name
6320    #[allow(clippy::too_many_arguments)]
6321    fn or_operator(
6322        &mut self,
6323        left: LogicalPlan,
6324        right: LogicalPlan,
6325        left_tag_cols_set: HashSet<String>,
6326        right_tag_cols_set: HashSet<String>,
6327        left_context: PromPlannerContext,
6328        right_context: PromPlannerContext,
6329        modifier: &Option<BinModifier>,
6330    ) -> Result<LogicalPlan> {
6331        let left_is_empty = Self::is_zero_row_empty_relation(&left);
6332        let right_is_empty = Self::is_zero_row_empty_relation(&right);
6333        match (left_is_empty, right_is_empty) {
6334            (true, false) => {
6335                self.ctx = right_context;
6336                return Ok(right);
6337            }
6338            (false, true) => {
6339                self.ctx = left_context;
6340                return Ok(left);
6341            }
6342            (true, true) => {
6343                self.ctx = left_context;
6344                return Ok(left);
6345            }
6346            (false, false) => {}
6347        }
6348
6349        ensure!(
6350            !left.schema().fields().is_empty() && !right.schema().fields().is_empty(),
6351            UnexpectedPlanExprSnafu {
6352                desc: "OR operator input has zero columns",
6353            }
6354        );
6355        let left_has_alternative_samples =
6356            Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns);
6357        let right_has_alternative_samples = Self::field_columns_are_alternative_samples(
6358            right.schema(),
6359            &right_context.field_columns,
6360        );
6361        ensure!(
6362            left_context.field_columns.len() == 1 || left_has_alternative_samples,
6363            MultiFieldsNotSupportedSnafu {
6364                operator: "OR operator"
6365            }
6366        );
6367        ensure!(
6368            right_context.field_columns.len() == 1 || right_has_alternative_samples,
6369            MultiFieldsNotSupportedSnafu {
6370                operator: "OR operator"
6371            }
6372        );
6373
6374        // prepare hash sets
6375        let all_tags = left_tag_cols_set
6376            .union(&right_tag_cols_set)
6377            .cloned()
6378            .collect::<HashSet<_>>();
6379        let left_qualifier = left.schema().qualified_field(0).0.cloned();
6380        let right_qualifier = right.schema().qualified_field(0).0.cloned();
6381        let left_qualifier_string = left_qualifier
6382            .as_ref()
6383            .map(|l| l.to_string())
6384            .unwrap_or_default();
6385        let right_qualifier_string = right_qualifier
6386            .as_ref()
6387            .map(|r| r.to_string())
6388            .unwrap_or_default();
6389        let left_time_index_column =
6390            left_context
6391                .time_index_column
6392                .clone()
6393                .with_context(|| TimeIndexNotFoundSnafu {
6394                    table: left_qualifier_string.clone(),
6395                })?;
6396        let right_time_index_column =
6397            right_context
6398                .time_index_column
6399                .clone()
6400                .with_context(|| TimeIndexNotFoundSnafu {
6401                    table: right_qualifier_string.clone(),
6402                })?;
6403        let native_histogram_type = Self::native_histogram_arrow_type();
6404        let is_numeric = |data_type: &ArrowDataType| {
6405            matches!(
6406                data_type,
6407                ArrowDataType::Int8
6408                    | ArrowDataType::Int16
6409                    | ArrowDataType::Int32
6410                    | ArrowDataType::Int64
6411                    | ArrowDataType::UInt8
6412                    | ArrowDataType::UInt16
6413                    | ArrowDataType::UInt32
6414                    | ArrowDataType::UInt64
6415                    | ArrowDataType::Float32
6416                    | ArrowDataType::Float64
6417            )
6418        };
6419        let left_fields = left_context
6420            .field_columns
6421            .iter()
6422            .map(|name| {
6423                left.schema()
6424                    .iter()
6425                    .find(|(_, field)| field.name() == name)
6426                    .map(|(qualifier, field)| {
6427                        (name.clone(), qualifier.cloned(), field.data_type().clone())
6428                    })
6429                    .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6430            })
6431            .collect::<Result<Vec<_>>>()?;
6432        let right_fields = right_context
6433            .field_columns
6434            .iter()
6435            .map(|name| {
6436                right
6437                    .schema()
6438                    .iter()
6439                    .find(|(_, field)| field.name() == name)
6440                    .map(|(qualifier, field)| {
6441                        (name.clone(), qualifier.cloned(), field.data_type().clone())
6442                    })
6443                    .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6444            })
6445            .collect::<Result<Vec<_>>>()?;
6446        let left_field = &left_fields[0];
6447        let right_field = &right_fields[0];
6448        let left_field_col = &left_field.0;
6449        let right_field_col = &right_field.0;
6450        let fields_are_samples = |fields: &[(String, Option<TableReference>, ArrowDataType)]| {
6451            fields.iter().all(|(_, _, data_type)| {
6452                is_numeric(data_type) || data_type == &native_histogram_type
6453            })
6454        };
6455        let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples {
6456            if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) {
6457                return UnexpectedPlanExprSnafu {
6458                    desc: format!(
6459                        "OR value fields have incompatible types: {:?} and {:?}",
6460                        left_fields
6461                            .iter()
6462                            .map(|(_, _, data_type)| data_type)
6463                            .collect::<Vec<_>>(),
6464                        right_fields
6465                            .iter()
6466                            .map(|(_, _, data_type)| data_type)
6467                            .collect::<Vec<_>>()
6468                    ),
6469                }
6470                .fail();
6471            }
6472            true
6473        } else {
6474            (left_field.2 == native_histogram_type && is_numeric(&right_field.2))
6475                || (right_field.2 == native_histogram_type && is_numeric(&left_field.2))
6476        };
6477        let target_field_type = if mixed_sample_types {
6478            // Mixed vectors use the existing response representation: one nullable float column
6479            // and one nullable native-histogram column.
6480            ArrowDataType::Float64
6481        } else if left_field.2 == right_field.2 {
6482            left_field.2.clone()
6483        } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) {
6484            ArrowDataType::Float64
6485        } else {
6486            return UnexpectedPlanExprSnafu {
6487                desc: format!(
6488                    "OR value fields have incompatible types: {:?} and {:?}",
6489                    left_field.2, right_field.2
6490                ),
6491            }
6492            .fail();
6493        };
6494        let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types {
6495            let mut reserved_names = left
6496                .schema()
6497                .fields()
6498                .iter()
6499                .chain(right.schema().fields().iter())
6500                .map(|field| field.name().clone())
6501                .collect::<HashSet<_>>();
6502            for (name, _, _) in left_fields.iter().chain(&right_fields) {
6503                reserved_names.remove(name);
6504            }
6505            reserved_names.extend(all_tags.iter().cloned());
6506            let unique_name = |prefix: &str, reserved_names: &mut HashSet<String>| {
6507                let mut index = 0;
6508                loop {
6509                    let name = format!("{prefix}{index}");
6510                    index += 1;
6511                    if reserved_names.insert(name.clone()) {
6512                        break name;
6513                    }
6514                }
6515            };
6516            let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names);
6517            let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names);
6518            (float_field, histogram_field)
6519        } else {
6520            (left_field_col.clone(), String::new())
6521        };
6522        let left_tag_types = left_tag_cols_set
6523            .iter()
6524            .map(|label| {
6525                left.schema()
6526                    .fields()
6527                    .iter()
6528                    .find(|field| field.name() == label)
6529                    .map(|field| (label.clone(), field.data_type().clone()))
6530                    .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
6531            })
6532            .collect::<Result<HashMap<_, _>>>()?;
6533        let right_tag_types = right_tag_cols_set
6534            .iter()
6535            .map(|label| {
6536                right
6537                    .schema()
6538                    .fields()
6539                    .iter()
6540                    .find(|field| field.name() == label)
6541                    .map(|field| (label.clone(), field.data_type().clone()))
6542                    .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
6543            })
6544            .collect::<Result<HashMap<_, _>>>()?;
6545        let mut target_tag_types = HashMap::with_capacity(all_tags.len());
6546        for label in &all_tags {
6547            let Some(data_type) =
6548                Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label))
6549            else {
6550                return UnexpectedPlanExprSnafu {
6551                    desc: format!(
6552                        "OR label {label} has incompatible types: {:?} and {:?}",
6553                        left_tag_types.get(label),
6554                        right_tag_types.get(label)
6555                    ),
6556                }
6557                .fail();
6558            };
6559            target_tag_types.insert(label.clone(), data_type);
6560        }
6561        let left_has_tsid = left
6562            .schema()
6563            .fields()
6564            .iter()
6565            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
6566        let right_has_tsid = right
6567            .schema()
6568            .fields()
6569            .iter()
6570            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
6571
6572        // step 0: fill all columns in output schema
6573        let mut all_columns_set = left
6574            .schema()
6575            .fields()
6576            .iter()
6577            .chain(right.schema().fields().iter())
6578            .map(|field| field.name().clone())
6579            .collect::<HashSet<_>>();
6580        // Keep `__tsid` only when both sides contain it, otherwise it may break schema alignment
6581        // (e.g. `unknown_metric or some_metric`).
6582        if !(left_has_tsid && right_has_tsid) {
6583            all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME);
6584        }
6585        // remove time index column
6586        all_columns_set.remove(&left_time_index_column);
6587        all_columns_set.remove(&right_time_index_column);
6588        if mixed_sample_types {
6589            for (name, _, _) in left_fields.iter().chain(&right_fields) {
6590                all_columns_set.remove(name);
6591            }
6592            all_columns_set.extend(all_tags.iter().cloned());
6593            all_columns_set.insert(mixed_float_field_col.clone());
6594            all_columns_set.insert(mixed_histogram_field_col.clone());
6595        } else if left_field_col != right_field_col {
6596            // remove field column in the right
6597            all_columns_set.remove(right_field_col);
6598        }
6599        let mut all_columns = all_columns_set.into_iter().collect::<Vec<_>>();
6600        // sort to ensure the generated schema is not volatile
6601        all_columns.sort_unstable();
6602        // use left time index column name as the result time index column name
6603        all_columns.insert(0, left_time_index_column.clone());
6604        let mut occupied_column_names = left
6605            .schema()
6606            .fields()
6607            .iter()
6608            .chain(right.schema().fields().iter())
6609            .map(|field| field.name().clone())
6610            .collect::<HashSet<_>>();
6611
6612        // step 1: align schema using project, fill non-exist columns with null
6613        let aligned_label_expr = |col: &String, source_types: &HashMap<String, ArrowDataType>| {
6614            let target_type = &target_tag_types[col];
6615            if let Some(source_type) = source_types.get(col) {
6616                let expr = DfExpr::Column(Column::new(None::<String>, col));
6617                if source_type == target_type {
6618                    expr
6619                } else {
6620                    DfExpr::Cast(Cast {
6621                        expr: Box::new(expr),
6622                        data_type: target_type.clone(),
6623                    })
6624                    .alias(col.clone())
6625                }
6626            } else {
6627                DfExpr::Literal(
6628                    Self::string_scalar_value(target_type, None)
6629                        .expect("target label type is a string"),
6630                    None,
6631                )
6632                .alias(col.clone())
6633            }
6634        };
6635        let null_histogram =
6636            ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?;
6637        let mixed_value_expr = |fields: &[(String, Option<TableReference>, ArrowDataType)],
6638                                output_col: &String| {
6639            if output_col == &mixed_float_field_col {
6640                if let Some((name, qualifier, data_type)) = fields
6641                    .iter()
6642                    .find(|(_, _, data_type)| is_numeric(data_type))
6643                {
6644                    let expr = DfExpr::Column(Column::new(qualifier.clone(), name));
6645                    if data_type == &ArrowDataType::Float64 {
6646                        expr.alias(output_col)
6647                    } else {
6648                        DfExpr::Cast(Cast {
6649                            expr: Box::new(expr),
6650                            data_type: ArrowDataType::Float64,
6651                        })
6652                        .alias(output_col)
6653                    }
6654                } else {
6655                    DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col)
6656                }
6657            } else {
6658                fields
6659                    .iter()
6660                    .find(|(_, _, data_type)| data_type == &native_histogram_type)
6661                    .map(|(name, qualifier, _)| {
6662                        DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col)
6663                    })
6664                    .unwrap_or_else(|| {
6665                        DfExpr::Literal(null_histogram.clone(), None).alias(output_col)
6666                    })
6667            }
6668        };
6669        let left_proj_exprs = all_columns.iter().map(|col| {
6670            if mixed_sample_types
6671                && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
6672            {
6673                mixed_value_expr(&left_fields, col)
6674            } else if !mixed_sample_types
6675                && col == left_field_col
6676                && left_field.2 != target_field_type
6677            {
6678                DfExpr::Cast(Cast {
6679                    expr: Box::new(DfExpr::Column(Column::new(
6680                        left_field.1.clone(),
6681                        left_field_col,
6682                    ))),
6683                    data_type: target_field_type.clone(),
6684                })
6685                .alias(left_field_col.clone())
6686            } else if target_tag_types.contains_key(col) {
6687                aligned_label_expr(col, &left_tag_types)
6688            } else {
6689                DfExpr::Column(Column::new(None::<String>, col))
6690            }
6691        });
6692        let right_time_index_expr = DfExpr::Column(Column::new(
6693            right_qualifier.clone(),
6694            right_time_index_column,
6695        ))
6696        .alias(left_time_index_column.clone());
6697        // The field column in right side may not have qualifier (it may be removed by join operation),
6698        // so we need to find it from the schema.
6699        // `skip(1)` to skip the time index column
6700        let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| {
6701            // expr
6702            if mixed_sample_types
6703                && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
6704            {
6705                mixed_value_expr(&right_fields, col)
6706            } else if !mixed_sample_types && col == left_field_col {
6707                let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col));
6708                if right_field.2 != target_field_type {
6709                    DfExpr::Cast(Cast {
6710                        expr: Box::new(expr),
6711                        data_type: target_field_type.clone(),
6712                    })
6713                    .alias(left_field_col.clone())
6714                } else if left_field_col != right_field_col {
6715                    expr.alias(left_field_col.clone())
6716                } else {
6717                    expr
6718                }
6719            } else if target_tag_types.contains_key(col) {
6720                aligned_label_expr(col, &right_tag_types)
6721            } else {
6722                DfExpr::Column(Column::new(None::<String>, col))
6723            }
6724        });
6725        let right_proj_exprs = [right_time_index_expr]
6726            .into_iter()
6727            .chain(right_proj_exprs_without_time_index);
6728
6729        let left_projected = LogicalPlanBuilder::from(left)
6730            .project(left_proj_exprs)
6731            .context(DataFusionPlanningSnafu)?
6732            .alias(left_qualifier_string.clone())
6733            .context(DataFusionPlanningSnafu)?
6734            .build()
6735            .context(DataFusionPlanningSnafu)?;
6736        let right_projected = LogicalPlanBuilder::from(right)
6737            .project(right_proj_exprs)
6738            .context(DataFusionPlanningSnafu)?
6739            .alias(right_qualifier_string.clone())
6740            .context(DataFusionPlanningSnafu)?
6741            .build()
6742            .context(DataFusionPlanningSnafu)?;
6743
6744        // step 2: compute match columns
6745        let mut match_columns = if let Some(modifier) = modifier
6746            && let Some(matching) = &modifier.matching
6747        {
6748            match matching {
6749                // keeps columns mentioned in `on`
6750                LabelModifier::Include(on) => on.labels.clone(),
6751                // removes columns memtioned in `ignoring`
6752                LabelModifier::Exclude(ignoring) => {
6753                    let ignoring = ignoring.labels.iter().cloned().collect::<HashSet<_>>();
6754                    all_tags.difference(&ignoring).cloned().collect()
6755                }
6756            }
6757        } else {
6758            all_tags.iter().cloned().collect()
6759        };
6760        // sort to ensure the generated plan is not volatile
6761        match_columns.sort_unstable();
6762        match_columns.dedup();
6763        occupied_column_names.extend(
6764            left_projected
6765                .schema()
6766                .fields()
6767                .iter()
6768                .chain(right_projected.schema().fields().iter())
6769                .map(|field| field.name().clone()),
6770        );
6771
6772        let visible_schema = left_projected.schema().clone();
6773        let visible_left_exprs = left_projected
6774            .schema()
6775            .iter()
6776            .map(|(qualifier, field)| {
6777                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6778            })
6779            .collect::<Vec<_>>();
6780        let visible_right_exprs = right_projected
6781            .schema()
6782            .iter()
6783            .map(|(qualifier, field)| {
6784                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6785            })
6786            .collect::<Vec<_>>();
6787        let mut left_match_exprs = Vec::with_capacity(match_columns.len());
6788        let mut right_match_exprs = Vec::with_capacity(match_columns.len());
6789        let mut next_internal_column = 0;
6790
6791        for label in &match_columns {
6792            let left_field = if left_tag_cols_set.contains(label) {
6793                Some(
6794                    left_projected
6795                        .schema()
6796                        .iter()
6797                        .find(|(_, field)| field.name() == label)
6798                        .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
6799                        .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
6800                )
6801            } else {
6802                None
6803            };
6804            let right_field = if right_tag_cols_set.contains(label) {
6805                Some(
6806                    right_projected
6807                        .schema()
6808                        .iter()
6809                        .find(|(_, field)| field.name() == label)
6810                        .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
6811                        .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
6812                )
6813            } else {
6814                None
6815            };
6816            let data_type = match (left_field.as_ref(), right_field.as_ref()) {
6817                (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => {
6818                    left_type.clone()
6819                }
6820                (Some((_, left_type)), Some((_, right_type))) => {
6821                    return UnexpectedPlanExprSnafu {
6822                        desc: format!(
6823                            "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}"
6824                        ),
6825                    }
6826                    .fail();
6827                }
6828                (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(),
6829                (None, None) => ArrowDataType::Utf8,
6830            };
6831            let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else {
6832                return UnexpectedPlanExprSnafu {
6833                    desc: format!("OR match label {label} must be a string"),
6834                }
6835                .fail();
6836            };
6837            let internal_name = loop {
6838                let name = format!("__promql_or_match_{next_internal_column}");
6839                next_internal_column += 1;
6840                if occupied_column_names.insert(name.clone()) {
6841                    break name;
6842                }
6843            };
6844            left_match_exprs.push(Self::normalized_match_key_expr(
6845                label,
6846                left_field,
6847                &value_type,
6848                &internal_name,
6849            ));
6850            right_match_exprs.push(Self::normalized_match_key_expr(
6851                label,
6852                right_field,
6853                &value_type,
6854                &internal_name,
6855            ));
6856        }
6857
6858        let left_augmented = LogicalPlanBuilder::from(left_projected)
6859            .project(visible_left_exprs.into_iter().chain(left_match_exprs))
6860            .context(DataFusionPlanningSnafu)?
6861            .build()
6862            .context(DataFusionPlanningSnafu)?;
6863        let right_augmented = LogicalPlanBuilder::from(right_projected)
6864            .project(visible_right_exprs.into_iter().chain(right_match_exprs))
6865            .context(DataFusionPlanningSnafu)?
6866            .build()
6867            .context(DataFusionPlanningSnafu)?;
6868
6869        // step 3: build `UnionDistinctOn` with normalized internal match keys.
6870        let visible_field_count = visible_schema.fields().len();
6871        let compare_key_indices =
6872            (visible_field_count..visible_field_count + match_columns.len()).collect::<Vec<_>>();
6873        let (time_qualifier, _) = visible_schema
6874            .iter()
6875            .find(|(_, field)| field.name() == &left_time_index_column)
6876            .with_context(|| TimeIndexNotFoundSnafu {
6877                table: left_qualifier_string.clone(),
6878            })?;
6879        let ts_col_idx = left_augmented
6880            .schema()
6881            .iter()
6882            .position(|(qualifier, field)| {
6883                qualifier == time_qualifier && field.name() == &left_time_index_column
6884            })
6885            .with_context(|| TimeIndexNotFoundSnafu {
6886                table: left_qualifier_string.clone(),
6887            })?;
6888        let union_distinct_on = UnionDistinctOn::try_new(
6889            left_augmented,
6890            right_augmented,
6891            compare_key_indices,
6892            ts_col_idx,
6893        )
6894        .context(DataFusionPlanningSnafu)?;
6895        let augmented_result = LogicalPlan::Extension(Extension {
6896            node: Arc::new(union_distinct_on),
6897        });
6898        let result = LogicalPlanBuilder::from(augmented_result)
6899            .project(visible_schema.iter().map(|(qualifier, field)| {
6900                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6901            }))
6902            .context(DataFusionPlanningSnafu)?
6903            .build()
6904            .context(DataFusionPlanningSnafu)?;
6905
6906        // step 4: update context
6907        let output_field_col = left_field_col.clone();
6908        let mut output_context = left_context;
6909        let mut visible_tags = all_tags.into_iter().collect::<Vec<_>>();
6910        visible_tags.sort_unstable();
6911        output_context.time_index_column = Some(left_time_index_column);
6912        output_context.tag_columns = visible_tags;
6913        output_context.field_columns = if mixed_sample_types {
6914            vec![mixed_float_field_col, mixed_histogram_field_col]
6915        } else {
6916            vec![output_field_col]
6917        };
6918        output_context.use_tsid = left_has_tsid && right_has_tsid;
6919        self.ctx = output_context;
6920
6921        Ok(result)
6922    }
6923
6924    /// Build a projection that project and perform operation expr for every value columns.
6925    /// Non-value columns (tag and timestamp) will be preserved in the projection.
6926    ///
6927    /// # Side effect
6928    ///
6929    /// This function will update the value columns in the context. Those new column names
6930    /// don't contains qualifier.
6931    fn projection_for_each_field_column<F>(
6932        &mut self,
6933        input: LogicalPlan,
6934        name_to_expr: F,
6935    ) -> Result<LogicalPlan>
6936    where
6937        F: FnMut(&String) -> Result<DfExpr>,
6938    {
6939        // Keep the generated float/histogram lane names while an element-wise operation
6940        // preserves both sample types, so downstream operators still recognize the pair.
6941        let preserve_field_names =
6942            Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
6943        let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
6944        // Derived labels can be unqualified even when the context still names the source table.
6945        let input_schema = input.schema().clone();
6946        let non_field_columns_iter = self
6947            .ctx
6948            .tag_columns
6949            .iter()
6950            .chain(self.ctx.time_index_column.iter())
6951            .map(|col| {
6952                input_schema
6953                    .qualified_field_with_name(table_ref.as_ref(), col)
6954                    .or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
6955                    .map(|field| DfExpr::Column(field.into()))
6956                    .context(DataFusionPlanningSnafu)
6957            });
6958        let tsid_iter =
6959            Self::optional_tsid_projection(input.schema(), table_ref.as_ref(), self.ctx.use_tsid)
6960                .into_iter()
6961                .map(Ok);
6962
6963        // build computation exprs
6964        let result_field_columns = self
6965            .ctx
6966            .field_columns
6967            .iter()
6968            .map(name_to_expr)
6969            .collect::<Result<Vec<_>>>()?;
6970
6971        // alias the computation exprs to remove qualifier
6972        if !preserve_field_names {
6973            self.ctx.field_columns = result_field_columns
6974                .iter()
6975                .map(|expr| expr.schema_name().to_string())
6976                .collect();
6977        }
6978        let field_columns_iter = result_field_columns
6979            .into_iter()
6980            .zip(self.ctx.field_columns.iter())
6981            .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
6982
6983        // chain non-field columns (unchanged) and field columns (applied computation then alias)
6984        let project_fields = non_field_columns_iter
6985            .chain(tsid_iter)
6986            .chain(field_columns_iter)
6987            .collect::<Result<Vec<_>>>()?;
6988
6989        LogicalPlanBuilder::from(input)
6990            .project(project_fields)
6991            .context(DataFusionPlanningSnafu)?
6992            .build()
6993            .context(DataFusionPlanningSnafu)
6994    }
6995
6996    /// Build a filter plan on one value column or a float/histogram alternative pair.
6997    fn filter_on_field_column<F>(&self, input: LogicalPlan, name_to_expr: F) -> Result<LogicalPlan>
6998    where
6999        F: FnMut(&String) -> Result<DfExpr>,
7000    {
7001        ensure!(
7002            self.ctx.field_columns.len() == 1
7003                || Self::field_columns_are_alternative_samples(
7004                    input.schema(),
7005                    &self.ctx.field_columns,
7006                ),
7007            UnsupportedExprSnafu {
7008                name: "filter on multi-value input"
7009            }
7010        );
7011
7012        let field_column_filters = self
7013            .ctx
7014            .field_columns
7015            .iter()
7016            .map(name_to_expr)
7017            .collect::<Result<Vec<_>>>()?;
7018        let field_column_filter =
7019            disjunction(field_column_filters).context(UnsupportedExprSnafu {
7020                name: "filter on empty input",
7021            })?;
7022
7023        LogicalPlanBuilder::from(input)
7024            .filter(field_column_filter)
7025            .context(DataFusionPlanningSnafu)?
7026            .build()
7027            .context(DataFusionPlanningSnafu)
7028    }
7029
7030    /// Generate an expr like `date_part("hour", <TIME_INDEX>)`. Caller should ensure the
7031    /// time index column in context is set
7032    fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
7033        let input_expr = datafusion::logical_expr::col(
7034            self.ctx
7035                .time_index_column
7036                .as_ref()
7037                // table name doesn't matters here
7038                .with_context(|| TimeIndexNotFoundSnafu {
7039                    table: "<doesn't matter>",
7040                })?,
7041        );
7042        let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
7043            func: datafusion_functions::datetime::date_part(),
7044            args: vec![date_part.lit(), input_expr],
7045        });
7046        Ok(fn_expr)
7047    }
7048
7049    fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
7050        let schema = plan.schema();
7051        if !schema
7052            .fields()
7053            .iter()
7054            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
7055        {
7056            return Ok(plan);
7057        }
7058
7059        // Preserve column qualifiers so downstream plan nodes can keep referencing
7060        // the columns by their original qualified names.
7061        let project_exprs = schema
7062            .iter()
7063            .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
7064            .map(|(qualifier, field)| {
7065                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7066            })
7067            .collect::<Vec<_>>();
7068
7069        LogicalPlanBuilder::from(plan)
7070            .project(project_exprs)
7071            .context(DataFusionPlanningSnafu)?
7072            .build()
7073            .context(DataFusionPlanningSnafu)
7074    }
7075
7076    /// Apply an alias to the query result by adding a projection with the alias name
7077    fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
7078        let fields_expr = self.create_field_column_exprs()?;
7079
7080        // TODO(dennis): how to support multi-value aliasing?
7081        ensure!(
7082            fields_expr.len() == 1,
7083            UnsupportedExprSnafu {
7084                name: "alias on multi-value result"
7085            }
7086        );
7087
7088        let project_fields = fields_expr
7089            .into_iter()
7090            .map(|expr| expr.alias(&alias_name))
7091            .chain(self.create_tag_column_exprs()?)
7092            .chain(Some(self.create_time_index_column_expr()?));
7093
7094        LogicalPlanBuilder::from(plan)
7095            .project(project_fields)
7096            .context(DataFusionPlanningSnafu)?
7097            .build()
7098            .context(DataFusionPlanningSnafu)
7099    }
7100}
7101
7102#[derive(Default, Debug)]
7103struct FunctionArgs {
7104    input: Option<PromExpr>,
7105    literals: Vec<DfExpr>,
7106}
7107
7108/// Represents different types of scalar functions supported in PromQL expressions.
7109/// Each variant defines how the function should be processed and what arguments it expects.
7110#[derive(Debug, Clone)]
7111enum ScalarFunc {
7112    /// DataFusion's registered(including built-in) scalar functions (e.g., abs, sqrt, round, clamp).
7113    /// These are passed through directly to DataFusion's execution engine.
7114    /// Processing: Simple argument insertion at the specified position.
7115    DataFusionBuiltin(Arc<ScalarUdfDef>),
7116    /// User-defined functions registered in DataFusion's function registry.
7117    /// Similar to DataFusionBuiltin but for custom functions not built into DataFusion.
7118    /// Processing: Direct pass-through with argument positioning.
7119    DataFusionUdf(Arc<ScalarUdfDef>),
7120    /// Native histogram helper UDFs. Non-histogram inputs are projected as NULL
7121    /// so the normal PromQL empty-value filter drops them.
7122    NativeHistogramUdf(Arc<ScalarUdfDef>),
7123    /// PromQL-specific functions that operate on time series data with temporal context.
7124    /// These functions require both timestamp ranges and values to perform calculations.
7125    /// Processing: Automatically injects timestamp_range and value columns as first arguments.
7126    /// Examples: idelta, irate, resets, changes, deriv, *_over_time function
7127    Udf(Arc<ScalarUdfDef>),
7128    /// PromQL functions requiring extrapolation calculations with explicit range information.
7129    /// These functions need to know the time range length to perform rate calculations.
7130    /// The second field contains the range length in milliseconds.
7131    /// Processing: Injects timestamp_range, value, time_index columns and appends range_length.
7132    /// Examples: increase, rate, delta
7133    // TODO(ruihang): maybe merge with Udf later
7134    ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
7135    /// Functions that generate expressions directly without external UDF calls.
7136    /// The expression is constructed during function matching and requires no additional processing.
7137    /// Examples: time(), minute(), hour(), month(), year() and other date/time extractors
7138    GeneratedExpr,
7139}
7140
7141#[cfg(test)]
7142mod test {
7143    use std::time::{Duration, UNIX_EPOCH};
7144
7145    use catalog::RegisterTableRequest;
7146    use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager};
7147    use common_base::Plugins;
7148    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
7149    use common_query::native_histogram::{
7150        CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array,
7151    };
7152    use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
7153    use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
7154    use common_query::test_util::DummyDecoder;
7155    use common_recordbatch::RecordBatch as GreptimeRecordBatch;
7156    use datafusion::arrow::array::{
7157        Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray,
7158    };
7159    use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema};
7160    use datafusion::arrow::record_batch::RecordBatch;
7161    use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider};
7162    use datafusion::datasource::memory::MemorySourceConfig;
7163    use datafusion::datasource::source::DataSourceExec;
7164    use datafusion::datasource::{MemTable, provider_as_source};
7165    use datafusion::execution::context::SessionContext;
7166    use datafusion::logical_expr::Extension;
7167    use datatypes::prelude::ConcreteDataType;
7168    use datatypes::schema::{ColumnSchema, Schema};
7169    use promql_parser::label::Labels;
7170    use promql_parser::parser;
7171    use session::context::QueryContext;
7172    use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
7173    use table::Table;
7174    use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder};
7175    use table::test_util::{EmptyTable, MemTable as GreptimeMemTable};
7176
7177    use super::*;
7178    use crate::QueryEngineContext;
7179    use crate::options::QueryOptions;
7180    use crate::parser::QueryLanguageParser;
7181    use crate::query_engine::DefaultSerializer;
7182
7183    mod delta;
7184
7185    fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> {
7186        if let LogicalPlan::Extension(Extension { node }) = plan
7187            && let Some(instant_manipulate) = node.as_any().downcast_ref::<InstantManipulate>()
7188        {
7189            return Some(instant_manipulate);
7190        }
7191
7192        plan.inputs().into_iter().find_map(find_instant_manipulate)
7193    }
7194
7195    fn build_query_engine_state() -> QueryEngineState {
7196        QueryEngineState::new(
7197            new_memory_catalog_manager().unwrap(),
7198            None,
7199            None,
7200            None,
7201            None,
7202            None,
7203            false,
7204            Plugins::default(),
7205            QueryOptions::default(),
7206        )
7207    }
7208
7209    #[test]
7210    fn common_label_type_preserves_only_shared_dictionary_encoding() {
7211        let dictionary = ArrowDataType::Dictionary(
7212            Box::new(ArrowDataType::UInt32),
7213            Box::new(ArrowDataType::Utf8),
7214        );
7215        let other_dictionary = ArrowDataType::Dictionary(
7216            Box::new(ArrowDataType::Int32),
7217            Box::new(ArrowDataType::Utf8),
7218        );
7219
7220        assert_eq!(
7221            Some(dictionary.clone()),
7222            PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary))
7223        );
7224        assert_eq!(
7225            Some(ArrowDataType::Utf8),
7226            PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8))
7227        );
7228        assert_eq!(
7229            Some(ArrowDataType::Utf8),
7230            PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary))
7231        );
7232        assert_eq!(
7233            Some(ArrowDataType::Utf8),
7234            PromPlanner::common_label_data_type(Some(&dictionary), None)
7235        );
7236    }
7237
7238    async fn build_optimized_promql_plan(
7239        table_provider: DfTableSourceProvider,
7240        eval_stmt: &EvalStmt,
7241    ) -> LogicalPlan {
7242        let state = build_query_engine_state();
7243        let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state)
7244            .await
7245            .unwrap();
7246        let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
7247        state
7248            .optimize_by_extension_rules(raw_plan, &context)
7249            .unwrap()
7250    }
7251
7252    async fn build_optimized_tsid_plan(
7253        query: &str,
7254        num_tag: usize,
7255        num_field: usize,
7256        end_secs: u64,
7257        lookback_secs: u64,
7258    ) -> String {
7259        let eval_stmt = EvalStmt {
7260            expr: parser::parse(query).unwrap(),
7261            start: UNIX_EPOCH,
7262            end: UNIX_EPOCH
7263                .checked_add(Duration::from_secs(end_secs))
7264                .unwrap(),
7265            interval: Duration::from_secs(5),
7266            lookback_delta: Duration::from_secs(lookback_secs),
7267        };
7268        let table_provider = build_test_table_provider_with_tsid(
7269            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7270            num_tag,
7271            num_field,
7272        )
7273        .await;
7274
7275        build_optimized_promql_plan(table_provider, &eval_stmt)
7276            .await
7277            .display_indent_schema()
7278            .to_string()
7279    }
7280
7281    async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) {
7282        let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
7283
7284        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
7285        assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0"));
7286        assert!(plan_str.contains("Distinct:"));
7287        assert!(plan_str.contains(expected_outer_agg), "{plan_str}");
7288        assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
7289    }
7290
7291    async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) {
7292        let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await;
7293        assert!(!plan_str.contains("Distinct:"), "{plan_str}");
7294    }
7295
7296    fn build_eval_stmt(expr: &str) -> EvalStmt {
7297        EvalStmt {
7298            expr: parser::parse(expr).unwrap(),
7299            start: UNIX_EPOCH,
7300            end: UNIX_EPOCH
7301                .checked_add(Duration::from_secs(100_000))
7302                .unwrap(),
7303            interval: Duration::from_secs(5),
7304            lookback_delta: Duration::from_secs(1),
7305        }
7306    }
7307
7308    enum DirectOrValue {
7309        Float64(f64),
7310        Int64(i64),
7311        NativeHistogram(NativeHistogram),
7312        Utf8(&'static str),
7313    }
7314
7315    impl DirectOrValue {
7316        fn data_type(&self) -> ArrowDataType {
7317            match self {
7318                Self::Float64(_) => ArrowDataType::Float64,
7319                Self::Int64(_) => ArrowDataType::Int64,
7320                Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(),
7321                Self::Utf8(_) => ArrowDataType::Utf8,
7322            }
7323        }
7324        fn array(&self) -> Arc<dyn Array> {
7325            match self {
7326                Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])),
7327                Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])),
7328                Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]),
7329                Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])),
7330            }
7331        }
7332    }
7333
7334    fn direct_or_histogram() -> NativeHistogram {
7335        NativeHistogram {
7336            schema: 0,
7337            zero_threshold: 0.0,
7338            sum: 1.0,
7339            reset_hint: CounterResetHint::Unknown,
7340            start_timestamp: None,
7341            custom_values: vec![],
7342            positive_spans: vec![],
7343            negative_spans: vec![],
7344            count: 1.0,
7345            zero_count: 1.0,
7346            positive_buckets: vec![],
7347            negative_buckets: vec![],
7348        }
7349    }
7350
7351    fn operator_metric_table(
7352        name: &str,
7353        table_id: u32,
7354        tag: &str,
7355        le: Option<&str>,
7356        value: DirectOrValue,
7357    ) -> table::TableRef {
7358        let value_type = match &value {
7359            DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(),
7360            DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(),
7361            DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(),
7362            DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(),
7363        };
7364        let tag_count = 1 + usize::from(le.is_some());
7365        let mut columns = vec![ColumnSchema::new(
7366            "tag".to_string(),
7367            ConcreteDataType::string_datatype(),
7368            false,
7369        )];
7370        if le.is_some() {
7371            columns.push(ColumnSchema::new(
7372                LE_COLUMN_NAME.to_string(),
7373                ConcreteDataType::string_datatype(),
7374                false,
7375            ));
7376        }
7377        columns.extend([
7378            ColumnSchema::new(
7379                "ts".to_string(),
7380                ConcreteDataType::timestamp_millisecond_datatype(),
7381                false,
7382            )
7383            .with_time_index(true),
7384            ColumnSchema::new("v".to_string(), value_type, true),
7385        ]);
7386        let schema = Arc::new(Schema::new(columns));
7387        let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc<dyn Array>];
7388        if let Some(le) = le {
7389            arrays.push(Arc::new(StringArray::from(vec![le])));
7390        }
7391        arrays.extend([
7392            Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc<dyn Array>,
7393            value.array(),
7394        ]);
7395        let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
7396        let backing = GreptimeMemTable::new_with_catalog(
7397            name,
7398            GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch),
7399            table_id,
7400            DEFAULT_CATALOG_NAME.to_string(),
7401            DEFAULT_SCHEMA_NAME.to_string(),
7402        );
7403        let value_index = tag_count + 1;
7404        let meta = TableMetaBuilder::empty()
7405            .schema(schema)
7406            .primary_key_indices((0..tag_count).collect())
7407            .value_indices(vec![value_index])
7408            .next_column_id((value_index + 1) as u32)
7409            .build()
7410            .unwrap();
7411        let info = Arc::new(
7412            TableInfoBuilder::default()
7413                .table_id(table_id)
7414                .name(name)
7415                .meta(meta)
7416                .build()
7417                .unwrap(),
7418        );
7419        Arc::new(Table::new(
7420            info,
7421            FilterPushDownType::Unsupported,
7422            backing.data_source(),
7423        ))
7424    }
7425
7426    fn operator_table_provider() -> DfTableSourceProvider {
7427        let catalog = MemoryCatalogManager::with_default_setup();
7428        let tables = [
7429            operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)),
7430            operator_metric_table(
7431                "lh",
7432                2_002,
7433                "b",
7434                None,
7435                DirectOrValue::NativeHistogram(direct_or_histogram()),
7436            ),
7437            operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)),
7438            operator_metric_table(
7439                "rh",
7440                2_004,
7441                "a",
7442                None,
7443                DirectOrValue::NativeHistogram(direct_or_histogram()),
7444            ),
7445            operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)),
7446            operator_metric_table(
7447                "bad_classic",
7448                2_006,
7449                "d",
7450                Some("broken"),
7451                DirectOrValue::Float64(1.0),
7452            ),
7453            operator_metric_table(
7454                "bad_native",
7455                2_007,
7456                "d",
7457                None,
7458                DirectOrValue::NativeHistogram(direct_or_histogram()),
7459            ),
7460        ];
7461        for table in tables {
7462            let info = table.table_info();
7463            catalog
7464                .register_table_sync(RegisterTableRequest {
7465                    catalog: DEFAULT_CATALOG_NAME.to_string(),
7466                    schema: DEFAULT_SCHEMA_NAME.to_string(),
7467                    table_name: info.name.clone(),
7468                    table_id: info.ident.table_id,
7469                    table,
7470                })
7471                .unwrap();
7472        }
7473        DfTableSourceProvider::new(
7474            catalog,
7475            false,
7476            QueryContext::arc(),
7477            DummyDecoder::arc(),
7478            false,
7479        )
7480    }
7481
7482    fn operator_eval_stmt(expr: &str) -> EvalStmt {
7483        let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap();
7484        EvalStmt {
7485            expr: parser::parse(expr).unwrap(),
7486            start: time,
7487            end: time,
7488            interval: Duration::from_secs(1),
7489            lookback_delta: Duration::from_secs(5),
7490        }
7491    }
7492
7493    struct DirectOrSource {
7494        name: &'static str,
7495        empty: bool,
7496        timestamp: i64,
7497        tags: Vec<(&'static str, Option<&'static str>)>,
7498        value: DirectOrValue,
7499    }
7500
7501    fn source(
7502        name: &'static str,
7503        empty: bool,
7504        timestamp: i64,
7505        tags: Vec<(&'static str, Option<&'static str>)>,
7506        value: DirectOrValue,
7507    ) -> DirectOrSource {
7508        DirectOrSource {
7509            name,
7510            empty,
7511            timestamp,
7512            tags,
7513            value,
7514        }
7515    }
7516
7517    fn tagged_source(
7518        name: &'static str,
7519        empty: bool,
7520        tag: (&'static str, Option<&'static str>),
7521        value: DirectOrValue,
7522    ) -> DirectOrSource {
7523        source(name, empty, 1, vec![("job", Some("job")), tag], value)
7524    }
7525
7526    fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource {
7527        source(name, true, 1, vec![("job", Some("job"))], value)
7528    }
7529
7530    fn table(source: &DirectOrSource) -> Arc<MemTable> {
7531        let mut fields = vec![Field::new(
7532            "ts",
7533            ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
7534            false,
7535        )];
7536        fields.extend(
7537            source
7538                .tags
7539                .iter()
7540                .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)),
7541        );
7542        fields.push(Field::new("v", source.value.data_type(), true));
7543        let schema = Arc::new(ArrowSchema::new(fields));
7544        let partitions = if source.empty {
7545            vec![vec![]]
7546        } else {
7547            let mut columns: Vec<Arc<dyn Array>> =
7548                vec![Arc::new(TimestampMillisecondArray::from(vec![
7549                    source.timestamp,
7550                ]))];
7551            columns.extend(
7552                source
7553                    .tags
7554                    .iter()
7555                    .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc<dyn Array>),
7556            );
7557            columns.push(source.value.array());
7558            vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]]
7559        };
7560        Arc::new(MemTable::try_new(schema, partitions).unwrap())
7561    }
7562
7563    fn scan(source: &DirectOrSource) -> LogicalPlan {
7564        LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None)
7565            .unwrap()
7566            .build()
7567            .unwrap()
7568    }
7569
7570    fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext {
7571        PromPlannerContext {
7572            table_name: Some(qualifier.to_string()),
7573            time_index_column: Some("ts".to_string()),
7574            field_columns: vec![field.to_string()],
7575            tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(),
7576            ..Default::default()
7577        }
7578    }
7579
7580    fn or_modifier(expr: &str) -> Option<BinModifier> {
7581        let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else {
7582            unreachable!()
7583        };
7584        expr.modifier
7585    }
7586
7587    async fn plan_direct_or(
7588        left: LogicalPlan,
7589        right: LogicalPlan,
7590        left_context: PromPlannerContext,
7591        right_context: PromPlannerContext,
7592        modifier: &Option<BinModifier>,
7593    ) -> LogicalPlan {
7594        let table_provider = build_test_table_provider_with_fields(
7595            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7596            &[],
7597        )
7598        .await;
7599        let mut planner = PromPlanner {
7600            table_provider,
7601            ctx: PromPlannerContext::default(),
7602            promql_annotations: None,
7603        };
7604        planner
7605            .or_operator(
7606                left,
7607                right,
7608                left_context.tag_columns.iter().cloned().collect(),
7609                right_context.tag_columns.iter().cloned().collect(),
7610                left_context,
7611                right_context,
7612                modifier,
7613            )
7614            .unwrap()
7615    }
7616
7617    async fn execute(
7618        plan: LogicalPlan,
7619        state: &QueryEngineState,
7620    ) -> (LogicalPlan, Vec<RecordBatch>) {
7621        let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
7622        let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
7623        let physical = state
7624            .session_state()
7625            .create_physical_plan(&optimized)
7626            .await
7627            .unwrap();
7628        let batches =
7629            datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
7630                .await
7631                .unwrap();
7632        (optimized, batches)
7633    }
7634
7635    async fn run(
7636        left: &DirectOrSource,
7637        right: &DirectOrSource,
7638        left_context: PromPlannerContext,
7639        right_context: PromPlannerContext,
7640        modifier: &Option<BinModifier>,
7641    ) -> (LogicalPlan, Vec<RecordBatch>) {
7642        let plan = plan_direct_or(
7643            scan(left),
7644            scan(right),
7645            left_context,
7646            right_context,
7647            modifier,
7648        )
7649        .await;
7650        execute(plan, &build_query_engine_state()).await
7651    }
7652
7653    async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) {
7654        let sample = |histogram: bool| {
7655            if histogram {
7656                DirectOrValue::NativeHistogram(direct_or_histogram())
7657            } else {
7658                DirectOrValue::Float64(1.25)
7659            }
7660        };
7661        let left = tagged_source(
7662            "lhs",
7663            false,
7664            (
7665                "k",
7666                Some(if histogram_on_left {
7667                    "histogram"
7668                } else {
7669                    "float"
7670                }),
7671            ),
7672            sample(histogram_on_left),
7673        );
7674        let right = tagged_source(
7675            "rhs",
7676            false,
7677            (
7678                "k",
7679                Some(if histogram_on_left {
7680                    "float"
7681                } else {
7682                    "histogram"
7683                }),
7684            ),
7685            sample(!histogram_on_left),
7686        );
7687        let table_provider = build_test_table_provider_with_fields(
7688            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7689            &[],
7690        )
7691        .await;
7692        let mut planner = PromPlanner {
7693            table_provider,
7694            ctx: PromPlannerContext::default(),
7695            promql_annotations: None,
7696        };
7697        let left_context = direct_or_context("lhs", &["job", "k"], "v");
7698        let right_context = direct_or_context("rhs", &["job", "k"], "v");
7699        let plan = planner
7700            .or_operator(
7701                scan(&left),
7702                scan(&right),
7703                left_context.tag_columns.iter().cloned().collect(),
7704                right_context.tag_columns.iter().cloned().collect(),
7705                left_context,
7706                right_context,
7707                &or_modifier("lhs or on(k) rhs"),
7708            )
7709            .unwrap();
7710        (planner, plan)
7711    }
7712
7713    async fn mixed_aggregate_input(histograms: Vec<NativeHistogram>) -> (PromPlanner, LogicalPlan) {
7714        let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0");
7715        let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0");
7716        let row_count = histograms.len() + 1;
7717        let schema = Arc::new(ArrowSchema::new(vec![
7718            Field::new(
7719                "ts",
7720                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
7721                false,
7722            ),
7723            Field::new("k", ArrowDataType::Utf8, false),
7724            Field::new(&float_field, ArrowDataType::Float64, true),
7725            Field::new(
7726                &histogram_field,
7727                native_histogram_value_type().as_arrow_type(),
7728                true,
7729            ),
7730        ]));
7731        let mut histogram_values = Vec::with_capacity(row_count);
7732        histogram_values.push(None);
7733        histogram_values.extend(histograms.into_iter().map(Some));
7734        let batch = RecordBatch::try_new(
7735            schema.clone(),
7736            vec![
7737                Arc::new(TimestampMillisecondArray::from(vec![1; row_count])),
7738                Arc::new(StringArray::from_iter_values(
7739                    (0..row_count).map(|row| format!("kind_{row}")),
7740                )),
7741                Arc::new(Float64Array::from_iter(
7742                    (0..row_count).map(|row| (row == 0).then_some(1.25)),
7743                )),
7744                build_histogram_array(&histogram_values),
7745            ],
7746        )
7747        .unwrap();
7748        let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
7749        let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
7750            .unwrap()
7751            .build()
7752            .unwrap();
7753        let table_provider = build_test_table_provider_with_fields(
7754            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7755            &[],
7756        )
7757        .await;
7758        let planner = PromPlanner {
7759            table_provider,
7760            ctx: PromPlannerContext {
7761                table_name: Some("mixed".to_string()),
7762                time_index_column: Some("ts".to_string()),
7763                field_columns: vec![float_field, histogram_field],
7764                tag_columns: vec!["k".to_string()],
7765                ..Default::default()
7766            },
7767            promql_annotations: None,
7768        };
7769        (planner, plan)
7770    }
7771
7772    fn assert_no_internal_or_keys(schema: &DFSchema) {
7773        assert!(
7774            schema
7775                .fields()
7776                .iter()
7777                .all(|field| !field.name().starts_with("__promql_or_match_")),
7778            "{schema:?}"
7779        );
7780    }
7781
7782    fn values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
7783        batches
7784            .iter()
7785            .flat_map(|batch| {
7786                batch
7787                    .column_by_name(column)
7788                    .unwrap()
7789                    .as_any()
7790                    .downcast_ref::<Float64Array>()
7791                    .unwrap()
7792                    .iter()
7793                    .flatten()
7794            })
7795            .collect()
7796    }
7797
7798    fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
7799        batches
7800            .iter()
7801            .flat_map(|batch| {
7802                let values = datafusion::arrow::compute::cast(
7803                    batch.column_by_name(column).unwrap(),
7804                    &ArrowDataType::Float64,
7805                )
7806                .unwrap();
7807                values
7808                    .as_any()
7809                    .downcast_ref::<Float64Array>()
7810                    .unwrap()
7811                    .iter()
7812                    .flatten()
7813                    .collect::<Vec<_>>()
7814            })
7815            .collect()
7816    }
7817
7818    fn histograms(batches: &[RecordBatch], column: &str) -> Vec<NativeHistogram> {
7819        batches
7820            .iter()
7821            .flat_map(|batch| {
7822                let values = batch
7823                    .column_by_name(column)
7824                    .unwrap()
7825                    .as_any()
7826                    .downcast_ref::<datafusion::arrow::array::StructArray>()
7827                    .unwrap();
7828                (0..values.len()).filter_map(|row| {
7829                    common_query::native_histogram::read_histogram(values, row).unwrap()
7830                })
7831            })
7832            .collect()
7833    }
7834
7835    fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option<String>)> {
7836        let mut rows = batches
7837            .iter()
7838            .flat_map(|batch| {
7839                let values = batch
7840                    .column_by_name("v")
7841                    .unwrap()
7842                    .as_any()
7843                    .downcast_ref::<Float64Array>()
7844                    .unwrap();
7845                let labels = batch
7846                    .column_by_name("k")
7847                    .map(|column| column.as_any().downcast_ref::<StringArray>().unwrap());
7848                (0..batch.num_rows()).map(move |i| {
7849                    (
7850                        values.value(i),
7851                        labels.and_then(|labels| {
7852                            (!labels.is_null(i)).then(|| labels.value(i).to_string())
7853                        }),
7854                    )
7855                })
7856            })
7857            .collect::<Vec<_>>();
7858        rows.sort_by(|left, right| left.0.total_cmp(&right.0));
7859        rows
7860    }
7861
7862    fn matrix_source(
7863        name: &'static str,
7864        k: Option<Option<&'static str>>,
7865        timestamp: i64,
7866        value: f64,
7867    ) -> DirectOrSource {
7868        let mut tags = vec![("job", Some("job"))];
7869        if let Some(k) = k {
7870            tags.push(("k", k));
7871        }
7872        source(name, false, timestamp, tags, DirectOrValue::Float64(value))
7873    }
7874
7875    fn matrix_context(name: &str, k: Option<Option<&str>>) -> PromPlannerContext {
7876        direct_or_context(
7877            name,
7878            if k.is_some() { &["job", "k"] } else { &["job"] },
7879            "v",
7880        )
7881    }
7882
7883    async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider {
7884        build_test_table_provider_with_fields(
7885            &[
7886                (
7887                    DEFAULT_SCHEMA_NAME.to_string(),
7888                    "non_existent_histogram_bucket".to_string(),
7889                ),
7890                (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()),
7891            ],
7892            &["pod", "instance"],
7893        )
7894        .await
7895    }
7896
7897    fn assert_normal_metric_schema(plan: &LogicalPlan) {
7898        let fields = plan.schema().fields();
7899        assert_eq!(fields.len(), 4, "{fields:?}");
7900        assert!(
7901            fields.iter().any(|field| field.name() == "pod"),
7902            "{fields:?}"
7903        );
7904        assert!(
7905            fields.iter().any(|field| field.name() == "instance"),
7906            "{fields:?}"
7907        );
7908        assert!(
7909            fields
7910                .iter()
7911                .any(|field| field.name() == greptime_timestamp()),
7912            "{fields:?}"
7913        );
7914        assert!(
7915            fields.iter().any(|field| {
7916                field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64
7917            }),
7918            "{fields:?}"
7919        );
7920    }
7921
7922    async fn build_test_table_provider_with_distinct_tags(
7923        table_tags: &[(&str, &[&str])],
7924    ) -> DfTableSourceProvider {
7925        let catalog_list = MemoryCatalogManager::with_default_setup();
7926        for (table_name, tags) in table_tags {
7927            let mut columns = tags
7928                .iter()
7929                .map(|tag| {
7930                    ColumnSchema::new(
7931                        (*tag).to_string(),
7932                        ConcreteDataType::string_datatype(),
7933                        false,
7934                    )
7935                })
7936                .collect::<Vec<_>>();
7937            columns.push(
7938                ColumnSchema::new(
7939                    greptime_timestamp().to_string(),
7940                    ConcreteDataType::timestamp_millisecond_datatype(),
7941                    false,
7942                )
7943                .with_time_index(true),
7944            );
7945            columns.push(ColumnSchema::new(
7946                greptime_value().to_string(),
7947                ConcreteDataType::float64_datatype(),
7948                true,
7949            ));
7950            let table_meta = TableMetaBuilder::empty()
7951                .schema(Arc::new(Schema::new(columns)))
7952                .primary_key_indices((0..tags.len()).collect())
7953                .next_column_id(1024)
7954                .build()
7955                .unwrap();
7956            let table_info = TableInfoBuilder::default()
7957                .name((*table_name).to_string())
7958                .meta(table_meta)
7959                .build()
7960                .unwrap();
7961
7962            assert!(
7963                catalog_list
7964                    .register_table_sync(RegisterTableRequest {
7965                        catalog: DEFAULT_CATALOG_NAME.to_string(),
7966                        schema: DEFAULT_SCHEMA_NAME.to_string(),
7967                        table_name: (*table_name).to_string(),
7968                        table_id: 1024,
7969                        table: EmptyTable::from_table_info(&table_info),
7970                    })
7971                    .is_ok()
7972            );
7973        }
7974
7975        DfTableSourceProvider::new(
7976            catalog_list,
7977            false,
7978            QueryContext::arc(),
7979            DummyDecoder::arc(),
7980            false,
7981        )
7982    }
7983
7984    fn contains_histogram_fold(plan: &LogicalPlan) -> bool {
7985        matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<HistogramFold>())
7986            || plan.inputs().into_iter().any(contains_histogram_fold)
7987    }
7988
7989    async fn build_set_op_context_table_provider() -> DfTableSourceProvider {
7990        build_test_table_provider_with_distinct_tags(&[
7991            ("bucket_metric", &["job", "le"]),
7992            ("normal_metric", &["job"]),
7993            ("fallback_metric", &["instance"]),
7994        ])
7995        .await
7996    }
7997
7998    async fn build_or_context_table_provider() -> DfTableSourceProvider {
7999        build_test_table_provider_with_distinct_tags(&[
8000            ("normal_metric", &["job"]),
8001            ("other_metric", &["instance"]),
8002            ("non_hist_metric", &["instance"]),
8003        ])
8004        .await
8005    }
8006
8007    async fn optimize_and_create_physical_plan(
8008        state: &QueryEngineState,
8009        plan: LogicalPlan,
8010    ) -> (
8011        LogicalPlan,
8012        Arc<dyn datafusion::physical_plan::ExecutionPlan>,
8013    ) {
8014        let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
8015        let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
8016        let physical = state
8017            .session_state()
8018            .create_physical_plan(&optimized)
8019            .await
8020            .unwrap();
8021        (optimized, physical)
8022    }
8023
8024    async fn build_test_table_provider(
8025        table_name_tuples: &[(String, String)],
8026        num_tag: usize,
8027        num_field: usize,
8028    ) -> DfTableSourceProvider {
8029        let catalog_list = MemoryCatalogManager::with_default_setup();
8030        for (schema_name, table_name) in table_name_tuples {
8031            let mut columns = vec![];
8032            for i in 0..num_tag {
8033                columns.push(ColumnSchema::new(
8034                    format!("tag_{i}"),
8035                    ConcreteDataType::string_datatype(),
8036                    false,
8037                ));
8038            }
8039            columns.push(
8040                ColumnSchema::new(
8041                    "timestamp".to_string(),
8042                    ConcreteDataType::timestamp_millisecond_datatype(),
8043                    false,
8044                )
8045                .with_time_index(true),
8046            );
8047            for i in 0..num_field {
8048                columns.push(ColumnSchema::new(
8049                    format!("field_{i}"),
8050                    ConcreteDataType::float64_datatype(),
8051                    true,
8052                ));
8053            }
8054            let schema = Arc::new(Schema::new(columns));
8055            let table_meta = TableMetaBuilder::empty()
8056                .schema(schema)
8057                .primary_key_indices((0..num_tag).collect())
8058                .value_indices((num_tag + 1..num_tag + 1 + num_field).collect())
8059                .next_column_id(1024)
8060                .build()
8061                .unwrap();
8062            let table_info = TableInfoBuilder::default()
8063                .name(table_name.clone())
8064                .meta(table_meta)
8065                .build()
8066                .unwrap();
8067            let table = EmptyTable::from_table_info(&table_info);
8068
8069            assert!(
8070                catalog_list
8071                    .register_table_sync(RegisterTableRequest {
8072                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8073                        schema: schema_name.clone(),
8074                        table_name: table_name.clone(),
8075                        table_id: 1024,
8076                        table,
8077                    })
8078                    .is_ok()
8079            );
8080        }
8081
8082        DfTableSourceProvider::new(
8083            catalog_list,
8084            false,
8085            QueryContext::arc(),
8086            DummyDecoder::arc(),
8087            false,
8088        )
8089    }
8090
8091    async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8092        build_test_native_histogram_table_provider_with_marker(table_name, false).await
8093    }
8094
8095    async fn build_test_native_histogram_table_provider_with_marker(
8096        table_name: &str,
8097        temporality_marker: bool,
8098    ) -> DfTableSourceProvider {
8099        let catalog_list = MemoryCatalogManager::with_default_setup();
8100        let mut columns = vec![
8101            ColumnSchema::new(
8102                "tag_0".to_string(),
8103                ConcreteDataType::string_datatype(),
8104                false,
8105            ),
8106            ColumnSchema::new(
8107                LE_COLUMN_NAME.to_string(),
8108                ConcreteDataType::string_datatype(),
8109                true,
8110            ),
8111        ];
8112        if temporality_marker {
8113            columns.push(ColumnSchema::new(
8114                OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8115                ConcreteDataType::string_datatype(),
8116                true,
8117            ));
8118        }
8119        let tag_count = columns.len();
8120        columns.extend([
8121            ColumnSchema::new(
8122                "timestamp".to_string(),
8123                ConcreteDataType::timestamp_millisecond_datatype(),
8124                false,
8125            )
8126            .with_time_index(true),
8127            ColumnSchema::new(
8128                greptime_native_histogram().to_string(),
8129                native_histogram_value_type().clone(),
8130                true,
8131            ),
8132        ]);
8133        let schema = Arc::new(Schema::new(columns));
8134        let table_meta = TableMetaBuilder::empty()
8135            .schema(schema)
8136            .primary_key_indices((0..tag_count).collect())
8137            .value_indices(vec![tag_count + 1])
8138            .next_column_id(1024)
8139            .build()
8140            .unwrap();
8141        let table_info = TableInfoBuilder::default()
8142            .name(table_name)
8143            .meta(table_meta)
8144            .build()
8145            .unwrap();
8146        let table = EmptyTable::from_table_info(&table_info);
8147
8148        assert!(
8149            catalog_list
8150                .register_table_sync(RegisterTableRequest {
8151                    catalog: DEFAULT_CATALOG_NAME.to_string(),
8152                    schema: DEFAULT_SCHEMA_NAME.to_string(),
8153                    table_name: table_name.to_string(),
8154                    table_id: 1024,
8155                    table,
8156                })
8157                .is_ok()
8158        );
8159
8160        DfTableSourceProvider::new(
8161            catalog_list,
8162            false,
8163            QueryContext::arc(),
8164            DummyDecoder::arc(),
8165            false,
8166        )
8167    }
8168
8169    async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8170        let catalog_list = MemoryCatalogManager::with_default_setup();
8171        let columns = vec![
8172            ColumnSchema::new(
8173                "tag_0".to_string(),
8174                ConcreteDataType::string_datatype(),
8175                false,
8176            ),
8177            ColumnSchema::new(
8178                "timestamp".to_string(),
8179                ConcreteDataType::timestamp_millisecond_datatype(),
8180                false,
8181            )
8182            .with_time_index(true),
8183            ColumnSchema::new(
8184                greptime_native_histogram().to_string(),
8185                native_histogram_value_type().clone(),
8186                true,
8187            ),
8188            ColumnSchema::new(
8189                "native_histogram_2".to_string(),
8190                native_histogram_value_type().clone(),
8191                true,
8192            ),
8193        ];
8194        let schema = Arc::new(Schema::new(columns));
8195        let table_meta = TableMetaBuilder::empty()
8196            .schema(schema)
8197            .primary_key_indices(vec![0])
8198            .value_indices(vec![2, 3])
8199            .next_column_id(1024)
8200            .build()
8201            .unwrap();
8202        let table_info = TableInfoBuilder::default()
8203            .name(table_name)
8204            .meta(table_meta)
8205            .build()
8206            .unwrap();
8207        let table = EmptyTable::from_table_info(&table_info);
8208
8209        assert!(
8210            catalog_list
8211                .register_table_sync(RegisterTableRequest {
8212                    catalog: DEFAULT_CATALOG_NAME.to_string(),
8213                    schema: DEFAULT_SCHEMA_NAME.to_string(),
8214                    table_name: table_name.to_string(),
8215                    table_id: 1024,
8216                    table,
8217                })
8218                .is_ok()
8219        );
8220
8221        DfTableSourceProvider::new(
8222            catalog_list,
8223            false,
8224            QueryContext::arc(),
8225            DummyDecoder::arc(),
8226            false,
8227        )
8228    }
8229
8230    async fn build_test_mixed_native_histogram_table_provider(
8231        table_name: &str,
8232    ) -> DfTableSourceProvider {
8233        build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await
8234    }
8235
8236    async fn build_test_mixed_native_histogram_table_provider_with_marker(
8237        table_name: &str,
8238        temporality_marker: bool,
8239    ) -> DfTableSourceProvider {
8240        let catalog_list = MemoryCatalogManager::with_default_setup();
8241        let mut columns = vec![ColumnSchema::new(
8242            "tag_0".to_string(),
8243            ConcreteDataType::string_datatype(),
8244            false,
8245        )];
8246        if temporality_marker {
8247            columns.push(ColumnSchema::new(
8248                OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8249                ConcreteDataType::string_datatype(),
8250                true,
8251            ));
8252        }
8253        let tag_count = columns.len();
8254        columns.extend([
8255            ColumnSchema::new(
8256                "timestamp".to_string(),
8257                ConcreteDataType::timestamp_millisecond_datatype(),
8258                false,
8259            )
8260            .with_time_index(true),
8261            ColumnSchema::new(
8262                greptime_native_histogram().to_string(),
8263                native_histogram_value_type().clone(),
8264                true,
8265            ),
8266            ColumnSchema::new(
8267                greptime_value().to_string(),
8268                ConcreteDataType::float64_datatype(),
8269                true,
8270            ),
8271        ]);
8272        let schema = Arc::new(Schema::new(columns));
8273        let table_meta = TableMetaBuilder::empty()
8274            .schema(schema.clone())
8275            .primary_key_indices((0..tag_count).collect())
8276            .value_indices(vec![tag_count + 1, tag_count + 2])
8277            .next_column_id(1024)
8278            .build()
8279            .unwrap();
8280        let table_info = Arc::new(
8281            TableInfoBuilder::default()
8282                .name(table_name)
8283                .meta(table_meta)
8284                .build()
8285                .unwrap(),
8286        );
8287        let mut arrays: Vec<Arc<dyn Array>> =
8288            vec![Arc::new(StringArray::from(vec!["float", "histogram"]))];
8289        if temporality_marker {
8290            arrays.push(Arc::new(StringArray::from(vec![
8291                Some(GREPTIME_TEMPORALITY_DELTA),
8292                Some(GREPTIME_TEMPORALITY_DELTA),
8293            ])));
8294        }
8295        arrays.extend([
8296            Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc<dyn Array>,
8297            build_histogram_array(&[None, Some(direct_or_histogram())]),
8298            Arc::new(Float64Array::from(vec![Some(2.0), None])),
8299        ]);
8300        let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
8301        let backing = GreptimeMemTable::new_with_catalog(
8302            table_name,
8303            GreptimeRecordBatch::from_df_record_batch(schema, batch),
8304            1024,
8305            DEFAULT_CATALOG_NAME.to_string(),
8306            DEFAULT_SCHEMA_NAME.to_string(),
8307        );
8308        let table = Arc::new(Table::new(
8309            table_info,
8310            FilterPushDownType::Unsupported,
8311            backing.data_source(),
8312        ));
8313
8314        assert!(
8315            catalog_list
8316                .register_table_sync(RegisterTableRequest {
8317                    catalog: DEFAULT_CATALOG_NAME.to_string(),
8318                    schema: DEFAULT_SCHEMA_NAME.to_string(),
8319                    table_name: table_name.to_string(),
8320                    table_id: 1024,
8321                    table,
8322                })
8323                .is_ok()
8324        );
8325
8326        DfTableSourceProvider::new(
8327            catalog_list,
8328            false,
8329            QueryContext::arc(),
8330            DummyDecoder::arc(),
8331            false,
8332        )
8333    }
8334
8335    fn classic_and_native_histogram_table_provider(
8336        native_tag: &str,
8337        native_le: Option<&str>,
8338        native_histogram: NativeHistogram,
8339    ) -> DfTableSourceProvider {
8340        let table_name = "mixed_histogram";
8341        let catalog = MemoryCatalogManager::with_default_setup();
8342        let schema = Arc::new(Schema::new(vec![
8343            ColumnSchema::new(
8344                "tag".to_string(),
8345                ConcreteDataType::string_datatype(),
8346                false,
8347            ),
8348            ColumnSchema::new(
8349                LE_COLUMN_NAME.to_string(),
8350                ConcreteDataType::string_datatype(),
8351                true,
8352            ),
8353            ColumnSchema::new(
8354                "timestamp".to_string(),
8355                ConcreteDataType::timestamp_millisecond_datatype(),
8356                false,
8357            )
8358            .with_time_index(true),
8359            ColumnSchema::new(
8360                greptime_native_histogram().to_string(),
8361                native_histogram_value_type().clone(),
8362                true,
8363            ),
8364            ColumnSchema::new(
8365                greptime_value().to_string(),
8366                ConcreteDataType::float64_datatype(),
8367                true,
8368            ),
8369        ]));
8370        let table_meta = TableMetaBuilder::empty()
8371            .schema(schema.clone())
8372            .primary_key_indices(vec![0, 1])
8373            .value_indices(vec![3, 4])
8374            .next_column_id(5)
8375            .build()
8376            .unwrap();
8377        let table_info = Arc::new(
8378            TableInfoBuilder::default()
8379                .name(table_name)
8380                .meta(table_meta)
8381                .build()
8382                .unwrap(),
8383        );
8384        let batch = RecordBatch::try_new(
8385            schema.arrow_schema().clone(),
8386            vec![
8387                Arc::new(StringArray::from(vec![
8388                    "classic", "classic", native_tag, "classic", "classic", native_tag,
8389                ])),
8390                Arc::new(StringArray::from(vec![
8391                    Some("1"),
8392                    Some("+Inf"),
8393                    native_le,
8394                    Some("1"),
8395                    Some("+Inf"),
8396                    native_le,
8397                ])),
8398                Arc::new(TimestampMillisecondArray::from(vec![
8399                    1_000, 1_000, 1_000, 2_000, 2_000, 2_000,
8400                ])),
8401                build_histogram_array(&[
8402                    None,
8403                    None,
8404                    Some(native_histogram.clone()),
8405                    None,
8406                    None,
8407                    Some(native_histogram),
8408                ]),
8409                Arc::new(Float64Array::from(vec![
8410                    Some(2.0),
8411                    Some(4.0),
8412                    None,
8413                    Some(2.0),
8414                    Some(4.0),
8415                    None,
8416                ])),
8417            ],
8418        )
8419        .unwrap();
8420        let backing = GreptimeMemTable::new_with_catalog(
8421            table_name,
8422            GreptimeRecordBatch::from_df_record_batch(schema, batch),
8423            2_200,
8424            DEFAULT_CATALOG_NAME.to_string(),
8425            DEFAULT_SCHEMA_NAME.to_string(),
8426        );
8427        let table = Arc::new(Table::new(
8428            table_info,
8429            FilterPushDownType::Unsupported,
8430            backing.data_source(),
8431        ));
8432        catalog
8433            .register_table_sync(RegisterTableRequest {
8434                catalog: DEFAULT_CATALOG_NAME.to_string(),
8435                schema: DEFAULT_SCHEMA_NAME.to_string(),
8436                table_name: table_name.to_string(),
8437                table_id: 2_200,
8438                table,
8439            })
8440            .unwrap();
8441
8442        DfTableSourceProvider::new(
8443            catalog,
8444            false,
8445            QueryContext::arc(),
8446            DummyDecoder::arc(),
8447            false,
8448        )
8449    }
8450
8451    async fn build_test_table_provider_with_tsid(
8452        table_name_tuples: &[(String, String)],
8453        num_tag: usize,
8454        num_field: usize,
8455    ) -> DfTableSourceProvider {
8456        let table_specs = table_name_tuples
8457            .iter()
8458            .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field))
8459            .collect::<Vec<_>>();
8460        build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await
8461    }
8462
8463    async fn build_test_table_provider_with_tsid_fields(
8464        table_specs: &[((String, String), usize)],
8465        num_tag: usize,
8466    ) -> DfTableSourceProvider {
8467        let table_specs = table_specs
8468            .iter()
8469            .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field))
8470            .collect::<Vec<_>>();
8471        build_test_table_provider_with_tsid_tag_fields(&table_specs).await
8472    }
8473
8474    async fn build_test_table_provider_with_tsid_tag_fields(
8475        table_specs: &[((String, String), usize, usize)],
8476    ) -> DfTableSourceProvider {
8477        let catalog_list = MemoryCatalogManager::with_default_setup();
8478
8479        let physical_table_name = "phy";
8480        let physical_table_id = 999u32;
8481        let physical_num_tag = table_specs
8482            .iter()
8483            .map(|(_, num_tag, _)| *num_tag)
8484            .max()
8485            .unwrap_or(0);
8486        let physical_num_field = table_specs
8487            .iter()
8488            .map(|(_, _, num_field)| *num_field)
8489            .max()
8490            .unwrap_or(0);
8491
8492        // Register a metric engine physical table with internal columns.
8493        {
8494            let mut columns = vec![
8495                ColumnSchema::new(
8496                    DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
8497                    ConcreteDataType::uint32_datatype(),
8498                    false,
8499                ),
8500                ColumnSchema::new(
8501                    DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
8502                    ConcreteDataType::uint64_datatype(),
8503                    false,
8504                ),
8505            ];
8506            for i in 0..physical_num_tag {
8507                columns.push(ColumnSchema::new(
8508                    format!("tag_{i}"),
8509                    ConcreteDataType::string_datatype(),
8510                    false,
8511                ));
8512            }
8513            columns.push(
8514                ColumnSchema::new(
8515                    "timestamp".to_string(),
8516                    ConcreteDataType::timestamp_millisecond_datatype(),
8517                    false,
8518                )
8519                .with_time_index(true),
8520            );
8521            for i in 0..physical_num_field {
8522                columns.push(ColumnSchema::new(
8523                    format!("field_{i}"),
8524                    ConcreteDataType::float64_datatype(),
8525                    true,
8526                ));
8527            }
8528
8529            let schema = Arc::new(Schema::new(columns));
8530            let primary_key_indices = (0..(2 + physical_num_tag)).collect::<Vec<_>>();
8531            let table_meta = TableMetaBuilder::empty()
8532                .schema(schema)
8533                .primary_key_indices(primary_key_indices)
8534                .value_indices(
8535                    (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(),
8536                )
8537                .engine(METRIC_ENGINE_NAME.to_string())
8538                .next_column_id(1024)
8539                .build()
8540                .unwrap();
8541            let table_info = TableInfoBuilder::default()
8542                .table_id(physical_table_id)
8543                .name(physical_table_name)
8544                .meta(table_meta)
8545                .build()
8546                .unwrap();
8547            let table = EmptyTable::from_table_info(&table_info);
8548
8549            assert!(
8550                catalog_list
8551                    .register_table_sync(RegisterTableRequest {
8552                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8553                        schema: DEFAULT_SCHEMA_NAME.to_string(),
8554                        table_name: physical_table_name.to_string(),
8555                        table_id: physical_table_id,
8556                        table,
8557                    })
8558                    .is_ok()
8559            );
8560        }
8561
8562        // Register metric engine logical tables without `__tsid`, referencing the physical table.
8563        for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate()
8564        {
8565            let mut columns = vec![];
8566            for i in 0..*num_tag {
8567                columns.push(ColumnSchema::new(
8568                    format!("tag_{i}"),
8569                    ConcreteDataType::string_datatype(),
8570                    false,
8571                ));
8572            }
8573            columns.push(
8574                ColumnSchema::new(
8575                    "timestamp".to_string(),
8576                    ConcreteDataType::timestamp_millisecond_datatype(),
8577                    false,
8578                )
8579                .with_time_index(true),
8580            );
8581            for i in 0..*num_field {
8582                columns.push(ColumnSchema::new(
8583                    format!("field_{i}"),
8584                    ConcreteDataType::float64_datatype(),
8585                    true,
8586                ));
8587            }
8588
8589            let schema = Arc::new(Schema::new(columns));
8590            let mut options = table::requests::TableOptions::default();
8591            options.extra_options.insert(
8592                LOGICAL_TABLE_METADATA_KEY.to_string(),
8593                physical_table_name.to_string(),
8594            );
8595            let table_id = 1024u32 + idx as u32;
8596            let table_meta = TableMetaBuilder::empty()
8597                .schema(schema)
8598                .primary_key_indices((0..*num_tag).collect())
8599                .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect())
8600                .engine(METRIC_ENGINE_NAME.to_string())
8601                .options(options)
8602                .next_column_id(1024)
8603                .build()
8604                .unwrap();
8605            let table_info = TableInfoBuilder::default()
8606                .table_id(table_id)
8607                .name(table_name.clone())
8608                .meta(table_meta)
8609                .build()
8610                .unwrap();
8611            let table = EmptyTable::from_table_info(&table_info);
8612
8613            assert!(
8614                catalog_list
8615                    .register_table_sync(RegisterTableRequest {
8616                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8617                        schema: schema_name.clone(),
8618                        table_name: table_name.clone(),
8619                        table_id,
8620                        table,
8621                    })
8622                    .is_ok()
8623            );
8624        }
8625
8626        DfTableSourceProvider::new(
8627            catalog_list,
8628            false,
8629            QueryContext::arc(),
8630            DummyDecoder::arc(),
8631            false,
8632        )
8633    }
8634
8635    async fn build_test_table_provider_with_fields(
8636        table_name_tuples: &[(String, String)],
8637        tags: &[&str],
8638    ) -> DfTableSourceProvider {
8639        let catalog_list = MemoryCatalogManager::with_default_setup();
8640        for (schema_name, table_name) in table_name_tuples {
8641            let mut columns = vec![];
8642            let num_tag = tags.len();
8643            for tag in tags {
8644                columns.push(ColumnSchema::new(
8645                    tag.to_string(),
8646                    ConcreteDataType::string_datatype(),
8647                    false,
8648                ));
8649            }
8650            columns.push(
8651                ColumnSchema::new(
8652                    greptime_timestamp().to_string(),
8653                    ConcreteDataType::timestamp_millisecond_datatype(),
8654                    false,
8655                )
8656                .with_time_index(true),
8657            );
8658            columns.push(ColumnSchema::new(
8659                greptime_value().to_string(),
8660                ConcreteDataType::float64_datatype(),
8661                true,
8662            ));
8663            let schema = Arc::new(Schema::new(columns));
8664            let table_meta = TableMetaBuilder::empty()
8665                .schema(schema)
8666                .primary_key_indices((0..num_tag).collect())
8667                .next_column_id(1024)
8668                .build()
8669                .unwrap();
8670            let table_info = TableInfoBuilder::default()
8671                .name(table_name.clone())
8672                .meta(table_meta)
8673                .build()
8674                .unwrap();
8675            let table = EmptyTable::from_table_info(&table_info);
8676
8677            assert!(
8678                catalog_list
8679                    .register_table_sync(RegisterTableRequest {
8680                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8681                        schema: schema_name.clone(),
8682                        table_name: table_name.clone(),
8683                        table_id: 1024,
8684                        table,
8685                    })
8686                    .is_ok()
8687            );
8688        }
8689
8690        DfTableSourceProvider::new(
8691            catalog_list,
8692            false,
8693            QueryContext::arc(),
8694            DummyDecoder::arc(),
8695            false,
8696        )
8697    }
8698
8699    // {
8700    //     input: `abs(some_metric{foo!="bar"})`,
8701    //     expected: &Call{
8702    //         Func: MustGetFunction("abs"),
8703    //         Args: Expressions{
8704    //             &VectorSelector{
8705    //                 Name: "some_metric",
8706    //                 LabelMatchers: []*labels.Matcher{
8707    //                     MustLabelMatcher(labels.MatchNotEqual, "foo", "bar"),
8708    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
8709    //                 },
8710    //             },
8711    //         },
8712    //     },
8713    // },
8714    async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) {
8715        let prom_expr =
8716            parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap();
8717        let eval_stmt = EvalStmt {
8718            expr: prom_expr,
8719            start: UNIX_EPOCH,
8720            end: UNIX_EPOCH
8721                .checked_add(Duration::from_secs(100_000))
8722                .unwrap(),
8723            interval: Duration::from_secs(5),
8724            lookback_delta: Duration::from_secs(1),
8725        };
8726
8727        let table_provider = build_test_table_provider(
8728            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8729            1,
8730            1,
8731        )
8732        .await;
8733        let plan =
8734            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8735                .await
8736                .unwrap();
8737
8738        let expected = String::from(
8739            "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\
8740            \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]\
8741            \n    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8742            \n      PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8743            \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]\
8744	            \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]\
8745            \n            TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"
8746        ).replace("TEMPLATE", plan_name);
8747
8748        assert_eq!(plan.display_indent_schema().to_string(), expected);
8749    }
8750
8751    #[tokio::test]
8752    async fn single_abs() {
8753        do_single_instant_function_call("abs", "abs").await;
8754    }
8755
8756    #[tokio::test]
8757    #[should_panic]
8758    async fn single_absent() {
8759        do_single_instant_function_call("absent", "").await;
8760    }
8761
8762    #[tokio::test]
8763    async fn single_ceil() {
8764        do_single_instant_function_call("ceil", "ceil").await;
8765    }
8766
8767    #[tokio::test]
8768    async fn single_exp() {
8769        do_single_instant_function_call("exp", "exp").await;
8770    }
8771
8772    #[tokio::test]
8773    async fn single_ln() {
8774        do_single_instant_function_call("ln", "ln").await;
8775    }
8776
8777    #[tokio::test]
8778    async fn single_log2() {
8779        do_single_instant_function_call("log2", "log2").await;
8780    }
8781
8782    #[tokio::test]
8783    async fn single_log10() {
8784        do_single_instant_function_call("log10", "log10").await;
8785    }
8786
8787    #[tokio::test]
8788    #[should_panic]
8789    async fn single_scalar() {
8790        do_single_instant_function_call("scalar", "").await;
8791    }
8792
8793    #[tokio::test]
8794    #[should_panic]
8795    async fn single_sgn() {
8796        do_single_instant_function_call("sgn", "").await;
8797    }
8798
8799    #[tokio::test]
8800    #[should_panic]
8801    async fn single_sort() {
8802        do_single_instant_function_call("sort", "").await;
8803    }
8804
8805    #[tokio::test]
8806    #[should_panic]
8807    async fn single_sort_desc() {
8808        do_single_instant_function_call("sort_desc", "").await;
8809    }
8810
8811    #[tokio::test]
8812    async fn single_sqrt() {
8813        do_single_instant_function_call("sqrt", "sqrt").await;
8814    }
8815
8816    #[tokio::test]
8817    async fn single_timestamp_plan_preserves_source_value() {
8818        let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#);
8819        let table_provider = build_test_table_provider(
8820            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8821            1,
8822            1,
8823        )
8824        .await;
8825        let plan =
8826            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8827                .await
8828                .unwrap();
8829
8830        let expected = String::from(
8831            "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8832            \n  Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8833            \n    Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8834            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
8835            \n        Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(some_metric.timestamp AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
8836            \n          PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8837            \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]\
8838            \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]\
8839            \n                TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
8840        );
8841
8842        assert_eq!(plan.display_indent_schema().to_string(), expected);
8843    }
8844
8845    #[tokio::test]
8846    async fn single_acos() {
8847        do_single_instant_function_call("acos", "acos").await;
8848    }
8849
8850    #[tokio::test]
8851    #[should_panic]
8852    async fn single_acosh() {
8853        do_single_instant_function_call("acosh", "").await;
8854    }
8855
8856    #[tokio::test]
8857    async fn single_asin() {
8858        do_single_instant_function_call("asin", "asin").await;
8859    }
8860
8861    #[tokio::test]
8862    #[should_panic]
8863    async fn single_asinh() {
8864        do_single_instant_function_call("asinh", "").await;
8865    }
8866
8867    #[tokio::test]
8868    async fn single_atan() {
8869        do_single_instant_function_call("atan", "atan").await;
8870    }
8871
8872    #[tokio::test]
8873    #[should_panic]
8874    async fn single_atanh() {
8875        do_single_instant_function_call("atanh", "").await;
8876    }
8877
8878    #[tokio::test]
8879    async fn single_cos() {
8880        do_single_instant_function_call("cos", "cos").await;
8881    }
8882
8883    #[tokio::test]
8884    #[should_panic]
8885    async fn single_cosh() {
8886        do_single_instant_function_call("cosh", "").await;
8887    }
8888
8889    #[tokio::test]
8890    async fn single_sin() {
8891        do_single_instant_function_call("sin", "sin").await;
8892    }
8893
8894    #[tokio::test]
8895    #[should_panic]
8896    async fn single_sinh() {
8897        do_single_instant_function_call("sinh", "").await;
8898    }
8899
8900    #[tokio::test]
8901    async fn single_tan() {
8902        do_single_instant_function_call("tan", "tan").await;
8903    }
8904
8905    #[tokio::test]
8906    #[should_panic]
8907    async fn single_tanh() {
8908        do_single_instant_function_call("tanh", "").await;
8909    }
8910
8911    #[tokio::test]
8912    #[should_panic]
8913    async fn single_deg() {
8914        do_single_instant_function_call("deg", "").await;
8915    }
8916
8917    #[tokio::test]
8918    #[should_panic]
8919    async fn single_rad() {
8920        do_single_instant_function_call("rad", "").await;
8921    }
8922
8923    // {
8924    //     input: "avg by (foo)(some_metric)",
8925    //     expected: &AggregateExpr{
8926    //         Op: AVG,
8927    //         Expr: &VectorSelector{
8928    //             Name: "some_metric",
8929    //             LabelMatchers: []*labels.Matcher{
8930    //                 MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
8931    //             },
8932    //             PosRange: PositionRange{
8933    //                 Start: 13,
8934    //                 End:   24,
8935    //             },
8936    //         },
8937    //         Grouping: []string{"foo"},
8938    //         PosRange: PositionRange{
8939    //             Start: 0,
8940    //             End:   25,
8941    //         },
8942    //     },
8943    // },
8944    async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) {
8945        let prom_expr = parser::parse(&format!(
8946            "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})",
8947        ))
8948        .unwrap();
8949        let mut eval_stmt = EvalStmt {
8950            expr: prom_expr,
8951            start: UNIX_EPOCH,
8952            end: UNIX_EPOCH
8953                .checked_add(Duration::from_secs(100_000))
8954                .unwrap(),
8955            interval: Duration::from_secs(5),
8956            lookback_delta: Duration::from_secs(1),
8957        };
8958
8959        // test group by
8960        let table_provider = build_test_table_provider(
8961            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8962            2,
8963            2,
8964        )
8965        .await;
8966        let plan =
8967            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8968                .await
8969                .unwrap();
8970        let expected_no_without = String::from(
8971            "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]\
8972            \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]\
8973            \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]\
8974            \n      PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
8975            \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]\
8976            \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]\
8977            \n            TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
8978        ).replace("TEMPLATE", plan_name);
8979        assert_eq!(
8980            plan.display_indent_schema().to_string(),
8981            expected_no_without
8982        );
8983
8984        // test group without
8985        if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr {
8986            *modifier = Some(LabelModifier::Exclude(Labels {
8987                labels: vec![String::from("tag_1")].into_iter().collect(),
8988            }));
8989        }
8990        let table_provider = build_test_table_provider(
8991            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8992            2,
8993            2,
8994        )
8995        .await;
8996        let plan =
8997            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8998                .await
8999                .unwrap();
9000        let expected_without = String::from(
9001            "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]\
9002            \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]\
9003            \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]\
9004            \n      PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9005            \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]\
9006            \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]\
9007            \n            TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
9008        ).replace("TEMPLATE", plan_name);
9009        assert_eq!(plan.display_indent_schema().to_string(), expected_without);
9010    }
9011
9012    #[tokio::test]
9013    async fn aggregate_sum() {
9014        do_aggregate_expr_plan("sum", "sum").await;
9015    }
9016
9017    #[tokio::test]
9018    async fn tsid_is_used_for_series_divide_when_available() {
9019        let prom_expr = parser::parse("some_metric").unwrap();
9020        let eval_stmt = EvalStmt {
9021            expr: prom_expr,
9022            start: UNIX_EPOCH,
9023            end: UNIX_EPOCH
9024                .checked_add(Duration::from_secs(100_000))
9025                .unwrap(),
9026            interval: Duration::from_secs(5),
9027            lookback_delta: Duration::from_secs(1),
9028        };
9029
9030        let table_provider = build_test_table_provider_with_tsid(
9031            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9032            1,
9033            1,
9034        )
9035        .await;
9036        let plan =
9037            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9038                .await
9039                .unwrap();
9040
9041        let plan_str = plan.display_indent_schema().to_string();
9042        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
9043        assert!(plan_str.contains("__tsid ASC NULLS FIRST"));
9044        assert!(
9045            !plan
9046                .schema()
9047                .fields()
9048                .iter()
9049                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
9050        );
9051
9052        let manipulate = find_instant_manipulate(&plan).unwrap();
9053        let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new(
9054            MemorySourceConfig::try_new(&[], Arc::new(ArrowSchema::empty()), None).unwrap(),
9055        ))));
9056        assert!(format!("{exec:?}").contains("reuse_tsid_column: true"));
9057    }
9058
9059    #[tokio::test]
9060    async fn default_binary_join_uses_tsid_when_available() {
9061        let eval_stmt = build_eval_stmt("some_metric / some_alt_metric");
9062
9063        let table_provider = build_test_table_provider_with_tsid(
9064            &[
9065                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9066                (
9067                    DEFAULT_SCHEMA_NAME.to_string(),
9068                    "some_alt_metric".to_string(),
9069                ),
9070            ],
9071            1,
9072            1,
9073        )
9074        .await;
9075        let plan =
9076            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9077                .await
9078                .unwrap();
9079
9080        let plan_str = plan.display_indent_schema().to_string();
9081        assert!(
9082            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9083            "{plan_str}"
9084        );
9085        assert!(
9086            !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9087            "{plan_str}"
9088        );
9089    }
9090
9091    #[tokio::test]
9092    async fn reject_binary_fill_modifiers() {
9093        let state = build_query_engine_state();
9094
9095        for query in [
9096            "some_metric + fill(0) some_alt_metric",
9097            "some_metric + fill_left(0) some_alt_metric",
9098            "some_metric + fill_right(0) some_alt_metric",
9099            "(some_metric + fill(0) some_alt_metric) + some_metric",
9100        ] {
9101            let eval_stmt = build_eval_stmt(query);
9102            let table_provider = build_test_table_provider(&[], 0, 0).await;
9103            let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state)
9104                .await
9105                .unwrap_err();
9106
9107            assert!(
9108                matches!(
9109                    &err,
9110                    crate::promql::error::Error::UnsupportedExpr { name, .. }
9111                        if name == "PromQL fill modifiers"
9112                ),
9113                "{err}"
9114            );
9115        }
9116    }
9117
9118    #[tokio::test]
9119    async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() {
9120        for query in [
9121            "timestamp(some_metric) / some_metric",
9122            "some_metric / timestamp(some_metric)",
9123        ] {
9124            let eval_stmt = build_eval_stmt(query);
9125
9126            let table_provider = build_test_table_provider_with_tsid(
9127                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9128                1,
9129                1,
9130            )
9131            .await;
9132            let plan =
9133                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9134                    .await
9135                    .unwrap();
9136
9137            let plan_str = plan.display_indent_schema().to_string();
9138            assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}");
9139            assert!(
9140                plan_str.contains("lhs.tag_0 = rhs.tag_0"),
9141                "{query}: {plan_str}"
9142            );
9143            assert!(
9144                !plan
9145                    .schema()
9146                    .fields()
9147                    .iter()
9148                    .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
9149                "{query}: {plan_str}"
9150            );
9151        }
9152    }
9153
9154    #[tokio::test]
9155    async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() {
9156        let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job");
9157
9158        let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
9159            (
9160                (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()),
9161                2,
9162                1,
9163            ),
9164            (
9165                (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()),
9166                1,
9167                1,
9168            ),
9169        ])
9170        .await;
9171        let plan =
9172            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9173                .await
9174                .unwrap();
9175        let plan_str = plan.display_indent_schema().to_string();
9176
9177        assert!(
9178            plan_str.contains("Boolean(false)") || plan_str.contains("false"),
9179            "{plan_str}"
9180        );
9181    }
9182
9183    #[tokio::test]
9184    async fn tsid_is_preserved_for_nested_default_binary_joins() {
9185        let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric");
9186
9187        let table_provider = build_test_table_provider_with_tsid(
9188            &[
9189                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9190                (
9191                    DEFAULT_SCHEMA_NAME.to_string(),
9192                    "some_alt_metric".to_string(),
9193                ),
9194                (
9195                    DEFAULT_SCHEMA_NAME.to_string(),
9196                    "some_third_metric".to_string(),
9197                ),
9198            ],
9199            1,
9200            1,
9201        )
9202        .await;
9203        let plan =
9204            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9205                .await
9206                .unwrap();
9207
9208        let plan_str = plan.display_indent_schema().to_string();
9209        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9210        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9211    }
9212
9213    #[tokio::test]
9214    async fn repeated_tsid_binary_operand_reuses_leaf_plan() {
9215        let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100");
9216
9217        let table_provider = build_test_table_provider_with_tsid(
9218            &[
9219                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9220                (
9221                    DEFAULT_SCHEMA_NAME.to_string(),
9222                    "some_alt_metric".to_string(),
9223                ),
9224            ],
9225            1,
9226            1,
9227        )
9228        .await;
9229        let plan =
9230            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9231                .await
9232                .unwrap();
9233
9234        let plan_str = plan.display_indent_schema().to_string();
9235        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9236        assert_eq!(
9237            plan_str
9238                .matches("Filter: phy.__table_id = UInt32(1024)")
9239                .count(),
9240            1,
9241            "{plan_str}"
9242        );
9243        assert_eq!(
9244            plan_str.matches("PromInstantManipulate").count(),
9245            2,
9246            "{plan_str}"
9247        );
9248        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9249    }
9250
9251    #[tokio::test]
9252    async fn repeated_tsid_binary_operand_reuses_shorter_field_side() {
9253        let eval_stmt =
9254            build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100");
9255
9256        let table_provider = build_test_table_provider_with_tsid_fields(
9257            &[
9258                (
9259                    (
9260                        DEFAULT_SCHEMA_NAME.to_string(),
9261                        "two_field_metric".to_string(),
9262                    ),
9263                    2,
9264                ),
9265                (
9266                    (
9267                        DEFAULT_SCHEMA_NAME.to_string(),
9268                        "one_field_metric".to_string(),
9269                    ),
9270                    1,
9271                ),
9272            ],
9273            1,
9274        )
9275        .await;
9276        let plan =
9277            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9278                .await
9279                .unwrap();
9280
9281        let field_names = plan
9282            .schema()
9283            .fields()
9284            .iter()
9285            .map(|field| field.name().clone())
9286            .collect::<Vec<_>>();
9287        let value_columns = field_names
9288            .iter()
9289            .filter(|name| {
9290                *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
9291            })
9292            .count();
9293        assert_eq!(value_columns, 1, "{field_names:?}");
9294        let plan_str = plan.display_indent_schema().to_string();
9295        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9296        assert_eq!(
9297            plan_str
9298                .matches("Filter: phy.__table_id = UInt32(1025)")
9299                .count(),
9300            1,
9301            "{plan_str}"
9302        );
9303        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9304    }
9305
9306    #[tokio::test]
9307    async fn binary_island_reuses_self_operand_without_join() {
9308        let eval_stmt = build_eval_stmt("some_metric / some_metric");
9309
9310        let table_provider = build_test_table_provider_with_tsid(
9311            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9312            1,
9313            1,
9314        )
9315        .await;
9316        let plan =
9317            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9318                .await
9319                .unwrap();
9320
9321        let plan_str = plan.display_indent_schema().to_string();
9322        assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}");
9323        assert_eq!(
9324            plan_str
9325                .matches("Filter: phy.__table_id = UInt32(1024)")
9326                .count(),
9327            1,
9328            "{plan_str}"
9329        );
9330        assert_eq!(
9331            plan_str.matches("PromInstantManipulate").count(),
9332            1,
9333            "{plan_str}"
9334        );
9335    }
9336
9337    #[tokio::test]
9338    async fn binary_island_reuses_leaf_across_two_branches() {
9339        let eval_stmt =
9340            build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)");
9341
9342        let table_provider = build_test_table_provider_with_tsid(
9343            &[
9344                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9345                (
9346                    DEFAULT_SCHEMA_NAME.to_string(),
9347                    "some_alt_metric".to_string(),
9348                ),
9349                (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()),
9350            ],
9351            1,
9352            1,
9353        )
9354        .await;
9355        let plan =
9356            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9357                .await
9358                .unwrap();
9359
9360        let plan_str = plan.display_indent_schema().to_string();
9361        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9362        assert_eq!(
9363            plan_str
9364                .matches("Filter: phy.__table_id = UInt32(1024)")
9365                .count(),
9366            1,
9367            "{plan_str}"
9368        );
9369        assert_eq!(
9370            plan_str.matches("PromInstantManipulate").count(),
9371            3,
9372            "{plan_str}"
9373        );
9374    }
9375
9376    #[tokio::test]
9377    async fn binary_island_generated_alias_avoids_user_column_names() {
9378        let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric");
9379
9380        let table_provider = build_test_table_provider_with_fields(
9381            &[
9382                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9383                (
9384                    DEFAULT_SCHEMA_NAME.to_string(),
9385                    "some_alt_metric".to_string(),
9386                ),
9387            ],
9388            &["prom_v0", "__prom_v0"],
9389        )
9390        .await;
9391        let plan =
9392            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9393                .await
9394                .unwrap();
9395
9396        let field_names = plan.schema().field_names();
9397        assert!(field_names.iter().any(|name| name.ends_with(".prom_v0")));
9398        assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0")));
9399
9400        let plan_str = plan.display_indent_schema().to_string();
9401        assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}");
9402        assert_eq!(
9403            plan_str.matches("PromInstantManipulate").count(),
9404            2,
9405            "{plan_str}"
9406        );
9407    }
9408
9409    #[tokio::test]
9410    async fn binary_island_clears_qualifier_for_nested_unary_projection() {
9411        let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)");
9412
9413        let table_provider = build_test_table_provider_with_tsid(
9414            &[
9415                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9416                (
9417                    DEFAULT_SCHEMA_NAME.to_string(),
9418                    "some_alt_metric".to_string(),
9419                ),
9420            ],
9421            1,
9422            1,
9423        )
9424        .await;
9425        let plan =
9426            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9427                .await
9428                .unwrap();
9429
9430        let plan_str = plan.display_indent_schema().to_string();
9431        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9432        assert_eq!(
9433            plan_str.matches("PromInstantManipulate").count(),
9434            2,
9435            "{plan_str}"
9436        );
9437    }
9438
9439    #[tokio::test]
9440    async fn binary_island_keeps_distinct_matcher_leaves() {
9441        let eval_stmt = build_eval_stmt(
9442            "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}",
9443        );
9444
9445        let table_provider = build_test_table_provider_with_tsid(
9446            &[
9447                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9448                (
9449                    DEFAULT_SCHEMA_NAME.to_string(),
9450                    "some_alt_metric".to_string(),
9451                ),
9452            ],
9453            1,
9454            1,
9455        )
9456        .await;
9457        let plan =
9458            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9459                .await
9460                .unwrap();
9461
9462        let plan_str = plan.display_indent_schema().to_string();
9463        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9464        assert_eq!(
9465            plan_str.matches("PromInstantManipulate").count(),
9466            3,
9467            "{plan_str}"
9468        );
9469    }
9470
9471    #[tokio::test]
9472    async fn binary_island_keeps_offset_leaves_distinct() {
9473        let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric");
9474
9475        let table_provider = build_test_table_provider_with_tsid(
9476            &[
9477                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9478                (
9479                    DEFAULT_SCHEMA_NAME.to_string(),
9480                    "some_alt_metric".to_string(),
9481                ),
9482            ],
9483            1,
9484            1,
9485        )
9486        .await;
9487        let plan =
9488            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9489                .await
9490                .unwrap();
9491
9492        let plan_str = plan.display_indent_schema().to_string();
9493        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9494        assert_eq!(
9495            plan_str.matches("PromInstantManipulate").count(),
9496            3,
9497            "{plan_str}"
9498        );
9499    }
9500
9501    #[tokio::test]
9502    async fn binary_island_falls_back_for_group_modifier() {
9503        let eval_stmt = build_eval_stmt(
9504            "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric",
9505        );
9506
9507        let table_provider = build_test_table_provider_with_tsid(
9508            &[
9509                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9510                (
9511                    DEFAULT_SCHEMA_NAME.to_string(),
9512                    "some_alt_metric".to_string(),
9513                ),
9514            ],
9515            1,
9516            1,
9517        )
9518        .await;
9519        let plan =
9520            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9521                .await
9522                .unwrap();
9523
9524        let plan_str = plan.display_indent_schema().to_string();
9525        assert_eq!(
9526            plan_str.matches("PromInstantManipulate").count(),
9527            3,
9528            "{plan_str}"
9529        );
9530    }
9531
9532    #[tokio::test]
9533    async fn binary_island_falls_back_for_comparison_filter() {
9534        let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric");
9535
9536        let table_provider = build_test_table_provider_with_tsid(
9537            &[
9538                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9539                (
9540                    DEFAULT_SCHEMA_NAME.to_string(),
9541                    "some_alt_metric".to_string(),
9542                ),
9543            ],
9544            1,
9545            1,
9546        )
9547        .await;
9548        let plan =
9549            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9550                .await
9551                .unwrap();
9552
9553        let plan_str = plan.display_indent_schema().to_string();
9554        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9555        assert_eq!(
9556            plan_str.matches("PromInstantManipulate").count(),
9557            3,
9558            "{plan_str}"
9559        );
9560    }
9561
9562    #[tokio::test]
9563    async fn tsid_binary_join_uses_shorter_field_side() {
9564        let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric");
9565
9566        let table_provider = build_test_table_provider_with_tsid_fields(
9567            &[
9568                (
9569                    (
9570                        DEFAULT_SCHEMA_NAME.to_string(),
9571                        "one_field_metric".to_string(),
9572                    ),
9573                    1,
9574                ),
9575                (
9576                    (
9577                        DEFAULT_SCHEMA_NAME.to_string(),
9578                        "two_field_metric".to_string(),
9579                    ),
9580                    2,
9581                ),
9582            ],
9583            1,
9584        )
9585        .await;
9586        let plan =
9587            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9588                .await
9589                .unwrap();
9590
9591        let field_names = plan
9592            .schema()
9593            .fields()
9594            .iter()
9595            .map(|field| field.name().clone())
9596            .collect::<Vec<_>>();
9597        let value_columns = field_names
9598            .iter()
9599            .filter(|name| {
9600                *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
9601            })
9602            .count();
9603        assert_eq!(value_columns, 1, "{field_names:?}");
9604    }
9605
9606    #[tokio::test]
9607    async fn comparison_binary_join_uses_shorter_field_side() {
9608        let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric");
9609
9610        let table_provider = build_test_table_provider_with_tsid_fields(
9611            &[
9612                (
9613                    (
9614                        DEFAULT_SCHEMA_NAME.to_string(),
9615                        "two_field_metric".to_string(),
9616                    ),
9617                    2,
9618                ),
9619                (
9620                    (
9621                        DEFAULT_SCHEMA_NAME.to_string(),
9622                        "one_field_metric".to_string(),
9623                    ),
9624                    1,
9625                ),
9626            ],
9627            1,
9628        )
9629        .await;
9630        let plan =
9631            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9632                .await
9633                .unwrap();
9634
9635        let field_names = plan
9636            .schema()
9637            .fields()
9638            .iter()
9639            .map(|field| field.name().clone())
9640            .collect::<Vec<_>>();
9641        assert!(
9642            field_names.iter().any(|name| name == "field_0"),
9643            "{field_names:?}"
9644        );
9645        assert!(
9646            !field_names.iter().any(|name| name == "field_1"),
9647            "{field_names:?}"
9648        );
9649    }
9650
9651    #[tokio::test]
9652    async fn label_matching_modifier_disables_tsid_binary_join() {
9653        let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric");
9654
9655        let table_provider = build_test_table_provider_with_tsid(
9656            &[
9657                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9658                (
9659                    DEFAULT_SCHEMA_NAME.to_string(),
9660                    "some_alt_metric".to_string(),
9661                ),
9662            ],
9663            2,
9664            1,
9665        )
9666        .await;
9667        let plan =
9668            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9669                .await
9670                .unwrap();
9671
9672        let plan_str = plan.display_indent_schema().to_string();
9673        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9674        assert!(
9675            plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"),
9676            "{plan_str}"
9677        );
9678    }
9679
9680    #[tokio::test]
9681    async fn ignoring_absent_label_keeps_tsid_binary_join() {
9682        let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric");
9683
9684        let table_provider = build_test_table_provider_with_tsid(
9685            &[
9686                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9687                (
9688                    DEFAULT_SCHEMA_NAME.to_string(),
9689                    "some_alt_metric".to_string(),
9690                ),
9691            ],
9692            2,
9693            1,
9694        )
9695        .await;
9696        let plan =
9697            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9698                .await
9699                .unwrap();
9700
9701        let plan_str = plan.display_indent_schema().to_string();
9702        assert!(
9703            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9704            "{plan_str}"
9705        );
9706        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9707        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9708    }
9709
9710    #[tokio::test]
9711    async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() {
9712        let eval_stmt =
9713            build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric");
9714
9715        let table_provider = build_test_table_provider_with_tsid(
9716            &[
9717                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9718                (
9719                    DEFAULT_SCHEMA_NAME.to_string(),
9720                    "some_alt_metric".to_string(),
9721                ),
9722            ],
9723            2,
9724            1,
9725        )
9726        .await;
9727        let plan =
9728            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9729                .await
9730                .unwrap();
9731
9732        let plan_str = plan.display_indent_schema().to_string();
9733        assert!(
9734            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9735            "{plan_str}"
9736        );
9737        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9738        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9739    }
9740
9741    #[tokio::test]
9742    async fn on_full_label_set_keeps_tsid_binary_join() {
9743        let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric");
9744
9745        let table_provider = build_test_table_provider_with_tsid(
9746            &[
9747                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9748                (
9749                    DEFAULT_SCHEMA_NAME.to_string(),
9750                    "some_alt_metric".to_string(),
9751                ),
9752            ],
9753            2,
9754            1,
9755        )
9756        .await;
9757        let plan =
9758            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9759                .await
9760                .unwrap();
9761
9762        let plan_str = plan.display_indent_schema().to_string();
9763        assert!(
9764            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9765            "{plan_str}"
9766        );
9767        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9768        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9769    }
9770
9771    #[tokio::test]
9772    async fn on_partial_label_set_disables_tsid_binary_join() {
9773        let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
9774
9775        let table_provider = build_test_table_provider_with_tsid(
9776            &[
9777                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9778                (
9779                    DEFAULT_SCHEMA_NAME.to_string(),
9780                    "some_alt_metric".to_string(),
9781                ),
9782            ],
9783            2,
9784            1,
9785        )
9786        .await;
9787        let plan =
9788            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9789                .await
9790                .unwrap();
9791
9792        let plan_str = plan.display_indent_schema().to_string();
9793        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9794        assert!(
9795            plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9796            "{plan_str}"
9797        );
9798        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9799    }
9800
9801    #[tokio::test]
9802    async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() {
9803        let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
9804
9805        let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
9806            (
9807                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9808                2,
9809                1,
9810            ),
9811            (
9812                (
9813                    DEFAULT_SCHEMA_NAME.to_string(),
9814                    "some_alt_metric".to_string(),
9815                ),
9816                1,
9817                1,
9818            ),
9819        ])
9820        .await;
9821        let plan =
9822            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9823                .await
9824                .unwrap();
9825
9826        let plan_str = plan.display_indent_schema().to_string();
9827        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9828        assert!(
9829            plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9830            "{plan_str}"
9831        );
9832        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9833    }
9834
9835    #[tokio::test]
9836    async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() {
9837        let eval_stmt = build_eval_stmt("some_metric > some_alt_metric");
9838
9839        let table_provider = build_test_table_provider_with_tsid(
9840            &[
9841                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9842                (
9843                    DEFAULT_SCHEMA_NAME.to_string(),
9844                    "some_alt_metric".to_string(),
9845                ),
9846            ],
9847            2,
9848            1,
9849        )
9850        .await;
9851        let mut planner = PromPlanner {
9852            table_provider,
9853            ctx: PromPlannerContext::from_eval_stmt(&eval_stmt),
9854            promql_annotations: None,
9855        };
9856        let plan = planner
9857            .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state())
9858            .await
9859            .unwrap();
9860
9861        let plan_str = plan.display_indent_schema().to_string();
9862        assert!(
9863            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9864            "{plan_str}"
9865        );
9866        assert!(
9867            plan.schema()
9868                .fields()
9869                .iter()
9870                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
9871            "{plan_str}"
9872        );
9873        assert!(planner.ctx.use_tsid, "{plan_str}");
9874    }
9875
9876    #[tokio::test]
9877    async fn comparison_bool_binary_join_uses_tsid_when_available() {
9878        let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric");
9879
9880        let table_provider = build_test_table_provider_with_tsid(
9881            &[
9882                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9883                (
9884                    DEFAULT_SCHEMA_NAME.to_string(),
9885                    "some_alt_metric".to_string(),
9886                ),
9887            ],
9888            2,
9889            1,
9890        )
9891        .await;
9892        let plan =
9893            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9894                .await
9895                .unwrap();
9896
9897        let plan_str = plan.display_indent_schema().to_string();
9898        assert!(
9899            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9900            "{plan_str}"
9901        );
9902        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9903        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9904    }
9905
9906    #[tokio::test]
9907    async fn scalar_count_count_range_keeps_full_window() {
9908        let plan_str = build_optimized_tsid_plan(
9909            "scalar(count(count(some_metric) by (tag_0)))",
9910            1,
9911            1,
9912            100_000,
9913            1,
9914        )
9915        .await;
9916        assert!(plan_str.contains("ScalarCalculate: tags=[]"));
9917        assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]"));
9918        assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]"));
9919    }
9920
9921    #[tokio::test]
9922    async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() {
9923        let plan_str = build_optimized_tsid_plan(
9924            "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))",
9925            2,
9926            1,
9927            10,
9928            300,
9929        )
9930        .await;
9931        assert!(plan_str.contains("Distinct:"), "{plan_str}");
9932    }
9933
9934    #[tokio::test]
9935    async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() {
9936        assert_nested_count_rewrite_applies(
9937            "count(count(some_metric) by (tag_0))",
9938            "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]"
9939        )
9940        .await;
9941    }
9942
9943    #[tokio::test]
9944    async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() {
9945        assert_nested_count_rewrite_applies(
9946            "count(sum(some_metric) by (tag_0))",
9947            "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]"
9948        )
9949        .await;
9950    }
9951
9952    #[tokio::test]
9953    async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() {
9954        for (query, expected_outer_agg) in [
9955            (
9956                "count(avg(some_metric) by (tag_0))",
9957                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]",
9958            ),
9959            (
9960                "count(min(some_metric) by (tag_0))",
9961                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]",
9962            ),
9963            (
9964                "count(max(some_metric) by (tag_0))",
9965                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]",
9966            ),
9967            (
9968                "count(stddev(some_metric) by (tag_0))",
9969                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]",
9970            ),
9971            (
9972                "count(stdvar(some_metric) by (tag_0))",
9973                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]",
9974            ),
9975        ] {
9976            assert_nested_count_rewrite_applies(query, expected_outer_agg).await;
9977        }
9978    }
9979
9980    #[tokio::test]
9981    async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() {
9982        let count_plan =
9983            build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1)
9984                .await;
9985        assert!(
9986            !count_plan.contains("some_metric.field_0 IS NOT NULL"),
9987            "{count_plan}"
9988        );
9989
9990        for query in [
9991            "count(sum(some_metric) by (tag_0))",
9992            "count(avg(some_metric) by (tag_0))",
9993            "count(min(some_metric) by (tag_0))",
9994            "count(max(some_metric) by (tag_0))",
9995            "count(stddev(some_metric) by (tag_0))",
9996            "count(stdvar(some_metric) by (tag_0))",
9997        ] {
9998            let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
9999            assert!(
10000                plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"),
10001                "{query}: {plan_str}"
10002            );
10003        }
10004    }
10005
10006    #[tokio::test]
10007    async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() {
10008        assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await;
10009        assert_nested_count_rewrite_missing(
10010            "count(sum(irate(some_metric[1h])) by (tag_0))",
10011            2,
10012            300,
10013        )
10014        .await;
10015    }
10016
10017    #[tokio::test]
10018    async fn physical_table_name_is_not_leaked_in_plan() {
10019        let prom_expr = parser::parse("some_metric").unwrap();
10020        let eval_stmt = EvalStmt {
10021            expr: prom_expr,
10022            start: UNIX_EPOCH,
10023            end: UNIX_EPOCH
10024                .checked_add(Duration::from_secs(100_000))
10025                .unwrap(),
10026            interval: Duration::from_secs(5),
10027            lookback_delta: Duration::from_secs(1),
10028        };
10029
10030        let table_provider = build_test_table_provider_with_tsid(
10031            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10032            1,
10033            1,
10034        )
10035        .await;
10036        let plan =
10037            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10038                .await
10039                .unwrap();
10040
10041        let plan_str = plan.display_indent_schema().to_string();
10042        assert!(plan_str.contains("TableScan: phy"), "{plan}");
10043        assert!(plan_str.contains("SubqueryAlias: some_metric"));
10044        assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)"));
10045        assert!(!plan_str.contains("TableScan: some_metric"));
10046    }
10047
10048    #[tokio::test]
10049    async fn sum_without_does_not_group_by_tsid() {
10050        let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap();
10051        let eval_stmt = EvalStmt {
10052            expr: prom_expr,
10053            start: UNIX_EPOCH,
10054            end: UNIX_EPOCH
10055                .checked_add(Duration::from_secs(100_000))
10056                .unwrap(),
10057            interval: Duration::from_secs(5),
10058            lookback_delta: Duration::from_secs(1),
10059        };
10060
10061        let table_provider = build_test_table_provider_with_tsid(
10062            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10063            1,
10064            1,
10065        )
10066        .await;
10067        let plan =
10068            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10069                .await
10070                .unwrap();
10071
10072        let plan_str = plan.display_indent_schema().to_string();
10073        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10074
10075        let aggr_line = plan_str
10076            .lines()
10077            .find(|line| line.contains("Aggregate: groupBy="))
10078            .unwrap();
10079        assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10080    }
10081
10082    #[tokio::test]
10083    async fn topk_without_does_not_partition_by_tsid() {
10084        let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap();
10085        let eval_stmt = EvalStmt {
10086            expr: prom_expr,
10087            start: UNIX_EPOCH,
10088            end: UNIX_EPOCH
10089                .checked_add(Duration::from_secs(100_000))
10090                .unwrap(),
10091            interval: Duration::from_secs(5),
10092            lookback_delta: Duration::from_secs(1),
10093        };
10094
10095        let table_provider = build_test_table_provider_with_tsid(
10096            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10097            1,
10098            1,
10099        )
10100        .await;
10101        let plan =
10102            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10103                .await
10104                .unwrap();
10105
10106        let plan_str = plan.display_indent_schema().to_string();
10107        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10108
10109        let window_line = plan_str
10110            .lines()
10111            .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10112            .unwrap();
10113        let partition_by = window_line
10114            .split("PARTITION BY [")
10115            .nth(1)
10116            .and_then(|s| s.split("] ORDER BY").next())
10117            .unwrap();
10118        assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10119    }
10120
10121    #[tokio::test]
10122    async fn sum_by_does_not_group_by_tsid() {
10123        let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap();
10124        let eval_stmt = EvalStmt {
10125            expr: prom_expr,
10126            start: UNIX_EPOCH,
10127            end: UNIX_EPOCH
10128                .checked_add(Duration::from_secs(100_000))
10129                .unwrap(),
10130            interval: Duration::from_secs(5),
10131            lookback_delta: Duration::from_secs(1),
10132        };
10133
10134        let table_provider = build_test_table_provider_with_tsid(
10135            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10136            1,
10137            1,
10138        )
10139        .await;
10140        let plan =
10141            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10142                .await
10143                .unwrap();
10144
10145        let plan_str = plan.display_indent_schema().to_string();
10146        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10147
10148        let aggr_line = plan_str
10149            .lines()
10150            .find(|line| line.contains("Aggregate: groupBy="))
10151            .unwrap();
10152        assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10153    }
10154
10155    #[tokio::test]
10156    async fn aggregate_over_binary_time_function_expr() {
10157        for op in ["sum", "min", "max", "avg"] {
10158            let prom_expr = parser::parse(&format!(
10159                "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)"
10160            ))
10161            .unwrap();
10162            let eval_stmt = EvalStmt {
10163                expr: prom_expr,
10164                start: UNIX_EPOCH,
10165                end: UNIX_EPOCH
10166                    .checked_add(Duration::from_secs(100_000))
10167                    .unwrap(),
10168                interval: Duration::from_secs(5),
10169                lookback_delta: Duration::from_secs(1),
10170            };
10171
10172            let table_provider = build_test_table_provider_with_tsid(
10173                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10174                3,
10175                1,
10176            )
10177            .await;
10178            let plan =
10179                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10180                    .await
10181                    .unwrap();
10182
10183            let plan_str = plan.display_indent_schema().to_string();
10184            let aggr_line = plan_str
10185                .lines()
10186                .find(|line| line.contains("Aggregate: groupBy="))
10187                .unwrap();
10188            assert!(aggr_line.contains(op), "{plan_str}");
10189            assert!(aggr_line.contains("first_value"), "{plan_str}");
10190            assert!(
10191                !plan
10192                    .schema()
10193                    .fields()
10194                    .iter()
10195                    .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME })
10196            );
10197        }
10198    }
10199
10200    #[tokio::test]
10201    async fn topk_by_does_not_partition_by_tsid() {
10202        let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap();
10203        let eval_stmt = EvalStmt {
10204            expr: prom_expr,
10205            start: UNIX_EPOCH,
10206            end: UNIX_EPOCH
10207                .checked_add(Duration::from_secs(100_000))
10208                .unwrap(),
10209            interval: Duration::from_secs(5),
10210            lookback_delta: Duration::from_secs(1),
10211        };
10212
10213        let table_provider = build_test_table_provider_with_tsid(
10214            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10215            1,
10216            1,
10217        )
10218        .await;
10219        let plan =
10220            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10221                .await
10222                .unwrap();
10223
10224        let plan_str = plan.display_indent_schema().to_string();
10225        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10226
10227        let window_line = plan_str
10228            .lines()
10229            .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10230            .unwrap();
10231        let partition_by = window_line
10232            .split("PARTITION BY [")
10233            .nth(1)
10234            .and_then(|s| s.split("] ORDER BY").next())
10235            .unwrap();
10236        assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10237    }
10238
10239    #[tokio::test]
10240    async fn selector_matcher_on_tsid_does_not_use_internal_column() {
10241        let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap();
10242        let eval_stmt = EvalStmt {
10243            expr: prom_expr,
10244            start: UNIX_EPOCH,
10245            end: UNIX_EPOCH
10246                .checked_add(Duration::from_secs(100_000))
10247                .unwrap(),
10248            interval: Duration::from_secs(5),
10249            lookback_delta: Duration::from_secs(1),
10250        };
10251
10252        let table_provider = build_test_table_provider_with_tsid(
10253            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10254            1,
10255            1,
10256        )
10257        .await;
10258        let plan =
10259            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10260                .await
10261                .unwrap();
10262
10263        fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet<Column>) {
10264            if let LogicalPlan::Filter(filter) = plan {
10265                datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap();
10266            }
10267            for input in plan.inputs() {
10268                collect_filter_cols(input, out);
10269            }
10270        }
10271
10272        let mut filter_cols = HashSet::new();
10273        collect_filter_cols(&plan, &mut filter_cols);
10274        assert!(
10275            !filter_cols
10276                .iter()
10277                .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME)
10278        );
10279    }
10280
10281    #[tokio::test]
10282    async fn tsid_is_not_used_when_physical_table_is_missing() {
10283        let prom_expr = parser::parse("some_metric").unwrap();
10284        let eval_stmt = EvalStmt {
10285            expr: prom_expr,
10286            start: UNIX_EPOCH,
10287            end: UNIX_EPOCH
10288                .checked_add(Duration::from_secs(100_000))
10289                .unwrap(),
10290            interval: Duration::from_secs(5),
10291            lookback_delta: Duration::from_secs(1),
10292        };
10293
10294        let catalog_list = MemoryCatalogManager::with_default_setup();
10295
10296        // Register a metric engine logical table referencing a missing physical table.
10297        let mut columns = vec![ColumnSchema::new(
10298            "tag_0".to_string(),
10299            ConcreteDataType::string_datatype(),
10300            false,
10301        )];
10302        columns.push(
10303            ColumnSchema::new(
10304                "timestamp".to_string(),
10305                ConcreteDataType::timestamp_millisecond_datatype(),
10306                false,
10307            )
10308            .with_time_index(true),
10309        );
10310        columns.push(ColumnSchema::new(
10311            "field_0".to_string(),
10312            ConcreteDataType::float64_datatype(),
10313            true,
10314        ));
10315        let schema = Arc::new(Schema::new(columns));
10316        let mut options = table::requests::TableOptions::default();
10317        options
10318            .extra_options
10319            .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string());
10320        let table_meta = TableMetaBuilder::empty()
10321            .schema(schema)
10322            .primary_key_indices(vec![0])
10323            .value_indices(vec![2])
10324            .engine(METRIC_ENGINE_NAME.to_string())
10325            .options(options)
10326            .next_column_id(1024)
10327            .build()
10328            .unwrap();
10329        let table_info = TableInfoBuilder::default()
10330            .table_id(1024)
10331            .name("some_metric")
10332            .meta(table_meta)
10333            .build()
10334            .unwrap();
10335        let table = EmptyTable::from_table_info(&table_info);
10336        catalog_list
10337            .register_table_sync(RegisterTableRequest {
10338                catalog: DEFAULT_CATALOG_NAME.to_string(),
10339                schema: DEFAULT_SCHEMA_NAME.to_string(),
10340                table_name: "some_metric".to_string(),
10341                table_id: 1024,
10342                table,
10343            })
10344            .unwrap();
10345
10346        let table_provider = DfTableSourceProvider::new(
10347            catalog_list,
10348            false,
10349            QueryContext::arc(),
10350            DummyDecoder::arc(),
10351            false,
10352        );
10353
10354        let plan =
10355            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10356                .await
10357                .unwrap();
10358
10359        let plan_str = plan.display_indent_schema().to_string();
10360        assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
10361        assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10362    }
10363
10364    #[tokio::test]
10365    async fn tsid_is_carried_only_when_aggregate_preserves_label_set() {
10366        let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap();
10367        let eval_stmt = EvalStmt {
10368            expr: prom_expr,
10369            start: UNIX_EPOCH,
10370            end: UNIX_EPOCH
10371                .checked_add(Duration::from_secs(100_000))
10372                .unwrap(),
10373            interval: Duration::from_secs(5),
10374            lookback_delta: Duration::from_secs(1),
10375        };
10376
10377        let table_provider = build_test_table_provider_with_tsid(
10378            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10379            1,
10380            1,
10381        )
10382        .await;
10383        let plan =
10384            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10385                .await
10386                .unwrap();
10387
10388        let plan_str = plan.display_indent_schema().to_string();
10389        assert!(plan_str.contains("first_value") && plan_str.contains("__tsid"));
10390        assert!(
10391            !plan
10392                .schema()
10393                .fields()
10394                .iter()
10395                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
10396        );
10397
10398        // Merging aggregate: label set is reduced, tsid should not be carried.
10399        let prom_expr = parser::parse("sum(some_metric)").unwrap();
10400        let eval_stmt = EvalStmt {
10401            expr: prom_expr,
10402            start: UNIX_EPOCH,
10403            end: UNIX_EPOCH
10404                .checked_add(Duration::from_secs(100_000))
10405                .unwrap(),
10406            interval: Duration::from_secs(5),
10407            lookback_delta: Duration::from_secs(1),
10408        };
10409        let table_provider = build_test_table_provider_with_tsid(
10410            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10411            1,
10412            1,
10413        )
10414        .await;
10415        let plan =
10416            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10417                .await
10418                .unwrap();
10419        let plan_str = plan.display_indent_schema().to_string();
10420        assert!(!plan_str.contains("first_value"));
10421    }
10422
10423    #[tokio::test]
10424    async fn or_operator_with_unknown_metric_does_not_require_tsid() {
10425        let prom_expr = parser::parse("unknown_metric or some_metric").unwrap();
10426        let eval_stmt = EvalStmt {
10427            expr: prom_expr,
10428            start: UNIX_EPOCH,
10429            end: UNIX_EPOCH
10430                .checked_add(Duration::from_secs(100_000))
10431                .unwrap(),
10432            interval: Duration::from_secs(5),
10433            lookback_delta: Duration::from_secs(1),
10434        };
10435
10436        let table_provider = build_test_table_provider_with_tsid(
10437            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10438            1,
10439            1,
10440        )
10441        .await;
10442
10443        let plan =
10444            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10445                .await
10446                .unwrap();
10447
10448        assert!(
10449            !plan
10450                .schema()
10451                .fields()
10452                .iter()
10453                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
10454        );
10455    }
10456
10457    #[tokio::test]
10458    async fn aggregate_avg() {
10459        do_aggregate_expr_plan("avg", "avg").await;
10460    }
10461
10462    #[tokio::test]
10463    #[should_panic] // output type doesn't match
10464    async fn aggregate_count() {
10465        do_aggregate_expr_plan("count", "count").await;
10466    }
10467
10468    #[tokio::test]
10469    async fn aggregate_min() {
10470        do_aggregate_expr_plan("min", "min").await;
10471    }
10472
10473    #[tokio::test]
10474    async fn aggregate_max() {
10475        do_aggregate_expr_plan("max", "max").await;
10476    }
10477
10478    #[tokio::test]
10479    async fn aggregate_group() {
10480        // Regression test for `group()` aggregator.
10481        // PromQL: sum(group by (cluster)(kubernetes_build_info{service="kubernetes",job="apiserver"}))
10482        // should be plannable, and `group()` should produce constant 1 for each group.
10483        let prom_expr = parser::parse(
10484            "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))",
10485        )
10486        .unwrap();
10487        let eval_stmt = EvalStmt {
10488            expr: prom_expr,
10489            start: UNIX_EPOCH,
10490            end: UNIX_EPOCH
10491                .checked_add(Duration::from_secs(100_000))
10492                .unwrap(),
10493            interval: Duration::from_secs(5),
10494            lookback_delta: Duration::from_secs(1),
10495        };
10496
10497        let table_provider = build_test_table_provider_with_fields(
10498            &[(
10499                DEFAULT_SCHEMA_NAME.to_string(),
10500                "kubernetes_build_info".to_string(),
10501            )],
10502            &["cluster", "service", "job"],
10503        )
10504        .await;
10505        let plan =
10506            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10507                .await
10508                .unwrap();
10509
10510        let plan_str = plan.display_indent_schema().to_string();
10511        assert!(plan_str.contains("max(Float64(1"));
10512    }
10513
10514    #[tokio::test]
10515    async fn aggregate_stddev() {
10516        do_aggregate_expr_plan("stddev", "stddev_pop").await;
10517    }
10518
10519    #[tokio::test]
10520    async fn aggregate_stdvar() {
10521        do_aggregate_expr_plan("stdvar", "var_pop").await;
10522    }
10523
10524    // TODO(ruihang): add range fn tests once exprs are ready.
10525
10526    // {
10527    //     input: "some_metric{tag_0="foo"} + some_metric{tag_0="bar"}",
10528    //     expected: &BinaryExpr{
10529    //         Op: ADD,
10530    //         LHS: &VectorSelector{
10531    //             Name: "a",
10532    //             LabelMatchers: []*labels.Matcher{
10533    //                     MustLabelMatcher(labels.MatchEqual, "tag_0", "foo"),
10534    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
10535    //             },
10536    //         },
10537    //         RHS: &VectorSelector{
10538    //             Name: "sum",
10539    //             LabelMatchers: []*labels.Matcher{
10540    //                     MustLabelMatcher(labels.MatchxEqual, "tag_0", "bar"),
10541    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
10542    //             },
10543    //         },
10544    //         VectorMatching: &VectorMatching{},
10545    //     },
10546    // },
10547    #[tokio::test]
10548    async fn binary_op_column_column() {
10549        let prom_expr =
10550            parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap();
10551        let eval_stmt = EvalStmt {
10552            expr: prom_expr,
10553            start: UNIX_EPOCH,
10554            end: UNIX_EPOCH
10555                .checked_add(Duration::from_secs(100_000))
10556                .unwrap(),
10557            interval: Duration::from_secs(5),
10558            lookback_delta: Duration::from_secs(1),
10559        };
10560
10561        let table_provider = build_test_table_provider(
10562            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10563            1,
10564            1,
10565        )
10566        .await;
10567        let plan =
10568            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10569                .await
10570                .unwrap();
10571
10572        let expected = String::from(
10573            "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]\
10574            \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]\
10575            \n    SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10576            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10577            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10578            \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]\
10579            \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]\
10580            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10581            \n    SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10582            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10583            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10584            \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]\
10585            \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]\
10586            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10587        );
10588
10589        assert_eq!(plan.display_indent_schema().to_string(), expected);
10590    }
10591
10592    async fn indie_query_plan_compare<T: AsRef<str>>(query: &str, expected: T) {
10593        let prom_expr = parser::parse(query).unwrap();
10594        let eval_stmt = EvalStmt {
10595            expr: prom_expr,
10596            start: UNIX_EPOCH,
10597            end: UNIX_EPOCH
10598                .checked_add(Duration::from_secs(100_000))
10599                .unwrap(),
10600            interval: Duration::from_secs(5),
10601            lookback_delta: Duration::from_secs(1),
10602        };
10603
10604        let table_provider = build_test_table_provider(
10605            &[
10606                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10607                (
10608                    "greptime_private".to_string(),
10609                    "some_alt_metric".to_string(),
10610                ),
10611            ],
10612            1,
10613            1,
10614        )
10615        .await;
10616        let plan =
10617            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10618                .await
10619                .unwrap();
10620
10621        assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref());
10622    }
10623
10624    #[tokio::test]
10625    async fn binary_op_literal_column() {
10626        let query = r#"1 + some_metric{tag_0="bar"}"#;
10627        let expected = String::from(
10628            "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]\
10629            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10630            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10631            \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]\
10632            \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]\
10633            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10634        );
10635
10636        indie_query_plan_compare(query, expected).await;
10637    }
10638
10639    #[tokio::test]
10640    async fn binary_op_literal_literal() {
10641        let query = r#"1 + 1"#;
10642        let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N]
10643  TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#;
10644        indie_query_plan_compare(query, expected).await;
10645    }
10646
10647    #[tokio::test]
10648    async fn simple_bool_grammar() {
10649        let query = "some_metric != bool 1.2345";
10650        let expected = String::from(
10651            "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]\
10652            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10653            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10654            \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]\
10655            \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]\
10656            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10657        );
10658
10659        indie_query_plan_compare(query, expected).await;
10660    }
10661
10662    #[tokio::test]
10663    async fn bool_with_additional_arithmetic() {
10664        let query = "some_metric + (1 == bool 2)";
10665        let expected = String::from(
10666            "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]\
10667            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10668            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10669            \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]\
10670            \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]\
10671            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10672        );
10673
10674        indie_query_plan_compare(query, expected).await;
10675    }
10676
10677    #[tokio::test]
10678    async fn simple_unary() {
10679        let query = "-some_metric";
10680        let expected = String::from(
10681            "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]\
10682            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10683            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10684            \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]\
10685            \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]\
10686            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10687        );
10688
10689        indie_query_plan_compare(query, expected).await;
10690    }
10691
10692    #[tokio::test]
10693    async fn increase_aggr() {
10694        let query = "increase(some_metric[5m])";
10695        let expected = String::from(
10696            "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]\
10697            \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]\
10698            \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))]\
10699            \n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10700            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10701            \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]\
10702            \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]\
10703            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10704        );
10705
10706        indie_query_plan_compare(query, expected).await;
10707    }
10708
10709    async fn native_histogram_plan(query: &str) -> String {
10710        let table_provider = build_test_native_histogram_table_provider("some_metric").await;
10711        let plan = PromPlanner::stmt_to_plan(
10712            table_provider,
10713            &build_eval_stmt(query),
10714            &build_query_engine_state(),
10715        )
10716        .await
10717        .unwrap();
10718        plan.display_indent_schema().to_string()
10719    }
10720
10721    #[tokio::test]
10722    async fn native_histogram_count_uses_native_udf() {
10723        let plan = native_histogram_plan("histogram_count(some_metric)").await;
10724
10725        assert!(plan.contains("prom_native_histogram_count"), "{plan}");
10726        assert!(!plan.contains("HistogramFold:"), "{plan}");
10727    }
10728
10729    #[tokio::test]
10730    async fn timestamp_filters_native_histogram_stale_marker_before_projection() {
10731        let mut stale = direct_or_histogram();
10732        stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
10733        let table = operator_metric_table(
10734            "stale_histogram",
10735            2_100,
10736            "a",
10737            None,
10738            DirectOrValue::NativeHistogram(stale),
10739        );
10740        let catalog = MemoryCatalogManager::with_default_setup();
10741        catalog
10742            .register_table_sync(RegisterTableRequest {
10743                catalog: DEFAULT_CATALOG_NAME.to_string(),
10744                schema: DEFAULT_SCHEMA_NAME.to_string(),
10745                table_name: "stale_histogram".to_string(),
10746                table_id: 2_100,
10747                table,
10748            })
10749            .unwrap();
10750        let provider = DfTableSourceProvider::new(
10751            catalog,
10752            false,
10753            QueryContext::arc(),
10754            DummyDecoder::arc(),
10755            false,
10756        );
10757        let state = build_query_engine_state();
10758        let plan = PromPlanner::stmt_to_plan(
10759            provider,
10760            &operator_eval_stmt("timestamp(stale_histogram)"),
10761            &state,
10762        )
10763        .await
10764        .unwrap();
10765        let plan_text = plan.display_indent_schema().to_string();
10766        assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}");
10767
10768        let (_, batches) = execute(plan, &state).await;
10769        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
10770    }
10771
10772    #[tokio::test]
10773    async fn timestamp_filters_stale_marker_from_mixed_sample_companion() {
10774        let histograms = build_histogram_array(&[None]);
10775        let schema = Arc::new(ArrowSchema::new(vec![
10776            Field::new(
10777                "timestamp",
10778                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
10779                false,
10780            ),
10781            Field::new(
10782                greptime_native_histogram(),
10783                histograms.data_type().clone(),
10784                true,
10785            ),
10786            Field::new(greptime_value(), ArrowDataType::Float64, true),
10787        ]));
10788        let batch = RecordBatch::try_new(
10789            schema.clone(),
10790            vec![
10791                Arc::new(TimestampMillisecondArray::from(vec![1_000])),
10792                histograms,
10793                Arc::new(Float64Array::from(vec![f64::from_bits(
10794                    PROMETHEUS_STALE_NAN_BITS,
10795                )])),
10796            ],
10797        )
10798        .unwrap();
10799        let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
10800        let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
10801            .unwrap()
10802            .build()
10803            .unwrap();
10804        let input = LogicalPlan::Extension(Extension {
10805            node: Arc::new(SeriesDivide::new(
10806                Vec::new(),
10807                "timestamp".to_string(),
10808                input,
10809            )),
10810        });
10811        let input = LogicalPlan::Extension(Extension {
10812            node: Arc::new(InstantManipulate::new(
10813                1_000,
10814                1_000,
10815                5_000,
10816                1_000,
10817                "timestamp".to_string(),
10818                Vec::new(),
10819                Some(greptime_native_histogram().to_string()),
10820                input,
10821            )),
10822        });
10823        // Match timestamp()'s parent projection, which otherwise prunes the companion lane.
10824        let plan = LogicalPlanBuilder::from(input)
10825            .project([col("timestamp")])
10826            .unwrap()
10827            .build()
10828            .unwrap();
10829
10830        let (_, batches) = execute(plan, &build_query_engine_state()).await;
10831        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
10832    }
10833
10834    #[tokio::test]
10835    async fn native_histogram_rate_can_feed_count() {
10836        let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await;
10837
10838        assert!(plan.contains("prom_native_histogram_rate"), "{plan}");
10839        assert!(plan.contains("prom_native_histogram_count"), "{plan}");
10840    }
10841
10842    #[tokio::test]
10843    async fn native_histogram_quantile_skips_classic_fold() {
10844        let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await;
10845
10846        assert!(plan.contains("prom_native_histogram_quantile"), "{plan}");
10847        assert!(!plan.contains("HistogramFold:"), "{plan}");
10848        assert!(plan.contains("some_metric.le"), "{plan}");
10849        // The phi literal is threaded into the native quantile UDF as its second argument.
10850        assert!(plan.contains("Float64(0.9)"), "{plan}");
10851        // The empty-values filter drops NULL quantile results so the output is empty
10852        // when all native histogram samples are dropped.
10853        assert!(plan.contains("IS NOT NULL"), "{plan}");
10854    }
10855
10856    #[tokio::test]
10857    async fn mixed_native_histogram_quantile_uses_histogram_field() {
10858        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
10859        let plan = PromPlanner::stmt_to_plan(
10860            table_provider,
10861            &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
10862            &build_query_engine_state(),
10863        )
10864        .await
10865        .unwrap()
10866        .display_indent_schema()
10867        .to_string();
10868
10869        assert!(
10870            plan.contains("prom_native_histogram_quantile(greptime_native_histogram"),
10871            "{plan}"
10872        );
10873        assert!(!plan.contains("EmptyRelation"), "{plan}");
10874    }
10875
10876    #[tokio::test]
10877    async fn mixed_histogram_helpers_execute_classic_and_native_samples() {
10878        let state = build_query_engine_state();
10879        for (query, expected) in [
10880            (
10881                "histogram_quantile(0.5, mixed_histogram)",
10882                vec![("classic", 1.0), ("native", 0.0)],
10883            ),
10884            (
10885                "histogram_fraction(-Inf, +Inf, mixed_histogram)",
10886                vec![("classic", 1.0), ("native", 1.0)],
10887            ),
10888        ] {
10889            let plan = PromPlanner::stmt_to_plan(
10890                classic_and_native_histogram_table_provider("native", None, direct_or_histogram()),
10891                &operator_eval_stmt(query),
10892                &state,
10893            )
10894            .await
10895            .unwrap();
10896            let plan_text = plan.display_indent_schema().to_string();
10897            assert!(plan_text.contains("HistogramFold:"), "{plan_text}");
10898            assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}");
10899            let value_field = plan
10900                .schema()
10901                .fields()
10902                .iter()
10903                .find(|field| field.data_type() == &ArrowDataType::Float64)
10904                .unwrap()
10905                .name()
10906                .clone();
10907
10908            let (_, batches) = execute(plan, &state).await;
10909            let mut actual = batches
10910                .iter()
10911                .flat_map(|batch| {
10912                    let tags = batch
10913                        .column_by_name("tag")
10914                        .unwrap()
10915                        .as_any()
10916                        .downcast_ref::<StringArray>()
10917                        .unwrap();
10918                    let values = batch
10919                        .column_by_name(&value_field)
10920                        .unwrap()
10921                        .as_any()
10922                        .downcast_ref::<Float64Array>()
10923                        .unwrap();
10924                    (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row)))
10925                })
10926                .collect::<Vec<_>>();
10927            actual.sort_by_key(|(tag, _)| *tag);
10928            assert_eq!(actual, expected, "{query}");
10929        }
10930    }
10931
10932    #[tokio::test]
10933    async fn mixed_histogram_helpers_report_annotations() {
10934        let state = build_query_engine_state();
10935        let mut native_histogram = direct_or_histogram();
10936        native_histogram.count = 2.0;
10937        native_histogram.sum = f64::NAN;
10938        for (native_tag, expected_rows, expected_warnings, expected_infos) in [
10939            (
10940                "classic",
10941                0,
10942                vec!["vector contains a mix of classic and native histograms"],
10943                vec![],
10944            ),
10945            (
10946                "native",
10947                2,
10948                vec![],
10949                vec!["input to histogram_quantile has NaN observations, result is skewed higher"],
10950            ),
10951        ] {
10952            let collector = PromqlAnnotationCollector::default();
10953            let plan = PromPlanner::stmt_to_plan_with_annotations(
10954                classic_and_native_histogram_table_provider(
10955                    native_tag,
10956                    None,
10957                    native_histogram.clone(),
10958                ),
10959                &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"),
10960                &state,
10961                Some(collector.clone()),
10962            )
10963            .await
10964            .unwrap();
10965
10966            let (_, batches) = execute(plan, &state).await;
10967            assert_eq!(
10968                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
10969                expected_rows
10970            );
10971            let mut warnings = vec![];
10972            let mut infos = vec![];
10973            collector.append_to(&mut warnings, &mut infos);
10974            assert_eq!(warnings, expected_warnings);
10975            assert_eq!(infos, expected_infos);
10976        }
10977    }
10978
10979    #[tokio::test]
10980    async fn mixed_histogram_helper_preserves_native_le_and_scans_once() {
10981        let state = build_query_engine_state();
10982        let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)");
10983        stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap();
10984        let plan = PromPlanner::stmt_to_plan(
10985            classic_and_native_histogram_table_provider(
10986                "classic",
10987                Some("native"),
10988                direct_or_histogram(),
10989            ),
10990            &stmt,
10991            &state,
10992        )
10993        .await
10994        .unwrap();
10995        let plan_text = plan.display_indent_schema().to_string();
10996        assert_eq!(
10997            plan_text.matches("TableScan: mixed_histogram").count(),
10998            1,
10999            "{plan_text}"
11000        );
11001
11002        let value_field = plan
11003            .schema()
11004            .fields()
11005            .iter()
11006            .find(|field| field.data_type() == &ArrowDataType::Float64)
11007            .unwrap()
11008            .name()
11009            .clone();
11010        let (_, batches) = execute(plan, &state).await;
11011        let mut actual = batches
11012            .iter()
11013            .flat_map(|batch| {
11014                let le = batch
11015                    .column_by_name(LE_COLUMN_NAME)
11016                    .unwrap()
11017                    .as_any()
11018                    .downcast_ref::<StringArray>()
11019                    .unwrap();
11020                let timestamps = batch
11021                    .column_by_name("timestamp")
11022                    .unwrap()
11023                    .as_any()
11024                    .downcast_ref::<TimestampMillisecondArray>()
11025                    .unwrap();
11026                let values = batch
11027                    .column_by_name(&value_field)
11028                    .unwrap()
11029                    .as_any()
11030                    .downcast_ref::<Float64Array>()
11031                    .unwrap();
11032                (0..batch.num_rows()).map(|row| {
11033                    (
11034                        timestamps.value(row),
11035                        (!le.is_null(row)).then(|| le.value(row).to_string()),
11036                        values.value(row),
11037                    )
11038                })
11039            })
11040            .collect::<Vec<_>>();
11041        actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1)));
11042        assert_eq!(
11043            actual,
11044            vec![
11045                (1_000, None, 1.0),
11046                (1_000, Some("native".to_string()), 0.0),
11047                (2_000, None, 1.0),
11048                (2_000, Some("native".to_string()), 0.0),
11049            ]
11050        );
11051    }
11052
11053    #[tokio::test]
11054    async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() {
11055        let state = build_query_engine_state();
11056        for native_le in [None, Some("native")] {
11057            for query in [
11058                "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))",
11059                "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))",
11060            ] {
11061                let plan = PromPlanner::stmt_to_plan(
11062                    classic_and_native_histogram_table_provider(
11063                        "native",
11064                        native_le,
11065                        direct_or_histogram(),
11066                    ),
11067                    &operator_eval_stmt(query),
11068                    &state,
11069                )
11070                .await
11071                .unwrap();
11072
11073                let (_, batches) = execute(plan, &state).await;
11074                assert_eq!(
11075                    batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
11076                    0,
11077                    "native_le={native_le:?}, query={query}"
11078                );
11079            }
11080        }
11081    }
11082
11083    #[tokio::test]
11084    async fn native_histogram_quantile_rejects_multi_field_input() {
11085        let table_provider = build_test_multi_histogram_table_provider("some_metric").await;
11086        let result = PromPlanner::stmt_to_plan(
11087            table_provider,
11088            &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
11089            &build_query_engine_state(),
11090        )
11091        .await;
11092
11093        let err = result.expect_err("histogram_quantile on two native histogram fields must fail");
11094        assert!(
11095            err.to_string()
11096                .contains("Multi fields calculation is not supported in histogram_quantile"),
11097            "{err}"
11098        );
11099    }
11100
11101    #[tokio::test]
11102    async fn native_histogram_topk_uses_drop_udf() {
11103        let plan = native_histogram_plan("topk(1, some_metric)").await;
11104
11105        assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}");
11106        assert!(
11107            plan.contains("Filter: prom_native_histogram_drop_float")
11108                && plan.contains("IS NOT NULL"),
11109            "{plan}"
11110        );
11111    }
11112
11113    #[tokio::test]
11114    async fn mixed_or_topk_bottomk_ignore_native_histograms() {
11115        for op in ["topk", "bottomk"] {
11116            let collector = PromqlAnnotationCollector::default();
11117            let state = build_query_engine_state();
11118            let plan = PromPlanner::stmt_to_plan_with_annotations(
11119                operator_table_provider(),
11120                &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")),
11121                &state,
11122                Some(collector.clone()),
11123            )
11124            .await
11125            .unwrap();
11126            let float_field = plan
11127                .schema()
11128                .fields()
11129                .iter()
11130                .find(|field| field.data_type() == &ArrowDataType::Float64)
11131                .unwrap()
11132                .name()
11133                .clone();
11134            assert!(
11135                plan.schema()
11136                    .fields()
11137                    .iter()
11138                    .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()),
11139                "{plan:?}"
11140            );
11141
11142            let (_, batches) = execute(plan, &state).await;
11143            assert_eq!(values(&batches, &float_field), vec![2.0], "{op}");
11144            let mut warnings = vec![];
11145            let mut infos = vec![];
11146            collector.append_to(&mut warnings, &mut infos);
11147            assert!(warnings.is_empty());
11148            assert_eq!(
11149                infos,
11150                vec![format!(
11151                    "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
11152                )]
11153            );
11154        }
11155    }
11156
11157    #[tokio::test]
11158    async fn native_histogram_scalar_is_ignored_before_scalar_calculate() {
11159        let plan = native_histogram_plan("scalar(some_metric)").await;
11160
11161        assert!(plan.contains("ScalarCalculate"), "{plan}");
11162        assert!(plan.contains("Filter: Boolean(false)"), "{plan}");
11163        assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11164    }
11165
11166    #[tokio::test]
11167    async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() {
11168        for function in ["sort", "sort_desc"] {
11169            let plan = native_histogram_plan(&format!("{function}(some_metric)")).await;
11170
11171            assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11172            assert!(
11173                !plan.contains(&format!("Sort: {}", greptime_native_histogram())),
11174                "{plan}"
11175            );
11176            assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11177        }
11178
11179        for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] {
11180            let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await;
11181
11182            assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}");
11183            assert!(plan.contains(greptime_native_histogram()), "{plan}");
11184            assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11185        }
11186    }
11187
11188    #[tokio::test]
11189    async fn unsupported_native_histogram_functions_use_drop_udf() {
11190        for query in [
11191            "deriv(some_metric[5m])",
11192            "min_over_time(some_metric[5m])",
11193            "quantile_over_time(0.9, some_metric[5m])",
11194            "predict_linear(some_metric[5m], 60)",
11195            "round(some_metric)",
11196            "abs(some_metric)",
11197        ] {
11198            let plan = native_histogram_plan(query).await;
11199
11200            assert!(
11201                plan.contains("prom_native_histogram_drop_float"),
11202                "{query}\n{plan}"
11203            );
11204        }
11205    }
11206
11207    #[tokio::test]
11208    async fn native_histogram_absent_over_time_uses_native_udf() {
11209        let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await;
11210
11211        assert!(
11212            plan.contains("prom_native_histogram_absent_over_time"),
11213            "{plan}"
11214        );
11215    }
11216
11217    #[tokio::test]
11218    async fn native_histogram_all_function_arms_route_correctly() {
11219        // Every native-histogram match arm in `create_function_expr` must route to the
11220        // expected UDF when all field columns are native histograms. `holt_winters` shares
11221        // the `double_exponential_smoothing` arm but is not registered in the promql
11222        // parser (0.10), so it cannot be exercised through a query string.
11223        let cases = [
11224            // Range functions routed to native histogram UDFs.
11225            (
11226                "increase(some_metric[5m])",
11227                "prom_native_histogram_increase",
11228            ),
11229            ("rate(some_metric[5m])", "prom_native_histogram_rate"),
11230            ("delta(some_metric[5m])", "prom_native_histogram_delta"),
11231            ("idelta(some_metric[5m])", "prom_native_histogram_idelta"),
11232            ("irate(some_metric[5m])", "prom_native_histogram_irate"),
11233            ("resets(some_metric[5m])", "prom_native_histogram_resets"),
11234            ("changes(some_metric[5m])", "prom_native_histogram_changes"),
11235            (
11236                "avg_over_time(some_metric[5m])",
11237                "prom_native_histogram_avg_over_time",
11238            ),
11239            (
11240                "sum_over_time(some_metric[5m])",
11241                "prom_native_histogram_sum_over_time",
11242            ),
11243            (
11244                "count_over_time(some_metric[5m])",
11245                "prom_native_histogram_count_over_time",
11246            ),
11247            (
11248                "last_over_time(some_metric[5m])",
11249                "prom_native_histogram_last_over_time",
11250            ),
11251            (
11252                "present_over_time(some_metric[5m])",
11253                "prom_native_histogram_present_over_time",
11254            ),
11255            // Unsupported functions dropped with the float-null UDF.
11256            ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"),
11257            (
11258                "min_over_time(some_metric[5m])",
11259                "prom_native_histogram_drop_float",
11260            ),
11261            (
11262                "max_over_time(some_metric[5m])",
11263                "prom_native_histogram_drop_float",
11264            ),
11265            (
11266                "stddev_over_time(some_metric[5m])",
11267                "prom_native_histogram_drop_float",
11268            ),
11269            (
11270                "stdvar_over_time(some_metric[5m])",
11271                "prom_native_histogram_drop_float",
11272            ),
11273            (
11274                "quantile_over_time(0.9, some_metric[5m])",
11275                "prom_native_histogram_drop_float",
11276            ),
11277            (
11278                "predict_linear(some_metric[5m], 60)",
11279                "prom_native_histogram_drop_float",
11280            ),
11281            (
11282                "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11283                "prom_native_histogram_drop_float",
11284            ),
11285            ("round(some_metric)", "prom_native_histogram_drop_float"),
11286            ("rad(some_metric)", "prom_native_histogram_drop_float"),
11287            ("deg(some_metric)", "prom_native_histogram_drop_float"),
11288            ("sgn(some_metric)", "prom_native_histogram_drop_float"),
11289            // Instant helper functions routed to native histogram UDFs.
11290            (
11291                "histogram_count(some_metric)",
11292                "prom_native_histogram_count",
11293            ),
11294            ("histogram_sum(some_metric)", "prom_native_histogram_sum"),
11295            ("histogram_avg(some_metric)", "prom_native_histogram_avg"),
11296            (
11297                "histogram_stddev(some_metric)",
11298                "prom_native_histogram_stddev",
11299            ),
11300            (
11301                "histogram_stdvar(some_metric)",
11302                "prom_native_histogram_stdvar",
11303            ),
11304            (
11305                "histogram_fraction(-2 + 1, 2 / 2, some_metric)",
11306                "prom_native_histogram_fraction",
11307            ),
11308        ];
11309
11310        for (query, expected_udf) in cases {
11311            let plan = native_histogram_plan(query).await;
11312            assert!(plan.contains(expected_udf), "{query}\n{plan}");
11313            if query.starts_with("histogram_fraction") {
11314                assert!(plan.contains("Float64(-1)"), "{query}\n{plan}");
11315            }
11316        }
11317    }
11318
11319    #[tokio::test]
11320    async fn mixed_native_histogram_ranges_use_coordinated_udfs() {
11321        let dual_output = [
11322            "increase(some_metric[5m])",
11323            "rate(some_metric[5m])",
11324            "delta(some_metric[5m])",
11325            "idelta(some_metric[5m])",
11326            "irate(some_metric[5m])",
11327            "avg_over_time(some_metric[5m])",
11328            "sum_over_time(some_metric[5m])",
11329            "last_over_time(some_metric[5m])",
11330        ];
11331        let float_output = [
11332            "resets(some_metric[5m])",
11333            "changes(some_metric[5m])",
11334            "deriv(some_metric[5m])",
11335            "min_over_time(some_metric[5m])",
11336            "max_over_time(some_metric[5m])",
11337            "count_over_time(some_metric[5m])",
11338            "absent_over_time(some_metric[5m])",
11339            "present_over_time(some_metric[5m])",
11340            "stddev_over_time(some_metric[5m])",
11341            "stdvar_over_time(some_metric[5m])",
11342            "quantile_over_time(0.9, some_metric[5m])",
11343            "predict_linear(some_metric[5m], 60)",
11344            "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11345        ];
11346
11347        for query in dual_output.iter().chain(float_output.iter()) {
11348            let plan = PromPlanner::stmt_to_plan(
11349                build_test_mixed_native_histogram_table_provider("some_metric").await,
11350                &build_eval_stmt(query),
11351                &build_query_engine_state(),
11352            )
11353            .await
11354            .unwrap()
11355            .display_indent_schema()
11356            .to_string();
11357            assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}");
11358            assert_eq!(
11359                plan.contains("prom_mixed_range_histogram"),
11360                dual_output.contains(query),
11361                "{query}\n{plan}"
11362            );
11363        }
11364
11365        let plan = PromPlanner::stmt_to_plan(
11366            build_test_mixed_native_histogram_table_provider("some_metric").await,
11367            &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"),
11368            &build_query_engine_state(),
11369        )
11370        .await
11371        .unwrap()
11372        .display_indent_schema()
11373        .to_string();
11374        let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11375  Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11376    PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))]
11377      PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11378        Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11379          Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11380            Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8]
11381              PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]
11382                PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11383                  PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11384                    Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11385                      Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]
11386                        TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#;
11387        assert_eq!(plan, expected);
11388    }
11389
11390    #[tokio::test]
11391    async fn mixed_native_histogram_rate_executes_real_ranges() {
11392        let schema = Arc::new(ArrowSchema::new(vec![
11393            Field::new(
11394                "timestamp",
11395                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
11396                false,
11397            ),
11398            Field::new(greptime_value(), ArrowDataType::Float64, true),
11399            Field::new(
11400                greptime_native_histogram(),
11401                native_histogram_value_type().as_arrow_type(),
11402                true,
11403            ),
11404        ]));
11405        let batch = RecordBatch::try_new(
11406            schema.clone(),
11407            vec![
11408                Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])),
11409                Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])),
11410                build_histogram_array(&[None, Some(direct_or_histogram()), None]),
11411            ],
11412        )
11413        .unwrap();
11414        let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
11415        let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
11416            .unwrap()
11417            .build()
11418            .unwrap();
11419        let collector = PromqlAnnotationCollector::default();
11420        let mut planner = PromPlanner {
11421            table_provider: build_test_table_provider_with_fields(
11422                &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
11423                &[],
11424            )
11425            .await,
11426            ctx: PromPlannerContext {
11427                start: 3000,
11428                end: 3000,
11429                interval: 1000,
11430                range: Some(3000),
11431                time_index_column: Some("timestamp".to_string()),
11432                field_columns: vec![
11433                    greptime_native_histogram().to_string(),
11434                    greptime_value().to_string(),
11435                ],
11436                ..Default::default()
11437            },
11438            promql_annotations: Some(collector.clone()),
11439        };
11440        let input = LogicalPlan::Extension(Extension {
11441            node: Arc::new(
11442                RangeManipulate::new(
11443                    3000,
11444                    3000,
11445                    1000,
11446                    3000,
11447                    "timestamp".to_string(),
11448                    planner.ctx.field_columns.clone(),
11449                    input,
11450                )
11451                .unwrap(),
11452            ),
11453        });
11454        let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else {
11455            unreachable!()
11456        };
11457        let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
11458            input.schema(),
11459            &planner.ctx.field_columns,
11460        );
11461        let state = build_query_engine_state();
11462        let (mut exprs, _) = planner
11463            .create_function_expr(&call.func, vec![], input.schema(), &state)
11464            .unwrap();
11465        exprs.insert(0, planner.create_time_index_column_expr().unwrap());
11466        let plan = LogicalPlanBuilder::from(input)
11467            .project(exprs)
11468            .unwrap()
11469            .filter(
11470                planner
11471                    .create_empty_values_filter_expr(preserve_any_value)
11472                    .unwrap(),
11473            )
11474            .unwrap()
11475            .build()
11476            .unwrap();
11477        let (_, batches) = execute(plan, &state).await;
11478        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
11479        let mut warnings = Vec::new();
11480        collector.append_to(&mut warnings, &mut Vec::new());
11481        assert!(
11482            warnings
11483                .iter()
11484                .any(|warning| warning.contains("mix of float and native histogram"))
11485        );
11486    }
11487
11488    #[tokio::test]
11489    async fn native_histogram_mixed_field_table_behaves() {
11490        // Exercise function planning after float and histogram samples have already been
11491        // represented as alternative nullable fields. Histogram functions must select the
11492        // histogram field without adding a NULL float field that would reject every row.
11493        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11494        let plan = PromPlanner::stmt_to_plan(
11495            table_provider,
11496            &build_eval_stmt("histogram_count(some_metric)"),
11497            &build_query_engine_state(),
11498        )
11499        .await
11500        .unwrap();
11501        let plan_str = plan.display_indent_schema().to_string();
11502        assert!(
11503            plan_str.contains("prom_native_histogram_count"),
11504            "{plan_str}"
11505        );
11506        assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}");
11507        assert!(
11508            plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"),
11509            "{plan_str}"
11510        );
11511
11512        // Value sorting keeps the float column and never sorts by the histogram column.
11513        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11514        let plan = PromPlanner::stmt_to_plan(
11515            table_provider,
11516            &build_eval_stmt("sort(some_metric)"),
11517            &build_query_engine_state(),
11518        )
11519        .await
11520        .unwrap();
11521        let plan_str = plan.display_indent_schema().to_string();
11522        assert!(
11523            plan_str.contains("greptime_value ASC NULLS FIRST"),
11524            "{plan_str}"
11525        );
11526        assert!(
11527            !plan_str.contains("greptime_native_histogram ASC"),
11528            "{plan_str}"
11529        );
11530
11531        // scalar() ignores histogram samples and evaluates only the float field.
11532        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11533        let plan = PromPlanner::stmt_to_plan(
11534            table_provider,
11535            &build_eval_stmt("scalar(some_metric)"),
11536            &build_query_engine_state(),
11537        )
11538        .await
11539        .unwrap();
11540        let plan_str = plan.display_indent_schema().to_string();
11541        assert!(plan_str.contains("ScalarCalculate"), "{plan_str}");
11542        assert!(
11543            plan_str.contains("greptime_value IS NOT NULL"),
11544            "{plan_str}"
11545        );
11546
11547        // Functions that preserve both alternative fields keep rows with either sample type.
11548        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11549        let plan = PromPlanner::stmt_to_plan(
11550            table_provider,
11551            &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#),
11552            &build_query_engine_state(),
11553        )
11554        .await
11555        .unwrap();
11556        let plan_str = plan.display_indent_schema().to_string();
11557        let filter = plan_str.lines().next().unwrap();
11558        assert!(
11559            filter.starts_with("Filter: ")
11560                && filter.contains("greptime_native_histogram IS NOT NULL")
11561                && filter.contains(" OR ")
11562                && filter.contains("greptime_value IS NOT NULL"),
11563            "{plan_str}"
11564        );
11565    }
11566
11567    #[tokio::test]
11568    async fn less_filter_on_value() {
11569        let query = "some_metric < 1.2345";
11570        let expected = String::from(
11571            "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11572            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11573            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11574            \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]\
11575            \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]\
11576            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11577        );
11578
11579        indie_query_plan_compare(query, expected).await;
11580    }
11581
11582    #[tokio::test]
11583    async fn count_over_time() {
11584        let query = "count_over_time(some_metric[5m])";
11585        let expected = String::from(
11586            "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]\
11587            \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]\
11588            \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))]\
11589            \n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11590            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11591            \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]\
11592            \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]\
11593            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11594        );
11595
11596        indie_query_plan_compare(query, expected).await;
11597    }
11598
11599    /// The outer `PromRangeManipulate` from a subquery must be preceded by
11600    /// `Sort` + `PromSeriesDivide`.
11601    #[tokio::test]
11602    async fn count_over_time_subquery() {
11603        let query = "count_over_time(some_metric[10m:1m])";
11604        let expected = String::from(
11605            "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]\
11606            \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]\
11607            \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))]\
11608            \n      PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11609            \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]\
11610            \n          PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11611            \n            PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11612            \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]\
11613            \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]\
11614            \n                  TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11615        );
11616        indie_query_plan_compare(query, expected).await;
11617    }
11618
11619    #[tokio::test]
11620    async fn test_hash_join() {
11621        let mut eval_stmt = EvalStmt {
11622            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11623            start: UNIX_EPOCH,
11624            end: UNIX_EPOCH
11625                .checked_add(Duration::from_secs(100_000))
11626                .unwrap(),
11627            interval: Duration::from_secs(5),
11628            lookback_delta: Duration::from_secs(1),
11629        };
11630
11631        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"}"#;
11632
11633        let prom_expr = parser::parse(case).unwrap();
11634        eval_stmt.expr = prom_expr;
11635        let table_provider = build_test_table_provider_with_fields(
11636            &[
11637                (
11638                    DEFAULT_SCHEMA_NAME.to_string(),
11639                    "http_server_requests_seconds_sum".to_string(),
11640                ),
11641                (
11642                    DEFAULT_SCHEMA_NAME.to_string(),
11643                    "http_server_requests_seconds_count".to_string(),
11644                ),
11645            ],
11646            &["uri", "kubernetes_namespace", "kubernetes_pod_name"],
11647        )
11648        .await;
11649        // Should be ok
11650        let plan =
11651            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11652                .await
11653                .unwrap();
11654        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\
11655            \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\
11656            \n    SubqueryAlias: http_server_requests_seconds_sum\
11657            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
11658            \n        PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
11659            \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\
11660            \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)\
11661            \n              TableScan: http_server_requests_seconds_sum\
11662            \n    SubqueryAlias: http_server_requests_seconds_count\
11663            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
11664            \n        PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
11665            \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\
11666            \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)\
11667            \n              TableScan: http_server_requests_seconds_count";
11668        assert_eq!(plan.to_string(), expected);
11669    }
11670
11671    #[tokio::test]
11672    async fn test_nested_histogram_quantile() {
11673        let mut eval_stmt = EvalStmt {
11674            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11675            start: UNIX_EPOCH,
11676            end: UNIX_EPOCH
11677                .checked_add(Duration::from_secs(100_000))
11678                .unwrap(),
11679            interval: Duration::from_secs(5),
11680            lookback_delta: Duration::from_secs(1),
11681        };
11682
11683        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]*-(.*)")"#;
11684
11685        let prom_expr = parser::parse(case).unwrap();
11686        eval_stmt.expr = prom_expr;
11687        let table_provider = build_test_table_provider_with_fields(
11688            &[(
11689                DEFAULT_SCHEMA_NAME.to_string(),
11690                "greptime_servers_grpc_requests_elapsed_bucket".to_string(),
11691            )],
11692            &["pod", "le", "path", "code", "container"],
11693        )
11694        .await;
11695        // Should be ok
11696        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11697            .await
11698            .unwrap();
11699    }
11700
11701    #[tokio::test]
11702    async fn test_histogram_quantile_binary_op() {
11703        let mut eval_stmt = EvalStmt {
11704            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11705            start: UNIX_EPOCH,
11706            end: UNIX_EPOCH
11707                .checked_add(Duration::from_secs(100_000))
11708                .unwrap(),
11709            interval: Duration::from_secs(5),
11710            lookback_delta: Duration::from_secs(1),
11711        };
11712
11713        // Arithmetic applied to a histogram_quantile() result. Regression for #8144:
11714        // HistogramFold used to drop the input column qualifiers, so the binary-op
11715        // projection failed to resolve the qualified tag column.
11716        let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#;
11717
11718        let prom_expr = parser::parse(case).unwrap();
11719        eval_stmt.expr = prom_expr;
11720        let table_provider = build_test_table_provider_with_fields(
11721            &[(
11722                DEFAULT_SCHEMA_NAME.to_string(),
11723                "http_request_duration_seconds_bucket".to_string(),
11724            )],
11725            &["pod", "le"],
11726        )
11727        .await;
11728        // Should plan without a "No field named ..." error.
11729        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11730            .await
11731            .unwrap();
11732    }
11733
11734    #[tokio::test]
11735    async fn test_parse_and_operator() {
11736        let mut eval_stmt = EvalStmt {
11737            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11738            start: UNIX_EPOCH,
11739            end: UNIX_EPOCH
11740                .checked_add(Duration::from_secs(100_000))
11741                .unwrap(),
11742            interval: Duration::from_secs(5),
11743            lookback_delta: Duration::from_secs(1),
11744        };
11745
11746        let cases = [
11747            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)"#,
11748            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)"#,
11749        ];
11750
11751        for case in cases {
11752            let prom_expr = parser::parse(case).unwrap();
11753            eval_stmt.expr = prom_expr;
11754            let table_provider = build_test_table_provider_with_fields(
11755                &[
11756                    (
11757                        DEFAULT_SCHEMA_NAME.to_string(),
11758                        "kubelet_volume_stats_used_bytes".to_string(),
11759                    ),
11760                    (
11761                        DEFAULT_SCHEMA_NAME.to_string(),
11762                        "kubelet_volume_stats_capacity_bytes".to_string(),
11763                    ),
11764                ],
11765                &["namespace", "persistentvolumeclaim"],
11766            )
11767            .await;
11768            // Should be ok
11769            let _ =
11770                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11771                    .await
11772                    .unwrap();
11773        }
11774    }
11775
11776    #[tokio::test]
11777    async fn test_nested_binary_op() {
11778        let mut eval_stmt = EvalStmt {
11779            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11780            start: UNIX_EPOCH,
11781            end: UNIX_EPOCH
11782                .checked_add(Duration::from_secs(100_000))
11783                .unwrap(),
11784            interval: Duration::from_secs(5),
11785            lookback_delta: Duration::from_secs(1),
11786        };
11787
11788        let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) -
11789        (
11790            sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m]))
11791            or
11792            vector(0)
11793        )"#;
11794
11795        let prom_expr = parser::parse(case).unwrap();
11796        eval_stmt.expr = prom_expr;
11797        let table_provider = build_test_table_provider_with_fields(
11798            &[(
11799                DEFAULT_SCHEMA_NAME.to_string(),
11800                "nginx_ingress_controller_requests".to_string(),
11801            )],
11802            &["namespace", "job"],
11803        )
11804        .await;
11805        // Should be ok
11806        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11807            .await
11808            .unwrap();
11809    }
11810
11811    #[tokio::test]
11812    async fn test_parse_or_operator() {
11813        let mut eval_stmt = EvalStmt {
11814            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11815            start: UNIX_EPOCH,
11816            end: UNIX_EPOCH
11817                .checked_add(Duration::from_secs(100_000))
11818                .unwrap(),
11819            interval: Duration::from_secs(5),
11820            lookback_delta: Duration::from_secs(1),
11821        };
11822
11823        let case = r#"
11824        sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) /
11825        (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100)
11826            or
11827        200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) /
11828        sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#;
11829
11830        let table_provider = build_test_table_provider_with_fields(
11831            &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
11832            &["tenant_name", "cluster_name"],
11833        )
11834        .await;
11835        eval_stmt.expr = parser::parse(case).unwrap();
11836        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11837            .await
11838            .unwrap();
11839
11840        let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11841            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) +
11842            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11843            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
11844            or
11845            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11846            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
11847            or
11848            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11849            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#;
11850        let table_provider = build_test_table_provider_with_fields(
11851            &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
11852            &["tenant_name", "cluster_name"],
11853        )
11854        .await;
11855        eval_stmt.expr = parser::parse(case).unwrap();
11856        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11857            .await
11858            .unwrap();
11859
11860        let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) +
11861            sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
11862            (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
11863            (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#;
11864        let table_provider = build_test_table_provider_with_fields(
11865            &[
11866                (
11867                    DEFAULT_SCHEMA_NAME.to_string(),
11868                    "background_waitevent_cnt".to_string(),
11869                ),
11870                (
11871                    DEFAULT_SCHEMA_NAME.to_string(),
11872                    "foreground_waitevent_cnt".to_string(),
11873                ),
11874            ],
11875            &["tenant_name", "cluster_name"],
11876        )
11877        .await;
11878        eval_stmt.expr = parser::parse(case).unwrap();
11879        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11880            .await
11881            .unwrap();
11882
11883        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)"#;
11884        let table_provider = build_test_table_provider_with_fields(
11885            &[
11886                (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()),
11887                (
11888                    DEFAULT_SCHEMA_NAME.to_string(),
11889                    "container_cpu_load_average_10s".to_string(),
11890                ),
11891                (
11892                    DEFAULT_SCHEMA_NAME.to_string(),
11893                    "container_spec_cpu_quota".to_string(),
11894                ),
11895            ],
11896            &["cluster_name", "host_name"],
11897        )
11898        .await;
11899        eval_stmt.expr = parser::parse(case).unwrap();
11900        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11901            .await
11902            .unwrap();
11903    }
11904
11905    #[tokio::test]
11906    async fn value_matcher() {
11907        // template
11908        let mut eval_stmt = EvalStmt {
11909            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11910            start: UNIX_EPOCH,
11911            end: UNIX_EPOCH
11912                .checked_add(Duration::from_secs(100_000))
11913                .unwrap(),
11914            interval: Duration::from_secs(5),
11915            lookback_delta: Duration::from_secs(1),
11916        };
11917
11918        let cases = [
11919            // single equal matcher
11920            (
11921                r#"some_metric{__field__="field_1"}"#,
11922                vec![
11923                    "some_metric.field_1",
11924                    "some_metric.tag_0",
11925                    "some_metric.tag_1",
11926                    "some_metric.tag_2",
11927                    "some_metric.timestamp",
11928                ],
11929            ),
11930            // two equal matchers
11931            (
11932                r#"some_metric{__field__="field_1", __field__="field_0"}"#,
11933                vec![
11934                    "some_metric.field_0",
11935                    "some_metric.field_1",
11936                    "some_metric.tag_0",
11937                    "some_metric.tag_1",
11938                    "some_metric.tag_2",
11939                    "some_metric.timestamp",
11940                ],
11941            ),
11942            // single not_eq matcher
11943            (
11944                r#"some_metric{__field__!="field_1"}"#,
11945                vec![
11946                    "some_metric.field_0",
11947                    "some_metric.field_2",
11948                    "some_metric.tag_0",
11949                    "some_metric.tag_1",
11950                    "some_metric.tag_2",
11951                    "some_metric.timestamp",
11952                ],
11953            ),
11954            // two not_eq matchers
11955            (
11956                r#"some_metric{__field__!="field_1", __field__!="field_2"}"#,
11957                vec![
11958                    "some_metric.field_0",
11959                    "some_metric.tag_0",
11960                    "some_metric.tag_1",
11961                    "some_metric.tag_2",
11962                    "some_metric.timestamp",
11963                ],
11964            ),
11965            // equal and not_eq matchers (no conflict)
11966            (
11967                r#"some_metric{__field__="field_1", __field__!="field_0"}"#,
11968                vec![
11969                    "some_metric.field_1",
11970                    "some_metric.tag_0",
11971                    "some_metric.tag_1",
11972                    "some_metric.tag_2",
11973                    "some_metric.timestamp",
11974                ],
11975            ),
11976            // equal and not_eq matchers (conflict)
11977            (
11978                r#"some_metric{__field__="field_2", __field__!="field_2"}"#,
11979                vec![
11980                    "some_metric.tag_0",
11981                    "some_metric.tag_1",
11982                    "some_metric.tag_2",
11983                    "some_metric.timestamp",
11984                ],
11985            ),
11986            // single regex eq matcher
11987            (
11988                r#"some_metric{__field__=~"field_1|field_2"}"#,
11989                vec![
11990                    "some_metric.field_1",
11991                    "some_metric.field_2",
11992                    "some_metric.tag_0",
11993                    "some_metric.tag_1",
11994                    "some_metric.tag_2",
11995                    "some_metric.timestamp",
11996                ],
11997            ),
11998            // single regex not_eq matcher
11999            (
12000                r#"some_metric{__field__!~"field_1|field_2"}"#,
12001                vec![
12002                    "some_metric.field_0",
12003                    "some_metric.tag_0",
12004                    "some_metric.tag_1",
12005                    "some_metric.tag_2",
12006                    "some_metric.timestamp",
12007                ],
12008            ),
12009        ];
12010
12011        for case in cases {
12012            let prom_expr = parser::parse(case.0).unwrap();
12013            eval_stmt.expr = prom_expr;
12014            let table_provider = build_test_table_provider(
12015                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12016                3,
12017                3,
12018            )
12019            .await;
12020            let plan =
12021                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12022                    .await
12023                    .unwrap();
12024            let mut fields = plan.schema().field_names();
12025            let mut expected = case.1.into_iter().map(String::from).collect::<Vec<_>>();
12026            fields.sort();
12027            expected.sort();
12028            assert_eq!(fields, expected, "case: {:?}", case.0);
12029        }
12030
12031        let bad_cases = [
12032            r#"some_metric{__field__="nonexistent"}"#,
12033            r#"some_metric{__field__!="nonexistent"}"#,
12034        ];
12035
12036        for case in bad_cases {
12037            let prom_expr = parser::parse(case).unwrap();
12038            eval_stmt.expr = prom_expr;
12039            let table_provider = build_test_table_provider(
12040                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12041                3,
12042                3,
12043            )
12044            .await;
12045            let plan =
12046                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12047                    .await;
12048            assert!(plan.is_err(), "case: {:?}", case);
12049        }
12050    }
12051
12052    #[tokio::test]
12053    async fn custom_schema() {
12054        let query = "some_alt_metric{__schema__=\"greptime_private\"}";
12055        let expected = String::from(
12056            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12057            \n  PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12058            \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]\
12059            \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]\
12060            \n        TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12061        );
12062
12063        indie_query_plan_compare(query, expected).await;
12064
12065        let query = "some_alt_metric{__database__=\"greptime_private\"}";
12066        let expected = String::from(
12067            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12068            \n  PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12069            \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]\
12070            \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]\
12071            \n        TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12072        );
12073
12074        indie_query_plan_compare(query, expected).await;
12075
12076        let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric";
12077        let expected = String::from(
12078            "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]\
12079            \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]\
12080            \n    SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12081            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12082            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12083            \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]\
12084            \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]\
12085            \n              TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12086            \n    SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12087            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12088            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12089            \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]\
12090            \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]\
12091            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12092        );
12093
12094        indie_query_plan_compare(query, expected).await;
12095    }
12096
12097    #[tokio::test]
12098    async fn only_equals_is_supported_for_special_matcher() {
12099        let queries = &[
12100            "some_alt_metric{__schema__!=\"greptime_private\"}",
12101            "some_alt_metric{__schema__=~\"lalala\"}",
12102            "some_alt_metric{__database__!=\"greptime_private\"}",
12103            "some_alt_metric{__database__=~\"lalala\"}",
12104        ];
12105
12106        for query in queries {
12107            let prom_expr = parser::parse(query).unwrap();
12108            let eval_stmt = EvalStmt {
12109                expr: prom_expr,
12110                start: UNIX_EPOCH,
12111                end: UNIX_EPOCH
12112                    .checked_add(Duration::from_secs(100_000))
12113                    .unwrap(),
12114                interval: Duration::from_secs(5),
12115                lookback_delta: Duration::from_secs(1),
12116            };
12117
12118            let table_provider = build_test_table_provider(
12119                &[
12120                    (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
12121                    (
12122                        "greptime_private".to_string(),
12123                        "some_alt_metric".to_string(),
12124                    ),
12125                ],
12126                1,
12127                1,
12128            )
12129            .await;
12130
12131            let plan =
12132                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12133                    .await;
12134            assert!(plan.is_err(), "query: {:?}", query);
12135        }
12136    }
12137
12138    #[tokio::test]
12139    async fn test_non_ms_precision() {
12140        let catalog_list = MemoryCatalogManager::with_default_setup();
12141        let columns = vec![
12142            ColumnSchema::new(
12143                "tag".to_string(),
12144                ConcreteDataType::string_datatype(),
12145                false,
12146            ),
12147            ColumnSchema::new(
12148                "timestamp".to_string(),
12149                ConcreteDataType::timestamp_nanosecond_datatype(),
12150                false,
12151            )
12152            .with_time_index(true),
12153            ColumnSchema::new(
12154                "field".to_string(),
12155                ConcreteDataType::float64_datatype(),
12156                true,
12157            ),
12158        ];
12159        let schema = Arc::new(Schema::new(columns));
12160        let table_meta = TableMetaBuilder::empty()
12161            .schema(schema)
12162            .primary_key_indices(vec![0])
12163            .value_indices(vec![2])
12164            .next_column_id(1024)
12165            .build()
12166            .unwrap();
12167        let table_info = TableInfoBuilder::default()
12168            .name("metrics".to_string())
12169            .meta(table_meta)
12170            .build()
12171            .unwrap();
12172        let table = EmptyTable::from_table_info(&table_info);
12173        assert!(
12174            catalog_list
12175                .register_table_sync(RegisterTableRequest {
12176                    catalog: DEFAULT_CATALOG_NAME.to_string(),
12177                    schema: DEFAULT_SCHEMA_NAME.to_string(),
12178                    table_name: "metrics".to_string(),
12179                    table_id: 1024,
12180                    table,
12181                })
12182                .is_ok()
12183        );
12184
12185        let plan = PromPlanner::stmt_to_plan(
12186            DfTableSourceProvider::new(
12187                catalog_list.clone(),
12188                false,
12189                QueryContext::arc(),
12190                DummyDecoder::arc(),
12191                true,
12192            ),
12193            &EvalStmt {
12194                expr: parser::parse("metrics{tag = \"1\"}").unwrap(),
12195                start: UNIX_EPOCH,
12196                end: UNIX_EPOCH
12197                    .checked_add(Duration::from_secs(100_000))
12198                    .unwrap(),
12199                interval: Duration::from_secs(5),
12200                lookback_delta: Duration::from_secs(1),
12201            },
12202            &build_query_engine_state(),
12203        )
12204        .await
12205        .unwrap();
12206        assert_eq!(
12207            plan.display_indent_schema().to_string(),
12208            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12209            \n  PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12210            \n    Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12211            \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)]\
12212            \n        Projection: metrics.field, metrics.tag, CAST(metrics.timestamp AS Timestamp(ms)) AS timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12213            \n          TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12214        );
12215        let plan = PromPlanner::stmt_to_plan(
12216            DfTableSourceProvider::new(
12217                catalog_list.clone(),
12218                false,
12219                QueryContext::arc(),
12220                DummyDecoder::arc(),
12221                true,
12222            ),
12223            &EvalStmt {
12224                expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(),
12225                start: UNIX_EPOCH,
12226                end: UNIX_EPOCH
12227                    .checked_add(Duration::from_secs(100_000))
12228                    .unwrap(),
12229                interval: Duration::from_secs(5),
12230                lookback_delta: Duration::from_secs(1),
12231            },
12232            &build_query_engine_state(),
12233        )
12234        .await
12235        .unwrap();
12236        assert_eq!(
12237            plan.display_indent_schema().to_string(),
12238            "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]\
12239            \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]\
12240            \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))]\
12241            \n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12242            \n        PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12243            \n          Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12244            \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)]\
12245            \n              Projection: metrics.field, metrics.tag, CAST(metrics.timestamp AS Timestamp(ms)) AS timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\
12246            \n                TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12247        );
12248    }
12249
12250    #[tokio::test]
12251    async fn test_nonexistent_label() {
12252        // template
12253        let mut eval_stmt = EvalStmt {
12254            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12255            start: UNIX_EPOCH,
12256            end: UNIX_EPOCH
12257                .checked_add(Duration::from_secs(100_000))
12258                .unwrap(),
12259            interval: Duration::from_secs(5),
12260            lookback_delta: Duration::from_secs(1),
12261        };
12262
12263        let case = r#"some_metric{nonexistent="hi"}"#;
12264        let prom_expr = parser::parse(case).unwrap();
12265        eval_stmt.expr = prom_expr;
12266        let table_provider = build_test_table_provider(
12267            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12268            3,
12269            3,
12270        )
12271        .await;
12272        // Should be ok
12273        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12274            .await
12275            .unwrap();
12276    }
12277
12278    #[tokio::test]
12279    async fn test_label_join() {
12280        let prom_expr = parser::parse(
12281            "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')",
12282        )
12283        .unwrap();
12284        let eval_stmt = EvalStmt {
12285            expr: prom_expr,
12286            start: UNIX_EPOCH,
12287            end: UNIX_EPOCH
12288                .checked_add(Duration::from_secs(100_000))
12289                .unwrap(),
12290            interval: Duration::from_secs(5),
12291            lookback_delta: Duration::from_secs(1),
12292        };
12293
12294        let table_provider =
12295            build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1)
12296                .await;
12297        let plan =
12298            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12299                .await
12300                .unwrap();
12301
12302        let expected = r#"
12303Filter: 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]
12304  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]
12305    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]
12306      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]
12307        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]
12308          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]
12309            TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
12310
12311        let ret = plan.display_indent_schema().to_string();
12312        assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
12313    }
12314
12315    #[tokio::test]
12316    async fn test_label_replace() {
12317        let prom_expr = parser::parse(
12318            "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")",
12319        )
12320        .unwrap();
12321        let eval_stmt = EvalStmt {
12322            expr: prom_expr,
12323            start: UNIX_EPOCH,
12324            end: UNIX_EPOCH
12325                .checked_add(Duration::from_secs(100_000))
12326                .unwrap(),
12327            interval: Duration::from_secs(5),
12328            lookback_delta: Duration::from_secs(1),
12329        };
12330
12331        let table_provider =
12332            build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1)
12333                .await;
12334        let plan =
12335            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12336                .await
12337                .unwrap();
12338
12339        let expected = r#"
12340Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8]
12341  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]
12342    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12343      PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12344        Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12345          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]
12346            TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
12347
12348        let ret = plan.display_indent_schema().to_string();
12349        assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
12350    }
12351
12352    #[tokio::test]
12353    async fn label_replace_aggregation_queries_plan_successfully() {
12354        let aggregate =
12355            r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#;
12356        let queries = [
12357            aggregate.to_string(),
12358            format!("{aggregate} <= 10"),
12359            format!("{aggregate} * 0.8"),
12360            format!("0.8 * {aggregate}"),
12361            format!("{aggregate} <= {aggregate} * 0.8"),
12362        ];
12363        let state = build_query_engine_state();
12364        let mut failures = Vec::new();
12365
12366        for query in queries {
12367            let table_provider = build_test_table_provider(
12368                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12369                1,
12370                1,
12371            )
12372            .await;
12373            if let Err(error) =
12374                PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await
12375            {
12376                failures.push(format!("{query}: {error:?}"));
12377            }
12378        }
12379
12380        assert!(failures.is_empty(), "{}", failures.join("\n"));
12381    }
12382
12383    #[tokio::test]
12384    async fn test_matchers_to_expr() {
12385        let mut eval_stmt = EvalStmt {
12386            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12387            start: UNIX_EPOCH,
12388            end: UNIX_EPOCH
12389                .checked_add(Duration::from_secs(100_000))
12390                .unwrap(),
12391            interval: Duration::from_secs(5),
12392            lookback_delta: Duration::from_secs(1),
12393        };
12394        let case =
12395            r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#;
12396
12397        let prom_expr = parser::parse(case).unwrap();
12398        eval_stmt.expr = prom_expr;
12399        let table_provider = build_test_table_provider(
12400            &[(
12401                DEFAULT_SCHEMA_NAME.to_string(),
12402                "prometheus_tsdb_head_series".to_string(),
12403            )],
12404            3,
12405            3,
12406        )
12407        .await;
12408        let plan =
12409            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12410                .await
12411                .unwrap();
12412        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]\
12413        \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]\
12414        \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]\
12415        \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]\
12416        \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]\
12417        \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]\
12418        \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]";
12419        assert_eq!(plan.display_indent_schema().to_string(), expected);
12420    }
12421
12422    #[tokio::test]
12423    async fn test_topk_expr() {
12424        let mut eval_stmt = EvalStmt {
12425            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12426            start: UNIX_EPOCH,
12427            end: UNIX_EPOCH
12428                .checked_add(Duration::from_secs(100_000))
12429                .unwrap(),
12430            interval: Duration::from_secs(5),
12431            lookback_delta: Duration::from_secs(1),
12432        };
12433        let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
12434
12435        let prom_expr = parser::parse(case).unwrap();
12436        eval_stmt.expr = prom_expr;
12437        let table_provider = build_test_table_provider_with_fields(
12438            &[
12439                (
12440                    DEFAULT_SCHEMA_NAME.to_string(),
12441                    "prometheus_tsdb_head_series".to_string(),
12442                ),
12443                (
12444                    DEFAULT_SCHEMA_NAME.to_string(),
12445                    "http_server_requests_seconds_count".to_string(),
12446                ),
12447            ],
12448            &["ip"],
12449        )
12450        .await;
12451
12452        let plan =
12453            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12454                .await
12455                .unwrap();
12456        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)]\
12457        \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]\
12458        \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]\
12459        \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]\
12460        \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]\
12461        \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]\
12462        \n            PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12463        \n              PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12464        \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]\
12465        \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]\
12466        \n                    TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12467
12468        assert_eq!(plan.display_indent_schema().to_string(), expected);
12469    }
12470
12471    #[tokio::test]
12472    async fn test_count_values_expr() {
12473        let mut eval_stmt = EvalStmt {
12474            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12475            start: UNIX_EPOCH,
12476            end: UNIX_EPOCH
12477                .checked_add(Duration::from_secs(100_000))
12478                .unwrap(),
12479            interval: Duration::from_secs(5),
12480            lookback_delta: Duration::from_secs(1),
12481        };
12482        let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
12483
12484        let prom_expr = parser::parse(case).unwrap();
12485        eval_stmt.expr = prom_expr;
12486        let table_provider = build_test_table_provider_with_fields(
12487            &[
12488                (
12489                    DEFAULT_SCHEMA_NAME.to_string(),
12490                    "prometheus_tsdb_head_series".to_string(),
12491                ),
12492                (
12493                    DEFAULT_SCHEMA_NAME.to_string(),
12494                    "http_server_requests_seconds_count".to_string(),
12495                ),
12496            ],
12497            &["ip"],
12498        )
12499        .await;
12500
12501        let plan =
12502            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12503                .await
12504                .unwrap();
12505        let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\
12506        \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 [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\
12507        \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]\
12508        \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12509        \n        PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12510        \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]\
12511        \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]\
12512        \n              TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12513
12514        assert_eq!(plan.display_indent_schema().to_string(), expected);
12515    }
12516
12517    #[tokio::test]
12518    async fn test_value_alias() {
12519        let mut eval_stmt = EvalStmt {
12520            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12521            start: UNIX_EPOCH,
12522            end: UNIX_EPOCH
12523                .checked_add(Duration::from_secs(100_000))
12524                .unwrap(),
12525            interval: Duration::from_secs(5),
12526            lookback_delta: Duration::from_secs(1),
12527        };
12528        let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
12529
12530        let prom_expr = parser::parse(case).unwrap();
12531        eval_stmt.expr = prom_expr;
12532        eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series");
12533        let table_provider = build_test_table_provider_with_fields(
12534            &[
12535                (
12536                    DEFAULT_SCHEMA_NAME.to_string(),
12537                    "prometheus_tsdb_head_series".to_string(),
12538                ),
12539                (
12540                    DEFAULT_SCHEMA_NAME.to_string(),
12541                    "http_server_requests_seconds_count".to_string(),
12542                ),
12543            ],
12544            &["ip"],
12545        )
12546        .await;
12547
12548        let plan =
12549            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12550                .await
12551                .unwrap();
12552        let expected = r#"
12553Projection: 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)]
12554  Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]
12555    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 [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]
12556      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]
12557        PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
12558          PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
12559            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]
12560              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]
12561                TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#;
12562        assert_eq!(format!("\n{}", plan.display_indent_schema()), expected);
12563    }
12564
12565    #[tokio::test]
12566    async fn test_quantile_expr() {
12567        let mut eval_stmt = EvalStmt {
12568            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12569            start: UNIX_EPOCH,
12570            end: UNIX_EPOCH
12571                .checked_add(Duration::from_secs(100_000))
12572                .unwrap(),
12573            interval: Duration::from_secs(5),
12574            lookback_delta: Duration::from_secs(1),
12575        };
12576        let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
12577
12578        let prom_expr = parser::parse(case).unwrap();
12579        eval_stmt.expr = prom_expr;
12580        let table_provider = build_test_table_provider_with_fields(
12581            &[
12582                (
12583                    DEFAULT_SCHEMA_NAME.to_string(),
12584                    "prometheus_tsdb_head_series".to_string(),
12585                ),
12586                (
12587                    DEFAULT_SCHEMA_NAME.to_string(),
12588                    "http_server_requests_seconds_count".to_string(),
12589                ),
12590            ],
12591            &["ip"],
12592        )
12593        .await;
12594
12595        let plan =
12596            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12597                .await
12598                .unwrap();
12599        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]\
12600        \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]\
12601        \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]\
12602        \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]\
12603        \n        PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12604        \n          PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12605        \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]\
12606        \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]\
12607        \n                TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12608
12609        assert_eq!(plan.display_indent_schema().to_string(), expected);
12610    }
12611
12612    #[tokio::test]
12613    async fn test_or_not_exists_table_label() {
12614        let state = build_query_engine_state();
12615        let provider = build_test_table_provider_with_fields(
12616            &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())],
12617            &["job"],
12618        )
12619        .await;
12620        let raw = PromPlanner::stmt_to_plan(
12621            provider,
12622            &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#),
12623            &state,
12624        )
12625        .await
12626        .unwrap();
12627        assert!(
12628            raw.display_indent_schema()
12629                .to_string()
12630                .contains("__promql_or_match_0@")
12631        );
12632        let (optimized, batches) = execute(raw, &state).await;
12633        assert_no_internal_or_keys(optimized.schema());
12634        assert!(batches.iter().all(|batch| {
12635            batch
12636                .schema()
12637                .fields()
12638                .iter()
12639                .all(|field| !field.name().starts_with("__promql_or_match_"))
12640        }));
12641    }
12642
12643    #[tokio::test]
12644    async fn test_histogram_quantile_missing_le_column() {
12645        let mut eval_stmt = EvalStmt {
12646            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12647            start: UNIX_EPOCH,
12648            end: UNIX_EPOCH
12649                .checked_add(Duration::from_secs(100_000))
12650                .unwrap(),
12651            interval: Duration::from_secs(5),
12652            lookback_delta: Duration::from_secs(1),
12653        };
12654
12655        // Test case: histogram_quantile with a table that doesn't have 'le' column
12656        let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#;
12657
12658        let prom_expr = parser::parse(case).unwrap();
12659        eval_stmt.expr = prom_expr;
12660
12661        // Create a table provider with a table that doesn't have 'le' column
12662        let table_provider = build_test_table_provider_with_fields(
12663            &[(
12664                DEFAULT_SCHEMA_NAME.to_string(),
12665                "non_existent_histogram_bucket".to_string(),
12666            )],
12667            &["pod", "instance"], // Note: no 'le' column
12668        )
12669        .await;
12670
12671        // Should return empty result instead of error
12672        let result =
12673            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12674                .await;
12675
12676        // This should succeed now (returning empty result) instead of failing with "Cannot find column le"
12677        assert!(
12678            result.is_ok(),
12679            "Expected successful plan creation with empty result, but got error: {:?}",
12680            result.err()
12681        );
12682
12683        // Verify that the result is an EmptyRelation
12684        let plan = result.unwrap();
12685        match plan {
12686            LogicalPlan::EmptyRelation(_) => {
12687                // This is what we expect
12688            }
12689            _ => panic!("Expected EmptyRelation, but got: {:?}", plan),
12690        }
12691    }
12692
12693    #[tokio::test]
12694    async fn test_direct_or_normalizes_missing_match_labels() {
12695        type Case<'a> = (
12696            Option<Option<&'a str>>,
12697            Option<Option<&'a str>>,
12698            i64,
12699            i64,
12700            &'a [(f64, Option<&'a str>)],
12701        );
12702
12703        let modifier = or_modifier("lhs or on(k) rhs");
12704        #[rustfmt::skip]
12705        let cases: &[Case<'_>] = &[
12706            (None, None, 1, 1, &[(1.0, None)]),
12707            (None, Some(Some("")), 1, 1, &[(1.0, None)]),
12708            (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]),
12709            (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
12710            (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]),
12711            (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]),
12712            (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
12713            (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]),
12714        ];
12715        for &(left, right, left_ts, right_ts, expected) in cases {
12716            let (optimized, batches) = run(
12717                &matrix_source("lhs", left, left_ts, 1.0),
12718                &matrix_source("rhs", right, right_ts, 2.0),
12719                matrix_context("lhs", left),
12720                matrix_context("rhs", right),
12721                &modifier,
12722            )
12723            .await;
12724            assert_no_internal_or_keys(optimized.schema());
12725            assert_eq!(
12726                rows(&batches),
12727                expected
12728                    .iter()
12729                    .map(|(value, label)| (*value, label.map(str::to_string)))
12730                    .collect::<Vec<_>>()
12731            );
12732        }
12733    }
12734
12735    #[tokio::test]
12736    async fn test_direct_or_match_modifiers() {
12737        for (modifier, left, right, expected) in [
12738            (None, "left", "right", 2),
12739            (or_modifier("lhs or on(k) rhs"), "same", "same", 1),
12740            (or_modifier("lhs or on() rhs"), "left", "right", 1),
12741            (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1),
12742        ] {
12743            let (_, batches) = run(
12744                &matrix_source("lhs", Some(Some(left)), 1, 1.0),
12745                &matrix_source("rhs", Some(Some(right)), 1, 2.0),
12746                direct_or_context("lhs", &["job", "k"], "v"),
12747                direct_or_context("rhs", &["job", "k"], "v"),
12748                &modifier,
12749            )
12750            .await;
12751            assert_eq!(
12752                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
12753                expected
12754            );
12755        }
12756    }
12757
12758    #[tokio::test]
12759    async fn test_direct_or_nested_projection_uses_left_context() {
12760        let left = matrix_source("lhs", Some(Some("k")), 1, 1.0);
12761        let right = matrix_source("rhs", Some(Some("k")), 1, 2.0);
12762        let raw = plan_direct_or(
12763            scan(&left),
12764            scan(&right),
12765            direct_or_context("lhs", &["job", "k"], "v"),
12766            direct_or_context("rhs", &["job", "k"], "v"),
12767            &or_modifier("lhs or on(k) rhs"),
12768        )
12769        .await;
12770        assert!(raw.schema().iter().any(|(qualifier, field)| {
12771            qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v"
12772        }));
12773        let nested = LogicalPlanBuilder::from(raw)
12774            .project(vec![
12775                DfExpr::BinaryExpr(BinaryExpr {
12776                    left: Box::new(DfExpr::Column(Column::new(
12777                        Some(TableReference::bare("lhs")),
12778                        "v",
12779                    ))),
12780                    op: Operator::Plus,
12781                    right: Box::new(lit(1.0)),
12782                })
12783                .alias("v_plus"),
12784            ])
12785            .unwrap()
12786            .build()
12787            .unwrap();
12788        let (_, batches) = execute(nested, &build_query_engine_state()).await;
12789        assert_eq!(values(&batches, "v_plus"), vec![2.0]);
12790    }
12791
12792    #[tokio::test]
12793    async fn test_direct_or_skips_user_internal_key_name() {
12794        const USER_TAG: &str = "__promql_or_match_0";
12795        let left = tagged_source(
12796            "lhs",
12797            false,
12798            (USER_TAG, Some("left")),
12799            DirectOrValue::Float64(1.0),
12800        );
12801        let right = tagged_source(
12802            "rhs",
12803            false,
12804            (USER_TAG, Some("right")),
12805            DirectOrValue::Float64(2.0),
12806        );
12807        let raw = plan_direct_or(
12808            scan(&left),
12809            scan(&right),
12810            direct_or_context("lhs", &["job", USER_TAG], "v"),
12811            direct_or_context("rhs", &["job", USER_TAG], "v"),
12812            &or_modifier("lhs or on(missing_label) rhs"),
12813        )
12814        .await;
12815        assert!(
12816            raw.display_indent_schema()
12817                .to_string()
12818                .contains("__promql_or_match_1@")
12819        );
12820        let (_, batches) = execute(raw, &build_query_engine_state()).await;
12821        assert!(
12822            batches
12823                .iter()
12824                .all(|batch| batch.column_by_name(USER_TAG).is_some())
12825        );
12826    }
12827
12828    #[tokio::test]
12829    async fn test_direct_or_substrait_round_trip_with_normalized_key() {
12830        let state = build_query_engine_state();
12831        let ctx = SessionContext::new_with_state(state.session_state());
12832        let catalog = Arc::new(MemoryCatalogProvider::new());
12833        catalog
12834            .register_schema("public", Arc::new(MemorySchemaProvider::new()))
12835            .unwrap();
12836        ctx.register_catalog("datafusion", catalog);
12837        let left = matrix_source("lhs", Some(Some("")), 1, 1.0);
12838        let right = matrix_source("rhs", None, 1, 2.0);
12839        ctx.register_table(
12840            TableReference::full("datafusion", "public", "lhs"),
12841            table(&left),
12842        )
12843        .unwrap();
12844        ctx.register_table(
12845            TableReference::full("datafusion", "public", "rhs"),
12846            table(&right),
12847        )
12848        .unwrap();
12849        let raw = plan_direct_or(
12850            ctx.table("datafusion.public.lhs")
12851                .await
12852                .unwrap()
12853                .into_unoptimized_plan(),
12854            ctx.table("datafusion.public.rhs")
12855                .await
12856                .unwrap()
12857                .into_unoptimized_plan(),
12858            direct_or_context("lhs", &["job", "k"], "v"),
12859            direct_or_context("rhs", &["job"], "v"),
12860            &or_modifier("lhs or on(k) rhs"),
12861        )
12862        .await;
12863        let decoded = DFLogicalSubstraitConvertor
12864            .decode(
12865                DFLogicalSubstraitConvertor
12866                    .encode(&raw, DefaultSerializer)
12867                    .unwrap(),
12868                ctx.state(),
12869            )
12870            .await
12871            .unwrap();
12872        let (optimized, batches) = execute(decoded, &state).await;
12873        assert_no_internal_or_keys(optimized.schema());
12874        assert!(batches.iter().all(|batch| {
12875            batch
12876                .schema()
12877                .fields()
12878                .iter()
12879                .all(|field| !field.name().starts_with("__promql_or_match_"))
12880        }));
12881        assert_eq!(values(&batches, "v"), vec![1.0]);
12882    }
12883
12884    #[tokio::test]
12885    async fn test_direct_or_numeric_value_types() {
12886        let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0));
12887        let right = tagged_source(
12888            "rhs",
12889            false,
12890            ("k", Some("rhs")),
12891            DirectOrValue::Float64(0.5),
12892        );
12893        let (optimized, batches) = run(
12894            &left,
12895            &right,
12896            direct_or_context("lhs", &["job", "k"], "v"),
12897            direct_or_context("rhs", &["job", "k"], "v"),
12898            &or_modifier("lhs or on(k) rhs"),
12899        )
12900        .await;
12901        assert_eq!(
12902            optimized
12903                .schema()
12904                .field_with_name(None, "v")
12905                .unwrap()
12906                .data_type(),
12907            &ArrowDataType::Float64
12908        );
12909        assert_eq!(values(&batches, "v"), vec![0.5]);
12910        let provider = build_test_table_provider_with_fields(
12911            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
12912            &[],
12913        )
12914        .await;
12915        let mut planner = PromPlanner {
12916            table_provider: provider,
12917            ctx: PromPlannerContext::default(),
12918            promql_annotations: None,
12919        };
12920        let left_context = direct_or_context("lhs", &["job"], "v");
12921        let right_context = direct_or_context("rhs", &["job"], "v");
12922        let error = planner
12923            .or_operator(
12924                scan(&job_source("lhs", DirectOrValue::Utf8("x"))),
12925                scan(&job_source("rhs", DirectOrValue::Float64(1.0))),
12926                left_context.tag_columns.iter().cloned().collect(),
12927                right_context.tag_columns.iter().cloned().collect(),
12928                left_context,
12929                right_context,
12930                &or_modifier("lhs or on() rhs"),
12931            )
12932            .unwrap_err();
12933        assert!(
12934            error
12935                .to_string()
12936                .contains("OR value fields have incompatible types")
12937        );
12938    }
12939
12940    #[tokio::test]
12941    async fn test_or_with_histogram_quantile_missing_le_column() {
12942        let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#;
12943        let eval_stmt = build_eval_stmt(case);
12944        let table_provider = build_missing_le_or_normal_metric_table_provider().await;
12945
12946        let plan =
12947            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12948                .await
12949                .unwrap();
12950        assert_normal_metric_schema(&plan);
12951    }
12952
12953    #[tokio::test]
12954    async fn test_or_with_right_empty_histogram_restores_left_context() {
12955        let eval_stmt = build_eval_stmt(
12956            r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#,
12957        );
12958        let table_provider = build_missing_le_or_normal_metric_table_provider().await;
12959
12960        PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12961            .await
12962            .unwrap();
12963    }
12964
12965    #[tokio::test]
12966    async fn test_or_with_both_empty_histograms() {
12967        let eval_stmt = build_eval_stmt(
12968            r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#,
12969        );
12970        let table_provider = build_test_table_provider_with_fields(
12971            &[
12972                (
12973                    DEFAULT_SCHEMA_NAME.to_string(),
12974                    "left_histogram_bucket".to_string(),
12975                ),
12976                (
12977                    DEFAULT_SCHEMA_NAME.to_string(),
12978                    "right_histogram_bucket".to_string(),
12979                ),
12980            ],
12981            &["pod", "instance"],
12982        )
12983        .await;
12984
12985        let plan =
12986            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12987                .await
12988                .unwrap();
12989        match plan {
12990            LogicalPlan::EmptyRelation(relation) => {
12991                assert!(!relation.produce_one_row);
12992                assert!(!relation.schema.fields().is_empty());
12993                assert!(
12994                    relation
12995                        .schema
12996                        .fields()
12997                        .iter()
12998                        .any(|field| field.data_type() == &ArrowDataType::Float64)
12999                );
13000                assert!(
13001                    relation
13002                        .schema
13003                        .fields()
13004                        .iter()
13005                        .any(|field| field.name() == "pod")
13006                );
13007                assert!(
13008                    !relation
13009                        .schema
13010                        .fields()
13011                        .iter()
13012                        .any(|field| field.name() == "instance")
13013                );
13014            }
13015            _ => panic!("Expected EmptyRelation, but got: {plan:?}"),
13016        }
13017    }
13018
13019    #[tokio::test]
13020    async fn test_nested_or_with_both_empty_histograms() {
13021        for case in [
13022            r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#,
13023            r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#,
13024        ] {
13025            let eval_stmt = build_eval_stmt(case);
13026            let table_provider = build_test_table_provider_with_fields(
13027                &[
13028                    (
13029                        DEFAULT_SCHEMA_NAME.to_string(),
13030                        "left_histogram_bucket".to_string(),
13031                    ),
13032                    (
13033                        DEFAULT_SCHEMA_NAME.to_string(),
13034                        "right_histogram_bucket".to_string(),
13035                    ),
13036                ],
13037                &["pod", "instance"],
13038            )
13039            .await;
13040
13041            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13042                .await
13043                .unwrap();
13044        }
13045    }
13046
13047    #[tokio::test]
13048    async fn test_or_with_empty_histogram_modifiers() {
13049        for case in [
13050            r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#,
13051            r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#,
13052        ] {
13053            let eval_stmt = build_eval_stmt(case);
13054            let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13055
13056            let plan =
13057                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13058                    .await
13059                    .unwrap();
13060            assert_normal_metric_schema(&plan);
13061        }
13062    }
13063
13064    #[tokio::test]
13065    async fn test_unless_preserves_left_context_for_histogram() {
13066        let eval_stmt = build_eval_stmt(
13067            r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#,
13068        );
13069        let state = build_query_engine_state();
13070        let plan = PromPlanner::stmt_to_plan(
13071            build_set_op_context_table_provider().await,
13072            &eval_stmt,
13073            &state,
13074        )
13075        .await
13076        .unwrap();
13077        assert!(contains_histogram_fold(&plan), "{plan:?}");
13078        let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await;
13079        assert!(contains_histogram_fold(&optimized), "{optimized:?}");
13080        let batches =
13081            datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
13082                .await
13083                .unwrap();
13084        assert!(batches.iter().all(|batch| batch.num_rows() == 0));
13085    }
13086
13087    #[tokio::test]
13088    async fn test_and_preserves_left_context_for_histogram() {
13089        let eval_stmt = build_eval_stmt(
13090            r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#,
13091        );
13092        let plan = PromPlanner::stmt_to_plan(
13093            build_set_op_context_table_provider().await,
13094            &eval_stmt,
13095            &build_query_engine_state(),
13096        )
13097        .await
13098        .unwrap();
13099        assert!(contains_histogram_fold(&plan), "{plan:?}");
13100    }
13101
13102    #[tokio::test]
13103    async fn test_and_preserves_left_context_when_le_is_missing() {
13104        let eval_stmt =
13105            build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#);
13106        let plan = PromPlanner::stmt_to_plan(
13107            build_set_op_context_table_provider().await,
13108            &eval_stmt,
13109            &build_query_engine_state(),
13110        )
13111        .await
13112        .unwrap();
13113        assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}");
13114        assert!(!plan.schema().fields().is_empty());
13115        assert!(!contains_histogram_fold(&plan), "{plan:?}");
13116    }
13117
13118    #[tokio::test]
13119    async fn test_or_context_uses_left_qualified_output() {
13120        let case = r#"(normal_metric or other_metric) + 1"#;
13121        let eval_stmt = build_eval_stmt(case);
13122        let state = build_query_engine_state();
13123        let plan =
13124            PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state)
13125                .await
13126                .unwrap();
13127        assert!(
13128            plan.schema()
13129                .fields()
13130                .iter()
13131                .any(|field| field.data_type() == &ArrowDataType::Float64),
13132            "{plan:?}"
13133        );
13134        let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await;
13135    }
13136
13137    #[tokio::test]
13138    async fn test_or_context_uses_left_qualified_empty_histogram_output() {
13139        let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#;
13140        let eval_stmt = build_eval_stmt(case);
13141        let plan = PromPlanner::stmt_to_plan(
13142            build_or_context_table_provider().await,
13143            &eval_stmt,
13144            &build_query_engine_state(),
13145        )
13146        .await
13147        .unwrap();
13148        assert!(
13149            plan.schema()
13150                .fields()
13151                .iter()
13152                .any(|field| field.data_type() == &ArrowDataType::Float64),
13153            "{plan:?}"
13154        );
13155    }
13156
13157    #[tokio::test]
13158    async fn test_direct_or_preserves_float_and_native_histogram_samples() {
13159        for histogram_on_left in [false, true] {
13160            let (planner, plan) = mixed_direct_or(histogram_on_left).await;
13161
13162            let float_field = &planner.ctx.field_columns[0];
13163            let histogram_field = &planner.ctx.field_columns[1];
13164            assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX));
13165            assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX));
13166            assert_eq!(
13167                plan.schema()
13168                    .field_with_name(None, float_field)
13169                    .unwrap()
13170                    .data_type(),
13171                &ArrowDataType::Float64
13172            );
13173            assert_eq!(
13174                plan.schema()
13175                    .field_with_name(None, histogram_field)
13176                    .unwrap()
13177                    .data_type(),
13178                &native_histogram_value_type().as_arrow_type()
13179            );
13180
13181            let (optimized, batches) = execute(plan, &build_query_engine_state()).await;
13182            assert_no_internal_or_keys(optimized.schema());
13183            let mut sample_kinds = batches
13184                .iter()
13185                .flat_map(|batch| {
13186                    let values = batch.column_by_name(float_field).unwrap();
13187                    let histograms = batch.column_by_name(histogram_field).unwrap();
13188                    (0..batch.num_rows())
13189                        .map(|row| (values.is_valid(row), histograms.is_valid(row)))
13190                })
13191                .collect::<Vec<_>>();
13192            sample_kinds.sort_unstable();
13193            assert_eq!(sample_kinds, vec![(false, true), (true, false)]);
13194        }
13195    }
13196
13197    #[tokio::test]
13198    async fn malformed_classic_bucket_does_not_drop_native_histogram() {
13199        let state = build_query_engine_state();
13200        let collector = PromqlAnnotationCollector::default();
13201        let plan = PromPlanner::stmt_to_plan_with_annotations(
13202            operator_table_provider(),
13203            &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"),
13204            &state,
13205            Some(collector.clone()),
13206        )
13207        .await
13208        .unwrap();
13209        let value_field = plan
13210            .schema()
13211            .fields()
13212            .iter()
13213            .find(|field| field.data_type() == &ArrowDataType::Float64)
13214            .unwrap()
13215            .name()
13216            .clone();
13217
13218        let (_, batches) = execute(plan, &state).await;
13219        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13220        assert_eq!(values(&batches, &value_field), vec![0.0]);
13221        let mut warnings = vec![];
13222        let mut infos = vec![];
13223        collector.append_to(&mut warnings, &mut infos);
13224        assert!(warnings.is_empty());
13225        assert!(infos.is_empty());
13226    }
13227
13228    #[tokio::test]
13229    async fn test_mixed_binary_operator_aligns_both_alternative_inputs() {
13230        let state = build_query_engine_state();
13231        let plan = PromPlanner::stmt_to_plan(
13232            operator_table_provider(),
13233            &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"),
13234            &state,
13235        )
13236        .await
13237        .unwrap();
13238        let plan_text = plan.display_indent_schema().to_string();
13239        assert!(
13240            plan_text.contains("prom_native_histogram_mul_scalar"),
13241            "{plan_text}"
13242        );
13243        assert!(
13244            plan_text.contains("prom_native_histogram_scalar_mul"),
13245            "{plan_text}"
13246        );
13247        let float_field = plan
13248            .schema()
13249            .fields()
13250            .iter()
13251            .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
13252            .unwrap()
13253            .name()
13254            .clone();
13255        let histogram_field = plan
13256            .schema()
13257            .fields()
13258            .iter()
13259            .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13260            .unwrap()
13261            .name()
13262            .clone();
13263
13264        let (_, batches) = execute(plan, &state).await;
13265        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
13266        assert!(values(&batches, &float_field).is_empty());
13267        let mut sums = histograms(&batches, &histogram_field)
13268            .into_iter()
13269            .map(|histogram| histogram.sum)
13270            .collect::<Vec<_>>();
13271        sums.sort_by(f64::total_cmp);
13272        assert_eq!(sums, vec![2.0, 3.0]);
13273    }
13274
13275    #[tokio::test]
13276    async fn test_mixed_binary_operator_reports_only_dropped_samples() {
13277        for (query, expected_rows, expected_infos) in [
13278            ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1),
13279            ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0),
13280            ("(lf or on(tag) lh) % on(tag) lh", 0, 1),
13281        ] {
13282            let state = build_query_engine_state();
13283            let annotations = PromqlAnnotationCollector::default();
13284            let plan = PromPlanner::stmt_to_plan_with_annotations(
13285                operator_table_provider(),
13286                &operator_eval_stmt(query),
13287                &state,
13288                Some(annotations.clone()),
13289            )
13290            .await
13291            .unwrap();
13292
13293            let (_, batches) = execute(plan, &state).await;
13294            assert_eq!(
13295                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
13296                expected_rows,
13297                "{query}"
13298            );
13299            let mut warnings = vec![];
13300            let mut infos = vec![];
13301            annotations.append_to(&mut warnings, &mut infos);
13302            assert!(warnings.is_empty(), "{query}: {warnings:?}");
13303            assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}");
13304        }
13305    }
13306
13307    #[tokio::test]
13308    async fn test_histogram_only_min_drops_empty_aggregate_group() {
13309        // `min` over native-histogram-only input drops every sample in the group, so the
13310        // NULL-valued aggregate row must be filtered out. Otherwise an outer expression
13311        // like `group()` resurrects the group Prometheus considers unseen.
13312        let state = build_query_engine_state();
13313        for query in ["min(lh)", "group(min(lh))"] {
13314            let plan = PromPlanner::stmt_to_plan(
13315                operator_table_provider(),
13316                &operator_eval_stmt(query),
13317                &state,
13318            )
13319            .await
13320            .unwrap();
13321            let (_, batches) = execute(plan, &state).await;
13322            assert_eq!(
13323                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
13324                0,
13325                "{query}"
13326            );
13327        }
13328    }
13329
13330    #[tokio::test]
13331    async fn test_mixed_min_drops_histogram_only_group() {
13332        // With alternative float/histogram fields, `min by (tag)` keeps float-only groups
13333        // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of
13334        // emitting a NULL-valued row for them.
13335        let state = build_query_engine_state();
13336        let plan = PromPlanner::stmt_to_plan(
13337            operator_table_provider(),
13338            &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"),
13339            &state,
13340        )
13341        .await
13342        .unwrap();
13343        let float_field = plan
13344            .schema()
13345            .fields()
13346            .iter()
13347            .find(|field| field.data_type() == &ArrowDataType::Float64)
13348            .unwrap()
13349            .name()
13350            .clone();
13351        let (_, batches) = execute(plan, &state).await;
13352        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13353        assert_eq!(values(&batches, &float_field), vec![2.0]);
13354    }
13355
13356    #[tokio::test]
13357    async fn test_mixed_or_can_feed_another_or() {
13358        let state = build_query_engine_state();
13359        let plan = PromPlanner::stmt_to_plan(
13360            operator_table_provider(),
13361            &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"),
13362            &state,
13363        )
13364        .await
13365        .unwrap();
13366        let float_field = plan
13367            .schema()
13368            .fields()
13369            .iter()
13370            .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
13371            .unwrap()
13372            .name()
13373            .clone();
13374        let histogram_field = plan
13375            .schema()
13376            .fields()
13377            .iter()
13378            .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13379            .unwrap()
13380            .name()
13381            .clone();
13382
13383        let (_, batches) = execute(plan, &state).await;
13384        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
13385        let mut float_values = values(&batches, &float_field);
13386        float_values.sort_by(f64::total_cmp);
13387        assert_eq!(float_values, vec![2.0, 7.0]);
13388        assert_eq!(histograms(&batches, &histogram_field).len(), 1);
13389    }
13390
13391    #[tokio::test]
13392    async fn test_mixed_fields_align_with_single_float_vector() {
13393        let (planner, mixed) = mixed_direct_or(false).await;
13394        let scale = tagged_source(
13395            "scale",
13396            false,
13397            ("k", Some("float")),
13398            DirectOrValue::Float64(2.0),
13399        );
13400        let scale = scan(&scale);
13401        let scale_fields = vec!["v".to_string()];
13402        let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else {
13403            unreachable!()
13404        };
13405
13406        let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
13407            mixed.schema(),
13408            scale.schema(),
13409            &planner.ctx.field_columns,
13410            &scale_fields,
13411            binary.op,
13412            false,
13413            false,
13414        );
13415        assert!(invalid_pairs.is_empty());
13416        assert_eq!(
13417            groups
13418                .iter()
13419                .map(|(output, _)| output.clone())
13420                .collect::<Vec<_>>(),
13421            planner.ctx.field_columns
13422        );
13423        assert_eq!(groups.len(), 2);
13424        assert!(
13425            groups
13426                .iter()
13427                .flat_map(|(_, pairs)| pairs)
13428                .all(|(_, right)| *right == &scale_fields[0])
13429        );
13430
13431        let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
13432            scale.schema(),
13433            mixed.schema(),
13434            &scale_fields,
13435            &planner.ctx.field_columns,
13436            binary.op,
13437            false,
13438            false,
13439        );
13440        assert!(invalid_pairs.is_empty());
13441        assert_eq!(
13442            groups
13443                .iter()
13444                .map(|(output, _)| output.clone())
13445                .collect::<Vec<_>>(),
13446            planner.ctx.field_columns
13447        );
13448        assert_eq!(groups.len(), 2);
13449        assert!(
13450            groups
13451                .iter()
13452                .flat_map(|(_, pairs)| pairs)
13453                .all(|(left, _)| *left == &scale_fields[0])
13454        );
13455    }
13456
13457    #[tokio::test]
13458    async fn test_non_bool_comparison_filters_mixed_sample_lanes() {
13459        let (planner, input) = mixed_direct_or(false).await;
13460        let input_schema = input.schema().clone();
13461        let plan = planner
13462            .filter_on_field_column(input, |field| {
13463                if PromPlanner::field_column_is_native_histogram(&input_schema, field) {
13464                    Ok(lit(false))
13465                } else {
13466                    Ok(col(field).gt(lit(0.0)))
13467                }
13468            })
13469            .unwrap();
13470        let float_field = planner.ctx.field_columns[0].clone();
13471
13472        let (_, batches) = execute(plan, &build_query_engine_state()).await;
13473        assert_eq!(values(&batches, &float_field), vec![1.25]);
13474        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13475    }
13476
13477    #[tokio::test]
13478    async fn test_mixed_left_and_unless_preserve_sample_lanes() {
13479        for (expression, expected_sample_kind) in [
13480            ("lhs and on(k) mask", (false, true)),
13481            ("lhs unless on(k) mask", (true, false)),
13482        ] {
13483            let (mut planner, left) = mixed_direct_or(false).await;
13484            let left_context = planner.ctx.clone();
13485            let float_field = left_context.field_columns[0].clone();
13486            let histogram_field = left_context.field_columns[1].clone();
13487            let mask = tagged_source(
13488                "mask",
13489                false,
13490                ("k", Some("histogram")),
13491                DirectOrValue::Float64(1.0),
13492            );
13493            let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else {
13494                unreachable!()
13495            };
13496            let plan = planner
13497                .set_op_on_non_field_columns(
13498                    left,
13499                    scan(&mask),
13500                    left_context,
13501                    direct_or_context("mask", &["job", "k"], "v"),
13502                    binary.op,
13503                    &binary.modifier,
13504                )
13505                .unwrap();
13506
13507            let (_, batches) = execute(plan, &build_query_engine_state()).await;
13508            let sample_kinds = batches
13509                .iter()
13510                .flat_map(|batch| {
13511                    let floats = batch.column_by_name(&float_field).unwrap();
13512                    let histograms = batch.column_by_name(&histogram_field).unwrap();
13513                    (0..batch.num_rows())
13514                        .map(|row| (floats.is_valid(row), histograms.is_valid(row)))
13515                })
13516                .collect::<Vec<_>>();
13517            assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}");
13518        }
13519    }
13520
13521    #[tokio::test]
13522    async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() {
13523        let plan = PromPlanner::stmt_to_plan(
13524            build_test_mixed_native_histogram_table_provider("some_metric").await,
13525            &build_eval_stmt("some_metric * scalar(vector(2))"),
13526            &build_query_engine_state(),
13527        )
13528        .await
13529        .unwrap();
13530        let schema = plan.schema();
13531        assert_eq!(
13532            schema
13533                .field_with_unqualified_name(greptime_value())
13534                .unwrap()
13535                .data_type(),
13536            &ArrowDataType::Float64
13537        );
13538        assert_eq!(
13539            schema
13540                .field_with_unqualified_name(greptime_native_histogram())
13541                .unwrap()
13542                .data_type(),
13543            &native_histogram_value_type().as_arrow_type()
13544        );
13545        assert!(
13546            plan.display_indent_schema()
13547                .to_string()
13548                .contains("prom_native_histogram_mul_scalar"),
13549            "{plan:?}"
13550        );
13551    }
13552
13553    #[tokio::test]
13554    async fn test_unsupported_histogram_binary_does_not_block_or_fallback() {
13555        let state = build_query_engine_state();
13556        let plan = PromPlanner::stmt_to_plan(
13557            operator_table_provider(),
13558            &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"),
13559            &state,
13560        )
13561        .await
13562        .unwrap();
13563        let float_field = plan
13564            .schema()
13565            .fields()
13566            .iter()
13567            .find(|field| field.data_type() == &ArrowDataType::Float64)
13568            .unwrap()
13569            .name()
13570            .clone();
13571        let histogram_field = plan
13572            .schema()
13573            .fields()
13574            .iter()
13575            .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type())
13576            .unwrap()
13577            .name()
13578            .clone();
13579
13580        let (_, batches) = execute(plan, &state).await;
13581        assert_eq!(values(&batches, &float_field), vec![0.0]);
13582        assert_eq!(histograms(&batches, &histogram_field).len(), 1);
13583    }
13584
13585    #[tokio::test]
13586    async fn test_unary_negates_mixed_float_and_native_histogram_samples() {
13587        for histogram_on_left in [false, true] {
13588            let (mut planner, input) = mixed_direct_or(histogram_on_left).await;
13589            let plan = planner.negate_field_columns(input).unwrap();
13590            assert!(PromPlanner::field_columns_are_alternative_samples(
13591                plan.schema(),
13592                &planner.ctx.field_columns
13593            ));
13594            let float_field = planner
13595                .ctx
13596                .field_columns
13597                .iter()
13598                .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX))
13599                .unwrap();
13600            let histogram_field = planner
13601                .ctx
13602                .field_columns
13603                .iter()
13604                .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13605                .unwrap();
13606
13607            let (_, batches) = execute(plan, &build_query_engine_state()).await;
13608            assert_eq!(values(&batches, float_field), vec![-1.25]);
13609            let histogram = batches
13610                .iter()
13611                .find_map(|batch| {
13612                    let values = batch
13613                        .column_by_name(histogram_field)
13614                        .unwrap()
13615                        .as_any()
13616                        .downcast_ref::<datafusion::arrow::array::StructArray>()
13617                        .unwrap();
13618                    (0..values.len()).find_map(|row| {
13619                        common_query::native_histogram::read_histogram(values, row).unwrap()
13620                    })
13621                })
13622                .unwrap();
13623            assert_eq!(histogram.count, -1.0);
13624            assert_eq!(histogram.sum, -1.0);
13625            assert_eq!(histogram.reset_hint, CounterResetHint::Gauge);
13626        }
13627    }
13628
13629    #[tokio::test]
13630    async fn test_native_histogram_sum_and_avg_execute_real_batches() {
13631        for op_name in ["sum", "avg"] {
13632            for incompatible in [false, true] {
13633                let mut second = direct_or_histogram();
13634                if incompatible {
13635                    second.schema = CUSTOM_BUCKETS_SCHEMA;
13636                    second.custom_values = vec![1.0];
13637                }
13638                let collector = PromqlAnnotationCollector::default();
13639                let (mut planner, input) =
13640                    mixed_aggregate_input(vec![direct_or_histogram(), second]).await;
13641                planner.promql_annotations = Some(collector.clone());
13642                let histogram_column = planner.ctx.field_columns[1].clone();
13643                planner.ctx.field_columns = vec![histogram_column.clone()];
13644                let input = LogicalPlanBuilder::from(input)
13645                    .project([col("ts"), col(&histogram_column)])
13646                    .unwrap()
13647                    .build()
13648                    .unwrap();
13649                let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13650                    parser::parse(&format!("{op_name}(mixed)")).unwrap()
13651                else {
13652                    unreachable!()
13653                };
13654                let (aggregate_exprs, _) =
13655                    planner.create_aggregate_exprs(op, &param, &input).unwrap();
13656                let plan = LogicalPlanBuilder::from(input)
13657                    .aggregate(vec![col("ts")], aggregate_exprs)
13658                    .unwrap()
13659                    .filter(planner.create_empty_values_filter_expr(false).unwrap())
13660                    .unwrap()
13661                    .build()
13662                    .unwrap();
13663
13664                let (_, batches) = execute(plan, &build_query_engine_state()).await;
13665                let mut warnings = vec![];
13666                let mut infos = vec![];
13667                collector.append_to(&mut warnings, &mut infos);
13668                assert!(infos.is_empty());
13669                if incompatible {
13670                    assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
13671                    assert!(warnings.iter().any(|warning| {
13672                        warning
13673                            == &format!(
13674                                "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas"
13675                            )
13676                    }));
13677                } else {
13678                    let histograms = histograms(&batches, &histogram_column);
13679                    assert_eq!(histograms.len(), 1);
13680                    let expected = if op_name == "sum" { 2.0 } else { 1.0 };
13681                    assert_eq!(histograms[0].count, expected);
13682                    assert_eq!(histograms[0].sum, expected);
13683                    assert!(warnings.is_empty());
13684                }
13685            }
13686        }
13687    }
13688
13689    #[tokio::test]
13690    async fn test_canonical_mixed_count_group_and_count_values_execute() {
13691        let state = build_query_engine_state();
13692        for (query, expected) in [
13693            ("count(some_metric)", vec![2.0]),
13694            ("group(some_metric)", vec![1.0]),
13695            (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]),
13696        ] {
13697            let plan = PromPlanner::stmt_to_plan(
13698                build_test_mixed_native_histogram_table_provider("some_metric").await,
13699                &operator_eval_stmt(query),
13700                &state,
13701            )
13702            .await
13703            .unwrap();
13704            assert!(
13705                plan.schema()
13706                    .fields()
13707                    .iter()
13708                    .all(|field| !field.name().starts_with("__promql_sample_count")),
13709                "{query}: {plan:?}"
13710            );
13711            let value_fields = plan
13712                .schema()
13713                .fields()
13714                .iter()
13715                .filter(|field| {
13716                    matches!(
13717                        field.data_type(),
13718                        ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64
13719                    ) || field.data_type() == &native_histogram_value_type().as_arrow_type()
13720                })
13721                .collect::<Vec<_>>();
13722            assert_eq!(value_fields.len(), 1, "{query}: {plan:?}");
13723            assert_ne!(
13724                value_fields[0].data_type(),
13725                &native_histogram_value_type().as_arrow_type(),
13726                "{query}: {plan:?}"
13727            );
13728            let value_column = value_fields[0].name().clone();
13729
13730            let (_, batches) = execute(plan, &state).await;
13731            let mut actual = numeric_values(&batches, &value_column);
13732            actual.sort_by(f64::total_cmp);
13733            assert_eq!(actual, expected, "{query}");
13734
13735            if query.starts_with("count_values") {
13736                let mut sample_labels = batches
13737                    .iter()
13738                    .flat_map(|batch| {
13739                        batch
13740                            .column_by_name("sample")
13741                            .unwrap()
13742                            .as_any()
13743                            .downcast_ref::<StringArray>()
13744                            .unwrap()
13745                            .iter()
13746                            .flatten()
13747                            .map(str::to_string)
13748                    })
13749                    .collect::<Vec<_>>();
13750                sample_labels.sort();
13751                let mut expected_labels =
13752                    vec!["2".to_string(), direct_or_histogram().promql_string()];
13753                expected_labels.sort();
13754                assert_eq!(sample_labels, expected_labels);
13755            }
13756        }
13757    }
13758
13759    #[tokio::test]
13760    async fn test_mixed_or_sum_aggregates_each_sample_type() {
13761        let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13762            parser::parse("sum(lhs)").unwrap()
13763        else {
13764            unreachable!()
13765        };
13766
13767        let collector = PromqlAnnotationCollector::default();
13768        let (mut planner, input) = mixed_direct_or(false).await;
13769        planner.promql_annotations = Some(collector.clone());
13770        let float_column = planner.ctx.field_columns[0].clone();
13771        let histogram_column = planner.ctx.field_columns[1].clone();
13772        let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
13773        let plan = LogicalPlanBuilder::from(input)
13774            .aggregate(vec![col("ts"), col("k")], aggregate_exprs)
13775            .unwrap()
13776            .filter(
13777                planner
13778                    .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
13779                    .unwrap(),
13780            )
13781            .unwrap()
13782            .project([
13783                col(&float_column),
13784                col(&histogram_column),
13785                col("ts"),
13786                col("k"),
13787            ])
13788            .unwrap()
13789            .build()
13790            .unwrap();
13791
13792        let (_, batches) = execute(plan, &build_query_engine_state()).await;
13793        assert_eq!(values(&batches, &float_column), vec![1.25]);
13794        let histogram = batches
13795            .iter()
13796            .find_map(|batch| {
13797                let values = batch
13798                    .column_by_name(&histogram_column)?
13799                    .as_any()
13800                    .downcast_ref::<datafusion::arrow::array::StructArray>()?;
13801                (0..values.len()).find_map(|row| {
13802                    common_query::native_histogram::read_histogram(values, row).unwrap()
13803                })
13804            })
13805            .unwrap();
13806        assert_eq!(histogram.count, 1.0);
13807        let mut warnings = vec![];
13808        let mut infos = vec![];
13809        collector.append_to(&mut warnings, &mut infos);
13810        assert!(warnings.is_empty());
13811
13812        let collector = PromqlAnnotationCollector::default();
13813        let (mut planner, input) = mixed_direct_or(false).await;
13814        planner.promql_annotations = Some(collector.clone());
13815        let float_column = planner.ctx.field_columns[0].clone();
13816        let histogram_column = planner.ctx.field_columns[1].clone();
13817        let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
13818        let plan = LogicalPlanBuilder::from(input)
13819            .aggregate(vec![col("ts")], aggregate_exprs)
13820            .unwrap()
13821            .filter(
13822                planner
13823                    .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
13824                    .unwrap(),
13825            )
13826            .unwrap()
13827            .project([col(&float_column), col(&histogram_column), col("ts")])
13828            .unwrap()
13829            .build()
13830            .unwrap();
13831
13832        let (_, batches) = execute(plan, &build_query_engine_state()).await;
13833        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
13834        let mut warnings = vec![];
13835        let mut infos = vec![];
13836        collector.append_to(&mut warnings, &mut infos);
13837        assert_eq!(
13838            warnings,
13839            vec![
13840                "sum: dropped aggregation result containing both float and native histogram samples"
13841            ]
13842        );
13843    }
13844
13845    #[tokio::test]
13846    async fn test_mixed_or_sum_drops_incompatible_mixed_group() {
13847        let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13848            parser::parse("sum(lhs)").unwrap()
13849        else {
13850            unreachable!()
13851        };
13852        let mut custom = direct_or_histogram();
13853        custom.schema = CUSTOM_BUCKETS_SCHEMA;
13854        custom.custom_values = vec![1.0];
13855        let collector = PromqlAnnotationCollector::default();
13856        let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await;
13857        planner.promql_annotations = Some(collector.clone());
13858        let float_column = planner.ctx.field_columns[0].clone();
13859        let histogram_column = planner.ctx.field_columns[1].clone();
13860        let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
13861        let plan = LogicalPlanBuilder::from(input)
13862            .aggregate(vec![col("ts")], aggregate_exprs)
13863            .unwrap()
13864            .filter(
13865                planner
13866                    .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
13867                    .unwrap(),
13868            )
13869            .unwrap()
13870            .project([col(&float_column), col(&histogram_column), col("ts")])
13871            .unwrap()
13872            .build()
13873            .unwrap();
13874
13875        let (_, batches) = execute(plan, &build_query_engine_state()).await;
13876        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
13877        let mut warnings = vec![];
13878        let mut infos = vec![];
13879        collector.append_to(&mut warnings, &mut infos);
13880        assert!(warnings.iter().any(|warning| {
13881            warning
13882                == "sum: dropped aggregation result containing both float and native histogram samples"
13883        }));
13884    }
13885
13886    #[tokio::test]
13887    async fn test_mixed_or_min_records_only_present_histograms() {
13888        let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13889            parser::parse("min(lhs)").unwrap()
13890        else {
13891            unreachable!()
13892        };
13893        let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms";
13894
13895        for (histograms, expected_infos) in [
13896            (vec![], vec![]),
13897            (vec![direct_or_histogram()], vec![expected_info]),
13898        ] {
13899            let collector = PromqlAnnotationCollector::default();
13900            let (mut planner, input) = mixed_aggregate_input(histograms).await;
13901            planner.promql_annotations = Some(collector.clone());
13902            let float_column = planner.ctx.field_columns[0].clone();
13903            let histogram_column = planner.ctx.field_columns[1].clone();
13904            let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
13905            let plan = LogicalPlanBuilder::from(input)
13906                .aggregate(vec![col("ts")], aggregate_exprs)
13907                .unwrap()
13908                .filter(
13909                    planner
13910                        .mixed_ignored_histogram_filter_expr(op, &histogram_column)
13911                        .unwrap(),
13912                )
13913                .unwrap()
13914                .project([col(&float_column), col("ts")])
13915                .unwrap()
13916                .build()
13917                .unwrap();
13918
13919            let (_, batches) = execute(plan, &build_query_engine_state()).await;
13920            assert_eq!(values(&batches, &float_column), vec![1.25]);
13921            let mut warnings = vec![];
13922            let mut infos = vec![];
13923            collector.append_to(&mut warnings, &mut infos);
13924            assert!(warnings.is_empty());
13925            assert_eq!(infos, expected_infos);
13926        }
13927    }
13928
13929    #[tokio::test]
13930    async fn test_mixed_or_value_aliases_do_not_replace_labels() {
13931        let left = source(
13932            "lhs",
13933            false,
13934            1,
13935            vec![("job", Some("job")), ("k", Some("float"))],
13936            DirectOrValue::Float64(1.0),
13937        );
13938        let right = source(
13939            "rhs",
13940            false,
13941            1,
13942            vec![
13943                ("job", Some("job")),
13944                ("k", Some("histogram")),
13945                (greptime_value(), Some("value-label")),
13946            ],
13947            DirectOrValue::NativeHistogram(direct_or_histogram()),
13948        );
13949        let table_provider = build_test_table_provider_with_fields(
13950            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
13951            &[],
13952        )
13953        .await;
13954        let mut planner = PromPlanner {
13955            table_provider,
13956            ctx: PromPlannerContext::default(),
13957            promql_annotations: None,
13958        };
13959        let left = LogicalPlanBuilder::from(scan(&left))
13960            .project(vec![
13961                col("ts"),
13962                col("job"),
13963                col("k"),
13964                col("v").alias(greptime_value()),
13965            ])
13966            .unwrap()
13967            .build()
13968            .unwrap();
13969        let left_context = direct_or_context("lhs", &["job", "k"], greptime_value());
13970        let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v");
13971        let plan = planner
13972            .or_operator(
13973                left,
13974                scan(&right),
13975                left_context.tag_columns.iter().cloned().collect(),
13976                right_context.tag_columns.iter().cloned().collect(),
13977                left_context,
13978                right_context,
13979                &or_modifier("lhs or on(k) rhs"),
13980            )
13981            .unwrap();
13982
13983        assert_eq!(
13984            plan.schema()
13985                .field_with_name(None, greptime_value())
13986                .unwrap()
13987                .data_type(),
13988            &ArrowDataType::Utf8
13989        );
13990        assert!(
13991            planner
13992                .ctx
13993                .field_columns
13994                .iter()
13995                .all(|field| { field != greptime_value() && field != greptime_native_histogram() })
13996        );
13997        assert!(PromPlanner::field_columns_are_alternative_samples(
13998            plan.schema(),
13999            &planner.ctx.field_columns
14000        ));
14001        let (_, batches) = execute(plan, &build_query_engine_state()).await;
14002        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
14003        let labels = batches
14004            .iter()
14005            .flat_map(|batch| {
14006                batch
14007                    .column_by_name(greptime_value())
14008                    .unwrap()
14009                    .as_any()
14010                    .downcast_ref::<StringArray>()
14011                    .unwrap()
14012                    .iter()
14013                    .flatten()
14014            })
14015            .collect::<Vec<_>>();
14016        assert_eq!(labels, vec!["value-label"]);
14017    }
14018
14019    #[tokio::test]
14020    async fn test_mixed_or_routes_float_histogram_and_label_functions() {
14021        for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] {
14022            let (mut planner, input) = mixed_direct_or(false).await;
14023            let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14024                input.schema(),
14025                &planner.ctx.field_columns,
14026            );
14027            let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else {
14028                unreachable!()
14029            };
14030            let state = build_query_engine_state();
14031            let (mut exprs, _) = planner
14032                .create_function_expr(&call.func, vec![], input.schema(), &state)
14033                .unwrap();
14034            exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14035            exprs.extend(planner.create_tag_column_exprs().unwrap());
14036            let plan = LogicalPlanBuilder::from(input)
14037                .project(exprs)
14038                .unwrap()
14039                .filter(
14040                    planner
14041                        .create_empty_values_filter_expr(preserve_any_value)
14042                        .unwrap(),
14043                )
14044                .unwrap()
14045                .build()
14046                .unwrap();
14047            let (_, batches) = execute(plan, &state).await;
14048            let values = batches
14049                .iter()
14050                .flat_map(|batch| {
14051                    batch
14052                        .schema()
14053                        .fields()
14054                        .iter()
14055                        .position(|field| field.data_type() == &ArrowDataType::Float64)
14056                        .map(|index| {
14057                            batch
14058                                .column(index)
14059                                .as_any()
14060                                .downcast_ref::<Float64Array>()
14061                                .unwrap()
14062                                .iter()
14063                                .flatten()
14064                        })
14065                        .into_iter()
14066                        .flatten()
14067                })
14068                .collect::<Vec<_>>();
14069            assert_eq!(values, vec![expected], "{function}");
14070        }
14071
14072        let (mut planner, input) = mixed_direct_or(false).await;
14073        let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14074            input.schema(),
14075            &planner.ctx.field_columns,
14076        );
14077        let PromExpr::Call(call) =
14078            parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap()
14079        else {
14080            unreachable!()
14081        };
14082        let args = planner.create_function_args(&call.args.args).unwrap();
14083        let state = build_query_engine_state();
14084        let (mut exprs, _) = planner
14085            .create_function_expr(&call.func, args.literals, input.schema(), &state)
14086            .unwrap();
14087        exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14088        exprs.extend(planner.create_tag_column_exprs().unwrap());
14089        let plan = LogicalPlanBuilder::from(input)
14090            .project(exprs)
14091            .unwrap()
14092            .filter(
14093                planner
14094                    .create_empty_values_filter_expr(preserve_any_value)
14095                    .unwrap(),
14096            )
14097            .unwrap()
14098            .build()
14099            .unwrap();
14100        let (_, batches) = execute(plan, &state).await;
14101        let sample_count = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
14102        assert_eq!(sample_count, 2);
14103    }
14104}