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::{NullTreatment, 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                // A shared final-row selector would discard earlier nonnull values.
417                || func.params.null_treatment == Some(NullTreatment::IgnoreNulls)
418            {
419                return None;
420            }
421            let order_by = &func.params.order_by;
422            if order_by.len() != 1 || !order_by[0].asc {
423                return None;
424            }
425            if let Some(existing) = &order_by_expr {
426                if existing != &order_by[0] {
427                    return None;
428                }
429            } else {
430                order_by_expr = Some(order_by[0].clone());
431            }
432        }
433        let Expr::Column(order_by_col) = order_by_expr?.expr else {
434            return None;
435        };
436        let mut group_by_cols = HashSet::with_capacity(aggregate.group_expr.len());
437        for expr in &aggregate.group_expr {
438            let Expr::Column(col) = expr else {
439                return None;
440            };
441            group_by_cols.insert(col.clone());
442        }
443        Some((group_by_cols, order_by_col))
444    }
445}
446
447#[cfg(test)]
448mod test {
449    use std::collections::HashMap;
450    use std::sync::Arc;
451
452    use datafusion::functions_aggregate::first_last::last_value_udaf;
453    use datafusion::functions_aggregate::min_max::max_udaf;
454    use datafusion::functions_window::row_number::RowNumber;
455    use datafusion::logical_expr::expr::WindowFunction;
456    use datafusion::logical_expr::{WindowFrame, WindowFunctionDefinition};
457    use datafusion::prelude::JoinType;
458    use datafusion_common::tree_node::TreeNodeRecursion;
459    use datafusion_expr::expr::{
460        AggregateFunction, AggregateFunctionParams, Cast, NullTreatment, WindowFunctionParams,
461    };
462    use datafusion_expr::expr_fn::scalar_subquery;
463    use datafusion_expr::{Extension, LogicalPlan, LogicalPlanBuilder, col, lit};
464    use datafusion_optimizer::OptimizerContext;
465    use datatypes::arrow::datatypes::DataType;
466    use datatypes::data_type::ConcreteDataType;
467    use datatypes::schema::ColumnSchema;
468    use promql::extension_plan::RangeManipulate;
469    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
470    use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
471    use store_api::storage::{RegionId, TimeSeriesRowSelector};
472
473    use super::*;
474
475    fn scan_requests(plan: &LogicalPlan) -> Vec<store_api::storage::ScanRequest> {
476        scan_requests_with_names(plan)
477            .into_iter()
478            .map(|(_, request)| request)
479            .collect()
480    }
481
482    fn scan_requests_with_names(
483        plan: &LogicalPlan,
484    ) -> Vec<(String, store_api::storage::ScanRequest)> {
485        let mut requests = Vec::new();
486        plan.apply_with_subqueries(|node| {
487            if let LogicalPlan::TableScan(scan) = node
488                && let Some(source) = scan.source.downcast_ref::<DefaultTableSource>()
489                && let Some(provider) = source.table_provider.downcast_ref::<DummyTableProvider>()
490            {
491                requests.push((scan.table_name.to_string(), provider.scan_request()));
492            }
493            Ok(TreeNodeRecursion::Continue)
494        })
495        .unwrap();
496        requests
497    }
498
499    fn instant_plan(provider: Arc<DummyTableProvider>, end: i64) -> LogicalPlan {
500        instant_plan_named(provider, "t", end)
501    }
502
503    fn instant_plan_with_filters(
504        provider: Arc<DummyTableProvider>,
505        filters: Vec<Expr>,
506    ) -> LogicalPlan {
507        let scan = scan_plan(provider, "t");
508        let LogicalPlan::TableScan(mut scan) = scan else {
509            unreachable!();
510        };
511        scan.filters = filters;
512        LogicalPlan::Extension(Extension {
513            node: Arc::new(InstantManipulate::new(
514                1000,
515                1000,
516                1000,
517                1000,
518                0,
519                "ts".to_string(),
520                vec![],
521                Some("v0".to_string()),
522                LogicalPlan::TableScan(scan),
523            )),
524        })
525    }
526
527    fn scan_plan(provider: Arc<DummyTableProvider>, table_name: &str) -> LogicalPlan {
528        LogicalPlanBuilder::scan(
529            table_name,
530            Arc::new(DefaultTableSource::new(provider)),
531            None,
532        )
533        .unwrap()
534        .build()
535        .unwrap()
536    }
537
538    fn mock_table_provider_with_timestamp(
539        region_id: RegionId,
540        timestamp_type: ConcreteDataType,
541    ) -> DummyTableProvider {
542        let mut builder = RegionMetadataBuilder::new(region_id);
543        builder
544            .push_column_metadata(ColumnMetadata {
545                column_schema: ColumnSchema::new("k0", ConcreteDataType::string_datatype(), true),
546                semantic_type: SemanticType::Tag,
547                column_id: 1,
548            })
549            .push_column_metadata(ColumnMetadata {
550                column_schema: ColumnSchema::new("ts", timestamp_type, false),
551                semantic_type: SemanticType::Timestamp,
552                column_id: 2,
553            })
554            .push_column_metadata(ColumnMetadata {
555                column_schema: ColumnSchema::new("v0", ConcreteDataType::float64_datatype(), false),
556                semantic_type: SemanticType::Field,
557                column_id: 3,
558            })
559            .primary_key(vec![1]);
560        let metadata = Arc::new(builder.build().unwrap());
561        let engine = Arc::new(MetaRegionEngine::with_metadata(metadata.clone()));
562        DummyTableProvider::new(region_id, engine, metadata)
563    }
564
565    fn last_value_aggregate(input: LogicalPlan) -> LogicalPlan {
566        last_value_aggregate_with_null_treatments(input, [None, None])
567    }
568
569    fn last_value_aggregate_with_null_treatments(
570        input: LogicalPlan,
571        null_treatments: [Option<NullTreatment>; 2],
572    ) -> LogicalPlan {
573        let aggregate = LogicalPlanBuilder::from(input)
574            .aggregate(
575                vec![col("k0")],
576                vec![
577                    Expr::AggregateFunction(AggregateFunction {
578                        func: last_value_udaf(),
579                        params: AggregateFunctionParams {
580                            args: vec![col("v0")],
581                            distinct: false,
582                            filter: None,
583                            order_by: vec![Sort {
584                                expr: col("ts"),
585                                asc: true,
586                                nulls_first: true,
587                            }],
588                            null_treatment: null_treatments[0],
589                        },
590                    }),
591                    Expr::AggregateFunction(AggregateFunction {
592                        func: last_value_udaf(),
593                        params: AggregateFunctionParams {
594                            args: vec![col("ts")],
595                            distinct: false,
596                            filter: None,
597                            order_by: vec![Sort {
598                                expr: col("ts"),
599                                asc: true,
600                                nulls_first: true,
601                            }],
602                            null_treatment: null_treatments[1],
603                        },
604                    }),
605                ],
606            )
607            .unwrap()
608            .build()
609            .unwrap();
610        let timestamp = aggregate.schema().field(2).name().clone();
611
612        LogicalPlanBuilder::from(aggregate)
613            .project(vec![col("k0"), col(timestamp).alias("ts")])
614            .unwrap()
615            .build()
616            .unwrap()
617    }
618
619    fn single_evaluation(input: LogicalPlan) -> LogicalPlan {
620        LogicalPlan::Extension(Extension {
621            node: Arc::new(InstantManipulate::new(
622                1000,
623                1000,
624                1000,
625                1000,
626                0,
627                "ts".to_string(),
628                vec![],
629                Some("v0".to_string()),
630                input,
631            )),
632        })
633    }
634
635    fn instant_with_expression_subquery(
636        outer_provider: Arc<DummyTableProvider>,
637        inner_plan: LogicalPlan,
638        outer_end: i64,
639    ) -> LogicalPlan {
640        let outer_scan = scan_plan(outer_provider, "outer");
641        let input = LogicalPlanBuilder::from(outer_scan)
642            .project(vec![
643                col("ts"),
644                col("v0"),
645                scalar_subquery(Arc::new(inner_plan)),
646            ])
647            .unwrap()
648            .build()
649            .unwrap();
650        LogicalPlan::Extension(Extension {
651            node: Arc::new(InstantManipulate::new(
652                1000,
653                outer_end,
654                1000,
655                1000,
656                0,
657                "ts".to_string(),
658                vec![],
659                Some("v0".to_string()),
660                input,
661            )),
662        })
663    }
664
665    fn instant_plan_named(
666        provider: Arc<DummyTableProvider>,
667        table_name: &str,
668        end: i64,
669    ) -> LogicalPlan {
670        let input = LogicalPlanBuilder::scan(
671            table_name,
672            Arc::new(DefaultTableSource::new(provider)),
673            None,
674        )
675        .unwrap()
676        .build()
677        .unwrap();
678        LogicalPlan::Extension(Extension {
679            node: Arc::new(InstantManipulate::new(
680                1000,
681                end,
682                1000,
683                1000,
684                0,
685                "ts".to_string(),
686                vec![],
687                Some("v0".to_string()),
688                input,
689            )),
690        })
691    }
692    use crate::optimizer::test_util::{
693        MetaRegionEngine, mock_table_provider, mock_table_provider_with_tsid,
694    };
695
696    #[test]
697    fn single_evaluation_sets_last_row_on_the_rewritten_scan() {
698        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
699        let plan = instant_plan(provider.clone(), 1000);
700        let rewritten = ScanHintRule
701            .rewrite(plan, &OptimizerContext::default())
702            .unwrap()
703            .data;
704
705        assert_eq!(
706            scan_requests(&rewritten)[0].series_row_selector,
707            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
708        );
709
710        assert_eq!(provider.scan_request().series_row_selector, None);
711    }
712
713    #[test]
714    fn single_evaluation_limit_sort_does_not_set_last_row_below_limit() {
715        let mut selectors = Vec::new();
716        for input in [
717            // Instant(1000, lookback=1000) -> Limit(0, 1) -> Sort(ts ASC) -> Scan.
718            LogicalPlanBuilder::from(scan_plan(
719                Arc::new(mock_table_provider(RegionId::new(1, 1))),
720                "t",
721            ))
722            .sort(vec![col("ts").sort(true, false)])
723            .unwrap()
724            .limit(0, Some(1))
725            .unwrap()
726            .build()
727            .unwrap(),
728            // Sort.fetch is a limit embedded in the Sort node and has the same boundary.
729            LogicalPlanBuilder::from(scan_plan(
730                Arc::new(mock_table_provider(RegionId::new(1, 1))),
731                "t",
732            ))
733            .sort_with_limit(vec![col("ts").sort(true, false)], Some(1))
734            .unwrap()
735            .build()
736            .unwrap(),
737        ] {
738            let rewritten = ScanHintRule
739                .rewrite(single_evaluation(input), &OptimizerContext::default())
740                .unwrap()
741                .data;
742            selectors.push(scan_requests(&rewritten)[0].series_row_selector);
743        }
744
745        // With samples (ts=100, v=10) and (ts=900, v=1), an unhinted ascending
746        // sort/limit pipeline returns 10. LastRow at the scan instead leaves only
747        // (900, 1) before the limit. A LastRow scan hint cannot cross either limit.
748        assert_eq!(selectors, vec![None, None]);
749    }
750
751    #[test]
752    fn single_evaluation_allows_controlled_promql_selector_chain() {
753        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
754        let projection = LogicalPlanBuilder::from(scan_plan(provider, "t"))
755            .project(vec![
756                col("k0").alias("k0"),
757                Expr::Cast(Cast::new(
758                    Box::new(col("ts")),
759                    DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
760                ))
761                .alias("ts"),
762                col("v0"),
763            ])
764            .unwrap()
765            .sort(vec![col("ts").sort(true, false)])
766            .unwrap()
767            .build()
768            .unwrap();
769        let divide = LogicalPlan::Extension(Extension {
770            node: Arc::new(SeriesDivide::new(
771                vec!["k0".to_string()],
772                "ts".to_string(),
773                projection,
774            )),
775        });
776        let normalize = LogicalPlan::Extension(Extension {
777            node: Arc::new(SeriesNormalize::new(
778                42,
779                "ts",
780                false,
781                vec!["k0".to_string()],
782                divide,
783            )),
784        });
785        let rewritten = ScanHintRule
786            .rewrite(single_evaluation(normalize), &OptimizerContext::default())
787            .unwrap()
788            .data;
789
790        assert_eq!(
791            scan_requests(&rewritten)[0].series_row_selector,
792            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
793        );
794    }
795
796    #[test]
797    fn single_evaluation_filtering_normalize_does_not_set_last_row() {
798        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
799        let normalize = LogicalPlan::Extension(Extension {
800            node: Arc::new(SeriesNormalize::new(
801                42,
802                "ts",
803                true,
804                vec!["k0".to_string()],
805                scan_plan(provider, "t"),
806            )),
807        });
808        let rewritten = ScanHintRule
809            .rewrite(single_evaluation(normalize), &OptimizerContext::default())
810            .unwrap()
811            .data;
812
813        // An older finite sample can precede a newest stale marker. LastRow
814        // would discard that sample before SeriesNormalize filters the marker.
815        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
816    }
817
818    #[test]
819    fn single_evaluation_rejects_row_changing_nodes() {
820        let provider = || Arc::new(mock_table_provider(RegionId::new(1, 1)));
821        let window = LogicalPlanBuilder::from(scan_plan(provider(), "window"))
822            .window(vec![Expr::WindowFunction(Box::new(WindowFunction {
823                fun: WindowFunctionDefinition::WindowUDF(Arc::new(RowNumber::new().into())),
824                params: WindowFunctionParams {
825                    args: vec![],
826                    partition_by: vec![col("k0")],
827                    order_by: vec![col("ts").sort(true, true)],
828                    window_frame: WindowFrame::new(Some(true)),
829                    filter: None,
830                    null_treatment: None,
831                    distinct: false,
832                },
833            }))])
834            .unwrap()
835            .build()
836            .unwrap();
837        let join = LogicalPlanBuilder::from(scan_plan(provider(), "left"))
838            .join(
839                scan_plan(provider(), "right"),
840                JoinType::Inner,
841                (Vec::<Column>::new(), Vec::<Column>::new()),
842                None,
843            )
844            .unwrap()
845            .project(vec![
846                Expr::Column(Column::new(Some("left"), "ts")),
847                Expr::Column(Column::new(Some("left"), "v0")),
848            ])
849            .unwrap()
850            .build()
851            .unwrap();
852        let nonlast_aggregate = LogicalPlanBuilder::from(scan_plan(provider(), "aggregate"))
853            .aggregate(
854                vec![col("k0")],
855                vec![
856                    Expr::AggregateFunction(AggregateFunction {
857                        func: max_udaf(),
858                        params: AggregateFunctionParams {
859                            args: vec![col("v0")],
860                            distinct: false,
861                            filter: None,
862                            order_by: vec![],
863                            null_treatment: None,
864                        },
865                    }),
866                    Expr::AggregateFunction(AggregateFunction {
867                        func: max_udaf(),
868                        params: AggregateFunctionParams {
869                            args: vec![col("ts")],
870                            distinct: false,
871                            filter: None,
872                            order_by: vec![],
873                            null_treatment: None,
874                        },
875                    })
876                    .alias("ts"),
877                ],
878            )
879            .unwrap()
880            .build()
881            .unwrap();
882        let range = LogicalPlan::Extension(Extension {
883            node: Arc::new(
884                RangeManipulate::new(
885                    1000,
886                    1000,
887                    1000,
888                    0,
889                    1000,
890                    "ts".to_string(),
891                    vec!["v0".to_string()],
892                    scan_plan(provider(), "range"),
893                )
894                .unwrap(),
895            ),
896        });
897
898        for (plan, scan_count) in [(window, 1), (join, 2), (nonlast_aggregate, 1), (range, 1)] {
899            let rewritten = ScanHintRule
900                .rewrite(single_evaluation(plan), &OptimizerContext::default())
901                .unwrap()
902                .data;
903            assert_eq!(
904                scan_requests(&rewritten)
905                    .into_iter()
906                    .map(|request| request.series_row_selector)
907                    .collect::<Vec<_>>(),
908                vec![None; scan_count]
909            );
910        }
911    }
912
913    #[test]
914    fn last_value_null_treatment_controls_selector() {
915        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
916        let context = OptimizerContext::default();
917
918        for treatment in [None, Some(NullTreatment::RespectNulls)] {
919            let rewritten = ScanHintRule
920                .rewrite(
921                    last_value_aggregate_with_null_treatments(
922                        scan_plan(provider.clone(), "t"),
923                        [treatment, treatment],
924                    ),
925                    &context,
926                )
927                .unwrap()
928                .data;
929            assert_eq!(
930                scan_requests(&rewritten)[0].series_row_selector,
931                Some(TimeSeriesRowSelector::LastRow { after_merge: false })
932            );
933        }
934
935        // One IGNORE NULLS member prevents the shared scan selector for every
936        // aggregate in the group, including the member that respects nulls.
937        let rewritten = ScanHintRule
938            .rewrite(
939                last_value_aggregate_with_null_treatments(
940                    scan_plan(provider, "t"),
941                    [None, Some(NullTreatment::IgnoreNulls)],
942                ),
943                &context,
944            )
945            .unwrap()
946            .data;
947        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
948    }
949
950    #[test]
951    fn single_evaluation_last_value_aggregate_keeps_legacy_selector() {
952        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
953        let rewritten = ScanHintRule
954            .rewrite(
955                single_evaluation(last_value_aggregate(scan_plan(provider, "t"))),
956                &OptimizerContext::default(),
957            )
958            .unwrap()
959            .data;
960
961        assert_eq!(
962            scan_requests(&rewritten)[0].series_row_selector,
963            Some(TimeSeriesRowSelector::LastRow { after_merge: false })
964        );
965    }
966
967    #[test]
968    fn single_evaluation_allows_second_to_millisecond_time_index_cast() {
969        let provider = Arc::new(mock_table_provider_with_timestamp(
970            RegionId::new(1, 1),
971            ConcreteDataType::timestamp_second_datatype(),
972        ));
973        let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
974            .project(vec![
975                Expr::Cast(Cast::new(
976                    Box::new(Expr::Column(Column::new(Some("t"), "ts"))),
977                    DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
978                ))
979                .alias("ts"),
980            ])
981            .unwrap()
982            .build()
983            .unwrap();
984        let rewritten = ScanHintRule
985            .rewrite(single_evaluation(input), &OptimizerContext::default())
986            .unwrap()
987            .data;
988
989        assert_eq!(
990            scan_requests(&rewritten)[0].series_row_selector,
991            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
992        );
993    }
994
995    #[test]
996    fn single_evaluation_uses_last_row_for_microsecond_and_nanosecond_time_indexes() {
997        for timestamp_type in [
998            ConcreteDataType::timestamp_microsecond_datatype(),
999            ConcreteDataType::timestamp_nanosecond_datatype(),
1000        ] {
1001            let direct_provider = Arc::new(mock_table_provider_with_timestamp(
1002                RegionId::new(1, 1),
1003                timestamp_type.clone(),
1004            ));
1005            let direct = ScanHintRule
1006                .rewrite(
1007                    single_evaluation(scan_plan(direct_provider, "direct")),
1008                    &OptimizerContext::default(),
1009                )
1010                .unwrap()
1011                .data;
1012            assert_eq!(
1013                scan_requests(&direct)[0].series_row_selector,
1014                Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1015            );
1016
1017            let projection_provider = Arc::new(mock_table_provider_with_timestamp(
1018                RegionId::new(1, 1),
1019                timestamp_type,
1020            ));
1021            let projection = LogicalPlanBuilder::from(scan_plan(projection_provider, "projection"))
1022                .project(vec![col("ts")])
1023                .unwrap()
1024                .build()
1025                .unwrap();
1026            let projected = ScanHintRule
1027                .rewrite(single_evaluation(projection), &OptimizerContext::default())
1028                .unwrap()
1029                .data;
1030            assert_eq!(
1031                scan_requests(&projected)[0].series_row_selector,
1032                Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1033            );
1034        }
1035    }
1036
1037    #[test]
1038    fn single_evaluation_rejects_lossy_microsecond_and_nanosecond_time_index_casts() {
1039        for timestamp_type in [
1040            ConcreteDataType::timestamp_microsecond_datatype(),
1041            ConcreteDataType::timestamp_nanosecond_datatype(),
1042        ] {
1043            let provider = Arc::new(mock_table_provider_with_timestamp(
1044                RegionId::new(1, 1),
1045                timestamp_type,
1046            ));
1047            let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
1048                .project(vec![
1049                    Expr::Cast(Cast::new(
1050                        Box::new(Expr::Column(Column::new(Some("t"), "ts"))),
1051                        DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
1052                    ))
1053                    .alias("ts"),
1054                ])
1055                .unwrap()
1056                .build()
1057                .unwrap();
1058            let rewritten = ScanHintRule
1059                .rewrite(single_evaluation(input), &OptimizerContext::default())
1060                .unwrap()
1061                .data;
1062
1063            assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1064        }
1065    }
1066
1067    #[test]
1068    fn single_evaluation_rejects_unresolved_time_index_cast_qualifier() {
1069        let provider = Arc::new(mock_table_provider_with_timestamp(
1070            RegionId::new(1, 1),
1071            ConcreteDataType::timestamp_second_datatype(),
1072        ));
1073        let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
1074            .project(vec![col("ts")])
1075            .unwrap()
1076            .build()
1077            .unwrap();
1078        let LogicalPlan::Projection(projection) = input else {
1079            unreachable!();
1080        };
1081        let unresolved_cast = Expr::Cast(Cast::new(
1082            Box::new(Expr::Column(Column::new(Some("missing"), "ts"))),
1083            DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
1084        ))
1085        .alias("ts");
1086
1087        assert!(!single_evaluation_projection_expr_allowed(
1088            &unresolved_cast,
1089            &projection
1090        ));
1091    }
1092
1093    #[test]
1094    fn single_evaluation_rejects_projection_expressions_that_change_rows() {
1095        let invalid_projections = [
1096            vec![col("ts").alias("renamed"), col("ts")],
1097            vec![
1098                Expr::BinaryExpr(datafusion_expr::expr::BinaryExpr::new(
1099                    Box::new(col("v0")),
1100                    datafusion_expr::Operator::Plus,
1101                    Box::new(lit(1.0_f64)),
1102                ))
1103                .alias("v0"),
1104                col("ts"),
1105            ],
1106            vec![
1107                Expr::Cast(Cast::new(
1108                    Box::new(col("ts")),
1109                    DataType::Timestamp(ArrowTimeUnit::Second, None),
1110                ))
1111                .alias("ts"),
1112            ],
1113            vec![
1114                Expr::Cast(Cast::new(Box::new(col("v0")), DataType::Int64)).alias("v0"),
1115                col("ts"),
1116            ],
1117            vec![
1118                Expr::Cast(Cast::new(
1119                    Box::new(col("ts")),
1120                    DataType::Timestamp(ArrowTimeUnit::Microsecond, None),
1121                ))
1122                .alias("ts"),
1123            ],
1124            vec![
1125                Expr::Cast(Cast::new(
1126                    Box::new(col("ts")),
1127                    DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
1128                ))
1129                .alias("ts"),
1130            ],
1131        ];
1132
1133        for expressions in invalid_projections {
1134            let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1135            let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
1136                .project(expressions)
1137                .unwrap()
1138                .build()
1139                .unwrap();
1140            let rewritten = ScanHintRule
1141                .rewrite(single_evaluation(input), &OptimizerContext::default())
1142                .unwrap()
1143                .data;
1144
1145            assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1146        }
1147    }
1148
1149    #[test]
1150    fn range_evaluation_does_not_set_last_row() {
1151        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1152        let plan = instant_plan(provider.clone(), 2000);
1153        let rewritten = ScanHintRule
1154            .rewrite(plan, &OptimizerContext::default())
1155            .unwrap()
1156            .data;
1157
1158        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1159
1160        assert_eq!(provider.scan_request().series_row_selector, None);
1161    }
1162
1163    #[test]
1164    fn residual_field_filter_does_not_set_last_row() {
1165        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1166        let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
1167            .filter(col("v0").gt(lit(1.0_f64)))
1168            .unwrap()
1169            .build()
1170            .unwrap();
1171        let rewritten = ScanHintRule
1172            .rewrite(single_evaluation(input), &OptimizerContext::default())
1173            .unwrap()
1174            .data;
1175
1176        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1177    }
1178
1179    #[test]
1180    fn outer_residual_filter_does_not_block_last_row() {
1181        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1182        let plan = LogicalPlanBuilder::from(single_evaluation(scan_plan(provider, "t")))
1183            .filter(col("v0").gt(lit(1.0_f64)))
1184            .unwrap()
1185            .build()
1186            .unwrap();
1187        let rewritten = ScanHintRule
1188            .rewrite(plan, &OptimizerContext::default())
1189            .unwrap()
1190            .data;
1191
1192        assert_eq!(
1193            scan_requests(&rewritten)[0].series_row_selector,
1194            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1195        );
1196    }
1197
1198    #[test]
1199    fn branch_local_residual_filter_isolation_in_both_orders() {
1200        for filtered_first in [true, false] {
1201            let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1202            let filtered = single_evaluation(
1203                LogicalPlanBuilder::from(scan_plan(provider.clone(), "filtered"))
1204                    .filter(col("v0").gt(lit(1.0_f64)))
1205                    .unwrap()
1206                    .build()
1207                    .unwrap(),
1208            );
1209            let plain = single_evaluation(scan_plan(provider.clone(), "plain"));
1210            let union = if filtered_first {
1211                LogicalPlanBuilder::from(filtered)
1212                    .union(plain)
1213                    .unwrap()
1214                    .build()
1215                    .unwrap()
1216            } else {
1217                LogicalPlanBuilder::from(plain)
1218                    .union(filtered)
1219                    .unwrap()
1220                    .build()
1221                    .unwrap()
1222            };
1223            let rewritten = ScanHintRule
1224                .rewrite(union, &OptimizerContext::default())
1225                .unwrap()
1226                .data;
1227            let requests = scan_requests_with_names(&rewritten)
1228                .into_iter()
1229                .collect::<HashMap<_, _>>();
1230
1231            assert_eq!(requests["filtered"].series_row_selector, None);
1232            assert_eq!(
1233                requests["plain"].series_row_selector,
1234                Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1235            );
1236            assert_eq!(provider.scan_request().series_row_selector, None);
1237        }
1238    }
1239
1240    #[test]
1241    fn union_inside_single_evaluation_blocks_both_branches() {
1242        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1243        let union = LogicalPlanBuilder::from(scan_plan(provider.clone(), "left"))
1244            .union(scan_plan(provider, "right"))
1245            .unwrap()
1246            .build()
1247            .unwrap();
1248        let rewritten = ScanHintRule
1249            .rewrite(single_evaluation(union), &OptimizerContext::default())
1250            .unwrap()
1251            .data;
1252
1253        assert_eq!(
1254            scan_requests(&rewritten)
1255                .into_iter()
1256                .map(|request| request.series_row_selector)
1257                .collect::<Vec<_>>(),
1258            vec![None, None]
1259        );
1260    }
1261
1262    #[test]
1263    fn residual_time_filters_do_not_set_last_row() {
1264        for predicate in [
1265            col("ts").lt(lit(1_i64)),
1266            Expr::Cast(Cast::new(Box::new(col("ts")), DataType::Int64)).gt(lit(1_i64)),
1267        ] {
1268            let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1269            let input = LogicalPlanBuilder::from(scan_plan(provider, "t"))
1270                .filter(predicate)
1271                .unwrap()
1272                .build()
1273                .unwrap();
1274            let rewritten = ScanHintRule
1275                .rewrite(single_evaluation(input), &OptimizerContext::default())
1276                .unwrap()
1277                .data;
1278
1279            assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1280        }
1281    }
1282
1283    #[test]
1284    fn inner_single_evaluation_resets_residual_filter() {
1285        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1286        let filtered = LogicalPlanBuilder::from(single_evaluation(scan_plan(provider, "t")))
1287            .filter(col("v0").gt(lit(1.0_f64)))
1288            .unwrap()
1289            .build()
1290            .unwrap();
1291        let rewritten = ScanHintRule
1292            .rewrite(single_evaluation(filtered), &OptimizerContext::default())
1293            .unwrap()
1294            .data;
1295
1296        assert_eq!(
1297            scan_requests(&rewritten)[0].series_row_selector,
1298            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1299        );
1300    }
1301
1302    #[test]
1303    fn single_evaluation_with_tag_filter_sets_last_row() {
1304        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1305        let plan = instant_plan_with_filters(provider, vec![col("k0").eq(lit("tag"))]);
1306        let rewritten = ScanHintRule
1307            .rewrite(plan, &OptimizerContext::default())
1308            .unwrap()
1309            .data;
1310
1311        assert_eq!(
1312            scan_requests(&rewritten)[0].series_row_selector,
1313            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1314        );
1315    }
1316
1317    #[test]
1318    fn single_evaluation_with_timestamp_filter_sets_last_row() {
1319        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1320        let plan = instant_plan_with_filters(provider, vec![col("ts").gt_eq(lit(1_i64))]);
1321        let rewritten = ScanHintRule
1322            .rewrite(plan, &OptimizerContext::default())
1323            .unwrap()
1324            .data;
1325
1326        assert_eq!(
1327            scan_requests(&rewritten)[0].series_row_selector,
1328            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1329        );
1330    }
1331
1332    #[test]
1333    fn single_evaluation_with_cast_timestamp_filter_does_not_set_last_row() {
1334        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1335        let filter = Expr::Cast(Cast::new(Box::new(col("ts")), DataType::Int64)).gt(lit(1_i64));
1336        let plan = instant_plan_with_filters(provider, vec![filter]);
1337        let rewritten = ScanHintRule
1338            .rewrite(plan, &OptimizerContext::default())
1339            .unwrap()
1340            .data;
1341
1342        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1343    }
1344
1345    #[test]
1346    fn single_evaluation_with_multi_column_timestamp_filter_does_not_set_last_row() {
1347        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1348        let plan = instant_plan_with_filters(provider, vec![col("ts").gt(col("k0"))]);
1349        let rewritten = ScanHintRule
1350            .rewrite(plan, &OptimizerContext::default())
1351            .unwrap()
1352            .data;
1353
1354        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1355    }
1356
1357    #[test]
1358    fn single_evaluation_with_field_filter_does_not_set_last_row() {
1359        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1360        let plan = instant_plan_with_filters(provider, vec![col("v0").gt(lit(1.0_f64))]);
1361        let rewritten = ScanHintRule
1362            .rewrite(plan, &OptimizerContext::default())
1363            .unwrap()
1364            .data;
1365
1366        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1367    }
1368
1369    #[test]
1370    fn single_evaluation_with_unknown_filter_does_not_set_last_row() {
1371        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1372        let plan = instant_plan_with_filters(provider, vec![col("unknown").eq(lit(1_i64))]);
1373        let rewritten = ScanHintRule
1374            .rewrite(plan, &OptimizerContext::default())
1375            .unwrap()
1376            .data;
1377
1378        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1379    }
1380
1381    #[test]
1382    fn expression_subquery_isolated_from_outer_single_evaluation() {
1383        let outer_provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1384        let inner_provider = Arc::new(mock_table_provider(RegionId::new(2, 1)));
1385        let plan = instant_with_expression_subquery(
1386            outer_provider.clone(),
1387            instant_plan_named(inner_provider, "inner", 2000),
1388            1000,
1389        );
1390        let rewritten = ScanHintRule
1391            .rewrite(plan, &OptimizerContext::default())
1392            .unwrap()
1393            .data;
1394
1395        let requests = scan_requests_with_names(&rewritten)
1396            .into_iter()
1397            .collect::<HashMap<_, _>>();
1398        assert_eq!(requests["outer"].series_row_selector, None);
1399        assert_eq!(requests["inner"].series_row_selector, None);
1400    }
1401
1402    #[test]
1403    fn expression_subquery_can_start_its_own_single_evaluation() {
1404        let outer_provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1405        let inner_provider = Arc::new(mock_table_provider(RegionId::new(2, 1)));
1406        let plan = instant_with_expression_subquery(
1407            outer_provider,
1408            instant_plan_named(inner_provider, "inner", 1000),
1409            1000,
1410        );
1411        let rewritten = ScanHintRule
1412            .rewrite(plan, &OptimizerContext::default())
1413            .unwrap()
1414            .data;
1415
1416        let requests = scan_requests_with_names(&rewritten)
1417            .into_iter()
1418            .collect::<HashMap<_, _>>();
1419        assert_eq!(requests["outer"].series_row_selector, None);
1420        assert_eq!(
1421            requests["inner"].series_row_selector,
1422            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1423        );
1424    }
1425
1426    #[test]
1427    fn nested_single_outer_range_inner_does_not_set_inner_last_row() {
1428        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1429        let plan = instant_plan_named(provider.clone(), "nested", 2000);
1430        let plan = LogicalPlan::Extension(Extension {
1431            node: Arc::new(InstantManipulate::new(
1432                1000,
1433                1000,
1434                1000,
1435                1000,
1436                0,
1437                "ts".to_string(),
1438                vec![],
1439                Some("v0".to_string()),
1440                plan,
1441            )),
1442        });
1443        let rewritten = ScanHintRule
1444            .rewrite(plan, &OptimizerContext::default())
1445            .unwrap()
1446            .data;
1447        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
1448        assert_eq!(provider.scan_request().series_row_selector, None);
1449    }
1450
1451    #[test]
1452    fn nested_range_outer_single_inner_sets_inner_last_row() {
1453        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1454        let plan = instant_plan_named(provider.clone(), "nested", 1000);
1455        let plan = LogicalPlan::Extension(Extension {
1456            node: Arc::new(InstantManipulate::new(
1457                1000,
1458                2000,
1459                1000,
1460                1000,
1461                0,
1462                "ts".to_string(),
1463                vec![],
1464                Some("v0".to_string()),
1465                plan,
1466            )),
1467        });
1468        let rewritten = ScanHintRule
1469            .rewrite(plan, &OptimizerContext::default())
1470            .unwrap()
1471            .data;
1472        assert_eq!(
1473            scan_requests(&rewritten)[0].series_row_selector,
1474            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1475        );
1476    }
1477
1478    #[test]
1479    fn shared_provider_isolated_between_single_and_range_scans() {
1480        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1481        let plan = LogicalPlanBuilder::from(instant_plan(provider.clone(), 1000))
1482            .union(instant_plan(provider.clone(), 2000))
1483            .unwrap()
1484            .build()
1485            .unwrap();
1486        let rewritten = ScanHintRule
1487            .rewrite(plan, &OptimizerContext::default())
1488            .unwrap()
1489            .data;
1490        let requests = scan_requests(&rewritten);
1491
1492        assert_eq!(requests.len(), 2);
1493        assert_eq!(
1494            requests[0].series_row_selector,
1495            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
1496        );
1497        assert_eq!(requests[1].series_row_selector, None);
1498        assert_eq!(provider.scan_request().series_row_selector, None);
1499    }
1500
1501    #[test]
1502    fn set_order_hint() {
1503        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1504        let table_source = Arc::new(DefaultTableSource::new(provider.clone()));
1505        let plan = LogicalPlanBuilder::scan("t", table_source, None)
1506            .unwrap()
1507            .sort(vec![col("ts").sort(true, false)])
1508            .unwrap()
1509            .sort(vec![col("ts").sort(false, true)])
1510            .unwrap()
1511            .build()
1512            .unwrap();
1513
1514        let context = OptimizerContext::default();
1515        let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data;
1516
1517        // should read the first (with `.sort(true, false)`) sort option
1518        let scan_req = scan_requests(&rewritten)[0].clone();
1519        assert_eq!(
1520            OrderOption {
1521                name: "ts".to_string(),
1522                options: SortOptions {
1523                    descending: false,
1524                    nulls_first: false
1525                }
1526            },
1527            scan_req.output_ordering.as_ref().unwrap()[0]
1528        );
1529    }
1530
1531    #[test]
1532    fn set_time_series_row_selector_hint() {
1533        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
1534        let plan = last_value_aggregate(scan_plan(provider.clone(), "t"));
1535
1536        let context = OptimizerContext::default();
1537        let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data;
1538
1539        let scan_req = scan_requests(&rewritten)[0].clone();
1540        assert_eq!(
1541            scan_req.series_row_selector,
1542            Some(TimeSeriesRowSelector::LastRow { after_merge: false })
1543        );
1544    }
1545
1546    #[test]
1547    fn set_order_hint_sets_per_series_distribution_for_tsid_sort() {
1548        let provider = Arc::new(mock_table_provider_with_tsid(RegionId::new(1, 1)));
1549        let table_source = Arc::new(DefaultTableSource::new(provider.clone()));
1550        let plan = LogicalPlanBuilder::scan("t", table_source, None)
1551            .unwrap()
1552            .sort(vec![
1553                col(DATA_SCHEMA_TSID_COLUMN_NAME).sort(true, true),
1554                col("ts").sort(true, true),
1555            ])
1556            .unwrap()
1557            .build()
1558            .unwrap();
1559
1560        let context = OptimizerContext::default();
1561        let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data;
1562
1563        let scan_req = scan_requests(&rewritten)[0].clone();
1564        assert_eq!(
1565            scan_req.distribution,
1566            Some(TimeSeriesDistribution::PerSeries)
1567        );
1568    }
1569}
query/optimizer/scan_hint.rs

query/optimizer/
scan_hint.rs
