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            0,
591            range_ms,
592            time_index_column,
593            self.ctx.field_columns.clone(),
594            divide_plan,
595        )
596        .context(DataFusionPlanningSnafu)?;
597
598        Ok(LogicalPlan::Extension(Extension {
599            node: Arc::new(manipulate),
600        }))
601    }
602
603    async fn prom_aggr_expr_to_plan(
604        &mut self,
605        query_engine_state: &QueryEngineState,
606        aggr_expr: &AggregateExpr,
607    ) -> Result<LogicalPlan> {
608        let AggregateExpr {
609            op,
610            expr,
611            modifier,
612            param,
613        } = aggr_expr;
614
615        let mut input = self.prom_expr_to_plan(expr, query_engine_state).await?;
616        let input_has_tsid = input.schema().fields().iter().any(|field| {
617            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
618                && field.data_type() == &ArrowDataType::UInt64
619        });
620
621        // `__tsid` based scan projection may prune tag columns. Ensure tags referenced in
622        // aggregation modifiers (`by`/`without`) are available before planning group keys.
623        let required_group_tags = match modifier {
624            None => BTreeSet::new(),
625            Some(LabelModifier::Include(labels)) => labels
626                .labels
627                .iter()
628                .filter(|label| !is_metric_engine_internal_column(label.as_str()))
629                .cloned()
630                .collect(),
631            Some(LabelModifier::Exclude(labels)) => {
632                let mut all_tags = self.collect_row_key_tag_columns_from_plan(&input)?;
633                for label in &labels.labels {
634                    let _ = all_tags.remove(label);
635                }
636                all_tags
637            }
638        };
639
640        if !required_group_tags.is_empty()
641            && required_group_tags
642                .iter()
643                .any(|tag| Self::find_case_sensitive_column(input.schema(), tag.as_str()).is_none())
644        {
645            input = self.ensure_tag_columns_available(input, &required_group_tags)?;
646            self.refresh_tag_columns_from_schema(input.schema());
647        }
648
649        match (*op).id() {
650            token::T_TOPK | token::T_BOTTOMK => {
651                self.prom_topk_bottomk_to_plan(aggr_expr, input).await
652            }
653            _ => {
654                // When `__tsid` is available, tag columns may have been pruned from the input plan.
655                // For `keep_tsid` decision we should compare against the full row-key label set,
656                // otherwise we may incorrectly treat label-reducing aggregates as preserving labels.
657                let input_tag_columns = if input_has_tsid {
658                    self.collect_row_key_tag_columns_from_plan(&input)?
659                        .into_iter()
660                        .collect::<Vec<_>>()
661                } else {
662                    self.ctx.tag_columns.clone()
663                };
664                // calculate columns to group by
665                // Need to append time index column into group by columns
666                let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
667                let mixed_sample_columns =
668                    Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
669                        .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
670                // Aggregates over native histogram inputs may drop every sample in a group
671                // (e.g. `min` over histogram-only samples, or `sum` over histograms with
672                // incompatible schemas) and leave a NULL-valued group row behind. Compute this
673                // before `create_aggregate_exprs` mutates `ctx.field_columns`.
674                let preserve_any_value = mixed_sample_columns.is_some();
675                let has_native_histogram = preserve_any_value
676                    || self.all_field_columns_are_native_histograms(input.schema());
677                // convert op and value columns to aggregate exprs
678                let (mut aggr_exprs, prev_field_exprs) =
679                    self.create_aggregate_exprs(*op, param, &input)?;
680                let prev_field_exprs =
681                    normalize_cols(prev_field_exprs, &input).context(DataFusionPlanningSnafu)?;
682
683                let keep_tsid = op.id() != token::T_COUNT_VALUES
684                    && input_has_tsid
685                    && input_tag_columns.iter().collect::<HashSet<_>>()
686                        == self.ctx.tag_columns.iter().collect::<HashSet<_>>();
687
688                if keep_tsid {
689                    aggr_exprs.push(
690                        first_value(
691                            DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)),
692                            vec![],
693                        )
694                        .alias(DATA_SCHEMA_TSID_COLUMN_NAME),
695                    );
696                }
697                self.ctx.use_tsid = keep_tsid;
698
699                // create plan
700                let builder = LogicalPlanBuilder::from(input);
701                let builder = if op.id() == token::T_COUNT_VALUES {
702                    let label = Self::get_param_value_as_str(*op, param)?;
703                    // `count_values` must be grouped by fields,
704                    // and project the fields to the new label.
705                    let count_value_exprs = prev_field_exprs.iter().map(|expr| {
706                        match expr {
707                            DfExpr::Column(column) => DfExpr::Column(column.clone()),
708                            _ => DfExpr::Column(Column::from_name(expr.schema_name().to_string())),
709                        }
710                        .alias(label)
711                    });
712                    let aggregate_group_exprs = group_exprs
713                        .iter()
714                        .cloned()
715                        .chain(prev_field_exprs.clone())
716                        .collect::<Vec<_>>();
717                    group_exprs.push(col(label));
718                    let project_fields = self
719                        .create_field_column_exprs()?
720                        .into_iter()
721                        .chain(self.create_tag_column_exprs()?)
722                        .chain(Some(self.create_time_index_column_expr()?))
723                        .chain(count_value_exprs);
724
725                    builder
726                        .aggregate(aggregate_group_exprs, aggr_exprs)
727                        .context(DataFusionPlanningSnafu)?
728                        .project(project_fields)
729                        .context(DataFusionPlanningSnafu)?
730                } else {
731                    builder
732                        .aggregate(group_exprs.clone(), aggr_exprs)
733                        .context(DataFusionPlanningSnafu)?
734                };
735
736                let builder = if let Some((float, histogram)) = mixed_sample_columns {
737                    let builder = match op.id() {
738                        token::T_SUM | token::T_AVG => builder
739                            .filter(self.mixed_aggregate_filter_expr(*op, &float, &histogram)?)
740                            .context(DataFusionPlanningSnafu)?,
741                        token::T_MIN
742                        | token::T_MAX
743                        | token::T_STDDEV
744                        | token::T_STDVAR
745                        | token::T_QUANTILE => builder
746                            .filter(self.mixed_ignored_histogram_filter_expr(*op, &histogram)?)
747                            .context(DataFusionPlanningSnafu)?,
748                        _ => builder,
749                    };
750
751                    match op.id() {
752                        token::T_SUM
753                        | token::T_AVG
754                        | token::T_MIN
755                        | token::T_MAX
756                        | token::T_STDDEV
757                        | token::T_STDVAR
758                        | token::T_QUANTILE => {
759                            let project_fields = self
760                                .create_field_column_exprs()?
761                                .into_iter()
762                                .chain(self.create_tag_column_exprs()?)
763                                .chain(self.ctx.use_tsid.then_some(DfExpr::Column(
764                                    Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME),
765                                )))
766                                .chain(Some(self.create_time_index_column_expr()?));
767                            builder
768                                .project(project_fields)
769                                .context(DataFusionPlanningSnafu)?
770                        }
771                        _ => builder,
772                    }
773                } else {
774                    builder
775                };
776
777                // Drop group rows whose every aggregated sample was discarded (NULL), so that
778                // e.g. `group(min(native_histogram))` doesn't resurrect groups Prometheus
779                // considers unseen. For alternative float/histogram fields keep the row if any
780                // field survived.
781                let builder = if has_native_histogram {
782                    builder
783                        .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
784                        .context(DataFusionPlanningSnafu)?
785                } else {
786                    builder
787                };
788
789                let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false));
790
791                builder
792                    .sort(sort_expr)
793                    .context(DataFusionPlanningSnafu)?
794                    .build()
795                    .context(DataFusionPlanningSnafu)
796            }
797        }
798    }
799
800    /// Create logical plan for PromQL topk and bottomk expr.
801    async fn prom_topk_bottomk_to_plan(
802        &mut self,
803        aggr_expr: &AggregateExpr,
804        input: LogicalPlan,
805    ) -> Result<LogicalPlan> {
806        let AggregateExpr {
807            op,
808            param,
809            modifier,
810            ..
811        } = aggr_expr;
812
813        let input_has_tsid = input.schema().fields().iter().any(|field| {
814            field.name() == DATA_SCHEMA_TSID_COLUMN_NAME
815                && field.data_type() == &ArrowDataType::UInt64
816        });
817        self.ctx.use_tsid = input_has_tsid;
818
819        let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?;
820
821        let mut input = input;
822        if let Some((float_column, histogram_column)) =
823            Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns)
824                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
825        {
826            let drop_histogram = DfExpr::ScalarFunction(ScalarFunction {
827                func: Arc::new(NativeHistogramDrop::bool_false_udf(
828                    format!(
829                        "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
830                        op
831                    ),
832                    self.promql_annotations.clone(),
833                )),
834                args: vec![col(&histogram_column)],
835            });
836            let keep_float = when(col(&histogram_column).is_not_null(), drop_histogram)
837                .otherwise(col(&float_column).is_not_null())
838                .context(DataFusionPlanningSnafu)?;
839            input = LogicalPlanBuilder::from(input)
840                .filter(keep_float)
841                .context(DataFusionPlanningSnafu)?
842                .build()
843                .context(DataFusionPlanningSnafu)?;
844            self.ctx.field_columns = vec![float_column];
845        }
846
847        if self.all_field_columns_are_native_histograms(input.schema()) {
848            let promql_annotations = self.promql_annotations.clone();
849            let input = self.projection_for_each_field_column(input, |col| {
850                Ok(DfExpr::ScalarFunction(ScalarFunction {
851                    func: Arc::new(NativeHistogramDrop::float_null_udf(
852                        format!(
853                            "{}: dropped native histogram samples because this aggregation is not supported for native histograms",
854                            op
855                        ),
856                        promql_annotations.clone(),
857                    )),
858                    args: vec![DfExpr::Column(Column::from_name(col))],
859                }))
860            })?;
861            return LogicalPlanBuilder::from(input)
862                .filter(self.create_empty_values_filter_expr(false)?)
863                .context(DataFusionPlanningSnafu)?
864                .build()
865                .context(DataFusionPlanningSnafu);
866        }
867
868        let val = Self::get_param_as_literal_expr(
869            param.as_deref(),
870            Some(*op),
871            Some(ArrowDataType::Float64),
872        )?;
873
874        // convert op and value columns to window exprs.
875        let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?;
876
877        let rank_columns: Vec<_> = window_exprs
878            .iter()
879            .map(|expr| expr.schema_name().to_string())
880            .collect();
881
882        // Create ranks filter with `Operator::Or`.
883        // Safety: at least one rank column
884        let filter: DfExpr = rank_columns
885            .iter()
886            .fold(None, |expr, rank| {
887                let predicate = DfExpr::BinaryExpr(BinaryExpr {
888                    left: Box::new(col(rank)),
889                    op: Operator::LtEq,
890                    right: Box::new(val.clone()),
891                });
892
893                match expr {
894                    None => Some(predicate),
895                    Some(expr) => Some(DfExpr::BinaryExpr(BinaryExpr {
896                        left: Box::new(expr),
897                        op: Operator::Or,
898                        right: Box::new(predicate),
899                    })),
900                }
901            })
902            .unwrap();
903
904        let rank_columns: Vec<_> = rank_columns.into_iter().map(col).collect();
905
906        let mut new_group_exprs = group_exprs.clone();
907        // Order by ranks
908        new_group_exprs.extend(rank_columns);
909
910        let group_sort_expr = new_group_exprs
911            .into_iter()
912            .map(|expr| expr.sort(true, false));
913
914        let project_fields = self
915            .create_field_column_exprs()?
916            .into_iter()
917            .chain(self.create_tag_column_exprs()?)
918            .chain(
919                self.ctx
920                    .use_tsid
921                    .then_some(DfExpr::Column(Column::from_name(
922                        DATA_SCHEMA_TSID_COLUMN_NAME,
923                    ))),
924            )
925            .chain(Some(self.create_time_index_column_expr()?));
926
927        LogicalPlanBuilder::from(input)
928            .window(window_exprs)
929            .context(DataFusionPlanningSnafu)?
930            .filter(filter)
931            .context(DataFusionPlanningSnafu)?
932            .sort(group_sort_expr)
933            .context(DataFusionPlanningSnafu)?
934            .project(project_fields)
935            .context(DataFusionPlanningSnafu)?
936            .build()
937            .context(DataFusionPlanningSnafu)
938    }
939
940    async fn prom_unary_expr_to_plan(
941        &mut self,
942        query_engine_state: &QueryEngineState,
943        unary_expr: &UnaryExpr,
944    ) -> Result<LogicalPlan> {
945        let UnaryExpr { expr } = unary_expr;
946        // Unary Expr in PromQL implys the `-` operator
947        let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
948        self.negate_field_columns(input)
949    }
950
951    fn negate_field_columns(&mut self, input: LogicalPlan) -> Result<LogicalPlan> {
952        let input_schema = input.schema().clone();
953        self.projection_for_each_field_column(input, |col| {
954            if Self::field_column_is_native_histogram(&input_schema, col) {
955                Ok(DfExpr::ScalarFunction(ScalarFunction {
956                    func: Arc::new(NativeHistogramNeg::scalar_udf()),
957                    args: vec![DfExpr::Column(col.into())],
958                }))
959            } else {
960                Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into()))))
961            }
962        })
963    }
964
965    async fn try_plan_binary_island(
966        &mut self,
967        binary_expr: &PromBinaryExpr,
968    ) -> Result<Option<LogicalPlan>> {
969        let original_ctx = self.ctx.clone();
970        let mut collect_env = IslandCollectEnv::default();
971        let Some(island_expr) =
972            IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env)
973        else {
974            return Ok(None);
975        };
976
977        if collect_env.leaves.is_empty()
978            || collect_env.vector_occurrences <= collect_env.leaves.len()
979        {
980            return Ok(None);
981        }
982
983        let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len());
984        for (idx, leaf) in collect_env.leaves.iter().enumerate() {
985            let plan = self
986                .prom_vector_selector_to_plan(&leaf.selector, false)
987                .await?;
988            let ctx = self.ctx.clone();
989            let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}"));
990            let plan = LogicalPlanBuilder::from(plan)
991                .alias(alias.clone())
992                .context(DataFusionPlanningSnafu)?
993                .build()
994                .context(DataFusionPlanningSnafu)?;
995            planned_leaves.push(PlannedIslandLeaf {
996                plan,
997                ctx,
998                alias,
999                display_table: leaf.display_table.clone(),
1000            });
1001        }
1002
1003        if planned_leaves.iter().any(|leaf| {
1004            Self::field_columns_contain_native_histogram(
1005                leaf.plan.schema(),
1006                &leaf.ctx.field_columns,
1007            )
1008        }) {
1009            self.ctx = original_ctx;
1010            return Ok(None);
1011        }
1012
1013        if !Self::binary_island_join_contexts_supported(&planned_leaves) {
1014            self.ctx = original_ctx;
1015            return Ok(None);
1016        }
1017
1018        let mut input = planned_leaves[0].plan.clone();
1019        for right_idx in 1..planned_leaves.len() {
1020            input = self.join_binary_island_leaf(
1021                input,
1022                &planned_leaves[0],
1023                &planned_leaves[right_idx],
1024            )?;
1025        }
1026
1027        let field_exprs =
1028            Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?;
1029        if field_exprs.scalar || field_exprs.exprs.is_empty() {
1030            self.ctx = original_ctx;
1031            return Ok(None);
1032        }
1033
1034        let plan = self.project_binary_island(
1035            input,
1036            &planned_leaves[0].alias,
1037            &planned_leaves[0].ctx,
1038            field_exprs,
1039        )?;
1040        Ok(Some(plan))
1041    }
1042
1043    fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool {
1044        if leaves
1045            .iter()
1046            .any(|leaf| leaf.ctx.time_index_column.is_none())
1047        {
1048            return false;
1049        }
1050
1051        if leaves.len() <= 1 {
1052            return true;
1053        }
1054
1055        let first_tags = leaves[0].ctx.tag_columns.iter().collect::<BTreeSet<_>>();
1056
1057        leaves.iter().skip(1).all(|leaf| {
1058            (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan))
1059                || leaf.ctx.tag_columns.iter().collect::<BTreeSet<_>>() == first_tags
1060        })
1061    }
1062
1063    fn join_binary_island_leaf(
1064        &self,
1065        left: LogicalPlan,
1066        first_leaf: &PlannedIslandLeaf,
1067        right_leaf: &PlannedIslandLeaf,
1068    ) -> Result<LogicalPlan> {
1069        let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty()
1070            || right_leaf.ctx.tag_columns.is_empty())
1071            && !first_leaf
1072                .ctx
1073                .tag_columns
1074                .iter()
1075                .chain(&right_leaf.ctx.tag_columns)
1076                .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL);
1077        let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns(
1078            left.schema(),
1079            right_leaf.plan.schema(),
1080            &first_leaf.ctx,
1081            &right_leaf.ctx,
1082            only_join_time_index,
1083            &None,
1084        )?;
1085
1086        if let (Some(left_time_index_column), Some(right_time_index_column)) = (
1087            first_leaf.ctx.time_index_column.clone(),
1088            right_leaf.ctx.time_index_column.clone(),
1089        ) {
1090            left_keys.insert(left_time_index_column);
1091            right_keys.insert(right_time_index_column);
1092        }
1093
1094        LogicalPlanBuilder::from(left)
1095            .join_detailed(
1096                right_leaf.plan.clone(),
1097                JoinType::Inner,
1098                (
1099                    left_keys
1100                        .into_iter()
1101                        .map(|name| Column::new(Some(first_leaf.alias.clone()), name))
1102                        .collect::<Vec<_>>(),
1103                    right_keys
1104                        .into_iter()
1105                        .map(|name| Column::new(Some(right_leaf.alias.clone()), name))
1106                        .collect::<Vec<_>>(),
1107                ),
1108                force_empty_join.then_some(lit(false)),
1109                NullEquality::NullEqualsNull,
1110            )
1111            .context(DataFusionPlanningSnafu)?
1112            .build()
1113            .context(DataFusionPlanningSnafu)
1114    }
1115
1116    fn build_binary_island_field_exprs(
1117        expr: &IslandExpr,
1118        leaves: &[PlannedIslandLeaf],
1119        schema: &DFSchemaRef,
1120    ) -> Result<IslandFieldExprs> {
1121        match expr {
1122            IslandExpr::VectorLeaf(id) => {
1123                let leaf = &leaves[*id];
1124                let exprs = leaf
1125                    .ctx
1126                    .field_columns
1127                    .iter()
1128                    .map(|field| {
1129                        schema
1130                            .qualified_field_with_name(Some(&leaf.alias), field)
1131                            .context(DataFusionPlanningSnafu)
1132                            .map(|field| DfExpr::Column(field.into()))
1133                    })
1134                    .collect::<Result<Vec<_>>>()?;
1135                let names = leaf
1136                    .ctx
1137                    .field_columns
1138                    .iter()
1139                    .map(|field| format!("{}.{}", leaf.display_table, field))
1140                    .collect();
1141                Ok(IslandFieldExprs {
1142                    exprs,
1143                    names,
1144                    scalar: false,
1145                })
1146            }
1147            IslandExpr::Scalar(expr) => Ok(IslandFieldExprs {
1148                exprs: vec![expr.clone()],
1149                names: vec![expr.schema_name().to_string()],
1150                scalar: true,
1151            }),
1152            IslandExpr::Unary { input } => {
1153                let input = Self::build_binary_island_field_exprs(input, leaves, schema)?;
1154                let mut exprs = Vec::with_capacity(input.exprs.len());
1155                let mut names = Vec::with_capacity(input.names.len());
1156                for (expr, name) in input.exprs.into_iter().zip(input.names) {
1157                    exprs.push(DfExpr::Negative(Box::new(expr)));
1158                    names.push(format!("-{name}"));
1159                }
1160                Ok(IslandFieldExprs {
1161                    exprs,
1162                    names,
1163                    scalar: input.scalar,
1164                })
1165            }
1166            IslandExpr::Binary { op, lhs, rhs } => {
1167                let same_leaf = match (&**lhs, &**rhs) {
1168                    (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right))
1169                        if left == right =>
1170                    {
1171                        Some(*left)
1172                    }
1173                    _ => None,
1174                };
1175                let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?;
1176                let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?;
1177                let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1178                let scalar = lhs.scalar && rhs.scalar;
1179                let op = op.to_string();
1180
1181                let (exprs, names) = match (lhs.scalar, rhs.scalar) {
1182                    (true, true) => {
1183                        let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?;
1184                        let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]);
1185                        (vec![expr], vec![name])
1186                    }
1187                    (true, false) => {
1188                        let mut exprs = Vec::with_capacity(rhs.exprs.len());
1189                        let mut names = Vec::with_capacity(rhs.names.len());
1190                        for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) {
1191                            exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?);
1192                            names.push(format!("{} {op} {rhs_name}", lhs.names[0]));
1193                        }
1194                        (exprs, names)
1195                    }
1196                    (false, true) => {
1197                        let mut exprs = Vec::with_capacity(lhs.exprs.len());
1198                        let mut names = Vec::with_capacity(lhs.names.len());
1199                        for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) {
1200                            exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?);
1201                            names.push(format!("{lhs_name} {op} {}", rhs.names[0]));
1202                        }
1203                        (exprs, names)
1204                    }
1205                    (false, false) => {
1206                        let mut exprs = Vec::new();
1207                        let mut names = Vec::new();
1208                        for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs
1209                            .exprs
1210                            .into_iter()
1211                            .zip(rhs.exprs)
1212                            .zip(lhs.names.into_iter().zip(rhs.names))
1213                            .enumerate()
1214                        {
1215                            if let Some(leaf) = same_leaf {
1216                                let field = leaves[leaf]
1217                                    .ctx
1218                                    .field_columns
1219                                    .get(idx)
1220                                    .cloned()
1221                                    .unwrap_or_else(|| lhs_name.clone());
1222                                lhs_name = format!("lhs.{field}");
1223                                rhs_name = format!("rhs.{field}");
1224                            }
1225                            exprs.push(expr_builder(lhs_expr, rhs_expr)?);
1226                            names.push(format!("{lhs_name} {op} {rhs_name}"));
1227                        }
1228                        (exprs, names)
1229                    }
1230                };
1231
1232                Ok(IslandFieldExprs {
1233                    exprs,
1234                    names,
1235                    scalar,
1236                })
1237            }
1238        }
1239    }
1240
1241    fn project_binary_island(
1242        &mut self,
1243        input: LogicalPlan,
1244        base_alias: &TableReference,
1245        base_ctx: &PromPlannerContext,
1246        field_exprs: IslandFieldExprs,
1247    ) -> Result<LogicalPlan> {
1248        self.ctx = base_ctx.clone();
1249
1250        let schema = input.schema();
1251        let non_field_exprs = base_ctx
1252            .tag_columns
1253            .iter()
1254            .chain(base_ctx.time_index_column.iter())
1255            .map(|column| {
1256                schema
1257                    .qualified_field_with_name(Some(base_alias), column)
1258                    .context(DataFusionPlanningSnafu)
1259                    .map(|field| DfExpr::Column(field.into()))
1260            });
1261        let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid)
1262            .into_iter()
1263            .map(Ok);
1264
1265        self.ctx.field_columns = field_exprs.names;
1266        let field_exprs = field_exprs
1267            .exprs
1268            .into_iter()
1269            .zip(self.ctx.field_columns.iter())
1270            .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
1271
1272        let project_exprs = non_field_exprs
1273            .chain(tsid_expr)
1274            .chain(field_exprs)
1275            .collect::<Result<Vec<_>>>()?;
1276
1277        let plan = LogicalPlanBuilder::from(input)
1278            .project(project_exprs)
1279            .context(DataFusionPlanningSnafu)?
1280            .build()
1281            .context(DataFusionPlanningSnafu)?;
1282
1283        self.ctx.table_name = None;
1284        self.ctx.schema_name = None;
1285
1286        Ok(plan)
1287    }
1288
1289    async fn prom_binary_expr_to_plan(
1290        &mut self,
1291        query_engine_state: &QueryEngineState,
1292        binary_expr: &PromBinaryExpr,
1293    ) -> Result<LogicalPlan> {
1294        // promql-parser accepts fill modifiers, but Greptime does not implement the
1295        // required outer joins and missing-value substitution. Reject them before the
1296        // binary-island fast path so they cannot silently behave like normal inner joins.
1297        if let Some(modifier) = &binary_expr.modifier {
1298            ensure!(
1299                modifier.fill_values.lhs.is_none() && modifier.fill_values.rhs.is_none(),
1300                UnsupportedExprSnafu {
1301                    name: "PromQL fill modifiers"
1302                }
1303            );
1304        }
1305
1306        if let Some(plan) = self.try_plan_binary_island(binary_expr).await? {
1307            return Ok(plan);
1308        }
1309
1310        let PromBinaryExpr {
1311            lhs,
1312            rhs,
1313            op,
1314            modifier,
1315        } = binary_expr;
1316
1317        // if set to true, comparison operator will return 0/1 (for true/false) instead of
1318        // filter on the result column
1319        let should_return_bool = if let Some(m) = modifier {
1320            m.return_bool
1321        } else {
1322            false
1323        };
1324        let is_comparison_op = Self::is_token_a_comparison_op(*op);
1325
1326        // we should build a filter plan here if the op is comparison op and need not
1327        // to return 0/1. Otherwise, we should build a projection plan
1328        match (
1329            Self::try_build_literal_expr(lhs),
1330            Self::try_build_literal_expr(rhs),
1331        ) {
1332            (Some(lhs), Some(rhs)) => {
1333                self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1334                self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1335                self.ctx.reset_table_name_and_schema();
1336                let field_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1337                let mut field_expr = field_expr_builder(lhs, rhs)?;
1338
1339                if is_comparison_op && should_return_bool {
1340                    field_expr = DfExpr::Cast(Cast {
1341                        expr: Box::new(field_expr),
1342                        data_type: ArrowDataType::Float64,
1343                    });
1344                }
1345
1346                Ok(LogicalPlan::Extension(Extension {
1347                    node: Arc::new(
1348                        EmptyMetric::new(
1349                            self.ctx.start,
1350                            self.ctx.end,
1351                            self.ctx.interval,
1352                            SPECIAL_TIME_FUNCTION.to_string(),
1353                            DEFAULT_FIELD_COLUMN.to_string(),
1354                            Some(field_expr),
1355                        )
1356                        .context(DataFusionPlanningSnafu)?,
1357                    ),
1358                }))
1359            }
1360            // lhs is a literal, rhs is a column
1361            (Some(mut expr), None) => {
1362                let input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1363                // check if the literal is a special time expr
1364                if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) {
1365                    expr = time_expr
1366                }
1367                let input_schema = input.schema().clone();
1368                let preserve_any_value = Self::field_columns_are_alternative_samples(
1369                    &input_schema,
1370                    &self.ctx.field_columns,
1371                );
1372                let has_native_histogram = Self::field_columns_contain_native_histogram(
1373                    &input_schema,
1374                    &self.ctx.field_columns,
1375                );
1376                let retain_field_columns = self
1377                    .ctx
1378                    .field_columns
1379                    .iter()
1380                    .map(|col| {
1381                        Self::binary_result_is_histogram(
1382                            *op,
1383                            false,
1384                            Self::field_column_is_native_histogram(&input_schema, col),
1385                        )
1386                        .is_some()
1387                    })
1388                    .collect();
1389                let promql_annotations = self.promql_annotations.clone();
1390                let bin_expr_builder = |col: &String| {
1391                    let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1392                    let rhs_is_histogram =
1393                        Self::field_column_is_native_histogram(&input_schema, col);
1394                    let rhs = DfExpr::Column(col.into());
1395                    let mut binary_expr = match Self::native_histogram_binary_expr(
1396                        *op,
1397                        expr.clone(),
1398                        false,
1399                        rhs.clone(),
1400                        rhs_is_histogram,
1401                        is_comparison_op && !should_return_bool,
1402                        promql_annotations.clone(),
1403                    )? {
1404                        Some(expr) => expr,
1405                        None => binary_expr_builder(expr.clone(), rhs)?,
1406                    };
1407
1408                    if is_comparison_op && should_return_bool {
1409                        binary_expr = DfExpr::Cast(Cast {
1410                            expr: Box::new(binary_expr),
1411                            data_type: ArrowDataType::Float64,
1412                        });
1413                    }
1414                    Ok(binary_expr)
1415                };
1416                if is_comparison_op && !should_return_bool {
1417                    self.filter_on_field_column(input, bin_expr_builder)
1418                } else {
1419                    let projected =
1420                        self.projection_for_each_field_column(input, bin_expr_builder)?;
1421                    self.filter_binary_projection(
1422                        projected,
1423                        has_native_histogram,
1424                        preserve_any_value,
1425                        retain_field_columns,
1426                    )
1427                }
1428            }
1429            // lhs is a column, rhs is a literal
1430            (None, Some(mut expr)) => {
1431                let input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1432                // check if the literal is a special time expr
1433                if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) {
1434                    expr = time_expr
1435                }
1436                let input_schema = input.schema().clone();
1437                let preserve_any_value = Self::field_columns_are_alternative_samples(
1438                    &input_schema,
1439                    &self.ctx.field_columns,
1440                );
1441                let has_native_histogram = Self::field_columns_contain_native_histogram(
1442                    &input_schema,
1443                    &self.ctx.field_columns,
1444                );
1445                let retain_field_columns = self
1446                    .ctx
1447                    .field_columns
1448                    .iter()
1449                    .map(|col| {
1450                        Self::binary_result_is_histogram(
1451                            *op,
1452                            Self::field_column_is_native_histogram(&input_schema, col),
1453                            false,
1454                        )
1455                        .is_some()
1456                    })
1457                    .collect();
1458                let promql_annotations = self.promql_annotations.clone();
1459                let bin_expr_builder = |col: &String| {
1460                    let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1461                    let lhs_is_histogram =
1462                        Self::field_column_is_native_histogram(&input_schema, col);
1463                    let lhs = DfExpr::Column(col.into());
1464                    let mut binary_expr = match Self::native_histogram_binary_expr(
1465                        *op,
1466                        lhs.clone(),
1467                        lhs_is_histogram,
1468                        expr.clone(),
1469                        false,
1470                        is_comparison_op && !should_return_bool,
1471                        promql_annotations.clone(),
1472                    )? {
1473                        Some(expr) => expr,
1474                        None => binary_expr_builder(lhs, expr.clone())?,
1475                    };
1476
1477                    if is_comparison_op && should_return_bool {
1478                        binary_expr = DfExpr::Cast(Cast {
1479                            expr: Box::new(binary_expr),
1480                            data_type: ArrowDataType::Float64,
1481                        });
1482                    }
1483                    Ok(binary_expr)
1484                };
1485                if is_comparison_op && !should_return_bool {
1486                    self.filter_on_field_column(input, bin_expr_builder)
1487                } else {
1488                    let projected =
1489                        self.projection_for_each_field_column(input, bin_expr_builder)?;
1490                    self.filter_binary_projection(
1491                        projected,
1492                        has_native_histogram,
1493                        preserve_any_value,
1494                        retain_field_columns,
1495                    )
1496                }
1497            }
1498            // both are columns. join them on time index
1499            (None, None) => {
1500                let left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?;
1501                let left_field_columns = self.ctx.field_columns.clone();
1502                let left_time_index_column = self.ctx.time_index_column.clone();
1503                let mut left_table_ref = self
1504                    .table_ref()
1505                    .unwrap_or_else(|_| TableReference::bare(""));
1506                let left_context = self.ctx.clone();
1507
1508                let right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?;
1509                let right_field_columns = self.ctx.field_columns.clone();
1510                let right_time_index_column = self.ctx.time_index_column.clone();
1511                let mut right_table_ref = self
1512                    .table_ref()
1513                    .unwrap_or_else(|_| TableReference::bare(""));
1514                let right_context = self.ctx.clone();
1515                let left_is_empty_metric = Self::is_empty_metric(&left_input);
1516                let right_is_empty_metric = Self::is_empty_metric(&right_input);
1517
1518                // TODO(ruihang): avoid join if left and right are the same table
1519
1520                // set op has "special" join semantics
1521                if Self::is_token_a_set_op(*op) {
1522                    return self.set_op_on_non_field_columns(
1523                        left_input,
1524                        right_input,
1525                        left_context,
1526                        right_context,
1527                        *op,
1528                        modifier,
1529                    );
1530                }
1531
1532                let has_native_histogram = Self::field_columns_contain_native_histogram(
1533                    left_input.schema(),
1534                    &left_field_columns,
1535                ) || Self::field_columns_contain_native_histogram(
1536                    right_input.schema(),
1537                    &right_field_columns,
1538                );
1539
1540                // normal join
1541                if left_table_ref == right_table_ref {
1542                    // rename table references to avoid ambiguity
1543                    left_table_ref = TableReference::bare("lhs");
1544                    right_table_ref = TableReference::bare("rhs");
1545                    // `self.ctx` have ctx in right plan, if right plan have no tag,
1546                    // we use left plan ctx as the ctx for subsequent calculations,
1547                    // to avoid case like `host + scalar(...)`
1548                    // we need preserve tag column on `host` table in subsequent projection,
1549                    // which only show in left plan ctx.
1550                    if self.ctx.tag_columns.is_empty() {
1551                        self.ctx = left_context.clone();
1552                        self.ctx.table_name = Some("lhs".to_string());
1553                    } else {
1554                        self.ctx.table_name = Some("rhs".to_string());
1555                    }
1556                } else if right_is_empty_metric && !left_is_empty_metric {
1557                    self.ctx = left_context.clone();
1558                }
1559                // Computed scalars reach this join path instead of the literal projection paths.
1560                // Broadcast them for arithmetic in the same way as literal scalars.
1561                let broadcast_scalar = !is_comparison_op;
1562                let (field_groups, invalid_field_pairs) = Self::align_binary_field_columns(
1563                    left_input.schema(),
1564                    right_input.schema(),
1565                    &left_field_columns,
1566                    &right_field_columns,
1567                    *op,
1568                    broadcast_scalar && lhs.value_type() == ValueType::Scalar,
1569                    broadcast_scalar && rhs.value_type() == ValueType::Scalar,
1570                );
1571                let left_aligned_field_columns = field_groups
1572                    .iter()
1573                    .flat_map(|(_, pairs)| {
1574                        pairs
1575                            .iter()
1576                            .map(|(left_col_name, _)| (*left_col_name).clone())
1577                    })
1578                    .collect::<Vec<_>>();
1579                let right_aligned_field_columns = field_groups
1580                    .iter()
1581                    .flat_map(|(_, pairs)| {
1582                        pairs
1583                            .iter()
1584                            .map(|(_, right_col_name)| (*right_col_name).clone())
1585                    })
1586                    .collect::<Vec<_>>();
1587                // Regular multi-field vectors combine their shared prefix. Alternative
1588                // float/histogram lanes instead align by valid PromQL sample combinations.
1589                self.ctx.field_columns = field_groups
1590                    .iter()
1591                    .map(|(output, _)| output.clone())
1592                    .collect();
1593                let mut field_groups = field_groups.into_iter();
1594                // `vector()` uses EmptyMetric and keeps GreptimeDB's timestamp broadcast.
1595                let has_empty_metric_operand = left_is_empty_metric || right_is_empty_metric;
1596
1597                let join_plan = self.join_on_non_field_columns(
1598                    left_input,
1599                    right_input,
1600                    left_table_ref.clone(),
1601                    right_table_ref.clone(),
1602                    left_time_index_column,
1603                    right_time_index_column,
1604                    lhs.value_type() == ValueType::Scalar
1605                        || rhs.value_type() == ValueType::Scalar
1606                        || has_empty_metric_operand
1607                        || ((left_context.tag_columns.is_empty()
1608                            || right_context.tag_columns.is_empty())
1609                            && !left_context
1610                                .tag_columns
1611                                .iter()
1612                                .chain(&right_context.tag_columns)
1613                                .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)),
1614                    modifier,
1615                    &left_context,
1616                    &right_context,
1617                )?;
1618                let join_plan_schema = join_plan.schema().clone();
1619                let promql_annotations = self.promql_annotations.clone();
1620                // These predicates always pass; they only evaluate otherwise-discarded pairs
1621                // while collecting annotations.
1622                let invalid_pair_predicates = invalid_field_pairs
1623                    .into_iter()
1624                    .filter(|_| promql_annotations.is_some())
1625                    .map(|(left_col_name, right_col_name)| {
1626                        let left_field = join_plan_schema
1627                            .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1628                            .context(DataFusionPlanningSnafu)?;
1629                        let right_field = join_plan_schema
1630                            .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1631                            .context(DataFusionPlanningSnafu)?;
1632                        let left_is_histogram =
1633                            left_field.1.data_type() == &Self::native_histogram_arrow_type();
1634                        let right_is_histogram =
1635                            right_field.1.data_type() == &Self::native_histogram_arrow_type();
1636                        let drop_expr = Self::native_histogram_binary_expr(
1637                            *op,
1638                            DfExpr::Column(left_field.into()),
1639                            left_is_histogram,
1640                            DfExpr::Column(right_field.into()),
1641                            right_is_histogram,
1642                            true,
1643                            promql_annotations.clone(),
1644                        )?
1645                        .with_context(|| UnexpectedPlanExprSnafu {
1646                            desc: "invalid native histogram pair produced no drop expression",
1647                        })?;
1648                        Ok(DfExpr::Not(Box::new(drop_expr)))
1649                    })
1650                    .collect::<Result<Vec<_>>>()?;
1651                let join_plan = if let Some(predicate) = conjunction(invalid_pair_predicates) {
1652                    LogicalPlanBuilder::from(join_plan)
1653                        .filter(predicate)
1654                        .context(DataFusionPlanningSnafu)?
1655                        .build()
1656                        .context(DataFusionPlanningSnafu)?
1657                } else {
1658                    join_plan
1659                };
1660
1661                let bin_expr_builder = |_: &String| {
1662                    let (_, field_pairs) =
1663                        field_groups
1664                            .next()
1665                            .with_context(|| UnexpectedPlanExprSnafu {
1666                                desc: "missing binary field group",
1667                            })?;
1668                    let binary_exprs = field_pairs
1669                        .into_iter()
1670                        .map(|(left_col_name, right_col_name)| {
1671                            let left_field = join_plan_schema
1672                                .qualified_field_with_name(Some(&left_table_ref), left_col_name)
1673                                .context(DataFusionPlanningSnafu)?;
1674                            let right_field = join_plan_schema
1675                                .qualified_field_with_name(Some(&right_table_ref), right_col_name)
1676                                .context(DataFusionPlanningSnafu)?;
1677                            let left_is_histogram =
1678                                left_field.1.data_type() == &Self::native_histogram_arrow_type();
1679                            let right_is_histogram =
1680                                right_field.1.data_type() == &Self::native_histogram_arrow_type();
1681                            let left_col = left_field.into();
1682                            let right_col = right_field.into();
1683
1684                            let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?;
1685                            let lhs = DfExpr::Column(left_col);
1686                            let rhs = DfExpr::Column(right_col);
1687                            let mut binary_expr = match Self::native_histogram_binary_expr(
1688                                *op,
1689                                lhs.clone(),
1690                                left_is_histogram,
1691                                rhs.clone(),
1692                                right_is_histogram,
1693                                is_comparison_op && !should_return_bool,
1694                                promql_annotations.clone(),
1695                            )? {
1696                                Some(expr) => expr,
1697                                None => binary_expr_builder(lhs, rhs)?,
1698                            };
1699                            if is_comparison_op && should_return_bool {
1700                                binary_expr = DfExpr::Cast(Cast {
1701                                    expr: Box::new(binary_expr),
1702                                    data_type: ArrowDataType::Float64,
1703                                });
1704                            }
1705                            Ok(binary_expr)
1706                        })
1707                        .collect::<Result<Vec<_>>>()?;
1708                    if let [binary_expr] = binary_exprs.as_slice() {
1709                        Ok(binary_expr.clone())
1710                    } else {
1711                        Ok(DfExpr::ScalarFunction(ScalarFunction {
1712                            func: coalesce(),
1713                            args: binary_exprs,
1714                        }))
1715                    }
1716                };
1717                if is_comparison_op && !should_return_bool {
1718                    // PromQL comparison operators without `bool` are filters:
1719                    //   - keep the instant-vector side sample values
1720                    //   - drop samples where the comparison is false
1721                    //
1722                    // So we filter on the join result and then project only the side that should
1723                    // be preserved according to PromQL semantics.
1724                    let filtered = self.filter_on_field_column(join_plan, bin_expr_builder)?;
1725                    let (project_table_ref, mut project_context, project_field_columns) =
1726                        match (lhs.value_type(), rhs.value_type()) {
1727                            (ValueType::Scalar, ValueType::Vector) => (
1728                                &right_table_ref,
1729                                right_context.clone(),
1730                                right_aligned_field_columns,
1731                            ),
1732                            _ => (
1733                                &left_table_ref,
1734                                left_context.clone(),
1735                                left_aligned_field_columns,
1736                            ),
1737                        };
1738                    project_context.field_columns = project_field_columns;
1739                    self.project_binary_join_side(filtered, project_table_ref, &project_context)
1740                } else {
1741                    let projected =
1742                        self.projection_for_each_field_column(join_plan, bin_expr_builder)?;
1743                    let preserve_any_value = Self::field_columns_are_alternative_samples(
1744                        projected.schema(),
1745                        &self.ctx.field_columns,
1746                    );
1747                    let retain_field_columns = vec![true; self.ctx.field_columns.len()];
1748                    self.filter_binary_projection(
1749                        projected,
1750                        has_native_histogram,
1751                        preserve_any_value,
1752                        retain_field_columns,
1753                    )
1754                }
1755            }
1756        }
1757    }
1758
1759    fn filter_binary_projection(
1760        &mut self,
1761        input: LogicalPlan,
1762        has_native_histogram: bool,
1763        preserve_any_value: bool,
1764        retain_field_columns: Vec<bool>,
1765    ) -> Result<LogicalPlan> {
1766        if !has_native_histogram {
1767            return Ok(input);
1768        }
1769
1770        ensure!(
1771            retain_field_columns.len() == self.ctx.field_columns.len(),
1772            UnexpectedPlanExprSnafu {
1773                desc: "binary output field count changed unexpectedly",
1774            }
1775        );
1776
1777        let filtered = LogicalPlanBuilder::from(input)
1778            .filter(self.create_empty_values_filter_expr(preserve_any_value)?)
1779            .context(DataFusionPlanningSnafu)?
1780            .build()
1781            .context(DataFusionPlanningSnafu)?;
1782        if retain_field_columns.iter().all(|retain| *retain) {
1783            return Ok(filtered);
1784        }
1785
1786        let retained = self
1787            .ctx
1788            .field_columns
1789            .iter()
1790            .zip(retain_field_columns)
1791            .filter(|(_, retain)| *retain)
1792            .map(|(field, _)| field.clone())
1793            .collect::<Vec<_>>();
1794        if retained.is_empty() {
1795            return Ok(filtered);
1796        }
1797        self.ctx.field_columns = retained;
1798
1799        let mut output_columns = self
1800            .ctx
1801            .field_columns
1802            .iter()
1803            .chain(&self.ctx.tag_columns)
1804            .cloned()
1805            .collect::<HashSet<_>>();
1806        output_columns.extend(self.ctx.time_index_column.iter().cloned());
1807        if self.ctx.use_tsid {
1808            output_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
1809        }
1810        let project_exprs = filtered
1811            .schema()
1812            .iter()
1813            .filter(|(_, field)| output_columns.contains(field.name()))
1814            .map(|(qualifier, field)| {
1815                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1816            })
1817            .collect::<Vec<_>>();
1818        LogicalPlanBuilder::from(filtered)
1819            .project(project_exprs)
1820            .context(DataFusionPlanningSnafu)?
1821            .build()
1822            .context(DataFusionPlanningSnafu)
1823    }
1824
1825    fn project_binary_join_side(
1826        &mut self,
1827        input: LogicalPlan,
1828        table_ref: &TableReference,
1829        context: &PromPlannerContext,
1830    ) -> Result<LogicalPlan> {
1831        let schema = input.schema();
1832
1833        let mut project_exprs =
1834            Vec::with_capacity(context.tag_columns.len() + context.field_columns.len() + 2);
1835
1836        // Project time index from the chosen side.
1837        if let Some(time_index_column) = &context.time_index_column {
1838            let time_index_col = schema
1839                .qualified_field_with_name(Some(table_ref), time_index_column)
1840                .context(DataFusionPlanningSnafu)?
1841                .into();
1842            project_exprs.push(DfExpr::Column(time_index_col));
1843        }
1844
1845        // Project field columns from the chosen side.
1846        for field_column in &context.field_columns {
1847            let field_col = schema
1848                .qualified_field_with_name(Some(table_ref), field_column)
1849                .context(DataFusionPlanningSnafu)?
1850                .into();
1851            project_exprs.push(DfExpr::Column(field_col));
1852        }
1853
1854        // Project tag columns from the chosen side.
1855        for tag_column in &context.tag_columns {
1856            let tag_col = schema
1857                .qualified_field_with_name(Some(table_ref), tag_column)
1858                .context(DataFusionPlanningSnafu)?
1859                .into();
1860            project_exprs.push(DfExpr::Column(tag_col));
1861        }
1862
1863        // Preserve `__tsid` if present, so it can still be used internally downstream. It's
1864        // stripped from the final output anyway.
1865        if let Some(tsid_col) =
1866            Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid)
1867        {
1868            project_exprs.push(tsid_col);
1869        }
1870
1871        let plan = LogicalPlanBuilder::from(input)
1872            .project(project_exprs)
1873            .context(DataFusionPlanningSnafu)?
1874            .build()
1875            .context(DataFusionPlanningSnafu)?;
1876
1877        // Update context to reflect the projected schema. Don't keep a table qualifier since
1878        // the result is a derived expression.
1879        self.ctx = context.clone();
1880        self.ctx.table_name = None;
1881        self.ctx.schema_name = None;
1882
1883        Ok(plan)
1884    }
1885
1886    fn prom_number_lit_to_plan(&mut self, number_literal: &NumberLiteral) -> Result<LogicalPlan> {
1887        let NumberLiteral { val } = number_literal;
1888        self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1889        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1890        self.ctx.reset_table_name_and_schema();
1891        let literal_expr = df_prelude::lit(*val);
1892
1893        let plan = LogicalPlan::Extension(Extension {
1894            node: Arc::new(
1895                EmptyMetric::new(
1896                    self.ctx.start,
1897                    self.ctx.end,
1898                    self.ctx.interval,
1899                    SPECIAL_TIME_FUNCTION.to_string(),
1900                    DEFAULT_FIELD_COLUMN.to_string(),
1901                    Some(literal_expr),
1902                )
1903                .context(DataFusionPlanningSnafu)?,
1904            ),
1905        });
1906        Ok(plan)
1907    }
1908
1909    fn prom_string_lit_to_plan(&mut self, string_literal: &StringLiteral) -> Result<LogicalPlan> {
1910        let StringLiteral { val } = string_literal;
1911        self.ctx.time_index_column = Some(DEFAULT_TIME_INDEX_COLUMN.to_string());
1912        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
1913        self.ctx.reset_table_name_and_schema();
1914        let literal_expr = df_prelude::lit(val.clone());
1915
1916        let plan = LogicalPlan::Extension(Extension {
1917            node: Arc::new(
1918                EmptyMetric::new(
1919                    self.ctx.start,
1920                    self.ctx.end,
1921                    self.ctx.interval,
1922                    SPECIAL_TIME_FUNCTION.to_string(),
1923                    DEFAULT_FIELD_COLUMN.to_string(),
1924                    Some(literal_expr),
1925                )
1926                .context(DataFusionPlanningSnafu)?,
1927            ),
1928        });
1929        Ok(plan)
1930    }
1931
1932    async fn prom_vector_selector_to_plan(
1933        &mut self,
1934        vector_selector: &VectorSelector,
1935        timestamp_fn: bool,
1936    ) -> Result<LogicalPlan> {
1937        let VectorSelector {
1938            name,
1939            offset,
1940            matchers,
1941            at: _,
1942        } = vector_selector;
1943        let matchers = self.preprocess_label_matchers(matchers, name)?;
1944        if let Some(empty_plan) = self.setup_context().await? {
1945            return Ok(empty_plan);
1946        }
1947        let offset_ms = match offset {
1948            Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
1949            Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
1950            None => 0,
1951        };
1952        let normalize = self
1953            .selector_to_series_normalize_plan(offset, matchers, false)
1954            .await?;
1955        let time_index_column =
1956            self.ctx
1957                .time_index_column
1958                .clone()
1959                .with_context(|| TimeIndexNotFoundSnafu {
1960                    table: self.ctx.table_name.clone().unwrap_or_default(),
1961                })?;
1962
1963        let (normalize, timestamp_value_column) = if timestamp_fn {
1964            // Keep the original sample for stale-marker detection while carrying
1965            // its timestamp through InstantManipulate in a private value column.
1966            let occupied = normalize
1967                .schema()
1968                .fields()
1969                .iter()
1970                .map(|field| field.name().as_str())
1971                .collect::<HashSet<_>>();
1972            let mut timestamp_value_column = TIMESTAMP_VALUE_PREFIX.to_string();
1973            while occupied.contains(timestamp_value_column.as_str()) {
1974                timestamp_value_column.push('_');
1975            }
1976            let mut project_exprs = normalize
1977                .schema()
1978                .iter()
1979                .map(|(qualifier, field)| {
1980                    DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
1981                })
1982                .collect::<Vec<_>>();
1983            // `timestamp()` preserves the shifted selector timeline even though
1984            // SeriesNormalize now retains raw native timestamp storage. Decimal
1985            // arithmetic shifts before truncating to milliseconds.
1986            let unit_factor = match col(&time_index_column)
1987                .get_type(normalize.schema())
1988                .context(DataFusionPlanningSnafu)?
1989            {
1990                ArrowDataType::Timestamp(ArrowTimeUnit::Second, _) => (1_000_i128, 4, 0),
1991                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, _) => (1, 1, 0),
1992                ArrowDataType::Timestamp(ArrowTimeUnit::Microsecond, _) => (1, 4, 3),
1993                ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, _) => (1, 7, 6),
1994                _ => unreachable!("time index is a timestamp"),
1995            };
1996            let sample_time = col(&time_index_column)
1997                .cast_to(&ArrowDataType::Int64, normalize.schema())
1998                .context(DataFusionPlanningSnafu)?
1999                .cast_to(&ArrowDataType::Decimal128(19, 0), normalize.schema())
2000                .context(DataFusionPlanningSnafu)?;
2001            let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2002                left: Box::new(sample_time),
2003                op: Operator::Multiply,
2004                right: Box::new(lit(ScalarValue::Decimal128(
2005                    Some(unit_factor.0),
2006                    unit_factor.1,
2007                    unit_factor.2,
2008                ))),
2009            });
2010            let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2011                left: Box::new(sample_time),
2012                op: Operator::Plus,
2013                right: Box::new(lit(ScalarValue::Decimal128(Some(offset_ms as i128), 19, 0))),
2014            })
2015            .cast_to(&ArrowDataType::Int64, normalize.schema())
2016            .context(DataFusionPlanningSnafu)?
2017            .cast_to(&ArrowDataType::Float64, normalize.schema())
2018            .context(DataFusionPlanningSnafu)?;
2019            let sample_time = DfExpr::BinaryExpr(BinaryExpr {
2020                left: Box::new(sample_time),
2021                op: Operator::Divide,
2022                right: Box::new(lit(1000.0)),
2023            });
2024            project_exprs.push(sample_time.alias(&timestamp_value_column));
2025            let normalize = LogicalPlanBuilder::from(normalize)
2026                .project(project_exprs)
2027                .context(DataFusionPlanningSnafu)?
2028                .build()
2029                .context(DataFusionPlanningSnafu)?;
2030            (normalize, Some(timestamp_value_column))
2031        } else {
2032            (normalize, None)
2033        };
2034
2035        let field_column = self.ctx.field_columns.first().cloned();
2036        let manipulate = InstantManipulate::new(
2037            self.ctx.start,
2038            self.ctx.end,
2039            self.ctx.lookback_delta,
2040            self.ctx.interval,
2041            offset_ms,
2042            time_index_column,
2043            if self.ctx.use_tsid {
2044                vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
2045            } else {
2046                self.ctx.tag_columns.clone()
2047            },
2048            field_column,
2049            normalize,
2050        );
2051        let manipulate = LogicalPlan::Extension(Extension {
2052            node: Arc::new(manipulate),
2053        });
2054        if let Some(timestamp_value_column) = timestamp_value_column {
2055            self.create_timestamp_func_plan(manipulate, &timestamp_value_column)
2056        } else {
2057            Ok(manipulate)
2058        }
2059    }
2060
2061    /// Builds a projection plan for the PromQL `timestamp()` function.
2062    /// Projects the time index column as the value column for each row.
2063    ///
2064    /// # Arguments
2065    /// * `input` - Input [`LogicalPlan`] after instant-vector selection.
2066    /// * `timestamp_value_column` - Private column containing each selected sample's timestamp.
2067    ///
2068    /// # Returns
2069    /// Returns a [`Result<LogicalPlan>`] where the resulting logical plan projects the timestamp
2070    /// column as the value column, along with the original tag and time index columns.
2071    ///
2072    /// # Timestamp vs. Time Function
2073    ///
2074    /// - **Timestamp Function (`timestamp()`)**: In PromQL, the `timestamp()` function returns the
2075    ///   timestamp (time index) of each sample as the value column.
2076    ///
2077    /// - **Time Function (`time()`)**: The `time()` function returns the evaluation time of the query
2078    ///   as a scalar value.
2079    ///
2080    /// # Side Effects
2081    /// Updates the planner context's field columns to the timestamp column name.
2082    ///
2083    fn create_timestamp_func_plan(
2084        &mut self,
2085        input: LogicalPlan,
2086        timestamp_value_column: &str,
2087    ) -> Result<LogicalPlan> {
2088        let time_expr = col(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN);
2089        self.ctx.field_columns = vec![time_expr.schema_name().to_string()];
2090        let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2);
2091        project_exprs.push(self.create_time_index_column_expr()?);
2092        project_exprs.push(time_expr);
2093        project_exprs.extend(self.create_tag_column_exprs()?);
2094
2095        LogicalPlanBuilder::from(input)
2096            .project(project_exprs)
2097            .context(DataFusionPlanningSnafu)?
2098            .build()
2099            .context(DataFusionPlanningSnafu)
2100    }
2101
2102    async fn prom_matrix_selector_to_plan(
2103        &mut self,
2104        matrix_selector: &MatrixSelector,
2105    ) -> Result<LogicalPlan> {
2106        let MatrixSelector { vs, range } = matrix_selector;
2107        let VectorSelector {
2108            name,
2109            offset,
2110            matchers,
2111            ..
2112        } = vs;
2113        let matchers = self.preprocess_label_matchers(matchers, name)?;
2114        ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
2115        let range_ms = range.as_millis() as _;
2116        self.ctx.range = Some(range_ms);
2117        let offset_ms = match offset {
2118            Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2119            Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2120            None => 0,
2121        };
2122
2123        // Some functions like rate may require special fields in the RangeManipulate plan
2124        // so we can't skip RangeManipulate.
2125        let normalize = match self.setup_context().await? {
2126            Some(empty_plan) => empty_plan,
2127            None => {
2128                self.selector_to_series_normalize_plan(offset, matchers, true)
2129                    .await?
2130            }
2131        };
2132        let manipulate = RangeManipulate::new(
2133            self.ctx.start,
2134            self.ctx.end,
2135            self.ctx.interval,
2136            offset_ms,
2137            // TODO(ruihang): convert via Timestamp datatypes to support different time units
2138            range_ms,
2139            self.ctx
2140                .time_index_column
2141                .clone()
2142                .expect("time index should be set in `setup_context`"),
2143            self.ctx.field_columns.clone(),
2144            normalize,
2145        )
2146        .context(DataFusionPlanningSnafu)?;
2147
2148        Ok(LogicalPlan::Extension(Extension {
2149            node: Arc::new(manipulate),
2150        }))
2151    }
2152
2153    async fn prom_call_expr_to_plan(
2154        &mut self,
2155        query_engine_state: &QueryEngineState,
2156        call_expr: &Call,
2157    ) -> Result<LogicalPlan> {
2158        let Call { func, args } = call_expr;
2159        // some special functions that are not expression but a plan
2160        match func.name {
2161            SPECIAL_HISTOGRAM_QUANTILE | SPECIAL_HISTOGRAM_FRACTION => {
2162                return self
2163                    .create_histogram_plan(func.name, args, query_engine_state)
2164                    .await;
2165            }
2166            SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await,
2167            SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await,
2168            SPECIAL_ABSENT_FUNCTION => {
2169                return self.create_absent_plan(args, query_engine_state).await;
2170            }
2171            _ => {}
2172        }
2173
2174        // transform function arguments
2175        let args = self.create_function_args(&args.args)?;
2176        let input = if let Some(prom_expr) = &args.input {
2177            self.prom_expr_to_plan_inner(prom_expr, func.name == "timestamp", query_engine_state)
2178                .await?
2179        } else {
2180            self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
2181            self.ctx.reset_table_name_and_schema();
2182            self.ctx.tag_columns = vec![];
2183            self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
2184            LogicalPlan::Extension(Extension {
2185                node: Arc::new(
2186                    EmptyMetric::new(
2187                        self.ctx.start,
2188                        self.ctx.end,
2189                        self.ctx.interval,
2190                        SPECIAL_TIME_FUNCTION.to_string(),
2191                        DEFAULT_FIELD_COLUMN.to_string(),
2192                        None,
2193                    )
2194                    .context(DataFusionPlanningSnafu)?,
2195                ),
2196            })
2197        };
2198        let (mut func_exprs, new_tags) = self.create_function_expr(
2199            func,
2200            args.literals.clone(),
2201            input.schema(),
2202            query_engine_state,
2203        )?;
2204        func_exprs.insert(0, self.create_time_index_column_expr()?);
2205        func_exprs.extend_from_slice(&self.create_tag_column_exprs()?);
2206        if let Some(tsid_col) =
2207            Self::optional_tsid_projection(input.schema(), None, self.ctx.use_tsid)
2208        {
2209            func_exprs.push(tsid_col);
2210        }
2211
2212        // A row survives as long as one field column produced a sample, and the fields without
2213        // one stay NULL, which is the shape a selector already emits. Requiring every field to
2214        // be non-NULL would drop one field's samples because another field has none in the same
2215        // window — the reason alternative float/histogram columns already needed this form. A
2216        // single field column reduces to the same predicate either way.
2217        let builder = LogicalPlanBuilder::from(input)
2218            .project(func_exprs)
2219            .context(DataFusionPlanningSnafu)?
2220            .filter(self.create_empty_values_filter_expr(true)?)
2221            .context(DataFusionPlanningSnafu)?;
2222
2223        let builder = match func.name {
2224            "sort" => builder
2225                .sort(self.create_field_columns_sort_exprs(true))
2226                .context(DataFusionPlanningSnafu)?,
2227            "sort_desc" => builder
2228                .sort(self.create_field_columns_sort_exprs(false))
2229                .context(DataFusionPlanningSnafu)?,
2230            "sort_by_label" => builder
2231                .sort(Self::create_sort_exprs_by_tags(
2232                    func.name,
2233                    args.literals,
2234                    true,
2235                )?)
2236                .context(DataFusionPlanningSnafu)?,
2237            "sort_by_label_desc" => builder
2238                .sort(Self::create_sort_exprs_by_tags(
2239                    func.name,
2240                    args.literals,
2241                    false,
2242                )?)
2243                .context(DataFusionPlanningSnafu)?,
2244
2245            _ => builder,
2246        };
2247
2248        // Update context tags after building plan
2249        // We can't push them before planning, because they won't exist until projection.
2250        for tag in new_tags {
2251            self.ctx.tag_columns.push(tag);
2252        }
2253
2254        let plan = builder.build().context(DataFusionPlanningSnafu)?;
2255        common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}");
2256
2257        Ok(plan)
2258    }
2259
2260    async fn prom_ext_expr_to_plan(
2261        &mut self,
2262        query_engine_state: &QueryEngineState,
2263        ext_expr: &promql_parser::parser::ast::Extension,
2264    ) -> Result<LogicalPlan> {
2265        // let promql_parser::parser::ast::Extension { expr } = ext_expr;
2266        let expr = &ext_expr.expr;
2267        let children = expr.children();
2268        let plan = self
2269            .prom_expr_to_plan(&children[0], query_engine_state)
2270            .await?;
2271        // Wrapper for the explanation/analyze of the existing plan
2272        // https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/builder/struct.LogicalPlanBuilder.html#method.explain
2273        // if `analyze` is true, runs the actual plan and produces
2274        // information about metrics during run.
2275        // if `verbose` is true, prints out additional details when VERBOSE keyword is specified
2276        match expr.name() {
2277            ANALYZE_NODE_NAME => LogicalPlanBuilder::from(plan)
2278                .explain(false, true)
2279                .unwrap()
2280                .build()
2281                .context(DataFusionPlanningSnafu),
2282            ANALYZE_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2283                .explain(true, true)
2284                .unwrap()
2285                .build()
2286                .context(DataFusionPlanningSnafu),
2287            EXPLAIN_NODE_NAME => LogicalPlanBuilder::from(plan)
2288                .explain(false, false)
2289                .unwrap()
2290                .build()
2291                .context(DataFusionPlanningSnafu),
2292            EXPLAIN_VERBOSE_NODE_NAME => LogicalPlanBuilder::from(plan)
2293                .explain(true, false)
2294                .unwrap()
2295                .build()
2296                .context(DataFusionPlanningSnafu),
2297            ALIAS_NODE_NAME => {
2298                let alias = expr
2299                    .as_any()
2300                    .downcast_ref::<AliasExpr>()
2301                    .context(UnexpectedPlanExprSnafu {
2302                        desc: "Expected AliasExpr",
2303                    })?
2304                    .alias
2305                    .clone();
2306                self.apply_alias(plan, alias)
2307            }
2308            _ => LogicalPlanBuilder::empty(true)
2309                .build()
2310                .context(DataFusionPlanningSnafu),
2311        }
2312    }
2313
2314    /// Extract metric name from `__name__` matcher and set it into [PromPlannerContext].
2315    /// Returns a new [Matchers] that doesn't contain metric name matcher.
2316    ///
2317    /// Each call to this function means new selector is started. Thus, the context will be reset
2318    /// at first.
2319    ///
2320    /// Name rule:
2321    /// - if `name` is some, then the matchers MUST NOT contain `__name__` matcher.
2322    /// - if `name` is none, then the matchers MAY contain NONE OR MULTIPLE `__name__` matchers.
2323    #[allow(clippy::mutable_key_type)]
2324    fn preprocess_label_matchers(
2325        &mut self,
2326        label_matchers: &Matchers,
2327        name: &Option<String>,
2328    ) -> Result<Matchers> {
2329        self.ctx.reset();
2330
2331        let metric_name;
2332        if let Some(name) = name.clone() {
2333            metric_name = Some(name);
2334            ensure!(
2335                label_matchers.find_matchers(METRIC_NAME).is_empty(),
2336                MultipleMetricMatchersSnafu
2337            );
2338        } else {
2339            let mut matches = label_matchers.find_matchers(METRIC_NAME);
2340            ensure!(!matches.is_empty(), NoMetricMatcherSnafu);
2341            ensure!(matches.len() == 1, MultipleMetricMatchersSnafu);
2342            ensure!(
2343                matches[0].op == MatchOp::Equal,
2344                UnsupportedMatcherOpSnafu {
2345                    matcher_op: matches[0].op.to_string(),
2346                    matcher: METRIC_NAME
2347                }
2348            );
2349            metric_name = matches.pop().map(|m| m.value);
2350        }
2351
2352        self.ctx.table_name = metric_name;
2353
2354        let mut matchers = HashSet::new();
2355        for matcher in &label_matchers.matchers {
2356            // TODO(ruihang): support other metric match ops
2357            if matcher.name == FIELD_COLUMN_MATCHER {
2358                self.ctx
2359                    .field_column_matcher
2360                    .get_or_insert_default()
2361                    .push(matcher.clone());
2362            } else if matcher.name == SCHEMA_COLUMN_MATCHER || matcher.name == DB_COLUMN_MATCHER {
2363                ensure!(
2364                    matcher.op == MatchOp::Equal,
2365                    UnsupportedMatcherOpSnafu {
2366                        matcher: matcher.name.clone(),
2367                        matcher_op: matcher.op.to_string(),
2368                    }
2369                );
2370                self.ctx.schema_name = Some(matcher.value.clone());
2371            } else if matcher.name != METRIC_NAME {
2372                self.ctx.selector_matcher.push(matcher.clone());
2373                let _ = matchers.insert(matcher.clone());
2374            }
2375        }
2376
2377        Ok(Matchers::new(matchers.into_iter().collect()))
2378    }
2379
2380    async fn selector_to_series_normalize_plan(
2381        &mut self,
2382        offset: &Option<Offset>,
2383        label_matchers: Matchers,
2384        is_range_selector: bool,
2385    ) -> Result<LogicalPlan> {
2386        // make table scan plan
2387        let table_ref = self.table_ref()?;
2388        let mut table_scan = self.create_table_scan_plan(table_ref.clone()).await?;
2389        let table_schema = table_scan.schema();
2390
2391        // make filter exprs
2392        let offset_duration = match offset {
2393            Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond,
2394            Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond),
2395            None => 0,
2396        };
2397        let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?;
2398        if let Some(time_index_filter) =
2399            self.build_time_index_filter(offset_duration, table_schema)?
2400        {
2401            scan_filters.push(time_index_filter);
2402        }
2403        if let Some(filter) = conjunction(scan_filters) {
2404            table_scan = LogicalPlanBuilder::from(table_scan)
2405                .filter(filter)
2406                .context(DataFusionPlanningSnafu)?
2407                .build()
2408                .context(DataFusionPlanningSnafu)?;
2409        }
2410
2411        // make a projection plan if there is any `__field__` matcher
2412        if let Some(field_matchers) = &self.ctx.field_column_matcher {
2413            let col_set = self.ctx.field_columns.iter().collect::<HashSet<_>>();
2414            // opt-in set
2415            let mut result_set = HashSet::new();
2416            // opt-out set
2417            let mut reverse_set = HashSet::new();
2418            for matcher in field_matchers {
2419                match &matcher.op {
2420                    MatchOp::Equal => {
2421                        if col_set.contains(&matcher.value) {
2422                            let _ = result_set.insert(matcher.value.clone());
2423                        } else {
2424                            return Err(ColumnNotFoundSnafu {
2425                                col: matcher.value.clone(),
2426                            }
2427                            .build());
2428                        }
2429                    }
2430                    MatchOp::NotEqual => {
2431                        if col_set.contains(&matcher.value) {
2432                            let _ = reverse_set.insert(matcher.value.clone());
2433                        } else {
2434                            return Err(ColumnNotFoundSnafu {
2435                                col: matcher.value.clone(),
2436                            }
2437                            .build());
2438                        }
2439                    }
2440                    MatchOp::Re(regex) => {
2441                        for col in &self.ctx.field_columns {
2442                            if regex.is_match(col) {
2443                                let _ = result_set.insert(col.clone());
2444                            }
2445                        }
2446                    }
2447                    MatchOp::NotRe(regex) => {
2448                        for col in &self.ctx.field_columns {
2449                            if regex.is_match(col) {
2450                                let _ = reverse_set.insert(col.clone());
2451                            }
2452                        }
2453                    }
2454                }
2455            }
2456            // merge two set
2457            if result_set.is_empty() {
2458                result_set = col_set.into_iter().cloned().collect();
2459            }
2460            for col in reverse_set {
2461                let _ = result_set.remove(&col);
2462            }
2463
2464            // mask the field columns in context using computed result set
2465            self.ctx.field_columns = self
2466                .ctx
2467                .field_columns
2468                .drain(..)
2469                .filter(|col| result_set.contains(col))
2470                .collect();
2471
2472            let exprs = result_set
2473                .into_iter()
2474                .map(|col| DfExpr::Column(Column::new_unqualified(col)))
2475                .chain(self.create_tag_column_exprs()?)
2476                .chain(
2477                    self.ctx
2478                        .use_tsid
2479                        .then_some(DfExpr::Column(Column::new_unqualified(
2480                            DATA_SCHEMA_TSID_COLUMN_NAME,
2481                        ))),
2482                )
2483                .chain(Some(self.create_time_index_column_expr()?))
2484                .collect::<Vec<_>>();
2485
2486            // reuse this variable for simplicity
2487            table_scan = LogicalPlanBuilder::from(table_scan)
2488                .project(exprs)
2489                .context(DataFusionPlanningSnafu)?
2490                .build()
2491                .context(DataFusionPlanningSnafu)?;
2492        }
2493
2494        // make sort plan
2495        let series_key_columns = if self.ctx.use_tsid {
2496            vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]
2497        } else {
2498            self.ctx.tag_columns.clone()
2499        };
2500
2501        let sort_exprs = if self.ctx.use_tsid {
2502            vec![
2503                DfExpr::Column(Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME)).sort(true, true),
2504                self.create_time_index_column_expr()?.sort(true, true),
2505            ]
2506        } else {
2507            self.create_tag_and_time_index_column_sort_exprs()?
2508        };
2509
2510        let sort_plan = LogicalPlanBuilder::from(table_scan)
2511            .sort(sort_exprs)
2512            .context(DataFusionPlanningSnafu)?
2513            .build()
2514            .context(DataFusionPlanningSnafu)?;
2515
2516        // make divide plan
2517        let time_index_column =
2518            self.ctx
2519                .time_index_column
2520                .clone()
2521                .with_context(|| TimeIndexNotFoundSnafu {
2522                    table: table_ref.to_string(),
2523                })?;
2524        let divide_plan = LogicalPlan::Extension(Extension {
2525            node: Arc::new(SeriesDivide::new(
2526                series_key_columns.clone(),
2527                time_index_column,
2528                sort_plan,
2529            )),
2530        });
2531
2532        // make series_normalize plan
2533        if !is_range_selector && offset_duration == 0 {
2534            return Ok(divide_plan);
2535        }
2536        let series_normalize = SeriesNormalize::new(
2537            offset_duration,
2538            self.ctx
2539                .time_index_column
2540                .clone()
2541                .with_context(|| TimeIndexNotFoundSnafu {
2542                    table: table_ref.to_quoted_string(),
2543                })?,
2544            is_range_selector,
2545            series_key_columns,
2546            divide_plan,
2547        );
2548        let logical_plan = LogicalPlan::Extension(Extension {
2549            node: Arc::new(series_normalize),
2550        });
2551
2552        Ok(logical_plan)
2553    }
2554
2555    /// Convert [LabelModifier] to [Column] exprs for aggregation.
2556    /// Timestamp column and tag columns will be included.
2557    ///
2558    /// # Side effect
2559    ///
2560    /// This method will also change the tag columns in ctx if `update_ctx` is true.
2561    fn agg_modifier_to_col(
2562        &mut self,
2563        input_schema: &DFSchemaRef,
2564        modifier: &Option<LabelModifier>,
2565        update_ctx: bool,
2566    ) -> Result<Vec<DfExpr>> {
2567        match modifier {
2568            None => {
2569                if update_ctx {
2570                    self.ctx.tag_columns.clear();
2571                }
2572                Ok(vec![self.create_time_index_column_expr()?])
2573            }
2574            Some(LabelModifier::Include(labels)) => {
2575                if update_ctx {
2576                    self.ctx.tag_columns.clear();
2577                }
2578                let mut exprs = Vec::with_capacity(labels.labels.len());
2579                for label in &labels.labels {
2580                    if is_metric_engine_internal_column(label) {
2581                        continue;
2582                    }
2583                    // nonexistence label will be ignored
2584                    if let Some(column_name) = Self::find_case_sensitive_column(input_schema, label)
2585                    {
2586                        exprs.push(DfExpr::Column(Column::from_name(column_name.clone())));
2587
2588                        if update_ctx {
2589                            // update the tag columns in context
2590                            self.ctx.tag_columns.push(column_name);
2591                        }
2592                    }
2593                }
2594                // add timestamp column
2595                exprs.push(self.create_time_index_column_expr()?);
2596
2597                Ok(exprs)
2598            }
2599            Some(LabelModifier::Exclude(labels)) => {
2600                let mut all_fields = input_schema
2601                    .fields()
2602                    .iter()
2603                    .map(|f| f.name())
2604                    .collect::<BTreeSet<_>>();
2605
2606                // Exclude metric engine internal columns (not PromQL labels) from the implicit
2607                // "without" label set.
2608                all_fields.retain(|col| !is_metric_engine_internal_column(col.as_str()));
2609
2610                // remove "without"-ed fields
2611                // nonexistence label will be ignored
2612                for label in &labels.labels {
2613                    let _ = all_fields.remove(label);
2614                }
2615
2616                // remove time index and value fields
2617                if let Some(time_index) = &self.ctx.time_index_column {
2618                    let _ = all_fields.remove(time_index);
2619                }
2620                for value in &self.ctx.field_columns {
2621                    let _ = all_fields.remove(value);
2622                }
2623
2624                if update_ctx {
2625                    // change the tag columns in context
2626                    self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect();
2627                }
2628
2629                // collect remaining fields and convert to col expr
2630                let mut exprs = all_fields
2631                    .into_iter()
2632                    .map(|c| DfExpr::Column(Column::from(c)))
2633                    .collect::<Vec<_>>();
2634
2635                // add timestamp column
2636                exprs.push(self.create_time_index_column_expr()?);
2637
2638                Ok(exprs)
2639            }
2640        }
2641    }
2642
2643    // TODO(ruihang): ignore `MetricNameLabel` (`__name__`) matcher
2644    pub fn matchers_to_expr(
2645        label_matchers: Matchers,
2646        table_schema: &DFSchemaRef,
2647    ) -> Result<Vec<DfExpr>> {
2648        let mut exprs = Vec::with_capacity(label_matchers.matchers.len());
2649        for matcher in label_matchers.matchers {
2650            if matcher.name == SCHEMA_COLUMN_MATCHER
2651                || matcher.name == DB_COLUMN_MATCHER
2652                || matcher.name == FIELD_COLUMN_MATCHER
2653            {
2654                continue;
2655            }
2656
2657            let accepts_empty = matcher.is_match("");
2658            let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str());
2659            let col = if let Some(column_name) = column_name {
2660                let column = DfExpr::Column(Column::from_name(&column_name));
2661                let field = table_schema
2662                    .index_of_column_by_name(None, &column_name)
2663                    .map(|index| table_schema.field(index));
2664                if accepts_empty
2665                    && column_name == OTLP_AGGREGATION_TEMPORALITY_LABEL
2666                    && let Some(data_type) = field
2667                        .filter(|field| {
2668                            field.is_nullable()
2669                                && Self::string_value_data_type(field.data_type()).is_some()
2670                        })
2671                        .map(|field| field.data_type())
2672                {
2673                    let empty = Self::string_scalar_value(data_type, Some(String::new()))
2674                        .expect("nullable label has a string type");
2675                    DfExpr::ScalarFunction(ScalarFunction {
2676                        func: coalesce(),
2677                        args: vec![column, DfExpr::Literal(empty, None)],
2678                    })
2679                } else {
2680                    column
2681                }
2682            } else {
2683                DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
2684                    .alias(matcher.name.clone())
2685            };
2686            let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None);
2687            let expr = match matcher.op {
2688                MatchOp::Equal => col.eq(lit),
2689                MatchOp::NotEqual => col.not_eq(lit),
2690                MatchOp::Re(re) => {
2691                    // TODO(ruihang): a more programmatic way to handle this in datafusion
2692
2693                    // This is a hack to handle `.+` and `.*`, and is not strictly correct
2694                    // `.` doesn't match newline (`\n`). Given this is in PromQL context,
2695                    // most of the time it's fine.
2696                    if re.as_str() == "^(?:.*)$" {
2697                        continue;
2698                    }
2699                    if re.as_str() == "^(?:.+)$" {
2700                        col.not_eq(DfExpr::Literal(
2701                            ScalarValue::Utf8(Some(String::new())),
2702                            None,
2703                        ))
2704                    } else {
2705                        DfExpr::BinaryExpr(BinaryExpr {
2706                            left: Box::new(col),
2707                            op: Operator::RegexMatch,
2708                            right: Box::new(DfExpr::Literal(
2709                                ScalarValue::Utf8(Some(re.as_str().to_string())),
2710                                None,
2711                            )),
2712                        })
2713                    }
2714                }
2715                MatchOp::NotRe(re) => {
2716                    if re.as_str() == "^(?:.*)$" {
2717                        DfExpr::Literal(ScalarValue::Boolean(Some(false)), None)
2718                    } else if re.as_str() == "^(?:.+)$" {
2719                        col.eq(DfExpr::Literal(
2720                            ScalarValue::Utf8(Some(String::new())),
2721                            None,
2722                        ))
2723                    } else {
2724                        DfExpr::BinaryExpr(BinaryExpr {
2725                            left: Box::new(col),
2726                            op: Operator::RegexNotMatch,
2727                            right: Box::new(DfExpr::Literal(
2728                                ScalarValue::Utf8(Some(re.as_str().to_string())),
2729                                None,
2730                            )),
2731                        })
2732                    }
2733                }
2734            };
2735            exprs.push(expr);
2736        }
2737
2738        Ok(exprs)
2739    }
2740
2741    fn find_case_sensitive_column(schema: &DFSchemaRef, column: &str) -> Option<String> {
2742        if is_metric_engine_internal_column(column) {
2743            return None;
2744        }
2745        schema
2746            .fields()
2747            .iter()
2748            .find(|field| field.name() == column)
2749            .map(|field| field.name().clone())
2750    }
2751
2752    fn table_from_source(&self, source: &Arc<dyn TableSource>) -> Result<table::TableRef> {
2753        Ok(source
2754            .as_any()
2755            .downcast_ref::<DefaultTableSource>()
2756            .context(UnknownTableSnafu)?
2757            .table_provider
2758            .as_any()
2759            .downcast_ref::<DfTableProviderAdapter>()
2760            .context(UnknownTableSnafu)?
2761            .table())
2762    }
2763
2764    fn table_ref(&self) -> Result<TableReference> {
2765        let table_name = self
2766            .ctx
2767            .table_name
2768            .clone()
2769            .context(TableNameNotFoundSnafu)?;
2770
2771        // set schema name if `__schema__` is given
2772        let table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2773            TableReference::partial(schema_name.as_str(), table_name.as_str())
2774        } else {
2775            TableReference::bare(table_name.as_str())
2776        };
2777
2778        Ok(table_ref)
2779    }
2780
2781    fn build_time_index_filter(
2782        &self,
2783        offset_duration: i64,
2784        schema: &DFSchemaRef,
2785    ) -> Result<Option<DfExpr>> {
2786        let start = self.ctx.start;
2787        let end = self.ctx.end;
2788        if end < start {
2789            return InvalidTimeRangeSnafu { start, end }.fail();
2790        }
2791        let time_index_expr = self.create_time_index_column_expr()?;
2792        let time_index_name = self.ctx.time_index_column.as_ref().unwrap();
2793        let unit = schema
2794            .index_of_column_by_name(None, time_index_name)
2795            .and_then(|index| match schema.field(index).data_type() {
2796                ArrowDataType::Timestamp(unit, _) => Some(*unit),
2797                _ => None,
2798            })
2799            .unwrap_or(ArrowTimeUnit::Millisecond);
2800        let native_value = |milliseconds: i128| match unit {
2801            ArrowTimeUnit::Second => milliseconds.div_euclid(1_000),
2802            ArrowTimeUnit::Millisecond => milliseconds,
2803            ArrowTimeUnit::Microsecond => milliseconds * 1_000,
2804            ArrowTimeUnit::Nanosecond => milliseconds * 1_000_000,
2805        };
2806        let scalar = |milliseconds: i128| -> Option<ScalarValue> {
2807            let value = i64::try_from(native_value(milliseconds)).ok()?;
2808            Some(match unit {
2809                ArrowTimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None),
2810                ArrowTimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None),
2811                ArrowTimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None),
2812                ArrowTimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None),
2813            })
2814        };
2815        let window = self.ctx.range.unwrap_or(self.ctx.lookback_delta);
2816        let filter = |lower_ms: i128, upper_ms: i128| {
2817            let lower_value = native_value(lower_ms);
2818            let upper_value = native_value(upper_ms);
2819            if lower_value > i128::from(i64::MAX) || upper_value < i128::from(i64::MIN) {
2820                return Some(lit(false));
2821            }
2822            let lower_filter = (lower_value >= i128::from(i64::MIN)).then(|| {
2823                let lower = DfExpr::Literal(scalar(lower_ms).unwrap(), None);
2824                if window == 0 {
2825                    time_index_expr.clone().gt_eq(lower)
2826                } else if unit == ArrowTimeUnit::Millisecond
2827                    && let Some(inclusive_lower) = lower_ms
2828                        .checked_add(1)
2829                        .and_then(|lower| i64::try_from(lower).ok())
2830                        .and_then(|lower| scalar(i128::from(lower)))
2831                {
2832                    time_index_expr
2833                        .clone()
2834                        .gt_eq(DfExpr::Literal(inclusive_lower, None))
2835                } else {
2836                    time_index_expr.clone().gt(lower)
2837                }
2838            });
2839            let upper_filter = (upper_value <= i128::from(i64::MAX)).then(|| {
2840                time_index_expr
2841                    .clone()
2842                    .lt_eq(DfExpr::Literal(scalar(upper_ms).unwrap(), None))
2843            });
2844
2845            // An underflowing lower bound must not discard a representable upper
2846            // bound: without it, LastRow could retain a future row and discard the
2847            // older eligible sample before the manipulator can check its time.
2848            match (lower_filter, upper_filter) {
2849                (Some(lower), Some(upper)) => Some(lower.and(upper)),
2850                (Some(filter), None) | (None, Some(filter)) => Some(filter),
2851                (None, None) => None,
2852            }
2853        };
2854        let bounds = |timestamp: i64| {
2855            let upper = i128::from(timestamp) - i128::from(offset_duration);
2856            (upper - i128::from(window), upper)
2857        };
2858        let num_points = (end as i128 - start as i128) / self.ctx.interval as i128;
2859        if num_points > MAX_SCATTER_POINTS as i128 || self.ctx.interval <= INTERVAL_1H {
2860            let (lower, _) = bounds(start);
2861            let (_, upper) = bounds(end);
2862            return Ok(filter(lower, upper));
2863        }
2864        let mut filters = Vec::new();
2865        for timestamp in (start..=end).step_by(self.ctx.interval as usize) {
2866            let (lower, upper) = bounds(timestamp);
2867            let Some(filter) = filter(lower, upper) else {
2868                // A point whose native bounds cannot be represented may cover the whole native
2869                // time domain, so its disjunct cannot be omitted.
2870                return Ok(None);
2871            };
2872            filters.push(filter);
2873        }
2874        Ok(filters.into_iter().reduce(DfExpr::or))
2875    }
2876
2877    /// Create a table scan plan and a filter plan with given filter.
2878    ///
2879    /// # Panic
2880    /// If the filter is empty
2881    async fn create_table_scan_plan(&mut self, table_ref: TableReference) -> Result<LogicalPlan> {
2882        let provider = self
2883            .table_provider
2884            .resolve_table(table_ref.clone())
2885            .await
2886            .context(CatalogSnafu)?;
2887
2888        let logical_table = self.table_from_source(&provider)?;
2889
2890        // Try to rewrite the table scan to physical table scan if possible.
2891        let mut maybe_phy_table_ref = table_ref.clone();
2892        let mut scan_provider = provider;
2893        let mut table_id_filter: Option<u32> = None;
2894
2895        // If it's a metric engine logical table, scan its physical table directly and filter by
2896        // `__table_id = logical_table_id` to get access to internal columns like `__tsid`.
2897        if logical_table.table_info().meta.engine == METRIC_ENGINE_NAME
2898            && let Some(physical_table_name) = logical_table
2899                .table_info()
2900                .meta
2901                .options
2902                .extra_options
2903                .get(LOGICAL_TABLE_METADATA_KEY)
2904        {
2905            let physical_table_ref = if let Some(schema_name) = &self.ctx.schema_name {
2906                TableReference::partial(schema_name.as_str(), physical_table_name.as_str())
2907            } else {
2908                TableReference::bare(physical_table_name.as_str())
2909            };
2910
2911            let physical_provider = match self
2912                .table_provider
2913                .resolve_table(physical_table_ref.clone())
2914                .await
2915            {
2916                Ok(provider) => provider,
2917                Err(e) if e.status_code() == StatusCode::TableNotFound => {
2918                    // Fall back to scanning the logical table. It still works, but without
2919                    // `__tsid` optimization.
2920                    scan_provider.clone()
2921                }
2922                Err(e) => return Err(e).context(CatalogSnafu),
2923            };
2924
2925            if !Arc::ptr_eq(&physical_provider, &scan_provider) {
2926                // Only rewrite when internal columns exist in physical schema.
2927                let physical_table = self.table_from_source(&physical_provider)?;
2928
2929                let has_table_id = physical_table
2930                    .schema()
2931                    .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
2932                    .is_some();
2933                let has_tsid = physical_table
2934                    .schema()
2935                    .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2936                    .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2937
2938                if has_table_id && has_tsid {
2939                    scan_provider = physical_provider;
2940                    maybe_phy_table_ref = physical_table_ref;
2941                    table_id_filter = Some(logical_table.table_info().ident.table_id);
2942                }
2943            }
2944        }
2945
2946        let scan_table = self.table_from_source(&scan_provider)?;
2947
2948        let use_tsid = table_id_filter.is_some()
2949            && scan_table
2950                .schema()
2951                .column_schema_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
2952                .is_some_and(|col| matches!(col.data_type, ConcreteDataType::UInt64(_)));
2953        self.ctx.use_tsid = use_tsid;
2954
2955        let all_table_tags = self.ctx.tag_columns.clone();
2956
2957        let scan_tag_columns = if use_tsid {
2958            let mut scan_tags = self.ctx.tag_columns.clone();
2959            for matcher in &self.ctx.selector_matcher {
2960                if is_metric_engine_internal_column(&matcher.name) {
2961                    continue;
2962                }
2963                if all_table_tags.iter().any(|tag| tag == &matcher.name) {
2964                    scan_tags.push(matcher.name.clone());
2965                }
2966            }
2967            scan_tags.sort_unstable();
2968            scan_tags.dedup();
2969            scan_tags
2970        } else {
2971            self.ctx.tag_columns.clone()
2972        };
2973
2974        let time_index_data_type = scan_table
2975            .schema()
2976            .timestamp_column()
2977            .with_context(|| TimeIndexNotFoundSnafu {
2978                table: maybe_phy_table_ref.to_quoted_string(),
2979            })?
2980            .data_type
2981            .clone();
2982        let is_time_index_second =
2983            time_index_data_type == ConcreteDataType::timestamp_second_datatype();
2984
2985        let scan_projection = if table_id_filter.is_some() {
2986            let mut required_columns = HashSet::new();
2987            required_columns.insert(DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string());
2988            required_columns.insert(self.ctx.time_index_column.clone().with_context(|| {
2989                TimeIndexNotFoundSnafu {
2990                    table: maybe_phy_table_ref.to_quoted_string(),
2991                }
2992            })?);
2993            for col in &scan_tag_columns {
2994                required_columns.insert(col.clone());
2995            }
2996            for col in &self.ctx.field_columns {
2997                required_columns.insert(col.clone());
2998            }
2999            if use_tsid {
3000                required_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
3001            }
3002
3003            let arrow_schema = scan_table.schema().arrow_schema().clone();
3004            Some(
3005                arrow_schema
3006                    .fields()
3007                    .iter()
3008                    .enumerate()
3009                    .filter(|(_, field)| required_columns.contains(field.name().as_str()))
3010                    .map(|(idx, _)| idx)
3011                    .collect::<Vec<_>>(),
3012            )
3013        } else {
3014            None
3015        };
3016
3017        let mut scan_plan =
3018            LogicalPlanBuilder::scan(maybe_phy_table_ref.clone(), scan_provider, scan_projection)
3019                .context(DataFusionPlanningSnafu)?
3020                .build()
3021                .context(DataFusionPlanningSnafu)?;
3022
3023        if let Some(table_id) = table_id_filter {
3024            scan_plan = LogicalPlanBuilder::from(scan_plan)
3025                .filter(
3026                    DfExpr::Column(Column::from_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME))
3027                        .eq(lit(table_id)),
3028                )
3029                .context(DataFusionPlanningSnafu)?
3030                .alias(table_ref.clone()) // rename the relation back to logical table's name after filtering
3031                .context(DataFusionPlanningSnafu)?
3032                .build()
3033                .context(DataFusionPlanningSnafu)?;
3034        }
3035
3036        if is_time_index_second {
3037            // Promote seconds so millisecond offsets remain exact; retain finer precision.
3038            // Later manipulators compare native sample ticks, while PromQL evaluation and
3039            // emitted timestamps remain millisecond-based, so this projection must not
3040            // silently truncate a finer-grained time index.
3041            let expr: Vec<_> = self
3042                .create_field_column_exprs()?
3043                .into_iter()
3044                .chain(
3045                    scan_tag_columns
3046                        .iter()
3047                        .map(|tag| DfExpr::Column(Column::from_name(tag))),
3048                )
3049                .chain(self.ctx.use_tsid.then_some(DfExpr::Column(Column::new(
3050                    Some(table_ref.clone()),
3051                    DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
3052                ))))
3053                .chain(Some(DfExpr::Alias(Alias {
3054                    expr: Box::new(DfExpr::Cast(Cast {
3055                        expr: Box::new(self.create_time_index_column_expr()?),
3056                        data_type: ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
3057                    })),
3058                    relation: Some(table_ref.clone()),
3059                    name: self
3060                        .ctx
3061                        .time_index_column
3062                        .as_ref()
3063                        .with_context(|| TimeIndexNotFoundSnafu {
3064                            table: table_ref.to_quoted_string(),
3065                        })?
3066                        .clone(),
3067                    metadata: None,
3068                })))
3069                .collect::<Vec<_>>();
3070            scan_plan = LogicalPlanBuilder::from(scan_plan)
3071                .project(expr)
3072                .context(DataFusionPlanningSnafu)?
3073                .build()
3074                .context(DataFusionPlanningSnafu)?;
3075        } else if table_id_filter.is_some()
3076            || time_index_data_type == ConcreteDataType::timestamp_microsecond_datatype()
3077            || time_index_data_type == ConcreteDataType::timestamp_nanosecond_datatype()
3078        {
3079            // Drop the internal `__table_id` column after filtering and preserve PromQL's
3080            // field/tag/timestamp column order for native microsecond/nanosecond timestamps.
3081            // Keeping the original time column also lets the existing ordering hints
3082            // use PerSeries scans without a cast, repartition, and sort. This benefits
3083            // multi-evaluation selectors too; only a single evaluation can use LastRow.
3084            let project_exprs = self
3085                .create_field_column_exprs()?
3086                .into_iter()
3087                .chain(
3088                    scan_tag_columns
3089                        .iter()
3090                        .map(|tag| DfExpr::Column(Column::from_name(tag))),
3091                )
3092                .chain(
3093                    self.ctx
3094                        .use_tsid
3095                        .then_some(DfExpr::Column(Column::from_name(
3096                            DATA_SCHEMA_TSID_COLUMN_NAME,
3097                        ))),
3098                )
3099                .chain(Some(self.create_time_index_column_expr()?))
3100                .collect::<Vec<_>>();
3101
3102            scan_plan = LogicalPlanBuilder::from(scan_plan)
3103                .project(project_exprs)
3104                .context(DataFusionPlanningSnafu)?
3105                .build()
3106                .context(DataFusionPlanningSnafu)?;
3107        }
3108
3109        let result = LogicalPlanBuilder::from(scan_plan)
3110            .build()
3111            .context(DataFusionPlanningSnafu)?;
3112        Ok(result)
3113    }
3114
3115    fn collect_row_key_tag_columns_from_plan(
3116        &self,
3117        plan: &LogicalPlan,
3118    ) -> Result<BTreeSet<String>> {
3119        fn walk(
3120            planner: &PromPlanner,
3121            plan: &LogicalPlan,
3122            out: &mut BTreeSet<String>,
3123        ) -> Result<()> {
3124            // Derived PromQL plans may contain non-Greptime scans without row-key metadata.
3125            if let LogicalPlan::TableScan(scan) = plan
3126                && let Ok(table) = planner.table_from_source(&scan.source)
3127            {
3128                for col in table.table_info().meta.row_key_column_names() {
3129                    if col != DATA_SCHEMA_TABLE_ID_COLUMN_NAME
3130                        && col != DATA_SCHEMA_TSID_COLUMN_NAME
3131                        && !is_metric_engine_internal_column(col)
3132                    {
3133                        out.insert(col.clone());
3134                    }
3135                }
3136            }
3137
3138            for input in plan.inputs() {
3139                walk(planner, input, out)?;
3140            }
3141            Ok(())
3142        }
3143
3144        let mut out = BTreeSet::new();
3145        walk(self, plan, &mut out)?;
3146        Ok(out)
3147    }
3148
3149    fn ensure_tag_columns_available(
3150        &self,
3151        plan: LogicalPlan,
3152        required_tags: &BTreeSet<String>,
3153    ) -> Result<LogicalPlan> {
3154        if required_tags.is_empty() {
3155            return Ok(plan);
3156        }
3157
3158        struct Rewriter {
3159            required_tags: BTreeSet<String>,
3160        }
3161
3162        impl TreeNodeRewriter for Rewriter {
3163            type Node = LogicalPlan;
3164
3165            fn f_up(
3166                &mut self,
3167                node: Self::Node,
3168            ) -> datafusion_common::Result<Transformed<Self::Node>> {
3169                match node {
3170                    LogicalPlan::TableScan(scan) => {
3171                        let schema = scan.source.schema();
3172                        let mut projection = match scan.projection.clone() {
3173                            Some(p) => p,
3174                            None => {
3175                                // Scanning all columns already covers required tags.
3176                                return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
3177                            }
3178                        };
3179
3180                        let mut changed = false;
3181                        for tag in &self.required_tags {
3182                            if let Some((idx, _)) = schema
3183                                .fields()
3184                                .iter()
3185                                .enumerate()
3186                                .find(|(_, field)| field.name() == tag)
3187                                && !projection.contains(&idx)
3188                            {
3189                                projection.push(idx);
3190                                changed = true;
3191                            }
3192                        }
3193
3194                        if !changed {
3195                            return Ok(Transformed::no(LogicalPlan::TableScan(scan)));
3196                        }
3197
3198                        projection.sort_unstable();
3199                        projection.dedup();
3200
3201                        let new_scan = TableScan::try_new(
3202                            scan.table_name.clone(),
3203                            scan.source.clone(),
3204                            Some(projection),
3205                            scan.filters,
3206                            scan.fetch,
3207                        )?;
3208                        Ok(Transformed::yes(LogicalPlan::TableScan(new_scan)))
3209                    }
3210                    LogicalPlan::Projection(proj) => {
3211                        let input_schema = proj.input.schema();
3212
3213                        let existing = proj
3214                            .schema
3215                            .fields()
3216                            .iter()
3217                            .map(|f| f.name().as_str())
3218                            .collect::<HashSet<_>>();
3219
3220                        let mut expr = proj.expr.clone();
3221                        let mut has_changed = false;
3222                        for tag in &self.required_tags {
3223                            if existing.contains(tag.as_str()) {
3224                                continue;
3225                            }
3226
3227                            if let Some(idx) = input_schema.index_of_column_by_name(None, tag) {
3228                                expr.push(DfExpr::Column(Column::from(
3229                                    input_schema.qualified_field(idx),
3230                                )));
3231                                has_changed = true;
3232                            }
3233                        }
3234
3235                        if !has_changed {
3236                            return Ok(Transformed::no(LogicalPlan::Projection(proj)));
3237                        }
3238
3239                        let new_proj = Projection::try_new(expr, proj.input)?;
3240                        Ok(Transformed::yes(LogicalPlan::Projection(new_proj)))
3241                    }
3242                    other => Ok(Transformed::no(other)),
3243                }
3244            }
3245        }
3246
3247        let mut rewriter = Rewriter {
3248            required_tags: required_tags.clone(),
3249        };
3250        let rewritten = plan
3251            .rewrite(&mut rewriter)
3252            .context(DataFusionPlanningSnafu)?;
3253        Ok(rewritten.data)
3254    }
3255
3256    fn refresh_tag_columns_from_schema(&mut self, schema: &DFSchemaRef) {
3257        let time_index = self.ctx.time_index_column.as_deref();
3258        let field_columns = self.ctx.field_columns.iter().collect::<HashSet<_>>();
3259
3260        let mut tags = schema
3261            .fields()
3262            .iter()
3263            .map(|f| f.name())
3264            .filter(|name| Some(name.as_str()) != time_index)
3265            .filter(|name| !field_columns.contains(name))
3266            .filter(|name| !is_metric_engine_internal_column(name))
3267            .cloned()
3268            .collect::<Vec<_>>();
3269        tags.sort_unstable();
3270        tags.dedup();
3271        self.ctx.tag_columns = tags;
3272    }
3273
3274    /// Setup [PromPlannerContext]'s state fields.
3275    ///
3276    /// Returns a logical plan for an empty metric.
3277    async fn setup_context(&mut self) -> Result<Option<LogicalPlan>> {
3278        let table_ref = self.table_ref()?;
3279        let source = match self.table_provider.resolve_table(table_ref.clone()).await {
3280            Err(e) if e.status_code() == StatusCode::TableNotFound => {
3281                let plan = self.setup_context_for_empty_metric()?;
3282                return Ok(Some(plan));
3283            }
3284            res => res.context(CatalogSnafu)?,
3285        };
3286        let table = self.table_from_source(&source)?;
3287
3288        // set time index column name
3289        let time_index = table
3290            .schema()
3291            .timestamp_column()
3292            .with_context(|| TimeIndexNotFoundSnafu {
3293                table: table_ref.to_quoted_string(),
3294            })?
3295            .name
3296            .clone();
3297        self.ctx.time_index_column = Some(time_index);
3298
3299        // set values columns
3300        let values = table
3301            .table_info()
3302            .meta
3303            .field_column_names()
3304            .cloned()
3305            .collect();
3306        self.ctx.field_columns = values;
3307
3308        // set primary key (tag) columns
3309        let tags = table
3310            .table_info()
3311            .meta
3312            .row_key_column_names()
3313            .filter(|col| {
3314                // remove metric engine's internal columns
3315                col != &DATA_SCHEMA_TABLE_ID_COLUMN_NAME && col != &DATA_SCHEMA_TSID_COLUMN_NAME
3316            })
3317            .cloned()
3318            .collect();
3319        self.ctx.tag_columns = tags;
3320
3321        self.ctx.use_tsid = false;
3322
3323        Ok(None)
3324    }
3325
3326    /// Setup [PromPlannerContext]'s state fields for a non existent table
3327    /// without any rows.
3328    fn setup_context_for_empty_metric(&mut self) -> Result<LogicalPlan> {
3329        self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
3330        self.ctx.reset_table_name_and_schema();
3331        self.ctx.tag_columns = vec![];
3332        self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()];
3333        self.ctx.use_tsid = false;
3334
3335        // The table doesn't have any data, so we set start to 0 and end to -1.
3336        let plan = LogicalPlan::Extension(Extension {
3337            node: Arc::new(
3338                EmptyMetric::new(
3339                    0,
3340                    -1,
3341                    self.ctx.interval,
3342                    SPECIAL_TIME_FUNCTION.to_string(),
3343                    DEFAULT_FIELD_COLUMN.to_string(),
3344                    Some(lit(0.0f64)),
3345                )
3346                .context(DataFusionPlanningSnafu)?,
3347            ),
3348        });
3349        Ok(plan)
3350    }
3351
3352    // TODO(ruihang): insert column expr
3353    fn create_function_args(&self, args: &[Box<PromExpr>]) -> Result<FunctionArgs> {
3354        let mut result = FunctionArgs::default();
3355
3356        for arg in args {
3357            // First try to parse as literal expression (including binary expressions like 100.0 + 3.0)
3358            if let Some(expr) = Self::try_build_literal_expr(arg) {
3359                result.literals.push(expr);
3360            } else {
3361                // If not a literal, treat as vector input
3362                match arg.as_ref() {
3363                    PromExpr::Subquery(_)
3364                    | PromExpr::VectorSelector(_)
3365                    | PromExpr::MatrixSelector(_)
3366                    | PromExpr::Extension(_)
3367                    | PromExpr::Aggregate(_)
3368                    | PromExpr::Paren(_)
3369                    | PromExpr::Call(_)
3370                    | PromExpr::Binary(_)
3371                    | PromExpr::Unary(_) => {
3372                        if result.input.replace(*arg.clone()).is_some() {
3373                            MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
3374                        }
3375                    }
3376
3377                    _ => {
3378                        let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?;
3379                        result.literals.push(expr);
3380                    }
3381                }
3382            }
3383        }
3384
3385        Ok(result)
3386    }
3387
3388    fn create_mixed_range_function_exprs(
3389        &mut self,
3390        func: &Function,
3391        mut other_input_exprs: VecDeque<DfExpr>,
3392        float_field: &str,
3393        histogram_field: &str,
3394        input_schema: &DFSchemaRef,
3395    ) -> Result<Option<Vec<DfExpr>>> {
3396        let returns_histogram = matches!(
3397            func.name,
3398            "rate"
3399                | "increase"
3400                | "delta"
3401                | "idelta"
3402                | "irate"
3403                | "avg_over_time"
3404                | "sum_over_time"
3405                | "last_over_time"
3406        );
3407        if !returns_histogram
3408            && !matches!(
3409                func.name,
3410                "changes"
3411                    | "resets"
3412                    | "deriv"
3413                    | "min_over_time"
3414                    | "max_over_time"
3415                    | "count_over_time"
3416                    | "absent_over_time"
3417                    | "present_over_time"
3418                    | "stddev_over_time"
3419                    | "stdvar_over_time"
3420                    | "quantile_over_time"
3421                    | "predict_linear"
3422                    | "double_exponential_smoothing"
3423                    | "holt_winters"
3424            )
3425        {
3426            return Ok(None);
3427        }
3428
3429        if func.name == "predict_linear" {
3430            other_input_exprs[0] = DfExpr::Cast(Cast {
3431                expr: Box::new(other_input_exprs[0].clone()),
3432                data_type: ArrowDataType::Int64,
3433            });
3434        }
3435
3436        let timestamp_range = DfExpr::Column(Column::from_name(
3437            RangeManipulate::build_timestamp_range_name(
3438                self.ctx.time_index_column.as_ref().unwrap(),
3439            ),
3440        ));
3441        let float_range = DfExpr::Column(Column::from_name(float_field));
3442        let histogram_range = DfExpr::Column(Column::from_name(histogram_field));
3443        let mut args = Vec::with_capacity(other_input_exprs.len() + 6);
3444        args.push(lit(func.name));
3445        args.push(timestamp_range.clone());
3446        args.push(float_range.clone());
3447        args.push(histogram_range.clone());
3448        args.extend(other_input_exprs);
3449        if matches!(func.name, "rate" | "increase" | "delta") {
3450            args.push(self.create_time_index_column_expr()?);
3451            args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?));
3452        }
3453
3454        let mut float_expr = DfExpr::ScalarFunction(ScalarFunction {
3455            func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3456            args: args.clone(),
3457        });
3458        if matches!(func.name, "rate" | "increase") {
3459            let raw_delta_function = if func.name == "rate" {
3460                "raw_delta_rate"
3461            } else {
3462                "raw_delta_increase"
3463            };
3464            let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
3465                func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())),
3466                args: vec![
3467                    lit(raw_delta_function),
3468                    timestamp_range,
3469                    float_range,
3470                    histogram_range,
3471                ],
3472            });
3473            float_expr = self.select_delta_range_math(
3474                func.name,
3475                input_schema,
3476                self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3477                delta_sum,
3478                float_expr,
3479            )?;
3480        }
3481        let exprs = if returns_histogram {
3482            self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()];
3483            vec![
3484                float_expr.alias(float_field),
3485                DfExpr::ScalarFunction(ScalarFunction {
3486                    func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())),
3487                    args,
3488                })
3489                .alias(histogram_field),
3490            ]
3491        } else {
3492            let display_name = float_expr.schema_name().to_string();
3493            self.ctx.field_columns = vec![display_name.clone()];
3494            vec![float_expr.alias(display_name)]
3495        };
3496        Ok(Some(exprs))
3497    }
3498
3499    /// Creates function expressions for projection and returns the expressions and new tags.
3500    ///
3501    /// # Side Effects
3502    ///
3503    /// This method will update [PromPlannerContext]'s fields and tags if needed.
3504    fn create_function_expr(
3505        &mut self,
3506        func: &Function,
3507        other_input_exprs: Vec<DfExpr>,
3508        input_schema: &DFSchemaRef,
3509        query_engine_state: &QueryEngineState,
3510    ) -> Result<(Vec<DfExpr>, Vec<String>)> {
3511        // TODO(ruihang): check function args list
3512        let mut other_input_exprs: VecDeque<DfExpr> = other_input_exprs.into();
3513        if let Some((float_field, histogram_field)) =
3514            Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns)
3515                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
3516            && let Some(exprs) = self.create_mixed_range_function_exprs(
3517                func,
3518                other_input_exprs.clone(),
3519                &float_field,
3520                &histogram_field,
3521                input_schema,
3522            )?
3523        {
3524            return Ok((exprs, vec![]));
3525        }
3526        let alternative_samples =
3527            Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns);
3528        let all_field_columns_are_native_histogram_ranges =
3529            self.all_field_columns_are_native_histogram_ranges(input_schema);
3530
3531        // TODO(ruihang): set this according to in-param list
3532        let field_column_pos = 0;
3533        let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
3534        // New labels after executing the function, e.g. `label_replace` etc.
3535        let mut new_tags = vec![];
3536        let promql_annotations = self.promql_annotations.clone();
3537        let native_histogram_drop_udf = |name: &str| {
3538            Arc::new(NativeHistogramDrop::float_null_udf(
3539                format!(
3540                    "{name}: dropped native histogram samples because this function is not supported for native histograms"
3541                ),
3542                promql_annotations.clone(),
3543            ))
3544        };
3545        let scalar_func = match func.name {
3546            "increase" => {
3547                if all_field_columns_are_native_histogram_ranges {
3548                    ScalarFunc::ExtrapolateUdf(
3549                        Arc::new(NativeHistogramIncrease::scalar_udf_with_collector(
3550                            self.promql_annotations.clone(),
3551                        )),
3552                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3553                    )
3554                } else {
3555                    ScalarFunc::ExtrapolateUdf(
3556                        Arc::new(Increase::scalar_udf()),
3557                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3558                    )
3559                }
3560            }
3561            "rate" => {
3562                if all_field_columns_are_native_histogram_ranges {
3563                    ScalarFunc::ExtrapolateUdf(
3564                        Arc::new(NativeHistogramRate::scalar_udf_with_collector(
3565                            self.promql_annotations.clone(),
3566                        )),
3567                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3568                    )
3569                } else {
3570                    ScalarFunc::ExtrapolateUdf(
3571                        Arc::new(Rate::scalar_udf()),
3572                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3573                    )
3574                }
3575            }
3576            "delta" => {
3577                if all_field_columns_are_native_histogram_ranges {
3578                    ScalarFunc::ExtrapolateUdf(
3579                        Arc::new(NativeHistogramDelta::scalar_udf_with_collector(
3580                            self.promql_annotations.clone(),
3581                        )),
3582                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3583                    )
3584                } else {
3585                    ScalarFunc::ExtrapolateUdf(
3586                        Arc::new(Delta::scalar_udf()),
3587                        self.ctx.range.context(ExpectRangeSelectorSnafu)?,
3588                    )
3589                }
3590            }
3591            "idelta" => {
3592                if all_field_columns_are_native_histogram_ranges {
3593                    ScalarFunc::Udf(Arc::new(NativeHistogramIDelta::scalar_udf_with_collector(
3594                        self.promql_annotations.clone(),
3595                    )))
3596                } else {
3597                    ScalarFunc::Udf(Arc::new(IDelta::<false>::scalar_udf()))
3598                }
3599            }
3600            "irate" => {
3601                if all_field_columns_are_native_histogram_ranges {
3602                    ScalarFunc::Udf(Arc::new(NativeHistogramIRate::scalar_udf_with_collector(
3603                        self.promql_annotations.clone(),
3604                    )))
3605                } else {
3606                    ScalarFunc::Udf(Arc::new(IDelta::<true>::scalar_udf()))
3607                }
3608            }
3609            "resets" => {
3610                if all_field_columns_are_native_histogram_ranges {
3611                    ScalarFunc::Udf(Arc::new(NativeHistogramResets::scalar_udf()))
3612                } else {
3613                    ScalarFunc::Udf(Arc::new(Resets::scalar_udf()))
3614                }
3615            }
3616            "changes" => {
3617                if all_field_columns_are_native_histogram_ranges {
3618                    ScalarFunc::Udf(Arc::new(NativeHistogramChanges::scalar_udf()))
3619                } else {
3620                    ScalarFunc::Udf(Arc::new(Changes::scalar_udf()))
3621                }
3622            }
3623            "deriv" => {
3624                if all_field_columns_are_native_histogram_ranges {
3625                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3626                } else {
3627                    ScalarFunc::Udf(Arc::new(Deriv::scalar_udf()))
3628                }
3629            }
3630            "avg_over_time" => {
3631                if all_field_columns_are_native_histogram_ranges {
3632                    ScalarFunc::Udf(Arc::new(
3633                        NativeHistogramAvgOverTime::scalar_udf_with_collector(
3634                            self.promql_annotations.clone(),
3635                        ),
3636                    ))
3637                } else {
3638                    ScalarFunc::Udf(Arc::new(AvgOverTime::scalar_udf()))
3639                }
3640            }
3641            "min_over_time" => {
3642                if all_field_columns_are_native_histogram_ranges {
3643                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3644                } else {
3645                    ScalarFunc::Udf(Arc::new(MinOverTime::scalar_udf()))
3646                }
3647            }
3648            "max_over_time" => {
3649                if all_field_columns_are_native_histogram_ranges {
3650                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3651                } else {
3652                    ScalarFunc::Udf(Arc::new(MaxOverTime::scalar_udf()))
3653                }
3654            }
3655            "sum_over_time" => {
3656                if all_field_columns_are_native_histogram_ranges {
3657                    ScalarFunc::Udf(Arc::new(
3658                        NativeHistogramSumOverTime::scalar_udf_with_collector(
3659                            self.promql_annotations.clone(),
3660                        ),
3661                    ))
3662                } else {
3663                    ScalarFunc::Udf(Arc::new(SumOverTime::scalar_udf()))
3664                }
3665            }
3666            "count_over_time" => {
3667                if all_field_columns_are_native_histogram_ranges {
3668                    ScalarFunc::Udf(Arc::new(NativeHistogramCountOverTime::scalar_udf()))
3669                } else {
3670                    ScalarFunc::Udf(Arc::new(CountOverTime::scalar_udf()))
3671                }
3672            }
3673            "last_over_time" => {
3674                if all_field_columns_are_native_histogram_ranges {
3675                    ScalarFunc::Udf(Arc::new(NativeHistogramLastOverTime::scalar_udf()))
3676                } else {
3677                    ScalarFunc::Udf(Arc::new(LastOverTime::scalar_udf()))
3678                }
3679            }
3680            "absent_over_time" => {
3681                if all_field_columns_are_native_histogram_ranges {
3682                    ScalarFunc::Udf(Arc::new(NativeHistogramAbsentOverTime::scalar_udf()))
3683                } else {
3684                    ScalarFunc::Udf(Arc::new(AbsentOverTime::scalar_udf()))
3685                }
3686            }
3687            "present_over_time" => {
3688                if all_field_columns_are_native_histogram_ranges {
3689                    ScalarFunc::Udf(Arc::new(NativeHistogramPresentOverTime::scalar_udf()))
3690                } else {
3691                    ScalarFunc::Udf(Arc::new(PresentOverTime::scalar_udf()))
3692                }
3693            }
3694            "stddev_over_time" => {
3695                if all_field_columns_are_native_histogram_ranges {
3696                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3697                } else {
3698                    ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf()))
3699                }
3700            }
3701            "stdvar_over_time" => {
3702                if all_field_columns_are_native_histogram_ranges {
3703                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3704                } else {
3705                    ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf()))
3706                }
3707            }
3708            "quantile_over_time" => {
3709                if all_field_columns_are_native_histogram_ranges {
3710                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3711                } else {
3712                    ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf()))
3713                }
3714            }
3715            "predict_linear" => {
3716                if all_field_columns_are_native_histogram_ranges {
3717                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3718                } else {
3719                    other_input_exprs[0] = DfExpr::Cast(Cast {
3720                        expr: Box::new(other_input_exprs[0].clone()),
3721                        data_type: ArrowDataType::Int64,
3722                    });
3723                    ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
3724                }
3725            }
3726            "double_exponential_smoothing" | "holt_winters" => {
3727                if all_field_columns_are_native_histogram_ranges {
3728                    ScalarFunc::Udf(native_histogram_drop_udf(func.name))
3729                } else {
3730                    ScalarFunc::Udf(Arc::new(DoubleExponentialSmoothing::scalar_udf()))
3731                }
3732            }
3733            "histogram_count" => {
3734                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramCount::scalar_udf()))
3735            }
3736            "histogram_sum" => {
3737                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramSum::scalar_udf()))
3738            }
3739            "histogram_avg" => {
3740                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramAvg::scalar_udf()))
3741            }
3742            "histogram_stddev" => {
3743                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStddev::scalar_udf()))
3744            }
3745            "histogram_stdvar" => {
3746                ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf()))
3747            }
3748            "time" => {
3749                exprs.push(build_special_time_expr(
3750                    self.ctx.time_index_column.as_ref().unwrap(),
3751                ));
3752                ScalarFunc::GeneratedExpr
3753            }
3754            "minute" => {
3755                // date_part('minute', time_index)
3756                let expr = self.date_part_on_time_index("minute")?;
3757                exprs.push(expr);
3758                ScalarFunc::GeneratedExpr
3759            }
3760            "hour" => {
3761                // date_part('hour', time_index)
3762                let expr = self.date_part_on_time_index("hour")?;
3763                exprs.push(expr);
3764                ScalarFunc::GeneratedExpr
3765            }
3766            "month" => {
3767                // date_part('month', time_index)
3768                let expr = self.date_part_on_time_index("month")?;
3769                exprs.push(expr);
3770                ScalarFunc::GeneratedExpr
3771            }
3772            "year" => {
3773                // date_part('year', time_index)
3774                let expr = self.date_part_on_time_index("year")?;
3775                exprs.push(expr);
3776                ScalarFunc::GeneratedExpr
3777            }
3778            "day_of_month" => {
3779                // date_part('day', time_index)
3780                let expr = self.date_part_on_time_index("day")?;
3781                exprs.push(expr);
3782                ScalarFunc::GeneratedExpr
3783            }
3784            "day_of_week" => {
3785                // date_part('dow', time_index)
3786                let expr = self.date_part_on_time_index("dow")?;
3787                exprs.push(expr);
3788                ScalarFunc::GeneratedExpr
3789            }
3790            "day_of_year" => {
3791                // date_part('doy', time_index)
3792                let expr = self.date_part_on_time_index("doy")?;
3793                exprs.push(expr);
3794                ScalarFunc::GeneratedExpr
3795            }
3796            "days_in_month" => {
3797                // date_part(
3798                //     'days',
3799                //     (date_trunc('month', <TIME INDEX>::date) + interval '1 month - 1 day')
3800                // );
3801                let day_lit_expr = "day".lit();
3802                let month_lit_expr = "month".lit();
3803                let interval_1month_lit_expr =
3804                    DfExpr::Literal(ScalarValue::IntervalYearMonth(Some(1)), None);
3805                let interval_1day_lit_expr = DfExpr::Literal(
3806                    ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(1, 0))),
3807                    None,
3808                );
3809                let the_1month_minus_1day_expr = DfExpr::BinaryExpr(BinaryExpr {
3810                    left: Box::new(interval_1month_lit_expr),
3811                    op: Operator::Minus,
3812                    right: Box::new(interval_1day_lit_expr),
3813                });
3814                let date_trunc_expr = DfExpr::ScalarFunction(ScalarFunction {
3815                    func: datafusion_functions::datetime::date_trunc(),
3816                    args: vec![month_lit_expr, self.create_time_index_column_expr()?],
3817                });
3818                let date_trunc_plus_interval_expr = DfExpr::BinaryExpr(BinaryExpr {
3819                    left: Box::new(date_trunc_expr),
3820                    op: Operator::Plus,
3821                    right: Box::new(the_1month_minus_1day_expr),
3822                });
3823                let date_part_expr = DfExpr::ScalarFunction(ScalarFunction {
3824                    func: datafusion_functions::datetime::date_part(),
3825                    args: vec![day_lit_expr, date_trunc_plus_interval_expr],
3826                });
3827
3828                exprs.push(date_part_expr);
3829                ScalarFunc::GeneratedExpr
3830            }
3831
3832            "label_join" => {
3833                self.ctx.use_tsid = false;
3834                let (concat_expr, dst_label) = Self::build_concat_labels_expr(
3835                    &mut other_input_exprs,
3836                    &self.ctx,
3837                    query_engine_state,
3838                )?;
3839
3840                // Reserve the current field columns except the `dst_label`.
3841                for value in &self.ctx.field_columns {
3842                    if *value != dst_label {
3843                        let expr = DfExpr::Column(Column::from_name(value));
3844                        exprs.push(expr);
3845                    }
3846                }
3847
3848                // Remove it from tag columns if exists to avoid duplicated column names
3849                self.ctx.tag_columns.retain(|tag| *tag != dst_label);
3850                new_tags.push(dst_label);
3851                // Add the new label expr to evaluate
3852                exprs.push(concat_expr);
3853
3854                ScalarFunc::GeneratedExpr
3855            }
3856            "label_replace" => {
3857                self.ctx.use_tsid = false;
3858                if let Some((replace_expr, dst_label)) = self
3859                    .build_regexp_replace_label_expr(&mut other_input_exprs, query_engine_state)?
3860                {
3861                    // Reserve the current field columns except the `dst_label`.
3862                    for value in &self.ctx.field_columns {
3863                        if *value != dst_label {
3864                            let expr = DfExpr::Column(Column::from_name(value));
3865                            exprs.push(expr);
3866                        }
3867                    }
3868
3869                    ensure!(
3870                        !self.ctx.tag_columns.contains(&dst_label),
3871                        SameLabelSetSnafu
3872                    );
3873                    new_tags.push(dst_label);
3874                    // Add the new label expr to evaluate
3875                    exprs.push(replace_expr);
3876                } else {
3877                    // Keep the current field columns
3878                    for value in &self.ctx.field_columns {
3879                        let expr = DfExpr::Column(Column::from_name(value));
3880                        exprs.push(expr);
3881                    }
3882                }
3883
3884                ScalarFunc::GeneratedExpr
3885            }
3886            "sort" | "sort_desc" => {
3887                // Value sorting silently ignores native histogram samples.
3888                for value in &self.ctx.field_columns {
3889                    if !Self::field_column_is_native_histogram(input_schema, value) {
3890                        exprs.push(DfExpr::Column(Column::from_name(value)));
3891                    }
3892                }
3893                // Keep a nullable float field so the normal empty-value filter produces an
3894                // empty vector when the input contains only histograms.
3895                if exprs.is_empty() {
3896                    exprs.push(DfExpr::Literal(ScalarValue::Float64(None), None));
3897                }
3898
3899                ScalarFunc::GeneratedExpr
3900            }
3901            "sort_by_label" | "sort_by_label_desc" | "timestamp" => {
3902                // These functions are not expression but a part of plan,
3903                // they are processed by `prom_call_expr_to_plan`.
3904                for value in &self.ctx.field_columns {
3905                    let expr = DfExpr::Column(Column::from_name(value));
3906                    exprs.push(expr);
3907                }
3908
3909                ScalarFunc::GeneratedExpr
3910            }
3911            "round" if self.all_field_columns_are_native_histograms(input_schema) => {
3912                if other_input_exprs.is_empty() {
3913                    other_input_exprs.push_front(0.0f64.lit());
3914                }
3915                ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3916            }
3917            "round" => {
3918                if other_input_exprs.is_empty() {
3919                    other_input_exprs.push_front(0.0f64.lit());
3920                }
3921                ScalarFunc::DataFusionUdf(Arc::new(Round::scalar_udf()))
3922            }
3923            "rad" | "deg" | "sgn" if self.all_field_columns_are_native_histograms(input_schema) => {
3924                ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3925            }
3926            "rad" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::radians()),
3927            "deg" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::degrees()),
3928            "sgn" => ScalarFunc::DataFusionBuiltin(datafusion::functions::math::signum()),
3929            "pi" => {
3930                // pi functions doesn't accepts any arguments, needs special processing
3931                let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3932                    func: datafusion::functions::math::pi(),
3933                    args: vec![],
3934                });
3935                exprs.push(fn_expr);
3936
3937                ScalarFunc::GeneratedExpr
3938            }
3939            _ => {
3940                if let Some(f) = query_engine_state
3941                    .session_state()
3942                    .scalar_functions()
3943                    .get(func.name)
3944                {
3945                    if self.all_field_columns_are_native_histograms(input_schema) {
3946                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3947                    } else {
3948                        ScalarFunc::DataFusionBuiltin(f.clone())
3949                    }
3950                } else if let Some(factory) = query_engine_state.scalar_function(func.name) {
3951                    if self.all_field_columns_are_native_histograms(input_schema) {
3952                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3953                    } else {
3954                        let func_state = query_engine_state.function_state();
3955                        let query_ctx = self.table_provider.query_ctx();
3956
3957                        ScalarFunc::DataFusionUdf(Arc::new(factory.provide(FunctionContext {
3958                            state: func_state,
3959                            query_ctx: query_ctx.clone(),
3960                        })))
3961                    }
3962                } else if let Some(f) = datafusion_functions::math::functions()
3963                    .iter()
3964                    .find(|f| f.name() == func.name)
3965                {
3966                    if self.all_field_columns_are_native_histograms(input_schema) {
3967                        ScalarFunc::DataFusionUdf(native_histogram_drop_udf(func.name))
3968                    } else {
3969                        ScalarFunc::DataFusionUdf(f.clone())
3970                    }
3971                } else {
3972                    return UnsupportedExprSnafu {
3973                        name: func.name.to_string(),
3974                    }
3975                    .fail();
3976                }
3977            }
3978        };
3979
3980        for value in &self.ctx.field_columns {
3981            let col_expr = DfExpr::Column(Column::from_name(value));
3982            let value_is_histogram = Self::field_column_is_native_histogram(input_schema, value);
3983
3984            match scalar_func.clone() {
3985                ScalarFunc::DataFusionBuiltin(func) => {
3986                    if alternative_samples && value_is_histogram {
3987                        continue;
3988                    }
3989                    other_input_exprs.insert(field_column_pos, col_expr);
3990                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
3991                        func,
3992                        args: other_input_exprs.clone().into(),
3993                    });
3994                    exprs.push(fn_expr);
3995                    let _ = other_input_exprs.remove(field_column_pos);
3996                }
3997                ScalarFunc::DataFusionUdf(func) => {
3998                    if alternative_samples && value_is_histogram {
3999                        continue;
4000                    }
4001                    let args = itertools::chain!(
4002                        other_input_exprs.iter().take(field_column_pos).cloned(),
4003                        std::iter::once(col_expr),
4004                        other_input_exprs.iter().skip(field_column_pos).cloned()
4005                    )
4006                    .collect_vec();
4007                    exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }))
4008                }
4009                ScalarFunc::NativeHistogramUdf(func) => {
4010                    if value_is_histogram {
4011                        let args = itertools::chain!(
4012                            other_input_exprs.iter().take(field_column_pos).cloned(),
4013                            std::iter::once(col_expr),
4014                            other_input_exprs.iter().skip(field_column_pos).cloned()
4015                        )
4016                        .collect_vec();
4017                        exprs.push(DfExpr::ScalarFunction(ScalarFunction { func, args }));
4018                    } else if !alternative_samples {
4019                        exprs.push(
4020                            DfExpr::Literal(ScalarValue::Float64(None), None).alias(format!(
4021                                "{}_{}",
4022                                func.name(),
4023                                value
4024                            )),
4025                        );
4026                    }
4027                }
4028                ScalarFunc::Udf(func) => {
4029                    let ts_range_expr = DfExpr::Column(Column::from_name(
4030                        RangeManipulate::build_timestamp_range_name(
4031                            self.ctx.time_index_column.as_ref().unwrap(),
4032                        ),
4033                    ));
4034                    other_input_exprs.insert(field_column_pos, ts_range_expr);
4035                    other_input_exprs.insert(field_column_pos + 1, col_expr);
4036                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4037                        func,
4038                        args: other_input_exprs.clone().into(),
4039                    });
4040                    exprs.push(fn_expr);
4041                    let _ = other_input_exprs.remove(field_column_pos + 1);
4042                    let _ = other_input_exprs.remove(field_column_pos);
4043                }
4044                ScalarFunc::ExtrapolateUdf(udf, range_length) => {
4045                    let ts_range_expr = DfExpr::Column(Column::from_name(
4046                        RangeManipulate::build_timestamp_range_name(
4047                            self.ctx.time_index_column.as_ref().unwrap(),
4048                        ),
4049                    ));
4050                    other_input_exprs.insert(field_column_pos, ts_range_expr.clone());
4051                    other_input_exprs.insert(field_column_pos + 1, col_expr.clone());
4052                    other_input_exprs
4053                        .insert(field_column_pos + 2, self.create_time_index_column_expr()?);
4054                    other_input_exprs.push_back(lit(range_length));
4055                    let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
4056                        func: udf,
4057                        args: other_input_exprs.clone().into(),
4058                    });
4059                    let fn_expr = if matches!(func.name, "rate" | "increase")
4060                        && !all_field_columns_are_native_histogram_ranges
4061                    {
4062                        let delta_sum = DfExpr::ScalarFunction(ScalarFunction {
4063                            func: Arc::new(SumOverTime::scalar_udf()),
4064                            args: vec![ts_range_expr, col_expr],
4065                        });
4066                        self.select_delta_range_math(
4067                            func.name,
4068                            input_schema,
4069                            range_length,
4070                            delta_sum,
4071                            fn_expr,
4072                        )?
4073                    } else {
4074                        fn_expr
4075                    };
4076                    exprs.push(fn_expr);
4077                    let _ = other_input_exprs.pop_back();
4078                    let _ = other_input_exprs.remove(field_column_pos + 2);
4079                    let _ = other_input_exprs.remove(field_column_pos + 1);
4080                    let _ = other_input_exprs.remove(field_column_pos);
4081                }
4082                ScalarFunc::GeneratedExpr => {}
4083            }
4084        }
4085
4086        // Update value columns' name, and alias them to remove qualifiers
4087        // For label functions such as `label_join`, `label_replace`, etc.,
4088        // we keep the fields unchanged.
4089        if !matches!(func.name, "label_join" | "label_replace") {
4090            let mut new_field_columns = Vec::with_capacity(exprs.len());
4091
4092            exprs = exprs
4093                .into_iter()
4094                .map(|expr| {
4095                    let display_name = expr.schema_name().to_string();
4096                    new_field_columns.push(display_name.clone());
4097                    Ok(expr.alias(display_name))
4098                })
4099                .collect::<std::result::Result<Vec<_>, _>>()
4100                .context(DataFusionPlanningSnafu)?;
4101
4102            self.ctx.field_columns = new_field_columns;
4103        }
4104
4105        Ok((exprs, new_tags))
4106    }
4107
4108    fn select_delta_range_math(
4109        &self,
4110        function: &str,
4111        input_schema: &DFSchemaRef,
4112        range_length: Millisecond,
4113        delta_sum: DfExpr,
4114        cumulative: DfExpr,
4115    ) -> Result<DfExpr> {
4116        let marker_is_delta = if self
4117            .ctx
4118            .tag_columns
4119            .iter()
4120            .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)
4121        {
4122            Self::field_column_type(input_schema, OTLP_AGGREGATION_TEMPORALITY_LABEL)
4123                .filter(|data_type| Self::string_value_data_type(data_type).is_some())
4124                .map(|_| {
4125                    DfExpr::Column(Column::from_name(OTLP_AGGREGATION_TEMPORALITY_LABEL))
4126                        .eq(lit(GREPTIME_TEMPORALITY_DELTA))
4127                })
4128        } else {
4129            None
4130        };
4131        let Some(marker_is_delta) = marker_is_delta else {
4132            return Ok(cumulative);
4133        };
4134
4135        let delta = if function == "rate" {
4136            DfExpr::BinaryExpr(BinaryExpr {
4137                left: Box::new(delta_sum),
4138                op: Operator::Divide,
4139                right: Box::new(lit(range_length as f64 / 1000.0)),
4140            })
4141        } else {
4142            delta_sum
4143        };
4144        let display_name = cumulative.schema_name().to_string();
4145        when(marker_is_delta, delta)
4146            .otherwise(cumulative)
4147            .context(DataFusionPlanningSnafu)
4148            .map(|expr| expr.alias(display_name))
4149    }
4150
4151    /// Validate label name according to Prometheus specification.
4152    /// Label names must match the regex: [a-zA-Z_][a-zA-Z0-9_]*
4153    /// Additionally, label names starting with double underscores are reserved for internal use.
4154    fn validate_label_name(label_name: &str) -> Result<()> {
4155        // Check if label name starts with double underscores (reserved)
4156        if label_name.starts_with("__") {
4157            return InvalidDestinationLabelNameSnafu { label_name }.fail();
4158        }
4159        // Check if label name matches the required pattern
4160        if !LABEL_NAME_REGEX.is_match(label_name) {
4161            return InvalidDestinationLabelNameSnafu { label_name }.fail();
4162        }
4163
4164        Ok(())
4165    }
4166
4167    /// Build expr for `label_replace` function
4168    fn build_regexp_replace_label_expr(
4169        &self,
4170        other_input_exprs: &mut VecDeque<DfExpr>,
4171        query_engine_state: &QueryEngineState,
4172    ) -> Result<Option<(DfExpr, String)>> {
4173        // label_replace(vector, dst_label, replacement, src_label, regex)
4174        let dst_label = match other_input_exprs.pop_front() {
4175            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4176            other => UnexpectedPlanExprSnafu {
4177                desc: format!("expected dst_label string literal, but found {:?}", other),
4178            }
4179            .fail()?,
4180        };
4181
4182        // Validate the destination label name
4183        Self::validate_label_name(&dst_label)?;
4184        let replacement = match other_input_exprs.pop_front() {
4185            Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), _)) => r,
4186            other => UnexpectedPlanExprSnafu {
4187                desc: format!("expected replacement string literal, but found {:?}", other),
4188            }
4189            .fail()?,
4190        };
4191        let src_label = match other_input_exprs.pop_front() {
4192            Some(DfExpr::Literal(ScalarValue::Utf8(Some(s)), None)) => s,
4193            other => UnexpectedPlanExprSnafu {
4194                desc: format!("expected src_label string literal, but found {:?}", other),
4195            }
4196            .fail()?,
4197        };
4198
4199        let regex = match other_input_exprs.pop_front() {
4200            Some(DfExpr::Literal(ScalarValue::Utf8(Some(r)), None)) => r,
4201            other => UnexpectedPlanExprSnafu {
4202                desc: format!("expected regex string literal, but found {:?}", other),
4203            }
4204            .fail()?,
4205        };
4206
4207        // Validate the regex before using it
4208        // doc: https://prometheus.io/docs/prometheus/latest/querying/functions/#label_replace
4209        regex::Regex::new(&regex).map_err(|_| {
4210            InvalidRegularExpressionSnafu {
4211                regex: regex.clone(),
4212            }
4213            .build()
4214        })?;
4215
4216        // If the src_label exists and regex is empty, keep everything unchanged.
4217        if self.ctx.tag_columns.contains(&src_label) && regex.is_empty() {
4218            return Ok(None);
4219        }
4220
4221        // If the src_label doesn't exists, and
4222        if !self.ctx.tag_columns.contains(&src_label) {
4223            if replacement.is_empty() {
4224                // the replacement is empty, keep everything unchanged.
4225                return Ok(None);
4226            } else {
4227                // the replacement is not empty, always adds dst_label with replacement value.
4228                return Ok(Some((
4229                    // alias literal `replacement` as dst_label
4230                    lit(replacement).alias(&dst_label),
4231                    dst_label,
4232                )));
4233            }
4234        }
4235
4236        // Preprocess the regex:
4237        // https://github.com/prometheus/prometheus/blob/d902abc50d6652ba8fe9a81ff8e5cce936114eba/promql/functions.go#L1575C32-L1575C37
4238        let regex = format!("^(?s:{regex})$");
4239
4240        let session_state = query_engine_state.session_state();
4241        let func = session_state
4242            .scalar_functions()
4243            .get("regexp_replace")
4244            .context(UnsupportedExprSnafu {
4245                name: "regexp_replace",
4246            })?;
4247
4248        // regexp_replace(src_label, regex, replacement)
4249        let args = vec![
4250            if src_label.is_empty() {
4251                DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
4252            } else {
4253                DfExpr::Column(Column::from_name(src_label))
4254            },
4255            DfExpr::Literal(ScalarValue::Utf8(Some(regex)), None),
4256            DfExpr::Literal(ScalarValue::Utf8(Some(replacement)), None),
4257        ];
4258
4259        Ok(Some((
4260            DfExpr::ScalarFunction(ScalarFunction {
4261                func: func.clone(),
4262                args,
4263            })
4264            .alias(&dst_label),
4265            dst_label,
4266        )))
4267    }
4268
4269    /// Build expr for `label_join` function
4270    fn build_concat_labels_expr(
4271        other_input_exprs: &mut VecDeque<DfExpr>,
4272        ctx: &PromPlannerContext,
4273        query_engine_state: &QueryEngineState,
4274    ) -> Result<(DfExpr, String)> {
4275        // label_join(vector, dst_label, separator, src_label_1, src_label_2, ...)
4276
4277        let dst_label = match other_input_exprs.pop_front() {
4278            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4279            other => UnexpectedPlanExprSnafu {
4280                desc: format!("expected dst_label string literal, but found {:?}", other),
4281            }
4282            .fail()?,
4283        };
4284        let separator = match other_input_exprs.pop_front() {
4285            Some(DfExpr::Literal(ScalarValue::Utf8(Some(d)), _)) => d,
4286            other => UnexpectedPlanExprSnafu {
4287                desc: format!("expected separator string literal, but found {:?}", other),
4288            }
4289            .fail()?,
4290        };
4291
4292        // Create a set of available columns (tag columns + field columns + time index column)
4293        let available_columns: HashSet<&str> = ctx
4294            .tag_columns
4295            .iter()
4296            .chain(ctx.field_columns.iter())
4297            .chain(ctx.time_index_column.as_ref())
4298            .map(|s| s.as_str())
4299            .collect();
4300
4301        let src_labels = other_input_exprs
4302            .iter()
4303            .map(|expr| {
4304                // Cast source label into column or null literal
4305                match expr {
4306                    DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => {
4307                        if label.is_empty() {
4308                            Ok(DfExpr::Literal(ScalarValue::Null, None))
4309                        } else if available_columns.contains(label.as_str()) {
4310                            // Label exists in the table schema
4311                            Ok(DfExpr::Column(Column::from_name(label)))
4312                        } else {
4313                            // Label doesn't exist, treat as empty string (null)
4314                            Ok(DfExpr::Literal(ScalarValue::Null, None))
4315                        }
4316                    }
4317                    other => UnexpectedPlanExprSnafu {
4318                        desc: format!(
4319                            "expected source label string literal, but found {:?}",
4320                            other
4321                        ),
4322                    }
4323                    .fail(),
4324                }
4325            })
4326            .collect::<Result<Vec<_>>>()?;
4327        ensure!(
4328            !src_labels.is_empty(),
4329            FunctionInvalidArgumentSnafu {
4330                fn_name: "label_join"
4331            }
4332        );
4333
4334        let session_state = query_engine_state.session_state();
4335        let func = session_state
4336            .scalar_functions()
4337            .get("concat_ws")
4338            .context(UnsupportedExprSnafu { name: "concat_ws" })?;
4339
4340        // concat_ws(separator, src_label_1, src_label_2, ...) as dst_label
4341        let mut args = Vec::with_capacity(1 + src_labels.len());
4342        args.push(DfExpr::Literal(ScalarValue::Utf8(Some(separator)), None));
4343        args.extend(src_labels);
4344
4345        Ok((
4346            DfExpr::ScalarFunction(ScalarFunction {
4347                func: func.clone(),
4348                args,
4349            })
4350            .alias(&dst_label),
4351            dst_label,
4352        ))
4353    }
4354
4355    fn create_time_index_column_expr(&self) -> Result<DfExpr> {
4356        Ok(DfExpr::Column(Column::from_name(
4357            self.ctx
4358                .time_index_column
4359                .clone()
4360                .with_context(|| TimeIndexNotFoundSnafu { table: "unknown" })?,
4361        )))
4362    }
4363
4364    fn create_tag_column_exprs(&self) -> Result<Vec<DfExpr>> {
4365        let mut result = Vec::with_capacity(self.ctx.tag_columns.len());
4366        for tag in &self.ctx.tag_columns {
4367            let expr = DfExpr::Column(Column::from_name(tag));
4368            result.push(expr);
4369        }
4370        Ok(result)
4371    }
4372
4373    fn create_field_column_exprs(&self) -> Result<Vec<DfExpr>> {
4374        let mut result = Vec::with_capacity(self.ctx.field_columns.len());
4375        for field in &self.ctx.field_columns {
4376            let expr = DfExpr::Column(Column::from_name(field));
4377            result.push(expr);
4378        }
4379        Ok(result)
4380    }
4381
4382    fn create_tag_and_time_index_column_sort_exprs(&self) -> Result<Vec<SortExpr>> {
4383        let mut result = self
4384            .ctx
4385            .tag_columns
4386            .iter()
4387            .map(|col| DfExpr::Column(Column::from_name(col)).sort(true, true))
4388            .collect::<Vec<_>>();
4389        result.push(self.create_time_index_column_expr()?.sort(true, true));
4390        Ok(result)
4391    }
4392
4393    fn create_field_columns_sort_exprs(&self, asc: bool) -> Vec<SortExpr> {
4394        self.ctx
4395            .field_columns
4396            .iter()
4397            .map(|col| DfExpr::Column(Column::from_name(col)).sort(asc, true))
4398            .collect::<Vec<_>>()
4399    }
4400
4401    fn create_sort_exprs_by_tags(
4402        func: &str,
4403        tags: Vec<DfExpr>,
4404        asc: bool,
4405    ) -> Result<Vec<SortExpr>> {
4406        ensure!(
4407            !tags.is_empty(),
4408            FunctionInvalidArgumentSnafu { fn_name: func }
4409        );
4410
4411        tags.iter()
4412            .map(|col| match col {
4413                DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => {
4414                    Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false))
4415                }
4416                other => UnexpectedPlanExprSnafu {
4417                    desc: format!("expected label string literal, but found {:?}", other),
4418                }
4419                .fail(),
4420            })
4421            .collect::<Result<Vec<_>>>()
4422    }
4423
4424    fn create_empty_values_filter_expr(&self, preserve_any_value: bool) -> Result<DfExpr> {
4425        let mut exprs = Vec::with_capacity(self.ctx.field_columns.len());
4426        for value in &self.ctx.field_columns {
4427            let expr = DfExpr::Column(Column::from_name(value)).is_not_null();
4428            exprs.push(expr);
4429        }
4430
4431        // This error context should be computed lazily: the planner may set `ctx.table_name` to
4432        // `None` for derived expressions (e.g. after projecting the LHS of a vector-vector
4433        // comparison filter). Eagerly calling `table_ref()?` here can turn a valid plan into
4434        // a `TableNameNotFound` error even when predicate construction succeeds.
4435        let predicate = if preserve_any_value {
4436            disjunction(exprs)
4437        } else {
4438            conjunction(exprs)
4439        };
4440        predicate.with_context(|| ValueNotFoundSnafu {
4441            table: self
4442                .table_ref()
4443                .map(|t| t.to_quoted_string())
4444                .unwrap_or_else(|_| "unknown".to_string()),
4445        })
4446    }
4447
4448    /// Creates a set of DataFusion `DfExpr::AggregateFunction` expressions for each value column using the specified aggregate function.
4449    ///
4450    /// # Side Effects
4451    ///
4452    /// This method modifies the value columns in the context by replacing them with the new columns
4453    /// created by the aggregate function application.
4454    ///
4455    /// # Returns
4456    ///
4457    /// Returns a tuple of `(aggregate_expressions, previous_field_expressions)` where:
4458    /// - `aggregate_expressions`: Expressions that apply the aggregate function to the original fields
4459    /// - `previous_field_expressions`: Original field expressions before aggregation. This is non-empty
4460    ///   only when the operation is `count_values`, as this operation requires preserving the original
4461    ///   values for grouping.
4462    ///
4463    fn create_aggregate_exprs(
4464        &mut self,
4465        op: TokenType,
4466        param: &Option<Box<PromExpr>>,
4467        input_plan: &LogicalPlan,
4468    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4469        let mixed_sample_columns =
4470            Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4471                .map(|(float, histogram)| (float.to_string(), histogram.to_string()));
4472        let is_group_agg = op.id() == token::T_GROUP;
4473        if is_group_agg && mixed_sample_columns.is_none() {
4474            ensure!(
4475                self.ctx.field_columns.len() == 1,
4476                MultiFieldsNotSupportedSnafu {
4477                    operator: "group()"
4478                }
4479            );
4480        }
4481
4482        if let Some((float, histogram)) = mixed_sample_columns {
4483            return self.create_mixed_aggregate_exprs(op, param, &float, &histogram);
4484        }
4485
4486        if self.all_field_columns_are_native_histograms(input_plan.schema()) {
4487            return self.create_native_histogram_aggregate_exprs(op, input_plan);
4488        }
4489
4490        // perform aggregate operation to each value column
4491        let exprs = self
4492            .ctx
4493            .field_columns
4494            .iter()
4495            .map(|col| {
4496                Self::create_numeric_aggregate_expr(
4497                    op,
4498                    param,
4499                    DfExpr::Column(Column::from_name(col)),
4500                )
4501            })
4502            .collect::<Result<Vec<_>>>()?;
4503
4504        // if the aggregator is `count_values`, it must be grouped by current fields.
4505        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4506            let prev_field_exprs: Vec<_> = self
4507                .ctx
4508                .field_columns
4509                .iter()
4510                .map(|col| DfExpr::Column(Column::from_name(col)))
4511                .collect();
4512
4513            ensure!(
4514                self.ctx.field_columns.len() == 1,
4515                UnsupportedExprSnafu {
4516                    name: "count_values on multi-value input"
4517                }
4518            );
4519
4520            prev_field_exprs
4521        } else {
4522            vec![]
4523        };
4524
4525        // update value column name according to the aggregators,
4526        let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
4527
4528        let normalized_exprs =
4529            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4530        for expr in normalized_exprs {
4531            new_field_columns.push(expr.schema_name().to_string());
4532        }
4533        self.ctx.field_columns = new_field_columns;
4534
4535        Ok((exprs, prev_field_exprs))
4536    }
4537
4538    fn create_numeric_aggregate_expr(
4539        op: TokenType,
4540        param: &Option<Box<PromExpr>>,
4541        input: DfExpr,
4542    ) -> Result<DfExpr> {
4543        let expr = match op.id() {
4544            token::T_SUM => sum_udaf().call(vec![input]),
4545            token::T_QUANTILE => {
4546                let q = Self::get_param_as_literal_expr(
4547                    param.as_deref(),
4548                    Some(op),
4549                    Some(ArrowDataType::Float64),
4550                )?;
4551                quantile_udaf().call(vec![q, input])
4552            }
4553            token::T_AVG => avg_udaf().call(vec![input]),
4554            token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]),
4555            token::T_MIN => min_udaf().call(vec![input]),
4556            token::T_MAX => max_udaf().call(vec![input]),
4557            // PromQL's `group()` aggregator produces 1 for each group.
4558            // Use `max(1.0)` (per-group) to match semantics and output type (Float64).
4559            token::T_GROUP => max_udaf().call(vec![lit(1_f64)]),
4560            token::T_STDDEV => stddev_pop_udaf().call(vec![input]),
4561            token::T_STDVAR => var_pop_udaf().call(vec![input]),
4562            token::T_TOPK | token::T_BOTTOMK => {
4563                return UnsupportedExprSnafu {
4564                    name: format!("{op:?}"),
4565                }
4566                .fail();
4567            }
4568            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4569        };
4570        Ok(expr)
4571    }
4572
4573    fn create_mixed_aggregate_exprs(
4574        &mut self,
4575        op: TokenType,
4576        param: &Option<Box<PromExpr>>,
4577        float_column: &str,
4578        histogram_column: &str,
4579    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4580        let float_input = DfExpr::Column(Column::from_name(float_column));
4581        let histogram_input = DfExpr::Column(Column::from_name(histogram_column));
4582        let float_count = count_udaf().call(vec![float_input.clone()]);
4583        let histogram_count = count_udaf().call(vec![histogram_input.clone()]);
4584        let mixed_sample_value = || {
4585            DfExpr::ScalarFunction(ScalarFunction {
4586                func: coalesce(),
4587                args: vec![
4588                    DfExpr::ScalarFunction(ScalarFunction {
4589                        func: Arc::new(PromqlFloatToString::scalar_udf()),
4590                        args: vec![float_input.clone()],
4591                    }),
4592                    DfExpr::ScalarFunction(ScalarFunction {
4593                        func: Arc::new(NativeHistogramToString::scalar_udf()),
4594                        args: vec![histogram_input.clone()],
4595                    }),
4596                ],
4597            })
4598        };
4599
4600        let (exprs, prev_field_exprs, field_columns) = match op.id() {
4601            token::T_SUM | token::T_AVG => (
4602                vec![
4603                    Self::create_numeric_aggregate_expr(op, param, float_input)?
4604                        .alias(float_column),
4605                    self.create_native_histogram_aggregate_expr(op, histogram_column)?,
4606                    float_count.alias(Self::mixed_sample_count_name(float_column)),
4607                    histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4608                ],
4609                vec![],
4610                vec![float_column.to_string(), histogram_column.to_string()],
4611            ),
4612            token::T_COUNT => {
4613                let present = when(
4614                    float_input
4615                        .clone()
4616                        .is_not_null()
4617                        .or(histogram_input.clone().is_not_null()),
4618                    lit(1_i64),
4619                )
4620                .otherwise(lit(ScalarValue::Int64(None)))
4621                .context(DataFusionPlanningSnafu)?;
4622                (
4623                    vec![count_udaf().call(vec![present]).alias(float_column)],
4624                    vec![],
4625                    vec![float_column.to_string()],
4626                )
4627            }
4628            token::T_GROUP => (
4629                vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)],
4630                vec![],
4631                vec![float_column.to_string()],
4632            ),
4633            token::T_COUNT_VALUES => {
4634                let value = mixed_sample_value();
4635                (
4636                    vec![count_udaf().call(vec![value.clone()]).alias(float_column)],
4637                    vec![value],
4638                    vec![float_column.to_string()],
4639                )
4640            }
4641            token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => {
4642                (
4643                    vec![
4644                        Self::create_numeric_aggregate_expr(op, param, float_input)?
4645                            .alias(float_column),
4646                        histogram_count.alias(Self::mixed_sample_count_name(histogram_column)),
4647                    ],
4648                    vec![],
4649                    vec![float_column.to_string()],
4650                )
4651            }
4652            token::T_TOPK | token::T_BOTTOMK => {
4653                return UnsupportedExprSnafu {
4654                    name: format!("{op:?}"),
4655                }
4656                .fail();
4657            }
4658            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4659        };
4660
4661        self.ctx.field_columns = field_columns;
4662        Ok((exprs, prev_field_exprs))
4663    }
4664
4665    fn mixed_sample_count_column(column: &str) -> DfExpr {
4666        DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column)))
4667    }
4668
4669    fn mixed_sample_count_name(column: &str) -> String {
4670        format!("__promql_sample_count({column})")
4671    }
4672
4673    fn mixed_aggregate_filter_expr(
4674        &self,
4675        op: TokenType,
4676        float_column: &str,
4677        histogram_column: &str,
4678    ) -> Result<DfExpr> {
4679        let float_count = Self::mixed_sample_count_column(float_column);
4680        let histogram_count = Self::mixed_sample_count_column(histogram_column);
4681        let mixed = float_count
4682            .clone()
4683            .gt(lit(0_i64))
4684            .and(histogram_count.clone().gt(lit(0_i64)));
4685        let drop_mixed = DfExpr::ScalarFunction(ScalarFunction {
4686            func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
4687                format!(
4688                    "{op}: dropped aggregation result containing both float and native histogram samples"
4689                ),
4690                self.promql_annotations.clone(),
4691            )),
4692            args: vec![float_count, histogram_count],
4693        });
4694
4695        when(mixed, drop_mixed)
4696            .otherwise(lit(true))
4697            .context(DataFusionPlanningSnafu)
4698    }
4699
4700    fn mixed_ignored_histogram_filter_expr(
4701        &self,
4702        op: TokenType,
4703        histogram_column: &str,
4704    ) -> Result<DfExpr> {
4705        let histogram_count = Self::mixed_sample_count_column(histogram_column);
4706        let has_histograms = histogram_count.clone().gt(lit(0_i64));
4707        let record_info = DfExpr::ScalarFunction(ScalarFunction {
4708            func: Arc::new(NativeHistogramDrop::bool_true_udf(
4709                format!(
4710                    "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4711                ),
4712                self.promql_annotations.clone(),
4713            )),
4714            args: vec![histogram_count],
4715        });
4716
4717        when(has_histograms, record_info)
4718            .otherwise(lit(true))
4719            .context(DataFusionPlanningSnafu)
4720    }
4721
4722    fn create_native_histogram_aggregate_expr(
4723        &self,
4724        op: TokenType,
4725        column: &str,
4726    ) -> Result<DfExpr> {
4727        let input = DfExpr::Column(Column::from_name(column));
4728        let expr = match op.id() {
4729            token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector(
4730                self.promql_annotations.clone(),
4731            ))
4732            .call(vec![input])
4733            .alias(column),
4734            token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector(
4735                self.promql_annotations.clone(),
4736            ))
4737            .call(vec![input])
4738            .alias(column),
4739            token::T_COUNT_VALUES | token::T_COUNT => {
4740                count_udaf().call(vec![input]).alias(column)
4741            }
4742            token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column),
4743            token::T_MIN
4744            | token::T_MAX
4745            | token::T_STDDEV
4746            | token::T_STDVAR
4747            | token::T_QUANTILE
4748            | token::T_TOPK
4749            | token::T_BOTTOMK => sum_udaf()
4750                .call(vec![DfExpr::ScalarFunction(ScalarFunction {
4751                    func: Arc::new(NativeHistogramDrop::float_null_udf(
4752                        format!(
4753                            "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
4754                        ),
4755                        self.promql_annotations.clone(),
4756                    )),
4757                    args: vec![input],
4758                })])
4759                .alias(column),
4760            _ => return UnexpectedTokenSnafu { token: op }.fail(),
4761        };
4762        Ok(expr)
4763    }
4764
4765    fn create_native_histogram_aggregate_exprs(
4766        &mut self,
4767        op: TokenType,
4768        input_plan: &LogicalPlan,
4769    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
4770        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
4771            ensure!(
4772                self.ctx.field_columns.len() == 1,
4773                UnsupportedExprSnafu {
4774                    name: "count_values on multi-value input"
4775                }
4776            );
4777            self.ctx
4778                .field_columns
4779                .iter()
4780                .map(|col| {
4781                    DfExpr::ScalarFunction(ScalarFunction {
4782                        func: Arc::new(NativeHistogramToString::scalar_udf()),
4783                        args: vec![DfExpr::Column(Column::from_name(col))],
4784                    })
4785                })
4786                .collect::<Vec<_>>()
4787        } else {
4788            vec![]
4789        };
4790
4791        let exprs = self
4792            .ctx
4793            .field_columns
4794            .iter()
4795            .map(|col| self.create_native_histogram_aggregate_expr(op, col))
4796            .collect::<Result<Vec<_>>>()?;
4797
4798        let normalized_exprs =
4799            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4800        self.ctx.field_columns = normalized_exprs
4801            .into_iter()
4802            .map(|expr| expr.schema_name().to_string())
4803            .collect();
4804
4805        Ok((exprs, prev_field_exprs))
4806    }
4807
4808    fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
4809        let param = param
4810            .as_deref()
4811            .with_context(|| FunctionInvalidArgumentSnafu {
4812                fn_name: op.to_string(),
4813            })?;
4814        let PromExpr::StringLiteral(StringLiteral { val }) = param else {
4815            return FunctionInvalidArgumentSnafu {
4816                fn_name: op.to_string(),
4817            }
4818            .fail();
4819        };
4820
4821        Ok(val)
4822    }
4823
4824    fn get_param_as_literal_expr(
4825        param: Option<&PromExpr>,
4826        op: Option<TokenType>,
4827        expected_type: Option<ArrowDataType>,
4828    ) -> Result<DfExpr> {
4829        let prom_param = param.with_context(|| {
4830            if let Some(op) = op {
4831                FunctionInvalidArgumentSnafu {
4832                    fn_name: op.to_string(),
4833                }
4834            } else {
4835                FunctionInvalidArgumentSnafu {
4836                    fn_name: "unknown".to_string(),
4837                }
4838            }
4839        })?;
4840
4841        let expr = Self::try_build_literal_expr(prom_param).with_context(|| {
4842            if let Some(op) = op {
4843                FunctionInvalidArgumentSnafu {
4844                    fn_name: op.to_string(),
4845                }
4846            } else {
4847                FunctionInvalidArgumentSnafu {
4848                    fn_name: "unknown".to_string(),
4849                }
4850            }
4851        })?;
4852
4853        // check if the type is expected
4854        if let Some(expected_type) = expected_type {
4855            // literal should not have reference to column
4856            let expr_type = expr
4857                .get_type(&DFSchema::empty())
4858                .context(DataFusionPlanningSnafu)?;
4859            if expected_type != expr_type {
4860                return FunctionInvalidArgumentSnafu {
4861                    fn_name: format!("expected {expected_type:?}, but found {expr_type:?}"),
4862                }
4863                .fail();
4864            }
4865        }
4866
4867        Ok(expr)
4868    }
4869
4870    /// Create [DfExpr::WindowFunction] expr for each value column with given window function.
4871    ///
4872    fn create_window_exprs(
4873        &mut self,
4874        op: TokenType,
4875        group_exprs: Vec<DfExpr>,
4876        input_plan: &LogicalPlan,
4877    ) -> Result<Vec<DfExpr>> {
4878        ensure!(
4879            self.ctx.field_columns.len() == 1,
4880            UnsupportedExprSnafu {
4881                name: "topk or bottomk on multi-value input"
4882            }
4883        );
4884
4885        assert!(matches!(op.id(), token::T_TOPK | token::T_BOTTOMK));
4886
4887        let asc = matches!(op.id(), token::T_BOTTOMK);
4888
4889        let tag_sort_exprs = self
4890            .create_tag_column_exprs()?
4891            .into_iter()
4892            .map(|expr| expr.sort(asc, true));
4893
4894        // perform window operation to each value column
4895        let exprs: Vec<DfExpr> = self
4896            .ctx
4897            .field_columns
4898            .iter()
4899            .map(|col| {
4900                let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1);
4901                // Order by value in the specific order
4902                sort_exprs.push(DfExpr::Column(Column::from(col)).sort(asc, true));
4903                // Then tags if the values are equal,
4904                // Try to ensure the relative stability of the output results.
4905                sort_exprs.extend(tag_sort_exprs.clone());
4906
4907                DfExpr::WindowFunction(Box::new(WindowFunction {
4908                    fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
4909                    params: WindowFunctionParams {
4910                        args: vec![],
4911                        partition_by: group_exprs.clone(),
4912                        order_by: sort_exprs,
4913                        window_frame: WindowFrame::new(Some(true)),
4914                        null_treatment: None,
4915                        distinct: false,
4916                        filter: None,
4917                    },
4918                }))
4919            })
4920            .collect();
4921
4922        let normalized_exprs =
4923            normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
4924        Ok(normalized_exprs)
4925    }
4926
4927    /// Create a classic, native, or mixed histogram helper plan.
4928    async fn create_histogram_plan(
4929        &mut self,
4930        function_name: &str,
4931        args: &PromFunctionArgs,
4932        query_engine_state: &QueryEngineState,
4933    ) -> Result<LogicalPlan> {
4934        let float_literal = |param: &PromExpr| -> Result<f64> {
4935            let value = (|| {
4936                let expr = Self::get_param_as_literal_expr(
4937                    Some(param),
4938                    None,
4939                    Some(ArrowDataType::Float64),
4940                )
4941                .ok()?;
4942                let simplifier = ExprSimplifier::new(SimplifyContext::default());
4943                let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?;
4944                let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else {
4945                    return None;
4946                };
4947                let ScalarValue::Float64(Some(value)) =
4948                    value.cast_to(&ArrowDataType::Float64).ok()?
4949                else {
4950                    return None;
4951                };
4952                Some(value)
4953            })()
4954            .with_context(|| FunctionInvalidArgumentSnafu {
4955                fn_name: function_name.to_string(),
4956            })?;
4957            Ok(value)
4958        };
4959        let (function, input) = match (function_name, args.args.as_slice()) {
4960            (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => (
4961                HistogramFoldOperation::Quantile(float_literal(quantile)?.into()),
4962                input.as_ref().clone(),
4963            ),
4964            (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => (
4965                HistogramFoldOperation::Fraction {
4966                    lower: float_literal(lower)?.into(),
4967                    upper: float_literal(upper)?.into(),
4968                },
4969                input.as_ref().clone(),
4970            ),
4971            _ => {
4972                return FunctionInvalidArgumentSnafu {
4973                    fn_name: function_name.to_string(),
4974                }
4975                .fail();
4976            }
4977        };
4978
4979        let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?;
4980        // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a
4981        // stable series identifier anymore. HistogramFold must not treat it as a label column.
4982        let input_plan = self.strip_tsid_column(input_plan)?;
4983        self.ctx.use_tsid = false;
4984
4985        if let Some((float_field, histogram_field)) =
4986            Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns)
4987                .map(|(float, histogram)| (float.to_string(), histogram.to_string()))
4988        {
4989            if self.ctx.has_le_tag() {
4990                return self.create_mixed_histogram_plan(
4991                    function,
4992                    input_plan,
4993                    float_field,
4994                    histogram_field,
4995                );
4996            }
4997            self.ctx.field_columns = vec![histogram_field];
4998        }
4999        if self.all_field_columns_are_native_histograms(input_plan.schema()) {
5000            return self.create_native_histogram_plan(function, input_plan);
5001        }
5002
5003        if !self.ctx.has_le_tag() {
5004            // Return empty result instead of error when 'le' column is not found
5005            // This handles the case when histogram metrics don't exist
5006            return Ok(LogicalPlan::EmptyRelation(
5007                datafusion::logical_expr::EmptyRelation {
5008                    produce_one_row: false,
5009                    schema: input_plan.schema().clone(),
5010                },
5011            ));
5012        }
5013        let time_index_column =
5014            self.ctx
5015                .time_index_column
5016                .clone()
5017                .with_context(|| TimeIndexNotFoundSnafu {
5018                    table: self.ctx.table_name.clone().unwrap_or_default(),
5019                })?;
5020        // FIXME(ruihang): support multi fields
5021        let field_column = self
5022            .ctx
5023            .field_columns
5024            .first()
5025            .with_context(|| FunctionInvalidArgumentSnafu {
5026                fn_name: function.function_name().to_string(),
5027            })?
5028            .clone();
5029        // remove le column from tag columns
5030        self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME);
5031
5032        let fold = HistogramFold::new_with_operation(
5033            LE_COLUMN_NAME.to_string(),
5034            field_column,
5035            time_index_column,
5036            function,
5037            None,
5038            input_plan,
5039        )
5040        .context(DataFusionPlanningSnafu)?;
5041        Ok(LogicalPlan::Extension(Extension {
5042            node: Arc::new(fold),
5043        }))
5044    }
5045
5046    fn create_native_histogram_expr(
5047        &self,
5048        function: HistogramFoldOperation,
5049        field_column: &str,
5050    ) -> DfExpr {
5051        let field = DfExpr::Column(Column::from_name(field_column));
5052        let (func, args) = match function {
5053            HistogramFoldOperation::Quantile(quantile) => (
5054                Arc::new(NativeHistogramQuantile::scalar_udf_with_collector(
5055                    self.promql_annotations.clone(),
5056                )),
5057                vec![field, lit(f64::from(quantile))],
5058            ),
5059            HistogramFoldOperation::Fraction { lower, upper } => (
5060                Arc::new(NativeHistogramFraction::scalar_udf_with_collector(
5061                    self.promql_annotations.clone(),
5062                )),
5063                vec![field, lit(f64::from(lower)), lit(f64::from(upper))],
5064            ),
5065        };
5066        DfExpr::ScalarFunction(ScalarFunction { func, args })
5067    }
5068
5069    fn create_native_histogram_plan(
5070        &mut self,
5071        function: HistogramFoldOperation,
5072        input_plan: LogicalPlan,
5073    ) -> Result<LogicalPlan> {
5074        ensure!(
5075            self.ctx.field_columns.len() == 1,
5076            MultiFieldsNotSupportedSnafu {
5077                operator: function.function_name()
5078            },
5079        );
5080
5081        let field_column = self.ctx.field_columns[0].clone();
5082        let function_expr = self.create_native_histogram_expr(function, &field_column);
5083        let display_name = function_expr.schema_name().to_string();
5084        self.ctx.field_columns = vec![display_name.clone()];
5085
5086        let project_exprs = std::iter::once(self.create_time_index_column_expr()?)
5087            .chain(std::iter::once(function_expr.alias(display_name)))
5088            .chain(self.create_tag_column_exprs()?)
5089            .collect::<Vec<_>>();
5090
5091        LogicalPlanBuilder::from(input_plan)
5092            .project(project_exprs)
5093            .context(DataFusionPlanningSnafu)?
5094            .filter(self.create_empty_values_filter_expr(false)?)
5095            .context(DataFusionPlanningSnafu)?
5096            .build()
5097            .context(DataFusionPlanningSnafu)
5098    }
5099
5100    fn create_mixed_histogram_plan(
5101        &mut self,
5102        function: HistogramFoldOperation,
5103        input_plan: LogicalPlan,
5104        float_field: String,
5105        histogram_field: String,
5106    ) -> Result<LogicalPlan> {
5107        let time_index_column =
5108            self.ctx
5109                .time_index_column
5110                .clone()
5111                .with_context(|| TimeIndexNotFoundSnafu {
5112                    table: self.ctx.table_name.clone().unwrap_or_default(),
5113                })?;
5114        let tag_columns = self.ctx.tag_columns.clone();
5115        let folded = HistogramFold::new_with_operation(
5116            LE_COLUMN_NAME.to_string(),
5117            float_field.clone(),
5118            time_index_column.clone(),
5119            function,
5120            Some(histogram_field.clone()),
5121            input_plan,
5122        )
5123        .context(DataFusionPlanningSnafu)?;
5124        let record_collision = DfExpr::ScalarFunction(ScalarFunction {
5125            func: Arc::new(NativeHistogramDrop::warning_bool_false_udf(
5126                "vector contains a mix of classic and native histograms".to_string(),
5127                self.promql_annotations.clone(),
5128            )),
5129            args: vec![col(&float_field), col(&histogram_field)],
5130        });
5131        let keep = when(
5132            col(&float_field)
5133                .is_not_null()
5134                .and(col(&histogram_field).is_not_null()),
5135            record_collision,
5136        )
5137        .otherwise(lit(true))
5138        .context(DataFusionPlanningSnafu)?;
5139
5140        let native_expr = self.create_native_histogram_expr(function, &histogram_field);
5141        let output_field = native_expr.schema_name().to_string();
5142        let value = DfExpr::ScalarFunction(ScalarFunction {
5143            func: coalesce(),
5144            args: vec![col(&float_field), native_expr],
5145        });
5146        self.ctx.field_columns = vec![output_field.clone()];
5147        LogicalPlanBuilder::from(LogicalPlan::Extension(Extension {
5148            node: Arc::new(folded),
5149        }))
5150        .filter(keep)
5151        .context(DataFusionPlanningSnafu)?
5152        .project(
5153            std::iter::once(col(&time_index_column))
5154                .chain(std::iter::once(value.alias(output_field)))
5155                .chain(tag_columns.iter().map(col)),
5156        )
5157        .context(DataFusionPlanningSnafu)?
5158        .build()
5159        .context(DataFusionPlanningSnafu)
5160    }
5161
5162    /// Create a [SPECIAL_VECTOR_FUNCTION] plan
5163    async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result<LogicalPlan> {
5164        if args.args.len() != 1 {
5165            return FunctionInvalidArgumentSnafu {
5166                fn_name: SPECIAL_VECTOR_FUNCTION.to_string(),
5167            }
5168            .fail();
5169        }
5170        let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?;
5171
5172        // reuse `SPECIAL_TIME_FUNCTION` as name of time index column
5173        self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string());
5174        self.ctx.reset_table_name_and_schema();
5175        self.ctx.tag_columns = vec![];
5176        self.ctx.field_columns = vec![greptime_value().to_string()];
5177        Ok(LogicalPlan::Extension(Extension {
5178            node: Arc::new(
5179                EmptyMetric::new(
5180                    self.ctx.start,
5181                    self.ctx.end,
5182                    self.ctx.interval,
5183                    SPECIAL_TIME_FUNCTION.to_string(),
5184                    greptime_value().to_string(),
5185                    Some(lit),
5186                )
5187                .context(DataFusionPlanningSnafu)?,
5188            ),
5189        }))
5190    }
5191
5192    /// Create a [SCALAR_FUNCTION] plan
5193    async fn create_scalar_plan(
5194        &mut self,
5195        args: &PromFunctionArgs,
5196        query_engine_state: &QueryEngineState,
5197    ) -> Result<LogicalPlan> {
5198        ensure!(
5199            args.len() == 1,
5200            FunctionInvalidArgumentSnafu {
5201                fn_name: SCALAR_FUNCTION
5202            }
5203        );
5204        let input = self
5205            .prom_expr_to_plan(&args.args[0], query_engine_state)
5206            .await?;
5207        let input_schema = input.schema().clone();
5208        let alternative_samples =
5209            Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns);
5210        let histogram_fields = self
5211            .ctx
5212            .field_columns
5213            .iter()
5214            .filter(|field| Self::field_column_is_native_histogram(&input_schema, field))
5215            .count();
5216        ensure!(
5217            self.ctx.field_columns.len() == 1 || alternative_samples,
5218            MultiFieldsNotSupportedSnafu {
5219                operator: SCALAR_FUNCTION
5220            },
5221        );
5222        let scalar_field = self
5223            .ctx
5224            .field_columns
5225            .iter()
5226            .find(|field| !Self::field_column_is_native_histogram(&input_schema, field))
5227            .or_else(|| self.ctx.field_columns.first())
5228            .cloned()
5229            .with_context(|| FunctionInvalidArgumentSnafu {
5230                fn_name: SCALAR_FUNCTION,
5231            })?;
5232        let input = if histogram_fields == self.ctx.field_columns.len() {
5233            // scalar() ignores histogram samples. An empty input makes ScalarCalculate emit NaN
5234            // for every evaluation timestamp without attempting a Struct-to-Float64 cast.
5235            LogicalPlanBuilder::from(input)
5236                .filter(lit(false))
5237                .context(DataFusionPlanningSnafu)?
5238                .build()
5239                .context(DataFusionPlanningSnafu)?
5240        } else if histogram_fields > 0 {
5241            // A mixed vector contributes only its float samples to scalar().
5242            LogicalPlanBuilder::from(input)
5243                .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null())
5244                .context(DataFusionPlanningSnafu)?
5245                .build()
5246                .context(DataFusionPlanningSnafu)?
5247        } else {
5248            input
5249        };
5250        let scalar_plan = LogicalPlan::Extension(Extension {
5251            node: Arc::new(
5252                ScalarCalculate::new(
5253                    self.ctx.start,
5254                    self.ctx.end,
5255                    self.ctx.interval,
5256                    input,
5257                    self.ctx.time_index_column.as_ref().unwrap(),
5258                    &self.ctx.tag_columns,
5259                    &scalar_field,
5260                    self.ctx.table_name.as_deref(),
5261                )
5262                .context(PromqlPlanNodeSnafu)?,
5263            ),
5264        });
5265        // scalar plan have no tag columns
5266        self.ctx.tag_columns.clear();
5267        self.ctx.field_columns.clear();
5268        self.ctx
5269            .field_columns
5270            .push(scalar_plan.schema().field(1).name().clone());
5271        Ok(scalar_plan)
5272    }
5273
5274    /// Create a [SPECIAL_ABSENT_FUNCTION] plan
5275    async fn create_absent_plan(
5276        &mut self,
5277        args: &PromFunctionArgs,
5278        query_engine_state: &QueryEngineState,
5279    ) -> Result<LogicalPlan> {
5280        if args.args.len() != 1 {
5281            return FunctionInvalidArgumentSnafu {
5282                fn_name: SPECIAL_ABSENT_FUNCTION.to_string(),
5283            }
5284            .fail();
5285        }
5286        let input = self
5287            .prom_expr_to_plan(&args.args[0], query_engine_state)
5288            .await?;
5289
5290        let time_index_expr = self.create_time_index_column_expr()?;
5291        let first_field_expr =
5292            self.create_field_column_exprs()?
5293                .pop()
5294                .with_context(|| ValueNotFoundSnafu {
5295                    table: self.ctx.table_name.clone().unwrap_or_default(),
5296                })?;
5297        let first_value_expr = first_value(first_field_expr, vec![]);
5298
5299        let ordered_aggregated_input = LogicalPlanBuilder::from(input)
5300            .aggregate(
5301                vec![time_index_expr.clone()],
5302                vec![first_value_expr.clone()],
5303            )
5304            .context(DataFusionPlanningSnafu)?
5305            .sort(vec![time_index_expr.sort(true, false)])
5306            .context(DataFusionPlanningSnafu)?
5307            .build()
5308            .context(DataFusionPlanningSnafu)?;
5309
5310        let fake_labels = self
5311            .ctx
5312            .selector_matcher
5313            .iter()
5314            .filter_map(|matcher| match matcher.op {
5315                MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())),
5316                _ => None,
5317            })
5318            .collect::<Vec<_>>();
5319
5320        // Create the absent plan
5321        let absent_plan = LogicalPlan::Extension(Extension {
5322            node: Arc::new(
5323                Absent::try_new(
5324                    self.ctx.start,
5325                    self.ctx.end,
5326                    self.ctx.interval,
5327                    self.ctx.time_index_column.as_ref().unwrap().clone(),
5328                    self.ctx.field_columns[0].clone(),
5329                    fake_labels,
5330                    ordered_aggregated_input,
5331                )
5332                .context(DataFusionPlanningSnafu)?,
5333            ),
5334        });
5335
5336        Ok(absent_plan)
5337    }
5338
5339    /// Try to build a DataFusion Literal Expression from PromQL Expr, return
5340    /// `None` if the input is not a literal expression.
5341    fn try_build_literal_expr(expr: &PromExpr) -> Option<DfExpr> {
5342        match expr {
5343            PromExpr::NumberLiteral(NumberLiteral { val }) => Some(val.lit()),
5344            PromExpr::StringLiteral(StringLiteral { val }) => Some(val.lit()),
5345            PromExpr::VectorSelector(_)
5346            | PromExpr::MatrixSelector(_)
5347            | PromExpr::Extension(_)
5348            | PromExpr::Aggregate(_)
5349            | PromExpr::Subquery(_) => None,
5350            PromExpr::Call(Call { func, .. }) => {
5351                if func.name == SPECIAL_TIME_FUNCTION {
5352                    // For time() function, don't treat it as a literal
5353                    // Let it be handled as a regular function call
5354                    None
5355                } else {
5356                    None
5357                }
5358            }
5359            PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr),
5360            PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new(
5361                Self::try_build_literal_expr(expr)?,
5362            ))),
5363            PromExpr::Binary(PromBinaryExpr {
5364                lhs,
5365                rhs,
5366                op,
5367                modifier,
5368            }) => {
5369                let lhs = Self::try_build_literal_expr(lhs)?;
5370                let rhs = Self::try_build_literal_expr(rhs)?;
5371                let is_comparison_op = Self::is_token_a_comparison_op(*op);
5372                let expr_builder = Self::prom_token_to_binary_expr_builder(*op).ok()?;
5373                let expr = expr_builder(lhs, rhs).ok()?;
5374
5375                let should_return_bool = if let Some(m) = modifier {
5376                    m.return_bool
5377                } else {
5378                    false
5379                };
5380                if is_comparison_op && should_return_bool {
5381                    Some(DfExpr::Cast(Cast {
5382                        expr: Box::new(expr),
5383                        data_type: ArrowDataType::Float64,
5384                    }))
5385                } else {
5386                    Some(expr)
5387                }
5388            }
5389        }
5390    }
5391
5392    fn try_build_special_time_expr_with_context(&self, expr: &PromExpr) -> Option<DfExpr> {
5393        match expr {
5394            PromExpr::Call(Call { func, .. }) => {
5395                if func.name == SPECIAL_TIME_FUNCTION
5396                    && let Some(time_index_col) = self.ctx.time_index_column.as_ref()
5397                {
5398                    Some(build_special_time_expr(time_index_col))
5399                } else {
5400                    None
5401                }
5402            }
5403            _ => None,
5404        }
5405    }
5406
5407    fn native_histogram_binary_expr(
5408        token: TokenType,
5409        lhs: DfExpr,
5410        lhs_is_histogram: bool,
5411        rhs: DfExpr,
5412        rhs_is_histogram: bool,
5413        filter_context: bool,
5414        promql_annotations: Option<PromqlAnnotationCollector>,
5415    ) -> Result<Option<DfExpr>> {
5416        if !lhs_is_histogram && !rhs_is_histogram {
5417            return Ok(None);
5418        }
5419
5420        let scalar_fn = |func: ScalarUdfDef, args| {
5421            DfExpr::ScalarFunction(ScalarFunction {
5422                func: Arc::new(func),
5423                args,
5424            })
5425        };
5426        let invalid_expr = || {
5427            let message = format!(
5428                "{}: dropped native histogram samples because this binary operation is not supported for native histograms",
5429                token
5430            );
5431            let func = if filter_context {
5432                NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone())
5433            } else {
5434                NativeHistogramDrop::float_null_udf(message, promql_annotations.clone())
5435            };
5436            let args = vec![lhs.clone(), rhs.clone()];
5437            scalar_fn(func, args)
5438        };
5439
5440        let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5441            (token::T_ADD, true, true) => scalar_fn(
5442                NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()),
5443                vec![lhs, rhs],
5444            ),
5445            (token::T_SUB, true, true) => scalar_fn(
5446                NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()),
5447                vec![lhs, rhs],
5448            ),
5449            (token::T_MUL, true, false) => {
5450                scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs])
5451            }
5452            (token::T_MUL, false, true) => {
5453                scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs])
5454            }
5455            (token::T_DIV, true, false) => {
5456                scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs])
5457            }
5458            (token::T_EQLC, true, true) => {
5459                scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs])
5460            }
5461            (token::T_NEQ, true, true) => {
5462                scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs])
5463            }
5464            _ => invalid_expr(),
5465        };
5466
5467        Ok(Some(expr))
5468    }
5469
5470    /// Return a lambda to build binary expression from token.
5471    /// Because some binary operator are function in DataFusion like `atan2` or `^`.
5472    #[allow(clippy::type_complexity)]
5473    fn prom_token_to_binary_expr_builder(
5474        token: TokenType,
5475    ) -> Result<Box<dyn Fn(DfExpr, DfExpr) -> Result<DfExpr>>> {
5476        let cast_float = |expr| {
5477            if matches!(
5478                &expr,
5479                DfExpr::Cast(Cast {
5480                    data_type: ArrowDataType::Float64,
5481                    ..
5482                })
5483            ) || matches!(&expr, DfExpr::Literal(ScalarValue::Float64(_), _))
5484            {
5485                expr
5486            } else {
5487                DfExpr::Cast(Cast {
5488                    expr: Box::new(expr),
5489                    data_type: ArrowDataType::Float64,
5490                })
5491            }
5492        };
5493        match token.id() {
5494            token::T_ADD => Ok(Box::new(move |lhs, rhs| {
5495                Ok(cast_float(lhs) + cast_float(rhs))
5496            })),
5497            token::T_SUB => Ok(Box::new(move |lhs, rhs| {
5498                Ok(cast_float(lhs) - cast_float(rhs))
5499            })),
5500            token::T_MUL => Ok(Box::new(move |lhs, rhs| {
5501                Ok(cast_float(lhs) * cast_float(rhs))
5502            })),
5503            token::T_DIV => Ok(Box::new(move |lhs, rhs| {
5504                Ok(cast_float(lhs) / cast_float(rhs))
5505            })),
5506            token::T_MOD => Ok(Box::new(move |lhs: DfExpr, rhs| {
5507                Ok(cast_float(lhs) % cast_float(rhs))
5508            })),
5509            token::T_EQLC => Ok(Box::new(|lhs, rhs| Ok(lhs.eq(rhs)))),
5510            token::T_NEQ => Ok(Box::new(|lhs, rhs| Ok(lhs.not_eq(rhs)))),
5511            token::T_GTR => Ok(Box::new(|lhs, rhs| Ok(lhs.gt(rhs)))),
5512            token::T_LSS => Ok(Box::new(|lhs, rhs| Ok(lhs.lt(rhs)))),
5513            token::T_GTE => Ok(Box::new(|lhs, rhs| Ok(lhs.gt_eq(rhs)))),
5514            token::T_LTE => Ok(Box::new(|lhs, rhs| Ok(lhs.lt_eq(rhs)))),
5515            token::T_POW => Ok(Box::new(move |lhs, rhs| {
5516                Ok(DfExpr::ScalarFunction(ScalarFunction {
5517                    func: datafusion_functions::math::power(),
5518                    args: vec![cast_float(lhs), cast_float(rhs)],
5519                }))
5520            })),
5521            token::T_ATAN2 => Ok(Box::new(move |lhs, rhs| {
5522                Ok(DfExpr::ScalarFunction(ScalarFunction {
5523                    func: datafusion_functions::math::atan2(),
5524                    args: vec![cast_float(lhs), cast_float(rhs)],
5525                }))
5526            })),
5527            _ => UnexpectedTokenSnafu { token }.fail(),
5528        }
5529    }
5530
5531    /// Check if the given op is a [comparison operator](https://prometheus.io/docs/prometheus/latest/querying/operators/#comparison-binary-operators).
5532    fn is_token_a_comparison_op(token: TokenType) -> bool {
5533        matches!(
5534            token.id(),
5535            token::T_EQLC
5536                | token::T_NEQ
5537                | token::T_GTR
5538                | token::T_LSS
5539                | token::T_GTE
5540                | token::T_LTE
5541        )
5542    }
5543
5544    /// Check if the given op is a set operator (UNION, INTERSECT and EXCEPT in SQL).
5545    fn is_token_a_set_op(token: TokenType) -> bool {
5546        matches!(
5547            token.id(),
5548            token::T_LAND // INTERSECT
5549                | token::T_LOR // UNION
5550                | token::T_LUNLESS // EXCEPT
5551        )
5552    }
5553
5554    fn align_binary_field_columns<'a>(
5555        left_schema: &DFSchemaRef,
5556        right_schema: &DFSchemaRef,
5557        left_field_columns: &'a [String],
5558        right_field_columns: &'a [String],
5559        op: TokenType,
5560        left_is_scalar: bool,
5561        right_is_scalar: bool,
5562    ) -> (
5563        Vec<(String, Vec<BinaryFieldPair<'a>>)>,
5564        Vec<BinaryFieldPair<'a>>,
5565    ) {
5566        // Mixed vectors store mutually exclusive float and histogram samples in two columns.
5567        // Retain each valid sample combination and group expressions by their output lane.
5568        let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns);
5569        let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns);
5570        let alternative_alignment = match (left_alternative, right_alternative) {
5571            (Some(output_names), Some(_)) => Some((
5572                output_names,
5573                left_field_columns
5574                    .iter()
5575                    .flat_map(|left| right_field_columns.iter().map(move |right| (left, right)))
5576                    .collect::<Vec<_>>(),
5577            )),
5578            (Some(output_names), None) if right_field_columns.len() == 1 => Some((
5579                output_names,
5580                left_field_columns
5581                    .iter()
5582                    .map(|left| (left, &right_field_columns[0]))
5583                    .collect::<Vec<_>>(),
5584            )),
5585            (None, Some(output_names)) if left_field_columns.len() == 1 => Some((
5586                output_names,
5587                right_field_columns
5588                    .iter()
5589                    .map(|right| (&left_field_columns[0], right))
5590                    .collect::<Vec<_>>(),
5591            )),
5592            _ => None,
5593        };
5594        let mut invalid_pairs = Vec::new();
5595        if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment {
5596            let mut float_pairs = Vec::new();
5597            let mut histogram_pairs = Vec::new();
5598            for (left, right) in field_pairs {
5599                let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left);
5600                let right_is_histogram =
5601                    Self::field_column_is_native_histogram(right_schema, right);
5602                match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) {
5603                    Some(false) => float_pairs.push((left, right)),
5604                    Some(true) => histogram_pairs.push((left, right)),
5605                    None => invalid_pairs.push((left, right)),
5606                }
5607            }
5608            if !float_pairs.is_empty() || !histogram_pairs.is_empty() {
5609                return (
5610                    [
5611                        (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)),
5612                        (!histogram_pairs.is_empty())
5613                            .then(|| (histogram_output.to_string(), histogram_pairs)),
5614                    ]
5615                    .into_iter()
5616                    .flatten()
5617                    .collect(),
5618                    invalid_pairs,
5619                );
5620            }
5621        }
5622
5623        if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 {
5624            return (
5625                right_field_columns
5626                    .iter()
5627                    .map(|right| (right.clone(), vec![(&left_field_columns[0], right)]))
5628                    .collect(),
5629                invalid_pairs,
5630            );
5631        }
5632        if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 {
5633            return (
5634                left_field_columns
5635                    .iter()
5636                    .map(|left| (left.clone(), vec![(left, &right_field_columns[0])]))
5637                    .collect(),
5638                invalid_pairs,
5639            );
5640        }
5641
5642        (
5643            left_field_columns
5644                .iter()
5645                .zip(right_field_columns.iter())
5646                .map(|(left, right)| (left.clone(), vec![(left, right)]))
5647                .collect(),
5648            invalid_pairs,
5649        )
5650    }
5651
5652    fn binary_result_is_histogram(
5653        token: TokenType,
5654        lhs_is_histogram: bool,
5655        rhs_is_histogram: bool,
5656    ) -> Option<bool> {
5657        match (token.id(), lhs_is_histogram, rhs_is_histogram) {
5658            (_, false, false) => Some(false),
5659            (token::T_ADD | token::T_SUB, true, true)
5660            | (token::T_MUL, true, false)
5661            | (token::T_MUL, false, true)
5662            | (token::T_DIV, true, false) => Some(true),
5663            (token::T_EQLC | token::T_NEQ, true, true) => Some(false),
5664            _ => None,
5665        }
5666    }
5667
5668    fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
5669        plan.schema()
5670            .fields()
5671            .iter()
5672            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5673    }
5674
5675    fn is_empty_metric(plan: &LogicalPlan) -> bool {
5676        matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<EmptyMetric>())
5677    }
5678
5679    fn native_histogram_arrow_type() -> ArrowDataType {
5680        native_histogram_value_type().as_arrow_type()
5681    }
5682
5683    fn field_column_type<'a>(
5684        schema: &'a DFSchemaRef,
5685        field_column: &str,
5686    ) -> Option<&'a ArrowDataType> {
5687        schema
5688            .index_of_column_by_name(None, field_column)
5689            .map(|idx| schema.field(idx).data_type())
5690    }
5691
5692    fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool {
5693        Self::field_column_type(schema, field_column)
5694            .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type())
5695    }
5696
5697    fn field_columns_contain_native_histogram(
5698        schema: &DFSchemaRef,
5699        field_columns: &[String],
5700    ) -> bool {
5701        field_columns
5702            .iter()
5703            .any(|field| Self::field_column_is_native_histogram(schema, field))
5704    }
5705
5706    fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5707        Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5708            matches!(
5709                data_type,
5710                ArrowDataType::Dictionary(key_type, value_type)
5711                    if key_type.as_ref() == &ArrowDataType::Int64
5712                        && value_type.as_ref() == &ArrowDataType::Float64
5713            )
5714        })
5715    }
5716
5717    fn field_columns_are_alternative_samples(
5718        schema: &DFSchemaRef,
5719        field_columns: &[String],
5720    ) -> bool {
5721        Self::alternative_sample_columns(schema, field_columns).is_some()
5722    }
5723
5724    fn alternative_sample_columns<'a>(
5725        schema: &DFSchemaRef,
5726        field_columns: &'a [String],
5727    ) -> Option<(&'a str, &'a str)> {
5728        if field_columns.len() != 2 {
5729            return None;
5730        }
5731
5732        let canonical_float = field_columns.iter().find(|field| {
5733            field.as_str() == greptime_value()
5734                && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5735                    || Self::field_column_is_float_range(schema, field))
5736        });
5737        let canonical_histogram = field_columns.iter().find(|field| {
5738            field.as_str() == greptime_native_histogram()
5739                && (Self::field_column_is_native_histogram(schema, field)
5740                    || Self::field_column_is_native_histogram_range(schema, field))
5741        });
5742        if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) {
5743            return Some((float, histogram));
5744        }
5745
5746        let float = field_columns.iter().find(|field| {
5747            field.starts_with(OR_FLOAT_FIELD_PREFIX)
5748                && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64)
5749                    || Self::field_column_is_float_range(schema, field))
5750        })?;
5751        let histogram = field_columns.iter().find(|field| {
5752            field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)
5753                && (Self::field_column_is_native_histogram(schema, field)
5754                    || Self::field_column_is_native_histogram_range(schema, field))
5755        })?;
5756        Some((float, histogram))
5757    }
5758
5759    fn alternative_sample_range_columns<'a>(
5760        schema: &DFSchemaRef,
5761        field_columns: &'a [String],
5762    ) -> Option<(&'a str, &'a str)> {
5763        Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| {
5764            Self::field_column_is_float_range(schema, float)
5765                && Self::field_column_is_native_histogram_range(schema, histogram)
5766        })
5767    }
5768
5769    fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool {
5770        Self::field_column_type(schema, field_column).is_some_and(|data_type| {
5771            matches!(
5772                data_type,
5773                ArrowDataType::Dictionary(key_type, value_type)
5774                    if key_type.as_ref() == &ArrowDataType::Int64
5775                        && value_type.as_ref() == &Self::native_histogram_arrow_type()
5776            )
5777        })
5778    }
5779
5780    fn all_field_columns_are_native_histograms(&self, schema: &DFSchemaRef) -> bool {
5781        !self.ctx.field_columns.is_empty()
5782            && self
5783                .ctx
5784                .field_columns
5785                .iter()
5786                .all(|field| Self::field_column_is_native_histogram(schema, field))
5787    }
5788
5789    fn all_field_columns_are_native_histogram_ranges(&self, schema: &DFSchemaRef) -> bool {
5790        !self.ctx.field_columns.is_empty()
5791            && self
5792                .ctx
5793                .field_columns
5794                .iter()
5795                .all(|field| Self::field_column_is_native_histogram_range(schema, field))
5796    }
5797
5798    fn optional_tsid_projection(
5799        schema: &DFSchemaRef,
5800        table_ref: Option<&TableReference>,
5801        keep_tsid: bool,
5802    ) -> Option<DfExpr> {
5803        keep_tsid.then_some(()).and_then(|_| {
5804            schema
5805                .qualified_field_with_name(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)
5806                .ok()
5807                .map(|field| DfExpr::Column(field.into()))
5808        })
5809    }
5810
5811    fn binary_join_key_columns(
5812        &self,
5813        left_schema: &DFSchemaRef,
5814        right_schema: &DFSchemaRef,
5815        left_context: &PromPlannerContext,
5816        right_context: &PromPlannerContext,
5817        only_join_time_index: bool,
5818        modifier: &Option<BinModifier>,
5819    ) -> Result<(BTreeSet<String>, BTreeSet<String>, bool)> {
5820        let has_tsid = |schema: &DFSchemaRef| {
5821            schema
5822                .fields()
5823                .iter()
5824                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
5825        };
5826        let use_tsid_join = !only_join_time_index
5827            && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
5828            && left_context.use_tsid
5829            && right_context.use_tsid
5830            && has_tsid(left_schema)
5831            && has_tsid(right_schema);
5832
5833        let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
5834            (
5835                BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5836                BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]),
5837            )
5838        } else {
5839            if only_join_time_index {
5840                (BTreeSet::new(), BTreeSet::new())
5841            } else {
5842                (
5843                    left_context
5844                        .tag_columns
5845                        .iter()
5846                        .cloned()
5847                        .collect::<BTreeSet<_>>(),
5848                    right_context
5849                        .tag_columns
5850                        .iter()
5851                        .cloned()
5852                        .collect::<BTreeSet<_>>(),
5853                )
5854            }
5855        };
5856
5857        if !use_tsid_join
5858            && let Some(modifier) = modifier
5859            && let Some(matching) = &modifier.matching
5860        {
5861            match matching {
5862                LabelModifier::Include(on) => {
5863                    let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5864                    left_tag_columns = left_tag_columns.intersection(&mask).cloned().collect();
5865                    right_tag_columns = right_tag_columns.intersection(&mask).cloned().collect();
5866                }
5867                LabelModifier::Exclude(ignoring) => {
5868                    for label in &ignoring.labels {
5869                        let _ = left_tag_columns.remove(label);
5870                        let _ = right_tag_columns.remove(label);
5871                    }
5872                }
5873            }
5874        }
5875
5876        let force_empty_join =
5877            !use_tsid_join && !only_join_time_index && left_tag_columns != right_tag_columns;
5878        if force_empty_join {
5879            let common_tag_columns = left_tag_columns
5880                .intersection(&right_tag_columns)
5881                .cloned()
5882                .collect::<BTreeSet<_>>();
5883            left_tag_columns = common_tag_columns.clone();
5884            right_tag_columns = common_tag_columns;
5885        }
5886
5887        Ok((left_tag_columns, right_tag_columns, force_empty_join))
5888    }
5889
5890    fn binary_modifier_preserves_tsid_join_key(
5891        &self,
5892        left_context: &PromPlannerContext,
5893        right_context: &PromPlannerContext,
5894        modifier: &Option<BinModifier>,
5895    ) -> bool {
5896        let Some(modifier) = modifier else {
5897            return true;
5898        };
5899
5900        if !matches!(modifier.card, VectorMatchCardinality::OneToOne) {
5901            return false;
5902        }
5903
5904        match &modifier.matching {
5905            None => true,
5906            Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
5907                !left_context.tag_columns.contains(label)
5908                    && !right_context.tag_columns.contains(label)
5909            }),
5910            Some(LabelModifier::Include(on)) => {
5911                let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
5912                let left_labels = left_context
5913                    .tag_columns
5914                    .iter()
5915                    .cloned()
5916                    .collect::<BTreeSet<_>>();
5917                let right_labels = right_context
5918                    .tag_columns
5919                    .iter()
5920                    .cloned()
5921                    .collect::<BTreeSet<_>>();
5922
5923                on_labels == left_labels && on_labels == right_labels
5924            }
5925        }
5926    }
5927
5928    /// Build a inner join on time index column and tag columns to concat two logical plans.
5929    /// When `only_join_time_index == true` we only join on the time index, because these two plan may not have the same tag columns
5930    #[allow(clippy::too_many_arguments)]
5931    fn join_on_non_field_columns(
5932        &self,
5933        left: LogicalPlan,
5934        right: LogicalPlan,
5935        left_table_ref: TableReference,
5936        right_table_ref: TableReference,
5937        left_time_index_column: Option<String>,
5938        right_time_index_column: Option<String>,
5939        only_join_time_index: bool,
5940        modifier: &Option<BinModifier>,
5941        left_context: &PromPlannerContext,
5942        right_context: &PromPlannerContext,
5943    ) -> Result<LogicalPlan> {
5944        let (mut left_tag_columns, mut right_tag_columns, mut force_empty_join) = self
5945            .binary_join_key_columns(
5946                left.schema(),
5947                right.schema(),
5948                left_context,
5949                right_context,
5950                only_join_time_index,
5951                modifier,
5952            )?;
5953        let use_tsid_join = !only_join_time_index
5954            && !force_empty_join
5955            && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()])
5956            && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]);
5957        let (left, right) = if !only_join_time_index
5958            && !use_tsid_join
5959            && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier)
5960        {
5961            let mut aligned_left_context = left_context.clone();
5962            let mut aligned_right_context = right_context.clone();
5963            let (left, right, _) = Self::align_temporality_match_column(
5964                left,
5965                right,
5966                &mut aligned_left_context,
5967                &mut aligned_right_context,
5968            )?;
5969            (left_tag_columns, right_tag_columns, force_empty_join) = self
5970                .binary_join_key_columns(
5971                    left.schema(),
5972                    right.schema(),
5973                    &aligned_left_context,
5974                    &aligned_right_context,
5975                    false,
5976                    modifier,
5977                )?;
5978            (left, right)
5979        } else {
5980            (left, right)
5981        };
5982
5983        // push time index column if it exists
5984        if let (Some(left_time_index_column), Some(right_time_index_column)) =
5985            (left_time_index_column, right_time_index_column)
5986        {
5987            left_tag_columns.insert(left_time_index_column);
5988            right_tag_columns.insert(right_time_index_column);
5989        }
5990
5991        let right = LogicalPlanBuilder::from(right)
5992            .alias(right_table_ref)
5993            .context(DataFusionPlanningSnafu)?
5994            .build()
5995            .context(DataFusionPlanningSnafu)?;
5996
5997        // Inner Join on time index column to concat two operator
5998        LogicalPlanBuilder::from(left)
5999            .alias(left_table_ref)
6000            .context(DataFusionPlanningSnafu)?
6001            .join_detailed(
6002                right,
6003                JoinType::Inner,
6004                (
6005                    left_tag_columns
6006                        .into_iter()
6007                        .map(Column::from_name)
6008                        .collect::<Vec<_>>(),
6009                    right_tag_columns
6010                        .into_iter()
6011                        .map(Column::from_name)
6012                        .collect::<Vec<_>>(),
6013                ),
6014                force_empty_join.then_some(lit(false)),
6015                NullEquality::NullEqualsNull,
6016            )
6017            .context(DataFusionPlanningSnafu)?
6018            .build()
6019            .context(DataFusionPlanningSnafu)
6020    }
6021
6022    fn selected_binary_match_labels(
6023        left_context: &PromPlannerContext,
6024        right_context: &PromPlannerContext,
6025        modifier: &Option<BinModifier>,
6026    ) -> BTreeSet<String> {
6027        let mut labels = left_context
6028            .tag_columns
6029            .iter()
6030            .chain(&right_context.tag_columns)
6031            .cloned()
6032            .collect::<BTreeSet<_>>();
6033        if let Some(matching) = modifier
6034            .as_ref()
6035            .and_then(|modifier| modifier.matching.as_ref())
6036        {
6037            match matching {
6038                LabelModifier::Include(on) => {
6039                    labels = on
6040                        .labels
6041                        .iter()
6042                        .filter(|label| {
6043                            left_context.tag_columns.contains(label)
6044                                || right_context.tag_columns.contains(label)
6045                        })
6046                        .cloned()
6047                        .collect();
6048                }
6049                LabelModifier::Exclude(ignoring) => {
6050                    for label in &ignoring.labels {
6051                        labels.remove(label);
6052                    }
6053                }
6054            }
6055        }
6056        labels
6057    }
6058
6059    fn only_temporality_match_label_mismatches(
6060        left_context: &PromPlannerContext,
6061        right_context: &PromPlannerContext,
6062        modifier: &Option<BinModifier>,
6063    ) -> bool {
6064        let mut mismatches =
6065            Self::selected_binary_match_labels(left_context, right_context, modifier)
6066                .into_iter()
6067                .filter(|label| {
6068                    left_context.tag_columns.contains(label)
6069                        != right_context.tag_columns.contains(label)
6070                });
6071        matches!(
6072            (mismatches.next(), mismatches.next()),
6073            (Some(label), None) if label == OTLP_AGGREGATION_TEMPORALITY_LABEL
6074        )
6075    }
6076
6077    fn align_temporality_match_column(
6078        mut left: LogicalPlan,
6079        mut right: LogicalPlan,
6080        left_context: &mut PromPlannerContext,
6081        right_context: &mut PromPlannerContext,
6082    ) -> Result<(LogicalPlan, LogicalPlan, bool)> {
6083        let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
6084        let left_has_marker = left_context.tag_columns.iter().any(|tag| tag == marker);
6085        let (present, add_to_left) = if left_has_marker {
6086            (&left, false)
6087        } else {
6088            (&right, true)
6089        };
6090        let data_type = present
6091            .schema()
6092            .fields()
6093            .iter()
6094            .find(|field| field.name() == marker)
6095            .map(|field| field.data_type().clone())
6096            .with_context(|| ColumnNotFoundSnafu {
6097                col: marker.to_string(),
6098            })?;
6099        let null = Self::string_scalar_value(&data_type, None).with_context(|| {
6100            UnexpectedPlanExprSnafu {
6101                desc: format!("temporality match label {marker} must be a string"),
6102            }
6103        })?;
6104        let add_marker = |plan: LogicalPlan| {
6105            let visible = plan
6106                .schema()
6107                .iter()
6108                .map(|(qualifier, field)| {
6109                    DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6110                })
6111                .collect::<Vec<_>>();
6112            LogicalPlanBuilder::from(plan)
6113                .project(
6114                    visible
6115                        .into_iter()
6116                        .chain([DfExpr::Literal(null, None).alias(marker)]),
6117                )
6118                .context(DataFusionPlanningSnafu)?
6119                .build()
6120                .context(DataFusionPlanningSnafu)
6121        };
6122
6123        if add_to_left {
6124            left = add_marker(left)?;
6125            left_context.tag_columns.push(marker.to_string());
6126        } else {
6127            right = add_marker(right)?;
6128            right_context.tag_columns.push(marker.to_string());
6129        }
6130        Ok((left, right, add_to_left))
6131    }
6132
6133    fn normalized_match_key_expr(
6134        label: &str,
6135        field: Option<(Option<TableReference>, ArrowDataType)>,
6136        value_type: &ArrowDataType,
6137        internal_name: &str,
6138    ) -> DfExpr {
6139        let empty = Self::string_scalar_value(value_type, Some(String::new()))
6140            .expect("match label value type is a string");
6141        let expr = if let Some((qualifier, data_type)) = field {
6142            let column = DfExpr::Column(Column::new(qualifier, label));
6143            let column = if &data_type == value_type {
6144                column
6145            } else {
6146                DfExpr::Cast(Cast {
6147                    expr: Box::new(column),
6148                    data_type: value_type.clone(),
6149                })
6150            };
6151            DfExpr::ScalarFunction(ScalarFunction {
6152                func: coalesce(),
6153                args: vec![column, DfExpr::Literal(empty, None)],
6154            })
6155        } else {
6156            DfExpr::Literal(empty, None)
6157        };
6158        expr.alias(internal_name)
6159    }
6160
6161    fn is_zero_row_empty_relation(plan: &LogicalPlan) -> bool {
6162        // `produce_one_row` is used for input-free plans that still emit one row;
6163        // only the false case is a statically proven empty vector.
6164        matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row)
6165    }
6166
6167    /// Build a set operator (AND/OR/UNLESS)
6168    fn set_op_on_non_field_columns(
6169        &mut self,
6170        mut left: LogicalPlan,
6171        mut right: LogicalPlan,
6172        left_context: PromPlannerContext,
6173        right_context: PromPlannerContext,
6174        op: TokenType,
6175        modifier: &Option<BinModifier>,
6176    ) -> Result<LogicalPlan> {
6177        let left_tag_col_set = left_context
6178            .tag_columns
6179            .iter()
6180            .cloned()
6181            .collect::<HashSet<_>>();
6182        let right_tag_col_set = right_context
6183            .tag_columns
6184            .iter()
6185            .cloned()
6186            .collect::<HashSet<_>>();
6187
6188        if matches!(op.id(), token::T_LOR) {
6189            return self.or_operator(
6190                left,
6191                right,
6192                left_tag_col_set,
6193                right_tag_col_set,
6194                left_context,
6195                right_context,
6196                modifier,
6197            );
6198        }
6199
6200        if let Some(modifier) = modifier {
6201            ensure!(
6202                matches!(
6203                    modifier.card,
6204                    VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany
6205                ),
6206                UnsupportedVectorMatchSnafu {
6207                    name: modifier.card.clone(),
6208                },
6209            );
6210        }
6211
6212        let output_context = left_context.clone();
6213        let visible_left_schema = left.schema().clone();
6214        let mut left_context = left_context;
6215        let mut right_context = right_context;
6216        let added_marker_to_left = if Self::only_temporality_match_label_mismatches(
6217            &left_context,
6218            &right_context,
6219            modifier,
6220        ) {
6221            let aligned = Self::align_temporality_match_column(
6222                left,
6223                right,
6224                &mut left_context,
6225                &mut right_context,
6226            )?;
6227            left = aligned.0;
6228            right = aligned.1;
6229            aligned.2
6230        } else {
6231            false
6232        };
6233
6234        let mut left_tag_col_set = left_context
6235            .tag_columns
6236            .iter()
6237            .cloned()
6238            .collect::<BTreeSet<_>>();
6239        let mut right_tag_col_set = right_context
6240            .tag_columns
6241            .iter()
6242            .cloned()
6243            .collect::<BTreeSet<_>>();
6244        if let Some(matching) = modifier
6245            .as_ref()
6246            .and_then(|modifier| modifier.matching.as_ref())
6247        {
6248            match matching {
6249                LabelModifier::Include(on) => {
6250                    let mask = on.labels.iter().cloned().collect::<BTreeSet<_>>();
6251                    left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect();
6252                    right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect();
6253                }
6254                LabelModifier::Exclude(ignoring) => {
6255                    for label in &ignoring.labels {
6256                        let _ = left_tag_col_set.remove(label);
6257                        let _ = right_tag_col_set.remove(label);
6258                    }
6259                }
6260            }
6261        }
6262        ensure!(
6263            left_tag_col_set == right_tag_col_set,
6264            CombineTableColumnMismatchSnafu {
6265                left: left_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6266                right: right_tag_col_set.iter().cloned().collect::<Vec<_>>(),
6267            }
6268        );
6269
6270        let left_time_index = left_context.time_index_column.clone().unwrap();
6271        let right_time_index = right_context.time_index_column.clone().unwrap();
6272
6273        // alias right time index column if necessary
6274        if left_context.time_index_column != right_context.time_index_column {
6275            let right_project_exprs = right
6276                .schema()
6277                .fields()
6278                .iter()
6279                .map(|field| {
6280                    if field.name() == &right_time_index {
6281                        DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index)
6282                    } else {
6283                        DfExpr::Column(Column::from_name(field.name()))
6284                    }
6285                })
6286                .collect::<Vec<_>>();
6287
6288            right = LogicalPlanBuilder::from(right)
6289                .project(right_project_exprs)
6290                .context(DataFusionPlanningSnafu)?
6291                .build()
6292                .context(DataFusionPlanningSnafu)?;
6293        }
6294
6295        let join_keys = left_tag_col_set
6296            .into_iter()
6297            .chain([left_time_index])
6298            .collect::<Vec<_>>();
6299
6300        ensure!(
6301            left_context.field_columns.len() == 1
6302                || Self::field_columns_are_alternative_samples(
6303                    left.schema(),
6304                    &left_context.field_columns,
6305                ),
6306            MultiFieldsNotSupportedSnafu {
6307                operator: "AND/UNLESS operator"
6308            }
6309        );
6310        // Generate join plan.
6311        // All set operations in PromQL are "distinct"
6312        let result = match op.id() {
6313            token::T_LAND => LogicalPlanBuilder::from(left)
6314                .distinct()
6315                .context(DataFusionPlanningSnafu)?
6316                .join_detailed(
6317                    right,
6318                    JoinType::LeftSemi,
6319                    (join_keys.clone(), join_keys),
6320                    None,
6321                    NullEquality::NullEqualsNull,
6322                )
6323                .context(DataFusionPlanningSnafu)?
6324                .build()
6325                .context(DataFusionPlanningSnafu),
6326            token::T_LUNLESS => LogicalPlanBuilder::from(left)
6327                .distinct()
6328                .context(DataFusionPlanningSnafu)?
6329                .join_detailed(
6330                    right,
6331                    JoinType::LeftAnti,
6332                    (join_keys.clone(), join_keys),
6333                    None,
6334                    NullEquality::NullEqualsNull,
6335                )
6336                .context(DataFusionPlanningSnafu)?
6337                .build()
6338                .context(DataFusionPlanningSnafu),
6339            token::T_LOR => {
6340                // OR is handled at the beginning of this function, as it cannot
6341                // be expressed using JOIN like AND and UNLESS.
6342                unreachable!()
6343            }
6344            _ => UnexpectedTokenSnafu { token: op }.fail(),
6345        }?;
6346        let result = if added_marker_to_left {
6347            LogicalPlanBuilder::from(result)
6348                .project(visible_left_schema.iter().map(|(qualifier, field)| {
6349                    DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6350                }))
6351                .context(DataFusionPlanningSnafu)?
6352                .build()
6353                .context(DataFusionPlanningSnafu)?
6354        } else {
6355            result
6356        };
6357
6358        // AND/UNLESS preserve the complete left operand schema and metadata.
6359        self.ctx = output_context;
6360        Ok(result)
6361    }
6362
6363    fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> {
6364        match data_type {
6365            data_type if data_type.is_string() => Some(data_type),
6366            ArrowDataType::Dictionary(_, value_type) if value_type.is_string() => Some(value_type),
6367            _ => None,
6368        }
6369    }
6370
6371    fn string_scalar_value(
6372        data_type: &ArrowDataType,
6373        value: Option<String>,
6374    ) -> Option<ScalarValue> {
6375        match data_type {
6376            ArrowDataType::Utf8 => Some(ScalarValue::Utf8(value)),
6377            ArrowDataType::LargeUtf8 => Some(ScalarValue::LargeUtf8(value)),
6378            ArrowDataType::Utf8View => Some(ScalarValue::Utf8View(value)),
6379            ArrowDataType::Dictionary(key_type, value_type) => Some(ScalarValue::Dictionary(
6380                key_type.clone(),
6381                Box::new(Self::string_scalar_value(value_type, value)?),
6382            )),
6383            _ => None,
6384        }
6385    }
6386
6387    fn common_label_data_type(
6388        left: Option<&ArrowDataType>,
6389        right: Option<&ArrowDataType>,
6390    ) -> Option<ArrowDataType> {
6391        match (left, right) {
6392            (Some(left), Some(right)) if left == right => {
6393                Self::string_value_data_type(left).map(|_| left.clone())
6394            }
6395            (Some(left), Some(right)) => {
6396                let left_value_type = Self::string_value_data_type(left)?;
6397                let right_value_type = Self::string_value_data_type(right)?;
6398                // DataFusion projections can decode dictionaries, but do not encode plain strings
6399                // as dictionaries. Preserve the encoding only when both inputs already share it.
6400                match (left_value_type, right_value_type) {
6401                    (left, right) if left == right => Some(left.clone()),
6402                    (ArrowDataType::LargeUtf8, _) | (_, ArrowDataType::LargeUtf8) => {
6403                        Some(ArrowDataType::LargeUtf8)
6404                    }
6405                    (ArrowDataType::Utf8View, ArrowDataType::Utf8View) => {
6406                        Some(ArrowDataType::Utf8View)
6407                    }
6408                    _ => Some(ArrowDataType::Utf8),
6409                }
6410            }
6411            (Some(data_type), None) | (None, Some(data_type)) => {
6412                Self::string_value_data_type(data_type).cloned()
6413            }
6414            (None, None) => Some(ArrowDataType::Utf8),
6415        }
6416    }
6417
6418    // TODO(ruihang): change function name
6419    #[allow(clippy::too_many_arguments)]
6420    fn or_operator(
6421        &mut self,
6422        left: LogicalPlan,
6423        right: LogicalPlan,
6424        left_tag_cols_set: HashSet<String>,
6425        right_tag_cols_set: HashSet<String>,
6426        left_context: PromPlannerContext,
6427        right_context: PromPlannerContext,
6428        modifier: &Option<BinModifier>,
6429    ) -> Result<LogicalPlan> {
6430        let left_is_empty = Self::is_zero_row_empty_relation(&left);
6431        let right_is_empty = Self::is_zero_row_empty_relation(&right);
6432        match (left_is_empty, right_is_empty) {
6433            (true, false) => {
6434                self.ctx = right_context;
6435                return Ok(right);
6436            }
6437            (false, true) => {
6438                self.ctx = left_context;
6439                return Ok(left);
6440            }
6441            (true, true) => {
6442                self.ctx = left_context;
6443                return Ok(left);
6444            }
6445            (false, false) => {}
6446        }
6447
6448        ensure!(
6449            !left.schema().fields().is_empty() && !right.schema().fields().is_empty(),
6450            UnexpectedPlanExprSnafu {
6451                desc: "OR operator input has zero columns",
6452            }
6453        );
6454        let left_has_alternative_samples =
6455            Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns);
6456        let right_has_alternative_samples = Self::field_columns_are_alternative_samples(
6457            right.schema(),
6458            &right_context.field_columns,
6459        );
6460        ensure!(
6461            left_context.field_columns.len() == 1 || left_has_alternative_samples,
6462            MultiFieldsNotSupportedSnafu {
6463                operator: "OR operator"
6464            }
6465        );
6466        ensure!(
6467            right_context.field_columns.len() == 1 || right_has_alternative_samples,
6468            MultiFieldsNotSupportedSnafu {
6469                operator: "OR operator"
6470            }
6471        );
6472
6473        // prepare hash sets
6474        let all_tags = left_tag_cols_set
6475            .union(&right_tag_cols_set)
6476            .cloned()
6477            .collect::<HashSet<_>>();
6478        let left_qualifier = left.schema().qualified_field(0).0.cloned();
6479        let right_qualifier = right.schema().qualified_field(0).0.cloned();
6480        let left_qualifier_string = left_qualifier
6481            .as_ref()
6482            .map(|l| l.to_string())
6483            .unwrap_or_default();
6484        let right_qualifier_string = right_qualifier
6485            .as_ref()
6486            .map(|r| r.to_string())
6487            .unwrap_or_default();
6488        let left_time_index_column =
6489            left_context
6490                .time_index_column
6491                .clone()
6492                .with_context(|| TimeIndexNotFoundSnafu {
6493                    table: left_qualifier_string.clone(),
6494                })?;
6495        let right_time_index_column =
6496            right_context
6497                .time_index_column
6498                .clone()
6499                .with_context(|| TimeIndexNotFoundSnafu {
6500                    table: right_qualifier_string.clone(),
6501                })?;
6502        let native_histogram_type = Self::native_histogram_arrow_type();
6503        let is_numeric = |data_type: &ArrowDataType| {
6504            matches!(
6505                data_type,
6506                ArrowDataType::Int8
6507                    | ArrowDataType::Int16
6508                    | ArrowDataType::Int32
6509                    | ArrowDataType::Int64
6510                    | ArrowDataType::UInt8
6511                    | ArrowDataType::UInt16
6512                    | ArrowDataType::UInt32
6513                    | ArrowDataType::UInt64
6514                    | ArrowDataType::Float32
6515                    | ArrowDataType::Float64
6516            )
6517        };
6518        let left_fields = left_context
6519            .field_columns
6520            .iter()
6521            .map(|name| {
6522                left.schema()
6523                    .iter()
6524                    .find(|(_, field)| field.name() == name)
6525                    .map(|(qualifier, field)| {
6526                        (name.clone(), qualifier.cloned(), field.data_type().clone())
6527                    })
6528                    .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6529            })
6530            .collect::<Result<Vec<_>>>()?;
6531        let right_fields = right_context
6532            .field_columns
6533            .iter()
6534            .map(|name| {
6535                right
6536                    .schema()
6537                    .iter()
6538                    .find(|(_, field)| field.name() == name)
6539                    .map(|(qualifier, field)| {
6540                        (name.clone(), qualifier.cloned(), field.data_type().clone())
6541                    })
6542                    .with_context(|| ColumnNotFoundSnafu { col: name.clone() })
6543            })
6544            .collect::<Result<Vec<_>>>()?;
6545        let left_field = &left_fields[0];
6546        let right_field = &right_fields[0];
6547        let left_field_col = &left_field.0;
6548        let right_field_col = &right_field.0;
6549        let fields_are_samples = |fields: &[(String, Option<TableReference>, ArrowDataType)]| {
6550            fields.iter().all(|(_, _, data_type)| {
6551                is_numeric(data_type) || data_type == &native_histogram_type
6552            })
6553        };
6554        let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples {
6555            if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) {
6556                return UnexpectedPlanExprSnafu {
6557                    desc: format!(
6558                        "OR value fields have incompatible types: {:?} and {:?}",
6559                        left_fields
6560                            .iter()
6561                            .map(|(_, _, data_type)| data_type)
6562                            .collect::<Vec<_>>(),
6563                        right_fields
6564                            .iter()
6565                            .map(|(_, _, data_type)| data_type)
6566                            .collect::<Vec<_>>()
6567                    ),
6568                }
6569                .fail();
6570            }
6571            true
6572        } else {
6573            (left_field.2 == native_histogram_type && is_numeric(&right_field.2))
6574                || (right_field.2 == native_histogram_type && is_numeric(&left_field.2))
6575        };
6576        let target_field_type = if mixed_sample_types {
6577            // Mixed vectors use the existing response representation: one nullable float column
6578            // and one nullable native-histogram column.
6579            ArrowDataType::Float64
6580        } else if left_field.2 == right_field.2 {
6581            left_field.2.clone()
6582        } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) {
6583            ArrowDataType::Float64
6584        } else {
6585            return UnexpectedPlanExprSnafu {
6586                desc: format!(
6587                    "OR value fields have incompatible types: {:?} and {:?}",
6588                    left_field.2, right_field.2
6589                ),
6590            }
6591            .fail();
6592        };
6593        let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types {
6594            let mut reserved_names = left
6595                .schema()
6596                .fields()
6597                .iter()
6598                .chain(right.schema().fields().iter())
6599                .map(|field| field.name().clone())
6600                .collect::<HashSet<_>>();
6601            for (name, _, _) in left_fields.iter().chain(&right_fields) {
6602                reserved_names.remove(name);
6603            }
6604            reserved_names.extend(all_tags.iter().cloned());
6605            let unique_name = |prefix: &str, reserved_names: &mut HashSet<String>| {
6606                let mut index = 0;
6607                loop {
6608                    let name = format!("{prefix}{index}");
6609                    index += 1;
6610                    if reserved_names.insert(name.clone()) {
6611                        break name;
6612                    }
6613                }
6614            };
6615            let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names);
6616            let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names);
6617            (float_field, histogram_field)
6618        } else {
6619            (left_field_col.clone(), String::new())
6620        };
6621        let left_tag_types = left_tag_cols_set
6622            .iter()
6623            .map(|label| {
6624                left.schema()
6625                    .fields()
6626                    .iter()
6627                    .find(|field| field.name() == label)
6628                    .map(|field| (label.clone(), field.data_type().clone()))
6629                    .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
6630            })
6631            .collect::<Result<HashMap<_, _>>>()?;
6632        let right_tag_types = right_tag_cols_set
6633            .iter()
6634            .map(|label| {
6635                right
6636                    .schema()
6637                    .fields()
6638                    .iter()
6639                    .find(|field| field.name() == label)
6640                    .map(|field| (label.clone(), field.data_type().clone()))
6641                    .with_context(|| ColumnNotFoundSnafu { col: label.clone() })
6642            })
6643            .collect::<Result<HashMap<_, _>>>()?;
6644        let mut target_tag_types = HashMap::with_capacity(all_tags.len());
6645        for label in &all_tags {
6646            let Some(data_type) =
6647                Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label))
6648            else {
6649                return UnexpectedPlanExprSnafu {
6650                    desc: format!(
6651                        "OR label {label} has incompatible types: {:?} and {:?}",
6652                        left_tag_types.get(label),
6653                        right_tag_types.get(label)
6654                    ),
6655                }
6656                .fail();
6657            };
6658            target_tag_types.insert(label.clone(), data_type);
6659        }
6660        let left_has_tsid = left
6661            .schema()
6662            .fields()
6663            .iter()
6664            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
6665        let right_has_tsid = right
6666            .schema()
6667            .fields()
6668            .iter()
6669            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME);
6670
6671        // step 0: fill all columns in output schema
6672        let mut all_columns_set = left
6673            .schema()
6674            .fields()
6675            .iter()
6676            .chain(right.schema().fields().iter())
6677            .map(|field| field.name().clone())
6678            .collect::<HashSet<_>>();
6679        // Keep `__tsid` only when both sides contain it, otherwise it may break schema alignment
6680        // (e.g. `unknown_metric or some_metric`).
6681        if !(left_has_tsid && right_has_tsid) {
6682            all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME);
6683        }
6684        // remove time index column
6685        all_columns_set.remove(&left_time_index_column);
6686        all_columns_set.remove(&right_time_index_column);
6687        if mixed_sample_types {
6688            for (name, _, _) in left_fields.iter().chain(&right_fields) {
6689                all_columns_set.remove(name);
6690            }
6691            all_columns_set.extend(all_tags.iter().cloned());
6692            all_columns_set.insert(mixed_float_field_col.clone());
6693            all_columns_set.insert(mixed_histogram_field_col.clone());
6694        } else if left_field_col != right_field_col {
6695            // remove field column in the right
6696            all_columns_set.remove(right_field_col);
6697        }
6698        let mut all_columns = all_columns_set.into_iter().collect::<Vec<_>>();
6699        // sort to ensure the generated schema is not volatile
6700        all_columns.sort_unstable();
6701        // use left time index column name as the result time index column name
6702        all_columns.insert(0, left_time_index_column.clone());
6703        let mut occupied_column_names = left
6704            .schema()
6705            .fields()
6706            .iter()
6707            .chain(right.schema().fields().iter())
6708            .map(|field| field.name().clone())
6709            .collect::<HashSet<_>>();
6710
6711        // step 1: align schema using project, fill non-exist columns with null
6712        let aligned_label_expr = |col: &String, source_types: &HashMap<String, ArrowDataType>| {
6713            let target_type = &target_tag_types[col];
6714            if let Some(source_type) = source_types.get(col) {
6715                let expr = DfExpr::Column(Column::new(None::<String>, col));
6716                if source_type == target_type {
6717                    expr
6718                } else {
6719                    DfExpr::Cast(Cast {
6720                        expr: Box::new(expr),
6721                        data_type: target_type.clone(),
6722                    })
6723                    .alias(col.clone())
6724                }
6725            } else {
6726                DfExpr::Literal(
6727                    Self::string_scalar_value(target_type, None)
6728                        .expect("target label type is a string"),
6729                    None,
6730                )
6731                .alias(col.clone())
6732            }
6733        };
6734        let null_histogram =
6735            ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?;
6736        let mixed_value_expr = |fields: &[(String, Option<TableReference>, ArrowDataType)],
6737                                output_col: &String| {
6738            if output_col == &mixed_float_field_col {
6739                if let Some((name, qualifier, data_type)) = fields
6740                    .iter()
6741                    .find(|(_, _, data_type)| is_numeric(data_type))
6742                {
6743                    let expr = DfExpr::Column(Column::new(qualifier.clone(), name));
6744                    if data_type == &ArrowDataType::Float64 {
6745                        expr.alias(output_col)
6746                    } else {
6747                        DfExpr::Cast(Cast {
6748                            expr: Box::new(expr),
6749                            data_type: ArrowDataType::Float64,
6750                        })
6751                        .alias(output_col)
6752                    }
6753                } else {
6754                    DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col)
6755                }
6756            } else {
6757                fields
6758                    .iter()
6759                    .find(|(_, _, data_type)| data_type == &native_histogram_type)
6760                    .map(|(name, qualifier, _)| {
6761                        DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col)
6762                    })
6763                    .unwrap_or_else(|| {
6764                        DfExpr::Literal(null_histogram.clone(), None).alias(output_col)
6765                    })
6766            }
6767        };
6768        let left_proj_exprs = all_columns.iter().map(|col| {
6769            if mixed_sample_types
6770                && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
6771            {
6772                mixed_value_expr(&left_fields, col)
6773            } else if !mixed_sample_types
6774                && col == left_field_col
6775                && left_field.2 != target_field_type
6776            {
6777                DfExpr::Cast(Cast {
6778                    expr: Box::new(DfExpr::Column(Column::new(
6779                        left_field.1.clone(),
6780                        left_field_col,
6781                    ))),
6782                    data_type: target_field_type.clone(),
6783                })
6784                .alias(left_field_col.clone())
6785            } else if target_tag_types.contains_key(col) {
6786                aligned_label_expr(col, &left_tag_types)
6787            } else {
6788                DfExpr::Column(Column::new(None::<String>, col))
6789            }
6790        });
6791        let right_time_index_expr = DfExpr::Column(Column::new(
6792            right_qualifier.clone(),
6793            right_time_index_column,
6794        ))
6795        .alias(left_time_index_column.clone());
6796        // The field column in right side may not have qualifier (it may be removed by join operation),
6797        // so we need to find it from the schema.
6798        // `skip(1)` to skip the time index column
6799        let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| {
6800            // expr
6801            if mixed_sample_types
6802                && (col == &mixed_float_field_col || col == &mixed_histogram_field_col)
6803            {
6804                mixed_value_expr(&right_fields, col)
6805            } else if !mixed_sample_types && col == left_field_col {
6806                let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col));
6807                if right_field.2 != target_field_type {
6808                    DfExpr::Cast(Cast {
6809                        expr: Box::new(expr),
6810                        data_type: target_field_type.clone(),
6811                    })
6812                    .alias(left_field_col.clone())
6813                } else if left_field_col != right_field_col {
6814                    expr.alias(left_field_col.clone())
6815                } else {
6816                    expr
6817                }
6818            } else if target_tag_types.contains_key(col) {
6819                aligned_label_expr(col, &right_tag_types)
6820            } else {
6821                DfExpr::Column(Column::new(None::<String>, col))
6822            }
6823        });
6824        let right_proj_exprs = [right_time_index_expr]
6825            .into_iter()
6826            .chain(right_proj_exprs_without_time_index);
6827
6828        let left_projected = LogicalPlanBuilder::from(left)
6829            .project(left_proj_exprs)
6830            .context(DataFusionPlanningSnafu)?
6831            .alias(left_qualifier_string.clone())
6832            .context(DataFusionPlanningSnafu)?
6833            .build()
6834            .context(DataFusionPlanningSnafu)?;
6835        let right_projected = LogicalPlanBuilder::from(right)
6836            .project(right_proj_exprs)
6837            .context(DataFusionPlanningSnafu)?
6838            .alias(right_qualifier_string.clone())
6839            .context(DataFusionPlanningSnafu)?
6840            .build()
6841            .context(DataFusionPlanningSnafu)?;
6842
6843        // step 2: compute match columns
6844        let mut match_columns = if let Some(modifier) = modifier
6845            && let Some(matching) = &modifier.matching
6846        {
6847            match matching {
6848                // keeps columns mentioned in `on`
6849                LabelModifier::Include(on) => on.labels.clone(),
6850                // removes columns memtioned in `ignoring`
6851                LabelModifier::Exclude(ignoring) => {
6852                    let ignoring = ignoring.labels.iter().cloned().collect::<HashSet<_>>();
6853                    all_tags.difference(&ignoring).cloned().collect()
6854                }
6855            }
6856        } else {
6857            all_tags.iter().cloned().collect()
6858        };
6859        // sort to ensure the generated plan is not volatile
6860        match_columns.sort_unstable();
6861        match_columns.dedup();
6862        occupied_column_names.extend(
6863            left_projected
6864                .schema()
6865                .fields()
6866                .iter()
6867                .chain(right_projected.schema().fields().iter())
6868                .map(|field| field.name().clone()),
6869        );
6870
6871        let visible_schema = left_projected.schema().clone();
6872        let visible_left_exprs = left_projected
6873            .schema()
6874            .iter()
6875            .map(|(qualifier, field)| {
6876                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6877            })
6878            .collect::<Vec<_>>();
6879        let visible_right_exprs = right_projected
6880            .schema()
6881            .iter()
6882            .map(|(qualifier, field)| {
6883                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
6884            })
6885            .collect::<Vec<_>>();
6886        let mut left_match_exprs = Vec::with_capacity(match_columns.len());
6887        let mut right_match_exprs = Vec::with_capacity(match_columns.len());
6888        let mut next_internal_column = 0;
6889
6890        for label in &match_columns {
6891            let left_field = if left_tag_cols_set.contains(label) {
6892                Some(
6893                    left_projected
6894                        .schema()
6895                        .iter()
6896                        .find(|(_, field)| field.name() == label)
6897                        .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
6898                        .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
6899                )
6900            } else {
6901                None
6902            };
6903            let right_field = if right_tag_cols_set.contains(label) {
6904                Some(
6905                    right_projected
6906                        .schema()
6907                        .iter()
6908                        .find(|(_, field)| field.name() == label)
6909                        .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone()))
6910                        .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?,
6911                )
6912            } else {
6913                None
6914            };
6915            let data_type = match (left_field.as_ref(), right_field.as_ref()) {
6916                (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => {
6917                    left_type.clone()
6918                }
6919                (Some((_, left_type)), Some((_, right_type))) => {
6920                    return UnexpectedPlanExprSnafu {
6921                        desc: format!(
6922                            "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}"
6923                        ),
6924                    }
6925                    .fail();
6926                }
6927                (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(),
6928                (None, None) => ArrowDataType::Utf8,
6929            };
6930            let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else {
6931                return UnexpectedPlanExprSnafu {
6932                    desc: format!("OR match label {label} must be a string"),
6933                }
6934                .fail();
6935            };
6936            let internal_name = loop {
6937                let name = format!("__promql_or_match_{next_internal_column}");
6938                next_internal_column += 1;
6939                if occupied_column_names.insert(name.clone()) {
6940                    break name;
6941                }
6942            };
6943            left_match_exprs.push(Self::normalized_match_key_expr(
6944                label,
6945                left_field,
6946                &value_type,
6947                &internal_name,
6948            ));
6949            right_match_exprs.push(Self::normalized_match_key_expr(
6950                label,
6951                right_field,
6952                &value_type,
6953                &internal_name,
6954            ));
6955        }
6956
6957        let left_augmented = LogicalPlanBuilder::from(left_projected)
6958            .project(visible_left_exprs.into_iter().chain(left_match_exprs))
6959            .context(DataFusionPlanningSnafu)?
6960            .build()
6961            .context(DataFusionPlanningSnafu)?;
6962        let right_augmented = LogicalPlanBuilder::from(right_projected)
6963            .project(visible_right_exprs.into_iter().chain(right_match_exprs))
6964            .context(DataFusionPlanningSnafu)?
6965            .build()
6966            .context(DataFusionPlanningSnafu)?;
6967
6968        // step 3: build `UnionDistinctOn` with normalized internal match keys.
6969        let visible_field_count = visible_schema.fields().len();
6970        let compare_key_indices =
6971            (visible_field_count..visible_field_count + match_columns.len()).collect::<Vec<_>>();
6972        let (time_qualifier, _) = visible_schema
6973            .iter()
6974            .find(|(_, field)| field.name() == &left_time_index_column)
6975            .with_context(|| TimeIndexNotFoundSnafu {
6976                table: left_qualifier_string.clone(),
6977            })?;
6978        let ts_col_idx = left_augmented
6979            .schema()
6980            .iter()
6981            .position(|(qualifier, field)| {
6982                qualifier == time_qualifier && field.name() == &left_time_index_column
6983            })
6984            .with_context(|| TimeIndexNotFoundSnafu {
6985                table: left_qualifier_string.clone(),
6986            })?;
6987        let union_distinct_on = UnionDistinctOn::try_new(
6988            left_augmented,
6989            right_augmented,
6990            compare_key_indices,
6991            ts_col_idx,
6992        )
6993        .context(DataFusionPlanningSnafu)?;
6994        let augmented_result = LogicalPlan::Extension(Extension {
6995            node: Arc::new(union_distinct_on),
6996        });
6997        let result = LogicalPlanBuilder::from(augmented_result)
6998            .project(visible_schema.iter().map(|(qualifier, field)| {
6999                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7000            }))
7001            .context(DataFusionPlanningSnafu)?
7002            .build()
7003            .context(DataFusionPlanningSnafu)?;
7004
7005        // step 4: update context
7006        let output_field_col = left_field_col.clone();
7007        let mut output_context = left_context;
7008        let mut visible_tags = all_tags.into_iter().collect::<Vec<_>>();
7009        visible_tags.sort_unstable();
7010        output_context.time_index_column = Some(left_time_index_column);
7011        output_context.tag_columns = visible_tags;
7012        output_context.field_columns = if mixed_sample_types {
7013            vec![mixed_float_field_col, mixed_histogram_field_col]
7014        } else {
7015            vec![output_field_col]
7016        };
7017        output_context.use_tsid = left_has_tsid && right_has_tsid;
7018        self.ctx = output_context;
7019
7020        Ok(result)
7021    }
7022
7023    /// Build a projection that project and perform operation expr for every value columns.
7024    /// Non-value columns (tag and timestamp) will be preserved in the projection.
7025    ///
7026    /// # Side effect
7027    ///
7028    /// This function will update the value columns in the context. Those new column names
7029    /// don't contains qualifier.
7030    fn projection_for_each_field_column<F>(
7031        &mut self,
7032        input: LogicalPlan,
7033        name_to_expr: F,
7034    ) -> Result<LogicalPlan>
7035    where
7036        F: FnMut(&String) -> Result<DfExpr>,
7037    {
7038        // Keep the generated float/histogram lane names while an element-wise operation
7039        // preserves both sample types, so downstream operators still recognize the pair.
7040        let preserve_field_names =
7041            Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
7042        let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
7043        // Derived labels can be unqualified even when the context still names the source table.
7044        let input_schema = input.schema().clone();
7045        let non_field_columns_iter = self
7046            .ctx
7047            .tag_columns
7048            .iter()
7049            .chain(self.ctx.time_index_column.iter())
7050            .map(|col| {
7051                input_schema
7052                    .qualified_field_with_name(table_ref.as_ref(), col)
7053                    .or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
7054                    .map(|field| DfExpr::Column(field.into()))
7055                    .context(DataFusionPlanningSnafu)
7056            });
7057        let tsid_iter =
7058            Self::optional_tsid_projection(input.schema(), table_ref.as_ref(), self.ctx.use_tsid)
7059                .into_iter()
7060                .map(Ok);
7061
7062        // build computation exprs
7063        let result_field_columns = self
7064            .ctx
7065            .field_columns
7066            .iter()
7067            .map(name_to_expr)
7068            .collect::<Result<Vec<_>>>()?;
7069
7070        // alias the computation exprs to remove qualifier
7071        if !preserve_field_names {
7072            self.ctx.field_columns = result_field_columns
7073                .iter()
7074                .map(|expr| expr.schema_name().to_string())
7075                .collect();
7076        }
7077        let field_columns_iter = result_field_columns
7078            .into_iter()
7079            .zip(self.ctx.field_columns.iter())
7080            .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::<String>, name))));
7081
7082        // chain non-field columns (unchanged) and field columns (applied computation then alias)
7083        let project_fields = non_field_columns_iter
7084            .chain(tsid_iter)
7085            .chain(field_columns_iter)
7086            .collect::<Result<Vec<_>>>()?;
7087
7088        LogicalPlanBuilder::from(input)
7089            .project(project_fields)
7090            .context(DataFusionPlanningSnafu)?
7091            .build()
7092            .context(DataFusionPlanningSnafu)
7093    }
7094
7095    /// Build a filter plan on one value column or a float/histogram alternative pair.
7096    fn filter_on_field_column<F>(&self, input: LogicalPlan, name_to_expr: F) -> Result<LogicalPlan>
7097    where
7098        F: FnMut(&String) -> Result<DfExpr>,
7099    {
7100        ensure!(
7101            self.ctx.field_columns.len() == 1
7102                || Self::field_columns_are_alternative_samples(
7103                    input.schema(),
7104                    &self.ctx.field_columns,
7105                ),
7106            UnsupportedExprSnafu {
7107                name: "filter on multi-value input"
7108            }
7109        );
7110
7111        let field_column_filters = self
7112            .ctx
7113            .field_columns
7114            .iter()
7115            .map(name_to_expr)
7116            .collect::<Result<Vec<_>>>()?;
7117        let field_column_filter =
7118            disjunction(field_column_filters).context(UnsupportedExprSnafu {
7119                name: "filter on empty input",
7120            })?;
7121
7122        LogicalPlanBuilder::from(input)
7123            .filter(field_column_filter)
7124            .context(DataFusionPlanningSnafu)?
7125            .build()
7126            .context(DataFusionPlanningSnafu)
7127    }
7128
7129    /// Generate an expr like `date_part("hour", <TIME_INDEX>)`. Caller should ensure the
7130    /// time index column in context is set
7131    fn date_part_on_time_index(&self, date_part: &str) -> Result<DfExpr> {
7132        let input_expr = datafusion::logical_expr::col(
7133            self.ctx
7134                .time_index_column
7135                .as_ref()
7136                // table name doesn't matters here
7137                .with_context(|| TimeIndexNotFoundSnafu {
7138                    table: "<doesn't matter>",
7139                })?,
7140        );
7141        let fn_expr = DfExpr::ScalarFunction(ScalarFunction {
7142            func: datafusion_functions::datetime::date_part(),
7143            args: vec![date_part.lit(), input_expr],
7144        });
7145        Ok(fn_expr)
7146    }
7147
7148    fn strip_tsid_column(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
7149        let schema = plan.schema();
7150        if !schema
7151            .fields()
7152            .iter()
7153            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
7154        {
7155            return Ok(plan);
7156        }
7157
7158        // Preserve column qualifiers so downstream plan nodes can keep referencing
7159        // the columns by their original qualified names.
7160        let project_exprs = schema
7161            .iter()
7162            .filter(|(_, field)| field.name() != DATA_SCHEMA_TSID_COLUMN_NAME)
7163            .map(|(qualifier, field)| {
7164                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
7165            })
7166            .collect::<Vec<_>>();
7167
7168        LogicalPlanBuilder::from(plan)
7169            .project(project_exprs)
7170            .context(DataFusionPlanningSnafu)?
7171            .build()
7172            .context(DataFusionPlanningSnafu)
7173    }
7174
7175    /// Apply an alias to the query result by adding a projection with the alias name
7176    fn apply_alias(&mut self, plan: LogicalPlan, alias_name: String) -> Result<LogicalPlan> {
7177        let fields_expr = self.create_field_column_exprs()?;
7178
7179        // TODO(dennis): how to support multi-value aliasing?
7180        ensure!(
7181            fields_expr.len() == 1,
7182            UnsupportedExprSnafu {
7183                name: "alias on multi-value result"
7184            }
7185        );
7186
7187        let project_fields = fields_expr
7188            .into_iter()
7189            .map(|expr| expr.alias(&alias_name))
7190            .chain(self.create_tag_column_exprs()?)
7191            .chain(Some(self.create_time_index_column_expr()?));
7192
7193        LogicalPlanBuilder::from(plan)
7194            .project(project_fields)
7195            .context(DataFusionPlanningSnafu)?
7196            .build()
7197            .context(DataFusionPlanningSnafu)
7198    }
7199}
7200
7201#[derive(Default, Debug)]
7202struct FunctionArgs {
7203    input: Option<PromExpr>,
7204    literals: Vec<DfExpr>,
7205}
7206
7207/// Represents different types of scalar functions supported in PromQL expressions.
7208/// Each variant defines how the function should be processed and what arguments it expects.
7209#[derive(Debug, Clone)]
7210enum ScalarFunc {
7211    /// DataFusion's registered(including built-in) scalar functions (e.g., abs, sqrt, round, clamp).
7212    /// These are passed through directly to DataFusion's execution engine.
7213    /// Processing: Simple argument insertion at the specified position.
7214    DataFusionBuiltin(Arc<ScalarUdfDef>),
7215    /// User-defined functions registered in DataFusion's function registry.
7216    /// Similar to DataFusionBuiltin but for custom functions not built into DataFusion.
7217    /// Processing: Direct pass-through with argument positioning.
7218    DataFusionUdf(Arc<ScalarUdfDef>),
7219    /// Native histogram helper UDFs. Non-histogram inputs are projected as NULL
7220    /// so the normal PromQL empty-value filter drops them.
7221    NativeHistogramUdf(Arc<ScalarUdfDef>),
7222    /// PromQL-specific functions that operate on time series data with temporal context.
7223    /// These functions require both timestamp ranges and values to perform calculations.
7224    /// Processing: Automatically injects timestamp_range and value columns as first arguments.
7225    /// Examples: idelta, irate, resets, changes, deriv, *_over_time function
7226    Udf(Arc<ScalarUdfDef>),
7227    /// PromQL functions requiring extrapolation calculations with explicit range information.
7228    /// These functions need to know the time range length to perform rate calculations.
7229    /// The second field contains the range length in milliseconds.
7230    /// Processing: Injects timestamp_range, value, time_index columns and appends range_length.
7231    /// Examples: increase, rate, delta
7232    // TODO(ruihang): maybe merge with Udf later
7233    ExtrapolateUdf(Arc<ScalarUdfDef>, i64),
7234    /// Functions that generate expressions directly without external UDF calls.
7235    /// The expression is constructed during function matching and requires no additional processing.
7236    /// Examples: time(), minute(), hour(), month(), year() and other date/time extractors
7237    GeneratedExpr,
7238}
7239
7240#[cfg(test)]
7241mod test {
7242    use std::time::{Duration, UNIX_EPOCH};
7243
7244    use catalog::RegisterTableRequest;
7245    use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager};
7246    use common_base::Plugins;
7247    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
7248    use common_query::native_histogram::{
7249        CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array,
7250    };
7251    use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
7252    use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
7253    use common_query::test_util::DummyDecoder;
7254    use common_recordbatch::RecordBatch as GreptimeRecordBatch;
7255    use datafusion::arrow::array::{
7256        Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray,
7257    };
7258    use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema};
7259    use datafusion::arrow::record_batch::RecordBatch;
7260    use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider};
7261    use datafusion::datasource::memory::MemorySourceConfig;
7262    use datafusion::datasource::source::DataSourceExec;
7263    use datafusion::datasource::{MemTable, provider_as_source};
7264    use datafusion::execution::context::SessionContext;
7265    use datafusion::logical_expr::Extension;
7266    use datatypes::prelude::ConcreteDataType;
7267    use datatypes::schema::{ColumnSchema, Schema};
7268    use promql_parser::label::Labels;
7269    use promql_parser::parser;
7270    use session::context::QueryContext;
7271    use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
7272    use table::Table;
7273    use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder};
7274    use table::test_util::{EmptyTable, MemTable as GreptimeMemTable};
7275
7276    use super::*;
7277    use crate::QueryEngineContext;
7278    use crate::options::QueryOptions;
7279    use crate::parser::QueryLanguageParser;
7280    use crate::query_engine::DefaultSerializer;
7281
7282    mod delta;
7283
7284    fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> {
7285        if let LogicalPlan::Extension(Extension { node }) = plan
7286            && let Some(instant_manipulate) = node.as_any().downcast_ref::<InstantManipulate>()
7287        {
7288            return Some(instant_manipulate);
7289        }
7290
7291        plan.inputs().into_iter().find_map(find_instant_manipulate)
7292    }
7293
7294    fn build_query_engine_state() -> QueryEngineState {
7295        QueryEngineState::new(
7296            new_memory_catalog_manager().unwrap(),
7297            None,
7298            None,
7299            None,
7300            None,
7301            None,
7302            false,
7303            Plugins::default(),
7304            QueryOptions::default(),
7305        )
7306    }
7307
7308    #[test]
7309    fn common_label_type_preserves_only_shared_dictionary_encoding() {
7310        let dictionary = ArrowDataType::Dictionary(
7311            Box::new(ArrowDataType::UInt32),
7312            Box::new(ArrowDataType::Utf8),
7313        );
7314        let other_dictionary = ArrowDataType::Dictionary(
7315            Box::new(ArrowDataType::Int32),
7316            Box::new(ArrowDataType::Utf8),
7317        );
7318
7319        assert_eq!(
7320            Some(dictionary.clone()),
7321            PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary))
7322        );
7323        assert_eq!(
7324            Some(ArrowDataType::Utf8),
7325            PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8))
7326        );
7327        assert_eq!(
7328            Some(ArrowDataType::Utf8),
7329            PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary))
7330        );
7331        assert_eq!(
7332            Some(ArrowDataType::Utf8),
7333            PromPlanner::common_label_data_type(Some(&dictionary), None)
7334        );
7335    }
7336
7337    async fn build_optimized_promql_plan(
7338        table_provider: DfTableSourceProvider,
7339        eval_stmt: &EvalStmt,
7340    ) -> LogicalPlan {
7341        let state = build_query_engine_state();
7342        let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state)
7343            .await
7344            .unwrap();
7345        let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
7346        state
7347            .optimize_by_extension_rules(raw_plan, &context)
7348            .unwrap()
7349    }
7350
7351    async fn build_optimized_tsid_plan(
7352        query: &str,
7353        num_tag: usize,
7354        num_field: usize,
7355        end_secs: u64,
7356        lookback_secs: u64,
7357    ) -> String {
7358        let eval_stmt = EvalStmt {
7359            expr: parser::parse(query).unwrap(),
7360            start: UNIX_EPOCH,
7361            end: UNIX_EPOCH
7362                .checked_add(Duration::from_secs(end_secs))
7363                .unwrap(),
7364            interval: Duration::from_secs(5),
7365            lookback_delta: Duration::from_secs(lookback_secs),
7366        };
7367        let table_provider = build_test_table_provider_with_tsid(
7368            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
7369            num_tag,
7370            num_field,
7371        )
7372        .await;
7373
7374        build_optimized_promql_plan(table_provider, &eval_stmt)
7375            .await
7376            .display_indent_schema()
7377            .to_string()
7378    }
7379
7380    async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) {
7381        let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
7382
7383        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
7384        assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0"));
7385        assert!(plan_str.contains("Distinct:"));
7386        assert!(plan_str.contains(expected_outer_agg), "{plan_str}");
7387        assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
7388    }
7389
7390    async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) {
7391        let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await;
7392        assert!(!plan_str.contains("Distinct:"), "{plan_str}");
7393    }
7394
7395    fn build_eval_stmt(expr: &str) -> EvalStmt {
7396        EvalStmt {
7397            expr: parser::parse(expr).unwrap(),
7398            start: UNIX_EPOCH,
7399            end: UNIX_EPOCH
7400                .checked_add(Duration::from_secs(100_000))
7401                .unwrap(),
7402            interval: Duration::from_secs(5),
7403            lookback_delta: Duration::from_secs(1),
7404        }
7405    }
7406
7407    enum DirectOrValue {
7408        Float64(f64),
7409        Int64(i64),
7410        NativeHistogram(NativeHistogram),
7411        Utf8(&'static str),
7412    }
7413
7414    impl DirectOrValue {
7415        fn data_type(&self) -> ArrowDataType {
7416            match self {
7417                Self::Float64(_) => ArrowDataType::Float64,
7418                Self::Int64(_) => ArrowDataType::Int64,
7419                Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(),
7420                Self::Utf8(_) => ArrowDataType::Utf8,
7421            }
7422        }
7423        fn array(&self) -> Arc<dyn Array> {
7424            match self {
7425                Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])),
7426                Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])),
7427                Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]),
7428                Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])),
7429            }
7430        }
7431    }
7432
7433    fn direct_or_histogram() -> NativeHistogram {
7434        NativeHistogram {
7435            schema: 0,
7436            zero_threshold: 0.0,
7437            sum: 1.0,
7438            reset_hint: CounterResetHint::Unknown,
7439            start_timestamp: None,
7440            custom_values: vec![],
7441            positive_spans: vec![],
7442            negative_spans: vec![],
7443            count: 1.0,
7444            zero_count: 1.0,
7445            positive_buckets: vec![],
7446            negative_buckets: vec![],
7447        }
7448    }
7449
7450    fn operator_metric_table(
7451        name: &str,
7452        table_id: u32,
7453        tag: &str,
7454        le: Option<&str>,
7455        value: DirectOrValue,
7456    ) -> table::TableRef {
7457        let value_type = match &value {
7458            DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(),
7459            DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(),
7460            DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(),
7461            DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(),
7462        };
7463        let tag_count = 1 + usize::from(le.is_some());
7464        let mut columns = vec![ColumnSchema::new(
7465            "tag".to_string(),
7466            ConcreteDataType::string_datatype(),
7467            false,
7468        )];
7469        if le.is_some() {
7470            columns.push(ColumnSchema::new(
7471                LE_COLUMN_NAME.to_string(),
7472                ConcreteDataType::string_datatype(),
7473                false,
7474            ));
7475        }
7476        columns.extend([
7477            ColumnSchema::new(
7478                "ts".to_string(),
7479                ConcreteDataType::timestamp_millisecond_datatype(),
7480                false,
7481            )
7482            .with_time_index(true),
7483            ColumnSchema::new("v".to_string(), value_type, true),
7484        ]);
7485        let schema = Arc::new(Schema::new(columns));
7486        let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc<dyn Array>];
7487        if let Some(le) = le {
7488            arrays.push(Arc::new(StringArray::from(vec![le])));
7489        }
7490        arrays.extend([
7491            Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc<dyn Array>,
7492            value.array(),
7493        ]);
7494        let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
7495        let backing = GreptimeMemTable::new_with_catalog(
7496            name,
7497            GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch),
7498            table_id,
7499            DEFAULT_CATALOG_NAME.to_string(),
7500            DEFAULT_SCHEMA_NAME.to_string(),
7501        );
7502        let value_index = tag_count + 1;
7503        let meta = TableMetaBuilder::empty()
7504            .schema(schema)
7505            .primary_key_indices((0..tag_count).collect())
7506            .value_indices(vec![value_index])
7507            .next_column_id((value_index + 1) as u32)
7508            .build()
7509            .unwrap();
7510        let info = Arc::new(
7511            TableInfoBuilder::default()
7512                .table_id(table_id)
7513                .name(name)
7514                .meta(meta)
7515                .build()
7516                .unwrap(),
7517        );
7518        Arc::new(Table::new(
7519            info,
7520            FilterPushDownType::Unsupported,
7521            backing.data_source(),
7522        ))
7523    }
7524
7525    fn operator_table_provider() -> DfTableSourceProvider {
7526        let catalog = MemoryCatalogManager::with_default_setup();
7527        let tables = [
7528            operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)),
7529            operator_metric_table(
7530                "lh",
7531                2_002,
7532                "b",
7533                None,
7534                DirectOrValue::NativeHistogram(direct_or_histogram()),
7535            ),
7536            operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)),
7537            operator_metric_table(
7538                "rh",
7539                2_004,
7540                "a",
7541                None,
7542                DirectOrValue::NativeHistogram(direct_or_histogram()),
7543            ),
7544            operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)),
7545            operator_metric_table(
7546                "bad_classic",
7547                2_006,
7548                "d",
7549                Some("broken"),
7550                DirectOrValue::Float64(1.0),
7551            ),
7552            operator_metric_table(
7553                "bad_native",
7554                2_007,
7555                "d",
7556                None,
7557                DirectOrValue::NativeHistogram(direct_or_histogram()),
7558            ),
7559        ];
7560        for table in tables {
7561            let info = table.table_info();
7562            catalog
7563                .register_table_sync(RegisterTableRequest {
7564                    catalog: DEFAULT_CATALOG_NAME.to_string(),
7565                    schema: DEFAULT_SCHEMA_NAME.to_string(),
7566                    table_name: info.name.clone(),
7567                    table_id: info.ident.table_id,
7568                    table,
7569                })
7570                .unwrap();
7571        }
7572        DfTableSourceProvider::new(
7573            catalog,
7574            false,
7575            QueryContext::arc(),
7576            DummyDecoder::arc(),
7577            false,
7578        )
7579    }
7580
7581    fn operator_eval_stmt(expr: &str) -> EvalStmt {
7582        let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap();
7583        EvalStmt {
7584            expr: parser::parse(expr).unwrap(),
7585            start: time,
7586            end: time,
7587            interval: Duration::from_secs(1),
7588            lookback_delta: Duration::from_secs(5),
7589        }
7590    }
7591
7592    struct DirectOrSource {
7593        name: &'static str,
7594        empty: bool,
7595        timestamp: i64,
7596        tags: Vec<(&'static str, Option<&'static str>)>,
7597        value: DirectOrValue,
7598    }
7599
7600    fn source(
7601        name: &'static str,
7602        empty: bool,
7603        timestamp: i64,
7604        tags: Vec<(&'static str, Option<&'static str>)>,
7605        value: DirectOrValue,
7606    ) -> DirectOrSource {
7607        DirectOrSource {
7608            name,
7609            empty,
7610            timestamp,
7611            tags,
7612            value,
7613        }
7614    }
7615
7616    fn tagged_source(
7617        name: &'static str,
7618        empty: bool,
7619        tag: (&'static str, Option<&'static str>),
7620        value: DirectOrValue,
7621    ) -> DirectOrSource {
7622        source(name, empty, 1, vec![("job", Some("job")), tag], value)
7623    }
7624
7625    fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource {
7626        source(name, true, 1, vec![("job", Some("job"))], value)
7627    }
7628
7629    fn table(source: &DirectOrSource) -> Arc<MemTable> {
7630        let mut fields = vec![Field::new(
7631            "ts",
7632            ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
7633            false,
7634        )];
7635        fields.extend(
7636            source
7637                .tags
7638                .iter()
7639                .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)),
7640        );
7641        fields.push(Field::new("v", source.value.data_type(), true));
7642        let schema = Arc::new(ArrowSchema::new(fields));
7643        let partitions = if source.empty {
7644            vec![vec![]]
7645        } else {
7646            let mut columns: Vec<Arc<dyn Array>> =
7647                vec![Arc::new(TimestampMillisecondArray::from(vec![
7648                    source.timestamp,
7649                ]))];
7650            columns.extend(
7651                source
7652                    .tags
7653                    .iter()
7654                    .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc<dyn Array>),
7655            );
7656            columns.push(source.value.array());
7657            vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]]
7658        };
7659        Arc::new(MemTable::try_new(schema, partitions).unwrap())
7660    }
7661
7662    fn scan(source: &DirectOrSource) -> LogicalPlan {
7663        LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None)
7664            .unwrap()
7665            .build()
7666            .unwrap()
7667    }
7668
7669    fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext {
7670        PromPlannerContext {
7671            table_name: Some(qualifier.to_string()),
7672            time_index_column: Some("ts".to_string()),
7673            field_columns: vec![field.to_string()],
7674            tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(),
7675            ..Default::default()
7676        }
7677    }
7678
7679    fn or_modifier(expr: &str) -> Option<BinModifier> {
7680        let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else {
7681            unreachable!()
7682        };
7683        expr.modifier
7684    }
7685
7686    async fn plan_direct_or(
7687        left: LogicalPlan,
7688        right: LogicalPlan,
7689        left_context: PromPlannerContext,
7690        right_context: PromPlannerContext,
7691        modifier: &Option<BinModifier>,
7692    ) -> LogicalPlan {
7693        let table_provider = build_test_table_provider_with_fields(
7694            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7695            &[],
7696        )
7697        .await;
7698        let mut planner = PromPlanner {
7699            table_provider,
7700            ctx: PromPlannerContext::default(),
7701            promql_annotations: None,
7702        };
7703        planner
7704            .or_operator(
7705                left,
7706                right,
7707                left_context.tag_columns.iter().cloned().collect(),
7708                right_context.tag_columns.iter().cloned().collect(),
7709                left_context,
7710                right_context,
7711                modifier,
7712            )
7713            .unwrap()
7714    }
7715
7716    async fn execute(
7717        plan: LogicalPlan,
7718        state: &QueryEngineState,
7719    ) -> (LogicalPlan, Vec<RecordBatch>) {
7720        let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
7721        let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
7722        let physical = state
7723            .session_state()
7724            .create_physical_plan(&optimized)
7725            .await
7726            .unwrap();
7727        let batches =
7728            datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
7729                .await
7730                .unwrap();
7731        (optimized, batches)
7732    }
7733
7734    async fn run(
7735        left: &DirectOrSource,
7736        right: &DirectOrSource,
7737        left_context: PromPlannerContext,
7738        right_context: PromPlannerContext,
7739        modifier: &Option<BinModifier>,
7740    ) -> (LogicalPlan, Vec<RecordBatch>) {
7741        let plan = plan_direct_or(
7742            scan(left),
7743            scan(right),
7744            left_context,
7745            right_context,
7746            modifier,
7747        )
7748        .await;
7749        execute(plan, &build_query_engine_state()).await
7750    }
7751
7752    async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) {
7753        let sample = |histogram: bool| {
7754            if histogram {
7755                DirectOrValue::NativeHistogram(direct_or_histogram())
7756            } else {
7757                DirectOrValue::Float64(1.25)
7758            }
7759        };
7760        let left = tagged_source(
7761            "lhs",
7762            false,
7763            (
7764                "k",
7765                Some(if histogram_on_left {
7766                    "histogram"
7767                } else {
7768                    "float"
7769                }),
7770            ),
7771            sample(histogram_on_left),
7772        );
7773        let right = tagged_source(
7774            "rhs",
7775            false,
7776            (
7777                "k",
7778                Some(if histogram_on_left {
7779                    "float"
7780                } else {
7781                    "histogram"
7782                }),
7783            ),
7784            sample(!histogram_on_left),
7785        );
7786        let table_provider = build_test_table_provider_with_fields(
7787            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7788            &[],
7789        )
7790        .await;
7791        let mut planner = PromPlanner {
7792            table_provider,
7793            ctx: PromPlannerContext::default(),
7794            promql_annotations: None,
7795        };
7796        let left_context = direct_or_context("lhs", &["job", "k"], "v");
7797        let right_context = direct_or_context("rhs", &["job", "k"], "v");
7798        let plan = planner
7799            .or_operator(
7800                scan(&left),
7801                scan(&right),
7802                left_context.tag_columns.iter().cloned().collect(),
7803                right_context.tag_columns.iter().cloned().collect(),
7804                left_context,
7805                right_context,
7806                &or_modifier("lhs or on(k) rhs"),
7807            )
7808            .unwrap();
7809        (planner, plan)
7810    }
7811
7812    async fn mixed_aggregate_input(histograms: Vec<NativeHistogram>) -> (PromPlanner, LogicalPlan) {
7813        let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0");
7814        let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0");
7815        let row_count = histograms.len() + 1;
7816        let schema = Arc::new(ArrowSchema::new(vec![
7817            Field::new(
7818                "ts",
7819                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
7820                false,
7821            ),
7822            Field::new("k", ArrowDataType::Utf8, false),
7823            Field::new(&float_field, ArrowDataType::Float64, true),
7824            Field::new(
7825                &histogram_field,
7826                native_histogram_value_type().as_arrow_type(),
7827                true,
7828            ),
7829        ]));
7830        let mut histogram_values = Vec::with_capacity(row_count);
7831        histogram_values.push(None);
7832        histogram_values.extend(histograms.into_iter().map(Some));
7833        let batch = RecordBatch::try_new(
7834            schema.clone(),
7835            vec![
7836                Arc::new(TimestampMillisecondArray::from(vec![1; row_count])),
7837                Arc::new(StringArray::from_iter_values(
7838                    (0..row_count).map(|row| format!("kind_{row}")),
7839                )),
7840                Arc::new(Float64Array::from_iter(
7841                    (0..row_count).map(|row| (row == 0).then_some(1.25)),
7842                )),
7843                build_histogram_array(&histogram_values),
7844            ],
7845        )
7846        .unwrap();
7847        let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
7848        let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
7849            .unwrap()
7850            .build()
7851            .unwrap();
7852        let table_provider = build_test_table_provider_with_fields(
7853            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
7854            &[],
7855        )
7856        .await;
7857        let planner = PromPlanner {
7858            table_provider,
7859            ctx: PromPlannerContext {
7860                table_name: Some("mixed".to_string()),
7861                time_index_column: Some("ts".to_string()),
7862                field_columns: vec![float_field, histogram_field],
7863                tag_columns: vec!["k".to_string()],
7864                ..Default::default()
7865            },
7866            promql_annotations: None,
7867        };
7868        (planner, plan)
7869    }
7870
7871    fn assert_no_internal_or_keys(schema: &DFSchema) {
7872        assert!(
7873            schema
7874                .fields()
7875                .iter()
7876                .all(|field| !field.name().starts_with("__promql_or_match_")),
7877            "{schema:?}"
7878        );
7879    }
7880
7881    fn values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
7882        batches
7883            .iter()
7884            .flat_map(|batch| {
7885                batch
7886                    .column_by_name(column)
7887                    .unwrap()
7888                    .as_any()
7889                    .downcast_ref::<Float64Array>()
7890                    .unwrap()
7891                    .iter()
7892                    .flatten()
7893            })
7894            .collect()
7895    }
7896
7897    fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec<f64> {
7898        batches
7899            .iter()
7900            .flat_map(|batch| {
7901                let values = datafusion::arrow::compute::cast(
7902                    batch.column_by_name(column).unwrap(),
7903                    &ArrowDataType::Float64,
7904                )
7905                .unwrap();
7906                values
7907                    .as_any()
7908                    .downcast_ref::<Float64Array>()
7909                    .unwrap()
7910                    .iter()
7911                    .flatten()
7912                    .collect::<Vec<_>>()
7913            })
7914            .collect()
7915    }
7916
7917    fn histograms(batches: &[RecordBatch], column: &str) -> Vec<NativeHistogram> {
7918        batches
7919            .iter()
7920            .flat_map(|batch| {
7921                let values = batch
7922                    .column_by_name(column)
7923                    .unwrap()
7924                    .as_any()
7925                    .downcast_ref::<datafusion::arrow::array::StructArray>()
7926                    .unwrap();
7927                (0..values.len()).filter_map(|row| {
7928                    common_query::native_histogram::read_histogram(values, row).unwrap()
7929                })
7930            })
7931            .collect()
7932    }
7933
7934    fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option<String>)> {
7935        let mut rows = batches
7936            .iter()
7937            .flat_map(|batch| {
7938                let values = batch
7939                    .column_by_name("v")
7940                    .unwrap()
7941                    .as_any()
7942                    .downcast_ref::<Float64Array>()
7943                    .unwrap();
7944                let labels = batch
7945                    .column_by_name("k")
7946                    .map(|column| column.as_any().downcast_ref::<StringArray>().unwrap());
7947                (0..batch.num_rows()).map(move |i| {
7948                    (
7949                        values.value(i),
7950                        labels.and_then(|labels| {
7951                            (!labels.is_null(i)).then(|| labels.value(i).to_string())
7952                        }),
7953                    )
7954                })
7955            })
7956            .collect::<Vec<_>>();
7957        rows.sort_by(|left, right| left.0.total_cmp(&right.0));
7958        rows
7959    }
7960
7961    fn matrix_source(
7962        name: &'static str,
7963        k: Option<Option<&'static str>>,
7964        timestamp: i64,
7965        value: f64,
7966    ) -> DirectOrSource {
7967        let mut tags = vec![("job", Some("job"))];
7968        if let Some(k) = k {
7969            tags.push(("k", k));
7970        }
7971        source(name, false, timestamp, tags, DirectOrValue::Float64(value))
7972    }
7973
7974    fn matrix_context(name: &str, k: Option<Option<&str>>) -> PromPlannerContext {
7975        direct_or_context(
7976            name,
7977            if k.is_some() { &["job", "k"] } else { &["job"] },
7978            "v",
7979        )
7980    }
7981
7982    async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider {
7983        build_test_table_provider_with_fields(
7984            &[
7985                (
7986                    DEFAULT_SCHEMA_NAME.to_string(),
7987                    "non_existent_histogram_bucket".to_string(),
7988                ),
7989                (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()),
7990            ],
7991            &["pod", "instance"],
7992        )
7993        .await
7994    }
7995
7996    fn assert_normal_metric_schema(plan: &LogicalPlan) {
7997        let fields = plan.schema().fields();
7998        assert_eq!(fields.len(), 4, "{fields:?}");
7999        assert!(
8000            fields.iter().any(|field| field.name() == "pod"),
8001            "{fields:?}"
8002        );
8003        assert!(
8004            fields.iter().any(|field| field.name() == "instance"),
8005            "{fields:?}"
8006        );
8007        assert!(
8008            fields
8009                .iter()
8010                .any(|field| field.name() == greptime_timestamp()),
8011            "{fields:?}"
8012        );
8013        assert!(
8014            fields.iter().any(|field| {
8015                field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64
8016            }),
8017            "{fields:?}"
8018        );
8019    }
8020
8021    async fn build_test_table_provider_with_distinct_tags(
8022        table_tags: &[(&str, &[&str])],
8023    ) -> DfTableSourceProvider {
8024        let catalog_list = MemoryCatalogManager::with_default_setup();
8025        for (table_name, tags) in table_tags {
8026            let mut columns = tags
8027                .iter()
8028                .map(|tag| {
8029                    ColumnSchema::new(
8030                        (*tag).to_string(),
8031                        ConcreteDataType::string_datatype(),
8032                        false,
8033                    )
8034                })
8035                .collect::<Vec<_>>();
8036            columns.push(
8037                ColumnSchema::new(
8038                    greptime_timestamp().to_string(),
8039                    ConcreteDataType::timestamp_millisecond_datatype(),
8040                    false,
8041                )
8042                .with_time_index(true),
8043            );
8044            columns.push(ColumnSchema::new(
8045                greptime_value().to_string(),
8046                ConcreteDataType::float64_datatype(),
8047                true,
8048            ));
8049            let table_meta = TableMetaBuilder::empty()
8050                .schema(Arc::new(Schema::new(columns)))
8051                .primary_key_indices((0..tags.len()).collect())
8052                .next_column_id(1024)
8053                .build()
8054                .unwrap();
8055            let table_info = TableInfoBuilder::default()
8056                .name((*table_name).to_string())
8057                .meta(table_meta)
8058                .build()
8059                .unwrap();
8060
8061            assert!(
8062                catalog_list
8063                    .register_table_sync(RegisterTableRequest {
8064                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8065                        schema: DEFAULT_SCHEMA_NAME.to_string(),
8066                        table_name: (*table_name).to_string(),
8067                        table_id: 1024,
8068                        table: EmptyTable::from_table_info(&table_info),
8069                    })
8070                    .is_ok()
8071            );
8072        }
8073
8074        DfTableSourceProvider::new(
8075            catalog_list,
8076            false,
8077            QueryContext::arc(),
8078            DummyDecoder::arc(),
8079            false,
8080        )
8081    }
8082
8083    fn contains_histogram_fold(plan: &LogicalPlan) -> bool {
8084        matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::<HistogramFold>())
8085            || plan.inputs().into_iter().any(contains_histogram_fold)
8086    }
8087
8088    async fn build_set_op_context_table_provider() -> DfTableSourceProvider {
8089        build_test_table_provider_with_distinct_tags(&[
8090            ("bucket_metric", &["job", "le"]),
8091            ("normal_metric", &["job"]),
8092            ("fallback_metric", &["instance"]),
8093        ])
8094        .await
8095    }
8096
8097    async fn build_or_context_table_provider() -> DfTableSourceProvider {
8098        build_test_table_provider_with_distinct_tags(&[
8099            ("normal_metric", &["job"]),
8100            ("other_metric", &["instance"]),
8101            ("non_hist_metric", &["instance"]),
8102        ])
8103        .await
8104    }
8105
8106    async fn optimize_and_create_physical_plan(
8107        state: &QueryEngineState,
8108        plan: LogicalPlan,
8109    ) -> (
8110        LogicalPlan,
8111        Arc<dyn datafusion::physical_plan::ExecutionPlan>,
8112    ) {
8113        let context = QueryEngineContext::new(state.session_state(), QueryContext::arc());
8114        let optimized = state.optimize_by_extension_rules(plan, &context).unwrap();
8115        let physical = state
8116            .session_state()
8117            .create_physical_plan(&optimized)
8118            .await
8119            .unwrap();
8120        (optimized, physical)
8121    }
8122
8123    async fn build_test_table_provider(
8124        table_name_tuples: &[(String, String)],
8125        num_tag: usize,
8126        num_field: usize,
8127    ) -> DfTableSourceProvider {
8128        let catalog_list = MemoryCatalogManager::with_default_setup();
8129        for (schema_name, table_name) in table_name_tuples {
8130            let mut columns = vec![];
8131            for i in 0..num_tag {
8132                columns.push(ColumnSchema::new(
8133                    format!("tag_{i}"),
8134                    ConcreteDataType::string_datatype(),
8135                    false,
8136                ));
8137            }
8138            columns.push(
8139                ColumnSchema::new(
8140                    "timestamp".to_string(),
8141                    ConcreteDataType::timestamp_millisecond_datatype(),
8142                    false,
8143                )
8144                .with_time_index(true),
8145            );
8146            for i in 0..num_field {
8147                columns.push(ColumnSchema::new(
8148                    format!("field_{i}"),
8149                    ConcreteDataType::float64_datatype(),
8150                    true,
8151                ));
8152            }
8153            let schema = Arc::new(Schema::new(columns));
8154            let table_meta = TableMetaBuilder::empty()
8155                .schema(schema)
8156                .primary_key_indices((0..num_tag).collect())
8157                .value_indices((num_tag + 1..num_tag + 1 + num_field).collect())
8158                .next_column_id(1024)
8159                .build()
8160                .unwrap();
8161            let table_info = TableInfoBuilder::default()
8162                .name(table_name.clone())
8163                .meta(table_meta)
8164                .build()
8165                .unwrap();
8166            let table = EmptyTable::from_table_info(&table_info);
8167
8168            assert!(
8169                catalog_list
8170                    .register_table_sync(RegisterTableRequest {
8171                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8172                        schema: schema_name.clone(),
8173                        table_name: table_name.clone(),
8174                        table_id: 1024,
8175                        table,
8176                    })
8177                    .is_ok()
8178            );
8179        }
8180
8181        DfTableSourceProvider::new(
8182            catalog_list,
8183            false,
8184            QueryContext::arc(),
8185            DummyDecoder::arc(),
8186            false,
8187        )
8188    }
8189
8190    async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8191        build_test_native_histogram_table_provider_with_marker(table_name, false).await
8192    }
8193
8194    async fn build_test_native_histogram_table_provider_with_marker(
8195        table_name: &str,
8196        temporality_marker: bool,
8197    ) -> DfTableSourceProvider {
8198        let catalog_list = MemoryCatalogManager::with_default_setup();
8199        let mut columns = vec![
8200            ColumnSchema::new(
8201                "tag_0".to_string(),
8202                ConcreteDataType::string_datatype(),
8203                false,
8204            ),
8205            ColumnSchema::new(
8206                LE_COLUMN_NAME.to_string(),
8207                ConcreteDataType::string_datatype(),
8208                true,
8209            ),
8210        ];
8211        if temporality_marker {
8212            columns.push(ColumnSchema::new(
8213                OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8214                ConcreteDataType::string_datatype(),
8215                true,
8216            ));
8217        }
8218        let tag_count = columns.len();
8219        columns.extend([
8220            ColumnSchema::new(
8221                "timestamp".to_string(),
8222                ConcreteDataType::timestamp_millisecond_datatype(),
8223                false,
8224            )
8225            .with_time_index(true),
8226            ColumnSchema::new(
8227                greptime_native_histogram().to_string(),
8228                native_histogram_value_type().clone(),
8229                true,
8230            ),
8231        ]);
8232        let schema = Arc::new(Schema::new(columns));
8233        let table_meta = TableMetaBuilder::empty()
8234            .schema(schema)
8235            .primary_key_indices((0..tag_count).collect())
8236            .value_indices(vec![tag_count + 1])
8237            .next_column_id(1024)
8238            .build()
8239            .unwrap();
8240        let table_info = TableInfoBuilder::default()
8241            .name(table_name)
8242            .meta(table_meta)
8243            .build()
8244            .unwrap();
8245        let table = EmptyTable::from_table_info(&table_info);
8246
8247        assert!(
8248            catalog_list
8249                .register_table_sync(RegisterTableRequest {
8250                    catalog: DEFAULT_CATALOG_NAME.to_string(),
8251                    schema: DEFAULT_SCHEMA_NAME.to_string(),
8252                    table_name: table_name.to_string(),
8253                    table_id: 1024,
8254                    table,
8255                })
8256                .is_ok()
8257        );
8258
8259        DfTableSourceProvider::new(
8260            catalog_list,
8261            false,
8262            QueryContext::arc(),
8263            DummyDecoder::arc(),
8264            false,
8265        )
8266    }
8267
8268    async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider {
8269        let catalog_list = MemoryCatalogManager::with_default_setup();
8270        let columns = vec![
8271            ColumnSchema::new(
8272                "tag_0".to_string(),
8273                ConcreteDataType::string_datatype(),
8274                false,
8275            ),
8276            ColumnSchema::new(
8277                "timestamp".to_string(),
8278                ConcreteDataType::timestamp_millisecond_datatype(),
8279                false,
8280            )
8281            .with_time_index(true),
8282            ColumnSchema::new(
8283                greptime_native_histogram().to_string(),
8284                native_histogram_value_type().clone(),
8285                true,
8286            ),
8287            ColumnSchema::new(
8288                "native_histogram_2".to_string(),
8289                native_histogram_value_type().clone(),
8290                true,
8291            ),
8292        ];
8293        let schema = Arc::new(Schema::new(columns));
8294        let table_meta = TableMetaBuilder::empty()
8295            .schema(schema)
8296            .primary_key_indices(vec![0])
8297            .value_indices(vec![2, 3])
8298            .next_column_id(1024)
8299            .build()
8300            .unwrap();
8301        let table_info = TableInfoBuilder::default()
8302            .name(table_name)
8303            .meta(table_meta)
8304            .build()
8305            .unwrap();
8306        let table = EmptyTable::from_table_info(&table_info);
8307
8308        assert!(
8309            catalog_list
8310                .register_table_sync(RegisterTableRequest {
8311                    catalog: DEFAULT_CATALOG_NAME.to_string(),
8312                    schema: DEFAULT_SCHEMA_NAME.to_string(),
8313                    table_name: table_name.to_string(),
8314                    table_id: 1024,
8315                    table,
8316                })
8317                .is_ok()
8318        );
8319
8320        DfTableSourceProvider::new(
8321            catalog_list,
8322            false,
8323            QueryContext::arc(),
8324            DummyDecoder::arc(),
8325            false,
8326        )
8327    }
8328
8329    async fn build_test_mixed_native_histogram_table_provider(
8330        table_name: &str,
8331    ) -> DfTableSourceProvider {
8332        build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await
8333    }
8334
8335    async fn build_test_mixed_native_histogram_table_provider_with_marker(
8336        table_name: &str,
8337        temporality_marker: bool,
8338    ) -> DfTableSourceProvider {
8339        let catalog_list = MemoryCatalogManager::with_default_setup();
8340        let mut columns = vec![ColumnSchema::new(
8341            "tag_0".to_string(),
8342            ConcreteDataType::string_datatype(),
8343            false,
8344        )];
8345        if temporality_marker {
8346            columns.push(ColumnSchema::new(
8347                OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(),
8348                ConcreteDataType::string_datatype(),
8349                true,
8350            ));
8351        }
8352        let tag_count = columns.len();
8353        columns.extend([
8354            ColumnSchema::new(
8355                "timestamp".to_string(),
8356                ConcreteDataType::timestamp_millisecond_datatype(),
8357                false,
8358            )
8359            .with_time_index(true),
8360            ColumnSchema::new(
8361                greptime_native_histogram().to_string(),
8362                native_histogram_value_type().clone(),
8363                true,
8364            ),
8365            ColumnSchema::new(
8366                greptime_value().to_string(),
8367                ConcreteDataType::float64_datatype(),
8368                true,
8369            ),
8370        ]);
8371        let schema = Arc::new(Schema::new(columns));
8372        let table_meta = TableMetaBuilder::empty()
8373            .schema(schema.clone())
8374            .primary_key_indices((0..tag_count).collect())
8375            .value_indices(vec![tag_count + 1, tag_count + 2])
8376            .next_column_id(1024)
8377            .build()
8378            .unwrap();
8379        let table_info = Arc::new(
8380            TableInfoBuilder::default()
8381                .name(table_name)
8382                .meta(table_meta)
8383                .build()
8384                .unwrap(),
8385        );
8386        let mut arrays: Vec<Arc<dyn Array>> =
8387            vec![Arc::new(StringArray::from(vec!["float", "histogram"]))];
8388        if temporality_marker {
8389            arrays.push(Arc::new(StringArray::from(vec![
8390                Some(GREPTIME_TEMPORALITY_DELTA),
8391                Some(GREPTIME_TEMPORALITY_DELTA),
8392            ])));
8393        }
8394        arrays.extend([
8395            Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc<dyn Array>,
8396            build_histogram_array(&[None, Some(direct_or_histogram())]),
8397            Arc::new(Float64Array::from(vec![Some(2.0), None])),
8398        ]);
8399        let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap();
8400        let backing = GreptimeMemTable::new_with_catalog(
8401            table_name,
8402            GreptimeRecordBatch::from_df_record_batch(schema, batch),
8403            1024,
8404            DEFAULT_CATALOG_NAME.to_string(),
8405            DEFAULT_SCHEMA_NAME.to_string(),
8406        );
8407        let table = Arc::new(Table::new(
8408            table_info,
8409            FilterPushDownType::Unsupported,
8410            backing.data_source(),
8411        ));
8412
8413        assert!(
8414            catalog_list
8415                .register_table_sync(RegisterTableRequest {
8416                    catalog: DEFAULT_CATALOG_NAME.to_string(),
8417                    schema: DEFAULT_SCHEMA_NAME.to_string(),
8418                    table_name: table_name.to_string(),
8419                    table_id: 1024,
8420                    table,
8421                })
8422                .is_ok()
8423        );
8424
8425        DfTableSourceProvider::new(
8426            catalog_list,
8427            false,
8428            QueryContext::arc(),
8429            DummyDecoder::arc(),
8430            false,
8431        )
8432    }
8433
8434    fn classic_and_native_histogram_table_provider(
8435        native_tag: &str,
8436        native_le: Option<&str>,
8437        native_histogram: NativeHistogram,
8438    ) -> DfTableSourceProvider {
8439        let table_name = "mixed_histogram";
8440        let catalog = MemoryCatalogManager::with_default_setup();
8441        let schema = Arc::new(Schema::new(vec![
8442            ColumnSchema::new(
8443                "tag".to_string(),
8444                ConcreteDataType::string_datatype(),
8445                false,
8446            ),
8447            ColumnSchema::new(
8448                LE_COLUMN_NAME.to_string(),
8449                ConcreteDataType::string_datatype(),
8450                true,
8451            ),
8452            ColumnSchema::new(
8453                "timestamp".to_string(),
8454                ConcreteDataType::timestamp_millisecond_datatype(),
8455                false,
8456            )
8457            .with_time_index(true),
8458            ColumnSchema::new(
8459                greptime_native_histogram().to_string(),
8460                native_histogram_value_type().clone(),
8461                true,
8462            ),
8463            ColumnSchema::new(
8464                greptime_value().to_string(),
8465                ConcreteDataType::float64_datatype(),
8466                true,
8467            ),
8468        ]));
8469        let table_meta = TableMetaBuilder::empty()
8470            .schema(schema.clone())
8471            .primary_key_indices(vec![0, 1])
8472            .value_indices(vec![3, 4])
8473            .next_column_id(5)
8474            .build()
8475            .unwrap();
8476        let table_info = Arc::new(
8477            TableInfoBuilder::default()
8478                .name(table_name)
8479                .meta(table_meta)
8480                .build()
8481                .unwrap(),
8482        );
8483        let batch = RecordBatch::try_new(
8484            schema.arrow_schema().clone(),
8485            vec![
8486                Arc::new(StringArray::from(vec![
8487                    "classic", "classic", native_tag, "classic", "classic", native_tag,
8488                ])),
8489                Arc::new(StringArray::from(vec![
8490                    Some("1"),
8491                    Some("+Inf"),
8492                    native_le,
8493                    Some("1"),
8494                    Some("+Inf"),
8495                    native_le,
8496                ])),
8497                Arc::new(TimestampMillisecondArray::from(vec![
8498                    1_000, 1_000, 1_000, 2_000, 2_000, 2_000,
8499                ])),
8500                build_histogram_array(&[
8501                    None,
8502                    None,
8503                    Some(native_histogram.clone()),
8504                    None,
8505                    None,
8506                    Some(native_histogram),
8507                ]),
8508                Arc::new(Float64Array::from(vec![
8509                    Some(2.0),
8510                    Some(4.0),
8511                    None,
8512                    Some(2.0),
8513                    Some(4.0),
8514                    None,
8515                ])),
8516            ],
8517        )
8518        .unwrap();
8519        let backing = GreptimeMemTable::new_with_catalog(
8520            table_name,
8521            GreptimeRecordBatch::from_df_record_batch(schema, batch),
8522            2_200,
8523            DEFAULT_CATALOG_NAME.to_string(),
8524            DEFAULT_SCHEMA_NAME.to_string(),
8525        );
8526        let table = Arc::new(Table::new(
8527            table_info,
8528            FilterPushDownType::Unsupported,
8529            backing.data_source(),
8530        ));
8531        catalog
8532            .register_table_sync(RegisterTableRequest {
8533                catalog: DEFAULT_CATALOG_NAME.to_string(),
8534                schema: DEFAULT_SCHEMA_NAME.to_string(),
8535                table_name: table_name.to_string(),
8536                table_id: 2_200,
8537                table,
8538            })
8539            .unwrap();
8540
8541        DfTableSourceProvider::new(
8542            catalog,
8543            false,
8544            QueryContext::arc(),
8545            DummyDecoder::arc(),
8546            false,
8547        )
8548    }
8549
8550    async fn build_test_table_provider_with_tsid(
8551        table_name_tuples: &[(String, String)],
8552        num_tag: usize,
8553        num_field: usize,
8554    ) -> DfTableSourceProvider {
8555        let table_specs = table_name_tuples
8556            .iter()
8557            .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field))
8558            .collect::<Vec<_>>();
8559        build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await
8560    }
8561
8562    async fn build_test_table_provider_with_tsid_fields(
8563        table_specs: &[((String, String), usize)],
8564        num_tag: usize,
8565    ) -> DfTableSourceProvider {
8566        let table_specs = table_specs
8567            .iter()
8568            .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field))
8569            .collect::<Vec<_>>();
8570        build_test_table_provider_with_tsid_tag_fields(&table_specs).await
8571    }
8572
8573    async fn build_test_table_provider_with_tsid_tag_fields(
8574        table_specs: &[((String, String), usize, usize)],
8575    ) -> DfTableSourceProvider {
8576        let catalog_list = MemoryCatalogManager::with_default_setup();
8577
8578        let physical_table_name = "phy";
8579        let physical_table_id = 999u32;
8580        let physical_num_tag = table_specs
8581            .iter()
8582            .map(|(_, num_tag, _)| *num_tag)
8583            .max()
8584            .unwrap_or(0);
8585        let physical_num_field = table_specs
8586            .iter()
8587            .map(|(_, _, num_field)| *num_field)
8588            .max()
8589            .unwrap_or(0);
8590
8591        // Register a metric engine physical table with internal columns.
8592        {
8593            let mut columns = vec![
8594                ColumnSchema::new(
8595                    DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
8596                    ConcreteDataType::uint32_datatype(),
8597                    false,
8598                ),
8599                ColumnSchema::new(
8600                    DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
8601                    ConcreteDataType::uint64_datatype(),
8602                    false,
8603                ),
8604            ];
8605            for i in 0..physical_num_tag {
8606                columns.push(ColumnSchema::new(
8607                    format!("tag_{i}"),
8608                    ConcreteDataType::string_datatype(),
8609                    false,
8610                ));
8611            }
8612            columns.push(
8613                ColumnSchema::new(
8614                    "timestamp".to_string(),
8615                    ConcreteDataType::timestamp_millisecond_datatype(),
8616                    false,
8617                )
8618                .with_time_index(true),
8619            );
8620            for i in 0..physical_num_field {
8621                columns.push(ColumnSchema::new(
8622                    format!("field_{i}"),
8623                    ConcreteDataType::float64_datatype(),
8624                    true,
8625                ));
8626            }
8627
8628            let schema = Arc::new(Schema::new(columns));
8629            let primary_key_indices = (0..(2 + physical_num_tag)).collect::<Vec<_>>();
8630            let table_meta = TableMetaBuilder::empty()
8631                .schema(schema)
8632                .primary_key_indices(primary_key_indices)
8633                .value_indices(
8634                    (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(),
8635                )
8636                .engine(METRIC_ENGINE_NAME.to_string())
8637                .next_column_id(1024)
8638                .build()
8639                .unwrap();
8640            let table_info = TableInfoBuilder::default()
8641                .table_id(physical_table_id)
8642                .name(physical_table_name)
8643                .meta(table_meta)
8644                .build()
8645                .unwrap();
8646            let table = EmptyTable::from_table_info(&table_info);
8647
8648            assert!(
8649                catalog_list
8650                    .register_table_sync(RegisterTableRequest {
8651                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8652                        schema: DEFAULT_SCHEMA_NAME.to_string(),
8653                        table_name: physical_table_name.to_string(),
8654                        table_id: physical_table_id,
8655                        table,
8656                    })
8657                    .is_ok()
8658            );
8659        }
8660
8661        // Register metric engine logical tables without `__tsid`, referencing the physical table.
8662        for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate()
8663        {
8664            let mut columns = vec![];
8665            for i in 0..*num_tag {
8666                columns.push(ColumnSchema::new(
8667                    format!("tag_{i}"),
8668                    ConcreteDataType::string_datatype(),
8669                    false,
8670                ));
8671            }
8672            columns.push(
8673                ColumnSchema::new(
8674                    "timestamp".to_string(),
8675                    ConcreteDataType::timestamp_millisecond_datatype(),
8676                    false,
8677                )
8678                .with_time_index(true),
8679            );
8680            for i in 0..*num_field {
8681                columns.push(ColumnSchema::new(
8682                    format!("field_{i}"),
8683                    ConcreteDataType::float64_datatype(),
8684                    true,
8685                ));
8686            }
8687
8688            let schema = Arc::new(Schema::new(columns));
8689            let mut options = table::requests::TableOptions::default();
8690            options.extra_options.insert(
8691                LOGICAL_TABLE_METADATA_KEY.to_string(),
8692                physical_table_name.to_string(),
8693            );
8694            let table_id = 1024u32 + idx as u32;
8695            let table_meta = TableMetaBuilder::empty()
8696                .schema(schema)
8697                .primary_key_indices((0..*num_tag).collect())
8698                .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect())
8699                .engine(METRIC_ENGINE_NAME.to_string())
8700                .options(options)
8701                .next_column_id(1024)
8702                .build()
8703                .unwrap();
8704            let table_info = TableInfoBuilder::default()
8705                .table_id(table_id)
8706                .name(table_name.clone())
8707                .meta(table_meta)
8708                .build()
8709                .unwrap();
8710            let table = EmptyTable::from_table_info(&table_info);
8711
8712            assert!(
8713                catalog_list
8714                    .register_table_sync(RegisterTableRequest {
8715                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8716                        schema: schema_name.clone(),
8717                        table_name: table_name.clone(),
8718                        table_id,
8719                        table,
8720                    })
8721                    .is_ok()
8722            );
8723        }
8724
8725        DfTableSourceProvider::new(
8726            catalog_list,
8727            false,
8728            QueryContext::arc(),
8729            DummyDecoder::arc(),
8730            false,
8731        )
8732    }
8733
8734    async fn build_test_table_provider_with_fields(
8735        table_name_tuples: &[(String, String)],
8736        tags: &[&str],
8737    ) -> DfTableSourceProvider {
8738        let catalog_list = MemoryCatalogManager::with_default_setup();
8739        for (schema_name, table_name) in table_name_tuples {
8740            let mut columns = vec![];
8741            let num_tag = tags.len();
8742            for tag in tags {
8743                columns.push(ColumnSchema::new(
8744                    tag.to_string(),
8745                    ConcreteDataType::string_datatype(),
8746                    false,
8747                ));
8748            }
8749            columns.push(
8750                ColumnSchema::new(
8751                    greptime_timestamp().to_string(),
8752                    ConcreteDataType::timestamp_millisecond_datatype(),
8753                    false,
8754                )
8755                .with_time_index(true),
8756            );
8757            columns.push(ColumnSchema::new(
8758                greptime_value().to_string(),
8759                ConcreteDataType::float64_datatype(),
8760                true,
8761            ));
8762            let schema = Arc::new(Schema::new(columns));
8763            let table_meta = TableMetaBuilder::empty()
8764                .schema(schema)
8765                .primary_key_indices((0..num_tag).collect())
8766                .next_column_id(1024)
8767                .build()
8768                .unwrap();
8769            let table_info = TableInfoBuilder::default()
8770                .name(table_name.clone())
8771                .meta(table_meta)
8772                .build()
8773                .unwrap();
8774            let table = EmptyTable::from_table_info(&table_info);
8775
8776            assert!(
8777                catalog_list
8778                    .register_table_sync(RegisterTableRequest {
8779                        catalog: DEFAULT_CATALOG_NAME.to_string(),
8780                        schema: schema_name.clone(),
8781                        table_name: table_name.clone(),
8782                        table_id: 1024,
8783                        table,
8784                    })
8785                    .is_ok()
8786            );
8787        }
8788
8789        DfTableSourceProvider::new(
8790            catalog_list,
8791            false,
8792            QueryContext::arc(),
8793            DummyDecoder::arc(),
8794            false,
8795        )
8796    }
8797
8798    // {
8799    //     input: `abs(some_metric{foo!="bar"})`,
8800    //     expected: &Call{
8801    //         Func: MustGetFunction("abs"),
8802    //         Args: Expressions{
8803    //             &VectorSelector{
8804    //                 Name: "some_metric",
8805    //                 LabelMatchers: []*labels.Matcher{
8806    //                     MustLabelMatcher(labels.MatchNotEqual, "foo", "bar"),
8807    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
8808    //                 },
8809    //             },
8810    //         },
8811    //     },
8812    // },
8813    async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) {
8814        let prom_expr =
8815            parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap();
8816        let eval_stmt = EvalStmt {
8817            expr: prom_expr,
8818            start: UNIX_EPOCH,
8819            end: UNIX_EPOCH
8820                .checked_add(Duration::from_secs(100_000))
8821                .unwrap(),
8822            interval: Duration::from_secs(5),
8823            lookback_delta: Duration::from_secs(1),
8824        };
8825
8826        let table_provider = build_test_table_provider(
8827            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8828            1,
8829            1,
8830        )
8831        .await;
8832        let plan =
8833            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8834                .await
8835                .unwrap();
8836
8837        let expected = String::from(
8838            "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\
8839            \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]\
8840            \n    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8841            \n      PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8842            \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]\
8843	            \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]\
8844            \n            TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"
8845        ).replace("TEMPLATE", plan_name);
8846
8847        assert_eq!(plan.display_indent_schema().to_string(), expected);
8848    }
8849
8850    #[tokio::test]
8851    async fn single_abs() {
8852        do_single_instant_function_call("abs", "abs").await;
8853    }
8854
8855    #[tokio::test]
8856    #[should_panic]
8857    async fn single_absent() {
8858        do_single_instant_function_call("absent", "").await;
8859    }
8860
8861    #[tokio::test]
8862    async fn single_ceil() {
8863        do_single_instant_function_call("ceil", "ceil").await;
8864    }
8865
8866    #[tokio::test]
8867    async fn single_exp() {
8868        do_single_instant_function_call("exp", "exp").await;
8869    }
8870
8871    #[tokio::test]
8872    async fn single_ln() {
8873        do_single_instant_function_call("ln", "ln").await;
8874    }
8875
8876    #[tokio::test]
8877    async fn single_log2() {
8878        do_single_instant_function_call("log2", "log2").await;
8879    }
8880
8881    #[tokio::test]
8882    async fn single_log10() {
8883        do_single_instant_function_call("log10", "log10").await;
8884    }
8885
8886    #[tokio::test]
8887    #[should_panic]
8888    async fn single_scalar() {
8889        do_single_instant_function_call("scalar", "").await;
8890    }
8891
8892    #[tokio::test]
8893    #[should_panic]
8894    async fn single_sgn() {
8895        do_single_instant_function_call("sgn", "").await;
8896    }
8897
8898    #[tokio::test]
8899    #[should_panic]
8900    async fn single_sort() {
8901        do_single_instant_function_call("sort", "").await;
8902    }
8903
8904    #[tokio::test]
8905    #[should_panic]
8906    async fn single_sort_desc() {
8907        do_single_instant_function_call("sort_desc", "").await;
8908    }
8909
8910    #[tokio::test]
8911    async fn single_sqrt() {
8912        do_single_instant_function_call("sqrt", "sqrt").await;
8913    }
8914
8915    #[tokio::test]
8916    async fn single_timestamp_plan_preserves_source_value() {
8917        let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#);
8918        let table_provider = build_test_table_provider(
8919            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
8920            1,
8921            1,
8922        )
8923        .await;
8924        let plan =
8925            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
8926                .await
8927                .unwrap();
8928
8929        let expected = String::from(
8930            "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8931            \n  Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8932            \n    Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\
8933            \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]\
8934            \n        Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(Some(1),1,0) + Decimal128(Some(0),19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\
8935            \n          PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
8936            \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]\
8937            \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]\
8938            \n                TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
8939        );
8940
8941        assert_eq!(plan.display_indent_schema().to_string(), expected);
8942    }
8943
8944    #[tokio::test]
8945    async fn single_acos() {
8946        do_single_instant_function_call("acos", "acos").await;
8947    }
8948
8949    #[tokio::test]
8950    #[should_panic]
8951    async fn single_acosh() {
8952        do_single_instant_function_call("acosh", "").await;
8953    }
8954
8955    #[tokio::test]
8956    async fn single_asin() {
8957        do_single_instant_function_call("asin", "asin").await;
8958    }
8959
8960    #[tokio::test]
8961    #[should_panic]
8962    async fn single_asinh() {
8963        do_single_instant_function_call("asinh", "").await;
8964    }
8965
8966    #[tokio::test]
8967    async fn single_atan() {
8968        do_single_instant_function_call("atan", "atan").await;
8969    }
8970
8971    #[tokio::test]
8972    #[should_panic]
8973    async fn single_atanh() {
8974        do_single_instant_function_call("atanh", "").await;
8975    }
8976
8977    #[tokio::test]
8978    async fn single_cos() {
8979        do_single_instant_function_call("cos", "cos").await;
8980    }
8981
8982    #[tokio::test]
8983    #[should_panic]
8984    async fn single_cosh() {
8985        do_single_instant_function_call("cosh", "").await;
8986    }
8987
8988    #[tokio::test]
8989    async fn single_sin() {
8990        do_single_instant_function_call("sin", "sin").await;
8991    }
8992
8993    #[tokio::test]
8994    #[should_panic]
8995    async fn single_sinh() {
8996        do_single_instant_function_call("sinh", "").await;
8997    }
8998
8999    #[tokio::test]
9000    async fn single_tan() {
9001        do_single_instant_function_call("tan", "tan").await;
9002    }
9003
9004    #[tokio::test]
9005    #[should_panic]
9006    async fn single_tanh() {
9007        do_single_instant_function_call("tanh", "").await;
9008    }
9009
9010    #[tokio::test]
9011    #[should_panic]
9012    async fn single_deg() {
9013        do_single_instant_function_call("deg", "").await;
9014    }
9015
9016    #[tokio::test]
9017    #[should_panic]
9018    async fn single_rad() {
9019        do_single_instant_function_call("rad", "").await;
9020    }
9021
9022    // {
9023    //     input: "avg by (foo)(some_metric)",
9024    //     expected: &AggregateExpr{
9025    //         Op: AVG,
9026    //         Expr: &VectorSelector{
9027    //             Name: "some_metric",
9028    //             LabelMatchers: []*labels.Matcher{
9029    //                 MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
9030    //             },
9031    //             PosRange: PositionRange{
9032    //                 Start: 13,
9033    //                 End:   24,
9034    //             },
9035    //         },
9036    //         Grouping: []string{"foo"},
9037    //         PosRange: PositionRange{
9038    //             Start: 0,
9039    //             End:   25,
9040    //         },
9041    //     },
9042    // },
9043    async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) {
9044        let prom_expr = parser::parse(&format!(
9045            "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})",
9046        ))
9047        .unwrap();
9048        let mut eval_stmt = EvalStmt {
9049            expr: prom_expr,
9050            start: UNIX_EPOCH,
9051            end: UNIX_EPOCH
9052                .checked_add(Duration::from_secs(100_000))
9053                .unwrap(),
9054            interval: Duration::from_secs(5),
9055            lookback_delta: Duration::from_secs(1),
9056        };
9057
9058        // test group by
9059        let table_provider = build_test_table_provider(
9060            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9061            2,
9062            2,
9063        )
9064        .await;
9065        let plan =
9066            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9067                .await
9068                .unwrap();
9069        let expected_no_without = String::from(
9070            "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]\
9071            \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]\
9072            \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]\
9073            \n      PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9074            \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]\
9075            \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]\
9076            \n            TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
9077        ).replace("TEMPLATE", plan_name);
9078        assert_eq!(
9079            plan.display_indent_schema().to_string(),
9080            expected_no_without
9081        );
9082
9083        // test group without
9084        if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr {
9085            *modifier = Some(LabelModifier::Exclude(Labels {
9086                labels: vec![String::from("tag_1")].into_iter().collect(),
9087            }));
9088        }
9089        let table_provider = build_test_table_provider(
9090            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9091            2,
9092            2,
9093        )
9094        .await;
9095        let plan =
9096            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9097                .await
9098                .unwrap();
9099        let expected_without = String::from(
9100            "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]\
9101            \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]\
9102            \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]\
9103            \n      PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\
9104            \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]\
9105            \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]\
9106            \n            TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]"
9107        ).replace("TEMPLATE", plan_name);
9108        assert_eq!(plan.display_indent_schema().to_string(), expected_without);
9109    }
9110
9111    #[tokio::test]
9112    async fn aggregate_sum() {
9113        do_aggregate_expr_plan("sum", "sum").await;
9114    }
9115
9116    #[tokio::test]
9117    async fn tsid_is_used_for_series_divide_when_available() {
9118        let prom_expr = parser::parse("some_metric").unwrap();
9119        let eval_stmt = EvalStmt {
9120            expr: prom_expr,
9121            start: UNIX_EPOCH,
9122            end: UNIX_EPOCH
9123                .checked_add(Duration::from_secs(100_000))
9124                .unwrap(),
9125            interval: Duration::from_secs(5),
9126            lookback_delta: Duration::from_secs(1),
9127        };
9128
9129        let table_provider = build_test_table_provider_with_tsid(
9130            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9131            1,
9132            1,
9133        )
9134        .await;
9135        let plan =
9136            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9137                .await
9138                .unwrap();
9139
9140        let plan_str = plan.display_indent_schema().to_string();
9141        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
9142        assert!(plan_str.contains("__tsid ASC NULLS FIRST"));
9143        assert!(
9144            !plan
9145                .schema()
9146                .fields()
9147                .iter()
9148                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
9149        );
9150
9151        let manipulate = find_instant_manipulate(&plan).unwrap();
9152        let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new(
9153            MemorySourceConfig::try_new(
9154                &[],
9155                Arc::new(
9156                    datafusion_expr::UserDefinedLogicalNodeCore::inputs(manipulate)[0]
9157                        .schema()
9158                        .as_arrow()
9159                        .clone(),
9160                ),
9161                None,
9162            )
9163            .unwrap(),
9164        ))));
9165        assert!(format!("{exec:?}").contains("reuse_tsid_column: true"));
9166    }
9167
9168    #[tokio::test]
9169    async fn default_binary_join_uses_tsid_when_available() {
9170        let eval_stmt = build_eval_stmt("some_metric / some_alt_metric");
9171
9172        let table_provider = build_test_table_provider_with_tsid(
9173            &[
9174                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9175                (
9176                    DEFAULT_SCHEMA_NAME.to_string(),
9177                    "some_alt_metric".to_string(),
9178                ),
9179            ],
9180            1,
9181            1,
9182        )
9183        .await;
9184        let plan =
9185            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9186                .await
9187                .unwrap();
9188
9189        let plan_str = plan.display_indent_schema().to_string();
9190        assert!(
9191            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9192            "{plan_str}"
9193        );
9194        assert!(
9195            !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9196            "{plan_str}"
9197        );
9198    }
9199
9200    #[tokio::test]
9201    async fn reject_binary_fill_modifiers() {
9202        let state = build_query_engine_state();
9203
9204        for query in [
9205            "some_metric + fill(0) some_alt_metric",
9206            "some_metric + fill_left(0) some_alt_metric",
9207            "some_metric + fill_right(0) some_alt_metric",
9208            "(some_metric + fill(0) some_alt_metric) + some_metric",
9209        ] {
9210            let eval_stmt = build_eval_stmt(query);
9211            let table_provider = build_test_table_provider(&[], 0, 0).await;
9212            let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state)
9213                .await
9214                .unwrap_err();
9215
9216            assert!(
9217                matches!(
9218                    &err,
9219                    crate::promql::error::Error::UnsupportedExpr { name, .. }
9220                        if name == "PromQL fill modifiers"
9221                ),
9222                "{err}"
9223            );
9224        }
9225    }
9226
9227    #[tokio::test]
9228    async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() {
9229        for query in [
9230            "timestamp(some_metric) / some_metric",
9231            "some_metric / timestamp(some_metric)",
9232        ] {
9233            let eval_stmt = build_eval_stmt(query);
9234
9235            let table_provider = build_test_table_provider_with_tsid(
9236                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
9237                1,
9238                1,
9239            )
9240            .await;
9241            let plan =
9242                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9243                    .await
9244                    .unwrap();
9245
9246            let plan_str = plan.display_indent_schema().to_string();
9247            assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}");
9248            assert!(
9249                plan_str.contains("lhs.tag_0 = rhs.tag_0"),
9250                "{query}: {plan_str}"
9251            );
9252            assert!(
9253                !plan
9254                    .schema()
9255                    .fields()
9256                    .iter()
9257                    .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
9258                "{query}: {plan_str}"
9259            );
9260        }
9261    }
9262
9263    #[tokio::test]
9264    async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() {
9265        let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job");
9266
9267        let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
9268            (
9269                (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()),
9270                2,
9271                1,
9272            ),
9273            (
9274                (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()),
9275                1,
9276                1,
9277            ),
9278        ])
9279        .await;
9280        let plan =
9281            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9282                .await
9283                .unwrap();
9284        let plan_str = plan.display_indent_schema().to_string();
9285
9286        assert!(
9287            plan_str.contains("Boolean(false)") || plan_str.contains("false"),
9288            "{plan_str}"
9289        );
9290    }
9291
9292    #[tokio::test]
9293    async fn tsid_is_preserved_for_nested_default_binary_joins() {
9294        let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric");
9295
9296        let table_provider = build_test_table_provider_with_tsid(
9297            &[
9298                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9299                (
9300                    DEFAULT_SCHEMA_NAME.to_string(),
9301                    "some_alt_metric".to_string(),
9302                ),
9303                (
9304                    DEFAULT_SCHEMA_NAME.to_string(),
9305                    "some_third_metric".to_string(),
9306                ),
9307            ],
9308            1,
9309            1,
9310        )
9311        .await;
9312        let plan =
9313            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9314                .await
9315                .unwrap();
9316
9317        let plan_str = plan.display_indent_schema().to_string();
9318        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9319        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9320    }
9321
9322    #[tokio::test]
9323    async fn repeated_tsid_binary_operand_reuses_leaf_plan() {
9324        let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100");
9325
9326        let table_provider = build_test_table_provider_with_tsid(
9327            &[
9328                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9329                (
9330                    DEFAULT_SCHEMA_NAME.to_string(),
9331                    "some_alt_metric".to_string(),
9332                ),
9333            ],
9334            1,
9335            1,
9336        )
9337        .await;
9338        let plan =
9339            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9340                .await
9341                .unwrap();
9342
9343        let plan_str = plan.display_indent_schema().to_string();
9344        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9345        assert_eq!(
9346            plan_str
9347                .matches("Filter: phy.__table_id = UInt32(1024)")
9348                .count(),
9349            1,
9350            "{plan_str}"
9351        );
9352        assert_eq!(
9353            plan_str.matches("PromInstantManipulate").count(),
9354            2,
9355            "{plan_str}"
9356        );
9357        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9358    }
9359
9360    #[tokio::test]
9361    async fn repeated_tsid_binary_operand_reuses_shorter_field_side() {
9362        let eval_stmt =
9363            build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100");
9364
9365        let table_provider = build_test_table_provider_with_tsid_fields(
9366            &[
9367                (
9368                    (
9369                        DEFAULT_SCHEMA_NAME.to_string(),
9370                        "two_field_metric".to_string(),
9371                    ),
9372                    2,
9373                ),
9374                (
9375                    (
9376                        DEFAULT_SCHEMA_NAME.to_string(),
9377                        "one_field_metric".to_string(),
9378                    ),
9379                    1,
9380                ),
9381            ],
9382            1,
9383        )
9384        .await;
9385        let plan =
9386            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9387                .await
9388                .unwrap();
9389
9390        let field_names = plan
9391            .schema()
9392            .fields()
9393            .iter()
9394            .map(|field| field.name().clone())
9395            .collect::<Vec<_>>();
9396        let value_columns = field_names
9397            .iter()
9398            .filter(|name| {
9399                *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
9400            })
9401            .count();
9402        assert_eq!(value_columns, 1, "{field_names:?}");
9403        let plan_str = plan.display_indent_schema().to_string();
9404        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9405        assert_eq!(
9406            plan_str
9407                .matches("Filter: phy.__table_id = UInt32(1025)")
9408                .count(),
9409            1,
9410            "{plan_str}"
9411        );
9412        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9413    }
9414
9415    #[tokio::test]
9416    async fn binary_island_reuses_self_operand_without_join() {
9417        let eval_stmt = build_eval_stmt("some_metric / some_metric");
9418
9419        let table_provider = build_test_table_provider_with_tsid(
9420            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
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(), 0, "{plan_str}");
9432        assert_eq!(
9433            plan_str
9434                .matches("Filter: phy.__table_id = UInt32(1024)")
9435                .count(),
9436            1,
9437            "{plan_str}"
9438        );
9439        assert_eq!(
9440            plan_str.matches("PromInstantManipulate").count(),
9441            1,
9442            "{plan_str}"
9443        );
9444    }
9445
9446    #[tokio::test]
9447    async fn binary_island_reuses_leaf_across_two_branches() {
9448        let eval_stmt =
9449            build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)");
9450
9451        let table_provider = build_test_table_provider_with_tsid(
9452            &[
9453                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9454                (
9455                    DEFAULT_SCHEMA_NAME.to_string(),
9456                    "some_alt_metric".to_string(),
9457                ),
9458                (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()),
9459            ],
9460            1,
9461            1,
9462        )
9463        .await;
9464        let plan =
9465            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9466                .await
9467                .unwrap();
9468
9469        let plan_str = plan.display_indent_schema().to_string();
9470        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9471        assert_eq!(
9472            plan_str
9473                .matches("Filter: phy.__table_id = UInt32(1024)")
9474                .count(),
9475            1,
9476            "{plan_str}"
9477        );
9478        assert_eq!(
9479            plan_str.matches("PromInstantManipulate").count(),
9480            3,
9481            "{plan_str}"
9482        );
9483    }
9484
9485    #[tokio::test]
9486    async fn binary_island_generated_alias_avoids_user_column_names() {
9487        let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric");
9488
9489        let table_provider = build_test_table_provider_with_fields(
9490            &[
9491                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9492                (
9493                    DEFAULT_SCHEMA_NAME.to_string(),
9494                    "some_alt_metric".to_string(),
9495                ),
9496            ],
9497            &["prom_v0", "__prom_v0"],
9498        )
9499        .await;
9500        let plan =
9501            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9502                .await
9503                .unwrap();
9504
9505        let field_names = plan.schema().field_names();
9506        assert!(field_names.iter().any(|name| name.ends_with(".prom_v0")));
9507        assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0")));
9508
9509        let plan_str = plan.display_indent_schema().to_string();
9510        assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}");
9511        assert_eq!(
9512            plan_str.matches("PromInstantManipulate").count(),
9513            2,
9514            "{plan_str}"
9515        );
9516    }
9517
9518    #[tokio::test]
9519    async fn binary_island_clears_qualifier_for_nested_unary_projection() {
9520        let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)");
9521
9522        let table_provider = build_test_table_provider_with_tsid(
9523            &[
9524                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9525                (
9526                    DEFAULT_SCHEMA_NAME.to_string(),
9527                    "some_alt_metric".to_string(),
9528                ),
9529            ],
9530            1,
9531            1,
9532        )
9533        .await;
9534        let plan =
9535            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9536                .await
9537                .unwrap();
9538
9539        let plan_str = plan.display_indent_schema().to_string();
9540        assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}");
9541        assert_eq!(
9542            plan_str.matches("PromInstantManipulate").count(),
9543            2,
9544            "{plan_str}"
9545        );
9546    }
9547
9548    #[tokio::test]
9549    async fn binary_island_keeps_distinct_matcher_leaves() {
9550        let eval_stmt = build_eval_stmt(
9551            "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}",
9552        );
9553
9554        let table_provider = build_test_table_provider_with_tsid(
9555            &[
9556                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9557                (
9558                    DEFAULT_SCHEMA_NAME.to_string(),
9559                    "some_alt_metric".to_string(),
9560                ),
9561            ],
9562            1,
9563            1,
9564        )
9565        .await;
9566        let plan =
9567            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9568                .await
9569                .unwrap();
9570
9571        let plan_str = plan.display_indent_schema().to_string();
9572        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9573        assert_eq!(
9574            plan_str.matches("PromInstantManipulate").count(),
9575            3,
9576            "{plan_str}"
9577        );
9578    }
9579
9580    #[tokio::test]
9581    async fn binary_island_keeps_offset_leaves_distinct() {
9582        let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric");
9583
9584        let table_provider = build_test_table_provider_with_tsid(
9585            &[
9586                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9587                (
9588                    DEFAULT_SCHEMA_NAME.to_string(),
9589                    "some_alt_metric".to_string(),
9590                ),
9591            ],
9592            1,
9593            1,
9594        )
9595        .await;
9596        let plan =
9597            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9598                .await
9599                .unwrap();
9600
9601        let plan_str = plan.display_indent_schema().to_string();
9602        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9603        assert_eq!(
9604            plan_str.matches("PromInstantManipulate").count(),
9605            3,
9606            "{plan_str}"
9607        );
9608    }
9609
9610    #[tokio::test]
9611    async fn binary_island_falls_back_for_group_modifier() {
9612        let eval_stmt = build_eval_stmt(
9613            "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric",
9614        );
9615
9616        let table_provider = build_test_table_provider_with_tsid(
9617            &[
9618                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9619                (
9620                    DEFAULT_SCHEMA_NAME.to_string(),
9621                    "some_alt_metric".to_string(),
9622                ),
9623            ],
9624            1,
9625            1,
9626        )
9627        .await;
9628        let plan =
9629            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9630                .await
9631                .unwrap();
9632
9633        let plan_str = plan.display_indent_schema().to_string();
9634        assert_eq!(
9635            plan_str.matches("PromInstantManipulate").count(),
9636            3,
9637            "{plan_str}"
9638        );
9639    }
9640
9641    #[tokio::test]
9642    async fn binary_island_falls_back_for_comparison_filter() {
9643        let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric");
9644
9645        let table_provider = build_test_table_provider_with_tsid(
9646            &[
9647                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9648                (
9649                    DEFAULT_SCHEMA_NAME.to_string(),
9650                    "some_alt_metric".to_string(),
9651                ),
9652            ],
9653            1,
9654            1,
9655        )
9656        .await;
9657        let plan =
9658            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9659                .await
9660                .unwrap();
9661
9662        let plan_str = plan.display_indent_schema().to_string();
9663        assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}");
9664        assert_eq!(
9665            plan_str.matches("PromInstantManipulate").count(),
9666            3,
9667            "{plan_str}"
9668        );
9669    }
9670
9671    #[tokio::test]
9672    async fn tsid_binary_join_uses_shorter_field_side() {
9673        let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric");
9674
9675        let table_provider = build_test_table_provider_with_tsid_fields(
9676            &[
9677                (
9678                    (
9679                        DEFAULT_SCHEMA_NAME.to_string(),
9680                        "one_field_metric".to_string(),
9681                    ),
9682                    1,
9683                ),
9684                (
9685                    (
9686                        DEFAULT_SCHEMA_NAME.to_string(),
9687                        "two_field_metric".to_string(),
9688                    ),
9689                    2,
9690                ),
9691            ],
9692            1,
9693        )
9694        .await;
9695        let plan =
9696            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9697                .await
9698                .unwrap();
9699
9700        let field_names = plan
9701            .schema()
9702            .fields()
9703            .iter()
9704            .map(|field| field.name().clone())
9705            .collect::<Vec<_>>();
9706        let value_columns = field_names
9707            .iter()
9708            .filter(|name| {
9709                *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME
9710            })
9711            .count();
9712        assert_eq!(value_columns, 1, "{field_names:?}");
9713    }
9714
9715    #[tokio::test]
9716    async fn comparison_binary_join_uses_shorter_field_side() {
9717        let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric");
9718
9719        let table_provider = build_test_table_provider_with_tsid_fields(
9720            &[
9721                (
9722                    (
9723                        DEFAULT_SCHEMA_NAME.to_string(),
9724                        "two_field_metric".to_string(),
9725                    ),
9726                    2,
9727                ),
9728                (
9729                    (
9730                        DEFAULT_SCHEMA_NAME.to_string(),
9731                        "one_field_metric".to_string(),
9732                    ),
9733                    1,
9734                ),
9735            ],
9736            1,
9737        )
9738        .await;
9739        let plan =
9740            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9741                .await
9742                .unwrap();
9743
9744        let field_names = plan
9745            .schema()
9746            .fields()
9747            .iter()
9748            .map(|field| field.name().clone())
9749            .collect::<Vec<_>>();
9750        assert!(
9751            field_names.iter().any(|name| name == "field_0"),
9752            "{field_names:?}"
9753        );
9754        assert!(
9755            !field_names.iter().any(|name| name == "field_1"),
9756            "{field_names:?}"
9757        );
9758    }
9759
9760    #[tokio::test]
9761    async fn label_matching_modifier_disables_tsid_binary_join() {
9762        let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric");
9763
9764        let table_provider = build_test_table_provider_with_tsid(
9765            &[
9766                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9767                (
9768                    DEFAULT_SCHEMA_NAME.to_string(),
9769                    "some_alt_metric".to_string(),
9770                ),
9771            ],
9772            2,
9773            1,
9774        )
9775        .await;
9776        let plan =
9777            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9778                .await
9779                .unwrap();
9780
9781        let plan_str = plan.display_indent_schema().to_string();
9782        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9783        assert!(
9784            plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"),
9785            "{plan_str}"
9786        );
9787    }
9788
9789    #[tokio::test]
9790    async fn ignoring_absent_label_keeps_tsid_binary_join() {
9791        let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric");
9792
9793        let table_provider = build_test_table_provider_with_tsid(
9794            &[
9795                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9796                (
9797                    DEFAULT_SCHEMA_NAME.to_string(),
9798                    "some_alt_metric".to_string(),
9799                ),
9800            ],
9801            2,
9802            1,
9803        )
9804        .await;
9805        let plan =
9806            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9807                .await
9808                .unwrap();
9809
9810        let plan_str = plan.display_indent_schema().to_string();
9811        assert!(
9812            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9813            "{plan_str}"
9814        );
9815        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9816        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9817    }
9818
9819    #[tokio::test]
9820    async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() {
9821        let eval_stmt =
9822            build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric");
9823
9824        let table_provider = build_test_table_provider_with_tsid(
9825            &[
9826                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9827                (
9828                    DEFAULT_SCHEMA_NAME.to_string(),
9829                    "some_alt_metric".to_string(),
9830                ),
9831            ],
9832            2,
9833            1,
9834        )
9835        .await;
9836        let plan =
9837            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9838                .await
9839                .unwrap();
9840
9841        let plan_str = plan.display_indent_schema().to_string();
9842        assert!(
9843            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9844            "{plan_str}"
9845        );
9846        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9847        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9848    }
9849
9850    #[tokio::test]
9851    async fn on_full_label_set_keeps_tsid_binary_join() {
9852        let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric");
9853
9854        let table_provider = build_test_table_provider_with_tsid(
9855            &[
9856                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9857                (
9858                    DEFAULT_SCHEMA_NAME.to_string(),
9859                    "some_alt_metric".to_string(),
9860                ),
9861            ],
9862            2,
9863            1,
9864        )
9865        .await;
9866        let plan =
9867            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9868                .await
9869                .unwrap();
9870
9871        let plan_str = plan.display_indent_schema().to_string();
9872        assert!(
9873            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9874            "{plan_str}"
9875        );
9876        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
9877        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9878    }
9879
9880    #[tokio::test]
9881    async fn on_partial_label_set_disables_tsid_binary_join() {
9882        let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
9883
9884        let table_provider = build_test_table_provider_with_tsid(
9885            &[
9886                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9887                (
9888                    DEFAULT_SCHEMA_NAME.to_string(),
9889                    "some_alt_metric".to_string(),
9890                ),
9891            ],
9892            2,
9893            1,
9894        )
9895        .await;
9896        let plan =
9897            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9898                .await
9899                .unwrap();
9900
9901        let plan_str = plan.display_indent_schema().to_string();
9902        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9903        assert!(
9904            plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9905            "{plan_str}"
9906        );
9907        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9908    }
9909
9910    #[tokio::test]
9911    async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() {
9912        let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
9913
9914        let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
9915            (
9916                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9917                2,
9918                1,
9919            ),
9920            (
9921                (
9922                    DEFAULT_SCHEMA_NAME.to_string(),
9923                    "some_alt_metric".to_string(),
9924                ),
9925                1,
9926                1,
9927            ),
9928        ])
9929        .await;
9930        let plan =
9931            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
9932                .await
9933                .unwrap();
9934
9935        let plan_str = plan.display_indent_schema().to_string();
9936        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
9937        assert!(
9938            plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
9939            "{plan_str}"
9940        );
9941        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
9942    }
9943
9944    #[tokio::test]
9945    async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() {
9946        let eval_stmt = build_eval_stmt("some_metric > some_alt_metric");
9947
9948        let table_provider = build_test_table_provider_with_tsid(
9949            &[
9950                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9951                (
9952                    DEFAULT_SCHEMA_NAME.to_string(),
9953                    "some_alt_metric".to_string(),
9954                ),
9955            ],
9956            2,
9957            1,
9958        )
9959        .await;
9960        let mut planner = PromPlanner {
9961            table_provider,
9962            ctx: PromPlannerContext::from_eval_stmt(&eval_stmt),
9963            promql_annotations: None,
9964        };
9965        let plan = planner
9966            .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state())
9967            .await
9968            .unwrap();
9969
9970        let plan_str = plan.display_indent_schema().to_string();
9971        assert!(
9972            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
9973            "{plan_str}"
9974        );
9975        assert!(
9976            plan.schema()
9977                .fields()
9978                .iter()
9979                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME),
9980            "{plan_str}"
9981        );
9982        assert!(planner.ctx.use_tsid, "{plan_str}");
9983    }
9984
9985    #[tokio::test]
9986    async fn comparison_bool_binary_join_uses_tsid_when_available() {
9987        let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric");
9988
9989        let table_provider = build_test_table_provider_with_tsid(
9990            &[
9991                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
9992                (
9993                    DEFAULT_SCHEMA_NAME.to_string(),
9994                    "some_alt_metric".to_string(),
9995                ),
9996            ],
9997            2,
9998            1,
9999        )
10000        .await;
10001        let plan =
10002            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10003                .await
10004                .unwrap();
10005
10006        let plan_str = plan.display_indent_schema().to_string();
10007        assert!(
10008            plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"),
10009            "{plan_str}"
10010        );
10011        assert!(!plan_str.contains("tag_0 ="), "{plan_str}");
10012        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
10013    }
10014
10015    #[tokio::test]
10016    async fn scalar_count_count_range_keeps_full_window() {
10017        let plan_str = build_optimized_tsid_plan(
10018            "scalar(count(count(some_metric) by (tag_0)))",
10019            1,
10020            1,
10021            100_000,
10022            1,
10023        )
10024        .await;
10025        assert!(plan_str.contains("ScalarCalculate: tags=[]"));
10026        assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]"));
10027        assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]"));
10028    }
10029
10030    #[tokio::test]
10031    async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() {
10032        let plan_str = build_optimized_tsid_plan(
10033            "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))",
10034            2,
10035            1,
10036            10,
10037            300,
10038        )
10039        .await;
10040        assert!(plan_str.contains("Distinct:"), "{plan_str}");
10041    }
10042
10043    #[tokio::test]
10044    async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() {
10045        assert_nested_count_rewrite_applies(
10046            "count(count(some_metric) by (tag_0))",
10047            "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]"
10048        )
10049        .await;
10050    }
10051
10052    #[tokio::test]
10053    async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() {
10054        assert_nested_count_rewrite_applies(
10055            "count(sum(some_metric) by (tag_0))",
10056            "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]"
10057        )
10058        .await;
10059    }
10060
10061    #[tokio::test]
10062    async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() {
10063        for (query, expected_outer_agg) in [
10064            (
10065                "count(avg(some_metric) by (tag_0))",
10066                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]",
10067            ),
10068            (
10069                "count(min(some_metric) by (tag_0))",
10070                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]",
10071            ),
10072            (
10073                "count(max(some_metric) by (tag_0))",
10074                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]",
10075            ),
10076            (
10077                "count(stddev(some_metric) by (tag_0))",
10078                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]",
10079            ),
10080            (
10081                "count(stdvar(some_metric) by (tag_0))",
10082                "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]",
10083            ),
10084        ] {
10085            assert_nested_count_rewrite_applies(query, expected_outer_agg).await;
10086        }
10087    }
10088
10089    #[tokio::test]
10090    async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() {
10091        let count_plan =
10092            build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1)
10093                .await;
10094        assert!(
10095            !count_plan.contains("some_metric.field_0 IS NOT NULL"),
10096            "{count_plan}"
10097        );
10098
10099        for query in [
10100            "count(sum(some_metric) by (tag_0))",
10101            "count(avg(some_metric) by (tag_0))",
10102            "count(min(some_metric) by (tag_0))",
10103            "count(max(some_metric) by (tag_0))",
10104            "count(stddev(some_metric) by (tag_0))",
10105            "count(stdvar(some_metric) by (tag_0))",
10106        ] {
10107            let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await;
10108            assert!(
10109                plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"),
10110                "{query}: {plan_str}"
10111            );
10112        }
10113    }
10114
10115    #[tokio::test]
10116    async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() {
10117        assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await;
10118        assert_nested_count_rewrite_missing(
10119            "count(sum(irate(some_metric[1h])) by (tag_0))",
10120            2,
10121            300,
10122        )
10123        .await;
10124    }
10125
10126    #[tokio::test]
10127    async fn physical_table_name_is_not_leaked_in_plan() {
10128        let prom_expr = parser::parse("some_metric").unwrap();
10129        let eval_stmt = EvalStmt {
10130            expr: prom_expr,
10131            start: UNIX_EPOCH,
10132            end: UNIX_EPOCH
10133                .checked_add(Duration::from_secs(100_000))
10134                .unwrap(),
10135            interval: Duration::from_secs(5),
10136            lookback_delta: Duration::from_secs(1),
10137        };
10138
10139        let table_provider = build_test_table_provider_with_tsid(
10140            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10141            1,
10142            1,
10143        )
10144        .await;
10145        let plan =
10146            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10147                .await
10148                .unwrap();
10149
10150        let plan_str = plan.display_indent_schema().to_string();
10151        assert!(plan_str.contains("TableScan: phy"), "{plan}");
10152        assert!(plan_str.contains("SubqueryAlias: some_metric"));
10153        assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)"));
10154        assert!(!plan_str.contains("TableScan: some_metric"));
10155    }
10156
10157    #[tokio::test]
10158    async fn sum_without_does_not_group_by_tsid() {
10159        let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap();
10160        let eval_stmt = EvalStmt {
10161            expr: prom_expr,
10162            start: UNIX_EPOCH,
10163            end: UNIX_EPOCH
10164                .checked_add(Duration::from_secs(100_000))
10165                .unwrap(),
10166            interval: Duration::from_secs(5),
10167            lookback_delta: Duration::from_secs(1),
10168        };
10169
10170        let table_provider = build_test_table_provider_with_tsid(
10171            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10172            1,
10173            1,
10174        )
10175        .await;
10176        let plan =
10177            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10178                .await
10179                .unwrap();
10180
10181        let plan_str = plan.display_indent_schema().to_string();
10182        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10183
10184        let aggr_line = plan_str
10185            .lines()
10186            .find(|line| line.contains("Aggregate: groupBy="))
10187            .unwrap();
10188        assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10189    }
10190
10191    #[tokio::test]
10192    async fn topk_without_does_not_partition_by_tsid() {
10193        let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap();
10194        let eval_stmt = EvalStmt {
10195            expr: prom_expr,
10196            start: UNIX_EPOCH,
10197            end: UNIX_EPOCH
10198                .checked_add(Duration::from_secs(100_000))
10199                .unwrap(),
10200            interval: Duration::from_secs(5),
10201            lookback_delta: Duration::from_secs(1),
10202        };
10203
10204        let table_provider = build_test_table_provider_with_tsid(
10205            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10206            1,
10207            1,
10208        )
10209        .await;
10210        let plan =
10211            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10212                .await
10213                .unwrap();
10214
10215        let plan_str = plan.display_indent_schema().to_string();
10216        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10217
10218        let window_line = plan_str
10219            .lines()
10220            .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10221            .unwrap();
10222        let partition_by = window_line
10223            .split("PARTITION BY [")
10224            .nth(1)
10225            .and_then(|s| s.split("] ORDER BY").next())
10226            .unwrap();
10227        assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10228    }
10229
10230    #[tokio::test]
10231    async fn sum_by_does_not_group_by_tsid() {
10232        let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap();
10233        let eval_stmt = EvalStmt {
10234            expr: prom_expr,
10235            start: UNIX_EPOCH,
10236            end: UNIX_EPOCH
10237                .checked_add(Duration::from_secs(100_000))
10238                .unwrap(),
10239            interval: Duration::from_secs(5),
10240            lookback_delta: Duration::from_secs(1),
10241        };
10242
10243        let table_provider = build_test_table_provider_with_tsid(
10244            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10245            1,
10246            1,
10247        )
10248        .await;
10249        let plan =
10250            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10251                .await
10252                .unwrap();
10253
10254        let plan_str = plan.display_indent_schema().to_string();
10255        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10256
10257        let aggr_line = plan_str
10258            .lines()
10259            .find(|line| line.contains("Aggregate: groupBy="))
10260            .unwrap();
10261        assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10262    }
10263
10264    #[tokio::test]
10265    async fn aggregate_over_binary_time_function_expr() {
10266        for op in ["sum", "min", "max", "avg"] {
10267            let prom_expr = parser::parse(&format!(
10268                "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)"
10269            ))
10270            .unwrap();
10271            let eval_stmt = EvalStmt {
10272                expr: prom_expr,
10273                start: UNIX_EPOCH,
10274                end: UNIX_EPOCH
10275                    .checked_add(Duration::from_secs(100_000))
10276                    .unwrap(),
10277                interval: Duration::from_secs(5),
10278                lookback_delta: Duration::from_secs(1),
10279            };
10280
10281            let table_provider = build_test_table_provider_with_tsid(
10282                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10283                3,
10284                1,
10285            )
10286            .await;
10287            let plan =
10288                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10289                    .await
10290                    .unwrap();
10291
10292            let plan_str = plan.display_indent_schema().to_string();
10293            let aggr_line = plan_str
10294                .lines()
10295                .find(|line| line.contains("Aggregate: groupBy="))
10296                .unwrap();
10297            assert!(aggr_line.contains(op), "{plan_str}");
10298            assert!(aggr_line.contains("first_value"), "{plan_str}");
10299            assert!(
10300                !plan
10301                    .schema()
10302                    .fields()
10303                    .iter()
10304                    .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME })
10305            );
10306        }
10307    }
10308
10309    #[tokio::test]
10310    async fn topk_by_does_not_partition_by_tsid() {
10311        let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap();
10312        let eval_stmt = EvalStmt {
10313            expr: prom_expr,
10314            start: UNIX_EPOCH,
10315            end: UNIX_EPOCH
10316                .checked_add(Duration::from_secs(100_000))
10317                .unwrap(),
10318            interval: Duration::from_secs(5),
10319            lookback_delta: Duration::from_secs(1),
10320        };
10321
10322        let table_provider = build_test_table_provider_with_tsid(
10323            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10324            1,
10325            1,
10326        )
10327        .await;
10328        let plan =
10329            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10330                .await
10331                .unwrap();
10332
10333        let plan_str = plan.display_indent_schema().to_string();
10334        assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10335
10336        let window_line = plan_str
10337            .lines()
10338            .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()"))
10339            .unwrap();
10340        let partition_by = window_line
10341            .split("PARTITION BY [")
10342            .nth(1)
10343            .and_then(|s| s.split("] ORDER BY").next())
10344            .unwrap();
10345        assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME));
10346    }
10347
10348    #[tokio::test]
10349    async fn selector_matcher_on_tsid_does_not_use_internal_column() {
10350        let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap();
10351        let eval_stmt = EvalStmt {
10352            expr: prom_expr,
10353            start: UNIX_EPOCH,
10354            end: UNIX_EPOCH
10355                .checked_add(Duration::from_secs(100_000))
10356                .unwrap(),
10357            interval: Duration::from_secs(5),
10358            lookback_delta: Duration::from_secs(1),
10359        };
10360
10361        let table_provider = build_test_table_provider_with_tsid(
10362            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10363            1,
10364            1,
10365        )
10366        .await;
10367        let plan =
10368            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10369                .await
10370                .unwrap();
10371
10372        fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet<Column>) {
10373            if let LogicalPlan::Filter(filter) = plan {
10374                datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap();
10375            }
10376            for input in plan.inputs() {
10377                collect_filter_cols(input, out);
10378            }
10379        }
10380
10381        let mut filter_cols = HashSet::new();
10382        collect_filter_cols(&plan, &mut filter_cols);
10383        assert!(
10384            !filter_cols
10385                .iter()
10386                .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME)
10387        );
10388    }
10389
10390    #[tokio::test]
10391    async fn tsid_is_not_used_when_physical_table_is_missing() {
10392        let prom_expr = parser::parse("some_metric").unwrap();
10393        let eval_stmt = EvalStmt {
10394            expr: prom_expr,
10395            start: UNIX_EPOCH,
10396            end: UNIX_EPOCH
10397                .checked_add(Duration::from_secs(100_000))
10398                .unwrap(),
10399            interval: Duration::from_secs(5),
10400            lookback_delta: Duration::from_secs(1),
10401        };
10402
10403        let catalog_list = MemoryCatalogManager::with_default_setup();
10404
10405        // Register a metric engine logical table referencing a missing physical table.
10406        let mut columns = vec![ColumnSchema::new(
10407            "tag_0".to_string(),
10408            ConcreteDataType::string_datatype(),
10409            false,
10410        )];
10411        columns.push(
10412            ColumnSchema::new(
10413                "timestamp".to_string(),
10414                ConcreteDataType::timestamp_millisecond_datatype(),
10415                false,
10416            )
10417            .with_time_index(true),
10418        );
10419        columns.push(ColumnSchema::new(
10420            "field_0".to_string(),
10421            ConcreteDataType::float64_datatype(),
10422            true,
10423        ));
10424        let schema = Arc::new(Schema::new(columns));
10425        let mut options = table::requests::TableOptions::default();
10426        options
10427            .extra_options
10428            .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string());
10429        let table_meta = TableMetaBuilder::empty()
10430            .schema(schema)
10431            .primary_key_indices(vec![0])
10432            .value_indices(vec![2])
10433            .engine(METRIC_ENGINE_NAME.to_string())
10434            .options(options)
10435            .next_column_id(1024)
10436            .build()
10437            .unwrap();
10438        let table_info = TableInfoBuilder::default()
10439            .table_id(1024)
10440            .name("some_metric")
10441            .meta(table_meta)
10442            .build()
10443            .unwrap();
10444        let table = EmptyTable::from_table_info(&table_info);
10445        catalog_list
10446            .register_table_sync(RegisterTableRequest {
10447                catalog: DEFAULT_CATALOG_NAME.to_string(),
10448                schema: DEFAULT_SCHEMA_NAME.to_string(),
10449                table_name: "some_metric".to_string(),
10450                table_id: 1024,
10451                table,
10452            })
10453            .unwrap();
10454
10455        let table_provider = DfTableSourceProvider::new(
10456            catalog_list,
10457            false,
10458            QueryContext::arc(),
10459            DummyDecoder::arc(),
10460            false,
10461        );
10462
10463        let plan =
10464            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10465                .await
10466                .unwrap();
10467
10468        let plan_str = plan.display_indent_schema().to_string();
10469        assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]"));
10470        assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]"));
10471    }
10472
10473    #[tokio::test]
10474    async fn tsid_is_carried_only_when_aggregate_preserves_label_set() {
10475        let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap();
10476        let eval_stmt = EvalStmt {
10477            expr: prom_expr,
10478            start: UNIX_EPOCH,
10479            end: UNIX_EPOCH
10480                .checked_add(Duration::from_secs(100_000))
10481                .unwrap(),
10482            interval: Duration::from_secs(5),
10483            lookback_delta: Duration::from_secs(1),
10484        };
10485
10486        let table_provider = build_test_table_provider_with_tsid(
10487            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10488            1,
10489            1,
10490        )
10491        .await;
10492        let plan =
10493            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10494                .await
10495                .unwrap();
10496
10497        let plan_str = plan.display_indent_schema().to_string();
10498        assert!(plan_str.contains("first_value") && plan_str.contains("__tsid"));
10499        assert!(
10500            !plan
10501                .schema()
10502                .fields()
10503                .iter()
10504                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
10505        );
10506
10507        // Merging aggregate: label set is reduced, tsid should not be carried.
10508        let prom_expr = parser::parse("sum(some_metric)").unwrap();
10509        let eval_stmt = EvalStmt {
10510            expr: prom_expr,
10511            start: UNIX_EPOCH,
10512            end: UNIX_EPOCH
10513                .checked_add(Duration::from_secs(100_000))
10514                .unwrap(),
10515            interval: Duration::from_secs(5),
10516            lookback_delta: Duration::from_secs(1),
10517        };
10518        let table_provider = build_test_table_provider_with_tsid(
10519            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10520            1,
10521            1,
10522        )
10523        .await;
10524        let plan =
10525            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10526                .await
10527                .unwrap();
10528        let plan_str = plan.display_indent_schema().to_string();
10529        assert!(!plan_str.contains("first_value"));
10530    }
10531
10532    #[tokio::test]
10533    async fn or_operator_with_unknown_metric_does_not_require_tsid() {
10534        let prom_expr = parser::parse("unknown_metric or some_metric").unwrap();
10535        let eval_stmt = EvalStmt {
10536            expr: prom_expr,
10537            start: UNIX_EPOCH,
10538            end: UNIX_EPOCH
10539                .checked_add(Duration::from_secs(100_000))
10540                .unwrap(),
10541            interval: Duration::from_secs(5),
10542            lookback_delta: Duration::from_secs(1),
10543        };
10544
10545        let table_provider = build_test_table_provider_with_tsid(
10546            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10547            1,
10548            1,
10549        )
10550        .await;
10551
10552        let plan =
10553            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10554                .await
10555                .unwrap();
10556
10557        assert!(
10558            !plan
10559                .schema()
10560                .fields()
10561                .iter()
10562                .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
10563        );
10564    }
10565
10566    #[tokio::test]
10567    async fn aggregate_avg() {
10568        do_aggregate_expr_plan("avg", "avg").await;
10569    }
10570
10571    #[tokio::test]
10572    #[should_panic] // output type doesn't match
10573    async fn aggregate_count() {
10574        do_aggregate_expr_plan("count", "count").await;
10575    }
10576
10577    #[tokio::test]
10578    async fn aggregate_min() {
10579        do_aggregate_expr_plan("min", "min").await;
10580    }
10581
10582    #[tokio::test]
10583    async fn aggregate_max() {
10584        do_aggregate_expr_plan("max", "max").await;
10585    }
10586
10587    #[tokio::test]
10588    async fn aggregate_group() {
10589        // Regression test for `group()` aggregator.
10590        // PromQL: sum(group by (cluster)(kubernetes_build_info{service="kubernetes",job="apiserver"}))
10591        // should be plannable, and `group()` should produce constant 1 for each group.
10592        let prom_expr = parser::parse(
10593            "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))",
10594        )
10595        .unwrap();
10596        let eval_stmt = EvalStmt {
10597            expr: prom_expr,
10598            start: UNIX_EPOCH,
10599            end: UNIX_EPOCH
10600                .checked_add(Duration::from_secs(100_000))
10601                .unwrap(),
10602            interval: Duration::from_secs(5),
10603            lookback_delta: Duration::from_secs(1),
10604        };
10605
10606        let table_provider = build_test_table_provider_with_fields(
10607            &[(
10608                DEFAULT_SCHEMA_NAME.to_string(),
10609                "kubernetes_build_info".to_string(),
10610            )],
10611            &["cluster", "service", "job"],
10612        )
10613        .await;
10614        let plan =
10615            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10616                .await
10617                .unwrap();
10618
10619        let plan_str = plan.display_indent_schema().to_string();
10620        assert!(plan_str.contains("max(Float64(1"));
10621    }
10622
10623    #[tokio::test]
10624    async fn aggregate_stddev() {
10625        do_aggregate_expr_plan("stddev", "stddev_pop").await;
10626    }
10627
10628    #[tokio::test]
10629    async fn aggregate_stdvar() {
10630        do_aggregate_expr_plan("stdvar", "var_pop").await;
10631    }
10632
10633    // TODO(ruihang): add range fn tests once exprs are ready.
10634
10635    // {
10636    //     input: "some_metric{tag_0="foo"} + some_metric{tag_0="bar"}",
10637    //     expected: &BinaryExpr{
10638    //         Op: ADD,
10639    //         LHS: &VectorSelector{
10640    //             Name: "a",
10641    //             LabelMatchers: []*labels.Matcher{
10642    //                     MustLabelMatcher(labels.MatchEqual, "tag_0", "foo"),
10643    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
10644    //             },
10645    //         },
10646    //         RHS: &VectorSelector{
10647    //             Name: "sum",
10648    //             LabelMatchers: []*labels.Matcher{
10649    //                     MustLabelMatcher(labels.MatchxEqual, "tag_0", "bar"),
10650    //                     MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"),
10651    //             },
10652    //         },
10653    //         VectorMatching: &VectorMatching{},
10654    //     },
10655    // },
10656    #[tokio::test]
10657    async fn binary_op_column_column() {
10658        let prom_expr =
10659            parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap();
10660        let eval_stmt = EvalStmt {
10661            expr: prom_expr,
10662            start: UNIX_EPOCH,
10663            end: UNIX_EPOCH
10664                .checked_add(Duration::from_secs(100_000))
10665                .unwrap(),
10666            interval: Duration::from_secs(5),
10667            lookback_delta: Duration::from_secs(1),
10668        };
10669
10670        let table_provider = build_test_table_provider(
10671            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
10672            1,
10673            1,
10674        )
10675        .await;
10676        let plan =
10677            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10678                .await
10679                .unwrap();
10680
10681        let expected = String::from(
10682            "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]\
10683            \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]\
10684            \n    SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10685            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10686            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10687            \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]\
10688            \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]\
10689            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10690            \n    SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10691            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10692            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10693            \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]\
10694            \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]\
10695            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10696        );
10697
10698        assert_eq!(plan.display_indent_schema().to_string(), expected);
10699    }
10700
10701    async fn indie_query_plan_compare<T: AsRef<str>>(query: &str, expected: T) {
10702        let prom_expr = parser::parse(query).unwrap();
10703        let eval_stmt = EvalStmt {
10704            expr: prom_expr,
10705            start: UNIX_EPOCH,
10706            end: UNIX_EPOCH
10707                .checked_add(Duration::from_secs(100_000))
10708                .unwrap(),
10709            interval: Duration::from_secs(5),
10710            lookback_delta: Duration::from_secs(1),
10711        };
10712
10713        let table_provider = build_test_table_provider(
10714            &[
10715                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
10716                (
10717                    "greptime_private".to_string(),
10718                    "some_alt_metric".to_string(),
10719                ),
10720            ],
10721            1,
10722            1,
10723        )
10724        .await;
10725        let plan =
10726            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
10727                .await
10728                .unwrap();
10729
10730        assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref());
10731    }
10732
10733    #[tokio::test]
10734    async fn binary_op_literal_column() {
10735        let query = r#"1 + some_metric{tag_0="bar"}"#;
10736        let expected = String::from(
10737            "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]\
10738            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10739            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10740            \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]\
10741            \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]\
10742            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10743        );
10744
10745        indie_query_plan_compare(query, expected).await;
10746    }
10747
10748    #[tokio::test]
10749    async fn binary_op_literal_literal() {
10750        let query = r#"1 + 1"#;
10751        let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N]
10752  TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#;
10753        indie_query_plan_compare(query, expected).await;
10754    }
10755
10756    #[tokio::test]
10757    async fn simple_bool_grammar() {
10758        let query = "some_metric != bool 1.2345";
10759        let expected = String::from(
10760            "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]\
10761            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10762            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10763            \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]\
10764            \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]\
10765            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10766        );
10767
10768        indie_query_plan_compare(query, expected).await;
10769    }
10770
10771    #[tokio::test]
10772    async fn bool_with_additional_arithmetic() {
10773        let query = "some_metric + (1 == bool 2)";
10774        let expected = String::from(
10775            "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]\
10776            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10777            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10778            \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]\
10779            \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]\
10780            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10781        );
10782
10783        indie_query_plan_compare(query, expected).await;
10784    }
10785
10786    #[tokio::test]
10787    async fn simple_unary() {
10788        let query = "-some_metric";
10789        let expected = String::from(
10790            "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]\
10791            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10792            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10793            \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]\
10794            \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]\
10795            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10796        );
10797
10798        indie_query_plan_compare(query, expected).await;
10799    }
10800
10801    #[tokio::test]
10802    async fn increase_aggr() {
10803        let query = "increase(some_metric[5m])";
10804        let expected = String::from(
10805            "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]\
10806            \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]\
10807            \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))]\
10808            \n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10809            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
10810            \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]\
10811            \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]\
10812            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
10813        );
10814
10815        indie_query_plan_compare(query, expected).await;
10816    }
10817
10818    async fn native_histogram_plan(query: &str) -> String {
10819        let table_provider = build_test_native_histogram_table_provider("some_metric").await;
10820        let plan = PromPlanner::stmt_to_plan(
10821            table_provider,
10822            &build_eval_stmt(query),
10823            &build_query_engine_state(),
10824        )
10825        .await
10826        .unwrap();
10827        plan.display_indent_schema().to_string()
10828    }
10829
10830    #[tokio::test]
10831    async fn native_histogram_count_uses_native_udf() {
10832        let plan = native_histogram_plan("histogram_count(some_metric)").await;
10833
10834        assert!(plan.contains("prom_native_histogram_count"), "{plan}");
10835        assert!(!plan.contains("HistogramFold:"), "{plan}");
10836    }
10837
10838    #[tokio::test]
10839    async fn timestamp_filters_native_histogram_stale_marker_before_projection() {
10840        let mut stale = direct_or_histogram();
10841        stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
10842        let table = operator_metric_table(
10843            "stale_histogram",
10844            2_100,
10845            "a",
10846            None,
10847            DirectOrValue::NativeHistogram(stale),
10848        );
10849        let catalog = MemoryCatalogManager::with_default_setup();
10850        catalog
10851            .register_table_sync(RegisterTableRequest {
10852                catalog: DEFAULT_CATALOG_NAME.to_string(),
10853                schema: DEFAULT_SCHEMA_NAME.to_string(),
10854                table_name: "stale_histogram".to_string(),
10855                table_id: 2_100,
10856                table,
10857            })
10858            .unwrap();
10859        let provider = DfTableSourceProvider::new(
10860            catalog,
10861            false,
10862            QueryContext::arc(),
10863            DummyDecoder::arc(),
10864            false,
10865        );
10866        let state = build_query_engine_state();
10867        let plan = PromPlanner::stmt_to_plan(
10868            provider,
10869            &operator_eval_stmt("timestamp(stale_histogram)"),
10870            &state,
10871        )
10872        .await
10873        .unwrap();
10874        let plan_text = plan.display_indent_schema().to_string();
10875        assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}");
10876
10877        let (_, batches) = execute(plan, &state).await;
10878        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
10879    }
10880
10881    #[tokio::test]
10882    async fn timestamp_filters_stale_marker_from_mixed_sample_companion() {
10883        let histograms = build_histogram_array(&[None]);
10884        let schema = Arc::new(ArrowSchema::new(vec![
10885            Field::new(
10886                "timestamp",
10887                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
10888                false,
10889            ),
10890            Field::new(
10891                greptime_native_histogram(),
10892                histograms.data_type().clone(),
10893                true,
10894            ),
10895            Field::new(greptime_value(), ArrowDataType::Float64, true),
10896        ]));
10897        let batch = RecordBatch::try_new(
10898            schema.clone(),
10899            vec![
10900                Arc::new(TimestampMillisecondArray::from(vec![1_000])),
10901                histograms,
10902                Arc::new(Float64Array::from(vec![f64::from_bits(
10903                    PROMETHEUS_STALE_NAN_BITS,
10904                )])),
10905            ],
10906        )
10907        .unwrap();
10908        let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
10909        let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
10910            .unwrap()
10911            .build()
10912            .unwrap();
10913        let input = LogicalPlan::Extension(Extension {
10914            node: Arc::new(SeriesDivide::new(
10915                Vec::new(),
10916                "timestamp".to_string(),
10917                input,
10918            )),
10919        });
10920        let input = LogicalPlan::Extension(Extension {
10921            node: Arc::new(InstantManipulate::new(
10922                1_000,
10923                1_000,
10924                5_000,
10925                1_000,
10926                0,
10927                "timestamp".to_string(),
10928                Vec::new(),
10929                Some(greptime_native_histogram().to_string()),
10930                input,
10931            )),
10932        });
10933        // Match timestamp()'s parent projection, which otherwise prunes the companion lane.
10934        let plan = LogicalPlanBuilder::from(input)
10935            .project([col("timestamp")])
10936            .unwrap()
10937            .build()
10938            .unwrap();
10939
10940        let (_, batches) = execute(plan, &build_query_engine_state()).await;
10941        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
10942    }
10943
10944    #[tokio::test]
10945    async fn native_histogram_rate_can_feed_count() {
10946        let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await;
10947
10948        assert!(plan.contains("prom_native_histogram_rate"), "{plan}");
10949        assert!(plan.contains("prom_native_histogram_count"), "{plan}");
10950    }
10951
10952    #[tokio::test]
10953    async fn native_histogram_quantile_skips_classic_fold() {
10954        let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await;
10955
10956        assert!(plan.contains("prom_native_histogram_quantile"), "{plan}");
10957        assert!(!plan.contains("HistogramFold:"), "{plan}");
10958        assert!(plan.contains("some_metric.le"), "{plan}");
10959        // The phi literal is threaded into the native quantile UDF as its second argument.
10960        assert!(plan.contains("Float64(0.9)"), "{plan}");
10961        // The empty-values filter drops NULL quantile results so the output is empty
10962        // when all native histogram samples are dropped.
10963        assert!(plan.contains("IS NOT NULL"), "{plan}");
10964    }
10965
10966    #[tokio::test]
10967    async fn mixed_native_histogram_quantile_uses_histogram_field() {
10968        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
10969        let plan = PromPlanner::stmt_to_plan(
10970            table_provider,
10971            &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
10972            &build_query_engine_state(),
10973        )
10974        .await
10975        .unwrap()
10976        .display_indent_schema()
10977        .to_string();
10978
10979        assert!(
10980            plan.contains("prom_native_histogram_quantile(greptime_native_histogram"),
10981            "{plan}"
10982        );
10983        assert!(!plan.contains("EmptyRelation"), "{plan}");
10984    }
10985
10986    #[tokio::test]
10987    async fn mixed_histogram_helpers_execute_classic_and_native_samples() {
10988        let state = build_query_engine_state();
10989        for (query, expected) in [
10990            (
10991                "histogram_quantile(0.5, mixed_histogram)",
10992                vec![("classic", 1.0), ("native", 0.0)],
10993            ),
10994            (
10995                "histogram_fraction(-Inf, +Inf, mixed_histogram)",
10996                vec![("classic", 1.0), ("native", 1.0)],
10997            ),
10998        ] {
10999            let plan = PromPlanner::stmt_to_plan(
11000                classic_and_native_histogram_table_provider("native", None, direct_or_histogram()),
11001                &operator_eval_stmt(query),
11002                &state,
11003            )
11004            .await
11005            .unwrap();
11006            let plan_text = plan.display_indent_schema().to_string();
11007            assert!(plan_text.contains("HistogramFold:"), "{plan_text}");
11008            assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}");
11009            let value_field = plan
11010                .schema()
11011                .fields()
11012                .iter()
11013                .find(|field| field.data_type() == &ArrowDataType::Float64)
11014                .unwrap()
11015                .name()
11016                .clone();
11017
11018            let (_, batches) = execute(plan, &state).await;
11019            let mut actual = batches
11020                .iter()
11021                .flat_map(|batch| {
11022                    let tags = batch
11023                        .column_by_name("tag")
11024                        .unwrap()
11025                        .as_any()
11026                        .downcast_ref::<StringArray>()
11027                        .unwrap();
11028                    let values = batch
11029                        .column_by_name(&value_field)
11030                        .unwrap()
11031                        .as_any()
11032                        .downcast_ref::<Float64Array>()
11033                        .unwrap();
11034                    (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row)))
11035                })
11036                .collect::<Vec<_>>();
11037            actual.sort_by_key(|(tag, _)| *tag);
11038            assert_eq!(actual, expected, "{query}");
11039        }
11040    }
11041
11042    #[tokio::test]
11043    async fn mixed_histogram_helpers_report_annotations() {
11044        let state = build_query_engine_state();
11045        let mut native_histogram = direct_or_histogram();
11046        native_histogram.count = 2.0;
11047        native_histogram.sum = f64::NAN;
11048        for (native_tag, expected_rows, expected_warnings, expected_infos) in [
11049            (
11050                "classic",
11051                0,
11052                vec!["vector contains a mix of classic and native histograms"],
11053                vec![],
11054            ),
11055            (
11056                "native",
11057                2,
11058                vec![],
11059                vec!["input to histogram_quantile has NaN observations, result is skewed higher"],
11060            ),
11061        ] {
11062            let collector = PromqlAnnotationCollector::default();
11063            let plan = PromPlanner::stmt_to_plan_with_annotations(
11064                classic_and_native_histogram_table_provider(
11065                    native_tag,
11066                    None,
11067                    native_histogram.clone(),
11068                ),
11069                &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"),
11070                &state,
11071                Some(collector.clone()),
11072            )
11073            .await
11074            .unwrap();
11075
11076            let (_, batches) = execute(plan, &state).await;
11077            assert_eq!(
11078                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
11079                expected_rows
11080            );
11081            let mut warnings = vec![];
11082            let mut infos = vec![];
11083            collector.append_to(&mut warnings, &mut infos);
11084            assert_eq!(warnings, expected_warnings);
11085            assert_eq!(infos, expected_infos);
11086        }
11087    }
11088
11089    #[tokio::test]
11090    async fn mixed_histogram_helper_preserves_native_le_and_scans_once() {
11091        let state = build_query_engine_state();
11092        let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)");
11093        stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap();
11094        let plan = PromPlanner::stmt_to_plan(
11095            classic_and_native_histogram_table_provider(
11096                "classic",
11097                Some("native"),
11098                direct_or_histogram(),
11099            ),
11100            &stmt,
11101            &state,
11102        )
11103        .await
11104        .unwrap();
11105        let plan_text = plan.display_indent_schema().to_string();
11106        assert_eq!(
11107            plan_text.matches("TableScan: mixed_histogram").count(),
11108            1,
11109            "{plan_text}"
11110        );
11111
11112        let value_field = plan
11113            .schema()
11114            .fields()
11115            .iter()
11116            .find(|field| field.data_type() == &ArrowDataType::Float64)
11117            .unwrap()
11118            .name()
11119            .clone();
11120        let (_, batches) = execute(plan, &state).await;
11121        let mut actual = batches
11122            .iter()
11123            .flat_map(|batch| {
11124                let le = batch
11125                    .column_by_name(LE_COLUMN_NAME)
11126                    .unwrap()
11127                    .as_any()
11128                    .downcast_ref::<StringArray>()
11129                    .unwrap();
11130                let timestamps = batch
11131                    .column_by_name("timestamp")
11132                    .unwrap()
11133                    .as_any()
11134                    .downcast_ref::<TimestampMillisecondArray>()
11135                    .unwrap();
11136                let values = batch
11137                    .column_by_name(&value_field)
11138                    .unwrap()
11139                    .as_any()
11140                    .downcast_ref::<Float64Array>()
11141                    .unwrap();
11142                (0..batch.num_rows()).map(|row| {
11143                    (
11144                        timestamps.value(row),
11145                        (!le.is_null(row)).then(|| le.value(row).to_string()),
11146                        values.value(row),
11147                    )
11148                })
11149            })
11150            .collect::<Vec<_>>();
11151        actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1)));
11152        assert_eq!(
11153            actual,
11154            vec![
11155                (1_000, None, 1.0),
11156                (1_000, Some("native".to_string()), 0.0),
11157                (2_000, None, 1.0),
11158                (2_000, Some("native".to_string()), 0.0),
11159            ]
11160        );
11161    }
11162
11163    #[tokio::test]
11164    async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() {
11165        let state = build_query_engine_state();
11166        for native_le in [None, Some("native")] {
11167            for query in [
11168                "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))",
11169                "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))",
11170            ] {
11171                let plan = PromPlanner::stmt_to_plan(
11172                    classic_and_native_histogram_table_provider(
11173                        "native",
11174                        native_le,
11175                        direct_or_histogram(),
11176                    ),
11177                    &operator_eval_stmt(query),
11178                    &state,
11179                )
11180                .await
11181                .unwrap();
11182
11183                let (_, batches) = execute(plan, &state).await;
11184                assert_eq!(
11185                    batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
11186                    0,
11187                    "native_le={native_le:?}, query={query}"
11188                );
11189            }
11190        }
11191    }
11192
11193    #[tokio::test]
11194    async fn native_histogram_quantile_rejects_multi_field_input() {
11195        let table_provider = build_test_multi_histogram_table_provider("some_metric").await;
11196        let result = PromPlanner::stmt_to_plan(
11197            table_provider,
11198            &build_eval_stmt("histogram_quantile(0.9, some_metric)"),
11199            &build_query_engine_state(),
11200        )
11201        .await;
11202
11203        let err = result.expect_err("histogram_quantile on two native histogram fields must fail");
11204        assert!(
11205            err.to_string()
11206                .contains("Multi fields calculation is not supported in histogram_quantile"),
11207            "{err}"
11208        );
11209    }
11210
11211    #[tokio::test]
11212    async fn native_histogram_topk_uses_drop_udf() {
11213        let plan = native_histogram_plan("topk(1, some_metric)").await;
11214
11215        assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}");
11216        assert!(
11217            plan.contains("Filter: prom_native_histogram_drop_float")
11218                && plan.contains("IS NOT NULL"),
11219            "{plan}"
11220        );
11221    }
11222
11223    #[tokio::test]
11224    async fn mixed_or_topk_bottomk_ignore_native_histograms() {
11225        for op in ["topk", "bottomk"] {
11226            let collector = PromqlAnnotationCollector::default();
11227            let state = build_query_engine_state();
11228            let plan = PromPlanner::stmt_to_plan_with_annotations(
11229                operator_table_provider(),
11230                &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")),
11231                &state,
11232                Some(collector.clone()),
11233            )
11234            .await
11235            .unwrap();
11236            let float_field = plan
11237                .schema()
11238                .fields()
11239                .iter()
11240                .find(|field| field.data_type() == &ArrowDataType::Float64)
11241                .unwrap()
11242                .name()
11243                .clone();
11244            assert!(
11245                plan.schema()
11246                    .fields()
11247                    .iter()
11248                    .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()),
11249                "{plan:?}"
11250            );
11251
11252            let (_, batches) = execute(plan, &state).await;
11253            assert_eq!(values(&batches, &float_field), vec![2.0], "{op}");
11254            let mut warnings = vec![];
11255            let mut infos = vec![];
11256            collector.append_to(&mut warnings, &mut infos);
11257            assert!(warnings.is_empty());
11258            assert_eq!(
11259                infos,
11260                vec![format!(
11261                    "{op}: dropped native histogram samples because this aggregation is not supported for native histograms"
11262                )]
11263            );
11264        }
11265    }
11266
11267    #[tokio::test]
11268    async fn native_histogram_scalar_is_ignored_before_scalar_calculate() {
11269        let plan = native_histogram_plan("scalar(some_metric)").await;
11270
11271        assert!(plan.contains("ScalarCalculate"), "{plan}");
11272        assert!(plan.contains("Filter: Boolean(false)"), "{plan}");
11273        assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11274    }
11275
11276    #[tokio::test]
11277    async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() {
11278        for function in ["sort", "sort_desc"] {
11279            let plan = native_histogram_plan(&format!("{function}(some_metric)")).await;
11280
11281            assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11282            assert!(
11283                !plan.contains(&format!("Sort: {}", greptime_native_histogram())),
11284                "{plan}"
11285            );
11286            assert!(!plan.contains("prom_native_histogram_drop"), "{plan}");
11287        }
11288
11289        for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] {
11290            let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await;
11291
11292            assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}");
11293            assert!(plan.contains(greptime_native_histogram()), "{plan}");
11294            assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}");
11295        }
11296    }
11297
11298    #[tokio::test]
11299    async fn unsupported_native_histogram_functions_use_drop_udf() {
11300        for query in [
11301            "deriv(some_metric[5m])",
11302            "min_over_time(some_metric[5m])",
11303            "quantile_over_time(0.9, some_metric[5m])",
11304            "predict_linear(some_metric[5m], 60)",
11305            "round(some_metric)",
11306            "abs(some_metric)",
11307        ] {
11308            let plan = native_histogram_plan(query).await;
11309
11310            assert!(
11311                plan.contains("prom_native_histogram_drop_float"),
11312                "{query}\n{plan}"
11313            );
11314        }
11315    }
11316
11317    #[tokio::test]
11318    async fn native_histogram_absent_over_time_uses_native_udf() {
11319        let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await;
11320
11321        assert!(
11322            plan.contains("prom_native_histogram_absent_over_time"),
11323            "{plan}"
11324        );
11325    }
11326
11327    #[tokio::test]
11328    async fn native_histogram_all_function_arms_route_correctly() {
11329        // Every native-histogram match arm in `create_function_expr` must route to the
11330        // expected UDF when all field columns are native histograms. `holt_winters` shares
11331        // the `double_exponential_smoothing` arm but is not registered in the promql
11332        // parser (0.10), so it cannot be exercised through a query string.
11333        let cases = [
11334            // Range functions routed to native histogram UDFs.
11335            (
11336                "increase(some_metric[5m])",
11337                "prom_native_histogram_increase",
11338            ),
11339            ("rate(some_metric[5m])", "prom_native_histogram_rate"),
11340            ("delta(some_metric[5m])", "prom_native_histogram_delta"),
11341            ("idelta(some_metric[5m])", "prom_native_histogram_idelta"),
11342            ("irate(some_metric[5m])", "prom_native_histogram_irate"),
11343            ("resets(some_metric[5m])", "prom_native_histogram_resets"),
11344            ("changes(some_metric[5m])", "prom_native_histogram_changes"),
11345            (
11346                "avg_over_time(some_metric[5m])",
11347                "prom_native_histogram_avg_over_time",
11348            ),
11349            (
11350                "sum_over_time(some_metric[5m])",
11351                "prom_native_histogram_sum_over_time",
11352            ),
11353            (
11354                "count_over_time(some_metric[5m])",
11355                "prom_native_histogram_count_over_time",
11356            ),
11357            (
11358                "last_over_time(some_metric[5m])",
11359                "prom_native_histogram_last_over_time",
11360            ),
11361            (
11362                "present_over_time(some_metric[5m])",
11363                "prom_native_histogram_present_over_time",
11364            ),
11365            // Unsupported functions dropped with the float-null UDF.
11366            ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"),
11367            (
11368                "min_over_time(some_metric[5m])",
11369                "prom_native_histogram_drop_float",
11370            ),
11371            (
11372                "max_over_time(some_metric[5m])",
11373                "prom_native_histogram_drop_float",
11374            ),
11375            (
11376                "stddev_over_time(some_metric[5m])",
11377                "prom_native_histogram_drop_float",
11378            ),
11379            (
11380                "stdvar_over_time(some_metric[5m])",
11381                "prom_native_histogram_drop_float",
11382            ),
11383            (
11384                "quantile_over_time(0.9, some_metric[5m])",
11385                "prom_native_histogram_drop_float",
11386            ),
11387            (
11388                "predict_linear(some_metric[5m], 60)",
11389                "prom_native_histogram_drop_float",
11390            ),
11391            (
11392                "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11393                "prom_native_histogram_drop_float",
11394            ),
11395            ("round(some_metric)", "prom_native_histogram_drop_float"),
11396            ("rad(some_metric)", "prom_native_histogram_drop_float"),
11397            ("deg(some_metric)", "prom_native_histogram_drop_float"),
11398            ("sgn(some_metric)", "prom_native_histogram_drop_float"),
11399            // Instant helper functions routed to native histogram UDFs.
11400            (
11401                "histogram_count(some_metric)",
11402                "prom_native_histogram_count",
11403            ),
11404            ("histogram_sum(some_metric)", "prom_native_histogram_sum"),
11405            ("histogram_avg(some_metric)", "prom_native_histogram_avg"),
11406            (
11407                "histogram_stddev(some_metric)",
11408                "prom_native_histogram_stddev",
11409            ),
11410            (
11411                "histogram_stdvar(some_metric)",
11412                "prom_native_histogram_stdvar",
11413            ),
11414            (
11415                "histogram_fraction(-2 + 1, 2 / 2, some_metric)",
11416                "prom_native_histogram_fraction",
11417            ),
11418        ];
11419
11420        for (query, expected_udf) in cases {
11421            let plan = native_histogram_plan(query).await;
11422            assert!(plan.contains(expected_udf), "{query}\n{plan}");
11423            if query.starts_with("histogram_fraction") {
11424                assert!(plan.contains("Float64(-1)"), "{query}\n{plan}");
11425            }
11426        }
11427    }
11428
11429    #[tokio::test]
11430    async fn mixed_native_histogram_ranges_use_coordinated_udfs() {
11431        let dual_output = [
11432            "increase(some_metric[5m])",
11433            "rate(some_metric[5m])",
11434            "delta(some_metric[5m])",
11435            "idelta(some_metric[5m])",
11436            "irate(some_metric[5m])",
11437            "avg_over_time(some_metric[5m])",
11438            "sum_over_time(some_metric[5m])",
11439            "last_over_time(some_metric[5m])",
11440        ];
11441        let float_output = [
11442            "resets(some_metric[5m])",
11443            "changes(some_metric[5m])",
11444            "deriv(some_metric[5m])",
11445            "min_over_time(some_metric[5m])",
11446            "max_over_time(some_metric[5m])",
11447            "count_over_time(some_metric[5m])",
11448            "absent_over_time(some_metric[5m])",
11449            "present_over_time(some_metric[5m])",
11450            "stddev_over_time(some_metric[5m])",
11451            "stdvar_over_time(some_metric[5m])",
11452            "quantile_over_time(0.9, some_metric[5m])",
11453            "predict_linear(some_metric[5m], 60)",
11454            "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)",
11455        ];
11456
11457        for query in dual_output.iter().chain(float_output.iter()) {
11458            let plan = PromPlanner::stmt_to_plan(
11459                build_test_mixed_native_histogram_table_provider("some_metric").await,
11460                &build_eval_stmt(query),
11461                &build_query_engine_state(),
11462            )
11463            .await
11464            .unwrap()
11465            .display_indent_schema()
11466            .to_string();
11467            assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}");
11468            assert_eq!(
11469                plan.contains("prom_mixed_range_histogram"),
11470                dual_output.contains(query),
11471                "{query}\n{plan}"
11472            );
11473        }
11474
11475        let plan = PromPlanner::stmt_to_plan(
11476            build_test_mixed_native_histogram_table_provider("some_metric").await,
11477            &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"),
11478            &build_query_engine_state(),
11479        )
11480        .await
11481        .unwrap()
11482        .display_indent_schema()
11483        .to_string();
11484        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]
11485  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]
11486    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))]
11487      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]
11488        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]
11489          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]
11490            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]
11491              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))]
11492                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]
11493                  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]
11494                    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]
11495                      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]
11496                        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]"#;
11497        assert_eq!(plan, expected);
11498    }
11499
11500    #[tokio::test]
11501    async fn mixed_native_histogram_rate_executes_real_ranges() {
11502        let schema = Arc::new(ArrowSchema::new(vec![
11503            Field::new(
11504                "timestamp",
11505                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
11506                false,
11507            ),
11508            Field::new(greptime_value(), ArrowDataType::Float64, true),
11509            Field::new(
11510                greptime_native_histogram(),
11511                native_histogram_value_type().as_arrow_type(),
11512                true,
11513            ),
11514        ]));
11515        let batch = RecordBatch::try_new(
11516            schema.clone(),
11517            vec![
11518                Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])),
11519                Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])),
11520                build_histogram_array(&[None, Some(direct_or_histogram()), None]),
11521            ],
11522        )
11523        .unwrap();
11524        let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
11525        let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None)
11526            .unwrap()
11527            .build()
11528            .unwrap();
11529        let collector = PromqlAnnotationCollector::default();
11530        let mut planner = PromPlanner {
11531            table_provider: build_test_table_provider_with_fields(
11532                &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
11533                &[],
11534            )
11535            .await,
11536            ctx: PromPlannerContext {
11537                start: 3000,
11538                end: 3000,
11539                interval: 1000,
11540                range: Some(3000),
11541                time_index_column: Some("timestamp".to_string()),
11542                field_columns: vec![
11543                    greptime_native_histogram().to_string(),
11544                    greptime_value().to_string(),
11545                ],
11546                ..Default::default()
11547            },
11548            promql_annotations: Some(collector.clone()),
11549        };
11550        let input = LogicalPlan::Extension(Extension {
11551            node: Arc::new(
11552                RangeManipulate::new(
11553                    3000,
11554                    3000,
11555                    1000,
11556                    0,
11557                    3000,
11558                    "timestamp".to_string(),
11559                    planner.ctx.field_columns.clone(),
11560                    input,
11561                )
11562                .unwrap(),
11563            ),
11564        });
11565        let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else {
11566            unreachable!()
11567        };
11568        let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
11569            input.schema(),
11570            &planner.ctx.field_columns,
11571        );
11572        let state = build_query_engine_state();
11573        let (mut exprs, _) = planner
11574            .create_function_expr(&call.func, vec![], input.schema(), &state)
11575            .unwrap();
11576        exprs.insert(0, planner.create_time_index_column_expr().unwrap());
11577        let plan = LogicalPlanBuilder::from(input)
11578            .project(exprs)
11579            .unwrap()
11580            .filter(
11581                planner
11582                    .create_empty_values_filter_expr(preserve_any_value)
11583                    .unwrap(),
11584            )
11585            .unwrap()
11586            .build()
11587            .unwrap();
11588        let (_, batches) = execute(plan, &state).await;
11589        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
11590        let mut warnings = Vec::new();
11591        collector.append_to(&mut warnings, &mut Vec::new());
11592        assert!(
11593            warnings
11594                .iter()
11595                .any(|warning| warning.contains("mix of float and native histogram"))
11596        );
11597    }
11598
11599    #[tokio::test]
11600    async fn native_histogram_mixed_field_table_behaves() {
11601        // Exercise function planning after float and histogram samples have already been
11602        // represented as alternative nullable fields. Histogram functions must select the
11603        // histogram field without adding a NULL float field that would reject every row.
11604        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11605        let plan = PromPlanner::stmt_to_plan(
11606            table_provider,
11607            &build_eval_stmt("histogram_count(some_metric)"),
11608            &build_query_engine_state(),
11609        )
11610        .await
11611        .unwrap();
11612        let plan_str = plan.display_indent_schema().to_string();
11613        assert!(
11614            plan_str.contains("prom_native_histogram_count"),
11615            "{plan_str}"
11616        );
11617        assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}");
11618        assert!(
11619            plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"),
11620            "{plan_str}"
11621        );
11622
11623        // Value sorting keeps the float column and never sorts by the histogram column.
11624        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11625        let plan = PromPlanner::stmt_to_plan(
11626            table_provider,
11627            &build_eval_stmt("sort(some_metric)"),
11628            &build_query_engine_state(),
11629        )
11630        .await
11631        .unwrap();
11632        let plan_str = plan.display_indent_schema().to_string();
11633        assert!(
11634            plan_str.contains("greptime_value ASC NULLS FIRST"),
11635            "{plan_str}"
11636        );
11637        assert!(
11638            !plan_str.contains("greptime_native_histogram ASC"),
11639            "{plan_str}"
11640        );
11641
11642        // scalar() ignores histogram samples and evaluates only the float field.
11643        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11644        let plan = PromPlanner::stmt_to_plan(
11645            table_provider,
11646            &build_eval_stmt("scalar(some_metric)"),
11647            &build_query_engine_state(),
11648        )
11649        .await
11650        .unwrap();
11651        let plan_str = plan.display_indent_schema().to_string();
11652        assert!(plan_str.contains("ScalarCalculate"), "{plan_str}");
11653        assert!(
11654            plan_str.contains("greptime_value IS NOT NULL"),
11655            "{plan_str}"
11656        );
11657
11658        // Functions that preserve both alternative fields keep rows with either sample type.
11659        let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await;
11660        let plan = PromPlanner::stmt_to_plan(
11661            table_provider,
11662            &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#),
11663            &build_query_engine_state(),
11664        )
11665        .await
11666        .unwrap();
11667        let plan_str = plan.display_indent_schema().to_string();
11668        let filter = plan_str.lines().next().unwrap();
11669        assert!(
11670            filter.starts_with("Filter: ")
11671                && filter.contains("greptime_native_histogram IS NOT NULL")
11672                && filter.contains(" OR ")
11673                && filter.contains("greptime_value IS NOT NULL"),
11674            "{plan_str}"
11675        );
11676    }
11677
11678    #[tokio::test]
11679    async fn less_filter_on_value() {
11680        let query = "some_metric < 1.2345";
11681        let expected = String::from(
11682            "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11683            \n  PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11684            \n    PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11685            \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]\
11686            \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]\
11687            \n          TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11688        );
11689
11690        indie_query_plan_compare(query, expected).await;
11691    }
11692
11693    #[tokio::test]
11694    async fn count_over_time() {
11695        let query = "count_over_time(some_metric[5m])";
11696        let expected = String::from(
11697            "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]\
11698            \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]\
11699            \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))]\
11700            \n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11701            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11702            \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]\
11703            \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]\
11704            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11705        );
11706
11707        indie_query_plan_compare(query, expected).await;
11708    }
11709
11710    /// The outer `PromRangeManipulate` from a subquery must be preceded by
11711    /// `Sort` + `PromSeriesDivide`.
11712    #[tokio::test]
11713    async fn count_over_time_subquery() {
11714        let query = "count_over_time(some_metric[10m:1m])";
11715        let expected = String::from(
11716            "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]\
11717            \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]\
11718            \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))]\
11719            \n      PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11720            \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]\
11721            \n          PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11722            \n            PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
11723            \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]\
11724            \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]\
11725            \n                  TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
11726        );
11727        indie_query_plan_compare(query, expected).await;
11728    }
11729
11730    #[tokio::test]
11731    async fn test_hash_join() {
11732        let mut eval_stmt = EvalStmt {
11733            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11734            start: UNIX_EPOCH,
11735            end: UNIX_EPOCH
11736                .checked_add(Duration::from_secs(100_000))
11737                .unwrap(),
11738            interval: Duration::from_secs(5),
11739            lookback_delta: Duration::from_secs(1),
11740        };
11741
11742        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"}"#;
11743
11744        let prom_expr = parser::parse(case).unwrap();
11745        eval_stmt.expr = prom_expr;
11746        let table_provider = build_test_table_provider_with_fields(
11747            &[
11748                (
11749                    DEFAULT_SCHEMA_NAME.to_string(),
11750                    "http_server_requests_seconds_sum".to_string(),
11751                ),
11752                (
11753                    DEFAULT_SCHEMA_NAME.to_string(),
11754                    "http_server_requests_seconds_count".to_string(),
11755                ),
11756            ],
11757            &["uri", "kubernetes_namespace", "kubernetes_pod_name"],
11758        )
11759        .await;
11760        // Should be ok
11761        let plan =
11762            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11763                .await
11764                .unwrap();
11765        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\
11766            \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\
11767            \n    SubqueryAlias: http_server_requests_seconds_sum\
11768            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
11769            \n        PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
11770            \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\
11771            \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)\
11772            \n              TableScan: http_server_requests_seconds_sum\
11773            \n    SubqueryAlias: http_server_requests_seconds_count\
11774            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\
11775            \n        PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\
11776            \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\
11777            \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)\
11778            \n              TableScan: http_server_requests_seconds_count";
11779        assert_eq!(plan.to_string(), expected);
11780    }
11781
11782    #[tokio::test]
11783    async fn test_nested_histogram_quantile() {
11784        let mut eval_stmt = EvalStmt {
11785            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11786            start: UNIX_EPOCH,
11787            end: UNIX_EPOCH
11788                .checked_add(Duration::from_secs(100_000))
11789                .unwrap(),
11790            interval: Duration::from_secs(5),
11791            lookback_delta: Duration::from_secs(1),
11792        };
11793
11794        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]*-(.*)")"#;
11795
11796        let prom_expr = parser::parse(case).unwrap();
11797        eval_stmt.expr = prom_expr;
11798        let table_provider = build_test_table_provider_with_fields(
11799            &[(
11800                DEFAULT_SCHEMA_NAME.to_string(),
11801                "greptime_servers_grpc_requests_elapsed_bucket".to_string(),
11802            )],
11803            &["pod", "le", "path", "code", "container"],
11804        )
11805        .await;
11806        // Should be ok
11807        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11808            .await
11809            .unwrap();
11810    }
11811
11812    #[tokio::test]
11813    async fn test_histogram_quantile_binary_op() {
11814        let mut eval_stmt = EvalStmt {
11815            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11816            start: UNIX_EPOCH,
11817            end: UNIX_EPOCH
11818                .checked_add(Duration::from_secs(100_000))
11819                .unwrap(),
11820            interval: Duration::from_secs(5),
11821            lookback_delta: Duration::from_secs(1),
11822        };
11823
11824        // Arithmetic applied to a histogram_quantile() result. Regression for #8144:
11825        // HistogramFold used to drop the input column qualifiers, so the binary-op
11826        // projection failed to resolve the qualified tag column.
11827        let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#;
11828
11829        let prom_expr = parser::parse(case).unwrap();
11830        eval_stmt.expr = prom_expr;
11831        let table_provider = build_test_table_provider_with_fields(
11832            &[(
11833                DEFAULT_SCHEMA_NAME.to_string(),
11834                "http_request_duration_seconds_bucket".to_string(),
11835            )],
11836            &["pod", "le"],
11837        )
11838        .await;
11839        // Should plan without a "No field named ..." error.
11840        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11841            .await
11842            .unwrap();
11843    }
11844
11845    #[tokio::test]
11846    async fn test_parse_and_operator() {
11847        let mut eval_stmt = EvalStmt {
11848            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11849            start: UNIX_EPOCH,
11850            end: UNIX_EPOCH
11851                .checked_add(Duration::from_secs(100_000))
11852                .unwrap(),
11853            interval: Duration::from_secs(5),
11854            lookback_delta: Duration::from_secs(1),
11855        };
11856
11857        let cases = [
11858            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)"#,
11859            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)"#,
11860        ];
11861
11862        for case in cases {
11863            let prom_expr = parser::parse(case).unwrap();
11864            eval_stmt.expr = prom_expr;
11865            let table_provider = build_test_table_provider_with_fields(
11866                &[
11867                    (
11868                        DEFAULT_SCHEMA_NAME.to_string(),
11869                        "kubelet_volume_stats_used_bytes".to_string(),
11870                    ),
11871                    (
11872                        DEFAULT_SCHEMA_NAME.to_string(),
11873                        "kubelet_volume_stats_capacity_bytes".to_string(),
11874                    ),
11875                ],
11876                &["namespace", "persistentvolumeclaim"],
11877            )
11878            .await;
11879            // Should be ok
11880            let _ =
11881                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11882                    .await
11883                    .unwrap();
11884        }
11885    }
11886
11887    #[tokio::test]
11888    async fn test_nested_binary_op() {
11889        let mut eval_stmt = EvalStmt {
11890            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11891            start: UNIX_EPOCH,
11892            end: UNIX_EPOCH
11893                .checked_add(Duration::from_secs(100_000))
11894                .unwrap(),
11895            interval: Duration::from_secs(5),
11896            lookback_delta: Duration::from_secs(1),
11897        };
11898
11899        let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) -
11900        (
11901            sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m]))
11902            or
11903            vector(0)
11904        )"#;
11905
11906        let prom_expr = parser::parse(case).unwrap();
11907        eval_stmt.expr = prom_expr;
11908        let table_provider = build_test_table_provider_with_fields(
11909            &[(
11910                DEFAULT_SCHEMA_NAME.to_string(),
11911                "nginx_ingress_controller_requests".to_string(),
11912            )],
11913            &["namespace", "job"],
11914        )
11915        .await;
11916        // Should be ok
11917        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11918            .await
11919            .unwrap();
11920    }
11921
11922    #[tokio::test]
11923    async fn test_parse_or_operator() {
11924        let mut eval_stmt = EvalStmt {
11925            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
11926            start: UNIX_EPOCH,
11927            end: UNIX_EPOCH
11928                .checked_add(Duration::from_secs(100_000))
11929                .unwrap(),
11930            interval: Duration::from_secs(5),
11931            lookback_delta: Duration::from_secs(1),
11932        };
11933
11934        let case = r#"
11935        sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) /
11936        (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100)
11937            or
11938        200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) /
11939        sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#;
11940
11941        let table_provider = build_test_table_provider_with_fields(
11942            &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
11943            &["tenant_name", "cluster_name"],
11944        )
11945        .await;
11946        eval_stmt.expr = parser::parse(case).unwrap();
11947        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11948            .await
11949            .unwrap();
11950
11951        let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11952            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) +
11953            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11954            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
11955            or
11956            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11957            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0
11958            or
11959            sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) /
11960            (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#;
11961        let table_provider = build_test_table_provider_with_fields(
11962            &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())],
11963            &["tenant_name", "cluster_name"],
11964        )
11965        .await;
11966        eval_stmt.expr = parser::parse(case).unwrap();
11967        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11968            .await
11969            .unwrap();
11970
11971        let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) +
11972            sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
11973            (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or
11974            (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#;
11975        let table_provider = build_test_table_provider_with_fields(
11976            &[
11977                (
11978                    DEFAULT_SCHEMA_NAME.to_string(),
11979                    "background_waitevent_cnt".to_string(),
11980                ),
11981                (
11982                    DEFAULT_SCHEMA_NAME.to_string(),
11983                    "foreground_waitevent_cnt".to_string(),
11984                ),
11985            ],
11986            &["tenant_name", "cluster_name"],
11987        )
11988        .await;
11989        eval_stmt.expr = parser::parse(case).unwrap();
11990        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
11991            .await
11992            .unwrap();
11993
11994        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)"#;
11995        let table_provider = build_test_table_provider_with_fields(
11996            &[
11997                (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()),
11998                (
11999                    DEFAULT_SCHEMA_NAME.to_string(),
12000                    "container_cpu_load_average_10s".to_string(),
12001                ),
12002                (
12003                    DEFAULT_SCHEMA_NAME.to_string(),
12004                    "container_spec_cpu_quota".to_string(),
12005                ),
12006            ],
12007            &["cluster_name", "host_name"],
12008        )
12009        .await;
12010        eval_stmt.expr = parser::parse(case).unwrap();
12011        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12012            .await
12013            .unwrap();
12014    }
12015
12016    #[tokio::test]
12017    async fn value_matcher() {
12018        // template
12019        let mut eval_stmt = EvalStmt {
12020            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12021            start: UNIX_EPOCH,
12022            end: UNIX_EPOCH
12023                .checked_add(Duration::from_secs(100_000))
12024                .unwrap(),
12025            interval: Duration::from_secs(5),
12026            lookback_delta: Duration::from_secs(1),
12027        };
12028
12029        let cases = [
12030            // single equal matcher
12031            (
12032                r#"some_metric{__field__="field_1"}"#,
12033                vec![
12034                    "some_metric.field_1",
12035                    "some_metric.tag_0",
12036                    "some_metric.tag_1",
12037                    "some_metric.tag_2",
12038                    "some_metric.timestamp",
12039                ],
12040            ),
12041            // two equal matchers
12042            (
12043                r#"some_metric{__field__="field_1", __field__="field_0"}"#,
12044                vec![
12045                    "some_metric.field_0",
12046                    "some_metric.field_1",
12047                    "some_metric.tag_0",
12048                    "some_metric.tag_1",
12049                    "some_metric.tag_2",
12050                    "some_metric.timestamp",
12051                ],
12052            ),
12053            // single not_eq matcher
12054            (
12055                r#"some_metric{__field__!="field_1"}"#,
12056                vec![
12057                    "some_metric.field_0",
12058                    "some_metric.field_2",
12059                    "some_metric.tag_0",
12060                    "some_metric.tag_1",
12061                    "some_metric.tag_2",
12062                    "some_metric.timestamp",
12063                ],
12064            ),
12065            // two not_eq matchers
12066            (
12067                r#"some_metric{__field__!="field_1", __field__!="field_2"}"#,
12068                vec![
12069                    "some_metric.field_0",
12070                    "some_metric.tag_0",
12071                    "some_metric.tag_1",
12072                    "some_metric.tag_2",
12073                    "some_metric.timestamp",
12074                ],
12075            ),
12076            // equal and not_eq matchers (no conflict)
12077            (
12078                r#"some_metric{__field__="field_1", __field__!="field_0"}"#,
12079                vec![
12080                    "some_metric.field_1",
12081                    "some_metric.tag_0",
12082                    "some_metric.tag_1",
12083                    "some_metric.tag_2",
12084                    "some_metric.timestamp",
12085                ],
12086            ),
12087            // equal and not_eq matchers (conflict)
12088            (
12089                r#"some_metric{__field__="field_2", __field__!="field_2"}"#,
12090                vec![
12091                    "some_metric.tag_0",
12092                    "some_metric.tag_1",
12093                    "some_metric.tag_2",
12094                    "some_metric.timestamp",
12095                ],
12096            ),
12097            // single regex eq matcher
12098            (
12099                r#"some_metric{__field__=~"field_1|field_2"}"#,
12100                vec![
12101                    "some_metric.field_1",
12102                    "some_metric.field_2",
12103                    "some_metric.tag_0",
12104                    "some_metric.tag_1",
12105                    "some_metric.tag_2",
12106                    "some_metric.timestamp",
12107                ],
12108            ),
12109            // single regex not_eq matcher
12110            (
12111                r#"some_metric{__field__!~"field_1|field_2"}"#,
12112                vec![
12113                    "some_metric.field_0",
12114                    "some_metric.tag_0",
12115                    "some_metric.tag_1",
12116                    "some_metric.tag_2",
12117                    "some_metric.timestamp",
12118                ],
12119            ),
12120        ];
12121
12122        for case in cases {
12123            let prom_expr = parser::parse(case.0).unwrap();
12124            eval_stmt.expr = prom_expr;
12125            let table_provider = build_test_table_provider(
12126                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12127                3,
12128                3,
12129            )
12130            .await;
12131            let plan =
12132                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12133                    .await
12134                    .unwrap();
12135            let mut fields = plan.schema().field_names();
12136            let mut expected = case.1.into_iter().map(String::from).collect::<Vec<_>>();
12137            fields.sort();
12138            expected.sort();
12139            assert_eq!(fields, expected, "case: {:?}", case.0);
12140        }
12141
12142        let bad_cases = [
12143            r#"some_metric{__field__="nonexistent"}"#,
12144            r#"some_metric{__field__!="nonexistent"}"#,
12145        ];
12146
12147        for case in bad_cases {
12148            let prom_expr = parser::parse(case).unwrap();
12149            eval_stmt.expr = prom_expr;
12150            let table_provider = build_test_table_provider(
12151                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12152                3,
12153                3,
12154            )
12155            .await;
12156            let plan =
12157                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12158                    .await;
12159            assert!(plan.is_err(), "case: {:?}", case);
12160        }
12161    }
12162
12163    #[tokio::test]
12164    async fn custom_schema() {
12165        let query = "some_alt_metric{__schema__=\"greptime_private\"}";
12166        let expected = String::from(
12167            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12168            \n  PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12169            \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]\
12170            \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]\
12171            \n        TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12172        );
12173
12174        indie_query_plan_compare(query, expected).await;
12175
12176        let query = "some_alt_metric{__database__=\"greptime_private\"}";
12177        let expected = String::from(
12178            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12179            \n  PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12180            \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]\
12181            \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]\
12182            \n        TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12183        );
12184
12185        indie_query_plan_compare(query, expected).await;
12186
12187        let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric";
12188        let expected = String::from(
12189            "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]\
12190            \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]\
12191            \n    SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12192            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12193            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12194            \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]\
12195            \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]\
12196            \n              TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12197            \n    SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12198            \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12199            \n        PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\
12200            \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]\
12201            \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]\
12202            \n              TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]",
12203        );
12204
12205        indie_query_plan_compare(query, expected).await;
12206    }
12207
12208    #[tokio::test]
12209    async fn only_equals_is_supported_for_special_matcher() {
12210        let queries = &[
12211            "some_alt_metric{__schema__!=\"greptime_private\"}",
12212            "some_alt_metric{__schema__=~\"lalala\"}",
12213            "some_alt_metric{__database__!=\"greptime_private\"}",
12214            "some_alt_metric{__database__=~\"lalala\"}",
12215        ];
12216
12217        for query in queries {
12218            let prom_expr = parser::parse(query).unwrap();
12219            let eval_stmt = EvalStmt {
12220                expr: prom_expr,
12221                start: UNIX_EPOCH,
12222                end: UNIX_EPOCH
12223                    .checked_add(Duration::from_secs(100_000))
12224                    .unwrap(),
12225                interval: Duration::from_secs(5),
12226                lookback_delta: Duration::from_secs(1),
12227            };
12228
12229            let table_provider = build_test_table_provider(
12230                &[
12231                    (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
12232                    (
12233                        "greptime_private".to_string(),
12234                        "some_alt_metric".to_string(),
12235                    ),
12236                ],
12237                1,
12238                1,
12239            )
12240            .await;
12241
12242            let plan =
12243                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12244                    .await;
12245            assert!(plan.is_err(), "query: {:?}", query);
12246        }
12247    }
12248
12249    #[tokio::test]
12250    async fn native_scan_bounds_preserve_zero_lookback_and_overflow() {
12251        let table_provider = build_test_table_provider(
12252            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12253            1,
12254            1,
12255        )
12256        .await;
12257        let mut planner = PromPlanner {
12258            table_provider,
12259            ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")),
12260            promql_annotations: None,
12261        };
12262        planner.ctx.time_index_column = Some("timestamp".to_string());
12263        planner.ctx.start = 1_000;
12264        planner.ctx.lookback_delta = 0;
12265        let schema = Arc::new(
12266            DFSchema::try_from(ArrowSchema::new(vec![Field::new(
12267                "timestamp",
12268                ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
12269                false,
12270            )]))
12271            .unwrap(),
12272        );
12273        for (end, interval, windows) in [
12274            (1_000, 1_000, 1),
12275            (2_000, 1_000, 1),
12276            (7_201_000, 7_200_000, 2),
12277        ] {
12278            planner.ctx.end = end;
12279            planner.ctx.interval = interval;
12280            let filter = planner
12281                .build_time_index_filter(0, &schema)
12282                .unwrap()
12283                .unwrap()
12284                .to_string();
12285            assert_eq!(filter.matches(">=").count(), windows, "{filter}");
12286            assert!(
12287                filter.contains("TimestampNanosecond(1000000000, None)"),
12288                "{filter}"
12289            );
12290        }
12291        planner.ctx.end = i64::MAX;
12292        let filter = planner
12293            .build_time_index_filter(0, &schema)
12294            .unwrap()
12295            .unwrap()
12296            .to_string();
12297        assert!(
12298            filter.contains("timestamp >= TimestampNanosecond(1000000000, None)"),
12299            "{filter}"
12300        );
12301
12302        // A lookback subtraction can underflow milliseconds while the upper bound remains
12303        // representable. Keep that upper bound so LastRow cannot select a future sample.
12304        let ms_schema = Arc::new(
12305            DFSchema::try_from(ArrowSchema::new(vec![Field::new(
12306                "timestamp",
12307                ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
12308                false,
12309            )]))
12310            .unwrap(),
12311        );
12312        planner.ctx.start = i64::MIN + 100;
12313        planner.ctx.end = planner.ctx.start;
12314        planner.ctx.lookback_delta = 200;
12315        let filter = planner
12316            .build_time_index_filter(0, &ms_schema)
12317            .unwrap()
12318            .unwrap()
12319            .to_string();
12320        assert_eq!(
12321            filter,
12322            format!(
12323                "timestamp <= TimestampMillisecond({}, None)",
12324                i64::MIN + 100
12325            )
12326        );
12327
12328        // The lower bound can also overflow while converting milliseconds to native nanoseconds.
12329        // Its representable upper bound still has to reach the scan.
12330        planner.ctx.start = 0;
12331        planner.ctx.end = 0;
12332        planner.ctx.lookback_delta = 300_000;
12333        let filter = planner
12334            .build_time_index_filter(9_223_372_036_854, &schema)
12335            .unwrap()
12336            .unwrap()
12337            .to_string();
12338        assert_eq!(
12339            filter,
12340            "timestamp <= TimestampNanosecond(-9223372036854000000, None)"
12341        );
12342    }
12343
12344    #[tokio::test]
12345    async fn test_non_ms_precision() {
12346        let catalog_list = MemoryCatalogManager::with_default_setup();
12347        let columns = vec![
12348            ColumnSchema::new(
12349                "tag".to_string(),
12350                ConcreteDataType::string_datatype(),
12351                false,
12352            ),
12353            ColumnSchema::new(
12354                "timestamp".to_string(),
12355                ConcreteDataType::timestamp_nanosecond_datatype(),
12356                false,
12357            )
12358            .with_time_index(true),
12359            ColumnSchema::new(
12360                "field".to_string(),
12361                ConcreteDataType::float64_datatype(),
12362                true,
12363            ),
12364        ];
12365        let schema = Arc::new(Schema::new(columns));
12366        let table_meta = TableMetaBuilder::empty()
12367            .schema(schema)
12368            .primary_key_indices(vec![0])
12369            .value_indices(vec![2])
12370            .next_column_id(1024)
12371            .build()
12372            .unwrap();
12373        let table_info = TableInfoBuilder::default()
12374            .name("metrics".to_string())
12375            .meta(table_meta)
12376            .build()
12377            .unwrap();
12378        let table = EmptyTable::from_table_info(&table_info);
12379        assert!(
12380            catalog_list
12381                .register_table_sync(RegisterTableRequest {
12382                    catalog: DEFAULT_CATALOG_NAME.to_string(),
12383                    schema: DEFAULT_SCHEMA_NAME.to_string(),
12384                    table_name: "metrics".to_string(),
12385                    table_id: 1024,
12386                    table,
12387                })
12388                .is_ok()
12389        );
12390
12391        let plan = PromPlanner::stmt_to_plan(
12392            DfTableSourceProvider::new(
12393                catalog_list.clone(),
12394                false,
12395                QueryContext::arc(),
12396                DummyDecoder::arc(),
12397                true,
12398            ),
12399            &EvalStmt {
12400                expr: parser::parse("metrics{tag = \"1\"}").unwrap(),
12401                start: UNIX_EPOCH,
12402                end: UNIX_EPOCH
12403                    .checked_add(Duration::from_secs(100_000))
12404                    .unwrap(),
12405                interval: Duration::from_secs(5),
12406                lookback_delta: Duration::from_secs(1),
12407            },
12408            &build_query_engine_state(),
12409        )
12410        .await
12411        .unwrap();
12412        assert_eq!(
12413            plan.display_indent_schema().to_string(),
12414            "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\n  PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n    Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n      Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-1000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n        Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n          TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12415        );
12416        let plan = PromPlanner::stmt_to_plan(
12417            DfTableSourceProvider::new(
12418                catalog_list.clone(),
12419                false,
12420                QueryContext::arc(),
12421                DummyDecoder::arc(),
12422                true,
12423            ),
12424            &EvalStmt {
12425                expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(),
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            &build_query_engine_state(),
12434        )
12435        .await
12436        .unwrap();
12437        assert_eq!(
12438            plan.display_indent_schema().to_string(),
12439            "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]\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]\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))]\n      PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n        PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n          Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n            Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-5000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n              Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n                TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]"
12440        );
12441    }
12442
12443    #[tokio::test]
12444    async fn test_nonexistent_label() {
12445        // template
12446        let mut eval_stmt = EvalStmt {
12447            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12448            start: UNIX_EPOCH,
12449            end: UNIX_EPOCH
12450                .checked_add(Duration::from_secs(100_000))
12451                .unwrap(),
12452            interval: Duration::from_secs(5),
12453            lookback_delta: Duration::from_secs(1),
12454        };
12455
12456        let case = r#"some_metric{nonexistent="hi"}"#;
12457        let prom_expr = parser::parse(case).unwrap();
12458        eval_stmt.expr = prom_expr;
12459        let table_provider = build_test_table_provider(
12460            &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12461            3,
12462            3,
12463        )
12464        .await;
12465        // Should be ok
12466        let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12467            .await
12468            .unwrap();
12469    }
12470
12471    #[tokio::test]
12472    async fn test_label_join() {
12473        let prom_expr = parser::parse(
12474            "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')",
12475        )
12476        .unwrap();
12477        let eval_stmt = EvalStmt {
12478            expr: prom_expr,
12479            start: UNIX_EPOCH,
12480            end: UNIX_EPOCH
12481                .checked_add(Duration::from_secs(100_000))
12482                .unwrap(),
12483            interval: Duration::from_secs(5),
12484            lookback_delta: Duration::from_secs(1),
12485        };
12486
12487        let table_provider =
12488            build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1)
12489                .await;
12490        let plan =
12491            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12492                .await
12493                .unwrap();
12494
12495        let expected = r#"
12496Filter: 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]
12497  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]
12498    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]
12499      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]
12500        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]
12501          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]
12502            TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
12503
12504        let ret = plan.display_indent_schema().to_string();
12505        assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
12506    }
12507
12508    #[tokio::test]
12509    async fn test_label_replace() {
12510        let prom_expr = parser::parse(
12511            "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")",
12512        )
12513        .unwrap();
12514        let eval_stmt = EvalStmt {
12515            expr: prom_expr,
12516            start: UNIX_EPOCH,
12517            end: UNIX_EPOCH
12518                .checked_add(Duration::from_secs(100_000))
12519                .unwrap(),
12520            interval: Duration::from_secs(5),
12521            lookback_delta: Duration::from_secs(1),
12522        };
12523
12524        let table_provider =
12525            build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1)
12526                .await;
12527        let plan =
12528            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12529                .await
12530                .unwrap();
12531
12532        let expected = r#"
12533Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8]
12534  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]
12535    PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12536      PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12537        Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]
12538          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]
12539            TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#;
12540
12541        let ret = plan.display_indent_schema().to_string();
12542        assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
12543    }
12544
12545    #[tokio::test]
12546    async fn label_replace_aggregation_queries_plan_successfully() {
12547        let aggregate =
12548            r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#;
12549        let queries = [
12550            aggregate.to_string(),
12551            format!("{aggregate} <= 10"),
12552            format!("{aggregate} * 0.8"),
12553            format!("0.8 * {aggregate}"),
12554            format!("{aggregate} <= {aggregate} * 0.8"),
12555        ];
12556        let state = build_query_engine_state();
12557        let mut failures = Vec::new();
12558
12559        for query in queries {
12560            let table_provider = build_test_table_provider(
12561                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
12562                1,
12563                1,
12564            )
12565            .await;
12566            if let Err(error) =
12567                PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await
12568            {
12569                failures.push(format!("{query}: {error:?}"));
12570            }
12571        }
12572
12573        assert!(failures.is_empty(), "{}", failures.join("\n"));
12574    }
12575
12576    #[tokio::test]
12577    async fn test_matchers_to_expr() {
12578        let mut eval_stmt = EvalStmt {
12579            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12580            start: UNIX_EPOCH,
12581            end: UNIX_EPOCH
12582                .checked_add(Duration::from_secs(100_000))
12583                .unwrap(),
12584            interval: Duration::from_secs(5),
12585            lookback_delta: Duration::from_secs(1),
12586        };
12587        let case =
12588            r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#;
12589
12590        let prom_expr = parser::parse(case).unwrap();
12591        eval_stmt.expr = prom_expr;
12592        let table_provider = build_test_table_provider(
12593            &[(
12594                DEFAULT_SCHEMA_NAME.to_string(),
12595                "prometheus_tsdb_head_series".to_string(),
12596            )],
12597            3,
12598            3,
12599        )
12600        .await;
12601        let plan =
12602            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12603                .await
12604                .unwrap();
12605        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]\
12606        \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]\
12607        \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]\
12608        \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]\
12609        \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]\
12610        \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]\
12611        \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]";
12612        assert_eq!(plan.display_indent_schema().to_string(), expected);
12613    }
12614
12615    #[tokio::test]
12616    async fn test_topk_expr() {
12617        let mut eval_stmt = EvalStmt {
12618            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12619            start: UNIX_EPOCH,
12620            end: UNIX_EPOCH
12621                .checked_add(Duration::from_secs(100_000))
12622                .unwrap(),
12623            interval: Duration::from_secs(5),
12624            lookback_delta: Duration::from_secs(1),
12625        };
12626        let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
12627
12628        let prom_expr = parser::parse(case).unwrap();
12629        eval_stmt.expr = prom_expr;
12630        let table_provider = build_test_table_provider_with_fields(
12631            &[
12632                (
12633                    DEFAULT_SCHEMA_NAME.to_string(),
12634                    "prometheus_tsdb_head_series".to_string(),
12635                ),
12636                (
12637                    DEFAULT_SCHEMA_NAME.to_string(),
12638                    "http_server_requests_seconds_count".to_string(),
12639                ),
12640            ],
12641            &["ip"],
12642        )
12643        .await;
12644
12645        let plan =
12646            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12647                .await
12648                .unwrap();
12649        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)]\
12650        \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]\
12651        \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]\
12652        \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]\
12653        \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]\
12654        \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]\
12655        \n            PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12656        \n              PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12657        \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]\
12658        \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]\
12659        \n                    TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12660
12661        assert_eq!(plan.display_indent_schema().to_string(), expected);
12662    }
12663
12664    #[tokio::test]
12665    async fn test_count_values_expr() {
12666        let mut eval_stmt = EvalStmt {
12667            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12668            start: UNIX_EPOCH,
12669            end: UNIX_EPOCH
12670                .checked_add(Duration::from_secs(100_000))
12671                .unwrap(),
12672            interval: Duration::from_secs(5),
12673            lookback_delta: Duration::from_secs(1),
12674        };
12675        let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
12676
12677        let prom_expr = parser::parse(case).unwrap();
12678        eval_stmt.expr = prom_expr;
12679        let table_provider = build_test_table_provider_with_fields(
12680            &[
12681                (
12682                    DEFAULT_SCHEMA_NAME.to_string(),
12683                    "prometheus_tsdb_head_series".to_string(),
12684                ),
12685                (
12686                    DEFAULT_SCHEMA_NAME.to_string(),
12687                    "http_server_requests_seconds_count".to_string(),
12688                ),
12689            ],
12690            &["ip"],
12691        )
12692        .await;
12693
12694        let plan =
12695            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12696                .await
12697                .unwrap();
12698        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]\
12699        \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]\
12700        \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]\
12701        \n      PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12702        \n        PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12703        \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]\
12704        \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]\
12705        \n              TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12706
12707        assert_eq!(plan.display_indent_schema().to_string(), expected);
12708    }
12709
12710    #[tokio::test]
12711    async fn test_value_alias() {
12712        let mut eval_stmt = EvalStmt {
12713            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12714            start: UNIX_EPOCH,
12715            end: UNIX_EPOCH
12716                .checked_add(Duration::from_secs(100_000))
12717                .unwrap(),
12718            interval: Duration::from_secs(5),
12719            lookback_delta: Duration::from_secs(1),
12720        };
12721        let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#;
12722
12723        let prom_expr = parser::parse(case).unwrap();
12724        eval_stmt.expr = prom_expr;
12725        eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series");
12726        let table_provider = build_test_table_provider_with_fields(
12727            &[
12728                (
12729                    DEFAULT_SCHEMA_NAME.to_string(),
12730                    "prometheus_tsdb_head_series".to_string(),
12731                ),
12732                (
12733                    DEFAULT_SCHEMA_NAME.to_string(),
12734                    "http_server_requests_seconds_count".to_string(),
12735                ),
12736            ],
12737            &["ip"],
12738        )
12739        .await;
12740
12741        let plan =
12742            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12743                .await
12744                .unwrap();
12745        let expected = r#"
12746Projection: 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)]
12747  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]
12748    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]
12749      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]
12750        PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
12751          PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]
12752            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]
12753              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]
12754                TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#;
12755        assert_eq!(format!("\n{}", plan.display_indent_schema()), expected);
12756    }
12757
12758    #[tokio::test]
12759    async fn test_quantile_expr() {
12760        let mut eval_stmt = EvalStmt {
12761            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12762            start: UNIX_EPOCH,
12763            end: UNIX_EPOCH
12764                .checked_add(Duration::from_secs(100_000))
12765                .unwrap(),
12766            interval: Duration::from_secs(5),
12767            lookback_delta: Duration::from_secs(1),
12768        };
12769        let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#;
12770
12771        let prom_expr = parser::parse(case).unwrap();
12772        eval_stmt.expr = prom_expr;
12773        let table_provider = build_test_table_provider_with_fields(
12774            &[
12775                (
12776                    DEFAULT_SCHEMA_NAME.to_string(),
12777                    "prometheus_tsdb_head_series".to_string(),
12778                ),
12779                (
12780                    DEFAULT_SCHEMA_NAME.to_string(),
12781                    "http_server_requests_seconds_count".to_string(),
12782                ),
12783            ],
12784            &["ip"],
12785        )
12786        .await;
12787
12788        let plan =
12789            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12790                .await
12791                .unwrap();
12792        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]\
12793        \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]\
12794        \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]\
12795        \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]\
12796        \n        PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12797        \n          PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\
12798        \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]\
12799        \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]\
12800        \n                TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]";
12801
12802        assert_eq!(plan.display_indent_schema().to_string(), expected);
12803    }
12804
12805    #[tokio::test]
12806    async fn test_or_not_exists_table_label() {
12807        let state = build_query_engine_state();
12808        let provider = build_test_table_provider_with_fields(
12809            &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())],
12810            &["job"],
12811        )
12812        .await;
12813        let raw = PromPlanner::stmt_to_plan(
12814            provider,
12815            &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#),
12816            &state,
12817        )
12818        .await
12819        .unwrap();
12820        assert!(
12821            raw.display_indent_schema()
12822                .to_string()
12823                .contains("__promql_or_match_0@")
12824        );
12825        let (optimized, batches) = execute(raw, &state).await;
12826        assert_no_internal_or_keys(optimized.schema());
12827        assert!(batches.iter().all(|batch| {
12828            batch
12829                .schema()
12830                .fields()
12831                .iter()
12832                .all(|field| !field.name().starts_with("__promql_or_match_"))
12833        }));
12834    }
12835
12836    #[tokio::test]
12837    async fn test_histogram_quantile_missing_le_column() {
12838        let mut eval_stmt = EvalStmt {
12839            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
12840            start: UNIX_EPOCH,
12841            end: UNIX_EPOCH
12842                .checked_add(Duration::from_secs(100_000))
12843                .unwrap(),
12844            interval: Duration::from_secs(5),
12845            lookback_delta: Duration::from_secs(1),
12846        };
12847
12848        // Test case: histogram_quantile with a table that doesn't have 'le' column
12849        let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#;
12850
12851        let prom_expr = parser::parse(case).unwrap();
12852        eval_stmt.expr = prom_expr;
12853
12854        // Create a table provider with a table that doesn't have 'le' column
12855        let table_provider = build_test_table_provider_with_fields(
12856            &[(
12857                DEFAULT_SCHEMA_NAME.to_string(),
12858                "non_existent_histogram_bucket".to_string(),
12859            )],
12860            &["pod", "instance"], // Note: no 'le' column
12861        )
12862        .await;
12863
12864        // Should return empty result instead of error
12865        let result =
12866            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
12867                .await;
12868
12869        // This should succeed now (returning empty result) instead of failing with "Cannot find column le"
12870        assert!(
12871            result.is_ok(),
12872            "Expected successful plan creation with empty result, but got error: {:?}",
12873            result.err()
12874        );
12875
12876        // Verify that the result is an EmptyRelation
12877        let plan = result.unwrap();
12878        match plan {
12879            LogicalPlan::EmptyRelation(_) => {
12880                // This is what we expect
12881            }
12882            _ => panic!("Expected EmptyRelation, but got: {:?}", plan),
12883        }
12884    }
12885
12886    #[tokio::test]
12887    async fn test_direct_or_normalizes_missing_match_labels() {
12888        type Case<'a> = (
12889            Option<Option<&'a str>>,
12890            Option<Option<&'a str>>,
12891            i64,
12892            i64,
12893            &'a [(f64, Option<&'a str>)],
12894        );
12895
12896        let modifier = or_modifier("lhs or on(k) rhs");
12897        #[rustfmt::skip]
12898        let cases: &[Case<'_>] = &[
12899            (None, None, 1, 1, &[(1.0, None)]),
12900            (None, Some(Some("")), 1, 1, &[(1.0, None)]),
12901            (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]),
12902            (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
12903            (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]),
12904            (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]),
12905            (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]),
12906            (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]),
12907        ];
12908        for &(left, right, left_ts, right_ts, expected) in cases {
12909            let (optimized, batches) = run(
12910                &matrix_source("lhs", left, left_ts, 1.0),
12911                &matrix_source("rhs", right, right_ts, 2.0),
12912                matrix_context("lhs", left),
12913                matrix_context("rhs", right),
12914                &modifier,
12915            )
12916            .await;
12917            assert_no_internal_or_keys(optimized.schema());
12918            assert_eq!(
12919                rows(&batches),
12920                expected
12921                    .iter()
12922                    .map(|(value, label)| (*value, label.map(str::to_string)))
12923                    .collect::<Vec<_>>()
12924            );
12925        }
12926    }
12927
12928    #[tokio::test]
12929    async fn test_direct_or_match_modifiers() {
12930        for (modifier, left, right, expected) in [
12931            (None, "left", "right", 2),
12932            (or_modifier("lhs or on(k) rhs"), "same", "same", 1),
12933            (or_modifier("lhs or on() rhs"), "left", "right", 1),
12934            (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1),
12935        ] {
12936            let (_, batches) = run(
12937                &matrix_source("lhs", Some(Some(left)), 1, 1.0),
12938                &matrix_source("rhs", Some(Some(right)), 1, 2.0),
12939                direct_or_context("lhs", &["job", "k"], "v"),
12940                direct_or_context("rhs", &["job", "k"], "v"),
12941                &modifier,
12942            )
12943            .await;
12944            assert_eq!(
12945                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
12946                expected
12947            );
12948        }
12949    }
12950
12951    #[tokio::test]
12952    async fn test_direct_or_nested_projection_uses_left_context() {
12953        let left = matrix_source("lhs", Some(Some("k")), 1, 1.0);
12954        let right = matrix_source("rhs", Some(Some("k")), 1, 2.0);
12955        let raw = plan_direct_or(
12956            scan(&left),
12957            scan(&right),
12958            direct_or_context("lhs", &["job", "k"], "v"),
12959            direct_or_context("rhs", &["job", "k"], "v"),
12960            &or_modifier("lhs or on(k) rhs"),
12961        )
12962        .await;
12963        assert!(raw.schema().iter().any(|(qualifier, field)| {
12964            qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v"
12965        }));
12966        let nested = LogicalPlanBuilder::from(raw)
12967            .project(vec![
12968                DfExpr::BinaryExpr(BinaryExpr {
12969                    left: Box::new(DfExpr::Column(Column::new(
12970                        Some(TableReference::bare("lhs")),
12971                        "v",
12972                    ))),
12973                    op: Operator::Plus,
12974                    right: Box::new(lit(1.0)),
12975                })
12976                .alias("v_plus"),
12977            ])
12978            .unwrap()
12979            .build()
12980            .unwrap();
12981        let (_, batches) = execute(nested, &build_query_engine_state()).await;
12982        assert_eq!(values(&batches, "v_plus"), vec![2.0]);
12983    }
12984
12985    #[tokio::test]
12986    async fn test_direct_or_skips_user_internal_key_name() {
12987        const USER_TAG: &str = "__promql_or_match_0";
12988        let left = tagged_source(
12989            "lhs",
12990            false,
12991            (USER_TAG, Some("left")),
12992            DirectOrValue::Float64(1.0),
12993        );
12994        let right = tagged_source(
12995            "rhs",
12996            false,
12997            (USER_TAG, Some("right")),
12998            DirectOrValue::Float64(2.0),
12999        );
13000        let raw = plan_direct_or(
13001            scan(&left),
13002            scan(&right),
13003            direct_or_context("lhs", &["job", USER_TAG], "v"),
13004            direct_or_context("rhs", &["job", USER_TAG], "v"),
13005            &or_modifier("lhs or on(missing_label) rhs"),
13006        )
13007        .await;
13008        assert!(
13009            raw.display_indent_schema()
13010                .to_string()
13011                .contains("__promql_or_match_1@")
13012        );
13013        let (_, batches) = execute(raw, &build_query_engine_state()).await;
13014        assert!(
13015            batches
13016                .iter()
13017                .all(|batch| batch.column_by_name(USER_TAG).is_some())
13018        );
13019    }
13020
13021    #[tokio::test]
13022    async fn test_direct_or_substrait_round_trip_with_normalized_key() {
13023        let state = build_query_engine_state();
13024        let ctx = SessionContext::new_with_state(state.session_state());
13025        let catalog = Arc::new(MemoryCatalogProvider::new());
13026        catalog
13027            .register_schema("public", Arc::new(MemorySchemaProvider::new()))
13028            .unwrap();
13029        ctx.register_catalog("datafusion", catalog);
13030        let left = matrix_source("lhs", Some(Some("")), 1, 1.0);
13031        let right = matrix_source("rhs", None, 1, 2.0);
13032        ctx.register_table(
13033            TableReference::full("datafusion", "public", "lhs"),
13034            table(&left),
13035        )
13036        .unwrap();
13037        ctx.register_table(
13038            TableReference::full("datafusion", "public", "rhs"),
13039            table(&right),
13040        )
13041        .unwrap();
13042        let raw = plan_direct_or(
13043            ctx.table("datafusion.public.lhs")
13044                .await
13045                .unwrap()
13046                .into_unoptimized_plan(),
13047            ctx.table("datafusion.public.rhs")
13048                .await
13049                .unwrap()
13050                .into_unoptimized_plan(),
13051            direct_or_context("lhs", &["job", "k"], "v"),
13052            direct_or_context("rhs", &["job"], "v"),
13053            &or_modifier("lhs or on(k) rhs"),
13054        )
13055        .await;
13056        let decoded = DFLogicalSubstraitConvertor
13057            .decode(
13058                DFLogicalSubstraitConvertor
13059                    .encode(&raw, DefaultSerializer)
13060                    .unwrap(),
13061                ctx.state(),
13062            )
13063            .await
13064            .unwrap();
13065        let (optimized, batches) = execute(decoded, &state).await;
13066        assert_no_internal_or_keys(optimized.schema());
13067        assert!(batches.iter().all(|batch| {
13068            batch
13069                .schema()
13070                .fields()
13071                .iter()
13072                .all(|field| !field.name().starts_with("__promql_or_match_"))
13073        }));
13074        assert_eq!(values(&batches, "v"), vec![1.0]);
13075    }
13076
13077    #[tokio::test]
13078    async fn test_direct_or_numeric_value_types() {
13079        let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0));
13080        let right = tagged_source(
13081            "rhs",
13082            false,
13083            ("k", Some("rhs")),
13084            DirectOrValue::Float64(0.5),
13085        );
13086        let (optimized, batches) = run(
13087            &left,
13088            &right,
13089            direct_or_context("lhs", &["job", "k"], "v"),
13090            direct_or_context("rhs", &["job", "k"], "v"),
13091            &or_modifier("lhs or on(k) rhs"),
13092        )
13093        .await;
13094        assert_eq!(
13095            optimized
13096                .schema()
13097                .field_with_name(None, "v")
13098                .unwrap()
13099                .data_type(),
13100            &ArrowDataType::Float64
13101        );
13102        assert_eq!(values(&batches, "v"), vec![0.5]);
13103        let provider = build_test_table_provider_with_fields(
13104            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
13105            &[],
13106        )
13107        .await;
13108        let mut planner = PromPlanner {
13109            table_provider: provider,
13110            ctx: PromPlannerContext::default(),
13111            promql_annotations: None,
13112        };
13113        let left_context = direct_or_context("lhs", &["job"], "v");
13114        let right_context = direct_or_context("rhs", &["job"], "v");
13115        let error = planner
13116            .or_operator(
13117                scan(&job_source("lhs", DirectOrValue::Utf8("x"))),
13118                scan(&job_source("rhs", DirectOrValue::Float64(1.0))),
13119                left_context.tag_columns.iter().cloned().collect(),
13120                right_context.tag_columns.iter().cloned().collect(),
13121                left_context,
13122                right_context,
13123                &or_modifier("lhs or on() rhs"),
13124            )
13125            .unwrap_err();
13126        assert!(
13127            error
13128                .to_string()
13129                .contains("OR value fields have incompatible types")
13130        );
13131    }
13132
13133    #[tokio::test]
13134    async fn test_or_with_histogram_quantile_missing_le_column() {
13135        let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#;
13136        let eval_stmt = build_eval_stmt(case);
13137        let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13138
13139        let plan =
13140            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13141                .await
13142                .unwrap();
13143        assert_normal_metric_schema(&plan);
13144    }
13145
13146    #[tokio::test]
13147    async fn test_or_with_right_empty_histogram_restores_left_context() {
13148        let eval_stmt = build_eval_stmt(
13149            r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#,
13150        );
13151        let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13152
13153        PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13154            .await
13155            .unwrap();
13156    }
13157
13158    #[tokio::test]
13159    async fn test_or_with_both_empty_histograms() {
13160        let eval_stmt = build_eval_stmt(
13161            r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#,
13162        );
13163        let table_provider = build_test_table_provider_with_fields(
13164            &[
13165                (
13166                    DEFAULT_SCHEMA_NAME.to_string(),
13167                    "left_histogram_bucket".to_string(),
13168                ),
13169                (
13170                    DEFAULT_SCHEMA_NAME.to_string(),
13171                    "right_histogram_bucket".to_string(),
13172                ),
13173            ],
13174            &["pod", "instance"],
13175        )
13176        .await;
13177
13178        let plan =
13179            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13180                .await
13181                .unwrap();
13182        match plan {
13183            LogicalPlan::EmptyRelation(relation) => {
13184                assert!(!relation.produce_one_row);
13185                assert!(!relation.schema.fields().is_empty());
13186                assert!(
13187                    relation
13188                        .schema
13189                        .fields()
13190                        .iter()
13191                        .any(|field| field.data_type() == &ArrowDataType::Float64)
13192                );
13193                assert!(
13194                    relation
13195                        .schema
13196                        .fields()
13197                        .iter()
13198                        .any(|field| field.name() == "pod")
13199                );
13200                assert!(
13201                    !relation
13202                        .schema
13203                        .fields()
13204                        .iter()
13205                        .any(|field| field.name() == "instance")
13206                );
13207            }
13208            _ => panic!("Expected EmptyRelation, but got: {plan:?}"),
13209        }
13210    }
13211
13212    #[tokio::test]
13213    async fn test_nested_or_with_both_empty_histograms() {
13214        for case in [
13215            r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#,
13216            r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#,
13217        ] {
13218            let eval_stmt = build_eval_stmt(case);
13219            let table_provider = build_test_table_provider_with_fields(
13220                &[
13221                    (
13222                        DEFAULT_SCHEMA_NAME.to_string(),
13223                        "left_histogram_bucket".to_string(),
13224                    ),
13225                    (
13226                        DEFAULT_SCHEMA_NAME.to_string(),
13227                        "right_histogram_bucket".to_string(),
13228                    ),
13229                ],
13230                &["pod", "instance"],
13231            )
13232            .await;
13233
13234            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13235                .await
13236                .unwrap();
13237        }
13238    }
13239
13240    #[tokio::test]
13241    async fn test_or_with_empty_histogram_modifiers() {
13242        for case in [
13243            r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#,
13244            r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#,
13245        ] {
13246            let eval_stmt = build_eval_stmt(case);
13247            let table_provider = build_missing_le_or_normal_metric_table_provider().await;
13248
13249            let plan =
13250                PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
13251                    .await
13252                    .unwrap();
13253            assert_normal_metric_schema(&plan);
13254        }
13255    }
13256
13257    #[tokio::test]
13258    async fn test_unless_preserves_left_context_for_histogram() {
13259        let eval_stmt = build_eval_stmt(
13260            r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#,
13261        );
13262        let state = build_query_engine_state();
13263        let plan = PromPlanner::stmt_to_plan(
13264            build_set_op_context_table_provider().await,
13265            &eval_stmt,
13266            &state,
13267        )
13268        .await
13269        .unwrap();
13270        assert!(contains_histogram_fold(&plan), "{plan:?}");
13271        let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await;
13272        assert!(contains_histogram_fold(&optimized), "{optimized:?}");
13273        let batches =
13274            datafusion::physical_plan::collect(physical, state.session_state().task_ctx())
13275                .await
13276                .unwrap();
13277        assert!(batches.iter().all(|batch| batch.num_rows() == 0));
13278    }
13279
13280    #[tokio::test]
13281    async fn test_and_preserves_left_context_for_histogram() {
13282        let eval_stmt = build_eval_stmt(
13283            r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#,
13284        );
13285        let plan = PromPlanner::stmt_to_plan(
13286            build_set_op_context_table_provider().await,
13287            &eval_stmt,
13288            &build_query_engine_state(),
13289        )
13290        .await
13291        .unwrap();
13292        assert!(contains_histogram_fold(&plan), "{plan:?}");
13293    }
13294
13295    #[tokio::test]
13296    async fn test_and_preserves_left_context_when_le_is_missing() {
13297        let eval_stmt =
13298            build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#);
13299        let plan = PromPlanner::stmt_to_plan(
13300            build_set_op_context_table_provider().await,
13301            &eval_stmt,
13302            &build_query_engine_state(),
13303        )
13304        .await
13305        .unwrap();
13306        assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}");
13307        assert!(!plan.schema().fields().is_empty());
13308        assert!(!contains_histogram_fold(&plan), "{plan:?}");
13309    }
13310
13311    #[tokio::test]
13312    async fn test_or_context_uses_left_qualified_output() {
13313        let case = r#"(normal_metric or other_metric) + 1"#;
13314        let eval_stmt = build_eval_stmt(case);
13315        let state = build_query_engine_state();
13316        let plan =
13317            PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state)
13318                .await
13319                .unwrap();
13320        assert!(
13321            plan.schema()
13322                .fields()
13323                .iter()
13324                .any(|field| field.data_type() == &ArrowDataType::Float64),
13325            "{plan:?}"
13326        );
13327        let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await;
13328    }
13329
13330    #[tokio::test]
13331    async fn test_or_context_uses_left_qualified_empty_histogram_output() {
13332        let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#;
13333        let eval_stmt = build_eval_stmt(case);
13334        let plan = PromPlanner::stmt_to_plan(
13335            build_or_context_table_provider().await,
13336            &eval_stmt,
13337            &build_query_engine_state(),
13338        )
13339        .await
13340        .unwrap();
13341        assert!(
13342            plan.schema()
13343                .fields()
13344                .iter()
13345                .any(|field| field.data_type() == &ArrowDataType::Float64),
13346            "{plan:?}"
13347        );
13348    }
13349
13350    #[tokio::test]
13351    async fn test_direct_or_preserves_float_and_native_histogram_samples() {
13352        for histogram_on_left in [false, true] {
13353            let (planner, plan) = mixed_direct_or(histogram_on_left).await;
13354
13355            let float_field = &planner.ctx.field_columns[0];
13356            let histogram_field = &planner.ctx.field_columns[1];
13357            assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX));
13358            assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX));
13359            assert_eq!(
13360                plan.schema()
13361                    .field_with_name(None, float_field)
13362                    .unwrap()
13363                    .data_type(),
13364                &ArrowDataType::Float64
13365            );
13366            assert_eq!(
13367                plan.schema()
13368                    .field_with_name(None, histogram_field)
13369                    .unwrap()
13370                    .data_type(),
13371                &native_histogram_value_type().as_arrow_type()
13372            );
13373
13374            let (optimized, batches) = execute(plan, &build_query_engine_state()).await;
13375            assert_no_internal_or_keys(optimized.schema());
13376            let mut sample_kinds = batches
13377                .iter()
13378                .flat_map(|batch| {
13379                    let values = batch.column_by_name(float_field).unwrap();
13380                    let histograms = batch.column_by_name(histogram_field).unwrap();
13381                    (0..batch.num_rows())
13382                        .map(|row| (values.is_valid(row), histograms.is_valid(row)))
13383                })
13384                .collect::<Vec<_>>();
13385            sample_kinds.sort_unstable();
13386            assert_eq!(sample_kinds, vec![(false, true), (true, false)]);
13387        }
13388    }
13389
13390    #[tokio::test]
13391    async fn malformed_classic_bucket_does_not_drop_native_histogram() {
13392        let state = build_query_engine_state();
13393        let collector = PromqlAnnotationCollector::default();
13394        let plan = PromPlanner::stmt_to_plan_with_annotations(
13395            operator_table_provider(),
13396            &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"),
13397            &state,
13398            Some(collector.clone()),
13399        )
13400        .await
13401        .unwrap();
13402        let value_field = plan
13403            .schema()
13404            .fields()
13405            .iter()
13406            .find(|field| field.data_type() == &ArrowDataType::Float64)
13407            .unwrap()
13408            .name()
13409            .clone();
13410
13411        let (_, batches) = execute(plan, &state).await;
13412        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13413        assert_eq!(values(&batches, &value_field), vec![0.0]);
13414        let mut warnings = vec![];
13415        let mut infos = vec![];
13416        collector.append_to(&mut warnings, &mut infos);
13417        assert!(warnings.is_empty());
13418        assert!(infos.is_empty());
13419    }
13420
13421    #[tokio::test]
13422    async fn test_mixed_binary_operator_aligns_both_alternative_inputs() {
13423        let state = build_query_engine_state();
13424        let plan = PromPlanner::stmt_to_plan(
13425            operator_table_provider(),
13426            &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"),
13427            &state,
13428        )
13429        .await
13430        .unwrap();
13431        let plan_text = plan.display_indent_schema().to_string();
13432        assert!(
13433            plan_text.contains("prom_native_histogram_mul_scalar"),
13434            "{plan_text}"
13435        );
13436        assert!(
13437            plan_text.contains("prom_native_histogram_scalar_mul"),
13438            "{plan_text}"
13439        );
13440        let float_field = plan
13441            .schema()
13442            .fields()
13443            .iter()
13444            .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
13445            .unwrap()
13446            .name()
13447            .clone();
13448        let histogram_field = plan
13449            .schema()
13450            .fields()
13451            .iter()
13452            .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13453            .unwrap()
13454            .name()
13455            .clone();
13456
13457        let (_, batches) = execute(plan, &state).await;
13458        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
13459        assert!(values(&batches, &float_field).is_empty());
13460        let mut sums = histograms(&batches, &histogram_field)
13461            .into_iter()
13462            .map(|histogram| histogram.sum)
13463            .collect::<Vec<_>>();
13464        sums.sort_by(f64::total_cmp);
13465        assert_eq!(sums, vec![2.0, 3.0]);
13466    }
13467
13468    #[tokio::test]
13469    async fn test_mixed_binary_operator_reports_only_dropped_samples() {
13470        for (query, expected_rows, expected_infos) in [
13471            ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1),
13472            ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0),
13473            ("(lf or on(tag) lh) % on(tag) lh", 0, 1),
13474        ] {
13475            let state = build_query_engine_state();
13476            let annotations = PromqlAnnotationCollector::default();
13477            let plan = PromPlanner::stmt_to_plan_with_annotations(
13478                operator_table_provider(),
13479                &operator_eval_stmt(query),
13480                &state,
13481                Some(annotations.clone()),
13482            )
13483            .await
13484            .unwrap();
13485
13486            let (_, batches) = execute(plan, &state).await;
13487            assert_eq!(
13488                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
13489                expected_rows,
13490                "{query}"
13491            );
13492            let mut warnings = vec![];
13493            let mut infos = vec![];
13494            annotations.append_to(&mut warnings, &mut infos);
13495            assert!(warnings.is_empty(), "{query}: {warnings:?}");
13496            assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}");
13497        }
13498    }
13499
13500    #[tokio::test]
13501    async fn test_histogram_only_min_drops_empty_aggregate_group() {
13502        // `min` over native-histogram-only input drops every sample in the group, so the
13503        // NULL-valued aggregate row must be filtered out. Otherwise an outer expression
13504        // like `group()` resurrects the group Prometheus considers unseen.
13505        let state = build_query_engine_state();
13506        for query in ["min(lh)", "group(min(lh))"] {
13507            let plan = PromPlanner::stmt_to_plan(
13508                operator_table_provider(),
13509                &operator_eval_stmt(query),
13510                &state,
13511            )
13512            .await
13513            .unwrap();
13514            let (_, batches) = execute(plan, &state).await;
13515            assert_eq!(
13516                batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
13517                0,
13518                "{query}"
13519            );
13520        }
13521    }
13522
13523    #[tokio::test]
13524    async fn test_mixed_min_drops_histogram_only_group() {
13525        // With alternative float/histogram fields, `min by (tag)` keeps float-only groups
13526        // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of
13527        // emitting a NULL-valued row for them.
13528        let state = build_query_engine_state();
13529        let plan = PromPlanner::stmt_to_plan(
13530            operator_table_provider(),
13531            &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"),
13532            &state,
13533        )
13534        .await
13535        .unwrap();
13536        let float_field = plan
13537            .schema()
13538            .fields()
13539            .iter()
13540            .find(|field| field.data_type() == &ArrowDataType::Float64)
13541            .unwrap()
13542            .name()
13543            .clone();
13544        let (_, batches) = execute(plan, &state).await;
13545        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13546        assert_eq!(values(&batches, &float_field), vec![2.0]);
13547    }
13548
13549    #[tokio::test]
13550    async fn test_mixed_or_can_feed_another_or() {
13551        let state = build_query_engine_state();
13552        let plan = PromPlanner::stmt_to_plan(
13553            operator_table_provider(),
13554            &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"),
13555            &state,
13556        )
13557        .await
13558        .unwrap();
13559        let float_field = plan
13560            .schema()
13561            .fields()
13562            .iter()
13563            .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX))
13564            .unwrap()
13565            .name()
13566            .clone();
13567        let histogram_field = plan
13568            .schema()
13569            .fields()
13570            .iter()
13571            .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13572            .unwrap()
13573            .name()
13574            .clone();
13575
13576        let (_, batches) = execute(plan, &state).await;
13577        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
13578        let mut float_values = values(&batches, &float_field);
13579        float_values.sort_by(f64::total_cmp);
13580        assert_eq!(float_values, vec![2.0, 7.0]);
13581        assert_eq!(histograms(&batches, &histogram_field).len(), 1);
13582    }
13583
13584    #[tokio::test]
13585    async fn test_mixed_fields_align_with_single_float_vector() {
13586        let (planner, mixed) = mixed_direct_or(false).await;
13587        let scale = tagged_source(
13588            "scale",
13589            false,
13590            ("k", Some("float")),
13591            DirectOrValue::Float64(2.0),
13592        );
13593        let scale = scan(&scale);
13594        let scale_fields = vec!["v".to_string()];
13595        let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else {
13596            unreachable!()
13597        };
13598
13599        let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
13600            mixed.schema(),
13601            scale.schema(),
13602            &planner.ctx.field_columns,
13603            &scale_fields,
13604            binary.op,
13605            false,
13606            false,
13607        );
13608        assert!(invalid_pairs.is_empty());
13609        assert_eq!(
13610            groups
13611                .iter()
13612                .map(|(output, _)| output.clone())
13613                .collect::<Vec<_>>(),
13614            planner.ctx.field_columns
13615        );
13616        assert_eq!(groups.len(), 2);
13617        assert!(
13618            groups
13619                .iter()
13620                .flat_map(|(_, pairs)| pairs)
13621                .all(|(_, right)| *right == &scale_fields[0])
13622        );
13623
13624        let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns(
13625            scale.schema(),
13626            mixed.schema(),
13627            &scale_fields,
13628            &planner.ctx.field_columns,
13629            binary.op,
13630            false,
13631            false,
13632        );
13633        assert!(invalid_pairs.is_empty());
13634        assert_eq!(
13635            groups
13636                .iter()
13637                .map(|(output, _)| output.clone())
13638                .collect::<Vec<_>>(),
13639            planner.ctx.field_columns
13640        );
13641        assert_eq!(groups.len(), 2);
13642        assert!(
13643            groups
13644                .iter()
13645                .flat_map(|(_, pairs)| pairs)
13646                .all(|(left, _)| *left == &scale_fields[0])
13647        );
13648    }
13649
13650    #[tokio::test]
13651    async fn test_non_bool_comparison_filters_mixed_sample_lanes() {
13652        let (planner, input) = mixed_direct_or(false).await;
13653        let input_schema = input.schema().clone();
13654        let plan = planner
13655            .filter_on_field_column(input, |field| {
13656                if PromPlanner::field_column_is_native_histogram(&input_schema, field) {
13657                    Ok(lit(false))
13658                } else {
13659                    Ok(col(field).gt(lit(0.0)))
13660                }
13661            })
13662            .unwrap();
13663        let float_field = planner.ctx.field_columns[0].clone();
13664
13665        let (_, batches) = execute(plan, &build_query_engine_state()).await;
13666        assert_eq!(values(&batches, &float_field), vec![1.25]);
13667        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
13668    }
13669
13670    #[tokio::test]
13671    async fn test_mixed_left_and_unless_preserve_sample_lanes() {
13672        for (expression, expected_sample_kind) in [
13673            ("lhs and on(k) mask", (false, true)),
13674            ("lhs unless on(k) mask", (true, false)),
13675        ] {
13676            let (mut planner, left) = mixed_direct_or(false).await;
13677            let left_context = planner.ctx.clone();
13678            let float_field = left_context.field_columns[0].clone();
13679            let histogram_field = left_context.field_columns[1].clone();
13680            let mask = tagged_source(
13681                "mask",
13682                false,
13683                ("k", Some("histogram")),
13684                DirectOrValue::Float64(1.0),
13685            );
13686            let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else {
13687                unreachable!()
13688            };
13689            let plan = planner
13690                .set_op_on_non_field_columns(
13691                    left,
13692                    scan(&mask),
13693                    left_context,
13694                    direct_or_context("mask", &["job", "k"], "v"),
13695                    binary.op,
13696                    &binary.modifier,
13697                )
13698                .unwrap();
13699
13700            let (_, batches) = execute(plan, &build_query_engine_state()).await;
13701            let sample_kinds = batches
13702                .iter()
13703                .flat_map(|batch| {
13704                    let floats = batch.column_by_name(&float_field).unwrap();
13705                    let histograms = batch.column_by_name(&histogram_field).unwrap();
13706                    (0..batch.num_rows())
13707                        .map(|row| (floats.is_valid(row), histograms.is_valid(row)))
13708                })
13709                .collect::<Vec<_>>();
13710            assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}");
13711        }
13712    }
13713
13714    #[tokio::test]
13715    async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() {
13716        let plan = PromPlanner::stmt_to_plan(
13717            build_test_mixed_native_histogram_table_provider("some_metric").await,
13718            &build_eval_stmt("some_metric * scalar(vector(2))"),
13719            &build_query_engine_state(),
13720        )
13721        .await
13722        .unwrap();
13723        let schema = plan.schema();
13724        assert_eq!(
13725            schema
13726                .field_with_unqualified_name(greptime_value())
13727                .unwrap()
13728                .data_type(),
13729            &ArrowDataType::Float64
13730        );
13731        assert_eq!(
13732            schema
13733                .field_with_unqualified_name(greptime_native_histogram())
13734                .unwrap()
13735                .data_type(),
13736            &native_histogram_value_type().as_arrow_type()
13737        );
13738        assert!(
13739            plan.display_indent_schema()
13740                .to_string()
13741                .contains("prom_native_histogram_mul_scalar"),
13742            "{plan:?}"
13743        );
13744    }
13745
13746    #[tokio::test]
13747    async fn test_unsupported_histogram_binary_does_not_block_or_fallback() {
13748        let state = build_query_engine_state();
13749        let plan = PromPlanner::stmt_to_plan(
13750            operator_table_provider(),
13751            &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"),
13752            &state,
13753        )
13754        .await
13755        .unwrap();
13756        let float_field = plan
13757            .schema()
13758            .fields()
13759            .iter()
13760            .find(|field| field.data_type() == &ArrowDataType::Float64)
13761            .unwrap()
13762            .name()
13763            .clone();
13764        let histogram_field = plan
13765            .schema()
13766            .fields()
13767            .iter()
13768            .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type())
13769            .unwrap()
13770            .name()
13771            .clone();
13772
13773        let (_, batches) = execute(plan, &state).await;
13774        assert_eq!(values(&batches, &float_field), vec![0.0]);
13775        assert_eq!(histograms(&batches, &histogram_field).len(), 1);
13776    }
13777
13778    #[tokio::test]
13779    async fn test_unary_negates_mixed_float_and_native_histogram_samples() {
13780        for histogram_on_left in [false, true] {
13781            let (mut planner, input) = mixed_direct_or(histogram_on_left).await;
13782            let plan = planner.negate_field_columns(input).unwrap();
13783            assert!(PromPlanner::field_columns_are_alternative_samples(
13784                plan.schema(),
13785                &planner.ctx.field_columns
13786            ));
13787            let float_field = planner
13788                .ctx
13789                .field_columns
13790                .iter()
13791                .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX))
13792                .unwrap();
13793            let histogram_field = planner
13794                .ctx
13795                .field_columns
13796                .iter()
13797                .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX))
13798                .unwrap();
13799
13800            let (_, batches) = execute(plan, &build_query_engine_state()).await;
13801            assert_eq!(values(&batches, float_field), vec![-1.25]);
13802            let histogram = batches
13803                .iter()
13804                .find_map(|batch| {
13805                    let values = batch
13806                        .column_by_name(histogram_field)
13807                        .unwrap()
13808                        .as_any()
13809                        .downcast_ref::<datafusion::arrow::array::StructArray>()
13810                        .unwrap();
13811                    (0..values.len()).find_map(|row| {
13812                        common_query::native_histogram::read_histogram(values, row).unwrap()
13813                    })
13814                })
13815                .unwrap();
13816            assert_eq!(histogram.count, -1.0);
13817            assert_eq!(histogram.sum, -1.0);
13818            assert_eq!(histogram.reset_hint, CounterResetHint::Gauge);
13819        }
13820    }
13821
13822    #[tokio::test]
13823    async fn test_native_histogram_sum_and_avg_execute_real_batches() {
13824        for op_name in ["sum", "avg"] {
13825            for incompatible in [false, true] {
13826                let mut second = direct_or_histogram();
13827                if incompatible {
13828                    second.schema = CUSTOM_BUCKETS_SCHEMA;
13829                    second.custom_values = vec![1.0];
13830                }
13831                let collector = PromqlAnnotationCollector::default();
13832                let (mut planner, input) =
13833                    mixed_aggregate_input(vec![direct_or_histogram(), second]).await;
13834                planner.promql_annotations = Some(collector.clone());
13835                let histogram_column = planner.ctx.field_columns[1].clone();
13836                planner.ctx.field_columns = vec![histogram_column.clone()];
13837                let input = LogicalPlanBuilder::from(input)
13838                    .project([col("ts"), col(&histogram_column)])
13839                    .unwrap()
13840                    .build()
13841                    .unwrap();
13842                let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13843                    parser::parse(&format!("{op_name}(mixed)")).unwrap()
13844                else {
13845                    unreachable!()
13846                };
13847                let (aggregate_exprs, _) =
13848                    planner.create_aggregate_exprs(op, &param, &input).unwrap();
13849                let plan = LogicalPlanBuilder::from(input)
13850                    .aggregate(vec![col("ts")], aggregate_exprs)
13851                    .unwrap()
13852                    .filter(planner.create_empty_values_filter_expr(false).unwrap())
13853                    .unwrap()
13854                    .build()
13855                    .unwrap();
13856
13857                let (_, batches) = execute(plan, &build_query_engine_state()).await;
13858                let mut warnings = vec![];
13859                let mut infos = vec![];
13860                collector.append_to(&mut warnings, &mut infos);
13861                assert!(infos.is_empty());
13862                if incompatible {
13863                    assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
13864                    assert!(warnings.iter().any(|warning| {
13865                        warning
13866                            == &format!(
13867                                "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas"
13868                            )
13869                    }));
13870                } else {
13871                    let histograms = histograms(&batches, &histogram_column);
13872                    assert_eq!(histograms.len(), 1);
13873                    let expected = if op_name == "sum" { 2.0 } else { 1.0 };
13874                    assert_eq!(histograms[0].count, expected);
13875                    assert_eq!(histograms[0].sum, expected);
13876                    assert!(warnings.is_empty());
13877                }
13878            }
13879        }
13880    }
13881
13882    #[tokio::test]
13883    async fn test_canonical_mixed_count_group_and_count_values_execute() {
13884        let state = build_query_engine_state();
13885        for (query, expected) in [
13886            ("count(some_metric)", vec![2.0]),
13887            ("group(some_metric)", vec![1.0]),
13888            (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]),
13889        ] {
13890            let plan = PromPlanner::stmt_to_plan(
13891                build_test_mixed_native_histogram_table_provider("some_metric").await,
13892                &operator_eval_stmt(query),
13893                &state,
13894            )
13895            .await
13896            .unwrap();
13897            assert!(
13898                plan.schema()
13899                    .fields()
13900                    .iter()
13901                    .all(|field| !field.name().starts_with("__promql_sample_count")),
13902                "{query}: {plan:?}"
13903            );
13904            let value_fields = plan
13905                .schema()
13906                .fields()
13907                .iter()
13908                .filter(|field| {
13909                    matches!(
13910                        field.data_type(),
13911                        ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64
13912                    ) || field.data_type() == &native_histogram_value_type().as_arrow_type()
13913                })
13914                .collect::<Vec<_>>();
13915            assert_eq!(value_fields.len(), 1, "{query}: {plan:?}");
13916            assert_ne!(
13917                value_fields[0].data_type(),
13918                &native_histogram_value_type().as_arrow_type(),
13919                "{query}: {plan:?}"
13920            );
13921            let value_column = value_fields[0].name().clone();
13922
13923            let (_, batches) = execute(plan, &state).await;
13924            let mut actual = numeric_values(&batches, &value_column);
13925            actual.sort_by(f64::total_cmp);
13926            assert_eq!(actual, expected, "{query}");
13927
13928            if query.starts_with("count_values") {
13929                let mut sample_labels = batches
13930                    .iter()
13931                    .flat_map(|batch| {
13932                        batch
13933                            .column_by_name("sample")
13934                            .unwrap()
13935                            .as_any()
13936                            .downcast_ref::<StringArray>()
13937                            .unwrap()
13938                            .iter()
13939                            .flatten()
13940                            .map(str::to_string)
13941                    })
13942                    .collect::<Vec<_>>();
13943                sample_labels.sort();
13944                let mut expected_labels =
13945                    vec!["2".to_string(), direct_or_histogram().promql_string()];
13946                expected_labels.sort();
13947                assert_eq!(sample_labels, expected_labels);
13948            }
13949        }
13950    }
13951
13952    #[tokio::test]
13953    async fn test_mixed_or_sum_aggregates_each_sample_type() {
13954        let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
13955            parser::parse("sum(lhs)").unwrap()
13956        else {
13957            unreachable!()
13958        };
13959
13960        let collector = PromqlAnnotationCollector::default();
13961        let (mut planner, input) = mixed_direct_or(false).await;
13962        planner.promql_annotations = Some(collector.clone());
13963        let float_column = planner.ctx.field_columns[0].clone();
13964        let histogram_column = planner.ctx.field_columns[1].clone();
13965        let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
13966        let plan = LogicalPlanBuilder::from(input)
13967            .aggregate(vec![col("ts"), col("k")], aggregate_exprs)
13968            .unwrap()
13969            .filter(
13970                planner
13971                    .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
13972                    .unwrap(),
13973            )
13974            .unwrap()
13975            .project([
13976                col(&float_column),
13977                col(&histogram_column),
13978                col("ts"),
13979                col("k"),
13980            ])
13981            .unwrap()
13982            .build()
13983            .unwrap();
13984
13985        let (_, batches) = execute(plan, &build_query_engine_state()).await;
13986        assert_eq!(values(&batches, &float_column), vec![1.25]);
13987        let histogram = batches
13988            .iter()
13989            .find_map(|batch| {
13990                let values = batch
13991                    .column_by_name(&histogram_column)?
13992                    .as_any()
13993                    .downcast_ref::<datafusion::arrow::array::StructArray>()?;
13994                (0..values.len()).find_map(|row| {
13995                    common_query::native_histogram::read_histogram(values, row).unwrap()
13996                })
13997            })
13998            .unwrap();
13999        assert_eq!(histogram.count, 1.0);
14000        let mut warnings = vec![];
14001        let mut infos = vec![];
14002        collector.append_to(&mut warnings, &mut infos);
14003        assert!(warnings.is_empty());
14004
14005        let collector = PromqlAnnotationCollector::default();
14006        let (mut planner, input) = mixed_direct_or(false).await;
14007        planner.promql_annotations = Some(collector.clone());
14008        let float_column = planner.ctx.field_columns[0].clone();
14009        let histogram_column = planner.ctx.field_columns[1].clone();
14010        let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
14011        let plan = LogicalPlanBuilder::from(input)
14012            .aggregate(vec![col("ts")], aggregate_exprs)
14013            .unwrap()
14014            .filter(
14015                planner
14016                    .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
14017                    .unwrap(),
14018            )
14019            .unwrap()
14020            .project([col(&float_column), col(&histogram_column), col("ts")])
14021            .unwrap()
14022            .build()
14023            .unwrap();
14024
14025        let (_, batches) = execute(plan, &build_query_engine_state()).await;
14026        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
14027        let mut warnings = vec![];
14028        let mut infos = vec![];
14029        collector.append_to(&mut warnings, &mut infos);
14030        assert_eq!(
14031            warnings,
14032            vec![
14033                "sum: dropped aggregation result containing both float and native histogram samples"
14034            ]
14035        );
14036    }
14037
14038    #[tokio::test]
14039    async fn test_mixed_or_sum_drops_incompatible_mixed_group() {
14040        let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14041            parser::parse("sum(lhs)").unwrap()
14042        else {
14043            unreachable!()
14044        };
14045        let mut custom = direct_or_histogram();
14046        custom.schema = CUSTOM_BUCKETS_SCHEMA;
14047        custom.custom_values = vec![1.0];
14048        let collector = PromqlAnnotationCollector::default();
14049        let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await;
14050        planner.promql_annotations = Some(collector.clone());
14051        let float_column = planner.ctx.field_columns[0].clone();
14052        let histogram_column = planner.ctx.field_columns[1].clone();
14053        let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
14054        let plan = LogicalPlanBuilder::from(input)
14055            .aggregate(vec![col("ts")], aggregate_exprs)
14056            .unwrap()
14057            .filter(
14058                planner
14059                    .mixed_aggregate_filter_expr(op, &float_column, &histogram_column)
14060                    .unwrap(),
14061            )
14062            .unwrap()
14063            .project([col(&float_column), col(&histogram_column), col("ts")])
14064            .unwrap()
14065            .build()
14066            .unwrap();
14067
14068        let (_, batches) = execute(plan, &build_query_engine_state()).await;
14069        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
14070        let mut warnings = vec![];
14071        let mut infos = vec![];
14072        collector.append_to(&mut warnings, &mut infos);
14073        assert!(warnings.iter().any(|warning| {
14074            warning
14075                == "sum: dropped aggregation result containing both float and native histogram samples"
14076        }));
14077    }
14078
14079    #[tokio::test]
14080    async fn test_mixed_or_min_records_only_present_histograms() {
14081        let PromExpr::Aggregate(AggregateExpr { op, param, .. }) =
14082            parser::parse("min(lhs)").unwrap()
14083        else {
14084            unreachable!()
14085        };
14086        let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms";
14087
14088        for (histograms, expected_infos) in [
14089            (vec![], vec![]),
14090            (vec![direct_or_histogram()], vec![expected_info]),
14091        ] {
14092            let collector = PromqlAnnotationCollector::default();
14093            let (mut planner, input) = mixed_aggregate_input(histograms).await;
14094            planner.promql_annotations = Some(collector.clone());
14095            let float_column = planner.ctx.field_columns[0].clone();
14096            let histogram_column = planner.ctx.field_columns[1].clone();
14097            let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, &param, &input).unwrap();
14098            let plan = LogicalPlanBuilder::from(input)
14099                .aggregate(vec![col("ts")], aggregate_exprs)
14100                .unwrap()
14101                .filter(
14102                    planner
14103                        .mixed_ignored_histogram_filter_expr(op, &histogram_column)
14104                        .unwrap(),
14105                )
14106                .unwrap()
14107                .project([col(&float_column), col("ts")])
14108                .unwrap()
14109                .build()
14110                .unwrap();
14111
14112            let (_, batches) = execute(plan, &build_query_engine_state()).await;
14113            assert_eq!(values(&batches, &float_column), vec![1.25]);
14114            let mut warnings = vec![];
14115            let mut infos = vec![];
14116            collector.append_to(&mut warnings, &mut infos);
14117            assert!(warnings.is_empty());
14118            assert_eq!(infos, expected_infos);
14119        }
14120    }
14121
14122    #[tokio::test]
14123    async fn test_mixed_or_value_aliases_do_not_replace_labels() {
14124        let left = source(
14125            "lhs",
14126            false,
14127            1,
14128            vec![("job", Some("job")), ("k", Some("float"))],
14129            DirectOrValue::Float64(1.0),
14130        );
14131        let right = source(
14132            "rhs",
14133            false,
14134            1,
14135            vec![
14136                ("job", Some("job")),
14137                ("k", Some("histogram")),
14138                (greptime_value(), Some("value-label")),
14139            ],
14140            DirectOrValue::NativeHistogram(direct_or_histogram()),
14141        );
14142        let table_provider = build_test_table_provider_with_fields(
14143            &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())],
14144            &[],
14145        )
14146        .await;
14147        let mut planner = PromPlanner {
14148            table_provider,
14149            ctx: PromPlannerContext::default(),
14150            promql_annotations: None,
14151        };
14152        let left = LogicalPlanBuilder::from(scan(&left))
14153            .project(vec![
14154                col("ts"),
14155                col("job"),
14156                col("k"),
14157                col("v").alias(greptime_value()),
14158            ])
14159            .unwrap()
14160            .build()
14161            .unwrap();
14162        let left_context = direct_or_context("lhs", &["job", "k"], greptime_value());
14163        let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v");
14164        let plan = planner
14165            .or_operator(
14166                left,
14167                scan(&right),
14168                left_context.tag_columns.iter().cloned().collect(),
14169                right_context.tag_columns.iter().cloned().collect(),
14170                left_context,
14171                right_context,
14172                &or_modifier("lhs or on(k) rhs"),
14173            )
14174            .unwrap();
14175
14176        assert_eq!(
14177            plan.schema()
14178                .field_with_name(None, greptime_value())
14179                .unwrap()
14180                .data_type(),
14181            &ArrowDataType::Utf8
14182        );
14183        assert!(
14184            planner
14185                .ctx
14186                .field_columns
14187                .iter()
14188                .all(|field| { field != greptime_value() && field != greptime_native_histogram() })
14189        );
14190        assert!(PromPlanner::field_columns_are_alternative_samples(
14191            plan.schema(),
14192            &planner.ctx.field_columns
14193        ));
14194        let (_, batches) = execute(plan, &build_query_engine_state()).await;
14195        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
14196        let labels = batches
14197            .iter()
14198            .flat_map(|batch| {
14199                batch
14200                    .column_by_name(greptime_value())
14201                    .unwrap()
14202                    .as_any()
14203                    .downcast_ref::<StringArray>()
14204                    .unwrap()
14205                    .iter()
14206                    .flatten()
14207            })
14208            .collect::<Vec<_>>();
14209        assert_eq!(labels, vec!["value-label"]);
14210    }
14211
14212    #[tokio::test]
14213    async fn test_mixed_or_routes_float_histogram_and_label_functions() {
14214        for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] {
14215            let (mut planner, input) = mixed_direct_or(false).await;
14216            let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14217                input.schema(),
14218                &planner.ctx.field_columns,
14219            );
14220            let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else {
14221                unreachable!()
14222            };
14223            let state = build_query_engine_state();
14224            let (mut exprs, _) = planner
14225                .create_function_expr(&call.func, vec![], input.schema(), &state)
14226                .unwrap();
14227            exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14228            exprs.extend(planner.create_tag_column_exprs().unwrap());
14229            let plan = LogicalPlanBuilder::from(input)
14230                .project(exprs)
14231                .unwrap()
14232                .filter(
14233                    planner
14234                        .create_empty_values_filter_expr(preserve_any_value)
14235                        .unwrap(),
14236                )
14237                .unwrap()
14238                .build()
14239                .unwrap();
14240            let (_, batches) = execute(plan, &state).await;
14241            let values = batches
14242                .iter()
14243                .flat_map(|batch| {
14244                    batch
14245                        .schema()
14246                        .fields()
14247                        .iter()
14248                        .position(|field| field.data_type() == &ArrowDataType::Float64)
14249                        .map(|index| {
14250                            batch
14251                                .column(index)
14252                                .as_any()
14253                                .downcast_ref::<Float64Array>()
14254                                .unwrap()
14255                                .iter()
14256                                .flatten()
14257                        })
14258                        .into_iter()
14259                        .flatten()
14260                })
14261                .collect::<Vec<_>>();
14262            assert_eq!(values, vec![expected], "{function}");
14263        }
14264
14265        let (mut planner, input) = mixed_direct_or(false).await;
14266        let preserve_any_value = PromPlanner::field_columns_are_alternative_samples(
14267            input.schema(),
14268            &planner.ctx.field_columns,
14269        );
14270        let PromExpr::Call(call) =
14271            parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap()
14272        else {
14273            unreachable!()
14274        };
14275        let args = planner.create_function_args(&call.args.args).unwrap();
14276        let state = build_query_engine_state();
14277        let (mut exprs, _) = planner
14278            .create_function_expr(&call.func, args.literals, input.schema(), &state)
14279            .unwrap();
14280        exprs.insert(0, planner.create_time_index_column_expr().unwrap());
14281        exprs.extend(planner.create_tag_column_exprs().unwrap());
14282        let plan = LogicalPlanBuilder::from(input)
14283            .project(exprs)
14284            .unwrap()
14285            .filter(
14286                planner
14287                    .create_empty_values_filter_expr(preserve_any_value)
14288                    .unwrap(),
14289            )
14290            .unwrap()
14291            .build()
14292            .unwrap();
14293        let (_, batches) = execute(plan, &state).await;
14294        let sample_count = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
14295        assert_eq!(sample_count, 2);
14296    }
14297}