query/optimizer/
scan_hint.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::HashSet;
16
17use api::v1::SemanticType;
18use arrow_schema::SortOptions;
19use common_function::aggrs::aggr_wrapper::aggr_state_func_name;
20use common_recordbatch::OrderOption;
21use common_recordbatch::filter::SimpleFilterEvaluator;
22use datafusion::datasource::DefaultTableSource;
23use datafusion_common::tree_node::{Transformed, TreeNodeRewriter};
24use datafusion_common::{Column, Result};
25use datafusion_expr::expr::Sort;
26use datafusion_expr::{Expr, LogicalPlan, utils};
27use datafusion_optimizer::{OptimizerConfig, OptimizerRule};
28use datatypes::arrow::datatypes::{DataType, TimeUnit as ArrowTimeUnit};
29use promql::extension_plan::{InstantManipulate, SeriesDivide, SeriesNormalize};
30use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
31use store_api::storage::{TimeSeriesDistribution, TimeSeriesRowSelector};
32
33use crate::dummy_catalog::DummyTableProvider;
34
35/// This rule will traverse the plan to collect necessary hints for leaf
36/// table scan node and set them in [`ScanRequest`]. Hints include:
37/// - the nearest order requirement to the leaf table scan node as ordering hint.
38/// - the group by columns when all aggregate functions are `last_value` as
39///   time series row selector hint.
40///
41/// [`ScanRequest`]: store_api::storage::ScanRequest
42#[derive(Debug)]
43pub struct ScanHintRule;
44
45impl OptimizerRule for ScanHintRule {
46    fn name(&self) -> &str {
47        "ScanHintRule"
48    }
49
50    fn rewrite(
51        &self,
52        plan: LogicalPlan,
53        _config: &dyn OptimizerConfig,
54    ) -> Result<Transformed<LogicalPlan>> {
55        Self::optimize(plan)
56    }
57}
58
59impl ScanHintRule {
60    fn optimize(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
61        let mut rewriter = ScanHintRewriter::default();
62        // The extension's input is included by DataFusion's normal TreeNode
63        // rewrite, so this is one scoped recursive walk (subquery expressions
64        // are included by the dedicated API as well).
65        plan.rewrite_with_subqueries(&mut rewriter)
66    }
67
68    fn set_hints(
69        plan: LogicalPlan,
70        rewriter: &mut ScanHintRewriter,
71    ) -> Result<Transformed<LogicalPlan>> {
72        let LogicalPlan::TableScan(mut table_scan) = plan else {
73            return Ok(Transformed::no(plan));
74        };
75        let Some(source) = table_scan.source.downcast_ref::<DefaultTableSource>() else {
76            return Ok(Transformed::no(LogicalPlan::TableScan(table_scan)));
77        };
78        // The provider in the region server is [DummyTableProvider].
79        let Some(original) = source.table_provider.downcast_ref::<DummyTableProvider>() else {
80            return Ok(Transformed::no(LogicalPlan::TableScan(table_scan)));
81        };
82
83        // Attached scan filters are checked below; residual Filter nodes are
84        // rejected by the single-evaluation path allowlist.
85        let filters_preserve_last_row = if rewriter.inside_single_evaluation {
86            Self::filters_preserve_last_row(&table_scan, original)
87        } else {
88            true
89        };
90        let use_last_row = rewriter.inside_single_evaluation && filters_preserve_last_row;
91
92        let has_hint =
93            rewriter.order_expr.is_some() || rewriter.ts_row_selector.is_some() || use_last_row;
94        if !has_hint {
95            return Ok(Transformed::no(LogicalPlan::TableScan(table_scan)));
96        }
97
98        // A provider can be used by several TableScan nodes. Fork its request
99        // for every hinted use-site before applying hints, rather than mutating
100        // the shared catalog provider. This keeps order and aggregate hints
101        // local as well as the new LastRow hint.
102        let adapter = original.clone_for_scan();
103        Self::apply_hints(&adapter, rewriter);
104        if use_last_row {
105            // Apply the instant-derived hint after the aggregate hint. Both
106            // select LastRow today, and this ordering preserves the existing
107            // aggregate selector when the instant guard rejects a scan.
108            adapter.with_time_series_selector_hint(TimeSeriesRowSelector::LastRow {
109                after_merge: true,
110            });
111        }
112        table_scan.source =
113            std::sync::Arc::new(DefaultTableSource::new(std::sync::Arc::new(adapter)));
114        Ok(Transformed::yes(LogicalPlan::TableScan(table_scan)))
115    }
116
117    /// Checks whether attached scan predicates permit instant-derived LastRow selection.
118    ///
119    /// Selecting the newest row too early can discard an older matching sample if a
120    /// predicate later rejects that row. Only recognized tag/time predicates are
121    /// allowed: tags select whole series, and supported time predicates constrain
122    /// the scan window before row selection. Field or unrecognized predicates are
123    /// conservatively rejected.
124    ///
125    /// This checks only attached predicates; the path allowlist separately rejects
126    /// residual Filter nodes between InstantManipulate and the scan.
127    fn filters_preserve_last_row(
128        table_scan: &datafusion_expr::logical_plan::TableScan,
129        provider: &DummyTableProvider,
130    ) -> bool {
131        let metadata = provider.region_metadata();
132        for filter in &table_scan.filters {
133            let Some(filter) = SimpleFilterEvaluator::try_new(filter) else {
134                return false;
135            };
136            let Some(column_metadata) = metadata.column_by_name(filter.column_name()) else {
137                return false;
138            };
139            if !matches!(
140                column_metadata.semantic_type,
141                SemanticType::Tag | SemanticType::Timestamp
142            ) {
143                return false;
144            }
145        }
146        true
147    }
148
149    fn apply_hints(adapter: &DummyTableProvider, rewriter: &mut ScanHintRewriter) {
150        if let Some(order_expr) = &rewriter.order_expr {
151            Self::set_order_hint(adapter, order_expr);
152        }
153        if let Some((group_by_cols, order_by_col)) = &rewriter.ts_row_selector {
154            Self::set_time_series_row_selector_hint(adapter, group_by_cols, order_by_col);
155        }
156    }
157
158    fn set_order_hint(adapter: &DummyTableProvider, order_expr: &Vec<Sort>) {
159        let mut opts = Vec::with_capacity(order_expr.len());
160        for sort in order_expr {
161            let name = match sort.expr.try_as_col() {
162                Some(col) => col.name.clone(),
163                None => return,
164            };
165            opts.push(OrderOption {
166                name,
167                options: SortOptions {
168                    descending: !sort.asc,
169                    nulls_first: sort.nulls_first,
170                },
171            });
172        }
173        adapter.with_ordering_hint(&opts);
174
175        let region_metadata = adapter.region_metadata();
176        let time_index_name = region_metadata
177            .time_index_column()
178            .column_schema
179            .name
180            .as_str();
181        let sort_cols = order_expr
182            .iter()
183            .filter_map(|s| s.expr.try_as_col())
184            .collect::<Vec<_>>();
185
186        // Special-case metric engine: when the nearest sort requirement is `__tsid, <time index>`,
187        // we can safely enable per-series distribution hint so the region can use `SeriesScan`.
188        //
189        // This pattern is produced by promql planning when `__tsid` is available and is used as the
190        // series identifier (instead of expanding to all tag columns).
191        if sort_cols.len() == 2
192            && sort_cols[0].name == DATA_SCHEMA_TSID_COLUMN_NAME
193            && sort_cols[1].name == time_index_name
194        {
195            adapter.with_distribution(TimeSeriesDistribution::PerSeries);
196            return;
197        }
198
199        let mut sort_expr_cursor = sort_cols.into_iter();
200        // ignore table without pk
201        if region_metadata.primary_key.is_empty() {
202            return;
203        }
204        let mut pk_column_iter = region_metadata.primary_key_columns();
205        let mut curr_sort_expr = sort_expr_cursor.next();
206        let mut curr_pk_col = pk_column_iter.next();
207
208        while let (Some(sort_expr), Some(pk_col)) = (curr_sort_expr, curr_pk_col) {
209            if sort_expr.name == pk_col.column_schema.name {
210                curr_sort_expr = sort_expr_cursor.next();
211                curr_pk_col = pk_column_iter.next();
212            } else {
213                return;
214            }
215        }
216
217        let next_remaining = sort_expr_cursor.next();
218        match (curr_sort_expr, next_remaining) {
219            (Some(expr), None)
220                if expr.name == region_metadata.time_index_column().column_schema.name =>
221            {
222                adapter.with_distribution(TimeSeriesDistribution::PerSeries);
223            }
224            (None, _) => adapter.with_distribution(TimeSeriesDistribution::PerSeries),
225            (Some(_), _) => {}
226        }
227    }
228
229    fn set_time_series_row_selector_hint(
230        adapter: &DummyTableProvider,
231        group_by_cols: &HashSet<Column>,
232        order_by_col: &Column,
233    ) {
234        let region_metadata = adapter.region_metadata();
235        let mut should_set_selector_hint = true;
236        // check if order_by column is time index
237        if let Some(column_metadata) = region_metadata.column_by_name(&order_by_col.name) {
238            if column_metadata.semantic_type != SemanticType::Timestamp {
239                should_set_selector_hint = false;
240            }
241        } else {
242            should_set_selector_hint = false;
243        }
244
245        // check if all group_by columns are primary key
246        for col in group_by_cols {
247            let Some(column_metadata) = region_metadata.column_by_name(&col.name) else {
248                should_set_selector_hint = false;
249                break;
250            };
251            if column_metadata.semantic_type != SemanticType::Tag {
252                should_set_selector_hint = false;
253                break;
254            }
255        }
256
257        if should_set_selector_hint {
258            adapter.with_time_series_selector_hint(TimeSeriesRowSelector::LastRow {
259                after_merge: false,
260            });
261        }
262    }
263}
264
265/// Traverse and apply hints with state scoped to the current logical-plan path.
266///
267/// Rewriting the scan while walking down the tree is important: the state then
268/// describes the actual parent path of that scan, and a shared provider is forked
269/// at that exact use-site. No traversal-order identity is involved.
270#[derive(Default)]
271struct ScanHintRewriter {
272    order_expr: Option<Vec<Sort>>,
273    order_stack: Vec<Option<Vec<Sort>>>,
274    ts_row_selector: Option<(HashSet<Column>, Column)>,
275    ts_stack: Vec<Option<(HashSet<Column>, Column)>>,
276    inside_single_evaluation: bool,
277    single_evaluation_stack: Vec<bool>,
278}
279
280impl TreeNodeRewriter for ScanHintRewriter {
281    type Node = LogicalPlan;
282
283    fn f_down(&mut self, node: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
284        self.order_stack.push(self.order_expr.clone());
285        self.ts_stack.push(self.ts_row_selector.clone());
286        self.single_evaluation_stack
287            .push(self.inside_single_evaluation);
288
289        if let LogicalPlan::Sort(sort) = &node {
290            self.order_expr = Some(sort.expr.clone());
291        }
292        if let LogicalPlan::Extension(extension) = &node
293            && let Some(instant) = extension.node.as_any().downcast_ref::<InstantManipulate>()
294        {
295            self.inside_single_evaluation = instant.is_single_evaluation();
296        } else if self.inside_single_evaluation {
297            // This allowlist is coupled to the controlled PromQL planner. It
298            // permits only nodes known to preserve the newest row per series;
299            // every other node is a sticky boundary until a nested instant
300            // extension establishes a new evaluation scope.
301            self.inside_single_evaluation = single_evaluation_node_allowed(&node);
302        }
303        if let LogicalPlan::Aggregate(aggregate) = &node {
304            self.ts_row_selector = Self::extract_last_value_selector(aggregate);
305        }
306
307        let is_branching = matches!(
308            node,
309            LogicalPlan::Subquery(_) | LogicalPlan::SubqueryAlias(_)
310        ) || node.inputs().len() > 1;
311        if is_branching {
312            self.ts_row_selector = None;
313        }
314        if let LogicalPlan::Filter(filter) = &node
315            && let Some(group_by_exprs) = &self.ts_row_selector
316        {
317            let mut referenced = HashSet::default();
318            utils::expr_to_columns(&filter.predicate, &mut referenced)?;
319            if !referenced.is_subset(&group_by_exprs.0) {
320                self.ts_row_selector = None;
321            }
322        }
323
324        ScanHintRule::set_hints(node, self)
325    }
326
327    fn f_up(&mut self, node: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
328        if let Some(previous) = self.order_stack.pop() {
329            self.order_expr = previous;
330        }
331        if let Some(previous) = self.ts_stack.pop() {
332            self.ts_row_selector = previous;
333        }
334        if let Some(previous) = self.single_evaluation_stack.pop() {
335            self.inside_single_evaluation = previous;
336        }
337        Ok(Transformed::no(node))
338    }
339}
340
341/// Returns whether a plan node can occur on the scan path of a controlled
342/// PromQL single evaluation without changing which row is newest per series.
343fn single_evaluation_node_allowed(node: &LogicalPlan) -> bool {
344    match node {
345        LogicalPlan::TableScan(_) | LogicalPlan::SubqueryAlias(_) => true,
346        LogicalPlan::Sort(sort) => sort.fetch.is_none(),
347        LogicalPlan::Projection(projection) => projection
348            .expr
349            .iter()
350            .all(|expr| single_evaluation_projection_expr_allowed(expr, projection)),
351        LogicalPlan::Extension(extension) => {
352            let extension = extension.node.as_any();
353            extension.is::<SeriesDivide>()
354                || extension
355                    .downcast_ref::<SeriesNormalize>()
356                    .is_some_and(|normalize| !normalize.filter_stale_markers())
357        }
358        _ => false,
359    }
360}
361
362/// This whitelist assumes the planner preserves time-index and series identity;
363/// it is not a proof that an arbitrary plan does so.
364///
365/// Identity projections preserve the selected samples. Only the planner's named
366/// seconds/milliseconds-to-milliseconds casts are accepted; a microsecond or
367/// nanosecond cast could collapse a future sample onto the evaluation boundary.
368fn single_evaluation_projection_expr_allowed(
369    expr: &Expr,
370    projection: &datafusion_expr::logical_plan::Projection,
371) -> bool {
372    match expr {
373        Expr::Column(_) => true,
374        Expr::Alias(alias) => match alias.expr.as_ref() {
375            Expr::Column(column) => alias.name == column.name,
376            Expr::Cast(cast) => {
377                let Expr::Column(column) = cast.expr.as_ref() else {
378                    return false;
379                };
380                alias.name == column.name
381                    && matches!(
382                        cast.field.data_type(),
383                        DataType::Timestamp(ArrowTimeUnit::Millisecond, None)
384                    )
385                    && matches!(
386                        projection
387                            .input
388                            .schema()
389                            .qualified_field_from_column(column),
390                        Ok((_, field))
391                            if matches!(field.data_type(), DataType::Timestamp(ArrowTimeUnit::Second | ArrowTimeUnit::Millisecond, None))
392                    )
393            }
394            _ => false,
395        },
396        _ => false,
397    }
398}
399
400impl ScanHintRewriter {
401    fn extract_last_value_selector(
402        aggregate: &datafusion_expr::logical_plan::Aggregate,
403    ) -> Option<(HashSet<Column>, Column)> {
404        let mut order_by_expr = None;
405        if aggregate.aggr_expr.is_empty() {
406            return None;
407        }
408        for expr in &aggregate.aggr_expr {
409            let Expr::AggregateFunction(func) = expr else {
410                return None;
411            };
412            if (func.func.name() != "last_value"
413                && func.func.name() != aggr_state_func_name("last_value"))
414                || func.params.filter.is_some()
415                || func.params.distinct
416            {
417                return None;
418            }
419            let order_by = &func.params.order_by;
420            if order_by.len() != 1 || !order_by[0].asc {
421                return None;
422            }
423            if let Some(existing) = &order_by_expr {
424                if existing != &order_by[0] {
425                    return None;
426                }
427            } else {
428                order_by_expr = Some(order_by[0].clone());
429            }
430        }
431        let Expr::Column(order_by_col) = order_by_expr?.expr else {
432            return None;
433        };
434        let mut group_by_cols = HashSet::with_capacity(aggregate.group_expr.len());
435        for expr in &aggregate.group_expr {
436            let Expr::Column(col) = expr else {
437                return None;
438            };
439            group_by_cols.insert(col.clone());
440        }
441        Some((group_by_cols, order_by_col))
442    }
443}
444
445#[cfg(test)]
446mod test {
447    use std::collections::HashMap;
448    use std::sync::Arc;
449
450    use datafusion::functions_aggregate::first_last::last_value_udaf;
451    use datafusion::functions_aggregate::min_max::max_udaf;
452    use datafusion::functions_window::row_number::RowNumber;
453    use datafusion::logical_expr::expr::WindowFunction;
454    use datafusion::logical_expr::{WindowFrame, WindowFunctionDefinition};
455    use datafusion::prelude::JoinType;
456    use datafusion_common::tree_node::TreeNodeRecursion;
457    use datafusion_expr::expr::{
458        AggregateFunction, AggregateFunctionParams, Cast, WindowFunctionParams,
459    };
460    use datafusion_expr::expr_fn::scalar_subquery;
461    use datafusion_expr::{Extension, LogicalPlan, LogicalPlanBuilder, col, lit};
462    use datafusion_optimizer::OptimizerContext;
463    use datatypes::arrow::datatypes::DataType;
464    use datatypes::data_type::ConcreteDataType;
465    use datatypes::schema::ColumnSchema;
466    use promql::extension_plan::RangeManipulate;
467    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
468    use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
469    use store_api::storage::{RegionId, TimeSeriesRowSelector};
470
471    use super::*;
472
473    fn scan_requests(plan: &LogicalPlan) -> Vec<store_api::storage::ScanRequest> {
474        scan_requests_with_names(plan)
475            .into_iter()
476            .map(|(_, request)| request)
477            .collect()
478    }
479
480    fn scan_requests_with_names(
481        plan: &LogicalPlan,
482    ) -> Vec<(String, store_api::storage::ScanRequest)> {
483        let mut requests = Vec::new();
484        plan.apply_with_subqueries(|node| {
485            if let LogicalPlan::TableScan(scan) = node
486                && let Some(source) = scan.source.downcast_ref::<DefaultTableSource>()
487                && let Some(provider) = source.table_provider.downcast_ref::<DummyTableProvider>()
488            {
489                requests.push((scan.table_name.to_string(), provider.scan_request()));
490            }
491            Ok(TreeNodeRecursion::Continue)
492        })
493        .unwrap();
494        requests
495    }
496
497    fn instant_plan(provider: Arc<DummyTableProvider>, end: i64) -> LogicalPlan {
498        instant_plan_named(provider, "t", end)
499    }
500
501    fn instant_plan_with_filters(
502        provider: Arc<DummyTableProvider>,
503        filters: Vec<Expr>,
504    ) -> LogicalPlan {
505        let scan = scan_plan(provider, "t");
506        let LogicalPlan::TableScan(mut scan) = scan else {
507            unreachable!();
508        };
509        scan.filters = filters;
510        LogicalPlan::Extension(Extension {
511            node: Arc::new(InstantManipulate::new(
512                1000,
513                1000,
514                1000,
515                1000,
516                0,
517                "ts".to_string(),
518                vec![],
519                Some("v0".to_string()),
520                LogicalPlan::TableScan(scan),
521            )),
522        })
523    }
524
525    fn scan_plan(provider: Arc<DummyTableProvider>, table_name: &str) -> LogicalPlan {
526        LogicalPlanBuilder::scan(
527            table_name,
528            Arc::new(DefaultTableSource::new(provider)),
529            None,
530        )
531        .unwrap()
532        .build()
533        .unwrap()
534    }
535
536    fn mock_table_provider_with_timestamp(
537        region_id: RegionId,
538        timestamp_type: ConcreteDataType,
539    ) -> DummyTableProvider {
540        let mut builder = RegionMetadataBuilder::new(region_id);
541        builder
542            .push_column_metadata(ColumnMetadata {
543                column_schema: ColumnSchema::new("k0", ConcreteDataType::string_datatype(), true),
544                semantic_type: SemanticType::Tag,
545                column_id: 1,
546            })
547            .push_column_metadata(ColumnMetadata {
548                column_schema: ColumnSchema::new("ts", timestamp_type, false),
549                semantic_type: SemanticType::Timestamp,
550                column_id: 2,
551            })
552            .push_column_metadata(ColumnMetadata {
553                column_schema: ColumnSchema::new("v0", ConcreteDataType::float64_datatype(), false),
554                semantic_type: SemanticType::Field,
555                column_id: 3,
556            })
557            .primary_key(vec![1]);
558        let metadata = Arc::new(builder.build().unwrap());
559        let engine = Arc::new(MetaRegionEngine::with_metadata(metadata.clone()));
560        DummyTableProvider::new(region_id, engine, metadata)
561    }
562
563    fn last_value_aggregate(input: LogicalPlan) -> LogicalPlan {
564        let aggregate = LogicalPlanBuilder::from(input)
565            .aggregate(
566                vec![col("k0")],
567                vec![
568                    Expr::AggregateFunction(AggregateFunction {
569                        func: last_value_udaf(),
570                        params: AggregateFunctionParams {
571                            args: vec![col("v0")],
572                            distinct: false,
573                            filter: None,
574                            order_by: vec![Sort {
575                                expr: col("ts"),
576                                asc: true,
577                                nulls_first: true,
578                            }],
579                            null_treatment: None,
580                        },
581                    }),
582                    Expr::AggregateFunction(AggregateFunction {
583                        func: last_value_udaf(),
584                        params: AggregateFunctionParams {
585                            args: vec![col("ts")],
586                            distinct: false,
587                            filter: None,
588                            order_by: vec![Sort {
589                                expr: col("ts"),
590                                asc: true,
591                                nulls_first: true,
592                            }],
593                            null_treatment: None,
594                        },
595                    }),
596                ],
597            )
598            .unwrap()
599            .build()
600            .unwrap();
601        let timestamp = aggregate.schema().field(2).name().clone();
602
603        LogicalPlanBuilder::from(aggregate)
604            .project(vec![col("k0"), col(timestamp).alias("ts")])
605            .unwrap()
606            .build()
607            .unwrap()
608    }
609
610    fn single_evaluation(input: LogicalPlan) -> LogicalPlan {
611        LogicalPlan::Extension(Extension {
612            node: Arc::new(InstantManipulate::new(
613                1000,
614                1000,
615                1000,
616                1000,
617                0,
618                "ts".to_string(),
619                vec![],
620                Some("v0".to_string()),
621                input,
622            )),
623        })
624    }
625
626    fn instant_with_expression_subquery(
627        outer_provider: Arc<DummyTableProvider>,
628        inner_plan: LogicalPlan,
629        outer_end: i64,
630    ) -> LogicalPlan {
631        let outer_scan = scan_plan(outer_provider, "outer");
632        let input = LogicalPlanBuilder::from(outer_scan)
633            .project(vec![
634                col("ts"),
635                col("v0"),
636                scalar_subquery(Arc::new(inner_plan)),
637            ])
638            .unwrap()
639            .build()
640            .unwrap();
641        LogicalPlan::Extension(Extension {
642            node: Arc::new(InstantManipulate::new(
643                1000,
644                outer_end,
645                1000,
646                1000,
647                0,
648                "ts".to_string(),
649                vec![],
650                Some("v0".to_string()),
651                input,
652            )),
653        })
654    }
655
656    fn instant_plan_named(
657        provider: Arc<DummyTableProvider>,
658        table_name: &str,
659        end: i64,
660    ) -> LogicalPlan {
661        let input = LogicalPlanBuilder::scan(
662            table_name,
663            Arc::new(DefaultTableSource::new(provider)),
664            None,
665        )
666        .unwrap()
667        .build()
668        .unwrap();
669        LogicalPlan::Extension(Extension {
670            node: Arc::new(InstantManipulate::new(
671                1000,
672                end,
673                1000,
674                1000,
675                0,
676                "ts".to_string(),
677                vec![],
678                Some("v0".to_string()),
679                input,
680            )),
681        })
682    }
683    use crate::optimizer::test_util::{
684        MetaRegionEngine, mock_table_provider, mock_table_provider_with_tsid,
685    };
686
687    #[test]
688    fn single_evaluation_sets_last_row_on_the_rewritten_scan() {
689        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
690        let plan = instant_plan(provider.clone(), 1000);
691        let rewritten = ScanHintRule
692            .rewrite(plan, &OptimizerContext::default())
693            .unwrap()
694            .data;
695
696        assert_eq!(
697            scan_requests(&rewritten)[0].series_row_selector,
698            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
699        );
700
701        assert_eq!(provider.scan_request().series_row_selector, None);
702    }
703
704    #[test]
705    fn single_evaluation_limit_sort_does_not_set_last_row_below_limit() {
706        let mut selectors = Vec::new();
707        for input in [
708            // Instant(1000, lookback=1000) -> Limit(0, 1) -> Sort(ts ASC) -> Scan.
709            LogicalPlanBuilder::from(scan_plan(
710                Arc::new(mock_table_provider(RegionId::new(1, 1))),
711                "t",
712            ))
713            .sort(vec![col("ts").sort(true, false)])
714            .unwrap()
715            .limit(0, Some(1))
716            .unwrap()
717            .build()
718            .unwrap(),
719            // Sort.fetch is a limit embedded in the Sort node and has the same boundary.
720            LogicalPlanBuilder::from(scan_plan(
721                Arc::new(mock_table_provider(RegionId::new(1, 1))),
722                "t",
723            ))
724            .sort_with_limit(vec![col("ts").sort(true, false)], Some(1))
725            .unwrap()
726            .build()
727            .unwrap(),
728        ] {
729            let rewritten = ScanHintRule
730                .rewrite(single_evaluation(input), &OptimizerContext::default())
731                .unwrap()
732                .data;
733            selectors.push(scan_requests(&rewritten)[0].series_row_selector);
734        }
735
736        // With samples (ts=100, v=10) and (ts=900, v=1), an unhinted ascending
737        // sort/limit pipeline returns 10. LastRow at the scan instead leaves only
738        // (900, 1) before the limit. A LastRow scan hint cannot cross either limit.
739        assert_eq!(selectors, vec![None, None]);
740    }
741
742    #[test]
743    fn single_evaluation_allows_controlled_promql_selector_chain() {
744        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
745        let projection = LogicalPlanBuilder::from(scan_plan(provider, "t"))
746            .project(vec![
747                col("k0").alias("k0"),
748                Expr::Cast(Cast::new(
749                    Box::new(col("ts")),
750                    DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
751                ))
752                .alias("ts"),
753                col("v0"),
754            ])
755            .unwrap()
756            .sort(vec![col("ts").sort(true, false)])
757            .unwrap()
758            .build()
759            .unwrap();
760        let divide = LogicalPlan::Extension(Extension {
761            node: Arc::new(SeriesDivide::new(
762                vec!["k0".to_string()],
763                "ts".to_string(),
764                projection,
765            )),
766        });
767        let normalize = LogicalPlan::Extension(Extension {
768            node: Arc::new(SeriesNormalize::new(
769                42,
770                "ts",
771                false,
772                vec!["k0".to_string()],
773                divide,
774            )),
775        });
776        let rewritten = ScanHintRule
777            .rewrite(single_evaluation(normalize), &OptimizerContext::default())
778            .unwrap()
779            .data;
780
781        assert_eq!(
782            scan_requests(&rewritten)[0].series_row_selector,
783            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
784        );
785    }
786
787    #[test]
788    fn single_evaluation_filtering_normalize_does_not_set_last_row() {
789        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
790        let normalize = LogicalPlan::Extension(Extension {
791            node: Arc::new(SeriesNormalize::new(
792                42,
793                "ts",
794                true,
795                vec!["k0".to_string()],
796                scan_plan(provider, "t"),
797            )),
798        });
799        let rewritten = ScanHintRule
800            .rewrite(single_evaluation(normalize), &OptimizerContext::default())
801            .unwrap()
802            .data;
803
804        // An older finite sample can precede a newest stale marker. LastRow
805        // would discard that sample before SeriesNormalize filters the marker.
806        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
807    }
808
809    #[test]
810    fn single_evaluation_rejects_row_changing_nodes() {
811        let provider = || Arc::new(mock_table_provider(RegionId::new(1, 1)));
812        let window = LogicalPlanBuilder::from(scan_plan(provider(), "window"))
813            .window(vec![Expr::WindowFunction(Box::new(WindowFunction {
814                fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
815                params: WindowFunctionParams {
816                    args: vec![],
817                    partition_by: vec![col("k0")],
818                    order_by: vec![col("ts").sort(true, true)],
819                    window_frame: WindowFrame::new(Some(true)),
820                    filter: None,
821                    null_treatment: None,
822                    distinct: false,
823                },
824            }))])
825            .unwrap()
826            .build()
827            .unwrap();
828        let join = LogicalPlanBuilder::from(scan_plan(provider(), "left"))
829            .join(
830                scan_plan(provider(), "right"),
831                JoinType::Inner,
832                (Vec::<Column>::new(), Vec::<Column>::new()),
833                None,
834            )
835            .unwrap()
836            .project(vec![
837                Expr::Column(Column::new(Some("left"), "ts")),
838                Expr::Column(Column::new(Some("left"), "v0")),
839            ])
840            .unwrap()
841            .build()
842            .unwrap();
843        let nonlast_aggregate = LogicalPlanBuilder::from(scan_plan(provider(), "aggregate"))
844            .aggregate(
845                vec![col("k0")],
846                vec![
847                    Expr::AggregateFunction(AggregateFunction {
848                        func: max_udaf(),
849                        params: AggregateFunctionParams {
850                            args: vec![col("v0")],
851                            distinct: false,
852                            filter: None,
853                            order_by: vec![],
854                            null_treatment: None,
855                        },
856                    }),
857                    Expr::AggregateFunction(AggregateFunction {
858                        func: max_udaf(),
859                        params: AggregateFunctionParams {
860                            args: vec![col("ts")],
861                            distinct: false,
862                            filter: None,
863                            order_by: vec![],
864                            null_treatment: None,
865                        },
866                    })
867                    .alias("ts"),
868                ],
869            )
870            .unwrap()
871            .build()
872            .unwrap();
873        let range = LogicalPlan::Extension(Extension {
874            node: Arc::new(
875                RangeManipulate::new(
876                    1000,
877                    1000,
878                    1000,
879                    0,
880                    1000,
881                    "ts".to_string(),
882                    vec!["v0".to_string()],
883                    scan_plan(provider(), "range"),
884                )
885                .unwrap(),
886            ),
887        });
888
889        for (plan, scan_count) in [(window, 1), (join, 2), (nonlast_aggregate, 1), (range, 1)] {
890            let rewritten = ScanHintRule
891                .rewrite(single_evaluation(plan), &OptimizerContext::default())
892                .unwrap()
893                .data;
894            assert_eq!(
895                scan_requests(&rewritten)
896                    .into_iter()
897                    .map(|request| request.series_row_selector)
898                    .collect::<Vec<_>>(),
899                vec![None; scan_count]
900            );
901        }
902    }
903
904    #[test]
905    fn single_evaluation_last_value_aggregate_keeps_legacy_selector() {
906        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
907        let rewritten = ScanHintRule
908            .rewrite(
909                single_evaluation(last_value_aggregate(scan_plan(provider, "t"))),
910                &OptimizerContext::default(),
911            )
912            .unwrap()
913            .data;
914
915        assert_eq!(
916            scan_requests(&rewritten)[0].series_row_selector,
917            Some(TimeSeriesRowSelector::LastRow { after_merge: false })
918        );
919    }
920
921    #[test]
922    fn single_evaluation_allows_second_to_millisecond_time_index_cast() {
923        let provider = Arc::new(mock_table_provider_with_timestamp(
924            RegionId::new(1, 1),
925            ConcreteDataType::timestamp_second_datatype(),
926        ));
927        let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
928            .project(vec![
929                Expr::Cast(Cast::new(
930                    Box::new(Expr::Column(Column::new(Some("t"), "ts"))),
931                    DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
932                ))
933                .alias("ts"),
934            ])
935            .unwrap()
936            .build()
937            .unwrap();
938        let rewritten = ScanHintRule
939            .rewrite(single_evaluation(input), &OptimizerContext::default())
940            .unwrap()
941            .data;
942
943        assert_eq!(
944            scan_requests(&rewritten)[0].series_row_selector,
945            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
946        );
947    }
948
949    #[test]
950    fn single_evaluation_uses_last_row_for_microsecond_and_nanosecond_time_indexes() {
951        for timestamp_type in [
952            ConcreteDataType::timestamp_microsecond_datatype(),
953            ConcreteDataType::timestamp_nanosecond_datatype(),
954        ] {
955            let direct_provider = Arc::new(mock_table_provider_with_timestamp(
956                RegionId::new(1, 1),
957                timestamp_type.clone(),
958            ));
959            let direct = ScanHintRule
960                .rewrite(
961                    single_evaluation(scan_plan(direct_provider, "direct")),
962                    &OptimizerContext::default(),
963                )
964                .unwrap()
965                .data;
966            assert_eq!(
967                scan_requests(&direct)[0].series_row_selector,
968                Some(TimeSeriesRowSelector::LastRow { after_merge: true })
969            );
970
971            let projection_provider = Arc::new(mock_table_provider_with_timestamp(
972                RegionId::new(1, 1),
973                timestamp_type,
974            ));
975            let projection = LogicalPlanBuilder::from(scan_plan(projection_provider, "projection"))
976                .project(vec![col("ts")])
977                .unwrap()
978                .build()
979                .unwrap();
980            let projected = ScanHintRule
981                .rewrite(single_evaluation(projection), &OptimizerContext::default())
982                .unwrap()
983                .data;
984            assert_eq!(
985                scan_requests(&projected)[0].series_row_selector,
986                Some(TimeSeriesRowSelector::LastRow { after_merge: true })
987            );
988        }
989    }
990
991    #[test]
992    fn single_evaluation_rejects_lossy_microsecond_and_nanosecond_time_index_casts() {
993        for timestamp_type in [
994            ConcreteDataType::timestamp_microsecond_datatype(),
995            ConcreteDataType::timestamp_nanosecond_datatype(),
996        ] {
997            let provider = Arc::new(mock_table_provider_with_timestamp(
998                RegionId::new(1, 1),
999                timestamp_type,
1000            ));
1001            let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
1002                .project(vec![
1003                    Expr::Cast(Cast::new(
1004                        Box::new(Expr::Column(Column::new(Some("t"), "ts"))),
1005                        DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
1006                    ))
1007                    .alias("ts"),
1008                ])
1009                .unwrap()
1010                .build()
1011                .unwrap();
1012            let rewritten = ScanHintRule
1013                .rewrite(single_evaluation(input), &OptimizerContext::default())
1014                .unwrap()
1015                .data;
1016
1017            assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1018        }
1019    }
1020
1021    #[test]
1022    fn single_evaluation_rejects_unresolved_time_index_cast_qualifier() {
1023        let provider = Arc::new(mock_table_provider_with_timestamp(
1024            RegionId::new(1, 1),
1025            ConcreteDataType::timestamp_second_datatype(),
1026        ));
1027        let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
1028            .project(vec![col("ts")])
1029            .unwrap()
1030            .build()
1031            .unwrap();
1032        let LogicalPlan::Projection(projection) = input else {
1033            unreachable!();
1034        };
1035        let unresolved_cast = Expr::Cast(Cast::new(
1036            Box::new(Expr::Column(Column::new(Some("missing"), "ts"))),
1037            DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
1038        ))
1039        .alias("ts");
1040
1041        assert!(!single_evaluation_projection_expr_allowed(
1042            &unresolved_cast,
1043            &projection
1044        ));
1045    }
1046
1047    #[test]
1048    fn single_evaluation_rejects_projection_expressions_that_change_rows() {
1049        let invalid_projections = [
1050            vec![col("ts").alias("renamed"), col("ts")],
1051            vec![
1052                Expr::BinaryExpr(datafusion_expr::expr::BinaryExpr::new(
1053                    Box::new(col("v0")),
1054                    datafusion_expr::Operator::Plus,
1055                    Box::new(lit(1.0_f64)),
1056                ))
1057                .alias("v0"),
1058                col("ts"),
1059            ],
1060            vec![
1061                Expr::Cast(Cast::new(
1062                    Box::new(col("ts")),
1063                    DataType::Timestamp(ArrowTimeUnit::Second, None),
1064                ))
1065                .alias("ts"),
1066            ],
1067            vec![
1068                Expr::Cast(Cast::new(Box::new(col("v0")), DataType::Int64)).alias("v0"),
1069                col("ts"),
1070            ],
1071            vec![
1072                Expr::Cast(Cast::new(
1073                    Box::new(col("ts")),
1074                    DataType::Timestamp(ArrowTimeUnit::Microsecond, None),
1075                ))
1076                .alias("ts"),
1077            ],
1078            vec![
1079                Expr::Cast(Cast::new(
1080                    Box::new(col("ts")),
1081                    DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1082                ))
1083                .alias("ts"),
1084            ],
1085        ];
1086
1087        for expressions in invalid_projections {
1088            let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1089            let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
1090                .project(expressions)
1091                .unwrap()
1092                .build()
1093                .unwrap();
1094            let rewritten = ScanHintRule
1095                .rewrite(single_evaluation(input), &OptimizerContext::default())
1096                .unwrap()
1097                .data;
1098
1099            assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1100        }
1101    }
1102
1103    #[test]
1104    fn range_evaluation_does_not_set_last_row() {
1105        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1106        let plan = instant_plan(provider.clone(), 2000);
1107        let rewritten = ScanHintRule
1108            .rewrite(plan, &OptimizerContext::default())
1109            .unwrap()
1110            .data;
1111
1112        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1113
1114        assert_eq!(provider.scan_request().series_row_selector, None);
1115    }
1116
1117    #[test]
1118    fn residual_field_filter_does_not_set_last_row() {
1119        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1120        let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
1121            .filter(col("v0").gt(lit(1.0_f64)))
1122            .unwrap()
1123            .build()
1124            .unwrap();
1125        let rewritten = ScanHintRule
1126            .rewrite(single_evaluation(input), &OptimizerContext::default())
1127            .unwrap()
1128            .data;
1129
1130        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1131    }
1132
1133    #[test]
1134    fn outer_residual_filter_does_not_block_last_row() {
1135        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1136        let plan = LogicalPlanBuilder::from(single_evaluation(scan_plan(provider, "t")))
1137            .filter(col("v0").gt(lit(1.0_f64)))
1138            .unwrap()
1139            .build()
1140            .unwrap();
1141        let rewritten = ScanHintRule
1142            .rewrite(plan, &OptimizerContext::default())
1143            .unwrap()
1144            .data;
1145
1146        assert_eq!(
1147            scan_requests(&rewritten)[0].series_row_selector,
1148            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1149        );
1150    }
1151
1152    #[test]
1153    fn branch_local_residual_filter_isolation_in_both_orders() {
1154        for filtered_first in [true, false] {
1155            let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1156            let filtered = single_evaluation(
1157                LogicalPlanBuilder::from(scan_plan(provider.clone(), "filtered"))
1158                    .filter(col("v0").gt(lit(1.0_f64)))
1159                    .unwrap()
1160                    .build()
1161                    .unwrap(),
1162            );
1163            let plain = single_evaluation(scan_plan(provider.clone(), "plain"));
1164            let union = if filtered_first {
1165                LogicalPlanBuilder::from(filtered)
1166                    .union(plain)
1167                    .unwrap()
1168                    .build()
1169                    .unwrap()
1170            } else {
1171                LogicalPlanBuilder::from(plain)
1172                    .union(filtered)
1173                    .unwrap()
1174                    .build()
1175                    .unwrap()
1176            };
1177            let rewritten = ScanHintRule
1178                .rewrite(union, &OptimizerContext::default())
1179                .unwrap()
1180                .data;
1181            let requests = scan_requests_with_names(&rewritten)
1182                .into_iter()
1183                .collect::<HashMap<_, _>>();
1184
1185            assert_eq!(requests["filtered"].series_row_selector, None);
1186            assert_eq!(
1187                requests["plain"].series_row_selector,
1188                Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1189            );
1190            assert_eq!(provider.scan_request().series_row_selector, None);
1191        }
1192    }
1193
1194    #[test]
1195    fn union_inside_single_evaluation_blocks_both_branches() {
1196        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1197        let union = LogicalPlanBuilder::from(scan_plan(provider.clone(), "left"))
1198            .union(scan_plan(provider, "right"))
1199            .unwrap()
1200            .build()
1201            .unwrap();
1202        let rewritten = ScanHintRule
1203            .rewrite(single_evaluation(union), &OptimizerContext::default())
1204            .unwrap()
1205            .data;
1206
1207        assert_eq!(
1208            scan_requests(&rewritten)
1209                .into_iter()
1210                .map(|request| request.series_row_selector)
1211                .collect::<Vec<_>>(),
1212            vec![None, None]
1213        );
1214    }
1215
1216    #[test]
1217    fn residual_time_filters_do_not_set_last_row() {
1218        for predicate in [
1219            col("ts").lt(lit(1_i64)),
1220            Expr::Cast(Cast::new(Box::new(col("ts")), DataType::Int64)).gt(lit(1_i64)),
1221        ] {
1222            let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1223            let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
1224                .filter(predicate)
1225                .unwrap()
1226                .build()
1227                .unwrap();
1228            let rewritten = ScanHintRule
1229                .rewrite(single_evaluation(input), &OptimizerContext::default())
1230                .unwrap()
1231                .data;
1232
1233            assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1234        }
1235    }
1236
1237    #[test]
1238    fn inner_single_evaluation_resets_residual_filter() {
1239        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1240        let filtered = LogicalPlanBuilder::from(single_evaluation(scan_plan(provider, "t")))
1241            .filter(col("v0").gt(lit(1.0_f64)))
1242            .unwrap()
1243            .build()
1244            .unwrap();
1245        let rewritten = ScanHintRule
1246            .rewrite(single_evaluation(filtered), &OptimizerContext::default())
1247            .unwrap()
1248            .data;
1249
1250        assert_eq!(
1251            scan_requests(&rewritten)[0].series_row_selector,
1252            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1253        );
1254    }
1255
1256    #[test]
1257    fn single_evaluation_with_tag_filter_sets_last_row() {
1258        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1259        let plan = instant_plan_with_filters(provider, vec![col("k0").eq(lit("tag"))]);
1260        let rewritten = ScanHintRule
1261            .rewrite(plan, &OptimizerContext::default())
1262            .unwrap()
1263            .data;
1264
1265        assert_eq!(
1266            scan_requests(&rewritten)[0].series_row_selector,
1267            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1268        );
1269    }
1270
1271    #[test]
1272    fn single_evaluation_with_timestamp_filter_sets_last_row() {
1273        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1274        let plan = instant_plan_with_filters(provider, vec![col("ts").gt_eq(lit(1_i64))]);
1275        let rewritten = ScanHintRule
1276            .rewrite(plan, &OptimizerContext::default())
1277            .unwrap()
1278            .data;
1279
1280        assert_eq!(
1281            scan_requests(&rewritten)[0].series_row_selector,
1282            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1283        );
1284    }
1285
1286    #[test]
1287    fn single_evaluation_with_cast_timestamp_filter_does_not_set_last_row() {
1288        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1289        let filter = Expr::Cast(Cast::new(Box::new(col("ts")), DataType::Int64)).gt(lit(1_i64));
1290        let plan = instant_plan_with_filters(provider, vec![filter]);
1291        let rewritten = ScanHintRule
1292            .rewrite(plan, &OptimizerContext::default())
1293            .unwrap()
1294            .data;
1295
1296        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1297    }
1298
1299    #[test]
1300    fn single_evaluation_with_multi_column_timestamp_filter_does_not_set_last_row() {
1301        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1302        let plan = instant_plan_with_filters(provider, vec![col("ts").gt(col("k0"))]);
1303        let rewritten = ScanHintRule
1304            .rewrite(plan, &OptimizerContext::default())
1305            .unwrap()
1306            .data;
1307
1308        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1309    }
1310
1311    #[test]
1312    fn single_evaluation_with_field_filter_does_not_set_last_row() {
1313        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1314        let plan = instant_plan_with_filters(provider, vec![col("v0").gt(lit(1.0_f64))]);
1315        let rewritten = ScanHintRule
1316            .rewrite(plan, &OptimizerContext::default())
1317            .unwrap()
1318            .data;
1319
1320        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1321    }
1322
1323    #[test]
1324    fn single_evaluation_with_unknown_filter_does_not_set_last_row() {
1325        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1326        let plan = instant_plan_with_filters(provider, vec![col("unknown").eq(lit(1_i64))]);
1327        let rewritten = ScanHintRule
1328            .rewrite(plan, &OptimizerContext::default())
1329            .unwrap()
1330            .data;
1331
1332        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1333    }
1334
1335    #[test]
1336    fn expression_subquery_isolated_from_outer_single_evaluation() {
1337        let outer_provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1338        let inner_provider = Arc::new(mock_table_provider(RegionId::new(2, 1)));
1339        let plan = instant_with_expression_subquery(
1340            outer_provider.clone(),
1341            instant_plan_named(inner_provider, "inner", 2000),
1342            1000,
1343        );
1344        let rewritten = ScanHintRule
1345            .rewrite(plan, &OptimizerContext::default())
1346            .unwrap()
1347            .data;
1348
1349        let requests = scan_requests_with_names(&rewritten)
1350            .into_iter()
1351            .collect::<HashMap<_, _>>();
1352        assert_eq!(requests["outer"].series_row_selector, None);
1353        assert_eq!(requests["inner"].series_row_selector, None);
1354    }
1355
1356    #[test]
1357    fn expression_subquery_can_start_its_own_single_evaluation() {
1358        let outer_provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1359        let inner_provider = Arc::new(mock_table_provider(RegionId::new(2, 1)));
1360        let plan = instant_with_expression_subquery(
1361            outer_provider,
1362            instant_plan_named(inner_provider, "inner", 1000),
1363            1000,
1364        );
1365        let rewritten = ScanHintRule
1366            .rewrite(plan, &OptimizerContext::default())
1367            .unwrap()
1368            .data;
1369
1370        let requests = scan_requests_with_names(&rewritten)
1371            .into_iter()
1372            .collect::<HashMap<_, _>>();
1373        assert_eq!(requests["outer"].series_row_selector, None);
1374        assert_eq!(
1375            requests["inner"].series_row_selector,
1376            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1377        );
1378    }
1379
1380    #[test]
1381    fn nested_single_outer_range_inner_does_not_set_inner_last_row() {
1382        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1383        let plan = instant_plan_named(provider.clone(), "nested", 2000);
1384        let plan = LogicalPlan::Extension(Extension {
1385            node: Arc::new(InstantManipulate::new(
1386                1000,
1387                1000,
1388                1000,
1389                1000,
1390                0,
1391                "ts".to_string(),
1392                vec![],
1393                Some("v0".to_string()),
1394                plan,
1395            )),
1396        });
1397        let rewritten = ScanHintRule
1398            .rewrite(plan, &OptimizerContext::default())
1399            .unwrap()
1400            .data;
1401        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1402        assert_eq!(provider.scan_request().series_row_selector, None);
1403    }
1404
1405    #[test]
1406    fn nested_range_outer_single_inner_sets_inner_last_row() {
1407        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1408        let plan = instant_plan_named(provider.clone(), "nested", 1000);
1409        let plan = LogicalPlan::Extension(Extension {
1410            node: Arc::new(InstantManipulate::new(
1411                1000,
1412                2000,
1413                1000,
1414                1000,
1415                0,
1416                "ts".to_string(),
1417                vec![],
1418                Some("v0".to_string()),
1419                plan,
1420            )),
1421        });
1422        let rewritten = ScanHintRule
1423            .rewrite(plan, &OptimizerContext::default())
1424            .unwrap()
1425            .data;
1426        assert_eq!(
1427            scan_requests(&rewritten)[0].series_row_selector,
1428            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1429        );
1430    }
1431
1432    #[test]
1433    fn shared_provider_isolated_between_single_and_range_scans() {
1434        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1435        let plan = LogicalPlanBuilder::from(instant_plan(provider.clone(), 1000))
1436            .union(instant_plan(provider.clone(), 2000))
1437            .unwrap()
1438            .build()
1439            .unwrap();
1440        let rewritten = ScanHintRule
1441            .rewrite(plan, &OptimizerContext::default())
1442            .unwrap()
1443            .data;
1444        let requests = scan_requests(&rewritten);
1445
1446        assert_eq!(requests.len(), 2);
1447        assert_eq!(
1448            requests[0].series_row_selector,
1449            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1450        );
1451        assert_eq!(requests[1].series_row_selector, None);
1452        assert_eq!(provider.scan_request().series_row_selector, None);
1453    }
1454
1455    #[test]
1456    fn set_order_hint() {
1457        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1458        let table_source = Arc::new(DefaultTableSource::new(provider.clone()));
1459        let plan = LogicalPlanBuilder::scan("t", table_source, None)
1460            .unwrap()
1461            .sort(vec![col("ts").sort(true, false)])
1462            .unwrap()
1463            .sort(vec![col("ts").sort(false, true)])
1464            .unwrap()
1465            .build()
1466            .unwrap();
1467
1468        let context = OptimizerContext::default();
1469        let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data;
1470
1471        // should read the first (with `.sort(true, false)`) sort option
1472        let scan_req = scan_requests(&rewritten)[0].clone();
1473        assert_eq!(
1474            OrderOption {
1475                name: "ts".to_string(),
1476                options: SortOptions {
1477                    descending: false,
1478                    nulls_first: false
1479                }
1480            },
1481            scan_req.output_ordering.as_ref().unwrap()[0]
1482        );
1483    }
1484
1485    #[test]
1486    fn set_time_series_row_selector_hint() {
1487        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1488        let plan = last_value_aggregate(scan_plan(provider.clone(), "t"));
1489
1490        let context = OptimizerContext::default();
1491        let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data;
1492
1493        let scan_req = scan_requests(&rewritten)[0].clone();
1494        assert_eq!(
1495            scan_req.series_row_selector,
1496            Some(TimeSeriesRowSelector::LastRow { after_merge: false })
1497        );
1498    }
1499
1500    #[test]
1501    fn set_order_hint_sets_per_series_distribution_for_tsid_sort() {
1502        let provider = Arc::new(mock_table_provider_with_tsid(RegionId::new(1, 1)));
1503        let table_source = Arc::new(DefaultTableSource::new(provider.clone()));
1504        let plan = LogicalPlanBuilder::scan("t", table_source, None)
1505            .unwrap()
1506            .sort(vec![
1507                col(DATA_SCHEMA_TSID_COLUMN_NAME).sort(true, true),
1508                col("ts").sort(true, true),
1509            ])
1510            .unwrap()
1511            .build()
1512            .unwrap();
1513
1514        let context = OptimizerContext::default();
1515        let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data;
1516
1517        let scan_req = scan_requests(&rewritten)[0].clone();
1518        assert_eq!(
1519            scan_req.distribution,
1520            Some(TimeSeriesDistribution::PerSeries)
1521        );
1522    }
1523}
query/optimizer/scan_hint.rs

query/optimizer/
scan_hint.rs
