From 53ef80a51685f9da4cbcc2889a07c7964dbbd006 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:27:18 +0800 Subject: [PATCH] feat(query): capture native raw-series routes and extend MergeScan for typed frames Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/merge_scan.rs | 653 ++++++++++++- src/query/src/dist_plan/planner.rs | 1210 +++++++++++++++++++++++-- src/query/src/query_engine/state.rs | 3 + typos.toml | 7 +- 4 files changed, 1758 insertions(+), 115 deletions(-) diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 041fea160cf..d6a853620cb 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -69,6 +69,7 @@ use crate::dist_plan::dyn_filter_bridge::{ query_context_with_refreshed_initial_dyn_filter_regs, register_dyn_filter_subscribers_for_region, register_remote_dyn_filters, }; +use crate::dist_plan::planner::CaptureNativeRawRoute; use crate::dist_plan::{ FilterId, RemoteDynFilterProducerId, RemoteDynFilterRegistryLease, Subscriber, }; @@ -211,20 +212,41 @@ fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool /// Validates the remote schema before positional column handling. /// -/// Field metadata (indexing/encoding/compression hints) is non-semantic at -/// this boundary and is ignored: two fields are equal when their name, data -/// type, and nullability match (see [`fields_semantically_equal`]). A -/// timestamp timezone difference is the only intentional exception, accepted -/// when the fields share the same timestamp unit, distinct timezones, and -/// equal name and nullability. Top-level Arrow schema metadata is -/// non-semantic at this boundary as well. JSON columns are compared -/// semantically (see [`json_fields_compatible`]) because their wire and -/// decoded representations use different physical arrow types. +/// Ordinary SQL schemas retain the lenient comparison below: auxiliary field +/// metadata and top-level metadata are ignored, timestamps may differ only by +/// timezone, and JSON wire/decoded types use [`json_fields_compatible`]. +/// Schemas carrying any `greptime.series.*` top-level marker instead use exact +/// frame identity for all reserved top-level metadata and every field, including +/// nested field metadata. This compares identity, not frame validity. fn validate_remote_schema( expected: &ArrowSchema, actual: &ArrowSchema, source: &str, ) -> Result<()> { + let has_series_frame_metadata = expected + .metadata() + .keys() + .chain(actual.metadata().keys()) + .any(|key| key.starts_with("greptime.series.")); + if has_series_frame_metadata { + let reserved_metadata_matches = expected + .metadata() + .iter() + .filter(|(key, _)| key.starts_with("greptime.series.")) + .all(|(key, value)| actual.metadata().get(key) == Some(value)) + && actual + .metadata() + .iter() + .filter(|(key, _)| key.starts_with("greptime.series.")) + .all(|(key, value)| expected.metadata().get(key) == Some(value)); + if !reserved_metadata_matches || expected.fields() != actual.fields() { + return Err(remote_schema_mismatch(format!( + "MergeScan {source} series frame schema mismatch" + ))); + } + return Ok(()); + } + if expected.fields().len() != actual.fields().len() { return Err(remote_schema_mismatch(format!( "MergeScan {source} schema field count mismatch: expected {}, actual {}", @@ -473,6 +495,32 @@ pub struct MergeScanExec { captured_remote_dyn_filters: Arc>>, target_partition: usize, enable_per_region_metrics: bool, + /// Optional route provenance captured by the distributed planner from the + /// same physical route that selected `regions`. Ordinary execution never + /// reads it, so `None` keeps every existing plan unchanged. + native_raw_route: Option, +} + +/// Borrowed handoff view over the ordinary plan state a native raw consumer +/// needs to read the captured route. +/// +/// It borrows the existing [`MergeScanExec`] fields instead of duplicating them; +/// the output schema is available from the public `ExecutionPlan::schema()` of +/// the plan this source came from. +// TODO(narthex): remove this allowance once the remote native PromQL planner +// (#9670) consumes the captured route; this PR only captures it. +#[allow(dead_code)] +pub(crate) struct NativeRawSource<'a> { + /// Frozen route provenance in ordinary region order. + pub(crate) route: &'a CaptureNativeRawRoute, + /// The exact ordinary input logical plan sent to every region. + pub(crate) plan: &'a LogicalPlan, + /// The table the ordinary plan scans. + pub(crate) table: &'a TableName, + /// The handler that performs the region reads. + pub(crate) region_query_handler: &'a RegionQueryHandlerRef, + /// The query context of the ordinary execution. + pub(crate) query_ctx: &'a QueryContextRef, } impl std::fmt::Debug for MergeScanExec { @@ -556,6 +604,43 @@ impl MergeScanExec { captured_remote_dyn_filters: Arc::default(), target_partition, enable_per_region_metrics, + native_raw_route: None, + }) + } + + /// Attaches the native raw route captured from the same physical route that + /// selected this plan's regions. + /// + /// The artifact is provenance for a native raw consumer; ordinary execution + /// never reads it. + pub(crate) fn with_native_raw_route(mut self, native_raw_route: CaptureNativeRawRoute) -> Self { + self.native_raw_route = Some(native_raw_route); + self + } + + /// Returns the frozen region targets captured for this plan together with + /// the ordinary execution state a native raw consumer needs to read them. + /// + /// Returns `None` unless the planner proved the route and ordinary region + /// dispatch stays unchanged: a scan that already captured remote dynamic + /// filters, or a flow sink scan, keeps its ordinary path, so a consumer must + /// fall back as well. An assigned producer id alone is not active intent, + /// because the analyzer assigns one to every merge scan. + // TODO(narthex): remove this allowance once the remote native PromQL planner + // (#9670) consumes the captured route; this PR only captures it. + #[allow(dead_code)] + pub(crate) fn native_raw_source(&self) -> Option> { + if !self.captured_remote_dyn_filters.lock().unwrap().is_empty() || self.is_flow_sink_scan() + { + return None; + } + let route = self.native_raw_route.as_ref()?; + Some(NativeRawSource { + route, + plan: &self.plan, + table: &self.table, + region_query_handler: &self.region_query_handler, + query_ctx: &self.query_ctx, }) } @@ -1310,7 +1395,8 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use std::task::{Context, Poll}; - use arrow::array::{Int64Array, TimestampMillisecondArray}; + use arrow::array::{Array, Float64Array, Int64Array, TimestampMillisecondArray}; + use arrow::buffer::{Buffer, ScalarBuffer}; use arrow_schema::{DataType as TestArrowDataType, Field, TimeUnit}; use async_trait::async_trait; use common_base::Plugins; @@ -1337,10 +1423,15 @@ mod tests { }; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::series::frame::{self, RawLabelMode}; + use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; + use datatypes::series::samples::{ + DenseSamplesRef, FieldSpec, SampleRef, Samples, SeriesFields, SeriesSchema, + }; use datatypes::vectors::{Int64Vector, StringVector, TimestampMillisecondVector}; use futures_util::{Stream, TryStreamExt}; use session::ReadPreference; - use session::context::QueryContext; + use session::context::{QueryContext, QueryContextBuilder}; use session::query_id::QueryId; use snafu::IntoError; use table::table::scan::REGION_SCAN_EXEC_NAME; @@ -1350,7 +1441,7 @@ mod tests { use super::*; use crate::dist_plan::DynFilterRegistryManager; - use crate::options::QueryOptions; + use crate::options::{FLOW_SINK_TABLE_ID, QueryOptions}; use crate::query_engine::{QueryEngineContext, QueryEngineState}; use crate::region_query::RegionQueryHandler; @@ -1519,6 +1610,217 @@ mod tests { ); } + #[test] + fn merge_scan_ordinary_construction_has_no_native_raw_route() { + // The artifact is provenance only: an ordinary plan carries none and + // executes exactly as before. + let exec = merge_scan_exec_with_sorted_input(1, 1); + + assert!(exec.native_raw_source().is_none()); + } + + #[test] + fn merge_scan_keeps_captured_native_raw_route() { + let region_id = RegionId::new(1024, 1); + let exec = + merge_scan_exec_with_sorted_input(1, 1).with_native_raw_route(CaptureNativeRawRoute { + targets: vec![(region_id, test_target(1))], + }); + + let source = exec.native_raw_source().unwrap(); + assert_eq!(vec![(region_id, test_target(1))], source.route.targets); + // The bundle borrows the ordinary plan state, not a copy of it. + assert_eq!(1, source.plan.schema().fields().len()); + assert_eq!("table", source.table.table_name); + assert_eq!(1, exec.partition_count()); + } + + #[test] + fn merge_scan_declines_native_raw_route_for_captured_dyn_filters() { + // An assigned remote dynamic-filter producer id is identity, not intent: + // the analyzer assigns one to every merge scan, so a plan that captured no + // filters keeps the handoff view. + let region_id = RegionId::new(1024, 1); + let exec = remote_dyn_filter_test_exec( + Arc::new(TestRegionQueryHandler::default()), + QueryContext::arc(), + ) + .with_native_raw_route(CaptureNativeRawRoute { + targets: vec![(region_id, test_target(1))], + }); + + assert!(exec.captured_remote_dyn_filters().is_empty()); + assert!(exec.native_raw_source().is_some()); + + // A filter this producer accepted is really captured, and only then does the + // plan keep its ordinary dispatch. + install_remote_dyn_filter(&exec); + assert_eq!(1, exec.captured_remote_dyn_filters().len()); + assert!(exec.native_raw_source().is_none()); + + // A rewrite of the same scan keeps the frozen route and the captured filters, + // so the rewritten plan still declines the handoff view. + let plan: Arc = Arc::new(exec); + let replaced = plan + .replace_children( + vec![], + datafusion::physical_plan::ReplaceChildrenOptions::new( + datafusion::physical_plan::ChildrenPropertiesMode::Recompute, + ), + ) + .unwrap(); + let exec = replaced.downcast_ref::().unwrap(); + assert_eq!( + vec![(region_id, test_target(1))], + exec.native_raw_route.as_ref().unwrap().targets + ); + assert_eq!(1, exec.captured_remote_dyn_filters().len()); + assert!(exec.native_raw_source().is_none()); + } + + #[tokio::test] + async fn merge_scan_execution_ignores_native_raw_route() { + // Attaching the artifact must not change what an ordinary execution + // reads: every region is still dispatched through the handler's own + // target selection, with the ordinary read preference, and the same + // batches come back in the same order. + let regions = vec![ + RegionId::new(1024, 2), + RegionId::new(1024, 1), + RegionId::new(1024, 3), + ]; + let batch_for = |region_number: u32| { + let base = i64::from(region_number) * 10; + record_batch( + int64_schema(&["a", "b"]), + vec![ + Arc::new(Int64Vector::from_slice([base + 1])) as _, + Arc::new(Int64Vector::from_slice([base + 2])) as _, + ], + ) + }; + let responses = || -> Vec<(RegionId, RecordBatch)> { + regions + .iter() + .map(|region_id| (*region_id, batch_for(region_id.region_number()))) + .collect() + }; + // Region striping keeps the selection order, so this is the whole output. + let expected = [(21, 22), (11, 12), (31, 32)]; + + let baseline_handler = Arc::new(TestRegionQueryHandler::new(responses())); + let baseline = collect_merge_scan(merge_scan_exec_with_handler( + regions.clone(), + expected_int64_schema(), + baseline_handler.clone(), + 1, + )) + .await + .unwrap(); + assert_eq!(expected.len(), baseline.len()); + for (batch, values) in baseline.iter().zip(expected) { + assert_int64_batch(batch, values); + } + + // Deliberately different peers, and in a deliberately different order: + // an execution that read the artifact would dispatch to these. + let handler = Arc::new(TestRegionQueryHandler::new(responses())); + let with_route = collect_merge_scan( + merge_scan_exec_with_handler( + regions.clone(), + expected_int64_schema(), + handler.clone(), + 1, + ) + .with_native_raw_route(CaptureNativeRawRoute { + targets: vec![ + (regions[2], test_target(7)), + (regions[0], test_target(8)), + (regions[1], test_target(9)), + ], + }), + ) + .await + .unwrap(); + + assert_eq!(baseline, with_route); + assert_eq!(expected.len(), with_route.len()); + for (batch, values) in with_route.iter().zip(expected) { + assert_int64_batch(batch, values); + } + for exec_handler in [&baseline_handler, &handler] { + assert_eq!(regions.len(), exec_handler.select_target_calls()); + assert_eq!( + vec![ReadPreference::Leader; regions.len()], + exec_handler.read_preferences() + ); + // The ordinary plan dispatches to the target its handler selected, + // never to a peer from the artifact. + assert_eq!(vec![1; regions.len()], exec_handler.do_get_targets()); + } + } + + #[test] + fn merge_scan_with_new_children_preserves_native_raw_route() { + let region_id = RegionId::new(1024, 1); + let plan: Arc = Arc::new( + merge_scan_exec_with_sorted_input(1, 1).with_native_raw_route(CaptureNativeRawRoute { + targets: vec![(region_id, test_target(3))], + }), + ); + + let replaced = plan + .replace_children( + vec![], + datafusion::physical_plan::ReplaceChildrenOptions::new( + datafusion::physical_plan::ChildrenPropertiesMode::Recompute, + ), + ) + .unwrap(); + let exec = replaced.downcast_ref::().unwrap(); + assert_eq!( + vec![(region_id, test_target(3))], + exec.native_raw_source().unwrap().route.targets + ); + } + + #[test] + fn merge_scan_declines_native_raw_route_for_flow_sink_scan() { + // A flow sink read keeps its ordinary dispatch, so the handoff view is + // not offered even though the artifact exists for these regions. + let region_id = RegionId::new(1024, 1); + let query_ctx = Arc::new( + QueryContextBuilder::default() + .set_extension( + FLOW_SINK_TABLE_ID.to_string(), + region_id.table_id().to_string(), + ) + .build(), + ); + let plan = LogicalPlanBuilder::empty(true).build().unwrap(); + let schema = plan.schema().as_arrow().clone(); + + let exec = MergeScanExec::new( + &SessionStateBuilder::new().build(), + TableName::new("catalog", "schema", "table"), + vec![region_id], + plan, + &schema, + Arc::new(TestRegionQueryHandler::default()), + query_ctx, + 1, + None, + false, + ) + .unwrap() + .with_native_raw_route(CaptureNativeRawRoute { + targets: vec![(region_id, test_target(1))], + }); + + assert!(exec.is_flow_sink_scan()); + assert!(exec.native_raw_source().is_none()); + } + #[test] fn merge_scan_advertises_ordering_when_each_partition_reads_at_most_one_region() { let exec = merge_scan_exec_with_sorted_input(3, 3); @@ -2045,6 +2347,9 @@ mod tests { #[derive(Default)] struct TestRegionQueryHandler { responses: HashMap, + select_target_calls: AtomicUsize, + read_preferences: Mutex>, + do_get_targets: Mutex>, } impl TestRegionQueryHandler { @@ -2061,7 +2366,10 @@ mod tests { ) }) .collect(); - Self { responses } + Self { + responses, + ..Default::default() + } } fn with_responses( @@ -2085,7 +2393,22 @@ mod tests { ) }) .collect(); - Self { responses } + Self { + responses, + ..Default::default() + } + } + + fn select_target_calls(&self) -> usize { + self.select_target_calls.load(Ordering::SeqCst) + } + + fn read_preferences(&self) -> Vec { + self.read_preferences.lock().unwrap().clone() + } + + fn do_get_targets(&self) -> Vec { + self.do_get_targets.lock().unwrap().clone() } } @@ -2469,17 +2792,20 @@ mod tests { impl RegionQueryHandler for TestRegionQueryHandler { async fn select_target( &self, - _read_preference: ReadPreference, + read_preference: ReadPreference, _region_id: RegionId, ) -> crate::error::Result { + self.select_target_calls.fetch_add(1, Ordering::SeqCst); + self.read_preferences.lock().unwrap().push(read_preference); Ok(test_target(1)) } async fn do_get( &self, - _target: &crate::region_query::RegionQueryTarget, + target: &crate::region_query::RegionQueryTarget, request: common_query::request::QueryRequest, ) -> crate::error::Result { + self.do_get_targets.lock().unwrap().push(target.peer().id); let response = self .responses .get(&request.region_id) @@ -3031,6 +3357,301 @@ mod tests { assert_eq!(merge_scan_schema_error_count_for_test(), errors_before + 1); } + #[tokio::test] + async fn merge_scan_remote_schema_identity_rejects_empty_raw_frame_metadata_mismatch() { + let region_id = RegionId::new(1024, 1); + let raw = raw_test_series_schema(); + let frame = frame::raw_series_frame_schema(&raw, RawLabelMode::Key).unwrap(); + let metadata = frame.metadata().clone(); + let mut cases = Vec::new(); + for (key, value) in [ + ("greptime.series.layout", "v0"), + ("greptime.series.content", "other"), + ("greptime.series.labels", "explicit"), + ("greptime.series.time", "other"), + ] { + let mut changed = metadata.clone(); + changed.insert(key.to_string(), value.to_string()); + cases.push((metadata.clone(), changed)); + } + let mut extra = metadata.clone(); + extra.insert("greptime.series.grid.future".to_string(), "v1".to_string()); + cases.push((metadata.clone(), extra)); + let mut removed = metadata.clone(); + removed.remove("greptime.series.layout"); + cases.push((metadata.clone(), removed)); + cases.push((metadata.clone(), StdHashMap::new())); + let mut actual_only = StdHashMap::new(); + actual_only.insert("greptime.series.content".to_string(), "raw".to_string()); + cases.push((StdHashMap::new(), actual_only)); + let mut expected_unknown = metadata.clone(); + expected_unknown.insert("greptime.series.future".to_string(), "v1".to_string()); + cases.push((expected_unknown, metadata.clone())); + + for (case_index, (expected_metadata, actual_metadata)) in cases.into_iter().enumerate() { + let expected = + ArrowSchema::new_with_metadata(frame.fields().clone(), expected_metadata); + let advertised = Arc::new( + Schema::try_from(Arc::new(ArrowSchema::new_with_metadata( + frame.fields().clone(), + actual_metadata, + ))) + .unwrap(), + ); + let errors_before = merge_scan_schema_error_count_for_test(); + let result = collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected, + Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + advertised, + vec![], + )])), + 1, + )) + .await; + assert!(result.is_err(), "metadata case {case_index} must fail"); + assert_eq!( + merge_scan_schema_error_count_for_test(), + errors_before + 1, + "metadata case {case_index}" + ); + } + } + + fn raw_test_series_schema() -> SeriesSchema { + SeriesSchema { + fields: vec![ + FieldSpec { + name: Arc::from("value"), + }, + FieldSpec { + name: Arc::from("other"), + }, + ] + .into(), + eval_grid: None, + } + } + + #[tokio::test] + async fn merge_scan_remote_schema_identity_rejects_later_raw_frame_inner_metadata_mismatch() { + let region_id = RegionId::new(1024, 1); + let raw = raw_test_series_schema(); + let mode = RawLabelMode::Key; + let frame = frame::raw_series_frame_schema(&raw, mode).unwrap(); + let mut changed_metadata = frame.metadata().clone(); + changed_metadata.insert("greptime.series.marker".to_string(), "changed".to_string()); + let changed_schema = Arc::new(ArrowSchema::new_with_metadata( + frame.fields().clone(), + changed_metadata, + )); + let chunk = |values: [f64; 2], times: [i64; 2], last| { + RawSeriesChunk::try_new( + &raw, + Buffer::from(vec![1u8, 2, 3]), + None, + RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).unwrap(), + SeriesFields::try_from_vec(vec![ + Samples::from_float(Float64Array::from(values.to_vec())), + Samples::from_float(Float64Array::from(vec![4.0, 5.0])), + ]) + .unwrap(), + last, + ) + .unwrap() + }; + let first = frame::pack_raw_series_chunk(&raw, mode, &chunk([1.0, 2.0], [10, 20], false)) + .unwrap() + .remove(0); + let mut second = + frame::pack_raw_series_chunk(&raw, mode, &chunk([3.0, 4.0], [30, 40], true)) + .unwrap() + .remove(0); + second = DfRecordBatch::try_new(changed_schema, second.columns().to_vec()).unwrap(); + let advertised = Arc::new(Schema::try_from(frame.clone()).unwrap()); + let first = RecordBatch::from_df_record_batch(advertised.clone(), first); + let second = RecordBatch::from_df_record_batch(advertised.clone(), second); + let errors_before = merge_scan_schema_error_count_for_test(); + let exec = merge_scan_exec_with_handler( + vec![region_id], + frame.as_ref().clone(), + Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + advertised, + vec![Ok(first), Ok(second)], + )])), + 1, + ); + let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); + let emitted = stream.next().await.unwrap().unwrap(); + let decoded = frame::unpack_raw_series_frame(&raw, mode, &emitted).unwrap(); + assert_eq!(decoded[0].time.timestamp(0), Some(10)); + assert!(matches!( + decoded[0].fields.get(0).unwrap().sample(0), + SampleRef::Float(1.0) + )); + assert!(stream.next().await.unwrap().is_err()); + assert!(stream.next().await.is_none()); + assert_eq!(merge_scan_schema_error_count_for_test(), errors_before + 1); + } + + #[tokio::test] + async fn merge_scan_remote_schema_identity_rejects_raw_frame_field_metadata_mismatch() { + let region_id = RegionId::new(1024, 1); + let raw = raw_test_series_schema(); + let expected = frame::raw_series_frame_schema(&raw, RawLabelMode::Key).unwrap(); + let frame_metadata = expected.metadata().clone(); + let fields = expected.fields().to_vec(); + let mut field_cases = Vec::new(); + let mut direct_metadata = fields.clone(); + direct_metadata[0] = Arc::new(fields[0].as_ref().clone().with_metadata(StdHashMap::from( + [("semantic".to_string(), "changed".to_string())], + ))); + field_cases.push(direct_metadata); + + let mut nested_metadata = fields.clone(); + let TestArrowDataType::Struct(children) = fields[2].data_type() else { + panic!("raw frame time field must be a struct"); + }; + let mut children = children.to_vec(); + children[0] = children[0] + .as_ref() + .clone() + .with_metadata(StdHashMap::from([( + "semantic".to_string(), + "changed".to_string(), + )])) + .into(); + nested_metadata[2] = fields[2] + .as_ref() + .clone() + .with_data_type(TestArrowDataType::Struct(children.into())) + .into(); + field_cases.push(nested_metadata); + + let mut changed_name = fields.clone(); + changed_name[0] = fields[0] + .as_ref() + .clone() + .with_name("other".to_string()) + .into(); + field_cases.push(changed_name); + let mut changed_type = fields.clone(); + changed_type[0] = fields[0] + .as_ref() + .clone() + .with_data_type(TestArrowDataType::LargeBinary) + .into(); + field_cases.push(changed_type); + let mut changed_nullability = fields; + changed_nullability[0] = changed_nullability[0] + .as_ref() + .clone() + .with_nullable(true) + .into(); + field_cases.push(changed_nullability); + + for (case_index, actual_fields) in field_cases.into_iter().enumerate() { + let advertised = Arc::new( + Schema::try_from(Arc::new(ArrowSchema::new_with_metadata( + actual_fields, + frame_metadata.clone(), + ))) + .unwrap(), + ); + let errors_before = merge_scan_schema_error_count_for_test(); + let result = collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected.as_ref().clone(), + Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + advertised, + vec![], + )])), + 1, + )) + .await; + assert!(result.is_err(), "field schema case {case_index} must fail"); + assert_eq!( + merge_scan_schema_error_count_for_test(), + errors_before + 1, + "field schema case {case_index}" + ); + } + } + + #[tokio::test] + async fn merge_scan_remote_schema_identity_preserves_valid_raw_frame_payload() { + let region_id = RegionId::new(1024, 1); + let raw = raw_test_series_schema(); + let expected_frame = frame::raw_series_frame_schema(&raw, RawLabelMode::Key).unwrap(); + let mut expected_metadata = expected_frame.metadata().clone(); + expected_metadata.insert("application.note".to_string(), "expected".to_string()); + let expected = Arc::new(ArrowSchema::new_with_metadata( + expected_frame.fields().clone(), + expected_metadata, + )); + let remote_schema = Arc::new(Schema::try_from(expected_frame.clone()).unwrap()); + let chunk = RawSeriesChunk::try_new( + &raw, + Buffer::from(vec![1u8, 2, 3]), + None, + RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap(), + SeriesFields::try_from_vec(vec![ + Samples::from_float(Float64Array::from(vec![Some(-0.0), None])), + Samples::from_float(Float64Array::from(vec![ + Some(f64::from_bits(0x7ff8_0000_0000_0042)), + Some(3.5), + ])), + ]) + .unwrap(), + true, + ) + .unwrap(); + let packed = frame::pack_raw_series_chunk(&raw, RawLabelMode::Key, &chunk).unwrap(); + let inner = packed.into_iter().next().unwrap(); + let original_columns = inner.columns().to_vec(); + let batch = RecordBatch::from_df_record_batch(remote_schema.clone(), inner); + let result = collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected.as_ref().clone(), + Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + remote_schema, + vec![Ok(batch)], + )])), + 1, + )) + .await + .unwrap(); + assert_eq!(result.len(), 1); + for (original, returned) in original_columns.iter().zip(result[0].columns()) { + assert!(Arc::ptr_eq(original, returned)); + } + let decoded = frame::unpack_raw_series_frame(&raw, RawLabelMode::Key, &result[0]); + let decoded = decoded.unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded[0].primary_key.as_slice(), &[1, 2, 3]); + assert!(decoded[0].is_last); + assert_eq!(decoded[0].time.timestamp(0), Some(10)); + assert_eq!(decoded[0].time.timestamp(1), Some(20)); + match decoded[0].fields.get(0).unwrap().dense().unwrap() { + DenseSamplesRef::Float(values) => { + assert_eq!(values.value(0).to_bits(), (-0.0f64).to_bits()); + assert!(values.is_null(1)); + } + _ => panic!("expected float samples"), + } + match decoded[0].fields.get(1).unwrap().dense().unwrap() { + DenseSamplesRef::Float(values) => { + assert_eq!(values.value(0).to_bits(), 0x7ff8_0000_0000_0042); + assert_eq!(values.value(1), 3.5); + } + _ => panic!("expected float samples"), + } + } + #[tokio::test] async fn merge_scan_remote_schema_identity_allows_top_level_metadata_version_mismatch() { let fields = expected_int64_schema().fields().clone(); diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index 34c0994f9e0..f67559ffa96 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -16,12 +16,14 @@ use std::sync::Arc; -use ahash::HashMap; +use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use arrow_schema::SortOptions; use async_trait::async_trait; use catalog::CatalogManagerRef; -use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, MITO_ENGINE}; +use common_meta::rpc::router::RegionRoute; use common_telemetry::debug; +use common_time::timestamp::{TimeUnit, Timestamp}; use datafusion::catalog::Session; use datafusion::common::Result; use datafusion::datasource::DefaultTableSource; @@ -34,8 +36,10 @@ use datafusion_common::{DataFusionError, TableReference}; use datafusion_expr::{LogicalPlan, UserDefinedLogicalNode}; use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; use datatypes::prelude::ConcreteDataType; -use partition::expr::PartitionExpr; +use datatypes::value::Value; +use partition::expr::{Operand, PartitionExpr, RestrictedOp}; use partition::manager::{PartitionRuleManagerRef, create_partitions_from_region_routes}; +use partition::overlap::associate_from_to; use session::context::QueryContext; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; @@ -51,7 +55,7 @@ use crate::dist_plan::merge_sort::{MergeSortExec, MergeSortLogicalPlan}; use crate::dist_plan::region_pruner::ConstraintPruner; use crate::error::{CatalogSnafu, PartitionRuleManagerSnafu, TableNotFoundSnafu}; use crate::optimizer::expr_depth; -use crate::region_query::RegionQueryHandlerRef; +use crate::region_query::{RegionQueryHandlerRef, RegionQueryTarget}; /// Planner for converting merge sort logical plan to physical plan. /// @@ -148,6 +152,9 @@ pub struct DistExtensionPlanner { partition_rule_manager: PartitionRuleManagerRef, region_query_handler: RegionQueryHandlerRef, enable_per_region_metrics: bool, + /// Gates the quadratic route-overlap proof with the existing native PromQL option. + /// Ordinary region selection runs identically either way. + capture_native_raw_route: bool, } impl DistExtensionPlanner { @@ -156,12 +163,14 @@ impl DistExtensionPlanner { partition_rule_manager: PartitionRuleManagerRef, region_query_handler: RegionQueryHandlerRef, enable_per_region_metrics: bool, + capture_native_raw_route: bool, ) -> Self { Self { catalog_manager, partition_rule_manager, region_query_handler, enable_per_region_metrics, + capture_native_raw_route, } } } @@ -201,7 +210,8 @@ impl ExtensionPlanner for DistExtensionPlanner { return fallback(optimized_plan).await; }; - let Ok(regions) = self.get_regions(&table_name, input_plan).await else { + let Ok((regions, native_raw_route)) = self.get_regions(&table_name, input_plan).await + else { // no peers found, going to execute them locally return fallback(optimized_plan).await; }; @@ -232,6 +242,12 @@ impl ExtensionPlanner for DistExtensionPlanner { merge_scan.remote_dyn_filter_producer_id(), self.enable_per_region_metrics, )?; + // Captured provenance only: ordinary execution of `MergeScanExec` never + // reads it, so a missing capture keeps the plan exactly as before. + let merge_scan_plan = match native_raw_route { + Some(native_raw_route) => merge_scan_plan.with_native_raw_route(native_raw_route), + None => merge_scan_plan, + }; Ok(Some(Arc::new(merge_scan_plan) as _)) } } @@ -244,16 +260,26 @@ impl DistExtensionPlanner { Ok(extractor.table_name) } + /// Selects the regions of the ordinary distributed plan and captures the + /// native raw route of the same physical route, if it can be proven. + /// + /// Every ordinary selection path (pruned and all fallbacks) returns the + /// regions a `MergeScanExec` would use, plus the frozen region targets + /// captured in that exact order. The capture is provenance for a native raw + /// consumer and is `None` when it is disabled or the route cannot be proven; + /// capturing never changes the ordinary region selection or its error + /// handling. async fn get_regions( &self, table_name: &TableName, logical_plan: &LogicalPlan, - ) -> Result> { + ) -> Result<(Vec, Option)> { let mut extractor = TableScanExtractor::default(); let _ = logical_plan.visit(&mut extractor)?; // Resolving by name again could bind an authorized scan to a replacement table. - let table = match extractor.captured_table { - Some(table) => table, + let captured_table = extractor.captured_table; + let table = match &captured_table { + Some(table) => table.clone(), None => self .catalog_manager .table( @@ -293,87 +319,107 @@ impl DistExtensionPlanner { partition_columns, all_regions, ); - if partition_columns.is_empty() { - return Ok(all_regions); - } - // Extract predicates from logical plan - let partition_expressions = match PredicateExtractor::extract_partition_expressions( - logical_plan, - &partition_columns, - ) { - Ok(expressions) => expressions, - Err(err) => { - common_telemetry::debug!( - "Failed to extract partition expressions for table {} (id: {}), using all regions: {:?}", - table_name, - table.table_info().table_id(), - err - ); - return Ok(all_regions); + // Ordinary region selection. Every fallback keeps all regions, exactly + // as before the native raw route existed. + let regions = 'pruning: { + if partition_columns.is_empty() { + break 'pruning all_regions; } - }; + // Extract predicates from logical plan + let partition_expressions = match PredicateExtractor::extract_partition_expressions( + logical_plan, + &partition_columns, + ) { + Ok(expressions) => expressions, + Err(err) => { + common_telemetry::debug!( + "Failed to extract partition expressions for table {} (id: {}), using all regions: {:?}", + table_name, + table.table_info().table_id(), + err + ); + break 'pruning all_regions; + } + }; - if partition_expressions.is_empty() { - return Ok(all_regions); - } + if partition_expressions.is_empty() { + break 'pruning all_regions; + } - let Some(partition_column_types) = self - .partition_column_types_for_pruning( - table_name, - table_info.as_ref(), - physical_table_id, + let Some(partition_column_types) = self + .partition_column_types_for_pruning( + table_name, + table_info.as_ref(), + physical_table_id, + &partition_expressions, + &all_regions, + ) + .await + else { + break 'pruning all_regions; + }; + + // Get partition information for the table if partition rule manager is available + let partitions = match create_partitions_from_region_routes( + table_info.table_id(), + &physical_table_route.region_routes, + ) { + Ok(partitions) => partitions, + Err(err) => { + common_telemetry::debug!( + "Failed to get partition information for table {}, using all regions: {:?}", + table_name, + err + ); + break 'pruning all_regions; + } + }; + if partitions.is_empty() { + break 'pruning all_regions; + } + // Apply region pruning based on partition rules + let pruned_regions = match ConstraintPruner::prune_regions( &partition_expressions, - &all_regions, + &partitions, + partition_column_types, + ) { + Ok(regions) => regions, + Err(err) => { + common_telemetry::debug!( + "Failed to prune regions for table {}, using all regions: {:?}", + table_name, + err + ); + break 'pruning all_regions; + } + }; + + common_telemetry::debug!( + "Region pruning for table {}: {} partition expressions applied, pruned from {} to {} regions", + table_name, + partition_expressions.len(), + all_regions.len(), + pruned_regions.len() + ); + + pruned_regions + }; + + // Captured from the same physical route that selected the regions above, + // in the same order. This never mutates or reorders `regions`; the + // experimental gate only skips the capture. + let native_raw_route = if self.capture_native_raw_route { + capture_native_raw_route( + captured_table.as_ref(), + &table_info, + physical_table_id, + &physical_table_route.region_routes, + ®ions, ) - .await - else { - return Ok(all_regions); + } else { + None }; - - // Get partition information for the table if partition rule manager is available - let partitions = match create_partitions_from_region_routes( - table_info.table_id(), - &physical_table_route.region_routes, - ) { - Ok(partitions) => partitions, - Err(err) => { - common_telemetry::debug!( - "Failed to get partition information for table {}, using all regions: {:?}", - table_name, - err - ); - return Ok(all_regions); - } - }; - if partitions.is_empty() { - return Ok(all_regions); - } - // Apply region pruning based on partition rules - let pruned_regions = match ConstraintPruner::prune_regions( - &partition_expressions, - &partitions, - partition_column_types, - ) { - Ok(regions) => regions, - Err(err) => { - common_telemetry::debug!( - "Failed to prune regions for table {}, using all regions: {:?}", - table_name, - err - ); - return Ok(all_regions); - } - }; - - common_telemetry::debug!( - "Region pruning for table {}: {} partition expressions applied, pruned from {} to {} regions", - table_name, - partition_expressions.len(), - all_regions.len(), - pruned_regions.len() - ); - - Ok(pruned_regions) + Ok((regions, native_raw_route)) } /// Resolves the partition-column types that are safe to use for region pruning. @@ -472,6 +518,199 @@ fn partition_column_types(table_info: &TableInfo) -> Vec<(String, ConcreteDataTy .collect() } +/// Frozen route provenance of an ordinary distributed plan for a native raw +/// consumer. +/// +/// `targets` pairs each region selected by the ordinary plan with the leader +/// target of the physical route that selected it, in the exact order the +/// ordinary `MergeScanExec` uses. It is only built when the routed partition +/// expressions express an explicit, non-empty, pairwise-disjoint millisecond +/// timestamp domain on the table captured from the authorized scan, and +/// ordinary execution never reads it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct CaptureNativeRawRoute { + /// `targets[i]` is the frozen leader target of the i-th selected region. + pub(crate) targets: Vec<(RegionId, RegionQueryTarget)>, +} + +/// Freezes the leader target of every `selected_regions` entry from +/// `region_routes`, the physical route that selected them. +/// +/// Returns `None` whenever that route cannot be proven to be an explicit, +/// non-empty, pairwise-disjoint timestamp domain. Declining is not an error and +/// never changes the ordinary region selection or its error handling: the +/// capture is provenance that the ordinary plan already ignores. +fn capture_native_raw_route( + captured_table: Option<&TableRef>, + table_info: &TableInfo, + physical_table_id: u32, + region_routes: &[RegionRoute], + selected_regions: &[RegionId], +) -> Option { + // Only the table captured from the authorized scan can prove the route: a + // name-only lookup may have bound the plan to a replacement table. + let captured_table = captured_table?; + if captured_table.table_type() != TableType::Base + || table_info.meta.engine != MITO_ENGINE + // A logical table shares another physical table's route, so its region + // ids do not identify the routes below. + || table_info.table_id() != physical_table_id + { + return None; + } + if selected_regions.is_empty() { + return None; + } + + let timestamp_column = table_info.meta.schema.timestamp_column()?; + // The native raw consumer currently supports millisecond timestamps only. + if timestamp_column.data_type.as_timestamp()?.unit() != TimeUnit::Millisecond { + return None; + } + // A nullable timestamp column can hold `NULL`s that belong to every + // partition domain, so the domains would not stay disjoint. + if timestamp_column.is_nullable() { + return None; + } + + let selected_expressions = + selected_route_expressions(table_info.table_id(), region_routes, selected_regions)?; + // The supported grammar proves every selected domain non-empty on its own; + // see [`explicit_timestamp_bound`] for the exact shapes that are accepted. + if !selected_expressions.iter().all(|expr| { + explicit_timestamp_bound(expr, ×tamp_column.name) + .is_some_and(|(op, value)| is_nonempty_timestamp_bound(&op, value)) + }) { + return None; + } + // `associate_from_to` proves only that the selected domains are pairwise + // disjoint: it normalizes the values to order-preserving ordinals and + // reports, per expression, which expressions it overlaps. Requiring each row + // to name exactly its own index is what makes the domains disjoint from one + // another, on top of every one of them being non-empty. + let overlaps = associate_from_to(&selected_expressions, &selected_expressions).ok()?; + if overlaps + .iter() + .enumerate() + .any(|(index, overlaps)| overlaps.len() != 1 || overlaps[0] != index) + { + return None; + } + + let mut route_by_region_number = HashMap::with_capacity(region_routes.len()); + for route in region_routes { + // Two routes for one region number cannot be mapped one-to-one. + if route_by_region_number + .insert(route.region.id.region_number(), route) + .is_some() + { + return None; + } + } + + let mut seen = HashSet::with_capacity(selected_regions.len()); + let mut targets = Vec::with_capacity(selected_regions.len()); + for region_id in selected_regions { + // A duplicated region would be dispatched twice. + if !seen.insert(*region_id) { + return None; + } + let route = route_by_region_number.get(®ion_id.region_number())?; + // The lookup is by region number, but only the full region id identifies + // this table's route: another table's route can carry the same region + // number, and freezing its target would dispatch a foreign region. + if route.region.id != *region_id { + return None; + } + // Reads go to the region leader, so a staging or downgrading leader + // must not be frozen as a query target. + if route.is_leader_staging() || route.is_leader_downgrading() { + return None; + } + let leader_peer = route.leader_peer.clone()?; + targets.push((*region_id, RegionQueryTarget::new(leader_peer))); + } + + Some(CaptureNativeRawRoute { targets }) +} + +/// Collects the explicit partition expression of every selected region. +/// +/// The expressions are decoded from the same router routes as the region ids and +/// are returned in the ordinary selection order, so the proof checks exactly the +/// regions the ordinary plan scans. A region without an explicit expression (for +/// example a default region) declines the capture. +fn selected_route_expressions( + table_id: u32, + region_routes: &[RegionRoute], + selected_regions: &[RegionId], +) -> Option> { + let partitions = create_partitions_from_region_routes(table_id, region_routes).ok()?; + let mut expressions = HashMap::with_capacity(partitions.len()); + for partition in partitions { + if let Some(expression) = partition.partition_expr { + expressions.insert(partition.id, expression); + } + } + selected_regions + .iter() + .map(|region_id| expressions.get(region_id).cloned()) + .collect() +} + +/// Returns the standalone `Column ↔ Timestamp(ms)` comparison of `expr` in +/// column-against-value order, or `None` if `expr` is not exactly one such +/// comparison. +/// +/// The route decoder canonicalizes reversed operands. Only one comparison is +/// supported: conjunctions can contain empty discrete domains that the shared +/// overlap helper cannot prove, and remain outside this proof grammar. +fn explicit_timestamp_bound( + expr: &PartitionExpr, + timestamp_column: &str, +) -> Option<(RestrictedOp, Timestamp)> { + let (Operand::Column(column), Operand::Value(value)) = (&*expr.lhs, &*expr.rhs) else { + return None; + }; + let op = expr.op.clone(); + if column.as_str() != timestamp_column { + return None; + } + if !matches!( + op, + RestrictedOp::Eq + | RestrictedOp::Lt + | RestrictedOp::LtEq + | RestrictedOp::Gt + | RestrictedOp::GtEq + ) { + return None; + } + match value { + Value::Timestamp(ts) if ts.unit() == TimeUnit::Millisecond => Some((op, *ts)), + _ => None, + } +} + +/// Returns whether `op` against `value` is satisfied by at least one +/// representable millisecond timestamp, i.e. the domain is not empty. +/// +/// Only the strict comparisons against the representable range bounds can be +/// unsatisfiable, and the bound check uses the timestamp value directly instead +/// of any interval arithmetic. +fn is_nonempty_timestamp_bound(op: &RestrictedOp, value: Timestamp) -> bool { + match op { + RestrictedOp::Eq => { + value >= Timestamp::MIN_MILLISECOND && value <= Timestamp::MAX_MILLISECOND + } + RestrictedOp::Lt => value > Timestamp::MIN_MILLISECOND, + RestrictedOp::LtEq => value >= Timestamp::MIN_MILLISECOND, + RestrictedOp::Gt => value < Timestamp::MAX_MILLISECOND, + RestrictedOp::GtEq => value <= Timestamp::MAX_MILLISECOND, + RestrictedOp::NotEq | RestrictedOp::And | RestrictedOp::Or => false, + } +} + /// Extract the scan name and captured table identity from a logical plan. #[derive(Default)] struct TableScanExtractor { @@ -547,22 +786,25 @@ mod tests { use async_trait::async_trait; use catalog::memory::MemoryCatalogManager; use catalog::{CatalogManagerRef, RegisterTableRequest}; - use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, MITO_ENGINE}; use common_meta::cache::new_table_route_cache; use common_meta::key::TableMetadataManager; use common_meta::key::table_route::TableRouteValue; use common_meta::kv_backend::memory::MemoryKvBackend; - use common_meta::rpc::router::{Region, RegionRoute}; + use common_meta::peer::Peer; + use common_meta::rpc::router::{LeaderState, Region, RegionRoute}; use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; + use common_time::timestamp::Timestamp; use datafusion::datasource::DefaultTableSource; - use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col as df_col, lit}; + use datafusion_common::ScalarValue; + use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, TableSource, col as df_col, lit}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::value::Value; use moka::future::CacheBuilder; use partition::cache::new_partition_info_cache; - use partition::expr::{PartitionExpr, col as partition_col}; + use partition::expr::{Operand, PartitionExpr, RestrictedOp, col as partition_col}; use partition::manager::PartitionRuleManager; use session::ReadPreference; use store_api::storage::RegionId; @@ -571,7 +813,7 @@ mod tests { use table::table_name::TableName; use table::test_util::EmptyTable; - use super::DistExtensionPlanner; + use super::{CaptureNativeRawRoute, DistExtensionPlanner, capture_native_raw_route}; use crate::region_query::{RegionQueryHandler, RegionQueryTarget}; const LOGICAL_TABLE_ID: u32 = 1024; @@ -749,6 +991,7 @@ mod tests { partition_rule_manager, Arc::new(UnusedRegionQueryHandler), false, + false, ); let table_source = Arc::new(DefaultTableSource::new(Arc::new( DfTableProviderAdapter::new(logical_table), @@ -786,10 +1029,11 @@ mod tests { async fn region_routing_uses_captured_table() { let (mut planner, plan, table_name) = planner_and_plan(vec![0], vec![None]).await; planner.catalog_manager = MemoryCatalogManager::with_default_setup(); - assert_eq!( - vec![RegionId::new(LOGICAL_TABLE_ID, 1)], - planner.get_regions(&table_name, &plan).await.unwrap() - ); + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + assert_eq!(vec![RegionId::new(LOGICAL_TABLE_ID, 1)], regions); + // The captured table is not the physical route owner, so the reason for + // using it is not provable and the capture is declined. + assert!(native_raw_route.is_none()); } #[tokio::test] @@ -797,10 +1041,9 @@ mod tests { let (planner, plan, table_name) = planner_and_plan(vec![0, 1], physical_partition_expressions()).await; - assert_eq!( - vec![RegionId::new(LOGICAL_TABLE_ID, 1)], - planner.get_regions(&table_name, &plan).await.unwrap() - ); + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + assert_eq!(vec![RegionId::new(LOGICAL_TABLE_ID, 1)], regions); + assert!(native_raw_route.is_none()); } #[tokio::test] @@ -808,7 +1051,9 @@ mod tests { let (planner, plan, table_name) = planner_and_plan(vec![0], physical_partition_expressions()).await; - assert_all_logical_regions(planner.get_regions(&table_name, &plan).await.unwrap()); + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + assert_all_logical_regions(regions); + assert!(native_raw_route.is_none()); } #[tokio::test] @@ -817,7 +1062,776 @@ mod tests { expressions[1] = None; let (planner, plan, table_name) = planner_and_plan(vec![0, 1], expressions).await; - assert_all_logical_regions(planner.get_regions(&table_name, &plan).await.unwrap()); + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + assert_all_logical_regions(regions); + assert!(native_raw_route.is_none()); + } + + // ------------------------- native raw route capture ---------------------- + + /// Table id of the base table used by the native raw route tests. + const NATIVE_TABLE_ID: u32 = 4096; + /// Timestamp split between the explicit domains of the native route tests. + const NATIVE_SPLIT_MS: i64 = 1_700_000_000_000; + + /// Table info of a base table with an explicit timestamp column. + fn native_table_info_with( + table_id: u32, + name: &str, + engine: &str, + timestamp_type: ConcreteDataType, + timestamp_nullable: bool, + partition_key_indices: Vec, + ) -> TableInfo { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("ts", timestamp_type, timestamp_nullable).with_time_index(true), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ])); + let meta = TableMeta { + schema, + primary_key_indices: vec![], + value_indices: vec![], + engine: engine.to_string(), + next_column_id: 2, + options: Default::default(), + created_on: Default::default(), + updated_on: Default::default(), + partition_key_indices, + column_ids: vec![0, 1], + }; + TableInfoBuilder::default() + .table_id(table_id) + .table_version(0) + .name(name.to_string()) + .catalog_name(DEFAULT_CATALOG_NAME.to_string()) + .schema_name(DEFAULT_SCHEMA_NAME.to_string()) + .desc(None) + .table_type(TableType::Base) + .meta(meta) + .build() + .unwrap() + } + + /// Table info of a base mito table partitioned on a millisecond timestamp. + fn native_table_info(table_id: u32, name: &str, engine: &str, ts_nullable: bool) -> TableInfo { + native_table_info_with( + table_id, + name, + engine, + ConcreteDataType::timestamp_millisecond_datatype(), + ts_nullable, + vec![0], + ) + } + + fn native_timestamp_value(ms: i64) -> Value { + Value::Timestamp(Timestamp::new_millisecond(ms)) + } + + fn native_ts_before(split_ms: i64) -> PartitionExpr { + partition_col("ts").lt(native_timestamp_value(split_ms)) + } + + fn native_ts_at(split_ms: i64) -> PartitionExpr { + partition_col("ts").eq(native_timestamp_value(split_ms)) + } + + fn native_ts_after(split_ms: i64) -> PartitionExpr { + partition_col("ts").gt(native_timestamp_value(split_ms)) + } + + fn native_ts_at_or_after(split_ms: i64) -> PartitionExpr { + partition_col("ts").gt_eq(native_timestamp_value(split_ms)) + } + + fn native_region_route( + region_number: u32, + expression: Option, + leader_peer_id: Option, + ) -> RegionRoute { + native_region_route_for_table(NATIVE_TABLE_ID, region_number, expression, leader_peer_id) + } + + /// A route whose region id belongs to `table_id`, which may differ from the + /// table the capture is asked about. + fn native_region_route_for_table( + table_id: u32, + region_number: u32, + expression: Option, + leader_peer_id: Option, + ) -> RegionRoute { + RegionRoute { + region: Region { + id: RegionId::new(table_id, region_number), + partition_expr: expression + .map(|expression| expression.as_json_str().unwrap()) + .unwrap_or_default(), + ..Default::default() + }, + leader_peer: leader_peer_id.map(|id| Peer::new(id, format!("127.0.0.1:{id}"))), + ..Default::default() + } + } + + fn native_target(peer_id: u64) -> RegionQueryTarget { + RegionQueryTarget::new(Peer::new(peer_id, format!("127.0.0.1:{peer_id}"))) + } + + /// Routes of the canonical explicit half-domain split. + fn native_half_domain_routes() -> Vec { + vec![ + native_region_route(1, Some(native_ts_before(NATIVE_SPLIT_MS)), Some(11)), + native_region_route(2, Some(native_ts_at_or_after(NATIVE_SPLIT_MS)), Some(22)), + ] + } + + fn native_region_ids() -> Vec { + vec![ + RegionId::new(NATIVE_TABLE_ID, 1), + RegionId::new(NATIVE_TABLE_ID, 2), + ] + } + + /// Builds a planner whose scan captures `info` over `routes`, the physical + /// region routes of that same table. + /// + /// `ts_filter` adds a filter on the timestamp column, which routes the + /// selection through pruning. `by_name_scan` scans a DataFusion provider + /// instead of a [`DfTableProviderAdapter`], which forces the ordinary name + /// lookup and leaves the scan without a captured table. + async fn native_planner_and_plan( + info: TableInfo, + routes: Vec, + ts_filter: Option, + by_name_scan: bool, + ) -> (DistExtensionPlanner, LogicalPlan, TableName) { + let table = EmptyTable::from_table_info(&info); + let catalog_manager = MemoryCatalogManager::with_default_setup(); + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: info.catalog_name.clone(), + schema: info.schema_name.clone(), + table_name: info.name.clone(), + table_id: info.table_id(), + table: table.clone(), + }) + .unwrap(); + + let backend = Arc::new(MemoryKvBackend::default()); + let metadata_manager = TableMetadataManager::new(backend.clone()); + metadata_manager + .create_table_metadata( + info.clone(), + TableRouteValue::physical(routes), + HashMap::new(), + ) + .await + .unwrap(); + + let table_route_cache = Arc::new(new_table_route_cache( + "native-raw-route-routes".to_string(), + CacheBuilder::new(16).build(), + backend.clone(), + )); + let partition_info_cache = Arc::new(new_partition_info_cache( + "native-raw-route-partitions".to_string(), + CacheBuilder::new(16).build(), + table_route_cache.clone(), + )); + let partition_rule_manager = Arc::new(PartitionRuleManager::new( + backend, + table_route_cache, + partition_info_cache, + )); + let catalog_manager: CatalogManagerRef = catalog_manager; + let planner = DistExtensionPlanner::new( + catalog_manager, + partition_rule_manager, + Arc::new(UnusedRegionQueryHandler), + false, + // The native raw route fixtures exercise the enabled capture. + true, + ); + let table_source: Arc = if by_name_scan { + Arc::new(DefaultTableSource::new(Arc::new( + datafusion::datasource::empty::EmptyTable::new( + info.meta.schema.arrow_schema().clone(), + ), + ))) + } else { + Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(table), + ))) + }; + let mut builder = + LogicalPlanBuilder::scan_with_filters(info.name.clone(), table_source, None, vec![]) + .unwrap(); + if let Some(split_ms) = ts_filter { + builder = builder + .filter( + df_col("ts").gt(lit(ScalarValue::TimestampMillisecond(Some(split_ms), None))), + ) + .unwrap(); + } + let plan = builder.build().unwrap(); + let table_name = + TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, info.name.clone()); + (planner, plan, table_name) + } + + /// Captures `selected` regions of `routes` without going through a logical + /// plan, to isolate one proof condition at a time. + fn capture_native_selected( + info: &TableInfo, + routes: &[RegionRoute], + selected: &[RegionId], + ) -> Option { + let table = EmptyTable::from_table_info(info); + capture_native_raw_route(Some(&table), info, info.table_id(), routes, selected) + } + + #[tokio::test] + async fn native_raw_route_captures_explicit_half_domains() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let (mut planner, plan, table_name) = + native_planner_and_plan(info, native_half_domain_routes(), None, false).await; + // The capture must come from the table captured by the scan: an empty + // catalog cannot rebind the authorized scan to a replacement table. + planner.catalog_manager = MemoryCatalogManager::with_default_setup(); + + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + assert_eq!(native_region_ids(), regions); + let native_raw_route = native_raw_route.expect("two explicit half domains are provable"); + assert_eq!( + vec![ + (RegionId::new(NATIVE_TABLE_ID, 1), native_target(11)), + (RegionId::new(NATIVE_TABLE_ID, 2), native_target(22)), + ], + native_raw_route.targets + ); + } + + #[tokio::test] + async fn native_raw_route_captures_on_pruning_path() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let (planner, plan, table_name) = native_planner_and_plan( + info, + native_half_domain_routes(), + Some(NATIVE_SPLIT_MS), + false, + ) + .await; + + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + // Timestamp predicates are not pruned yet, so the ordinary plan keeps + // every region; the capture still describes exactly those regions. + assert_eq!(native_region_ids(), regions); + let native_raw_route = native_raw_route.expect("the pruning fallback captures the route"); + assert_eq!(2, native_raw_route.targets.len()); + assert_eq!( + (RegionId::new(NATIVE_TABLE_ID, 1), native_target(11)), + native_raw_route.targets[0] + ); + assert_eq!( + (RegionId::new(NATIVE_TABLE_ID, 2), native_target(22)), + native_raw_route.targets[1] + ); + } + + #[tokio::test] + async fn native_raw_route_declines_pruned_subset_without_changing_it() { + // A table partitioned on `host`: the pruner returns a proper subset, + // whose domains are not timestamp domains, so only the ordinary ids + // survive. + let info = native_table_info_with( + NATIVE_TABLE_ID, + "native", + MITO_ENGINE, + ConcreteDataType::timestamp_millisecond_datatype(), + false, + vec![1], + ); + let routes = vec![ + native_region_route( + 1, + Some(partition_col("host").lt(Value::String("m".into()))), + Some(11), + ), + native_region_route( + 2, + Some(partition_col("host").gt_eq(Value::String("m".into()))), + Some(22), + ), + ]; + let (planner, plan, table_name) = native_planner_and_plan(info, routes, None, false).await; + let plan = LogicalPlanBuilder::from(plan) + .filter(df_col("host").eq(lit("a"))) + .unwrap() + .build() + .unwrap(); + + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + assert_eq!(vec![RegionId::new(NATIVE_TABLE_ID, 1)], regions); + assert!(native_raw_route.is_none()); + } + + #[test] + fn native_raw_route_preserves_selected_subset_and_order() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let routes = vec![ + native_region_route(1, Some(native_ts_before(NATIVE_SPLIT_MS)), Some(11)), + native_region_route(2, Some(native_ts_at(NATIVE_SPLIT_MS)), Some(22)), + native_region_route(3, Some(native_ts_after(NATIVE_SPLIT_MS)), Some(33)), + ]; + + // A pruned selection in pruning order, which is not route order. + let selected = vec![ + RegionId::new(NATIVE_TABLE_ID, 3), + RegionId::new(NATIVE_TABLE_ID, 1), + ]; + let captured = capture_native_selected(&info, &routes, &selected) + .expect("three disjoint timestamp domains are provable"); + assert_eq!( + vec![ + (RegionId::new(NATIVE_TABLE_ID, 3), native_target(33)), + (RegionId::new(NATIVE_TABLE_ID, 1), native_target(11)), + ], + captured.targets + ); + + // The complete selection keeps route order. + let all = vec![ + RegionId::new(NATIVE_TABLE_ID, 1), + RegionId::new(NATIVE_TABLE_ID, 2), + RegionId::new(NATIVE_TABLE_ID, 3), + ]; + let captured = capture_native_selected(&info, &routes, &all) + .expect("the complete selection is provable"); + assert_eq!( + vec![ + (RegionId::new(NATIVE_TABLE_ID, 1), native_target(11)), + (RegionId::new(NATIVE_TABLE_ID, 2), native_target(22)), + (RegionId::new(NATIVE_TABLE_ID, 3), native_target(33)), + ], + captured.targets + ); + } + + #[tokio::test] + async fn native_raw_route_declines_overlapping_boundary() { + // `ts <= split` and `ts >= split` share the split, so the domains are + // not disjoint even though both are explicit. + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let routes = vec![ + native_region_route( + 1, + Some(partition_col("ts").lt_eq(native_timestamp_value(NATIVE_SPLIT_MS))), + Some(11), + ), + native_region_route(2, Some(native_ts_at_or_after(NATIVE_SPLIT_MS)), Some(22)), + ]; + let (planner, plan, table_name) = native_planner_and_plan(info, routes, None, false).await; + + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + assert_eq!(native_region_ids(), regions); + assert!(native_raw_route.is_none()); + } + + #[test] + fn native_raw_route_declines_disjunction() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let disjunction = PartitionExpr::new( + Operand::Expr(native_ts_before(NATIVE_SPLIT_MS)), + RestrictedOp::Or, + Operand::Expr(native_ts_at_or_after(NATIVE_SPLIT_MS)), + ); + let routes = vec![native_region_route(1, Some(disjunction), Some(11))]; + + assert!( + capture_native_selected(&info, &routes, &[RegionId::new(NATIVE_TABLE_ID, 1)]).is_none() + ); + } + + #[test] + fn native_raw_route_declines_mixed_columns() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let mixed = native_ts_before(NATIVE_SPLIT_MS) + .and(partition_col("host").lt(Value::String("m".into()))); + let routes = vec![ + native_region_route(1, Some(mixed), Some(11)), + native_region_route(2, Some(native_ts_at_or_after(NATIVE_SPLIT_MS)), Some(22)), + ]; + + assert!(capture_native_selected(&info, &routes, &native_region_ids()).is_none()); + } + + #[test] + fn native_raw_route_declines_spaced_ordinal_conjunction() { + // `ts > 0 AND ts < 1` is empty over millisecond integers, but the + // ordinal normalization spaces 0 and 1 apart and would call them + // overlapping, so the conjunction must decline before the overlap proof. + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let empty_conjunction = partition_col("ts") + .gt(native_timestamp_value(0)) + .and(native_ts_before(1)); + let routes = vec![native_region_route(1, Some(empty_conjunction), Some(11))]; + + assert!( + capture_native_selected(&info, &routes, &[RegionId::new(NATIVE_TABLE_ID, 1)]).is_none() + ); + } + + #[test] + fn native_raw_route_declines_adjacent_exclusive_conjunction() { + // Adjacent exclusive domains are disjoint, but they are still expressed + // as conjunctions, which the supported grammar does not accept. + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let lower = + native_ts_at_or_after(NATIVE_SPLIT_MS).and(native_ts_before(NATIVE_SPLIT_MS + 1_000)); + let routes = vec![ + native_region_route(1, Some(lower), Some(11)), + native_region_route( + 2, + Some(native_ts_at_or_after(NATIVE_SPLIT_MS + 1_000)), + Some(22), + ), + ]; + + assert!( + capture_native_selected(&info, &routes, &[RegionId::new(NATIVE_TABLE_ID, 1)]).is_none() + ); + } + + #[test] + fn native_raw_route_declines_contradictory_equality_conjunction() { + // `ts = 5 AND ts = 6` is empty, but an equality short-circuits the range + // the overlap helper derives, so the conjunction must decline first. + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let contradictory = native_ts_at(5).and(native_ts_at(6)); + let routes = vec![native_region_route(1, Some(contradictory), Some(11))]; + + assert!( + capture_native_selected(&info, &routes, &[RegionId::new(NATIVE_TABLE_ID, 1)]).is_none() + ); + } + + #[test] + fn native_raw_route_declines_empty_strict_bounds() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let selected = [RegionId::new(NATIVE_TABLE_ID, 1)]; + + // No millisecond timestamp is below the minimum or above the maximum. + let below_min = vec![native_region_route( + 1, + Some(partition_col("ts").lt(Value::Timestamp(Timestamp::MIN_MILLISECOND))), + Some(11), + )]; + assert!(capture_native_selected(&info, &below_min, &selected).is_none()); + + let above_max = vec![native_region_route( + 1, + Some(partition_col("ts").gt(Value::Timestamp(Timestamp::MAX_MILLISECOND))), + Some(11), + )]; + assert!(capture_native_selected(&info, &above_max, &selected).is_none()); + + // The same bounds with the operands flipped describe the same empty + // domains. + let flipped_below_min = vec![native_region_route( + 1, + Some(PartitionExpr::new( + Operand::Value(Value::Timestamp(Timestamp::MIN_MILLISECOND)), + RestrictedOp::Gt, + Operand::Column("ts".to_string()), + )), + Some(11), + )]; + assert!(capture_native_selected(&info, &flipped_below_min, &selected).is_none()); + + let flipped_above_max = vec![native_region_route( + 1, + Some(PartitionExpr::new( + Operand::Value(Value::Timestamp(Timestamp::MAX_MILLISECOND)), + RestrictedOp::Lt, + Operand::Column("ts".to_string()), + )), + Some(11), + )]; + assert!(capture_native_selected(&info, &flipped_above_max, &selected).is_none()); + + // Inclusive bounds at the same edges are non-empty. + let at_max = vec![native_region_route( + 1, + Some(partition_col("ts").gt_eq(Value::Timestamp(Timestamp::MAX_MILLISECOND))), + Some(11), + )]; + let captured = capture_native_selected(&info, &at_max, &selected) + .expect("the maximum millisecond timestamp is representable"); + assert_eq!(vec![(selected[0], native_target(11))], captured.targets); + } + + #[test] + fn native_raw_route_declines_foreign_table_route_with_same_region_number() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let foreign_table_id = NATIVE_TABLE_ID + 1; + // The routes belong to another table but carry the same region numbers + // and provable domains, so only the full region id can reject them. + let routes = vec![ + native_region_route_for_table( + foreign_table_id, + 1, + Some(native_ts_before(NATIVE_SPLIT_MS)), + Some(11), + ), + native_region_route_for_table( + foreign_table_id, + 2, + Some(native_ts_at_or_after(NATIVE_SPLIT_MS)), + Some(22), + ), + ]; + + assert!(capture_native_selected(&info, &routes, &native_region_ids()).is_none()); + } + + #[tokio::test] + async fn native_raw_route_declines_foreign_table_route_without_changing_ordinary_ids() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let foreign_table_id = NATIVE_TABLE_ID + 1; + let routes = vec![ + native_region_route_for_table( + foreign_table_id, + 1, + Some(native_ts_before(NATIVE_SPLIT_MS)), + Some(11), + ), + native_region_route_for_table( + foreign_table_id, + 2, + Some(native_ts_at_or_after(NATIVE_SPLIT_MS)), + Some(22), + ), + ]; + let (planner, plan, table_name) = native_planner_and_plan(info, routes, None, false).await; + + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + // The ordinary plan still selects its own region ids from the route, and + // only the capture declines. + assert_eq!(native_region_ids(), regions); + assert!(native_raw_route.is_none()); + } + + #[test] + fn native_raw_route_declines_unsupported_timestamp_bounds() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let selected = [RegionId::new(NATIVE_TABLE_ID, 1)]; + + // Milliseconds are the only supported unit. + let seconds = vec![native_region_route( + 1, + Some(partition_col("ts").lt(Value::Timestamp(Timestamp::new_second(1_700_000_000)))), + Some(11), + )]; + assert!(capture_native_selected(&info, &seconds, &selected).is_none()); + + // The value must be a timestamp, not another type. + let integer = vec![native_region_route( + 1, + Some(partition_col("ts").lt(Value::Int64(NATIVE_SPLIT_MS))), + Some(11), + )]; + assert!(capture_native_selected(&info, &integer, &selected).is_none()); + + // `!=` does not describe a domain. + let not_eq = vec![native_region_route( + 1, + Some(partition_col("ts").not_eq(native_timestamp_value(NATIVE_SPLIT_MS))), + Some(11), + )]; + assert!(capture_native_selected(&info, ¬_eq, &selected).is_none()); + + // The timestamp column type must be millisecond as well. + let second_column_info = native_table_info_with( + NATIVE_TABLE_ID, + "native", + MITO_ENGINE, + ConcreteDataType::timestamp_second_datatype(), + false, + vec![0], + ); + let second_column = vec![native_region_route( + 1, + Some(partition_col("ts").lt(native_timestamp_value(NATIVE_SPLIT_MS))), + Some(11), + )]; + assert!(capture_native_selected(&second_column_info, &second_column, &selected).is_none()); + } + + #[test] + fn native_raw_route_accepts_flipped_operands() { + // The shared collider flips `value op column` the same way, so the + // capture accepts it too. + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let routes = vec![ + native_region_route( + 1, + Some(PartitionExpr::new( + Operand::Value(native_timestamp_value(NATIVE_SPLIT_MS)), + RestrictedOp::Gt, + Operand::Column("ts".to_string()), + )), + Some(11), + ), + native_region_route(2, Some(native_ts_at_or_after(NATIVE_SPLIT_MS)), Some(22)), + ]; + + // `split > ts` is `ts < split`, so region 1 and region 2 are disjoint. + let captured = capture_native_selected(&info, &routes, &native_region_ids()) + .expect("flipped operands describe the same domains"); + assert_eq!(2, captured.targets.len()); + assert_eq!( + (RegionId::new(NATIVE_TABLE_ID, 1), native_target(11)), + captured.targets[0] + ); + } + + #[tokio::test] + async fn native_raw_route_declines_missing_expression() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let routes = vec![ + native_region_route(1, Some(native_ts_before(NATIVE_SPLIT_MS)), Some(11)), + // The default region has no explicit expression at all. + native_region_route(2, None, Some(22)), + ]; + let (planner, plan, table_name) = native_planner_and_plan(info, routes, None, false).await; + + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + assert_eq!(native_region_ids(), regions); + assert!(native_raw_route.is_none()); + } + + #[test] + fn native_raw_route_declines_missing_leader() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let routes = vec![ + native_region_route(1, Some(native_ts_before(NATIVE_SPLIT_MS)), Some(11)), + native_region_route(2, Some(native_ts_at_or_after(NATIVE_SPLIT_MS)), None), + ]; + + assert!(capture_native_selected(&info, &routes, &native_region_ids()).is_none()); + } + + #[test] + fn native_raw_route_declines_staging_or_downgrading_leader() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let mut routes = native_half_domain_routes(); + assert!(capture_native_selected(&info, &routes, &native_region_ids()).is_some()); + + routes[1].leader_state = Some(LeaderState::Staging); + assert!(capture_native_selected(&info, &routes, &native_region_ids()).is_none()); + + routes[1].leader_state = Some(LeaderState::Downgrading); + assert!(capture_native_selected(&info, &routes, &native_region_ids()).is_none()); + } + + #[test] + fn native_raw_route_declines_duplicate_regions() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let routes = native_half_domain_routes(); + + let duplicated_selection = [ + RegionId::new(NATIVE_TABLE_ID, 1), + RegionId::new(NATIVE_TABLE_ID, 1), + ]; + assert!(capture_native_selected(&info, &routes, &duplicated_selection).is_none()); + + // Two routes for one region number cannot be mapped one-to-one. + let duplicated_routes = vec![ + native_region_route(1, Some(native_ts_before(NATIVE_SPLIT_MS)), Some(11)), + native_region_route(1, Some(native_ts_at_or_after(NATIVE_SPLIT_MS)), Some(22)), + ]; + assert!( + capture_native_selected( + &info, + &duplicated_routes, + &[RegionId::new(NATIVE_TABLE_ID, 1)] + ) + .is_none() + ); + } + + #[test] + fn native_raw_route_declines_unprovable_tables() { + let routes = native_half_domain_routes(); + let selected = native_region_ids(); + + // A logical table whose route resolves to another physical table. + let logical_info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let table = EmptyTable::from_table_info(&logical_info); + assert!( + capture_native_raw_route( + Some(&table), + &logical_info, + PHYSICAL_TABLE_ID, + &routes, + &selected, + ) + .is_none() + ); + + // A non-mito engine. + let metric_info = native_table_info(NATIVE_TABLE_ID, "native", "metric", false); + assert!(capture_native_selected(&metric_info, &routes, &selected).is_none()); + + // A nullable timestamp column cannot prove disjointness. + let nullable_info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, true); + assert!(capture_native_selected(&nullable_info, &routes, &selected).is_none()); + + // Without a captured table the name lookup may bind a replacement. + assert!( + capture_native_raw_route(None, &logical_info, NATIVE_TABLE_ID, &routes, &selected,) + .is_none() + ); + } + + #[tokio::test] + async fn native_raw_route_declines_name_lookup_without_captured_table() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let (planner, plan, table_name) = + native_planner_and_plan(info, native_half_domain_routes(), None, true).await; + + // The ordinary plan still resolves the table by name and keeps every + // region; only the capture is skipped. + let (regions, native_raw_route) = planner.get_regions(&table_name, &plan).await.unwrap(); + assert_eq!(native_region_ids(), regions); + assert!(native_raw_route.is_none()); + } + + #[tokio::test] + async fn native_raw_route_capture_follows_experimental_flag() { + let info = native_table_info(NATIVE_TABLE_ID, "native", MITO_ENGINE, false); + let (mut planner, plan, table_name) = + native_planner_and_plan(info, native_half_domain_routes(), None, false).await; + + // Enabled: the same provable half domains capture completely, in + // ordinary region order. + let (enabled_regions, enabled_route) = + planner.get_regions(&table_name, &plan).await.unwrap(); + assert_eq!(native_region_ids(), enabled_regions); + assert_eq!( + vec![ + (RegionId::new(NATIVE_TABLE_ID, 1), native_target(11)), + (RegionId::new(NATIVE_TABLE_ID, 2), native_target(22)), + ], + enabled_route.expect("the enabled fixture captures").targets + ); + + // Disabled: the same regions in the same order, and no capture at all. + planner.capture_native_raw_route = false; + let (disabled_regions, disabled_route) = + planner.get_regions(&table_name, &plan).await.unwrap(); + assert_eq!(enabled_regions, disabled_regions); + assert!(disabled_route.is_none()); } fn assert_all_logical_regions(mut regions: Vec) { diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index be681cfcf57..e2095c33530 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -274,6 +274,7 @@ impl QueryEngineState { partition_rule_manager, region_query_handler.clone(), options.enable_per_region_metrics, + options.experimental_native_promql, ))) .with_optimizer_rules(optimizer.rules) .with_physical_optimizer_rules(physical_optimizer.rules) @@ -575,6 +576,7 @@ impl DfQueryPlanner { partition_rule_manager: Option, region_query_handler: Option, enable_per_region_metrics: bool, + capture_native_raw_route: bool, ) -> Self { let mut planners: Vec> = vec![ Arc::new(PromExtensionPlanner), @@ -589,6 +591,7 @@ impl DfQueryPlanner { partition_rule_manager, region_query_handler, enable_per_region_metrics, + capture_native_raw_route, ))); planners.push(Arc::new(MergeSortExtensionPlanner {})); } diff --git a/typos.toml b/typos.toml index 44dcda61687..c6f1a1b0178 100644 --- a/typos.toml +++ b/typos.toml @@ -11,6 +11,8 @@ consts = "consts" # Upstream crate rskafka defines a `DeadlineExceded` error variant, but it is misspelled. # We need to extend the word list to avoid typos in our codebase. Exceded = "Exceded" +# Unicode general category `Nd` (decimal digit) used in the Go group-name regex. +Nd = "Nd" [default.extend-identifiers] allocatedp = "allocatedp" @@ -23,5 +25,8 @@ extend-exclude = [ "*.result", "src/pipeline/benches/data.log", "cyborg/pnpm-lock.yaml", - "tests-integration/resources/*" + "tests-integration/resources/*", + # Generated promql-compat golden fixtures: sample points are hex IEEE-754 + # bit patterns, which the spell checker reads as words. + "src/promql/**/*_cases.txt" ]