diff --git a/Cargo.lock b/Cargo.lock index 2ef45f7114a..606322e47b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11378,6 +11378,7 @@ dependencies = [ "regex", "snafu 0.8.6", "store-api", + "table", "tokio", "twox-hash", ] diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs index 9fbde16fb7b..87a8281e170 100644 --- a/src/cmd/src/datanode/scanbench.rs +++ b/src/cmd/src/datanode/scanbench.rs @@ -47,7 +47,7 @@ use mito2::sst::file_ref::FileReferenceManager; use moka::future::CacheBuilder; use object_store::manager::ObjectStoreManager; use object_store::util::normalize_dir; -use query::optimizer::parallelize_scan::ParallelizeScan; +use query::optimizer::parallelize_scan::assign_partition_ranges; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt}; use sqlparser::ast::ExprWithAlias as SqlExprWithAlias; @@ -995,8 +995,7 @@ impl ScanbenchCommand { let all_ranges: Vec<_> = original_partitions.into_iter().flatten().collect(); // Distribute ranges across partitions - let mut partitions = - ParallelizeScan::assign_partition_range(all_ranges, self.parallelism); + let mut partitions = assign_partition_ranges(all_ranges, self.parallelism); // Sort ranges within each partition by start time ascending for partition in &mut partitions { diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index ca0f3b74057..ec8305efa1c 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -734,6 +734,9 @@ impl StartCommand { procedure_executor.clone(), process_manager, ) + // The region servers this frontend reads are the ones built above, in this process, so a + // region input can travel as a logical plan instead of Substrait bytes. + .with_in_process_region_transport() .with_local_file_access(local_file_access); plugins::setup_frontend_plugins_post_build(&mut plugins, &plugin_opts, &fe_instance) diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 624ccb2f9cd..918655d0c5e 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -78,6 +78,9 @@ pub struct FrontendBuilder { procedure_executor: ProcedureExecutorRef, process_manager: ProcessManagerRef, local_file_access: LocalFileAccess, + /// Whether this frontend reaches the regions of its datanodes in process. Declared by whoever + /// assembles the frontend, because only it knows what its node manager talks to. + in_process_region_transport: bool, } impl FrontendBuilder { @@ -105,6 +108,7 @@ impl FrontendBuilder { procedure_executor, process_manager, local_file_access: LocalFileAccess::Disabled, + in_process_region_transport: false, } } @@ -157,6 +161,20 @@ impl FrontendBuilder { self } + /// Declares that every region this frontend queries is read in this very process. + /// + /// The caller must own the node manager this frontend was built with, and know that its + /// `DatanodeManager` hands region reads to a `RegionServer` in this process — as standalone + /// startup does, since it passes the datanode's own region server to the node manager it + /// installs. The distributed planner then hands a region a `LogicalPlan` instead of Substrait + /// bytes, which `RegionServer::handle_read` accepts. A frontend whose node manager reaches + /// remote datanodes must leave it off: those datanodes only accept Substrait bytes + /// (`RegionServer::handle_remote_read`), so a plan that does not encode would fail there. + pub fn with_in_process_region_transport(mut self) -> Self { + self.in_process_region_transport = true; + self + } + pub fn options(&self) -> &FrontendOptions { &self.options } @@ -315,7 +333,7 @@ impl FrontendBuilder { let mut query_options = self.options.query.clone(); query_options.enable_per_region_metrics = self.options.logging.enable_per_region_metrics; - let query_engine = QueryEngineFactory::try_new_with_plugins( + let query_engine = QueryEngineFactory::try_new_with_plugins_and_region_transport( self.catalog_manager.clone(), Some(partition_manager.clone()), Some(region_query_handler.clone()), @@ -323,6 +341,7 @@ impl FrontendBuilder { Some(procedure_service_handler), Some(Arc::new(flow_service)), true, + self.in_process_region_transport, plugins.clone(), query_options, ) diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index 433fac8c7f1..8fb06782240 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -256,10 +256,41 @@ impl SeriesScan { .is_some_and(|context| context.format == crate::config::SeriesIndexFormat::Columnar) } + /// Returns whether raw series come from the columnar prepared source. + /// + /// Derived only from the pinned series index version and every configured partition range, so + /// preflight, the ordering guarantee, the partition reader and scanner-wide preparation agree. + /// `false` means the stored sources are read by the ordinary scan. + fn uses_prepared_columnar(&self) -> Result { + if !self.uses_columnar() { + return Ok(false); + } + let ranges = self + .properties + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + prepared::validate_inputs(&self.stream_ctx, &ranges).map_err(BoxedError::new) + } + + /// Returns the raw ordering guarantee for a scan whose columnar readiness is already known. + fn raw_ordering(&self, prepared_columnar: bool) -> SeriesOrdering { + if prepared_columnar { + return SeriesOrdering::KeyThenTime; + } + match self.mode { + SeriesScanMode::Legacy => SeriesOrdering::KeyThenTime, + SeriesScanMode::TwoPhase => SeriesOrdering::Contiguous, + } + } + + /// Prepares the raw scan and reports whether its data comes from the columnar prepared source. fn prepare_raw_scan( &self, label_source: &RawLabelSource, - ) -> Result, BoxedError> { + ) -> Result, BoxedError> { let input = &self.stream_ctx.input; let metadata = input.mapper.metadata(); match label_source { @@ -386,25 +417,19 @@ impl SeriesScan { ); return Ok(None); } - if self.uses_columnar() { - let ranges = self - .properties - .partitions - .iter() - .flatten() - .copied() - .collect::>(); - prepared::validate_inputs(&self.stream_ctx, &ranges).map_err(BoxedError::new)?; - } + // Structural eligibility errors still reject the raw path; ordinary scans cover sources + // without columnar artifacts, so missing coverage is not an error here. + let prepared_columnar = self.uses_prepared_columnar()?; let schema = Arc::new(SeriesSchema { fields: specs.into(), eval_grid: None, }); let typed_schema = input.mapper.input_arrow_schema(false); let mut predicates = Vec::new(); - let predicate = if self.uses_columnar() { + let predicate = if prepared_columnar { input.predicate_group().predicate() } else { + // The ordinary scan applies region predicates before this residual filtering. input.predicate_group().predicate_without_region() }; if let Some(predicate) = predicate { @@ -452,11 +477,14 @@ impl SeriesScan { predicates.push(physical); } } - Ok(Some(RawScanPreparation { - schema, - fields, - predicates, - })) + Ok(Some(( + RawScanPreparation { + schema, + fields, + predicates, + }, + prepared_columnar, + ))) } fn supports_two_phase(input: &ScanInput) -> bool { @@ -844,12 +872,11 @@ impl RegionScanner for SeriesScan { } fn raw_series_ordering(&self) -> SeriesOrdering { - if self.uses_columnar() { - return SeriesOrdering::KeyThenTime; - } - match self.mode { - SeriesScanMode::Legacy => SeriesOrdering::KeyThenTime, - SeriesScanMode::TwoPhase => SeriesOrdering::Contiguous, + // Readiness is fallible here, but every fallible entry point rejects the same structural + // errors before execution, so the `false` fallback below is never observed in a query. + match self.uses_prepared_columnar() { + Ok(prepared_columnar) => self.raw_ordering(prepared_columnar), + Err(_) => self.raw_ordering(false), } } @@ -859,7 +886,7 @@ impl RegionScanner for SeriesScan { ) -> Result>, BoxedError> { Ok(self .prepare_raw_scan(label_source)? - .map(|preparation| preparation.schema)) + .map(|(preparation, _)| preparation.schema)) } fn scan_series_partition( @@ -869,10 +896,10 @@ impl RegionScanner for SeriesScan { partition: usize, label_source: &RawLabelSource, ) -> Result, BoxedError> { - let Some(preparation) = self.prepare_raw_scan(label_source)? else { + let Some((preparation, prepared_columnar)) = self.prepare_raw_scan(label_source)? else { return Ok(None); }; - if self.uses_columnar() { + if prepared_columnar { let partitions = self.properties.num_partitions().max(1); if partition >= partitions { return Err(BoxedError::new( @@ -897,10 +924,12 @@ impl RegionScanner for SeriesScan { return Ok(Some(raw::from_stream( schema, labels, - SeriesOrdering::KeyThenTime, + self.raw_ordering(prepared_columnar), Box::pin(stream), ))); } + // Stored sources without columnar artifacts (memtables, unindexed SSTs) are read by the + // ordinary scan, so the raw path stays available. let metrics = new_partition_metrics( &self.stream_ctx, ctx.explain_verbose, @@ -911,7 +940,7 @@ impl RegionScanner for SeriesScan { let batches = self .scan_batch_in_partition(ctx, partition, metrics.clone(), metrics_set) .map_err(BoxedError::new)?; - let ordering = self.raw_series_ordering(); + let ordering = self.raw_ordering(prepared_columnar); Ok(Some(raw::make_stream( batches, preparation.schema, @@ -957,9 +986,13 @@ impl RegionScanner for SeriesScan { &self.metrics_list, ); Box::pin(async move { - let Some(preparation) = preparation? else { + let Some((preparation, prepared_columnar)) = preparation? else { return Ok(None); }; + if !prepared_columnar { + // Known unsupported coverage: the caller keeps the ordinary scan instead. + return Ok(None); + } state .prepare(stream_ctx, ranges, labels, preparation, metrics) .await diff --git a/src/mito2/src/read/series_scan/prepared.rs b/src/mito2/src/read/series_scan/prepared.rs index 8dc002f9b23..0c87d6f18d4 100644 --- a/src/mito2/src/read/series_scan/prepared.rs +++ b/src/mito2/src/read/series_scan/prepared.rs @@ -137,20 +137,30 @@ fn same_labels(left: &RawLabelSource, right: &RawLabelSource) -> bool { } /// Validates the pinned assignments and coverage without opening artifacts. -pub(crate) fn validate_inputs(context: &StreamContext, ranges: &[PartitionRange]) -> Result<()> { - collect_inputs(context, ranges).map(|_| ()) +/// +/// Returns whether every assigned source can be read from columnar artifacts. `false` is not an +/// error: the caller scans those sources with the ordinary reader instead. +pub(crate) fn validate_inputs(context: &StreamContext, ranges: &[PartitionRange]) -> Result { + collect_inputs(context, ranges).map(|inputs| inputs.is_some()) } +/// Collects the assigned SSTs and their covering columnar series indexes. +/// +/// `None` means the assignments are structurally valid but not ready for columnar reads: an +/// assigned memtable, an SST without columnar series coverage, or an SST without columnar +/// samples. Structural corruption is still an error, and every assignment is validated before +/// either answer is returned, so a later invalid assignment is never masked by missing artifacts. fn collect_inputs( context: &StreamContext, ranges: &[PartitionRange], -) -> Result<(Vec, HashSet)> { +) -> Result, HashSet)>> { let input = &context.input; let metadata = input.region_metadata(); let index = input .series_index .as_ref() .ok_or_else(|| invalid("missing columnar index context"))?; + let mut ready = true; let mut files = BTreeMap::>::new(); for range in ranges { let range = context @@ -159,7 +169,9 @@ fn collect_inputs( .ok_or_else(|| invalid("invalid prepared scan range"))?; for row_group in &range.row_group_indices { if context.is_mem_range_index(*row_group) { - return Err(invalid("columnar prepared reads require flushed memtables")); + // Artifacts exist for flushed SSTs only. + ready = false; + continue; } if !context.is_file_range_index(*row_group) { return Err(invalid( @@ -203,24 +215,18 @@ fn collect_inputs( } } if !covered { - return Err(invalid(format!( - "SST {} has no columnar series coverage; run build_series_index()", - file.file_id() - ))); + ready = false; } - index + if index .version .samples .get(&file.file_id().file_id()) - .filter(|handle| handle.entry().artifact.format == ArtifactFormat::ColumnarSamples) - .ok_or_else(|| { - invalid(format!( - "SST {} has no columnar samples; run build_series_index()", - file.file_id() - )) - })?; + .is_none_or(|handle| handle.entry().artifact.format != ArtifactFormat::ColumnarSamples) + { + ready = false; + } } - Ok((file_indices, indexes)) + Ok(ready.then_some((file_indices, indexes))) } struct Candidate { @@ -254,7 +260,12 @@ impl ColumnarPreparedSource { .series_index .as_ref() .ok_or_else(|| invalid("missing columnar index context"))?; - let (files, indexes) = collect_inputs(&context, &ranges)?; + let Some((files, indexes)) = collect_inputs(&context, &ranges)? else { + // Readiness was checked before preparation, so this is a real capability change. + return Err(invalid( + "columnar series coverage changed after series preparation", + )); + }; let mut samples = Vec::>::new(); for file_idx in files { let file = &input.files[file_idx]; diff --git a/src/mito2/src/read/series_scan/prepared/tests.rs b/src/mito2/src/read/series_scan/prepared/tests.rs index a0fbee4c14c..2d47eb36676 100644 --- a/src/mito2/src/read/series_scan/prepared/tests.rs +++ b/src/mito2/src/read/series_scan/prepared/tests.rs @@ -25,13 +25,14 @@ use datafusion::execution::TaskContext; use datafusion::physical_plan::{ExecutionPlan, metrics::ExecutionPlanMetricsSet}; use datafusion_expr::{col, lit}; use datatypes::prelude::ConcreteDataType; +use datatypes::series::SeriesOrdering; use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::SampleRef; use object_store::ObjectStore; use object_store::layers::mock::{self, MockLayerBuilder, oio}; use store_api::codec::PrimaryKeyEncoding; use store_api::region_engine::{ - BoundTag, PrepareRequest, QueryScanContext, RegionEngine, RegionScanner, + BoundTag, PartitionRange, PrepareRequest, QueryScanContext, RegionEngine, RegionScanner, }; use store_api::region_request::{ RegionBuildIndexRequest, RegionDeleteRequest, RegionPutRequest, RegionRequest, @@ -51,6 +52,10 @@ use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_wit use crate::test_util::{TestEnv, flush_region}; async fn fixture() -> (TestEnv, MitoEngine, MitoRegionRef) { + fixture_with(false).await +} + +async fn fixture_with(experimental_series_scan_v2: bool) -> (TestEnv, MitoEngine, MitoRegionRef) { let mut env = TestEnv::with_prefix("prepared-columnar").await; let mut metadata = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse); metadata @@ -66,6 +71,7 @@ async fn fixture() -> (TestEnv, MitoEngine, MitoRegionRef) { MitoConfig { experimental_enable_series_index: true, experimental_series_index_format: SelectedFormat::Columnar, + experimental_series_scan_v2, experimental_series_index_maintenance_interval: Duration::from_secs(3600), min_compaction_interval: Duration::from_secs(3600), ..Default::default() @@ -192,6 +198,41 @@ impl mock::Read for ReadRecorder { } } +/// Fails reads of the selected artifact paths so tests can prove that object-store failures +/// surface as errors instead of being classified as not-ready fallbacks. +struct FailingReader { + inner: oio::Reader, + fail: bool, +} +impl mock::Read for FailingReader { + async fn read( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, object_store::Buffer)> { + if self.fail { + return Err(injected_failure()); + } + self.inner.read(range).await + } + + async fn open( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, Box)> { + if self.fail { + return Err(injected_failure()); + } + self.inner.open(range).await + } +} + +fn injected_failure() -> mock::Error { + mock::Error::new( + mock::ErrorKind::Unexpected, + "injected series manifest read failure", + ) +} + fn scanner(region: &MitoRegionRef, store: ObjectStore) -> SeriesScan { scanner_with_format(region, store, SelectedFormat::Columnar) } @@ -269,6 +310,35 @@ fn rows(chunks: &[RawSeriesChunk]) -> Vec<(Vec, i64, f64)> { .collect() } +fn samples(chunks: &[RawSeriesChunk]) -> Vec<(i64, f64)> { + rows(chunks) + .into_iter() + .map(|(_, time, value)| (time, value)) + .collect() +} + +/// Reads the raw path of every partition, checks each stream declares `ordering`, and returns the +/// samples ordered by time so a check does not depend on how range grouping spread a memtable. +async fn all_raw_samples( + scanner: &dyn RegionScanner, + ctx: &QueryScanContext, + metrics: &ExecutionPlanMetricsSet, + binding: &RawLabelSource, + ordering: SeriesOrdering, +) -> Vec<(i64, f64)> { + let mut collected = Vec::new(); + for partition in 0..scanner.properties().num_partitions().max(1) { + let stream = scanner + .scan_series_partition(ctx, metrics, partition, binding) + .unwrap() + .unwrap(); + assert_eq!(ordering, stream.ordering()); + collected.extend(samples(&stream.try_collect::>().await.unwrap())); + } + collected.sort_unstable_by_key(|(time, _)| *time); + collected +} + #[tokio::test] async fn prepared_sources_share_discovery_and_read_independently() { let (_env, engine, region) = fixture().await; @@ -386,89 +456,133 @@ async fn prepared_sources_share_discovery_and_read_independently() { engine.stop().await.unwrap(); } +#[rstest::rstest] +#[case::legacy(false)] +#[case::two_phase(true)] #[tokio::test] -async fn native_scan_reports_incomplete_coverage_and_uses_prepared_source() { - let (_env, engine, region) = fixture().await; +async fn native_scan_falls_back_to_raw_sources_until_columnar_coverage(#[case] v2: bool) { + let (_env, engine, region) = fixture_with(v2).await; let request = ScanRequest { distribution: Some(TimeSeriesDistribution::PerSeries), ..Default::default() }; - // Three files are below the automatic index threshold. + let ctx = QueryScanContext::default(); + let metrics = ExecutionPlanMetricsSet::new(); + // The three flushed SSTs only have columnar artifacts after `build`. + let fallback_ordering = if v2 { + SeriesOrdering::Contiguous + } else { + SeriesOrdering::KeyThenTime + }; let raw_scanner = engine .handle_query(region.region_id, request.clone()) .await .unwrap(); let binding = labels(raw_scanner.metadata()); - assert!(raw_scanner.raw_series_schema(&binding).is_err()); - assert!( - raw_scanner - .scan_series_partition( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - 0, - &binding, - ) - .is_err() + assert_eq!(fallback_ordering, raw_scanner.raw_series_ordering()); + assert!(raw_scanner.raw_series_schema(&binding).unwrap().is_some()); + assert_eq!( + vec![(1000, 1.0), (2000, 1.0)], + all_raw_samples( + raw_scanner.as_ref(), + &ctx, + &metrics, + &binding, + fallback_ordering + ) + .await ); + // Known unready coverage selects the ordinary scan instead of a prepared source. assert!( raw_scanner - .prepare_series( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - &binding - ) + .prepare_series(&ctx, &metrics, &binding) .await - .is_err() + .unwrap() + .is_none() ); + build(&engine, ®ion).await; - let raw_scanner = engine + let mut raw_scanner = engine .handle_query(region.region_id, request.clone()) .await .unwrap(); + // Move every assigned range into a new last partition: readiness, the ordering guarantee and + // the reader must all follow the new assignment, not the pre-prepare partition list. + let assigned = raw_scanner + .properties() + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + raw_scanner + .prepare(PrepareRequest::default().with_ranges(vec![vec![], assigned])) + .unwrap(); + assert_eq!(2, raw_scanner.properties().num_partitions()); let binding = labels(raw_scanner.metadata()); + assert_eq!( + SeriesOrdering::KeyThenTime, + raw_scanner.raw_series_ordering() + ); + let source = raw_scanner + .prepare_series(&ctx, &metrics, &binding) + .await + .unwrap() + .unwrap(); + assert_eq!(1, source.num_series()); + // The first partition owns no range and the second owns all of them. + assert_eq!( + vec![(1000, 1.0), (2000, 1.0)], + all_raw_samples( + raw_scanner.as_ref(), + &ctx, + &metrics, + &binding, + SeriesOrdering::KeyThenTime + ) + .await + ); let plan = Arc::new(RegionScanExec::new(raw_scanner, request.clone(), None).unwrap()) .try_to_series_scan(binding.clone()) .unwrap() .unwrap(); - let frame = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); - let chunks = unpack_raw_series(frame, Some(plan.series_schema().clone())) - .unwrap() - .try_collect::>() - .await - .unwrap(); + assert_eq!(SeriesOrdering::KeyThenTime, plan.raw_series_ordering()); + let frame = plan.execute(1, Arc::new(TaskContext::default())).unwrap(); + let unpacked = unpack_raw_series(frame, Some(plan.series_schema().clone())).unwrap(); + // The packed frame declares the same guarantee the plan advertises. + assert_eq!(plan.raw_series_ordering(), unpacked.ordering()); assert_eq!( - vec![1000, 2000], - rows(&chunks).iter().map(|(_, t, _)| *t).collect::>() + vec![(1000, 1.0), (2000, 1.0)], + samples(&unpacked.try_collect::>().await.unwrap()) ); + + // An assigned memtable is known unready, so the raw path keeps the ordinary scan and must + // still return the new row. write(&engine, ®ion, 0, 3000, Some(4.0)).await; let raw_scanner = engine .handle_query(region.region_id, request) .await .unwrap(); let binding = labels(raw_scanner.metadata()); - assert!(raw_scanner.raw_series_schema(&binding).is_err()); + assert_eq!(fallback_ordering, raw_scanner.raw_series_ordering()); + assert!(raw_scanner.raw_series_schema(&binding).unwrap().is_some()); assert!( raw_scanner - .scan_series_partition( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - 0, - &binding, - ) - .is_err() - ); - assert!( - raw_scanner - .prepare_series( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - &binding - ) + .prepare_series(&ctx, &metrics, &binding) .await - .err() .unwrap() - .output_msg() - .contains("memtables") + .is_none() + ); + assert_eq!( + vec![(1000, 1.0), (2000, 1.0), (3000, 4.0)], + all_raw_samples( + raw_scanner.as_ref(), + &ctx, + &metrics, + &binding, + fallback_ordering + ) + .await ); engine.stop().await.unwrap(); } @@ -1046,7 +1160,8 @@ async fn preflight_rejects_unsupported_assignments_without_artifact_io() { }; let mut context = make_context(); let ranges = context.partition_ranges(); - validate_inputs(&context, &ranges).unwrap(); + // Complete columnar coverage is the only ready answer. + assert!(validate_inputs(&context, &ranges).unwrap()); context.ranges[ranges[0].identifier].row_group_indices = vec![RowGroupIndex { index: context.input.num_files(), row_group_index: -1, @@ -1083,30 +1198,30 @@ async fn preflight_rejects_unsupported_assignments_without_artifact_io() { .contains("foreign SST") ); + // Missing columnar samples are a not-ready answer, not an error. let mut context = make_context(); let index = context.input.series_index.as_mut().unwrap(); index.version = Arc::new(SeriesIndexVersion { series_indexes: index.version.series_indexes.clone(), ..Default::default() }); + assert!(!validate_inputs(&context, &ranges).unwrap()); + // Invalid assignments remain errors even when sample coverage is incomplete. + let mut invalid = ranges.clone(); + invalid.push(PartitionRange { + identifier: ranges[0].identifier + ranges.len(), + ..ranges[0] + }); assert!( - validate_inputs(&context, &ranges) + validate_inputs(&context, &invalid) .unwrap_err() .to_string() - .contains("no columnar samples") + .contains("invalid prepared scan range") ); let scan = SeriesScan::new(context.input, false); let binding = labels(scan.metadata()); - assert!(scan.raw_series_schema(&binding).is_err()); - assert!( - scan.scan_series_partition( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - 0, - &binding - ) - .is_err() - ); + // The raw path stays available for the ordinary scan and preparation is skipped. + assert!(scan.raw_series_schema(&binding).unwrap().is_some()); assert!( scan.prepare_series( &QueryScanContext::default(), @@ -1114,8 +1229,49 @@ async fn preflight_rejects_unsupported_assignments_without_artifact_io() { &binding ) .await - .is_err() + .unwrap() + .is_none() ); assert!(reads.lock().unwrap().is_empty()); engine.stop().await.unwrap(); } + +#[tokio::test] +async fn series_artifact_read_failure_propagates() { + let (_env, engine, region) = fixture().await; + build(&engine, ®ion).await; + // The published catalog is complete; the injected manifest read failure must still propagate. + let store = region.series_index_store.clone().unwrap().layer( + MockLayerBuilder::default() + .reader_factory(Arc::new(|path, _, inner| { + Box::new(FailingReader { + inner, + fail: path.contains("/series/") && path.ends_with("/manifest.json"), + }) + })) + .build() + .unwrap(), + ); + let context = StreamContext::seq_scan_ctx(scan_input( + ®ion, + store, + SelectedFormat::Columnar, + (0..region.version().metadata.column_metadatas.len()).collect(), + )); + let ranges = context.partition_ranges(); + assert!(validate_inputs(&context, &ranges).unwrap()); + let scan = SeriesScan::new(context.input, false); + let binding = labels(scan.metadata()); + assert!(scan.raw_series_schema(&binding).unwrap().is_some()); + let error = scan + .prepare_series( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + &binding, + ) + .await + .err() + .expect("a failed artifact read must not report a not-ready source"); + assert!(error.output_msg().contains("injected"), "{error}"); + engine.stop().await.unwrap(); +} diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index fa9da9f322e..22c4a4cf25c 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -75,6 +75,7 @@ use partition::expr::{Operand, PartitionExpr, RestrictedOp}; use partition::multi_dim::MultiDimPartitionRule; use query::parser::QueryStatement; use query::plan::extract_and_rewrite_full_table_names; +use query::promql::planner::NATIVE_PROMQL_EXTENSION_KEY; use query::query_engine::DefaultSerializer; use query::sql::create_table_stmt; use session::context::QueryContextRef; @@ -1164,7 +1165,12 @@ impl StatementExecutor { ) .await? } - Statement::Tql(query) => self.plan_tql(query.clone(), &ctx).await?, + Statement::Tql(query) => { + // The serialized view plan must not depend on the native execution extension. + let mut view_ctx = ctx.as_ref().clone(); + view_ctx.set_extension(NATIVE_PROMQL_EXTENSION_KEY, "false"); + self.plan_tql(query.clone(), &Arc::new(view_ctx)).await? + } _ => { return InvalidViewStmtSnafu {}.fail(); } diff --git a/src/promql/AGENTS.md b/src/promql/AGENTS.md index dc90c5a0eec..53ad3c61169 100644 --- a/src/promql/AGENTS.md +++ b/src/promql/AGENTS.md @@ -16,38 +16,57 @@ operators. Statement-to-plan translation lives in `src/query/src/promql/`. | Kernels | `src/promql/src/functions.rs`, `src/promql/src/functions/` | PromQL functions | | Matching | `src/promql/src/vector_match.rs` | Binary vector matching | | Extension operators | `src/promql/src/extension_plan.rs`, `src/promql/src/extension_plan/` | Logical nodes and execution plans | -| Native model | `src/promql/src/native/` | Compiled tree, chunks, display, metrics | +| Native model | `src/promql/src/native/` | Compiled tree, chunks, display, metrics and execution | -## Two evaluation paths, and their status +## Evaluation paths -The easiest thing to get wrong here. Both exist; only one runs. - -- **Extension operators are live.** `PromPlanner` builds `SeriesDivide`, - `SeriesNormalize`, `RangeManipulate`, `InstantManipulate`, `EmptyMetric` and - `UnionDistinctOn`; `PromExtensionPlanner` - (`src/promql/src/extension_plan/planner.rs`) maps each logical node to its - execution plan, one branch per node, `None` for anything else. -- **Nothing plans into the native model yet.** `src/query/src/promql/native_compile.rs` - compiles a statement into a `NativeTree`, but no planner calls it, so a tree - built here reaches no execution. `src/promql/src/native/` holds the tree, the - chunk protocol, display and metrics, and — once a runtime lands — the driver - that runs a tree node by node. Compiling and running are separate concerns - owned by separate files; do not assume a tree is executable because it - compiles. +- **Extension operators:** the existing PromQL path plans operators such as + `SeriesDivide`, `SeriesNormalize`, `RangeManipulate`, `InstantManipulate`, + `EmptyMetric` and `UnionDistinctOn`. `PromExtensionPlanner` + (`src/promql/src/extension_plan/planner.rs`) maps the logical nodes to their + execution plans. +- **Experimental fused native-series POC (default off):** the query planner offers + a PromQL statement only when the engine's `in_process_region_transport` and + `QueryOptions::experimental_native_promql` options are enabled and the per-request + `query.experimental_native_promql` extension is true. Admission is limited to a + bare vector selector (parentheses allowed) and unary `rate`, `increase`, + `sum_over_time`, and `abs` calls over a selector; range selectors occur only + inside those calls. It is a standalone route: unsupported expressions and all + other cases use the existing planner. Accepted plans use process-local + `FusedSeries`, with ordinary `RecordBatch` output collected through DataFusion. + Scan N→Fused N retains region/partition parallelism, but this POC does not merge + one logical series across regions or scan partitions, guarantee canonical final + labels across them, or use global label-set occupancy; `NativeExecutor` retains + its partition-local `Emitted` state. Repartitioning redistributes rows but does not perform + semantic series merging. A future gather/repartition plus merge/finalize is + outside this POC. Metric scans retain the static `__table_id` predicate. +- **Admission and SQL output:** the guards live in + `src/query/src/promql/planner/native.rs`. Non-`Float64` fields, alternative + (`or`) matchers and `rate`/`increase` on temporality-tagged tables keep the + existing plan; the latter preserves legacy per-series delta summation. + Allowing query fallback through the engine option, session variable or + `query_fallback` hint also disables native admission. Native SQL/TQL output + has column order `time index, fields, tags`, raw field names with nullable + `Float64` values, and nullable `Utf8` tags (including integer table tags). + Unlike a legacy raw selector, it does not preserve tag types; legacy function + expressions can also use different field names and column order. The native + time index is a non-nullable millisecond timestamp without a time zone. ## Change coupling -- **An extension node reaches a datanode only if three things line up**: a payload - branch in `src/common/substrait/src/extension_serializer.rs`, a physical branch - in `PromExtensionPlanner`, and a `Commutativity::Commutative` arm in - `check_extension_plan` (`src/query/src/dist_plan/commutativity.rs`). Missing the - third is silent: the node falls through to `Unsupported`, the distributed - boundary is placed above it, and it never reaches a region — with no error. +- Distributed extension operators need matching handling in the Substrait + serializer (`src/common/substrait/src/extension_serializer.rs`), the physical + `PromExtensionPlanner`, and `check_extension_plan` + (`src/query/src/dist_plan/commutativity.rs`). Process-local `FusedSeries` is + the exception: it is not serialized to a datanode. Its distributed placement + is restricted to a direct table scan, a bound metric name and trees that read + one series at a time. Native admission requires partition columns to be tags, + but region/partition-local execution does not merge a logical series across + regions or scan partitions; do not treat that placement guard as a global + series-merge or final-label guarantee. - A new logical node also needs its physical planner reachable from `DfQueryPlanner::create_physical_plan`, assembled in `src/query/src/query_engine/state.rs`. -- User-visible changes need planner coverage plus a sqlness case under - `tests/cases/standalone/common/promql/`. ## Gotchas @@ -72,9 +91,9 @@ The easiest thing to get wrong here. Both exist; only one runs. the producer ordering the adapter adds. Nothing in that metadata binds the key to a decode — the frame decoder compares those format keys and the field list, and never a label identity. `unpack_raw_series` therefore requires a - caller-supplied `RawLabelSource`: decode with the source that was actually bound - when the stream was produced, because the frame schema does not replace that - binding. The metadata describes format, not label identity. + caller-supplied `RawLabelSource`: decode with the source that was actually + bound when the stream was produced, because the frame schema does not replace + that binding. The metadata describes format, not label identity. - **A wrong label binding is silently wrong, not an error.** Storage validates the binding it is given (metadata identity, tag-ness, name matches, no duplicates) but cannot tell whether that set is the one the query meant, and the metric name @@ -85,26 +104,5 @@ The easiest thing to get wrong here. Both exist; only one runs. ```bash cargo nextest run -p promql cargo nextest run -p query # planner and native-compile changes -cargo sqlness bare -t # PromQL cases live under tests/cases/standalone/common/promql/ +cargo sqlness bare -t # PromQL cases under tests/cases/standalone/common/promql/ ``` - -## Not wired up yet - -The storage side of the raw path exists; what is missing is planner wiring. Both -gaps below are structural, not bugs. - -- **The native tree has no planner entry point.** `native_compile::compile` has no - caller, and a `NativeTree` becomes a physical plan only once something builds it - while it still holds the AST and the selector binding — the point where the - matchers, `__field__` and the lookback are still available. Compiling later, from - a plan, loses them. -- **No planner reaches a raw scan.** mito2 implements - `RegionScanner::raw_series_schema` (`mito2/src/read/series_scan.rs`), so - `RegionScanExec::try_to_series_scan` can build `SeriesScanExec` from a prepared - scanner; the trait's `Ok(None)` default only refuses engines that have not - implemented it. What is missing is the planner side: nothing selects that - conversion, and `SeriesScanExec::children()` is empty, so wrapping a - `RegionScanExec` before partition preparation leaves `ParallelizeScan` unable to - assign partitions. Either give the partitioner a child it can see, or convert to - the native leaf only after partition preparation — the second is the contract - `RegionScanExec::try_to_series_scan` documents. diff --git a/src/promql/Cargo.toml b/src/promql/Cargo.toml index a3d8f5c4e71..31e4d08a9fb 100644 --- a/src/promql/Cargo.toml +++ b/src/promql/Cargo.toml @@ -28,6 +28,7 @@ prost.workspace = true regex.workspace = true snafu.workspace = true store-api.workspace = true +table.workspace = true twox-hash.workspace = true [dev-dependencies] diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index d957851902a..fae78c06390 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -37,18 +37,15 @@ //! was not compiled against, or an output schema the conversion cannot classify. use std::cmp::Ordering; -use std::collections::hash_map::Entry; -use std::collections::{HashMap, VecDeque}; +use std::collections::VecDeque; use std::fmt::{self, Formatter}; use std::hash::{Hash, Hasher}; use std::pin::Pin; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use std::task::{Context, Poll}; use common_query::series_frame::{ensure_frame_matches_label_source, unpack_raw_series}; -use datafusion::arrow::array::{ - Array, ArrayRef, Float64Array, StringArray, TimestampMillisecondArray, -}; +use datafusion::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; use datafusion::arrow::buffer::ScalarBuffer; use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; use datafusion::arrow::record_batch::RecordBatch; @@ -56,9 +53,7 @@ use datafusion::common::DFSchemaRef; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::context::TaskContext; -use datafusion::execution::memory_pool::{ - MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, -}; +use datafusion::execution::memory_pool::UnboundedMemoryPool; use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; @@ -69,17 +64,15 @@ use datafusion::physical_plan::{ }; use datatypes::series::SeriesOrdering; use datatypes::series::raw::RawSeriesChunk; -use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, Series, SeriesSchema}; +use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, SeriesSchema}; use datatypes::series::time::EvalGrid; use futures::Stream; use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream}; use crate::evaluated::EvaluatedChunk; -use crate::grid_series::{LabelSet, label_set}; use crate::native::execute::{NativeExecutionState, NativeExecutor}; use crate::native::metrics::{self, NodePath}; use crate::native::plan::{NativeExpr, NativeTree}; -use crate::native::runs::{Run, absorb}; use crate::source::materialize_raw_labels; #[cfg(test)] @@ -456,89 +449,6 @@ fn tree_input(tree: &NativeTree) -> DataFusionResult { Ok(TreeInput { path, schema }) } -/// Refuses a tree whose logical series can be split across the partitions this node runs. -/// -/// Keeping the metric name is not what keeps two raw series apart. An empty tag value and an absent -/// tag are the same canonical label set, so `{host="h", zone=""}` and `{host="h"}` are one series to -/// Prometheus while the scan encodes them as two raw keys. A window that evaluated its input without -/// merging it therefore passes such a pair through as two series under one label set, exactly as the -/// functions that drop `__name__` do by merging them, and a map drops `__name__` outright. -/// -/// A stateful node above one of those needs both halves together, and the runtime does provide that -/// within one execution: a window whose input can hold two series under one label set collects the -/// fragments, merges each label set once, and only then evaluates. That merge is per execution, and -/// this node runs one execution per partition. Two raw series that become one label set can land in -/// different partitions - the scan keeps a *series* in one partition, which says nothing about two -/// different series that later become indistinguishable - so each execution merges only the half it -/// was given. The consumer above then answers from one half and reports nothing for the other, which -/// is an empty result rather than an error. -/// -/// Whether a window merged its input is runtime state this check does not have: it runs over a bare -/// tree, where a window's input is only something that produces samples. A window below another -/// window is therefore treated as one that did not merge, whatever function it runs, which refuses -/// `rate` under `rate` (which drops the name) and `last_over_time` under `rate` (which keeps it) -/// alike. -/// -/// Reading such a tree needs an exchange above this node that merges a logical series across -/// partitions, which does not exist yet. Until it does the tree is refused where it can hurt: over -/// one partition there is nothing to split, so only a tree run over more than one is checked. -/// -/// The compiler cannot build a nested window today - [`compile`](crate::promql::native_compile::compile) -/// requires a range function to apply directly to a selector, `([])`, -/// and a selector is raw samples - so no tree it produces has a window above a window, and this -/// check costs no executable query. It is checked because [`FusedSeriesExec::try_new`] accepts any -/// well-formed tree, and a tree it accepts is a tree it must answer correctly. -fn reject_a_split_logical_series( - tree: &NativeTree, - partition_count: usize, -) -> DataFusionResult<()> { - if partition_count == 1 { - return Ok(()); - } - // The walk is outermost first, so it meets nodes in the reverse of the order data flows: a - // consumer met before a rewriter is one that the rewriter feeds. - let mut consumer: Option = None; - let mut rejection: Option<(String, String)> = None; - metrics::walk(tree.root(), &mut |node, path| { - if rejection.is_some() { - return; - } - // A node can leave two raw series on one canonical label set. A map always can: it drops - // `__name__` and every empty-valued tag. An instant selection preserves its input labels, - // rather than dropping them, but can still leave two raw series on one canonical label set: - // `{zone=""}` and the absent `zone` are one label set and two raw keys. A window can - // whenever it did not merge its input, which a bare tree says nothing about: a check on the - // function alone accepts `last_over_time`, whose output keeps the name and can still hold - // both raw series. Every map, instant selection, and window is therefore taken to be a - // possible rewriter, and the refusal below is what narrows that to the shape that costs - // correctness - a window with a rewriter below it, a window being the only consumer here. - let rewrites = matches!( - node.expr, - NativeExpr::Map { .. } | NativeExpr::Instant { .. } | NativeExpr::Window { .. } - ); - if rewrites && let Some(above) = &consumer { - rejection = Some((path.as_str().to_string(), above.clone())); - return; - } - let consumes = matches!(node.expr, NativeExpr::Window { .. }); - if consumes && consumer.is_none() { - consumer = Some(path.as_str().to_string()); - } - }); - let Some((rewriter, above)) = rejection else { - return Ok(()); - }; - Err(DataFusionError::Plan(format!( - "FusedSeries runs a tree where `{rewriter}` can leave two raw series on one canonical label \ - set and `{above}` above it reads a whole series, over {partition_count} partitions: an \ - empty tag value and an absent tag are the same label set, the runtime merges such \ - fragments once per execution, and this node runs one execution per partition, so `{above}` \ - would read one half and answer an empty result. Reading this tree needs an exchange between \ - `{rewriter}` and `{above}` that merges a logical series across partitions, which is not \ - implemented" - ))) -} - /// Checks that `tree` reads the schema `binding` declares, on the grid it declares. /// /// The output grid is compared by [`EvalGrid::spec`], not by allocation: that is the rule @@ -762,11 +672,6 @@ pub struct FusedSeriesExec { input: Arc, /// The columns this node produces. output_schema: SchemaRef, - /// The grid positions every label set in this plan's output already occupies, and what they - /// keep charged on the memory pool. - /// - /// Plan-level rather than per-partition: see [`LabelSets`]. - labels: SharedLabelSets, /// Per-node metrics, registered by each execution of the tree. metric: ExecutionPlanMetricsSet, /// This node's partitioning, emission and boundedness. @@ -807,8 +712,6 @@ impl FusedSeriesExec { )?; let classification = classify_output(&output_schema, &tree_input.schema, &tags, &time_column)?; - // Read here rather than below, because whether the split below can happen decides whether - // the tree shape is one this node may run at all. let input_properties = input.properties(); let partition_count = input_properties.partitioning.partition_count(); if partition_count == 0 { @@ -816,7 +719,6 @@ impl FusedSeriesExec { "FusedSeriesExec input has no partitions".into(), )); } - reject_a_split_logical_series(&tree, partition_count)?; // The frame states its own label layout and series ordering; this plan's binding is an // expectation checked against that declaration, not an input to decoding. ensure_frame_matches_label_source(input.schema().as_ref(), &label_source)?; @@ -857,10 +759,6 @@ impl FusedSeriesExec { input_properties.emission_type, input_properties.boundedness, )); - // Shared by every partition's stream: a label set colliding is a property of the whole - // output vector, not of one partition. Building it needs no pool; the owner's reservation is - // attached in `execute`, which is the first point that has one. - let labels: SharedLabelSets = Arc::new(Mutex::new(LabelSets::default())); Ok(Self { promql, tree, @@ -873,7 +771,6 @@ impl FusedSeriesExec { tags: tags.into(), input, output_schema, - labels, metric: ExecutionPlanMetricsSet::new(), properties, }) @@ -967,14 +864,6 @@ impl ExecutionPlan for FusedSeriesExec { partition: usize, context: Arc, ) -> DataFusionResult { - // The label-set state is the plan's, so the reservation that pays for it is attached exactly - // once, to the pool of the first execution. Every partition calls this and only the first - // takes effect: re-registering per stream would move a charge the other partitions keep - // growing, and would give it back when that stream ends while the entries stay. - self.labels - .lock() - .unwrap_or_else(|e| e.into_inner()) - .attach(&context.runtime_env().memory_pool); let input = self.input.execute(partition, Arc::clone(&context))?; // Batches become raw series chunks through the shared frame decoder, which is the inverse of // what the producer packed: one chunk per frame row, each series closed exactly once by the @@ -990,8 +879,6 @@ impl ExecutionPlan for FusedSeriesExec { &self.metric, partition, )?; - // What the label sets hold is charged to the reservation attached above, so nothing is - // reserved per stream: the entries outlive every stream and are released with the plan. Ok(Box::pin(FusedSeriesStream { input: Some(input), executor: Some(NativeExecutor::new(state)), @@ -999,7 +886,6 @@ impl ExecutionPlan for FusedSeriesExec { classification: Arc::clone(&self.classification), tags: Arc::clone(&self.tags), output_schema: self.output_schema.clone(), - labels: Arc::clone(&self.labels), ready: VecDeque::new(), done: false, })) @@ -1100,263 +986,6 @@ fn field_column(samples: &Samples) -> DataFusionResult { } } -/// The bytes one label set holds, for the memory reservation. -/// -/// It counts each label's name and value with the pair that refers to them, which is the same -/// convention [`crate::native::execute`]'s own label-set table keeps. -fn labels_bytes(labels: &LabelSet) -> usize { - labels - .iter() - .map(|(name, value)| size_of::<(Arc, Arc)>() + name.len() + value.len()) - .sum() -} - -/// Renders the label set Prometheus compares the way it prints one, for the duplicate-label-set -/// error. -fn render_labels(labels: &LabelSet) -> String { - let parts = labels - .iter() - .map(|(name, value)| format!("{name}={value:?}")) - .collect::>(); - format!("{{{}}}", parts.join(", ")) -} - -/// The grid positions each final label set already occupies, per field and as merged, sorted runs. -/// -/// Shared by every partition's stream, because a label-set collision is a property of the whole -/// output vector and not of one partition. Two series that carry the same labels after the range -/// function and its maps drop a label land in different partitions, and the runtime's own per-node -/// `Emitted` table never sees both: it belongs to one execution. Checking only within a partition -/// would miss exactly those collisions. -/// -/// The key is [`crate::grid_series::label_set`]'s normalization of the final labels rather than -/// [`datatypes::series::MetricName`] itself: a label with an empty value is no label to -/// Prometheus, so `{host="h", zone=""}` and `{host="h"}` are one label set and have to be compared -/// as one. Encoding the tags as they are would file them apart and let one point hold two values. -/// -/// Occupancy is per field and per position, not per chunk: two series whose label sets became equal -/// while holding values in *different* fields merge cleanly in Prometheus, so only the same field -/// at the same position collides. That is the rule -/// [`crate::grid_series::merge_same_labelset`] applies. -/// -/// The map belongs to the plan, not to one stream: a partition that ends leaves what it claimed -/// for the partitions still running, so the entries are only released with the plan itself. The -/// charge has the same lifetime, which is what makes the bytes the map holds - one entry per label -/// set of the output, kept until the plan dies - visible to the memory pool. Charging each stream -/// for what it inserted instead returned those bytes when that stream ended, while the entries the -/// later partitions are checked against stayed: everything the map held after the first partition -/// was unaccounted for. -#[derive(Default)] -struct LabelSets { - /// The grid positions each final label set already occupies, one run list per field of the - /// series schema. Each list is sorted, disjoint and has the runs that touch joined, which is - /// what lets a claim look a position up instead of walking the history. - seen: HashMap>>, - /// What `seen` keeps charged on the task's memory pool. `None` until the first execution, - /// since construction has no `TaskContext` and therefore no pool to register against. - reservation: Option, -} - -impl LabelSets { - /// Attaches the reservation to the pool of the first execution. - /// - /// Later executions keep the first reservation: it already pays for what `seen` holds, and - /// moving it to another pool would drop the bytes other partitions are still growing. - fn attach(&mut self, pool: &Arc) { - if self.reservation.is_none() { - self.reservation = - Some(MemoryConsumer::new("FusedSeriesExec label sets").register(pool)); - } - } - - /// Claims the occupied grid positions of one chunk of a final series for its label set. - /// - /// A collision is an error only when both series hold a sample of *one field* at the same - /// evaluation timestamp, which is Prometheus' range-query rule. Series with disjoint points are - /// legal, and so are series that fill the same points field by field: Prometheus merges both - /// into one output series, while this node emits them as separate batches and defers that merge - /// to matrix assembly, so every sample survives. - fn claim(&mut self, series: &Series) -> DataFusionResult<()> { - let runs = occupied_runs(series)?; - // A chunk holding no value occupies nothing, so there is no collision to remember and no - // entry to write. Such a chunk is not rare: a selection over a series that staleness or the - // range left empty still closes it, and hands over a closing chunk with no rows. Recording - // one entry per such series would grow the table - and its charge, which the plan holds for - // its whole life - for a query that never answers a value, until an empty result exhausted - // the pool. The executor's own `Emitted::occupy` takes the same early return. - if runs.iter().all(Vec::is_empty) { - return Ok(()); - } - let key = label_set(&series.metric); - let Self { seen, reservation } = self; - // Nothing reaches a claim without an `execute` first, and an execution that ran attached the - // reservation. Claiming without one would leave the map growing uncharged, so it is refused - // rather than silently skipped. - let Some(reservation) = reservation.as_mut() else { - return Err(DataFusionError::Execution( - "FusedSeriesExec label-set state is not attached to a memory pool".to_string(), - )); - }; - match seen.entry(key) { - Entry::Occupied(entry) => { - if let Some(field) = colliding_field(entry.get(), &runs) { - return Err(DataFusionError::Execution(format!( - "vector cannot contain metrics with the same labelset: {} (field \ - {field})", - render_labels(entry.key()) - ))); - } - let stored = entry.into_mut(); - // An entry recorded from a chunk of fewer fields is padded, so every field of this - // chunk has a slot of its own: a slot is a field index of the schema, so both sides - // count the same field the same way. A chunk of *fewer* fields than the entry - // leaves the remaining slots to the chunks that fill them. - // Stage only the affected run lists, and reserve before committing. A forward - // append whose existing capacity suffices needs no copy of its stored history. - let mut growth = 0usize; - let mut shrink = 0usize; - let mut merged_runs = Vec::with_capacity(runs.len()); - for (index, incoming) in runs.iter().enumerate() { - let Some(old) = stored.get(index) else { - let copy = incoming.clone(); - growth += copy.capacity(); - merged_runs.push(Some(copy)); - continue; - }; - let append_fits = incoming.is_empty() - || (old.last().is_none_or(|last| incoming[0].start >= last.end) - && old.capacity() - >= old.len() + incoming.len() - - usize::from(old.last().is_some_and(|last| { - incoming - .first() - .is_some_and(|first| first.start == last.end) - }))); - if append_fits { - merged_runs.push(None); - } else { - let mut merged = old.clone(); - absorb(&mut merged, incoming); - if merged.capacity() >= old.capacity() { - growth += merged.capacity() - old.capacity(); - } else { - shrink += old.capacity() - merged.capacity(); - } - merged_runs.push(Some(merged)); - } - } - reservation.try_grow(growth * size_of::())?; - for (index, (incoming, merged)) in runs.iter().zip(merged_runs).enumerate() { - if let Some(merged) = merged { - if index < stored.len() { - stored[index] = merged; - } else { - stored.push(merged); - } - } else if index < stored.len() { - absorb(&mut stored[index], incoming); - } - } - reservation.try_shrink(shrink * size_of::())?; - } - Entry::Vacant(entry) => { - // Only the label-set state is this node's own cross-poll state that grows with - // the output, so it is the only thing reserved here. The runs are charged by the - // capacity they were gathered with rather than by their length, so that a later - // claim appending into that slack is charged nothing for it. - let bytes = labels_bytes(entry.key()) - + runs.iter().map(|runs| runs.capacity()).sum::() * size_of::(); - reservation.try_grow(bytes)?; - entry.insert(runs); - } - } - Ok(()) - } -} - -type SharedLabelSets = Arc>; - -/// The grid positions the rows of one chunk hold a value at, one run list per field. -/// -/// Each [`Run`] is a half-open range of absolute grid indices, and the runs of one field are sorted -/// and disjoint. Occupancy is kept per field because it is what a collision is: two series whose -/// label sets became equal merge field by field in Prometheus, so only two values of the same field -/// at one position collide. Marking the whole position as taken as soon as *any* field held a value -/// would refuse the complementary pair that merges cleanly. -fn occupied_runs(series: &Series) -> DataFusionResult>> { - // Rows are addressed on the statement's grid, so a row sits at the absolute position two - // series can collide on. An explicit axis has no such position, and a chunk never carries - // one: this is the refusal `EvaluatedChunk` already makes when the chunk is built. - if series.time.grid().is_none() { - return Err(DataFusionError::Execution( - "FusedSeriesExec requires evaluated series on a time grid".to_string(), - )); - } - let rows = series.time.len(); - let range = series.time.as_grid_range(); - let mut runs: Vec> = vec![Vec::new(); series.fields.len()]; - // A field that is one dense, null-free float run holds a value on every row the chunk carries, - // so it occupies exactly the range the chunk reports and no row of it has to be inspected. This - // is the shape the fused evaluator produces for a range function over ordinary float samples. - // Every other field is walked below, and a field filled here is skipped there. - let mut walked: Vec<(usize, &Samples)> = Vec::new(); - for (index, samples) in series.fields.iter().enumerate() { - match (range, samples.dense()) { - (Some(range), Some(DenseSamplesRef::Float(values))) - if rows > 0 && values.len() == rows && values.null_count() == 0 => - { - runs[index] = vec![Run { - start: range.start(), - end: range.start() + rows, - }]; - } - _ => walked.push((index, samples)), - } - } - if walked.is_empty() { - return Ok(runs); - } - // `positions` walks a chunk that forms one run and one whose rows form several the same way, - // in position order, so the holes between two runs are never read as rows. - let mut cursors: Vec<_> = walked.iter().map(|(_, samples)| samples.cursor()).collect(); - for (position, row) in series.time.positions() { - for ((index, _), cursor) in walked.iter().zip(&mut cursors) { - if !matches!(cursor.at(row), Some(sample) if !matches!(sample, SampleRef::Missing)) { - continue; - } - match runs[*index].last_mut() { - Some(run) if run.end == position => run.end += 1, - _ => runs[*index].push(Run { - start: position, - end: position + 1, - }), - } - } - } - Ok(runs) -} - -/// The field of `runs` that holds a value where `stored` already holds one, if any. -/// -/// A collision is two values of one field at one position, so each field is checked on its own. -/// Each incoming run is looked up rather than the stored history walked: the stored runs are sorted -/// and disjoint, so the only run that can share a position with `run` is the first one that does not -/// end before it, which is the look-up the runtime's own `Emitted::occupy` makes. Walking the -/// history instead costs a sparse series one pass over everything it stored per chunk - a query -/// whose chunks move forward one position at a time, a hole apart, is legal and would cost the -/// square of its chunk count under the lock every partition shares. -fn colliding_field(stored: &[Vec], runs: &[Vec]) -> Option { - for (field, (stored, new)) in stored.iter().zip(runs).enumerate() { - for run in new { - let at = stored.partition_point(|stored| stored.end <= run.start); - if stored.get(at).is_some_and(|stored| stored.overlaps(*run)) { - return Some(field); - } - } - } - None -} - /// The batches one execution of [`FusedSeriesExec`] yields. pub struct FusedSeriesStream { /// The input frames, read as raw series chunks, or `None` once the stream ended in an error. @@ -1375,10 +1004,6 @@ pub struct FusedSeriesStream { tags: Arc<[String]>, /// The schema of every batch this stream yields. output_schema: SchemaRef, - /// The label-set occupancy shared with the other partitions' streams. - /// - /// What it holds is charged to the plan's own reservation, attached by the first execution. - labels: SharedLabelSets, /// The batches the executor settled since the last poll of the input. ready: VecDeque, /// Whether the input ended and everything it settled has been handed out. @@ -1447,9 +1072,6 @@ impl FusedSeriesStream { self.done = true; self.executor = None; self.input = None; - // The label-set map and its charge are the plan's, so a stream stopping early gives back - // nothing: the entries another partition is still checking against have to stay charged, and - // are released when the plan itself is dropped. // A batch queued before the error was settled by a run that has now stopped. self.ready.clear(); } @@ -1457,12 +1079,6 @@ impl FusedSeriesStream { /// Converts the chunks the executor settled into batches, in the order it produced them. fn convert(&mut self, chunks: Vec) -> DataFusionResult<()> { for chunk in &chunks { - // The label set is claimed before its batch is queued, so a collision surfaces at the - // chunk that reached it and nothing this plan would refuse is handed out after it. - self.labels - .lock() - .unwrap_or_else(|e| e.into_inner()) - .claim(&chunk.series)?; let batch = series_to_record_batch( &self.classification, &self.tags, diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index c3837187474..2d98ad0fb38 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -47,8 +47,7 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::context::TaskContext; -use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; -use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::execution::memory_pool::UnboundedMemoryPool; use datafusion::logical_expr::{EmptyRelation, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_optimizer::PhysicalOptimizerRule; @@ -68,17 +67,13 @@ use datatypes::series::{Label, MetricName, SeriesOrdering}; use futures::{Stream, StreamExt, TryStreamExt}; use store_api::region_engine::RawLabelSource; -use super::{ - FusedSeries, FusedSeriesBinding, LabelSets, OutputColumn, POLL_BUDGET, label_set, - occupied_runs, series_to_record_batch, -}; +use super::{FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET, series_to_record_batch}; use crate::evaluated::EvaluatedChunk; use crate::native::execute::{NativeExecutionState, NativeExecutor}; use crate::native::plan::{ - AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundWindow, Grouping, - NativeExpr, NativeNode, NativeTree, NativeType, + AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundWindow, Grouping, NativeExpr, + NativeNode, NativeTree, NativeType, }; -use crate::native::runs::absorb_rebuilds; use crate::range_eval::{RangeFunction, RangeWindow}; /// The query every tree here is compiled from. @@ -111,8 +106,6 @@ const OTHER: &str = "other"; const HOST: &str = "host"; const ZONE: &str = "zone"; const NAME: &str = "__name__"; -/// A metadata label, which the window drops along with the metric name. -const TYPE: &str = "__type__"; const METRIC: &str = "http_requests_total"; fn grid() -> Arc { @@ -129,7 +122,7 @@ fn raw_schema() -> Arc { }) } -/// A raw schema of two float fields, so one series can fill one of them and another the other. +/// A raw schema of two float fields, used to verify a frame with different fields is rejected. fn two_field_schema() -> Arc { Arc::new(SeriesSchema { fields: Arc::from([ @@ -182,102 +175,6 @@ fn rate_tree() -> NativeTree { rate_tree_of(raw_schema()) } -/// An instant selection feeding a range window. -fn instant_window_tree() -> NativeTree { - let instant = instant_tree().root().clone(); - let grid = instant - .output - .samples() - .expect("an instant selection produces samples") - .eval_grid - .clone() - .expect("an instant selection's output is evaluated on a grid"); - let evaluated = instant.output.clone(); - NativeTree::try_new(NativeNode { - expr: NativeExpr::Window { - input: Box::new(instant), - op: BoundWindow { - function: RangeFunction::Rate, - window: RangeWindow { - range_ms: RANGE_MS, - offset_ms: 0, - anchor_ms: None, - modifier: None, - }, - input_grid: Some(Arc::clone(&grid)), - output_grid: grid, - }, - }, - output: evaluated, - }) - .expect("the bound tree is legal") -} - -/// Two windows stacked, `rate(last_over_time(x[1s])[2s])`, which no compiler produces. -/// -/// A range function has to apply directly to a selector, and a selector is raw samples, so a tree -/// with a window above a window is only reachable by writing it down by hand - which is a tree -/// [`FusedSeriesExec`](super::FusedSeriesExec) accepts, and therefore one it has to answer correctly. -/// -/// The inner window is the interesting one: `last_over_time` keeps `__name__`, so the pair it can -/// split is not two series whose names it dropped but `{host="h", zone=""}` and `{host="h"}`, one -/// canonical label set and two raw keys. Its fragments are not merged unless the node above it -/// merges them, and that merge happens once per execution: with one execution per partition each -/// partition sees one fragment, so the outer window reads one half of the pair. -fn nested_window_tree() -> NativeTree { - let grid = grid(); - // The inner window's output, which is also the outer window's output: a window reports one value - // per input field on its own grid. - let evaluated = NativeType::Samples(Arc::new(SeriesSchema { - fields: raw_schema().fields.clone(), - eval_grid: Some(Arc::clone(&grid)), - })); - let raw = NativeType::Samples(raw_schema()); - NativeTree::try_new(NativeNode { - expr: NativeExpr::Window { - input: Box::new(NativeNode { - expr: NativeExpr::Window { - input: Box::new(NativeNode { - expr: NativeExpr::Input(BoundInput { - output: raw.clone(), - ordering: SeriesOrdering::KeyThenTime, - }), - output: raw, - }), - op: BoundWindow { - function: RangeFunction::LastOverTime, - window: RangeWindow { - range_ms: 1_000, - offset_ms: 0, - anchor_ms: None, - modifier: None, - }, - // The inner window's input is raw samples, so it has no grid of its own. - input_grid: None, - output_grid: Arc::clone(&grid), - }, - }, - output: evaluated.clone(), - }), - op: BoundWindow { - function: RangeFunction::Rate, - window: RangeWindow { - range_ms: 2_000, - offset_ms: 0, - anchor_ms: None, - modifier: None, - }, - // The outer window reads the inner one's evaluated output, on that output's grid, - // and reports one value per field there like the inner one does. - input_grid: Some(Arc::clone(&grid)), - output_grid: Arc::clone(&grid), - }, - }, - output: evaluated, - }) - .expect("the bound tree is legal") -} - /// The same tree asking for `quantile`, whose state the runtime does not keep yet. fn quantile_aggregate_tree() -> NativeTree { let complete = aggregate_tree(); @@ -329,37 +226,6 @@ fn aggregate_tree() -> NativeTree { .expect("the bound tree is legal") } -/// The tree an instant selection compiles to: the last raw sample at or before each grid point. -/// -/// Unlike a window or a map, a selection rewrites no label: the series it answers carries the labels -/// its input did. -fn instant_tree() -> NativeTree { - let grid = grid(); - let evaluated = NativeType::Samples(Arc::new(SeriesSchema { - fields: raw_schema().fields.clone(), - eval_grid: Some(Arc::clone(&grid)), - })); - let raw = NativeType::Samples(raw_schema()); - NativeTree::try_new(NativeNode { - expr: NativeExpr::Instant { - input: Box::new(NativeNode { - expr: NativeExpr::Input(BoundInput { - output: raw.clone(), - ordering: SeriesOrdering::KeyThenTime, - }), - output: raw, - }), - op: BoundInstant { - lookback_ms: RANGE_MS, - offset_ms: 0, - output_grid: grid, - }, - }, - output: evaluated, - }) - .expect("the bound tree is legal") -} - /// What the tree's input was compiled against, for the input schema and output tags given. fn binding_of( tree: &NativeTree, @@ -411,6 +277,66 @@ fn output_schema() -> SchemaRef { ])) } +/// The output of [`zone_node`]: the field, `host`, `zone`, the dropped name and the grid times. +fn zone_output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(ZONE, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +/// `rate(x[5m])` over frames that carry `zone` as well as `host`, with both in the output. +fn zone_node() -> FusedSeries { + let tree = rate_tree(); + let binding = binding_of( + &tree, + raw_schema(), + vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], + ); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + zone_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The two chunks of the collision: one series once with an empty `zone` and once with no `zone`. +/// +/// They are one identity to Prometheus — an empty label value is no label — but two different raw +/// series: the scan encodes the empty value and the absent one differently, so they reach the +/// runtime as two series that close on the same label set. +fn zone_chunks() -> Vec { + let rows = series_samples(1.0); + vec![ + chunk_of( + &raw_schema(), + "zone-empty", + metric(&[(HOST, "web-1"), (ZONE, "")], Some(METRIC)), + ×_of(&rows), + vec![values_of(&rows)], + true, + ), + chunk_of( + &raw_schema(), + "zone-absent", + metric(&[(HOST, "web-1")], Some(METRIC)), + ×_of(&rows), + vec![values_of(&rows)], + true, + ), + ] +} + fn input_plan(schema: SchemaRef) -> LogicalPlan { LogicalPlan::EmptyRelation(EmptyRelation { produce_one_row: false, @@ -447,15 +373,6 @@ fn memory_exec_partitioned(partitions: Vec>) -> Arc) -> Arc { - let runtime = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::clone(pool) as Arc) - .build() - .expect("the runtime environment builds"); - Arc::new(TaskContext::default().with_runtime(Arc::new(runtime))) -} - /// A physical input whose one partition's stream `make` builds. fn staged( schema: SchemaRef, @@ -664,11 +581,6 @@ fn values_of(rows: &[(i64, f64)]) -> Vec> { rows.iter().map(|(_, value)| Some(*value)).collect() } -/// One missing sample per timestamp of `rows`, which is a field the series holds no value in. -fn no_values(rows: &[(i64, f64)]) -> Vec> { - rows.iter().map(|_| None).collect() -} - /// One chunk of `metric`, keyed `key`, holding `times` and one sample per field of `schema`. fn chunk_of( schema: &SeriesSchema, @@ -829,87 +741,6 @@ fn rows(batches: &[RecordBatch]) -> Vec { rows } -/// The output of [`zone_node`]: the field, `host`, `zone`, the dropped name and the grid times. -fn zone_output_schema() -> SchemaRef { - Arc::new(Schema::new(vec![ - Field::new(VALUE, DataType::Float64, true), - Field::new(HOST, DataType::Utf8, true), - Field::new(ZONE, DataType::Utf8, true), - Field::new(NAME, DataType::Utf8, true), - Field::new( - TIME, - DataType::Timestamp(TimeUnit::Millisecond, None), - false, - ), - ])) -} - -/// `rate(x[5m])` over frames that carry `zone` as well as `host`, with both in the output. -fn zone_node() -> FusedSeries { - let tree = rate_tree(); - let binding = binding_of( - &tree, - raw_schema(), - vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], - ); - FusedSeries::try_new( - Arc::from(PROMQL), - tree, - binding, - input_plan(frame_schema()), - zone_output_schema().to_dfschema_ref().unwrap(), - ) - .expect("the binding matches the tree") -} - -/// The instant-selection node over frames that carry `zone` as well as `host`. -/// -/// Its output labels are the input's, so an empty-valued label reaches this node's own state as the -/// frame wrote it. -fn instant_node() -> FusedSeries { - let tree = instant_tree(); - let binding = binding_of( - &tree, - raw_schema(), - vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], - ); - FusedSeries::try_new( - Arc::from(PROMQL), - tree, - binding, - input_plan(frame_schema()), - zone_output_schema().to_dfschema_ref().unwrap(), - ) - .expect("the binding matches the tree") -} - -/// The two chunks of the collision: one series once with an empty `zone` and once with no `zone`. -/// -/// They are one identity to Prometheus — an empty label value is no label — but two different raw -/// series: the scan encodes the empty value and the absent one differently, so they reach the -/// runtime as two series that close on the same label set. -fn zone_chunks() -> Vec { - let rows = series_samples(1.0); - vec![ - chunk_of( - &raw_schema(), - "zone-empty", - metric(&[(HOST, "web-1"), (ZONE, "")], Some(METRIC)), - ×_of(&rows), - vec![values_of(&rows)], - true, - ), - chunk_of( - &raw_schema(), - "zone-absent", - metric(&[(HOST, "web-1")], Some(METRIC)), - ×_of(&rows), - vec![values_of(&rows)], - true, - ), - ] -} - /// The canonical label set these tests collide on, as the runtime renders it in its refusal. const DUPLICATE_LABEL_SET: &str = "the node at root produces two series under the label set [(\"host\", \"web-1\")]"; @@ -984,7 +815,7 @@ async fn a_runtime_error_after_rows_releases_the_input() { ); } -/// Two series with *different* names collapse once `rate` drops the name. +/// A collision produced inside one native execution is still refused by its local runtime state. /// /// This is the collision the canonical label set exists for, at the layer that sees it: /// `http_requests_total{host="web-1"}` and `http_errors_total{host="web-1"}` reach the window as two @@ -1489,85 +1320,6 @@ fn an_input_with_no_partitions_is_rejected_when_the_plan_is_built() { ); } -/// Two windows stacked are refused over more than one partition, because the inner one's fragments -/// need not have been merged. -/// -/// The counterexample is legal PromQL once nesting is expressible: the inner `last_over_time` runs -/// over two raw series that are one canonical label set, `{host="h", zone=""}` holding `(1s, 1)` and -/// `{host="h"}` holding `(2s, 2)`, and the outer `rate` needs both in one window. In a single -/// partition the outer window merges the two fragments and answers `(2s, 1)`; in two partitions each -/// executor is handed one fragment and has one sample to rate, so the outer window answers nothing. -/// That is an empty result rather than an error, which is why no check further down can see it. -/// -/// `last_over_time` keeps `__name__`, so a check that looked only for a dropped metric name accepts -/// this tree. What makes it unsafe is that a window's output can hold two series under one label set -/// whenever it did not merge its input, which a bare tree cannot say. -/// -/// One partition stays accepted: there is nothing to split there, which is what makes the refusal -/// about the split rather than about the shape of the tree. -#[test] -fn a_nested_window_over_two_partitions_is_refused() { - let error = plan( - nested_window_tree(), - memory_exec_partitioned(vec![vec![], vec![]]), - ) - .expect_err("each of two partitions would merge only the fragment it was given"); - let message = error.to_string(); - assert!( - message.contains("`root.input`"), - "the refusal must name the inner window: {message}" - ); - assert!( - message.contains("`root`"), - "the refusal must name the window above it: {message}" - ); - assert!( - message.contains("2 partitions"), - "the refusal must name the partition count: {message}" - ); - - // The same tree over one partition is accepted, because one execution sees every fragment and - // merges them into one series before the window above it evaluates. - plan(nested_window_tree(), memory_exec(Vec::new())) - .expect("a single partition cannot split a logical series"); -} - -/// An instant selection can leave distinct raw labels on one canonical label set for a window. -#[test] -fn an_instant_then_window_over_two_partitions_is_refused() { - let error = plan( - instant_window_tree(), - memory_exec_partitioned(vec![vec![], vec![]]), - ) - .expect_err("each partition could hand the window only one fragment"); - let message = error.to_string(); - assert!( - message.contains("`root.input`"), - "the refusal must name the instant selection: {message}" - ); - assert!( - message.contains("`root`"), - "the refusal must name the consuming window: {message}" - ); - assert!( - message.contains("2 partitions"), - "the refusal must name the partition count: {message}" - ); -} - -/// A single partition cannot split a logical series, even when an instant feeds a window. -#[test] -fn an_instant_then_window_over_one_partition_is_accepted() { - plan(instant_window_tree(), memory_exec(Vec::new())).expect("one partition sees all fragments"); -} - -/// A window directly over raw input remains valid across partitions without a preceding rewriter. -#[test] -fn a_window_over_raw_input_on_two_partitions_is_accepted() { - plan(rate_tree(), memory_exec_partitioned(vec![vec![], vec![]])) - .expect("ordinary input-to-window trees do not split a logical series"); -} - /// One poll reads at most [`POLL_BUDGET`] chunks of input, and stopping there loses nothing. /// /// `rate(x[5m])` decides nothing until a sample reaches the grid's first point, so a series can @@ -1753,165 +1505,7 @@ fn error_chain(error: &DataFusionError) -> String { message } -/// Two partitions, one series each, and both series leave the window on the same label set. -/// -/// `rate` drops `__name__`, so `http_requests_total{host="web-1"}` reads as `{host="web-1"}` in -/// the output. The key column keeps each of the two whole inside its own partition, so a check that -/// lived inside one execution would see one series each and report nothing: only state shared by -/// the partitions sees that the output vector holds the label set twice. -#[tokio::test] -async fn a_duplicate_label_set_across_partitions_is_refused() { - let series = frame(&series_chunk("web-1", &series_samples(1.0), true)); - let exec = plan( - rate_tree(), - memory_exec_partitioned(vec![vec![series.clone()], vec![series]]), - ) - .expect("the runtime executes a rate tree"); - assert_eq!( - 2, - exec.properties().output_partitioning().partition_count(), - "this test needs one series in each of two partitions" - ); - - let error = collect(exec, SessionContext::default().task_ctx()) - .await - .expect_err("the second partition lands on the label set the first one claimed"); - let message = error_chain(&error); - assert!( - message.contains("vector cannot contain metrics with the same labelset"), - "refused for another reason: {message}" - ); - assert!( - message.contains("host=\"web-1\""), - "the label set must be readable: {message}" - ); -} - -/// The same label set on points that do not overlap is legal, and every point survives. -/// -/// A rate needs two samples in its window, so the first series only has a value at the first grid -/// point and the second only at the last: the two share a final label set but no evaluation -/// timestamp. Prometheus merges such series into one; this node emits both batches and leaves the -/// merge to matrix assembly, so the check has to allow what it would otherwise refuse. -#[tokio::test] -async fn a_duplicate_label_set_on_disjoint_grid_points_is_emitted() { - let first = vec![frame(&series_chunk( - "web-1", - &[(60_000, 60.0), (GRID_START, GRID_START as f64 / 1_000.0)], - true, - ))]; - let last = GRID_START + (GRID_LEN as i64 - 1) * GRID_STEP; - let second = vec![frame(&series_chunk( - "web-1", - &[ - (GRID_START, GRID_START as f64 / 1_000.0), - (last, last as f64 / 1_000.0), - ], - true, - ))]; - let exec = plan(rate_tree(), memory_exec_partitioned(vec![first, second])) - .expect("the runtime executes a rate tree"); - - let batches = collect(exec, SessionContext::default().task_ctx()) - .await - .expect("disjoint grid points are legal"); - assert_eq!(2, batches.len(), "each partition emits its own series"); - let mut rows = rows(&batches); - rows.sort_by_key(|row| row.0); - assert_eq!( - vec![GRID_START, last], - rows.iter().map(|row| row.0).collect::>(), - "both series must survive with the one point each holds" - ); - for ((_, value, host, _), rate) in rows.iter().zip([1.0, 1.0]) { - assert_eq!(Some("web-1"), host.as_deref()); - let value = value.expect("the point each series holds has a value"); - assert!((rate - value).abs() < 1e-9, "{value} is not {rate}"); - } -} - -/// What a label set holds stays charged after the stream that first claimed it ends, and what a -/// later partition adds to it is charged as well. -/// -/// The map belongs to the plan: it outlives every partition and is released with the plan. Charging -/// each stream for what it inserted returned the bytes of a label set when *that* stream ended, -/// while the entry the next partition is checked against stayed in the map, so everything the map -/// held after the first partition was invisible to the pool. What is charged here is therefore read -/// between the partitions: it has to survive the first stream, and the runs the second partition -/// merges into the entry have to cost what they hold. -#[tokio::test] -async fn a_label_set_stays_charged_after_its_first_stream_ends() { - // The first series holds a value only at the first grid point and the second only at the last, - // so the two share a final label set without sharing an evaluation timestamp: the second - // partition is legal and has to be answered, and it has to find the first one's entry there. - let first = vec![frame(&series_chunk( - "web-0", - &[(60_000, 60.0), (GRID_START, GRID_START as f64 / 1_000.0)], - true, - ))]; - let last = GRID_START + (GRID_LEN as i64 - 1) * GRID_STEP; - let second = vec![frame(&series_chunk( - "web-0", - &[ - (GRID_START, GRID_START as f64 / 1_000.0), - (last, last as f64 / 1_000.0), - ], - true, - ))]; - let exec = plan(rate_tree(), memory_exec_partitioned(vec![first, second])) - .expect("the runtime executes a rate tree"); - assert_eq!( - 2, - exec.properties().output_partitioning().partition_count(), - "this test needs one series in each of two partitions" - ); - let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); - let context = task_context(&pool); - - // The first partition claims the label set, answers it and ends. - let mut stream = exec - .execute(0, Arc::clone(&context)) - .expect("the node executes"); - let mut answered = Vec::new(); - while let Some(batch) = stream.next().await { - answered.push(batch.expect("the first partition runs")); - } - assert_eq!(1, answered.len(), "one series over one grid point"); - drop(stream); - let charged = pool.reserved(); - assert!( - charged > 0, - "the first partition's label set is still in the plan's map, but nothing is charged for it" - ); - - // The second partition lands on that same label set on points the first one does not hold: it - // is answered, and the run it merges into the entry is what its charge grows by. - let mut stream = exec - .execute(1, Arc::clone(&context)) - .expect("the node executes"); - let mut answered = Vec::new(); - while let Some(batch) = stream.next().await { - answered.push(batch.expect("the second partition runs")); - } - assert_eq!(1, answered.len(), "disjoint points are legal"); - drop(stream); - assert!( - pool.reserved() > charged, - "the second partition's run grew the entry, but nothing was charged for it: {} is not \ - above {charged}", - pool.reserved() - ); - - // The entries live as long as the plan, so dropping it is what gives their bytes back. - drop(exec); - assert_eq!( - 0, - pool.reserved(), - "the plan's drop releases what its label sets held" - ); -} - -/// Each partition runs its own series and answers it, with no label set claimed twice. +/// Each partition runs its own series and answers it independently. #[tokio::test] async fn every_partition_evaluates_its_own_series() { let partitions = [1.0, 2.0, 3.0] @@ -1998,7 +1592,7 @@ fn the_optimizer_inserts_no_exchange_between_the_input_and_this_node() { reorders a series' chunks: the node would close a series it has not finished" ); // No exchange is not the same as no change: a rule that raised the count some other way would - // still split a series. The count is the input's, which is what the node reports and what it + // The count is the input's, which is what the node reports and what it // runs one evaluator per. assert_eq!( 2, @@ -2118,654 +1712,6 @@ async fn a_partition_past_the_inputs_count_is_refused() { ); } -/// The `rate(x[5m])` tree over two fields, with both fields and the tags in the output. -fn two_field_node() -> FusedSeries { - let tree = rate_tree_of(two_field_schema()); - let binding = binding_of( - &tree, - two_field_schema(), - vec![HOST.to_string(), NAME.to_string()], - ); - FusedSeries::try_new( - Arc::from(PROMQL), - tree, - binding, - input_plan(frame_schema_of(&two_field_schema())), - two_field_output_schema().to_dfschema_ref().unwrap(), - ) - .expect("the binding matches the tree") -} - -/// The output of [`two_field_node`]: both fields, the tags and the grid times. -fn two_field_output_schema() -> SchemaRef { - Arc::new(Schema::new(vec![ - Field::new(VALUE, DataType::Float64, true), - Field::new(OTHER, DataType::Float64, true), - Field::new(HOST, DataType::Utf8, true), - Field::new(NAME, DataType::Utf8, true), - Field::new( - TIME, - DataType::Timestamp(TimeUnit::Millisecond, None), - false, - ), - ])) -} - -/// One chunk of `{__name__=METRIC, host=, __type__=}` whose `field` holds `rows` and -/// whose other field is missing everywhere. -fn one_field_chunk(field: usize, host: &str, kind: &str, last: bool) -> RawSeriesChunk { - let rows = series_samples(1.0); - let (value, other) = if field == 0 { - (values_of(&rows), no_values(&rows)) - } else { - (no_values(&rows), values_of(&rows)) - }; - chunk_of( - &two_field_schema(), - &format!("{host}-{kind}"), - metric(&[(HOST, host), (TYPE, kind)], Some(METRIC)), - ×_of(&rows), - vec![value, other], - last, - ) -} - -/// Two series that close on one label set may still fill it, one field each. -/// -/// `rate` drops `__name__` and `__type__`, so the two series below leave the window as -/// `{host="web-1"}` and read as one label set. They hold values in *different* fields, which is the -/// pair [`crate::grid_series::merge_same_labelset`] merges: one value of one field at each position, -/// never two. Occupancy is therefore per field, and refusing this pair would refuse a series -/// Prometheus answers. -#[tokio::test] -async fn a_duplicate_label_set_of_complementary_fields_is_emitted() { - let exec = plan_of( - two_field_node(), - memory_exec_partitioned_with( - frame_schema_of(&two_field_schema()), - vec![ - frames_of( - &two_field_schema(), - &[one_field_chunk(0, "web-1", "counter", true)], - ), - frames_of( - &two_field_schema(), - &[one_field_chunk(1, "web-1", "gauge", true)], - ), - ], - ), - ) - .expect("the runtime executes a rate tree"); - assert_eq!( - 2, - exec.properties().output_partitioning().partition_count(), - "this test needs one series in each of two partitions" - ); - - let batches = collect(exec, SessionContext::default().task_ctx()) - .await - .expect("complementary fields are legal"); - assert_eq!(2, batches.len(), "each partition emits its own series"); - // Each batch answers one field of the shared label set and leaves the other missing: the two - // together are the merged series, with no value of one field holding a position twice. - let answered = batches - .iter() - .map(|batch| { - let host = batch - .column_by_name(HOST) - .and_then(|host| host.as_any().downcast_ref::()) - .expect("every batch has its tag column"); - assert!( - (0..host.len()).all(|row| host.value(row) == "web-1"), - "both batches answer the one label set" - ); - let present = [VALUE, OTHER] - .into_iter() - .filter(|field| { - let column = batch - .column_by_name(field) - .unwrap_or_else(|| panic!("every batch has a {field} column")) - .as_any() - .downcast_ref::() - .expect("the field column is Float64"); - column.null_count() < column.len() - }) - .collect::>(); - assert_eq!( - 1, - present.len(), - "each series fills exactly one field: {present:?}" - ); - present[0] - }) - .collect::>(); - let mut answered = answered; - answered.sort_unstable(); - assert_eq!(vec![OTHER, VALUE], answered, "the two fields are filled"); -} - -/// Two series that close on one label set may not hold the same field at the same position. -/// -/// This is the pair no merge resolves: at one grid position one field of the label set would carry -/// two values. It is refused even though the two series name different `__type__` labels, because -/// the window drops that label and what reaches the output is one label set twice. -#[tokio::test] -async fn a_duplicate_label_set_of_one_field_is_refused() { - let exec = plan_of( - two_field_node(), - memory_exec_partitioned_with( - frame_schema_of(&two_field_schema()), - vec![ - frames_of( - &two_field_schema(), - &[one_field_chunk(0, "web-1", "counter", true)], - ), - frames_of( - &two_field_schema(), - &[one_field_chunk(0, "web-1", "gauge", true)], - ), - ], - ), - ) - .expect("the runtime executes a rate tree"); - - let error = collect(exec, SessionContext::default().task_ctx()) - .await - .expect_err("both series hold a value of one field at one position"); - let message = error_chain(&error); - assert!( - message.contains("vector cannot contain metrics with the same labelset"), - "refused for another reason: {message}" - ); - assert!( - message.contains("host=\"web-1\""), - "the label set must be readable: {message}" - ); -} - -/// Two partitions whose label sets differ only in an empty value hold one label set. -/// -/// `{host="web-1", zone=""}` and `{host="web-1"}` are one series to Prometheus: an empty label value -/// is no label. An instant selection rewrites nothing, so both label sets reach this node's own -/// state as their frames wrote them, and a state that compared the tags as they arrived would file -/// them apart. Neither partition sees the other's series, so nothing else in the plan can refuse -/// the pair: this node's state is what compares them, and what it has to compare is the label set -/// Prometheus compares. -#[tokio::test] -async fn partitions_whose_label_sets_differ_only_in_an_empty_value_are_refused() { - let [empty, absent] = <[_; 2]>::try_from(zone_chunks()).expect("the fixture has two chunks"); - let exec = plan_of( - instant_node(), - memory_exec_partitioned(vec![frames(&[empty]), frames(&[absent])]), - ) - .expect("the runtime executes an instant selection"); - assert_eq!( - 2, - exec.properties().output_partitioning().partition_count(), - "this test needs one series in each of two partitions" - ); - - let error = collect(exec, SessionContext::default().task_ctx()) - .await - .expect_err("both partitions hold one value of one field at the same position"); - let message = error_chain(&error); - assert!( - message.contains("vector cannot contain metrics with the same labelset"), - "refused for another reason: {message}" - ); - assert!( - message.contains("host=\"web-1\""), - "the label set must be readable: {message}" - ); -} - -/// A chunk holding no value records nothing, so a query that answers nothing holds nothing. -/// -/// A selection over a series that staleness or the range left empty still closes that series, and -/// hands over a closing chunk with no rows - [`crate::native::execute::no_row_chunk`] is where it is -/// built. Every chunk reaches [`LabelSets::claim`], so an implementation that wrote an entry per -/// chunk would record one per empty series and charge the plan's reservation for it, for a query -/// whose answer is empty: enough of them would exhaust a finite pool without a single row produced. -/// -/// The entries are checked rather than the charge, because a charge of zero is what a mistake here -/// would still show if the label bytes happened to be free: what must not grow is the table. -#[test] -fn a_chunk_holding_no_value_is_claimed_without_recording_it() { - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - let mut labels = LabelSets::default(); - labels.attach(&pool); - let axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 0)]) - .expect("an empty run over the grid"); - let empty_series = |index: usize| Series { - metric: metric(&[(HOST, &format!("web-{index}"))], Some(METRIC)), - time: axis.clone(), - // Five positions were never filled, so no field occupies one. - fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![f64::NAN; 0]))), - }; - - for index in 0..64 { - labels - .claim(&empty_series(index)) - .expect("an empty chunk cannot collide with anything"); - } - - assert_eq!( - 0, - labels.seen.len(), - "64 empty series were recorded, and the table outlives the query: {:?}", - labels.seen.keys().collect::>() - ); - assert_eq!( - 0, - pool.reserved(), - "nothing is held, so nothing may be charged" - ); -} - -/// An empty chunk does not stop a later chunk of the same label set from being recorded. -/// -/// Skipping the empty chunk is only safe while it also records nothing a later one needs: a series -/// that answers nothing holds no position, so a second series under one label set has to be free to -/// claim the positions it really holds. -#[test] -fn a_later_chunk_still_claims_the_label_set_an_empty_one_skipped() { - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - let mut labels = LabelSets::default(); - labels.attach(&pool); - let empty_axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 0)]) - .expect("an empty run over the grid"); - let full_axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 5)]) - .expect("one run over the whole grid"); - - labels - .claim(&Series { - metric: metric(&[(HOST, "h")], Some(METRIC)), - time: empty_axis, - fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![f64::NAN; 0]))), - }) - .expect("an empty chunk cannot collide with anything"); - labels - .claim(&Series { - metric: metric(&[(HOST, "h")], Some(METRIC)), - time: full_axis, - fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; 5]))), - }) - .expect("the first chunk to hold a value claims the positions it holds"); - - assert_eq!(1, labels.seen.len(), "the positions are recorded once"); - assert!(pool.reserved() > 0, "the recorded positions are charged"); -} - -/// The claim compares the label set Prometheus compares, so an empty value is no label. -/// -/// Driven directly at the layer that keeps the state: a window and a map both drop an empty-valued -/// label before a chunk reaches this node, so no tree this crate can build produces the pair below. -/// The state has to key on the same normalization they apply all the same, or the equality it -/// enforces would be one Prometheus does not have. -#[test] -fn a_claim_compares_label_sets_without_empty_values() { - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - let mut labels = LabelSets::default(); - labels.attach(&pool); - let axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 5)]) - .expect("one run over the whole grid"); - let series_of = |metric: MetricName| Series { - metric, - time: axis.clone(), - fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; 5]))), - }; - - labels - .claim(&series_of(metric(&[(HOST, "h"), (ZONE, "")], Some(METRIC)))) - .expect("the first label set is claimed"); - let error = labels - .claim(&series_of(metric(&[(HOST, "h")], Some(METRIC)))) - .expect_err("an empty label value is no label"); - assert!( - error - .to_string() - .contains("vector cannot contain metrics with the same labelset"), - "refused for another reason: {error}" - ); - assert_eq!( - vec![(Arc::from(HOST), Arc::from("h"))], - label_set(&metric(&[(HOST, "h"), (ZONE, "")], None)), - "the empty value is not part of the label set" - ); -} - -/// A series whose chunks move forward, a hole apart, appends to its stored runs rather than -/// re-merging them. -/// -/// One grid point per chunk with a hole after it (`0, 2, 4, ...`) is a legal stream and the shape a -/// short `count_over_time` over sub-sampled, chunked input produces: the chunks of one label set -/// hold their values at disjoint positions, so they merge into one series and no position ever -/// carries two values. What it used to cost was quadratic - `claim` walked the whole stored history -/// to look for a collision and then rebuilt all of it with a merge, for each chunk - and all of it -/// ran under the label-set lock every partition shares, so a sparse series of `K` chunks made the -/// other partitions wait for `K(K+1)/2` run moves. -/// -/// The counter is the runtime's own, so this test measures the append path the native table takes -/// as well: one fix, not two that can drift apart again. -#[test] -fn a_forward_chunk_does_not_rebuild_the_stored_runs() { - // How many one-position chunks the series arrives in: enough that the quadratic cost it used to - // pay would be unmistakable, since that cost was the square of this figure. - const CHUNKS: usize = 64; - - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - let mut labels = LabelSets::default(); - labels.attach(&pool); - // Two positions per chunk, so a hole follows every point the series holds. - let grid = Arc::new( - EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 2 * CHUNKS as u64).unwrap()) - .unwrap(), - ); - let chunk = |position: usize| Series { - metric: metric(&[(HOST, "h")], Some(METRIC)), - time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(position, 1)]) - .expect("one position of the grid"), - fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), - }; - - let before = absorb_rebuilds(); - for index in 0..CHUNKS { - labels - .claim(&chunk(2 * index)) - .expect("the chunks of one label set hold values at disjoint positions"); - } - - assert_eq!( - 0, - absorb_rebuilds() - before, - "the {CHUNKS} forward chunks rebuilt the stored runs instead of appending to them" - ); - assert_eq!( - 1, - labels.seen.len(), - "every chunk belongs to one label set, so the history is one entry" - ); -} - -/// A backward disjoint claim after a spare-capacity run merge updates the exact memory charge. -#[test] -fn a_backward_disjoint_claim_after_spare_capacity_is_charged_correctly() { - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - let mut labels = LabelSets::default(); - labels.attach(&pool); - let grid = - Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); - let chunk = |start, len| Series { - metric: metric(&[(HOST, "h")], Some(METRIC)), - time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, len)]) - .unwrap(), - fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; len]))), - }; - labels.claim(&chunk(10, 2)).unwrap(); - let runs = labels - .seen - .values_mut() - .next() - .unwrap() - .first_mut() - .unwrap(); - let old_capacity = runs.capacity(); - runs.reserve(16); - let new_capacity = runs.capacity(); - labels - .reservation - .as_mut() - .unwrap() - .try_grow((new_capacity - old_capacity) * std::mem::size_of::()) - .unwrap(); - let before = pool.reserved(); - - labels - .claim(&chunk(0, 1)) - .expect("the position is disjoint and earlier than stored history"); - let runs = labels.seen.values().next().unwrap().first().unwrap(); - assert_eq!( - vec![(0, 1), (10, 12)], - runs.iter() - .map(|run| (run.start, run.end)) - .collect::>() - ); - assert_eq!( - before - (new_capacity - runs.capacity()) * std::mem::size_of::(), - pool.reserved(), - "the accounting follows the replacement vector's actual capacity" - ); -} - -/// A backward claim that joins two spare-capacity runs charges a replacement vector's shrinkage. -#[test] -fn a_backward_claim_merging_spare_capacity_runs_accounts_for_shrinkage() { - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - let mut labels = LabelSets::default(); - labels.attach(&pool); - let grid = - Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 3).unwrap()).unwrap()); - let chunk = |values: Vec>| Series { - metric: metric(&[(HOST, "shrink")], Some(METRIC)), - time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 3)]) - .expect("the chunk covers the small grid"), - fields: SeriesFields::one(Samples::from_float(Float64Array::from(values))), - }; - - labels - .claim(&chunk(vec![Some(1.0), None, Some(1.0)])) - .expect("the first claim stores runs at positions 0 and 2"); - let runs = labels.seen.values().next().unwrap().first().unwrap(); - assert_eq!( - vec![(0, 1), (2, 3)], - runs.iter() - .map(|run| (run.start, run.end)) - .collect::>() - ); - let before_capacity = runs.capacity(); - let before_reserved = pool.reserved(); - - labels - .claim(&chunk(vec![None, Some(1.0), None])) - .expect("the disjoint middle position joins both existing runs"); - - let runs = labels.seen.values().next().unwrap().first().unwrap(); - let after_capacity = runs.capacity(); - assert_eq!( - vec![(0, 3)], - runs.iter() - .map(|run| (run.start, run.end)) - .collect::>() - ); - assert!( - before_capacity > after_capacity, - "the regression requires replacement-vector shrinkage (before={before_capacity}, after={after_capacity})" - ); - let run_size = std::mem::size_of::(); - assert_eq!( - before_reserved - (before_capacity - after_capacity) * run_size, - pool.reserved(), - "the reservation tracks the replacement vector's actual capacity decrease" - ); -} - -/// A newly retained field is charged for the cloned run vector's capacity, not incoming slack. -#[test] -fn a_new_field_claim_is_charged_for_its_retained_capacity() { - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - let mut labels = LabelSets::default(); - labels.attach(&pool); - let grid = - Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 4).unwrap()).unwrap()); - let first = Series { - metric: metric(&[(HOST, "new-field")], Some(METRIC)), - time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 1)]).unwrap(), - fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), - }; - labels.claim(&first).unwrap(); - - let incoming = Series { - metric: metric(&[(HOST, "new-field")], Some(METRIC)), - time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 4)]).unwrap(), - fields: SeriesFields::try_from_vec(vec![ - Samples::from_float(Float64Array::from(vec![None, None, None, None])), - Samples::from_float(Float64Array::from(vec![Some(1.0), None, Some(2.0), None])), - ]) - .unwrap(), - }; - let incoming_runs = occupied_runs(&incoming).unwrap(); - assert_eq!(2, incoming_runs[1].len()); - assert!( - incoming_runs[1].capacity() > incoming_runs[1].len(), - "the incoming run vector has spare capacity" - ); - let before = pool.reserved(); - labels.claim(&incoming).unwrap(); - - let retained = &labels.seen.values().next().unwrap()[1]; - assert_eq!(2, retained.len()); - assert_eq!( - retained.capacity() * std::mem::size_of::(), - pool.reserved() - before, - "the reservation charges the retained cloned vector capacity" - ); - assert!( - retained.capacity() < incoming_runs[1].capacity(), - "the incoming spare capacity was not retained or charged" - ); -} - -/// Chunked fields are traversed sequentially while occupancy is converted into grid runs. -#[test] -fn chunked_samples_occupancy_preserves_sparse_rows() { - let values = Samples::try_from_chunks(vec![ - Samples::from_float(Float64Array::from(vec![Some(1.0), None])), - Samples::from_float(Float64Array::from(vec![None, Some(2.0)])), - ]) - .unwrap(); - let series = Series { - metric: metric(&[(HOST, "chunked")], Some(METRIC)), - time: TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 4)]).unwrap(), - fields: SeriesFields::one(values), - }; - let runs = occupied_runs(&series).expect("chunked occupancy is read through a forward cursor"); - assert_eq!( - vec![(0, 1), (3, 4)], - runs[0] - .iter() - .map(|run| (run.start, run.end)) - .collect::>() - ); -} - -/// Appending a run that touches stored history reuses its last slot instead of rebuilding it. -#[test] -fn a_touching_append_with_full_capacity_does_not_rebuild_history() { - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - let mut labels = LabelSets::default(); - labels.attach(&pool); - let grid = - Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); - let chunk = |start| Series { - metric: metric(&[(HOST, "touch")], Some(METRIC)), - time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, 1)]).unwrap(), - fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), - }; - labels.claim(&chunk(0)).unwrap(); - let stored = labels - .seen - .values_mut() - .next() - .unwrap() - .first_mut() - .unwrap(); - stored.shrink_to_fit(); - let cap = stored.capacity(); - let allocation = stored.as_ptr(); - assert_eq!( - stored.len(), - cap, - "the test requires a full-capacity run list" - ); - let before = absorb_rebuilds(); - labels - .claim(&chunk(1)) - .expect("the touching position extends the same run"); - let stored = labels.seen.values().next().unwrap().first().unwrap(); - assert_eq!( - vec![(0, 2)], - stored - .iter() - .map(|run| (run.start, run.end)) - .collect::>() - ); - assert_eq!( - cap, - stored.capacity(), - "extending a run did not allocate or copy its history" - ); - assert_eq!( - allocation, - stored.as_ptr(), - "extending a run retained the stored vector allocation" - ); - assert_eq!( - before, - absorb_rebuilds(), - "touching append uses the append path" - ); -} - -/// A failed claim does not mutate the shared entry, while its reservation owner remains alive. -#[test] -fn a_refused_growth_keeps_shared_label_entry_unchanged() { - use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; - let limit = 4096; - let pool = Arc::new(GreedyMemoryPool::new(limit)); - let erased: Arc = pool.clone(); - let mut labels = LabelSets::default(); - labels.attach(&erased); - let grid = - Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); - let chunk = |host: &str, start| Series { - metric: metric(&[(HOST, host)], Some(METRIC)), - time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, 1)]).unwrap(), - fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), - }; - labels - .claim(&chunk("shared", 0)) - .expect("initial label fits"); - let key = label_set(&chunk("shared", 0).metric); - let before_runs = labels.seen[&key][0] - .iter() - .map(|run| (run.start, run.end)) - .collect::>(); - let before_reserved = pool.reserved(); - let blocker: MemoryReservation = MemoryConsumer::new("test blocker").register(&erased); - blocker - .try_grow(limit - before_reserved) - .expect("the blocker occupies remaining capacity"); - let error = labels - .claim(&chunk("shared", 2)) - .expect_err("growth exceeds the pool"); - assert!(!error.to_string().is_empty(), "pool refusal is reported"); - assert_eq!( - before_runs, - labels.seen[&key][0] - .iter() - .map(|run| (run.start, run.end)) - .collect::>(), - "the persistent entry is unchanged" - ); - assert_eq!(limit, pool.reserved(), "both reservations remain live"); - drop(blocker); - assert_eq!( - before_reserved, - pool.reserved(), - "the plan-owned entry remains charged" - ); -} - /// Binding ordering includes field names, while schema ordering refuses unequal DF schemas. #[test] fn unequal_binding_and_output_schema_are_not_ordered_equal() { @@ -2806,81 +1752,3 @@ fn unequal_binding_and_output_schema_are_not_ordered_equal() { left_node.partial_cmp(&right_node) ); } - -/// A chunk reaching back before the last stored run still merges into the history. -/// -/// The forward path only takes a chunk that starts at or after the end of the last stored run; a -/// chunk that starts before it can interleave with what is stored, so it has to go through the full -/// merge, and the result still has to be sorted, disjoint and with the runs that touch joined - -/// which is what the look-up in [`LabelSets::claim`] relies on. -/// -/// A chunk reaching back over a stored run is the other case, and it stays refused: `[0, 11)` shares -/// position 10 with the stored `[10, 12)`, so it is a collision and the stored history is left -/// as it was. -#[test] -fn a_chunk_reaching_back_still_merges() { - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - let mut labels = LabelSets::default(); - labels.attach(&pool); - let grid = - Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); - // One chunk holding `len` positions from `start`, every one of them with a value. - let chunk = |start: usize, len: usize| Series { - metric: metric(&[(HOST, "h")], Some(METRIC)), - time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, len)]) - .expect("one run over the grid"), - fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; len]))), - }; - // The runs stored under the one label set, which one field of the output holds. - let stored = |labels: &LabelSets| -> Vec<(usize, usize)> { - labels - .seen - .values() - .next() - .expect("the chunk was recorded under one label set") - .first() - .expect("the series has one field") - .iter() - .map(|run| (run.start, run.end)) - .collect() - }; - - labels - .claim(&chunk(10, 2)) - .expect("the first chunk claims the positions it holds"); - assert_eq!( - vec![(10, 12)], - stored(&labels), - "the run `[10, 12)` is stored" - ); - - let error = labels - .claim(&chunk(0, 11)) - .expect_err("`[0, 11)` holds a value at position 10, where `[10, 12)` holds one"); - assert!( - error - .to_string() - .contains("vector cannot contain metrics with the same labelset"), - "refused for another reason: {error}" - ); - assert_eq!( - vec![(10, 12)], - stored(&labels), - "a refused claim left the stored history as it was" - ); - - let before = absorb_rebuilds(); - labels - .claim(&chunk(0, 10)) - .expect("`[0, 10)` reaches back before the last stored run without sharing a position"); - assert_eq!( - vec![(0, 12)], - stored(&labels), - "the run reaching back did not merge into the stored one and join it" - ); - assert_eq!( - 1, - absorb_rebuilds() - before, - "a chunk starting before the last stored run was appended rather than merged" - ); -} diff --git a/src/promql/src/extension_plan/planner.rs b/src/promql/src/extension_plan/planner.rs index 50cf4543d59..701ca4be01d 100644 --- a/src/promql/src/extension_plan/planner.rs +++ b/src/promql/src/extension_plan/planner.rs @@ -16,15 +16,16 @@ use std::sync::Arc; use async_trait::async_trait; use datafusion::catalog::Session; -use datafusion::error::Result as DfResult; +use datafusion::error::{DataFusionError, Result as DfResult}; use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner}; +use table::table::scan::RegionScanExec; use crate::extension_plan::{ - Absent, EmptyMetric, HistogramFold, InstantManipulate, RangeManipulate, ScalarCalculate, - SeriesDivide, SeriesNormalize, UnionDistinctOn, + Absent, EmptyMetric, FusedSeries, HistogramFold, InstantManipulate, RangeManipulate, + ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, }; pub struct PromExtensionPlanner; @@ -65,8 +66,102 @@ impl ExtensionPlanner for PromExtensionPlanner { ))) } else if let Some(node) = node.as_any().downcast_ref::() { Ok(Some(node.to_execution_plan(physical_inputs[0].clone()))) + } else if let Some(node) = node.as_any().downcast_ref::() { + plan_fused_series( + node, + physical_inputs, + session.config_options().execution.target_partitions, + ) } else { Ok(None) } } } + +/// The label `__name__` is reported under, which is a column of this node's output but not a tag +/// column of any region. +const METRIC_NAME: &str = "__name__"; + +/// Builds the physical node of a [`FusedSeries`] over the raw series leaf its input can provide. +/// +/// The node reads a raw series frame, so what it needs below it is the region scan itself: the +/// scan's own raw-series entrypoint is what packs a series' chunks contiguously and declares their +/// ordering. Anything else in between - a filter, a projection, an exchange - produces a different +/// frame, or no frame at all, and this layer cannot look under it. +/// +/// `Ok(None)` is therefore the answer for a child that is not a region scan: the logical plan is +/// already fixed and DataFusion reports an unplanned extension node normally. A region scan that is +/// found but refuses the raw path is a different case — the logical route chose this node for this +/// scan and nothing here can undo that choice, because the logical plan is not available — so it is +/// reported as an error that names the reason instead of silently producing no plan. +/// The fused leaf needs a prepared scan: range assignment runs after physical planning and cannot +/// see a `SeriesScanExec`. Preparation keeps the scan's parallel layout, and the series-producer +/// contract guarantees that one series stays in one scan partition. +fn plan_fused_series( + fused: &FusedSeries, + physical_inputs: &[Arc], + target_partitions: usize, +) -> DfResult>> { + let Some(input) = physical_inputs.first() else { + return Ok(None); + }; + let Some(scan) = input.downcast_ref::() else { + return Ok(None); + }; + + let binding = fused.binding(); + let Some(metric_name) = binding.metric_name.as_deref() else { + return Err(DataFusionError::Plan( + "the FusedSeries node binds no metric name, so the region scan below it cannot be \ + asked for the series of one metric" + .to_string(), + )); + }; + // `__name__` is the metric name rather than a column of the region, so it is not one of the + // logical tags the scan is asked to decode; the frame names each series' own metric. + let logical_tags = binding + .tags + .iter() + .filter(|tag| tag.as_str() != METRIC_NAME) + .cloned() + .collect::>(); + let field_names = binding + .series_schema + .fields + .iter() + .map(|field| field.name.to_string()) + .collect::>(); + + let scan = if scan.is_partition_set() { + Arc::new(scan.clone()) + } else { + let partitions = table::table::scan::assign_partition_ranges( + scan.get_partition_ranges(), + target_partitions.max(1), + ); + Arc::new( + scan.with_new_partitions(partitions, target_partitions.max(1)) + .map_err(|error| DataFusionError::External(error.into_inner()))?, + ) + }; + // The scan resolves the label source and the ordering it will declare itself, against its own + // metadata snapshot, and hands both back through the leaf it builds. Naming them here instead + // would be a second copy of a binding only the scanner can make: the physical metadata and + // column ids a `RawLabelSource::PrimaryKey` carries are not reachable from a region scan's + // public surface, so the layer that resolves them is the scanner. + let Some(series_scan) = + scan.try_to_named_series_scan(metric_name, &logical_tags, &field_names)? + else { + return Err(DataFusionError::Plan(format!( + "the region scan under FusedSeries could not serve a raw series frame for metric \ + {metric_name:?}: the scanner refused the raw series path for this scan, so this node \ + has no input to read" + ))); + }; + let label_source = series_scan.label_source().clone(); + let ordering = series_scan.raw_series_ordering(); + + fused + .to_execution_plan(series_scan, label_source, ordering) + .map(Some) +} diff --git a/src/promql/src/native/plan.rs b/src/promql/src/native/plan.rs index fdff0450cfa..96e24d516c4 100644 --- a/src/promql/src/native/plan.rs +++ b/src/promql/src/native/plan.rs @@ -349,6 +349,28 @@ pub struct NativeNode { pub output: NativeType, } +/// Whether `node` reads the whole input vector at a point rather than one series at a time. +/// +/// Three kinds do: an `Aggregate` folds every series of its group, a `Binary` pairs two vectors' +/// series by label, and a `Vector` operation (`sort`, `absent`, `topk`, `limitk`, +/// `histogram_quantile`, ...) reorders or reads the whole vector at a grid point. Everything else - +/// an input, a scalar, a selection, a window, a map - answers one series from one series' rows. +/// +/// Both a region-side placement and the native route's admission ask this, and they must agree: a +/// tree admitted here but refused there becomes a node the physical layer cannot build, and one +/// refused here but admitted there would be pushed to a region that holds part of the metric. +pub fn reads_across_series(node: &NativeNode) -> bool { + match &node.expr { + NativeExpr::Input(_) | NativeExpr::Scalar(_) => false, + NativeExpr::Instant { input, .. } + | NativeExpr::Window { input, .. } + | NativeExpr::Map { input, .. } => reads_across_series(input), + NativeExpr::Aggregate { .. } | NativeExpr::Binary { .. } | NativeExpr::Vector { .. } => { + true + } + } +} + /// A bound evaluation tree. /// /// The tree itself carries no run-time state: every execution builds its own, so one compiled plan diff --git a/src/promql/src/native/runs.rs b/src/promql/src/native/runs.rs index 8de325a5386..5322221ffab 100644 --- a/src/promql/src/native/runs.rs +++ b/src/promql/src/native/runs.rs @@ -21,12 +21,9 @@ //! merged away, while a chunk that fills positions the set left empty merges cleanly, as //! [`crate::grid_series::merge_same_labelset`] does. //! -//! The state is kept per label set, per field, as runs, so a chunk of many positions costs one run -//! rather than one entry per position. Keeping it in one place is what keeps the two layers that -//! check it — the native runtime's own [`crate::native::execute::Emitted`] table and the -//! cross-partition [`crate::extension_plan::fused_series::LabelSets`] — from drifting apart: both -//! ask the same functions for the same answers, so the fast path one of them earned is the fast path -//! the other gets. +//! The native runtime's partition-local [`crate::native::execute::Emitted`] table keeps this state +//! per label set and field as runs, so a chunk of many positions costs one run rather than one +//! entry per position. /// One run of grid positions, half-open, that hold a value in one field. #[derive(Clone, Copy, Debug)] diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index a39514a0e43..92ec8502fab 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -37,7 +37,7 @@ use datafusion_optimizer::push_down_filter::PushDownFilter; use datafusion_optimizer::rewrite_set_comparison::RewriteSetComparison; use datafusion_optimizer::scalar_subquery_to_join::ScalarSubqueryToJoin; use datafusion_optimizer::simplify_expressions::SimplifyExpressions; -use promql::extension_plan::SeriesDivide; +use promql::extension_plan::{FusedSeries, SeriesDivide}; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use table::metadata::TableType; use table::table::adapter::DfTableProviderAdapter; @@ -106,8 +106,35 @@ impl ExtensionOptions for DistPlannerOptions { } } -#[derive(Debug)] -pub struct DistPlannerAnalyzer; +/// Rewrites a logical plan into a distributed plan: the part above a merge scan runs on the +/// frontend, and every `MergeScan` carries the sub-plan that runs in a region. +#[derive(Debug, Default)] +pub struct DistPlannerAnalyzer { + /// Whether every region a rewritten plan reads is reached in process. + /// + /// A region input normally has to be encoded to Substrait before it can leave the frontend, and + /// [`PlanRewriter::should_expand`] cuts the plan at the first node that does not encode. In + /// process, a root `FusedSeries` plan travels as a `LogicalPlan` instead and skips that gate; + /// other plans retain the ordinary encoding checks. + /// + /// Whoever builds the query engine sets this, once, from the transport it wired. It is neither a + /// query-time decision nor a property of a region. + in_process_region_transport: bool, +} + +impl DistPlannerAnalyzer { + /// Creates the analyzer. + /// + /// `in_process_region_transport` must only be `true` for an engine whose region reads are served + /// by region servers in the same process. Its plans are then handed to + /// `RegionServer::handle_read` as `LogicalPlan`s; the remote read path + /// (`RegionServer::handle_remote_read`) decodes Substrait bytes and cannot do that. + pub fn new(in_process_region_transport: bool) -> Self { + Self { + in_process_region_transport, + } + } +} impl AnalyzerRule for DistPlannerAnalyzer { fn name(&self) -> &str { @@ -250,12 +277,15 @@ fn unwrap_dictionary_literals(plan: LogicalPlan) -> DfResult DfResult { + let in_process_region_transport = self.in_process_region_transport; // Use the subquery-aware transform so expression subqueries (including // nested scalar subqueries left in place by DataFusion 55's // `enable_physical_uncorrelated_scalar_subquery`) are visited at every // depth and their table scans get wrapped in `MergeScan`. - let plan = plan.transform_with_subqueries(&Self::inspect_plan_with_subquery)?; - let mut rewriter = PlanRewriter::new(&plan.data); + let plan = plan.transform_with_subqueries(&|plan| { + Self::inspect_plan_with_subquery(plan, in_process_region_transport) + })?; + let mut rewriter = PlanRewriter::new(&plan.data, in_process_region_transport); let result = plan.data.rewrite(&mut rewriter)?.data; Self::assign_merge_scan_remote_dyn_filter_producer_ids(result) } @@ -267,7 +297,10 @@ impl DistPlannerAnalyzer { Self::assign_merge_scan_remote_dyn_filter_producer_ids(result) } - fn inspect_plan_with_subquery(plan: LogicalPlan) -> DfResult> { + fn inspect_plan_with_subquery( + plan: LogicalPlan, + in_process_region_transport: bool, + ) -> DfResult> { // Workaround for https://github.com/GreptimeTeam/greptimedb/issues/5469 and https://github.com/GreptimeTeam/greptimedb/issues/5799 // FIXME(yingwen): Remove the `Limit` plan once we update DataFusion. if let LogicalPlan::Limit(_) | LogicalPlan::Distinct(_) = &plan { @@ -277,7 +310,12 @@ impl DistPlannerAnalyzer { let exprs = plan .expressions_consider_join() .into_iter() - .map(|e| e.transform(&Self::transform_subquery).map(|x| x.data)) + .map(|e| { + e.transform(|subquery_expr| { + Self::transform_subquery(subquery_expr, in_process_region_transport) + }) + .map(|x| x.data) + }) .collect::>>()?; // Some plans that are special treated (should not call `with_new_exprs` on them) @@ -289,27 +327,33 @@ impl DistPlannerAnalyzer { } } - fn transform_subquery(expr: Expr) -> DfResult> { + fn transform_subquery( + expr: Expr, + in_process_region_transport: bool, + ) -> DfResult> { match expr { Expr::Exists(exists) => Ok(Transformed::yes(Expr::Exists(Exists { - subquery: Self::handle_subquery(exists.subquery)?, + subquery: Self::handle_subquery(exists.subquery, in_process_region_transport)?, negated: exists.negated, }))), Expr::InSubquery(in_subquery) => Ok(Transformed::yes(Expr::InSubquery(InSubquery { expr: in_subquery.expr, - subquery: Self::handle_subquery(in_subquery.subquery)?, + subquery: Self::handle_subquery(in_subquery.subquery, in_process_region_transport)?, negated: in_subquery.negated, }))), Expr::ScalarSubquery(scalar_subquery) => Ok(Transformed::yes(Expr::ScalarSubquery( - Self::handle_subquery(scalar_subquery)?, + Self::handle_subquery(scalar_subquery, in_process_region_transport)?, ))), _ => Ok(Transformed::no(expr)), } } - fn handle_subquery(subquery: Subquery) -> DfResult { - let mut rewriter = PlanRewriter::new(&subquery.subquery); + fn handle_subquery( + subquery: Subquery, + in_process_region_transport: bool, + ) -> DfResult { + let mut rewriter = PlanRewriter::new(&subquery.subquery, in_process_region_transport); let mut rewrote_subquery = subquery .subquery .as_ref() @@ -400,7 +444,12 @@ enum RewriterStatus { struct PlanRewriter { /// Whether the whole plan this rewriter walks can be encoded to Substrait. /// Used by [`PlanRewriter::should_expand`] to skip the per-node encoding check. + /// + /// Used for the ordinary standalone fast path when the root plan is encodable. It is not + /// needed when the root itself is the in-process `FusedSeries` exception. whole_plan_encodable: bool, + /// In-process transport skips the encodability gate only when the rewritten root itself is FusedSeries. + skip_encodability_gate: bool, /// Current level in the tree level: usize, /// Simulated stack for the `rewrite` recursion @@ -456,14 +505,28 @@ struct PlanRewriter { new_child_plan: Option, } +fn is_fused_series_plan(plan: &LogicalPlan) -> bool { + matches!(plan, LogicalPlan::Extension(extension) if extension.node.as_any().is::()) +} + impl PlanRewriter { /// `plan` must be the root of the tree that is about to be rewritten, otherwise /// [`PlanRewriter::should_expand`] may skip a check it has to perform. - fn new(plan: &LogicalPlan) -> Self { + fn new(plan: &LogicalPlan, in_process_region_transport: bool) -> Self { + let native_root = match plan { + LogicalPlan::Analyze(analyze) => analyze.input.as_ref(), + _ => plan, + }; + let skip_encodability_gate = + in_process_region_transport && is_fused_series_plan(native_root); Self { - whole_plan_encodable: DFLogicalSubstraitConvertor - .encode(plan, DefaultSerializer) - .is_ok(), + // Preserve the ordinary whole-plan fast path unless this is the in-process root + // exception: no node in that tree needs an encoding check. + skip_encodability_gate, + whole_plan_encodable: !skip_encodability_gate + && DFLogicalSubstraitConvertor + .encode(plan, DefaultSerializer) + .is_ok(), ..Default::default() } } @@ -495,7 +558,15 @@ impl PlanRewriter { // is the dominant cost on deep PromQL plans. // When the root does not encode, the plan holds at least one node that has to stay // on the frontend and only the per-node check locates it. - if !self.whole_plan_encodable + // + // For the in-process root `FusedSeries` exception, encodability says nothing about where + // a node may run, and asking it would cut the plan at the first node with no serializer + // branch before `check_plan` could decide. Skipping it only removes a reason to expand: every + // check below still runs, and a node they call `NonCommutative`/`Unimplemented`/ + // `Unsupported` still expands, so a node no serializer lists only rides to a region when + // the categorizer places it there. + if !self.skip_encodability_gate + && !self.whole_plan_encodable && let Err(e) = DFLogicalSubstraitConvertor.encode(plan, DefaultSerializer) { debug!( diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 694f2bcde59..b512655dd2c 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -68,6 +68,38 @@ fn collect_merge_scan_remote_dyn_filter_producer_ids( producer_ids.extend(producer_id_list); } +/// The analyzer as an engine whose regions are reached over the wire builds it: every plan a region +/// receives from it has to be encodable. +fn analyzer() -> DistPlannerAnalyzer { + DistPlannerAnalyzer::default() +} + +#[test] +fn ordinary_encodable_plan_keeps_in_process_encodability_gate() { + let table = TestTable::table_with_name(0, "t".to_string()); + let table_source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(table), + ))); + let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) + .unwrap() + .build() + .unwrap(); + + assert!( + DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .is_ok() + ); + + let in_process_rewriter = PlanRewriter::new(&plan, true); + assert!(in_process_rewriter.whole_plan_encodable); + assert!(!in_process_rewriter.skip_encodability_gate); + + let wire_rewriter = PlanRewriter::new(&plan, false); + assert!(wire_rewriter.whole_plan_encodable); + assert!(!wire_rewriter.skip_encodability_gate); +} + struct MergeScanRemoteDynFilterProducerIdCollector<'a> { producer_ids: &'a mut Vec, } @@ -395,24 +427,29 @@ fn frontend_only_histogram_folds_stay_above_merge_scan() { "{name}" ); - let result = DistPlannerAnalyzer {} - .analyze(plan, &ConfigOptions::default()) - .unwrap(); - let result_text = result.to_string(); - assert!( - result_text.contains("HistogramFold:"), - "{name}\n{result_text}" - ); + for analyzer in [ + DistPlannerAnalyzer::default(), + DistPlannerAnalyzer::new(true), + ] { + let result = analyzer + .analyze(plan.clone(), &ConfigOptions::default()) + .unwrap(); + let result_text = result.to_string(); + assert!( + result_text.contains("HistogramFold:"), + "{name}\n{result_text}" + ); - let remote_input = find_merge_scan(&result).unwrap().input(); - let remote_text = remote_input.to_string(); - assert!( - !remote_text.contains("HistogramFold:"), - "{name}\n{remote_text}" - ); - DFLogicalSubstraitConvertor - .encode(remote_input, DefaultSerializer) - .unwrap(); + let remote_input = find_merge_scan(&result).unwrap().input(); + let remote_text = remote_input.to_string(); + assert!( + !remote_text.contains("HistogramFold:"), + "{name}\n{remote_text}" + ); + DFLogicalSubstraitConvertor + .encode(remote_input, DefaultSerializer) + .unwrap(); + } } } @@ -442,19 +479,24 @@ fn unencodable_unnest_stays_above_merge_scan() { .is_err() ); - let result = DistPlannerAnalyzer {} - .analyze(plan, &ConfigOptions::default()) - .unwrap(); - let result_text = result.to_string(); - assert!(result_text.contains("Unnest:"), "{result_text}"); + for analyzer in [ + DistPlannerAnalyzer::default(), + DistPlannerAnalyzer::new(true), + ] { + let result = analyzer + .analyze(plan.clone(), &ConfigOptions::default()) + .unwrap(); + let result_text = result.to_string(); + assert!(result_text.contains("Unnest:"), "{result_text}"); - let remote_input = find_merge_scan(&result).unwrap().input(); - let remote_text = remote_input.to_string(); - assert!(!remote_text.contains("Unnest:"), "{remote_text}"); - assert!(remote_text.contains("make_array"), "{remote_text}"); - DFLogicalSubstraitConvertor - .encode(remote_input, DefaultSerializer) - .unwrap(); + let remote_input = find_merge_scan(&result).unwrap().input(); + let remote_text = remote_input.to_string(); + assert!(!remote_text.contains("Unnest:"), "{remote_text}"); + assert!(remote_text.contains("make_array"), "{remote_text}"); + DFLogicalSubstraitConvertor + .encode(remote_input, DefaultSerializer) + .unwrap(); + } } #[test] @@ -481,9 +523,7 @@ fn dictionary_literals_are_unwrapped_before_pushdown() { .build() .unwrap(); - let result = DistPlannerAnalyzer {} - .analyze(plan, &ConfigOptions::default()) - .unwrap(); + let result = analyzer().analyze(plan, &ConfigOptions::default()).unwrap(); assert!( result .to_string() @@ -578,7 +618,7 @@ fn expand_proj_sort_proj() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.number", @@ -625,7 +665,7 @@ fn expand_proj_sort_partial_proj() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.number", @@ -661,7 +701,7 @@ fn expand_sort_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1, t.pk2, t.pk3, t.ts, t.number", @@ -699,7 +739,7 @@ fn expand_sort_alias_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: something", @@ -738,7 +778,7 @@ fn expand_sort_alias_conflict_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan.clone(), &config); + let result = analyzer().analyze(plan.clone(), &config); assert!(result.is_err(), "Expected error for ambiguous alias"); assert!(format!("{result:?}").contains("AmbiguousReference")); @@ -746,7 +786,7 @@ fn expand_sort_alias_conflict_limit() { config.extensions.insert(DistPlannerOptions { allow_query_fallback: true, }); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=10", @@ -780,7 +820,7 @@ fn expand_sort_alias_conflict_but_not_really_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1", @@ -818,7 +858,7 @@ fn expand_limit_sort() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Sort: t.pk1 ASC NULLS LAST", @@ -853,7 +893,7 @@ fn expand_sort_limit_sort() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Sort: t.pk1 ASC NULLS LAST", @@ -898,7 +938,7 @@ fn expand_proj_step_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: min(t.number)", @@ -943,7 +983,7 @@ fn expand_proj_alias_fake_part_col_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: pk1, pk2, min(t.number)", @@ -987,7 +1027,7 @@ fn expand_proj_alias_aliased_part_col_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: pk42, pk43, min(t.number)", @@ -1022,7 +1062,7 @@ fn expand_expr_over_part_col_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: substr(t.pk1,Int64(1),Int64(1)), t.pk2, min(t.number)", @@ -1057,7 +1097,7 @@ fn expand_part_col_aggr_step_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: min(max(t.number))", @@ -1090,7 +1130,7 @@ fn expand_step_aggr_step_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Aggregate: groupBy=[[]], aggr=[[min(max(t.number))]]", @@ -1135,7 +1175,7 @@ fn expand_part_col_aggr_part_col_aggr() { assert_eq!(expected_original, plan.to_string()); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1, t.pk2, min(max(t.number))", @@ -1167,7 +1207,7 @@ fn expand_step_aggr_proj() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: min(t.number)", @@ -1225,7 +1265,7 @@ fn expand_complex_col_req_sort_pql() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1, t.pk2", @@ -1271,7 +1311,7 @@ fn expand_proj_sort_step_aggr_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=10", " Aggregate: groupBy=[[]], aggr=[[min(t.number)]]", @@ -1312,7 +1352,7 @@ fn expand_proj_sort_limit_step_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Aggregate: groupBy=[[]], aggr=[[min(t.number)]]", " Projection: t.number", @@ -1354,7 +1394,7 @@ fn expand_proj_limit_step_aggr_sort() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Sort: min(t.number) ASC NULLS LAST", " Aggregate: groupBy=[[]], aggr=[[min(t.number)]]", @@ -1396,7 +1436,7 @@ fn expand_proj_sort_part_col_aggr_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=10", " Aggregate: groupBy=[[t.pk1, t.pk2]], aggr=[[min(t.number)]]", @@ -1438,7 +1478,7 @@ fn expand_proj_sort_limit_part_col_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Aggregate: groupBy=[[t.pk1, t.pk2]], aggr=[[min(t.number)]]", " Projection: t.number, t.pk1, t.pk2", @@ -1480,7 +1520,7 @@ fn expand_proj_part_col_aggr_limit_sort() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Sort: t.pk2 ASC NULLS LAST", " Projection: t.pk1, t.pk2, min(t.number)", @@ -1522,7 +1562,7 @@ fn expand_proj_part_col_aggr_sort_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1, t.pk2, min(t.number)", " Limit: skip=0, fetch=10", @@ -1557,9 +1597,7 @@ fn test_simplify_select_now_expression() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); common_telemetry::info!("Analyzed plan: {}", result); @@ -1615,7 +1653,7 @@ fn sibling_merge_scans_have_unique_remote_dyn_filter_producer_ids() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let mut producer_ids = Vec::new(); collect_merge_scan_remote_dyn_filter_producer_id_list(&result, &mut producer_ids); @@ -1670,7 +1708,7 @@ fn pre_merge_scan_optimizer_eliminates_projected_false_filter() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); assert_eq!("EmptyRelation: rows=0", result.to_string()); } @@ -1718,7 +1756,7 @@ fn test_simplify_now_expression() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let plan_str = result.to_string(); common_telemetry::info!("Analyzed plan: {}", plan_str); @@ -1757,7 +1795,7 @@ fn expand_proj_limit_part_col_aggr_sort() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); // Pre-MergeScan optimizer intentionally excludes PushDownLimit, so the // remote plan shows an explicit Limit node instead of `fetch=10` on // TableScan. @@ -1802,7 +1840,7 @@ fn expand_proj_limit_sort_part_col_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); // Pre-MergeScan optimizer intentionally excludes PushDownLimit, so the // remote plan shows an explicit Limit node instead of `fetch=10` on // TableScan. @@ -1840,7 +1878,7 @@ fn expand_step_aggr_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=10", " Projection: t.pk1, min(t.number)", @@ -1873,7 +1911,7 @@ fn expand_step_aggr_avg_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=10", " Projection: t.pk1, avg(t.number)", @@ -1906,7 +1944,7 @@ fn expand_part_col_aggr_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1, t.pk2, min(t.number)", " Limit: skip=0, fetch=10", @@ -1940,7 +1978,7 @@ fn transform_simple_projection_filter() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Distinct:", " MergeScan [is_placeholder=false]", @@ -1968,7 +2006,7 @@ fn transform_aggregator() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = "Projection: avg(t.number)\ \n MergeScan [is_placeholder=false, remote_input=[\ \nAggregate: groupBy=[[]], aggr=[[avg(t.number)]]\ @@ -1994,7 +2032,7 @@ fn transform_distinct_order() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.number", " MergeScan [is_placeholder=false, remote_input=[ @@ -2022,7 +2060,7 @@ fn transform_single_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = "Projection: t.number\ \n MergeScan [is_placeholder=false, remote_input=[ Limit: skip=0, fetch=1 @@ -2063,7 +2101,7 @@ fn transform_unalighed_join_with_alias() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=1", " LeftSemi Join: t.number = right.number", @@ -2099,7 +2137,7 @@ fn transform_subquery_sort_alias() { .build() .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: a.pk1, a.pk2, a.pk3, a.ts, a.number", " MergeSort: a.number ASC NULLS LAST", @@ -2130,7 +2168,7 @@ fn transform_sort_subquery_alias() { .build() .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: a.pk1, a.pk2, a.pk3, a.ts, a.number", " MergeSort: a.number ASC NULLS LAST", @@ -2169,7 +2207,7 @@ fn remote_dyn_filter_producer_ids_do_not_collide_between_subquery_and_outer_plan .unwrap() .build() .unwrap(); - let rewritten = DistPlannerAnalyzer {}.try_push_down(outer_plan).unwrap(); + let rewritten = analyzer().try_push_down(outer_plan).unwrap(); let mut producer_ids = BTreeSet::new(); collect_merge_scan_remote_dyn_filter_producer_ids(&rewritten, &mut producer_ids); @@ -2201,7 +2239,7 @@ fn date_bin_ts_group_by() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ r#"Projection: date_bin(IntervalDayTime("IntervalDayTime { days: 0, milliseconds: 60000 }"),t.ts), min(t.number)"#, @@ -2249,9 +2287,7 @@ fn test_last_value_order_by() { try_encode_decode_substrait(&plan, ctx.state()); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let expected = [ "Projection: last_value(t.ts) ORDER BY [t.ts ASC NULLS FIRST]", @@ -2312,9 +2348,7 @@ fn test_last_value_no_order_by() { try_encode_decode_substrait(&split.lower_state, ctx.state()); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let expected = [ "Projection: last_value(t.ts)", @@ -2344,9 +2378,7 @@ fn test_table_scan_projection() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let expected = [ "Projection: t.ts", " MergeScan [is_placeholder=false, remote_input=[", @@ -2380,9 +2412,7 @@ fn test_table_scan_cast_projection_pushdown() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let expected = [ "Projection: ts", " MergeScan [is_placeholder=false, remote_input=[", @@ -2426,9 +2456,7 @@ fn test_cast_filter_simplified_after_pushdown() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let expected = [ "Projection: ts", " MergeScan [is_placeholder=false, remote_input=[", @@ -2483,7 +2511,7 @@ fn test_join_side_local_filter_pushdown_into_merge_scan() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); assert_remote_table_scan_filters_are_safe(&result); let predicate = col("t1.pk1").eq(lit("v")); @@ -2555,7 +2583,7 @@ fn test_left_join_left_side_filter_pushdown_into_merge_scan() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); assert_remote_table_scan_filters_are_safe(&result); let plan_str = result.to_string(); @@ -2611,7 +2639,7 @@ fn test_join_cross_table_predicate_not_pushed_to_single_side() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); assert_remote_table_scan_filters_are_safe(&result); let plan_str = result.to_string(); @@ -2668,9 +2696,7 @@ fn scheduled_now_yields_stable_literal_in_remote_plan() { scheduled_time: Some(scheduled_time), }); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let result_str = result.to_string(); common_telemetry::info!("Analyzed plan with scheduled time: {}", result_str); @@ -2706,9 +2732,7 @@ fn scheduled_now_yields_stable_literal_in_remote_plan() { // Absent the extension (default config), the same plan simplifies to a // wall-clock literal — the remote_input contains a variable timestamp. let default_config = ConfigOptions::default(); - let wall_result = DistPlannerAnalyzer {} - .analyze(plan, &default_config) - .unwrap(); + let wall_result = analyzer().analyze(plan, &default_config).unwrap(); let wall_str = wall_result.to_string(); // The wall-clock result must NOT contain the scheduled literal. assert!( @@ -2769,7 +2793,7 @@ fn scheduled_current_timestamp_filter_folds_before_remote_pushdown() { scheduled_time: Some(scheduled_time), }); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); assert_remote_table_scan_filters_are_safe(&result); let result_str = result.to_string(); @@ -2816,7 +2840,7 @@ fn scheduled_none_falls_back_to_wall_clock() { scheduled_time: None, }); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let result_str = result.to_string(); // Must still simplify now() — remotely a literal, top is now(). @@ -2884,7 +2908,7 @@ fn expand_nested_scalar_subquery() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let result_str = result.to_string(); // The main scan and both subquery levels must each be wrapped in a MergeScan. diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index d02b746a8bf..4221b009390 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -21,8 +21,9 @@ use datafusion::error::Result as DfResult; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode}; use promql::extension_plan::{ - EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize, + EmptyMetric, FusedSeries, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize, }; +use promql::native::plan::reads_across_series; use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use crate::dist_plan::MergeScanLogicalPlan; @@ -240,6 +241,38 @@ impl Categorizer { Commutativity::Commutative } + name if name == FusedSeries::name() => { + let fused = plan.as_any().downcast_ref::().unwrap(); + // A region builds this node from the region scan itself: the conversion takes that + // scan's raw-series entrypoint and cannot look through anything in between, so a + // node whose input is not that scan has nothing to read where it would be placed + // and stays where it was planned. + if !matches!(fused.inputs().as_slice(), [LogicalPlan::TableScan(_)]) { + return Commutativity::NonCommutative; + } + // The region asks the scan for the series of one named metric, so a node that binds + // no name cannot be built there at all. + if fused.binding().metric_name.is_none() { + return Commutativity::NonCommutative; + } + // A region evaluates the tree over whatever series its own scan returns, so this + // is only answerable when every node of the tree reads one series at a time. An + // aggregate or a binary operation would fold or pair across series, and one region + // holds only part of them. + if reads_across_series(fused.tree().root()) { + return Commutativity::NonCommutative; + } + + // A region evaluates the whole tree over the series its scan returns, so this is + // only correct while each series reaches exactly one region: a series split + // across regions would be ranged per fragment, and a range over fragments is not + // the range over the series. Native admission + // (`promql::planner::native::try_plan_native_fused`) only builds the node when + // every partition column of the table is one of the binding's tag columns, so + // the partition key is a function of the label set and a series cannot straddle + // regions. + Commutativity::Commutative + } name if name == SeriesNormalize::name() || name == InstantManipulate::name() || name == RangeManipulate::name() => @@ -349,16 +382,39 @@ impl Categorizer { #[cfg(test)] mod tests { use std::collections::{BTreeMap, BTreeSet}; + use std::time::{Duration, UNIX_EPOCH}; + use datafusion::arrow::datatypes::{ + DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit, + }; + use datafusion::common::{DFSchema, DFSchemaRef}; + use datafusion::datasource::DefaultTableSource; use datafusion_common::Column; + use datafusion_common::config::ConfigOptions; use datafusion_expr::LogicalPlanBuilder; use datafusion_expr::expr::ScalarFunction; + use datafusion_expr::{Extension, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion_functions::core::coalesce; + use datafusion_optimizer::analyzer::AnalyzerRule; + use datatypes::series::SeriesOrdering; + use datatypes::series::samples::{FieldSpec, SeriesSchema}; + use promql::extension_plan::{FusedSeries, FusedSeriesBinding}; use promql::functions::{ NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile, }; + use promql::native::plan::{ + AggregateKind, AggregateStage, AggregateStateSchema, BoundAggregate, BoundInput, Grouping, + NativeExpr, NativeNode, NativeType, + }; + use promql_parser::parser; + use store_api::storage::RegionId; + use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use super::*; + use crate::dist_plan::DistPlannerAnalyzer; + use crate::optimizer::test_util::mock_table_provider; + use crate::promql::native_compile::{compile, statement_grid}; + use crate::query_engine::DefaultSerializer; #[test] fn series_divide_by_tsid_is_commutative() { @@ -393,6 +449,388 @@ mod tests { )); } } + + /// The statement one `FusedSeries` test node is bound to. + fn native_statement() -> promql_parser::parser::EvalStmt { + promql_parser::parser::EvalStmt { + expr: parser::parse("some_metric").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + } + } + + /// A `FusedSeries` node over `input`, bound to the tree a bare selector compiles to. + /// + /// The tree is the simplest shape the compiler accepts, because what these tests change is + /// where the node may be placed, not what it computes. + fn fused_series_plan(input: LogicalPlan) -> LogicalPlan { + fused_series_plan_named(input, Some("some_metric")) + } + + /// The same node, with the metric name it binds chosen by the caller. + fn fused_series_plan_named(input: LogicalPlan, metric_name: Option<&str>) -> LogicalPlan { + let stmt = native_statement(); + let series_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }); + let tree = compile( + &stmt, + Arc::clone(&series_schema), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let binding = FusedSeriesBinding { + series_schema, + grid: statement_grid(&stmt).unwrap(), + metric_name: metric_name.map(Arc::from), + time_index: "ts".to_string(), + tags: Vec::new(), + }; + let output_schema: DFSchemaRef = Arc::new( + DFSchema::new_with_metadata( + vec![ + ( + None, + Arc::new(Field::new("v0", ArrowDataType::Float64, true)), + ), + ( + None, + Arc::new(Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )), + ), + ], + Default::default(), + ) + .unwrap(), + ); + + LogicalPlan::Extension(Extension { + node: Arc::new( + FusedSeries::try_new( + Arc::from("some_metric"), + tree, + binding, + input, + output_schema, + ) + .unwrap(), + ), + }) + } + + /// The region's own view of the metric's table: a scan whose provider is the region provider. + fn table_scan_plan() -> LogicalPlan { + let provider = mock_table_provider(RegionId::new(1, 1)); + LogicalPlanBuilder::scan( + "some_metric", + Arc::new(DefaultTableSource::new(Arc::new(provider))), + None, + ) + .unwrap() + .build() + .unwrap() + } + + fn holds_fused_series(plan: &LogicalPlan) -> bool { + if let LogicalPlan::Extension(extension) = plan + && extension.node.as_any().is::() + { + return true; + } + plan.inputs().into_iter().any(holds_fused_series) + } + + /// An extension node with no `ExtensionSerializer` branch: neither one under its own name nor + /// one it shares, since the serializer lists names and not shapes. It stands in for the nodes + /// that cannot reach a region as bytes for their own reasons, which the encoding gate must keep + /// honouring. The categorizer refuses it, as it refuses every name it does not know. + #[derive(Debug, PartialEq, Eq, Hash)] + struct UnserializableNode { + input: LogicalPlan, + schema: DFSchemaRef, + } + + impl PartialOrd for UnserializableNode { + fn partial_cmp(&self, other: &Self) -> Option { + let ordering = self.input.partial_cmp(&other.input); + match ordering { + Some(std::cmp::Ordering::Equal) if self.schema != other.schema => None, + ordering => ordering, + } + } + } + + impl UnserializableNode { + fn new(input: LogicalPlan) -> Self { + let schema = input.schema().clone(); + Self { input, schema } + } + } + + impl UserDefinedLogicalNodeCore for UnserializableNode { + fn name(&self) -> &str { + "UnserializableNode" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "UnserializableNode") + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> DfResult { + Ok(Self::new(inputs.swap_remove(0))) + } + } + + /// The plan that node builds over `input`. + fn unserializable_plan(input: LogicalPlan) -> LogicalPlan { + LogicalPlan::Extension(Extension { + node: Arc::new(UnserializableNode::new(input)), + }) + } + + fn holds_unserializable_node(plan: &LogicalPlan) -> bool { + if let LogicalPlan::Extension(extension) = plan + && extension.node.as_any().is::() + { + return true; + } + plan.inputs().into_iter().any(holds_unserializable_node) + } + + fn holds_table_scan(plan: &LogicalPlan) -> bool { + if matches!(plan, LogicalPlan::TableScan(_)) { + return true; + } + plan.inputs().into_iter().any(holds_table_scan) + } + + /// The merge scan `plan` is, if it is one. + fn as_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { + if let LogicalPlan::Extension(extension) = plan { + extension + .node + .as_any() + .downcast_ref::() + } else { + None + } + } + + /// The first merge scan in `plan`'s visible tree. A merge scan's remote input is not one of its + /// children, so this does not look into one. + fn first_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { + if let Some(merge_scan) = as_merge_scan(plan) { + return Some(merge_scan); + } + plan.inputs().into_iter().find_map(first_merge_scan) + } + + fn analyze(analyzer: DistPlannerAnalyzer, plan: LogicalPlan) -> LogicalPlan { + analyzer.analyze(plan, &ConfigOptions::default()).unwrap() + } + + /// Whether the remote input of any merge scan in `plan` holds a node `wanted` accepts. + fn any_remote_input_holds(plan: &LogicalPlan, wanted: fn(&LogicalPlan) -> bool) -> bool { + let mut holds = false; + plan.apply(|node| { + if let Some(merge_scan) = as_merge_scan(node) { + holds |= wanted(merge_scan.input()); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + holds + } + + #[test] + fn fused_series_over_its_scan_is_commutative() { + let plan = fused_series_plan(table_scan_plan()); + + assert!(matches!( + Categorizer::check_plan(&plan, None).unwrap(), + Commutativity::Commutative + )); + } + + /// The region builds the node from the scan's own raw-series entrypoint, so an input that is + /// not that scan leaves the node with nothing to read where it would be placed. + #[test] + fn fused_series_over_anything_but_a_scan_is_not_commutative() { + let scan = table_scan_plan(); + let input = LogicalPlanBuilder::from(scan) + .filter(datafusion_expr::col("k0").eq(datafusion_expr::lit("a"))) + .unwrap() + .build() + .unwrap(); + let plan = fused_series_plan(input); + + assert!(matches!( + Categorizer::check_plan(&plan, None).unwrap(), + Commutativity::NonCommutative + )); + } + + /// The region asks the scan for one named metric's series, so a node binding no name cannot + /// be built there. + #[test] + fn fused_series_that_binds_no_metric_name_is_not_commutative() { + let plan = fused_series_plan_named(table_scan_plan(), None); + + assert!(matches!( + Categorizer::check_plan(&plan, None).unwrap(), + Commutativity::NonCommutative + )); + } + + /// A tree built by something other than the compiler can hold a node that folds several + /// series. The check is a property of the tree, so it is exercised on the node itself: the + /// compiler emits no such node today, and nothing else would build a tree containing one. + #[test] + fn a_tree_that_folds_across_series_is_not_per_series() { + let raw = NativeType::Samples(Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + })); + let input = NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }; + let group = Grouping::By(Arc::from(["k0".into()])); + let aggregate = NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(input), + op: BoundAggregate { + kind: AggregateKind::Sum, + group: group.clone(), + stage: AggregateStage::Partial, + }, + }, + output: NativeType::AggregateState(Arc::new(AggregateStateSchema { + group, + accumulators: Arc::from([]), + grid: statement_grid(&native_statement()).unwrap(), + })), + }; + assert!(reads_across_series(&aggregate)); + + // The same walk over what the compiler emits stays per-series. + let single_selector = compile( + &native_statement(), + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + assert!(!reads_across_series(single_selector.root())); + } + + /// The encoding gate is what kept the node on the frontend, and it only exists for a plan that + /// has to be encoded. + /// + /// `PlanRewriter::should_expand` asked two questions about a node before it let the descent + /// continue, and the Substrait one came first: unless the *whole* root plan encoded, every + /// subtree that failed to encode was expanded on, whichever way `check_plan` answered. A + /// `FusedSeries` node has no serializer branch — `ExtensionSerializer` answers `NotImplemented` + /// for every name it does not list — so over the wire the boundary is still drawn at the node + /// and the region receives the bare scan. + /// + /// A plan handed to `RegionServer::handle_read` as a `LogicalPlan` is never encoded, so there + /// the gate has nothing to say and the commutativity arm decides: the node rides inside the + /// merge scan. + #[test] + fn fused_series_rides_to_the_region_only_when_regions_are_reached_in_process() { + let plan = fused_series_plan(table_scan_plan()); + + assert!( + DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .is_err(), + "the fused node encodes now, so the encoding gate decides nothing" + ); + assert!( + DFLogicalSubstraitConvertor + .encode(&table_scan_plan(), DefaultSerializer) + .is_ok(), + "the scan under the fused node encodes, so the fused node is the only reason this plan does not" + ); + + // Over the wire the plan has to encode, so the boundary is drawn at the node. + let encoded = analyze(DistPlannerAnalyzer::new(false), plan.clone()); + assert!( + !any_remote_input_holds(&encoded, holds_fused_series), + "the analyzer placed the node although its plan has to encode: {encoded}" + ); + + // In process nothing is encoded, so the categorizer places the node. + let in_process = analyze(DistPlannerAnalyzer::new(true), plan); + assert!( + any_remote_input_holds(&in_process, holds_fused_series), + "the analyzer kept the node on the frontend although regions are reached in process: {in_process}" + ); + } + + /// Skipping the encoding gate removes one reason to expand and adds none: a node that does not + /// encode for its own reasons, and that the categorizer refuses, still draws the boundary at + /// itself when regions are reached in process. + /// + /// Without this, the skip would amount to placing every unencodable node region-side. + #[test] + fn a_plan_that_does_not_encode_for_another_reason_still_expands_at_that_node() { + let plan = unserializable_plan(table_scan_plan()); + assert!( + DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .is_err(), + "the stand-in node has a serializer branch now, so it no longer proves anything" + ); + + let analyzed = analyze(DistPlannerAnalyzer::new(true), plan); + + assert!( + !any_remote_input_holds(&analyzed, holds_unserializable_node), + "the analyzer placed a node the categorizer refuses: {analyzed}" + ); + let inputs = analyzed.inputs(); + let [input] = inputs.as_slice() else { + panic!("expected the node to keep one input in: {analyzed}"); + }; + // `expand` keeps a projection over the merge scan to restore the schema of what it wrapped. + let merge_scan = first_merge_scan(input) + .unwrap_or_else(|| panic!("expected the boundary right below the node in: {analyzed}")); + assert!( + holds_table_scan(merge_scan.input()), + "expected the scan under the node to be the region input in: {analyzed}" + ); + } } pub type Transformer = Arc Option>; diff --git a/src/query/src/optimizer/parallelize_scan.rs b/src/query/src/optimizer/parallelize_scan.rs index 8e0113e69df..a0b0545d4e0 100644 --- a/src/query/src/optimizer/parallelize_scan.rs +++ b/src/query/src/optimizer/parallelize_scan.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::BinaryHeap; use std::sync::Arc; use common_telemetry::debug; @@ -22,9 +21,10 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{DataFusionError, Result}; -use store_api::region_engine::PartitionRange; use table::table::scan::RegionScanExec; +pub use table::table::scan::assign_partition_ranges; + #[derive(Debug)] pub struct ParallelizeScan; @@ -71,7 +71,7 @@ impl ParallelizeScan { // assign ranges to each partition let mut partition_ranges = - Self::assign_partition_range(ranges, expected_partition_num); + assign_partition_ranges(ranges, expected_partition_num); debug!( "Assign {total_range_num} ranges to {expected_partition_num} partitions" ); @@ -116,176 +116,15 @@ impl ParallelizeScan { Ok(result) } - - /// Distribute [`PartitionRange`]s to each partition. - /// - /// Currently we assign ranges to partitions according to their rows so each partition - /// has similar number of rows. This method always return `expected_partition_num` partitions. - pub fn assign_partition_range( - mut ranges: Vec, - expected_partition_num: usize, - ) -> Vec> { - if ranges.is_empty() { - // Returns a single partition with no range. - return vec![vec![]; expected_partition_num]; - } - - if ranges.len() == 1 { - let mut vec = vec![vec![]; expected_partition_num]; - vec[0] = ranges; - return vec; - } - - // Sort ranges by number of rows in descending order. - ranges.sort_by_key(|b| std::cmp::Reverse(b.num_rows)); - let mut partition_ranges = vec![vec![]; expected_partition_num]; - - #[derive(Eq, PartialEq)] - struct HeapNode { - num_rows: usize, - partition_idx: usize, - } - - impl Ord for HeapNode { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - // Reverse for min-heap. - self.num_rows.cmp(&other.num_rows).reverse() - } - } - - impl PartialOrd for HeapNode { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } - } - - let mut part_heap = - BinaryHeap::from_iter((0..expected_partition_num).map(|partition_idx| HeapNode { - num_rows: 0, - partition_idx, - })); - - // Assigns the range to the partition with the smallest number of rows. - for range in ranges { - let mut node = part_heap.pop().unwrap(); - let partition_idx = node.partition_idx; - node.num_rows += range.num_rows; - partition_ranges[partition_idx].push(range); - part_heap.push(node); - } - - partition_ranges - } } #[cfg(test)] mod test { use common_time::Timestamp; use common_time::timestamp::TimeUnit; + use store_api::region_engine::PartitionRange; - use super::*; - - #[test] - fn test_assign_partition_range() { - let ranges = vec![ - PartitionRange { - start: Timestamp::new(0, TimeUnit::Second), - end: Timestamp::new(10, TimeUnit::Second), - num_rows: 100, - identifier: 1, - }, - PartitionRange { - start: Timestamp::new(10, TimeUnit::Second), - end: Timestamp::new(20, TimeUnit::Second), - num_rows: 200, - identifier: 2, - }, - PartitionRange { - start: Timestamp::new(20, TimeUnit::Second), - end: Timestamp::new(30, TimeUnit::Second), - num_rows: 150, - identifier: 3, - }, - PartitionRange { - start: Timestamp::new(30, TimeUnit::Second), - end: Timestamp::new(40, TimeUnit::Second), - num_rows: 250, - identifier: 4, - }, - ]; - - // assign to 2 partitions - let expected_partition_num = 2; - let result = - ParallelizeScan::assign_partition_range(ranges.clone(), expected_partition_num); - let expected = vec![ - vec![ - PartitionRange { - start: Timestamp::new(30, TimeUnit::Second), - end: Timestamp::new(40, TimeUnit::Second), - num_rows: 250, - identifier: 4, - }, - PartitionRange { - start: Timestamp::new(0, TimeUnit::Second), - end: Timestamp::new(10, TimeUnit::Second), - num_rows: 100, - identifier: 1, - }, - ], - vec![ - PartitionRange { - start: Timestamp::new(10, TimeUnit::Second), - end: Timestamp::new(20, TimeUnit::Second), - num_rows: 200, - identifier: 2, - }, - PartitionRange { - start: Timestamp::new(20, TimeUnit::Second), - end: Timestamp::new(30, TimeUnit::Second), - num_rows: 150, - identifier: 3, - }, - ], - ]; - assert_eq!(result, expected); - - // assign 4 ranges to 5 partitions. - let expected_partition_num = 5; - let result = ParallelizeScan::assign_partition_range(ranges, expected_partition_num); - let expected = vec![ - vec![PartitionRange { - start: Timestamp::new(30, TimeUnit::Second), - end: Timestamp::new(40, TimeUnit::Second), - num_rows: 250, - identifier: 4, - }], - vec![PartitionRange { - start: Timestamp::new(0, TimeUnit::Second), - end: Timestamp::new(10, TimeUnit::Second), - num_rows: 100, - identifier: 1, - }], - vec![PartitionRange { - start: Timestamp::new(10, TimeUnit::Second), - end: Timestamp::new(20, TimeUnit::Second), - num_rows: 200, - identifier: 2, - }], - vec![], - vec![PartitionRange { - start: Timestamp::new(20, TimeUnit::Second), - end: Timestamp::new(30, TimeUnit::Second), - num_rows: 150, - identifier: 3, - }], - ]; - assert_eq!(result, expected); - - // assign 0 ranges to 5 partitions. Should return 5 empty ranges. - let result = ParallelizeScan::assign_partition_range(vec![], 5); - assert_eq!(result.len(), 5); - } + use super::assign_partition_ranges; #[test] fn test_assign_unbalance_partition_range() { @@ -318,8 +157,7 @@ mod test { // assign to 2 partitions let expected_partition_num = 2; - let result = - ParallelizeScan::assign_partition_range(ranges.clone(), expected_partition_num); + let result = assign_partition_ranges(ranges.clone(), expected_partition_num); let expected = vec![ vec![PartitionRange { start: Timestamp::new(30, TimeUnit::Second), diff --git a/src/query/src/optimizer/scan_hint.rs b/src/query/src/optimizer/scan_hint.rs index 4685889c998..0a1334dd9c0 100644 --- a/src/query/src/optimizer/scan_hint.rs +++ b/src/query/src/optimizer/scan_hint.rs @@ -26,7 +26,7 @@ use datafusion_expr::expr::Sort; use datafusion_expr::{Expr, LogicalPlan, utils}; use datafusion_optimizer::{OptimizerConfig, OptimizerRule}; use datatypes::arrow::datatypes::{DataType, TimeUnit as ArrowTimeUnit}; -use promql::extension_plan::{InstantManipulate, SeriesDivide, SeriesNormalize}; +use promql::extension_plan::{FusedSeries, InstantManipulate, SeriesDivide, SeriesNormalize}; use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use store_api::storage::{TimeSeriesDistribution, TimeSeriesRowSelector}; @@ -37,6 +37,9 @@ use crate::dummy_catalog::DummyTableProvider; /// - the nearest order requirement to the leaf table scan node as ordering hint. /// - the group by columns when all aggregate functions are `last_value` as /// time series row selector hint. +/// - the per-series distribution when the leaf table scan node is the input of a +/// [`FusedSeries`] node, which reads the raw series frames a region only packs +/// for that distribution. /// /// [`ScanRequest`]: store_api::storage::ScanRequest #[derive(Debug)] @@ -68,6 +71,7 @@ impl ScanHintRule { fn set_hints( plan: LogicalPlan, rewriter: &mut ScanHintRewriter, + fused_series_input: bool, ) -> Result> { let LogicalPlan::TableScan(mut table_scan) = plan else { return Ok(Transformed::no(plan)); @@ -89,8 +93,10 @@ impl ScanHintRule { }; let use_last_row = rewriter.inside_single_evaluation && filters_preserve_last_row; - let has_hint = - rewriter.order_expr.is_some() || rewriter.ts_row_selector.is_some() || use_last_row; + let has_hint = rewriter.order_expr.is_some() + || rewriter.ts_row_selector.is_some() + || use_last_row + || fused_series_input; if !has_hint { return Ok(Transformed::no(LogicalPlan::TableScan(table_scan))); } @@ -101,6 +107,14 @@ impl ScanHintRule { // local as well as the new LastRow hint. let adapter = original.clone_for_scan(); Self::apply_hints(&adapter, rewriter); + if fused_series_input { + // The scan below a `FusedSeries` node reads packs of raw series, which a region only + // produces for a per-series request; the existing PromQL pipeline asks for the same + // distribution through a sort on `(series key, time)`, which the native node does not + // have below it. Applying it after the order hints makes the plan shape the deciding + // one: the node's own input is what reads the frames. + adapter.with_distribution(TimeSeriesDistribution::PerSeries); + } if use_last_row { // Apply the instant-derived hint after the aggregate hint. Both // select LastRow today, and this ordering preserves the existing @@ -275,6 +289,12 @@ struct ScanHintRewriter { ts_stack: Vec, Column)>>, inside_single_evaluation: bool, single_evaluation_stack: Vec, + /// Whether the node being visited is the input of a [`FusedSeries`] node. + /// + /// The flag only describes that one edge: it is set while a `FusedSeries` node is visited, so + /// the node visited next — that node's own input — sees it, and a node further down does not. + fused_series_input: bool, + fused_series_input_stack: Vec, } impl TreeNodeRewriter for ScanHintRewriter { @@ -285,6 +305,13 @@ impl TreeNodeRewriter for ScanHintRewriter { self.ts_stack.push(self.ts_row_selector.clone()); self.single_evaluation_stack .push(self.inside_single_evaluation); + self.fused_series_input_stack.push(self.fused_series_input); + // Read the edge this node sits on before recording what this node is, so a `FusedSeries` + // node hands the flag to its own input and to nothing below it. + let fused_series_input = std::mem::replace( + &mut self.fused_series_input, + matches!(&node, LogicalPlan::Extension(extension) if extension.node.as_any().is::()), + ); if let LogicalPlan::Sort(sort) = &node { self.order_expr = Some(sort.expr.clone()); @@ -321,7 +348,7 @@ impl TreeNodeRewriter for ScanHintRewriter { } } - ScanHintRule::set_hints(node, self) + ScanHintRule::set_hints(node, self, fused_series_input) } fn f_up(&mut self, node: LogicalPlan) -> Result> { @@ -334,6 +361,9 @@ impl TreeNodeRewriter for ScanHintRewriter { if let Some(previous) = self.single_evaluation_stack.pop() { self.inside_single_evaluation = previous; } + if let Some(previous) = self.fused_series_input_stack.pop() { + self.fused_series_input = previous; + } Ok(Transformed::no(node)) } } @@ -446,6 +476,7 @@ impl ScanHintRewriter { mod test { use std::collections::HashMap; use std::sync::Arc; + use std::time::{Duration, UNIX_EPOCH}; use datafusion::functions_aggregate::first_last::last_value_udaf; use datafusion::functions_aggregate::min_max::max_udaf; @@ -454,21 +485,26 @@ mod test { use datafusion::logical_expr::{WindowFrame, WindowFunctionDefinition}; use datafusion::prelude::JoinType; use datafusion_common::tree_node::TreeNodeRecursion; + use datafusion_common::{DFSchema, DFSchemaRef}; use datafusion_expr::expr::{ AggregateFunction, AggregateFunctionParams, Cast, WindowFunctionParams, }; use datafusion_expr::expr_fn::scalar_subquery; use datafusion_expr::{Extension, LogicalPlan, LogicalPlanBuilder, col, lit}; use datafusion_optimizer::OptimizerContext; - use datatypes::arrow::datatypes::DataType; + use datatypes::arrow::datatypes::{DataType, Field, TimeUnit}; use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; - use promql::extension_plan::RangeManipulate; + use datatypes::series::SeriesOrdering; + use datatypes::series::samples::{FieldSpec, SeriesSchema}; + use promql::extension_plan::{FusedSeries, FusedSeriesBinding, RangeManipulate}; + use promql_parser::parser; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use store_api::storage::{RegionId, TimeSeriesRowSelector}; use super::*; + use crate::promql::native_compile::{compile, statement_grid}; fn scan_requests(plan: &LogicalPlan) -> Vec { scan_requests_with_names(plan) @@ -1520,4 +1556,123 @@ mod test { Some(TimeSeriesDistribution::PerSeries) ); } + + /// A `FusedSeries` node over `input`, bound to the tree a bare selector compiles to. + /// + /// The rule reads the plan shape, not what the node computes, so the simplest shape the + /// compiler accepts is enough for these tests. + fn fused_series(input: LogicalPlan) -> LogicalPlan { + let stmt = promql_parser::parser::EvalStmt { + expr: parser::parse("some_metric").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let series_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }); + let tree = compile( + &stmt, + Arc::clone(&series_schema), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let binding = FusedSeriesBinding { + series_schema, + grid: statement_grid(&stmt).unwrap(), + metric_name: Some("some_metric".into()), + time_index: "ts".to_string(), + tags: Vec::new(), + }; + let output_schema: DFSchemaRef = Arc::new( + DFSchema::new_with_metadata( + vec![ + (None, Arc::new(Field::new("v0", DataType::Float64, true))), + ( + None, + Arc::new(Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + )), + ), + ], + Default::default(), + ) + .unwrap(), + ); + + LogicalPlan::Extension(Extension { + node: Arc::new( + FusedSeries::try_new( + Arc::from("some_metric"), + tree, + binding, + input, + output_schema, + ) + .unwrap(), + ), + }) + } + + /// The region only packs raw series frames for a per-series request, and the node below is + /// what reads them, so the node's own input asks for that distribution. + #[test] + fn a_scan_under_fused_series_asks_for_the_per_series_distribution() { + let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); + let plan = fused_series(scan_plan(provider.clone(), "t")); + + let rewritten = ScanHintRule + .rewrite(plan, &OptimizerContext::default()) + .unwrap() + .data; + + assert_eq!( + scan_requests(&rewritten)[0].distribution, + Some(TimeSeriesDistribution::PerSeries) + ); + // The hint belongs to this use-site only, never to the shared catalog provider. + assert_eq!(provider.scan_request().distribution, None); + } + + /// The ordinary path: a scan with no fused node above it is left exactly as it was. + #[test] + fn a_scan_not_under_fused_series_keeps_its_request() { + let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); + let plan = scan_plan(provider.clone(), "t"); + + let rewritten = ScanHintRule + .rewrite(plan, &OptimizerContext::default()) + .unwrap() + .data; + + assert_eq!(scan_requests(&rewritten)[0].distribution, None); + assert_eq!(provider.scan_request().distribution, None); + } + + /// Only the node's own input is a raw series reader; a scan further down feeds something else + /// and keeps the request it had. + #[test] + fn a_scan_below_fused_series_input_keeps_its_request() { + let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); + let scan = scan_plan(provider.clone(), "t"); + let input = LogicalPlanBuilder::from(scan) + .project(vec![col("ts"), col("v0")]) + .unwrap() + .build() + .unwrap(); + let plan = fused_series(input); + + let rewritten = ScanHintRule + .rewrite(plan, &OptimizerContext::default()) + .unwrap() + .data; + + assert_eq!(scan_requests(&rewritten)[0].distribution, None); + } } diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index 44cd1653cee..a70e0108817 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -50,6 +50,8 @@ use sql::statements::statement::Statement; use sql::statements::tql::Tql; use sqlparser::ast::{AccessExpr, Value, visit_expressions_mut}; +use crate::datafusion::QUERY_FALLBACK_HINT; +use crate::dist_plan::DistPlannerOptions; use crate::error::{ CteColumnSchemaMismatchSnafu, PlanSqlSnafu, QueryPlanSnafu, Result, SqlSnafu, UnimplementedSnafu, @@ -316,8 +318,27 @@ impl DfLogicalPlanner { } #[tracing::instrument(skip_all)] - async fn plan_pql(&self, stmt: &EvalStmt, query_ctx: QueryContextRef) -> Result { + async fn plan_pql( + &self, + stmt: &EvalStmt, + query_ctx: QueryContextRef, + native_route_allowed: bool, + ) -> Result { let mut scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; + // Whether the distributed planner may fall back to the fallback plan rewriter. The + // execution engine resolves the same three sources the same way: the engine option, which + // it copies into the session state it plans with, and then, per query, the session + // variable or the `query_fallback` hint. Each source only ever turns the fallback on, so + // the value they end up with is the disjunction of them. + let allow_query_fallback = scheduled_state + .config_options() + .extensions + .get::() + .is_some_and(|options| options.allow_query_fallback) + || query_ctx.configuration_parameter().allow_query_fallback() + || query_ctx + .extension(QUERY_FALLBACK_HINT) + .is_some_and(|fallback| fallback.to_lowercase().parse::().unwrap_or(false)); let promql_annotations = query_ctx.remote_query_id().map(promql_annotation_collector); if let Some(collector) = &promql_annotations { scheduled_state @@ -345,6 +366,16 @@ impl DfLogicalPlanner { stmt, &self.engine_state, promql_annotations, + // A query allowed to fall back keeps the existing plan. The native node's physical + // planner needs the bare region scan as its input; the fallback rewriter instead wraps + // every table scan in `MergeScan`, so the node above it would have no physical plan and + // the query would fail after the fallback was already taken. Declining here keeps the + // whole query on the legacy path. + native_route_allowed + && !allow_query_fallback + && self.engine_state.experimental_native_promql() + && self.engine_state.in_process_region_transport() + && crate::promql::planner::native::native_promql_requested(&query_ctx), ) .await .map_err(BoxedError::new) @@ -461,8 +492,16 @@ impl DfLogicalPlanner { alias: eval.alias.clone(), }; let stmt = QueryLanguageParser::parse_promql(&prom_query, &query_ctx)?; - - self.plan(&stmt, query_ctx).await + // A hybrid CTE registers this plan as a table the outer SQL query reads, so the + // `FusedSeries` node would sit under a SQL root. The in-process encodability + // exception is scoped to a plan whose root is the node, so here the node would be + // cut off and left in the frontend, where the physical planner cannot build it - + // a query that used to answer would fail instead. The route is therefore not + // offered for an embedded statement. + let QueryStatement::Promql(eval, _alias) = &stmt else { + unreachable!("a TQL eval parses to a PromQL statement"); + }; + self.plan_pql(eval, query_ctx, false).await } Tql::Explain(_) => UnimplementedSnafu { operation: "TQL EXPLAIN in CTEs", @@ -699,7 +738,8 @@ impl LogicalPlanner for DfLogicalPlanner { async fn plan(&self, stmt: &QueryStatement, query_ctx: QueryContextRef) -> Result { match stmt { QueryStatement::Sql(stmt) => self.plan_sql(stmt, query_ctx).await, - QueryStatement::Promql(stmt, _alias) => self.plan_pql(stmt, query_ctx).await, + // The statement's own plan is what a region receives, so the native route is offered. + QueryStatement::Promql(stmt, _alias) => self.plan_pql(stmt, query_ctx, true).await, } } @@ -760,7 +800,9 @@ mod tests { use table::test_util::EmptyTable; use super::*; + use crate::datafusion::QUERY_FALLBACK_HINT; use crate::parser::{PromQuery, QueryLanguageParser}; + use crate::promql::planner::NATIVE_PROMQL_EXTENSION_KEY; use crate::{QueryEngineFactory, QueryEngineRef}; async fn create_test_engine() -> QueryEngineRef { @@ -819,6 +861,14 @@ mod tests { } fn create_promql_test_engine() -> QueryEngineRef { + create_promql_test_engine_with(crate::options::QueryOptions::default(), false) + } + + /// The PromQL fixture, additionally stating the engine options and how its regions are reached. + fn create_promql_test_engine_with( + query_options: crate::options::QueryOptions, + in_process_region_transport: bool, + ) -> QueryEngineRef { let catalog_manager = MemoryCatalogManager::with_default_setup(); let physical_table_name = "phy"; let physical_table_id = 999u32; @@ -909,15 +959,19 @@ mod tests { }) .unwrap(); - QueryEngineFactory::new( + QueryEngineFactory::try_new_with_plugins_and_region_transport( catalog_manager, None, None, None, None, + None, false, - crate::options::QueryOptions::default(), + in_process_region_transport, + Default::default(), + query_options, ) + .unwrap() .query_engine() } @@ -950,6 +1004,87 @@ mod tests { engine.planner().plan(&stmt, query_ctx).await.unwrap() } + /// Plans `rate(some_metric[5m])` on the fixture above with the fallback sources the caller + /// states — the engine option, the session variable and the `query_fallback` hint — while the + /// native route is enabled and requested, so only the fallback guard decides. + async fn plan_native_promql_with_fallback( + engine_fallback: bool, + session_fallback: bool, + hint: Option<&str>, + ) -> String { + let engine = create_promql_test_engine_with( + crate::options::QueryOptions { + allow_query_fallback: engine_fallback, + experimental_native_promql: true, + ..Default::default() + }, + // The native route plans a region input the region server in this process accepts. + true, + ); + // The native route request and the fallback hint both reach the planner as query context + // extensions, the way a client sends them. + let mut builder = QueryContextBuilder::default() + .set_extension(NATIVE_PROMQL_EXTENSION_KEY.to_string(), "true".to_string()); + if let Some(hint) = hint { + builder = builder.set_extension(QUERY_FALLBACK_HINT.to_string(), hint.to_string()); + } + let query_ctx = builder.build(); + query_ctx + .configuration_parameter() + .set_allow_query_fallback(session_fallback); + let query_ctx = Arc::new(query_ctx); + let stmt = QueryLanguageParser::parse_promql( + &PromQuery { + query: "rate(some_metric[5m])".to_string(), + start: "0".to_string(), + end: "10".to_string(), + step: "5s".to_string(), + lookback: "300s".to_string(), + alias: None, + }, + &query_ctx, + ) + .unwrap(); + + engine + .planner() + .plan(&stmt, query_ctx) + .await + .unwrap() + .display_indent() + .to_string() + } + + /// The native route is admitted only when the query cannot fall back, and every source the + /// execution engine resolves the fallback from decides that at plan time too. + #[tokio::test] + async fn test_plan_pql_declines_native_route_when_query_may_fall_back() { + // (engine option, session variable, hint, native route admitted) + for (engine, session, hint, native) in [ + // Nothing asks for the fallback, so the route is planned as before. + (false, false, None, true), + // A hint that does not parse as `true` leaves the fallback off, as it does at execution. + (false, false, Some("false"), true), + (false, false, Some("yes"), true), + // Either per-query source turns the fallback on, so the route is declined. + (false, true, None, false), + (false, false, Some("true"), false), + (false, false, Some("TRUE"), false), + // Each source only turns the fallback on: neither a `false` hint nor a `false` session + // variable clears the engine option, and a `false` hint does not clear the variable. + (true, false, None, false), + (true, false, Some("false"), false), + (true, true, Some("false"), false), + ] { + let plan = plan_native_promql_with_fallback(engine, session, hint).await; + assert_eq!( + plan.contains("FusedSeries"), + native, + "engine={engine}, session={session}, hint={hint:?}: {plan}" + ); + } + } + /// Plans `sql` and runs the DataFusion analyzer, which is where /// `InsertAssignmentRule` sits. Planning alone stops short of it, so these /// assertions would not see the assignment rewrite at all. diff --git a/src/query/src/promql/native_compile.rs b/src/query/src/promql/native_compile.rs index 373f3ac2294..7e50cc3f8ae 100644 --- a/src/query/src/promql/native_compile.rs +++ b/src/query/src/promql/native_compile.rs @@ -1116,7 +1116,11 @@ fn instant_selection(context: &Context, selector: &VectorSelector) -> Result Result> { +/// +/// Visible to the planner because the binding a [`NativeTree`] is planned against has to name the +/// same grid the tree itself was compiled on, and [`compile`]'s callers therefore read it here +/// rather than deriving the grid a second time. +pub(crate) fn statement_grid(stmt: &EvalStmt) -> Result> { let start_ms = signed_millis_since_epoch(stmt.start)?; let end_ms = signed_millis_since_epoch(stmt.end)?; // The step is narrowed before [`GridSpec`] sees it, so an interval to which no millisecond diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index afe52c7b560..18bb5f2c316 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -16,6 +16,8 @@ mod at_modifier; mod function_plans; mod island; mod matching_filters; +pub(crate) mod native; +pub use native::NATIVE_PROMQL_EXTENSION_KEY; mod set_operator; @@ -292,6 +294,16 @@ pub struct PromPlanner { ctx: PromPlannerContext, /// Optional collector passed to native histogram UDFs. promql_annotations: Option, + /// Whether a statement is first offered to the native fused-series compiler. + /// + /// Read per query from the context extension + /// [`NATIVE_PROMQL_EXTENSION_KEY`](super::planner::native::NATIVE_PROMQL_EXTENSION_KEY), and + /// `false` when the query says nothing: with it off, planning takes the existing path unchanged + /// and nothing below is reached. + /// + /// The engine-wide `QueryOptions::experimental_native_promql` is applied while `QueryEngineState` + /// is built; this planner reads the per-query request extension instead. + native_fused: bool, } type BinaryFieldPair<'a> = (&'a String, &'a String); @@ -302,22 +314,38 @@ impl PromPlanner { stmt: &EvalStmt, query_engine_state: &QueryEngineState, ) -> Result { - Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None).await + Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None, false) + .await } /// Plans a PromQL statement and passes the optional collector to histogram UDFs. + /// + /// `native_fused` enables the native fused-series route: the statement is offered to + /// [`compile`](crate::promql::native_compile::compile) first, and whatever that compiler does + /// not accept is planned by the existing path, unchanged. See + /// [`Self::try_plan_native_fused`]. pub async fn stmt_to_plan_with_annotations( table_provider: DfTableSourceProvider, stmt: &EvalStmt, query_engine_state: &QueryEngineState, promql_annotations: Option, + native_fused: bool, ) -> Result { let mut planner = Self { table_provider, ctx: PromPlannerContext::from_eval_stmt(stmt)?, promql_annotations, + native_fused, }; + // The whole statement is offered to the native compiler first, because that compiler folds + // a range function and the math functions over it into one tree: offered any later, the + // expression has already been split into the nodes those functions would have shared. A + // shape it does not accept is planned by the existing path exactly as if it were not here. + if let Some(plan) = planner.try_plan_native_fused(&stmt.expr).await? { + return planner.strip_tsid_column(plan); + } + let plan = planner .prom_root_expr_to_plan(&stmt.expr, query_engine_state) .await?; @@ -2394,9 +2422,20 @@ impl PromPlanner { // let promql_parser::parser::ast::Extension { expr } = ext_expr; let expr = &ext_expr.expr; let children = expr.children(); - let plan = self - .prom_root_expr_to_plan(&children[0], query_engine_state) - .await?; + let native_plan = match expr.name() { + ANALYZE_NODE_NAME + | ANALYZE_VERBOSE_NODE_NAME + | EXPLAIN_NODE_NAME + | EXPLAIN_VERBOSE_NODE_NAME => self.try_plan_native_fused(&children[0]).await?, + _ => None, + }; + let plan = match native_plan { + Some(plan) => plan, + None => { + self.prom_root_expr_to_plan(&children[0], query_engine_state) + .await? + } + }; // Wrapper for the explanation/analyze of the existing plan // https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/builder/struct.LogicalPlanBuilder.html#method.explain // if `analyze` is true, runs the actual plan and produces diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs new file mode 100644 index 00000000000..9cb5bdacba0 --- /dev/null +++ b/src/query/src/promql/planner/native.rs @@ -0,0 +1,528 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The native fused-series route of the PromQL planner. +//! +//! [`PromPlanner::try_plan_native_fused`] offers a small experimental subset to the native +//! compiler: a vector selector (optionally parenthesized), or nested unary `rate`, `increase`, +//! `sum_over_time`, and `abs` calls over a selector. It replaces that expression with one +//! [`FusedSeries`] node over the selector's table scan. Other expression shapes keep the existing +//! planner's route. +//! +//! What the node needs from here is three things the existing planner already resolves: the raw +//! series schema the tree's input was compiled against (the selector's selected field columns), the +//! statement's evaluation grid, and the logical tag names the output carries. The output schema is +//! built to the shape the existing unary pipeline produces — the time index, then the series +//! fields, then the tags — so a node planned here can be consumed by the same code that consumes a +//! `SeriesNormalize` output. + +use std::sync::Arc; +use std::time::{Duration, UNIX_EPOCH}; + +use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE, MITO2_ENGINE}; +use common_query::prelude::OTLP_AGGREGATION_TEMPORALITY_LABEL; +use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; +use datafusion::common::{DFSchema, DFSchemaRef}; +use datafusion::error::DataFusionError; +use datafusion::logical_expr::{Extension, LogicalPlan, TableScanBuilder}; +use datafusion_expr::Expr; +use datatypes::series::SeriesOrdering; +use datatypes::series::samples::{FieldSpec, SeriesSchema}; +use promql::extension_plan::{FusedSeries, FusedSeriesBinding}; +use promql_parser::parser::value::ValueType; +use promql_parser::parser::{ + EvalStmt, Expr as PromExpr, MatrixSelector, ParenExpr, VectorSelector, +}; +use session::context::QueryContextRef; +use snafu::ResultExt; +use store_api::metric_engine_consts::{DATA_SCHEMA_TABLE_ID_COLUMN_NAME, METRIC_ENGINE_NAME}; + +use crate::promql::error::{DataFusionPlanningSnafu, Result}; +use crate::promql::native_compile::{compile, statement_grid}; +use crate::promql::planner::PromPlanner; + +/// The query-context key that turns the native fused-series route on for one query. +/// +/// Named after the engine-wide option it mirrors ([`QueryOptions::experimental_native_promql`]), +/// and read exactly the way the other per-query override in this crate is read +/// ([`remote_dyn_filter_pushdown_enabled_from_extensions`]): a query-context extension, which a +/// client can set per request over the HTTP hints header or as a gRPC context key. The engine-wide +/// option is applied while building query-engine state; this planner reads the per-query request +/// extension at the call site; see [`native_promql_requested`]. +/// +/// [`QueryOptions::experimental_native_promql`]: crate::options::QueryOptions::experimental_native_promql +/// [`remote_dyn_filter_pushdown_enabled_from_extensions`]: +/// crate::options::remote_dyn_filter_pushdown_enabled_from_extensions +pub const NATIVE_PROMQL_EXTENSION_KEY: &str = "query.experimental_native_promql"; + +/// Whether `query_ctx` asks for the native fused-series route. +/// +/// The route is off unless the query asks for it, so a query that says nothing keeps the plan the +/// existing planner builds. The value is compared literally (`true` in any case) rather than through +/// a shared boolean parser: an unrecognized value leaves the route off, which is the same plan as +/// the flag never having been read, so a typo costs a slow query and never a wrong one. +pub(crate) fn native_promql_requested(query_ctx: &QueryContextRef) -> bool { + query_ctx + .extension(NATIVE_PROMQL_EXTENSION_KEY) + .is_some_and(|value| value.eq_ignore_ascii_case("true")) +} + +/// The ordering this route compiles its trees against. +/// +/// Which ordering a raw-series scan declares is decided when the region scanner is prepared, during +/// physical planning, and the logical planner cannot read it: a legacy scan declares `KeyThenTime` +/// and a two-phase scan `Contiguous`. So the route names the weakest ordering both provide — +/// `Contiguous`, every series contiguous and the series themselves unordered — rather than +/// overclaiming the legacy scan's stronger `KeyThenTime`. A native evaluation only refuses +/// `Interleaved`, so a tree compiled here runs on either scanner. +const NATIVE_ORDERING: SeriesOrdering = SeriesOrdering::Contiguous; + +impl PromPlanner { + /// Plans `prom_expr` as one [`FusedSeries`] node, or returns `None` to keep the existing plan. + /// + /// The route is offered the whole expression so its tree computes the requested expression + /// over the planner's resolved selector context: table name, time index, tag columns, and + /// field columns (after any `__field__` matcher). These are the raw series schema the input + /// tree is compiled against. + /// + /// `None` is not a failure: unsupported shapes and queries this route cannot bind faithfully + /// continue through the existing planner. Admission is limited to selectors and unary calls + /// named above, plus the scan, schema, matcher, and grid constraints checked below. + pub(super) async fn try_plan_native_fused( + &mut self, + prom_expr: &PromExpr, + ) -> Result> { + if !self.native_fused { + return Ok(None); + } + // Only the statement root is offered: the tree is compiled against the statement's grid, + // and a subquery's expression is evaluated on its own, smaller one. A range-vector root is + // the instant query's raw sample result, which the existing path reports as a matrix rather + // than as one point per grid position. + if prom_expr.value_type() == ValueType::Matrix { + return Ok(None); + } + + // Admit only a direct selector or the selected unary calls applied to one selector. + let Some(selector) = admitted_selector(prom_expr) else { + return Ok(None); + }; + + // These matchers are alternatives the one scan below the node cannot represent, so the + // query keeps the plan the existing path builds for it. + if !selector.matchers.or_matchers.is_empty() { + return Ok(None); + } + // Resolve the selector the way the existing planner does, and take what is left as the + // rows it selects. The scan below the node is the selector's own scan, unfiltered, so a + // matcher that survives preprocessing selects fewer rows here than the scan reads. + let matchers = self.preprocess_label_matchers(&selector.matchers, &selector.name)?; + if !matchers.matchers.is_empty() { + return Ok(None); + } + // An `__field__` matcher narrows the field columns the existing planner builds its scan + // with; folding every field of the table instead would evaluate a different window. + if self.ctx.field_column_matcher.is_some() { + return Ok(None); + } + // A metric without a table reads as an empty metric, which the existing planner plans. + if self.setup_context().await?.is_some() { + return Ok(None); + } + + let Some(table_name) = self.ctx.table_name.clone() else { + return Ok(None); + }; + let Some(time_index) = self.ctx.time_index_column.clone() else { + return Ok(None); + }; + let fields = self.ctx.field_columns.clone(); + let tags = self.ctx.tag_columns.clone(); + // A series with no field has nothing to fold, and the tree's input would be an empty + // schema, which no window can read. + if fields.is_empty() { + return Ok(None); + } + // On a delta-temporality table a series' samples are increments, which the existing + // planner sums over the window for `increase` and divides by the range for `rate` + // ([`PromPlanner::select_delta_range_math`]), per series because the tag can differ between + // them. A compiled tree computes both for a cumulative series, so any `rate`/`increase` + // over such a table keeps the existing plan; a bare selector still folds. + if tags + .iter() + .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL) + && calls_rate_or_increase(prom_expr) + { + return Ok(None); + } + + // The selector's own table scan, built by the same method the existing path uses: the + // physical planner turns it into the region scan the node reads, and building a scan here + // instead would let the two paths disagree about the columns they read. + let table_ref = self.table_ref()?; + let mut input = self.create_table_scan_plan(table_ref.clone()).await?; + let logical_scan = metric_scan_and_table_id(&input); + let has_verified_metric_rewrite = logical_scan.is_some(); + let scan = match logical_scan { + Some((mut scan, predicate)) => { + // Preserve logical field order and the table restriction on the direct physical scan. + let source_schema = scan.source.schema(); + let Some(mut ordered_projection) = fields + .iter() + .map(|field| source_schema.index_of(field).ok()) + .collect::>>() + else { + return Ok(None); + }; + let original_projection = scan.projection.as_ref(); + let selected_indices = original_projection + .cloned() + .unwrap_or_else(|| (0..source_schema.fields().len()).collect()); + for index in selected_indices { + let is_field = fields + .iter() + .any(|field| field == source_schema.field(index).name()); + if !is_field && !ordered_projection.contains(&index) { + ordered_projection.push(index); + } + } + scan.filters.push(predicate); + scan = TableScanBuilder::from(scan) + .with_projection(Some(ordered_projection)) + .build() + .context(DataFusionPlanningSnafu)?; + input = LogicalPlan::TableScan(scan.clone()); + scan + } + None => match &input { + LogicalPlan::TableScan(scan) => scan.clone(), + _ => return Ok(None), + }, + }; + let table = self.table_from_source(&scan.source)?; + let table_meta = &table.table_info().meta; + // An engine without a raw-series scanner has no frame to give this node, and by the time + // the physical planner asks, the logical plan is already this node: the refusal there is a + // planning error, not a fallback. So the engine is checked here, where the existing plan is + // still the answer. A file-engine table is the concrete case: its scanner is a + // `SinglePartitionScanner`, which leaves `raw_series_schema` at the default `None`. + if !supports_native_raw_series(&table_meta.engine) + || (table_meta.engine == METRIC_ENGINE_NAME && !has_verified_metric_rewrite) + { + return Ok(None); + } + // The region's raw-series preflight refuses these, and the physical layer turns that refusal into an error, so decline this route here instead. + let time_index_type = table_meta + .schema + .timestamp_column() + .map(|column| &column.data_type); + if time_index_type + != Some(&datatypes::prelude::ConcreteDataType::timestamp_millisecond_datatype()) + || table_meta + .options + .extra_options + .get(store_api::mito_engine_options::APPEND_MODE_KEY) + .is_some_and(|value| value.eq_ignore_ascii_case("true")) + { + return Ok(None); + } + // FusedSeries evaluates each region independently, so partition keys must identify a series. + if !table + .table_info() + .meta + .partition_column_names() + .all(|column| tags.iter().any(|tag| tag == column)) + { + return Ok(None); + } + // The legacy selector converts an offset the same way; sharing that conversion keeps the + // two scan paths from disagreeing about an extreme duration. + let offset = Self::offset_millis(&selector.offset); + let (saved_start, saved_end, saved_range) = (self.ctx.start, self.ctx.end, self.ctx.range); + self.ctx.range = selector_range_millis(prom_expr); + // Keep the fused scan's bounds identical to the legacy selector scan, including @ and offset. + let time_filter = (|| { + let anchored = self.at_modifier_offset(&selector.at, &selector.offset)?; + if anchored.is_some() { + self.ctx.end = saved_start; + } + self.build_time_index_filter(anchored.unwrap_or(offset), input.schema()) + })(); + self.ctx.start = saved_start; + self.ctx.end = saved_end; + self.ctx.range = saved_range; + let time_filter = time_filter?; + if let Some(filter) = time_filter { + let LogicalPlan::TableScan(scan) = &mut input else { + unreachable!(); + }; + scan.filters.push(filter); + } + // The region's columnar series-index samples, one of the readers this scan may be served + // by, store `Float64` fields only, and the logical plan cannot tell which reader arrives. + // So every field has to be `Float64`: an integer, a `Float32`, a decimal, a string or a + // native histogram field keeps the existing plan, which the physical layer still serves. + if !fields.iter().all(|field| { + Self::field_column_type(input.schema(), field) == Some(&ArrowDataType::Float64) + }) { + return Ok(None); + } + let Some(stmt) = self.native_statement(prom_expr) else { + return Ok(None); + }; + let Ok(grid) = statement_grid(&stmt) else { + // The statement's grid is not one a tree can be evaluated on, so there is no grid to + // address the node's output by. + return Ok(None); + }; + + let series_schema = Arc::new(SeriesSchema { + fields: fields + .iter() + .map(|name| FieldSpec { + name: Arc::from(name.as_str()), + }) + .collect::>() + .into(), + eval_grid: None, + }); + let tree = match compile(&stmt, Arc::clone(&series_schema), NATIVE_ORDERING) { + Ok(tree) => tree, + // The compiler refusing a shape is this route declining it: the existing planner gets + // the query exactly as it would have without this flag. + Err(DataFusionError::NotImplemented(_)) => return Ok(None), + Err(error) => return Err(error).context(DataFusionPlanningSnafu), + }; + let binding = FusedSeriesBinding { + series_schema, + // The tree's own output is evaluated on the statement's grid, and the binding names + // the same one, so a chunk's absolute positions are read on the grid it was computed + // on. + grid, + metric_name: Some(Arc::from(table_name.as_str())), + time_index: time_index.clone(), + tags: tags.clone(), + }; + let output_schema = fused_output_schema(&fields, &tags, &time_index)?; + let node = FusedSeries::try_new( + Arc::from(prom_expr.to_string()), + tree, + binding, + input, + output_schema, + ) + .context(DataFusionPlanningSnafu)?; + + Ok(Some(LogicalPlan::Extension(Extension { + node: Arc::new(node), + }))) + } + + /// The statement being planned, rebuilt from the evaluation numbers the planner holds. + /// + /// The compiler reads the statement's grid, lookback and `@` anchors rather than the planner's + /// mutable state. The route is only offered at the statement root, where `ctx.start` and + /// `ctx.end` are the statement's own bounds: they are rewritten while an anchored window or a + /// subquery is folded, so a nested expression would need the bounds it was actually entered + /// with, which is why no nested expression is offered here. + /// + /// `None` means the bounds are not representable as a [`SystemTime`], which is not a grid a + /// tree can be evaluated on either. + fn native_statement(&self, prom_expr: &PromExpr) -> Option { + let start = UNIX_EPOCH.checked_add(Duration::from_millis(self.ctx.stmt_start as u64))?; + let end = UNIX_EPOCH.checked_add(Duration::from_millis(self.ctx.stmt_end as u64))?; + Some(EvalStmt { + expr: prom_expr.clone(), + start, + end, + interval: Duration::from_millis(self.ctx.interval as u64), + lookback_delta: Duration::from_millis(self.ctx.lookback_delta as u64), + }) + } +} + +/// The columns a [`FusedSeries`] node produces. +/// +/// The shape is the one the existing unary pipeline consumes — the time index, then the series +/// fields, then the tags — and each column is declared the way the node's own conversion reads it +/// back: +/// the time column is a non-nullable millisecond timestamp with no timezone, because every +/// evaluated chunk is addressed on the statement's grid and carries one; a field is a nullable +/// `Float64`, because a grid point inside the statement's interval can have no sample; and a tag is +/// a nullable `Utf8`, because a range function drops the metric name and a series can be missing a +/// label entirely. +fn fused_output_schema( + fields: &[String], + tags: &[String], + time_index: &str, +) -> Result { + let mut columns = Vec::with_capacity(fields.len() + tags.len() + 1); + columns.push(( + None, + Arc::new(Field::new( + time_index, + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )), + )); + for name in fields { + columns.push(( + None, + Arc::new(Field::new(name, ArrowDataType::Float64, true)), + )); + } + for name in tags { + columns.push((None, Arc::new(Field::new(name, ArrowDataType::Utf8, true)))); + } + + Ok(Arc::new( + DFSchema::new_with_metadata(columns, Default::default()) + .context(DataFusionPlanningSnafu)?, + )) +} + +/// The one selector `expr` folds, if it names exactly one. +/// +/// The walk follows the grammar the native compiler accepts — a call, a range selector, a selector +/// and parentheses — and stops at every other shape, because a shape the compiler cannot compile +/// needs no selector resolved. Two distinct selectors are not one input, and a tree built over the +/// wrong one of them would evaluate a window the query did not write. +fn selector_range_millis(expr: &PromExpr) -> Option { + match expr { + PromExpr::MatrixSelector(matrix) => Some(matrix.range.as_millis() as i64), + PromExpr::Paren(paren) => selector_range_millis(&paren.expr), + PromExpr::Call(call) => call + .args + .args + .iter() + .find_map(|arg| selector_range_millis(arg)), + _ => None, + } +} + +/// Whether `engine` can serve raw series frames to the fused-series node. +/// +/// A whitelist avoids probing a region from logical planning; declining only costs the native route. +fn supports_native_raw_series(engine: &str) -> bool { + matches!(engine, MITO_ENGINE | MITO2_ENGINE | METRIC_ENGINE) +} + +/// Recognizes only the metric logical-table scan emitted by `create_table_scan_plan`. +fn metric_scan_and_table_id( + plan: &LogicalPlan, +) -> Option<(datafusion::logical_expr::TableScan, Expr)> { + let plan = match plan { + LogicalPlan::Projection(projection) => projection.input.as_ref(), + plan => plan, + }; + let plan = match plan { + LogicalPlan::SubqueryAlias(alias) => alias.input.as_ref(), + plan => plan, + }; + let LogicalPlan::Filter(filter) = plan else { + return None; + }; + let Expr::BinaryExpr(binary) = &filter.predicate else { + return None; + }; + if binary.op != datafusion::logical_expr::Operator::Eq + || !matches!(binary.left.as_ref(), Expr::Column(column) if column.name == DATA_SCHEMA_TABLE_ID_COLUMN_NAME) + { + return None; + } + if !matches!( + binary.right.as_ref(), + Expr::Literal(datafusion::scalar::ScalarValue::UInt32(Some(_)), _) + ) { + return None; + } + let LogicalPlan::TableScan(scan) = filter.input.as_ref() else { + return None; + }; + Some((scan.clone(), filter.predicate.clone())) +} + +fn admitted_selector(expr: &PromExpr) -> Option<&VectorSelector> { + match expr { + PromExpr::VectorSelector(selector) => Some(selector), + // Range selectors are admitted only as arguments within one of these supported calls. + PromExpr::MatrixSelector(MatrixSelector { vs, .. }) => Some(vs), + PromExpr::Paren(ParenExpr { expr }) => admitted_selector(expr), + PromExpr::Call(call) + if matches!( + call.func.name, + "rate" | "increase" | "sum_over_time" | "abs" + ) && call.args.args.len() == 1 => + { + admitted_selector(&call.args.args[0]) + } + _ => None, + } +} + +/// Whether `expr` calls `rate`/`increase`, the functions the existing planner rewrites per series +/// for a delta-temporality table. Only nested calls and parentheses can hold one in an expression +/// this route accepts; every other shape names no single selector or reads across series. +fn calls_rate_or_increase(expr: &PromExpr) -> bool { + match expr { + PromExpr::Call(call) => { + matches!(call.func.name, "rate" | "increase") + || call.args.args.iter().any(|arg| calls_rate_or_increase(arg)) + } + PromExpr::Paren(ParenExpr { expr }) => calls_rate_or_increase(expr), + _ => false, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use session::context::QueryContextBuilder; + + use super::{NATIVE_PROMQL_EXTENSION_KEY, native_promql_requested}; + + #[test] + fn native_promql_requested_is_false_when_extension_is_absent() { + let query_ctx = Arc::new(QueryContextBuilder::default().build()); + + assert!(!native_promql_requested(&query_ctx)); + } + + #[test] + fn native_promql_requested_accepts_true_case_insensitively() { + for value in ["true", "TRUE", "TrUe"] { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .set_extension(NATIVE_PROMQL_EXTENSION_KEY.to_string(), value.to_string()) + .build(), + ); + + assert!(native_promql_requested(&query_ctx), "value: {value}"); + } + } + + #[test] + fn native_promql_requested_rejects_false_and_unrecognized_values() { + for value in ["false", "FALSE", "yes", "1", " true ", ""] { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .set_extension(NATIVE_PROMQL_EXTENSION_KEY.to_string(), value.to_string()) + .build(), + ); + + assert!(!native_promql_requested(&query_ctx), "value: {value:?}"); + } + } +} diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index ae40640b50b..13229b11474 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -18,7 +18,9 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use catalog::RegisterTableRequest; use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; use common_base::Plugins; -use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_catalog::consts::{ + DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE, +}; use common_query::native_histogram::{ CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, }; @@ -36,10 +38,12 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::datasource::{MemTable, provider_as_source}; use datafusion::execution::context::SessionContext; -use datafusion::logical_expr::{Extension, col}; +use datafusion::logical_expr::{Extension, UserDefinedLogicalNodeCore, col}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; -use promql::extension_plan::HistogramFold; +use datatypes::series::SeriesOrdering; +use promql::extension_plan::{FusedSeries, HistogramFold}; +use promql::native::plan::NativeExpr; use promql_parser::label::Labels; use promql_parser::parser; use session::context::QueryContext; @@ -474,6 +478,7 @@ async fn plan_direct_or( table_provider, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; planner .or_operator( @@ -563,6 +568,7 @@ async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) table_provider, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let left_context = direct_or_context("lhs", &["job", "k"], "v"); let right_context = direct_or_context("rhs", &["job", "k"], "v"); @@ -635,6 +641,7 @@ async fn mixed_aggregate_input(histograms: Vec) -> (PromPlanner ..Default::default() }, promql_annotations: None, + native_fused: false, }; (planner, plan) } @@ -891,10 +898,157 @@ async fn optimize_and_create_physical_plan( (optimized, physical) } +/// The engine a fixture gets when the test does not care where the scan is served. +/// +/// `TableMeta::empty()` leaves it blank, and the existing plan assertions are written against that +/// blank engine; naming it here keeps the native-route fixtures explicit instead. +const DEFAULT_ENGINE: &str = ""; + +/// The engine the native route accepts, because its scanner serves raw series frames. +/// +/// It also decides whether a tag column is dictionary encoded, so a fixture that sets it sees +/// `Dictionary(UInt32, Utf8)` where a blank-engine fixture sees `Utf8`. +const NATIVE_ENGINE: &str = MITO_ENGINE; + async fn build_test_table_provider( table_name_tuples: &[(String, String)], num_tag: usize, num_field: usize, +) -> DfTableSourceProvider { + build_test_table_provider_with_partition_columns(table_name_tuples, num_tag, num_field, vec![]) + .await +} + +async fn build_test_table_provider_with_partition_columns( + table_name_tuples: &[(String, String)], + num_tag: usize, + num_field: usize, + partition_key_indices: Vec, +) -> DfTableSourceProvider { + build_test_table_provider_with_options( + table_name_tuples, + num_tag, + num_field, + partition_key_indices, + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + DEFAULT_ENGINE, + ) + .await +} + +/// A table whose single value field has the given type. +async fn build_test_table_provider_with_field_type( + field_type: ConcreteDataType, +) -> DfTableSourceProvider { + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + field_type, + // The native engine, so the field guard below is what refuses this table: on the blank + // engine the engine gate would answer first and this test would prove nothing. + NATIVE_ENGINE, + ) + .await +} + +/// A one-tag, one-field fixture on the engine the native route accepts. +async fn native_engine_provider(partition_key_indices: Vec) -> DfTableSourceProvider { + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + partition_key_indices, + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, + ) + .await +} + +/// A delta-temporality table on the engine the native route accepts. +/// +/// The OTLP receiver writes the aggregation temporality as a tag of every series, so a table with +/// that tag can hold delta and cumulative series side by side: the existing planner sums the +/// increments of a delta series instead of rating its counter. The value field is the `Float64` +/// this route folds. +async fn native_engine_temporality_marker_provider() -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), + ConcreteDataType::string_datatype(), + true, + ), + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ])); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .engine(NATIVE_ENGINE.to_string()) + .primary_key_indices(vec![0, 1]) + .value_indices(vec![3]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name("some_metric") + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "some_metric".to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +#[allow(clippy::too_many_arguments)] +async fn build_test_table_provider_with_options( + table_name_tuples: &[(String, String)], + num_tag: usize, + num_field: usize, + partition_key_indices: Vec, + timestamp_type: ConcreteDataType, + append_mode: bool, + field_type: ConcreteDataType, + engine: &str, ) -> DfTableSourceProvider { let catalog_list = MemoryCatalogManager::with_default_setup(); for (schema_name, table_name) in table_name_tuples { @@ -907,26 +1061,35 @@ async fn build_test_table_provider( )); } columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), + ColumnSchema::new("timestamp".to_string(), timestamp_type.clone(), false) + .with_time_index(true), ); for i in 0..num_field { columns.push(ColumnSchema::new( format!("field_{i}"), - ConcreteDataType::float64_datatype(), + field_type.clone(), true, )); } let schema = Arc::new(Schema::new(columns)); let table_meta = TableMetaBuilder::empty() .schema(schema) + .engine(engine.to_string()) .primary_key_indices((0..num_tag).collect()) + .partition_key_indices(partition_key_indices.clone()) .value_indices((num_tag + 1..num_tag + 1 + num_field).collect()) .next_column_id(1024) + .options(table::requests::TableOptions { + extra_options: if append_mode { + HashMap::from([( + store_api::mito_engine_options::APPEND_MODE_KEY.to_string(), + "true".to_string(), + )]) + } else { + HashMap::new() + }, + ..Default::default() + }) .build() .unwrap(); let table_info = TableInfoBuilder::default() @@ -1345,6 +1508,13 @@ async fn build_test_table_provider_with_tsid_fields( async fn build_test_table_provider_with_tsid_tag_fields( table_specs: &[((String, String), usize, usize)], +) -> DfTableSourceProvider { + build_test_table_provider_with_tsid_tag_fields_and_field_order(table_specs, None).await +} + +async fn build_test_table_provider_with_tsid_tag_fields_and_field_order( + table_specs: &[((String, String), usize, usize)], + logical_field_order: Option<&[usize]>, ) -> DfTableSourceProvider { let catalog_list = MemoryCatalogManager::with_default_setup(); @@ -1449,7 +1619,16 @@ async fn build_test_table_provider_with_tsid_tag_fields( ) .with_time_index(true), ); - for i in 0..*num_field { + let field_order = logical_field_order + .map(|order| { + order + .iter() + .copied() + .filter(|i| i < num_field) + .collect::>() + }) + .unwrap_or_else(|| (0..*num_field).collect()); + for i in field_order { columns.push(ColumnSchema::new( format!("field_{i}"), ConcreteDataType::float64_datatype(), @@ -1617,6 +1796,693 @@ async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) assert_eq!(plan.display_indent_schema().to_string(), expected); } +/// Plans `query` with the native fused-series route on or off. +/// +/// The table is the one [`build_test_table_provider`] builds — one tag `tag_0`, the millisecond +/// time index `timestamp`, one field `field_0` — which is the shape the selector context resolves +/// to a raw series schema of one field. +async fn build_native_fused_plan(query: &str, native_fused: bool) -> LogicalPlan { + let table_provider = build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, + ) + .await; + build_native_fused_plan_with_provider(query, native_fused, table_provider).await +} + +async fn build_native_fused_plan_with_provider( + query: &str, + native_fused: bool, + table_provider: DfTableSourceProvider, +) -> LogicalPlan { + let eval_stmt = build_eval_stmt(query); + PromPlanner::stmt_to_plan_with_annotations( + table_provider, + &eval_stmt, + &build_query_engine_state(), + None, + native_fused, + ) + .await + .unwrap() +} + +fn find_metric_physical_scan(plan: &LogicalPlan) -> Option { + match plan { + LogicalPlan::Extension(Extension { node }) => node + .as_any() + .downcast_ref::() + .and_then(|fused| match fused.inputs().first() { + Some(LogicalPlan::TableScan(scan)) => Some(scan.clone()), + _ => None, + }), + _ => plan + .inputs() + .into_iter() + .find_map(find_metric_physical_scan), + } +} + +fn find_fused_series(plan: &LogicalPlan) -> Option<&FusedSeries> { + if let LogicalPlan::Extension(Extension { node }) = plan + && let Some(fused) = node.as_any().downcast_ref::() + { + return Some(fused); + } + + plan.inputs().into_iter().find_map(find_fused_series) +} + +/// The flag on and a shape the compiler accepts: the statement plans to one native node. +#[tokio::test] +async fn native_fused_rate_plans_to_a_fused_series_node() { + let plan = build_native_fused_plan("rate(some_metric[5m])", true).await; + + let display = plan.display_indent().to_string(); + assert!( + display.contains("FusedSeries: rate(some_metric[5m])"), + "{display}" + ); + // The node reads the selector's own scan; nothing between them packs the frame it reads. + assert!(display.contains("TableScan: some_metric"), "{display}"); + // The node carries the query a reader wrote, not the shape it compiled to. + let fused = find_fused_series(&plan).expect("the plan holds a native node"); + assert_eq!("rate(some_metric[5m])", fused.promql()); + assert_eq!( + vec!["timestamp", "field_0", "tag_0"], + plan.schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>() + ); + // The scan below is chosen during physical planning, so the route cannot know whether it is + // the legacy scan (`KeyThenTime`) or the two-phase one (`Contiguous`) and must bind the weaker + // guarantee both provide. `rate` compiles to one window over that raw input. + let NativeExpr::Window { input, .. } = &fused.tree().root().expr else { + panic!("the root is not a window: {:?}", fused.tree().root().expr); + }; + let NativeExpr::Input(bound) = &input.expr else { + panic!("the window's input is not an input: {:?}", input.expr); + }; + assert_eq!( + SeriesOrdering::Contiguous, + bound.ordering, + "the tree claimed an ordering a raw-series scan need not guarantee" + ); +} + +#[tokio::test] +async fn native_fused_time_partitioned_table_falls_back() { + let table_provider = native_engine_provider(vec![1]).await; + let native = + build_native_fused_plan_with_provider("rate(some_metric[2m])", true, table_provider).await; + let fallback_provider = native_engine_provider(vec![1]).await; + let without_flag = + build_native_fused_plan_with_provider("rate(some_metric[2m])", false, fallback_provider) + .await; + let display = native.display_indent().to_string(); + assert!(!display.contains("FusedSeries"), "{display}"); + assert_eq!( + display, + without_flag.display_indent().to_string(), + "a time-partitioned table must keep the existing plan, not an approximation of it" + ); +} + +#[tokio::test] +async fn native_fused_non_millisecond_time_index_falls_back() { + let table_provider = build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_microsecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, + ) + .await; + let native = + build_native_fused_plan_with_provider("rate(some_metric[5m])", true, table_provider).await; + let fallback_provider = build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_microsecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, + ) + .await; + let without_flag = + build_native_fused_plan_with_provider("rate(some_metric[5m])", false, fallback_provider) + .await; + assert!(!native.display_indent().to_string().contains("FusedSeries")); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string() + ); +} + +#[tokio::test] +async fn native_fused_append_mode_table_falls_back() { + let table_provider = build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + true, + ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, + ) + .await; + let native = + build_native_fused_plan_with_provider("rate(some_metric[5m])", true, table_provider).await; + let fallback_provider = build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + true, + ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, + ) + .await; + let without_flag = + build_native_fused_plan_with_provider("rate(some_metric[5m])", false, fallback_provider) + .await; + assert!(!native.display_indent().to_string().contains("FusedSeries")); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string() + ); +} + +#[tokio::test] +async fn native_fused_tag_partitioned_table_still_plans_fused_series() { + let table_provider = native_engine_provider(vec![0]).await; + let plan = + build_native_fused_plan_with_provider("rate(some_metric[5m])", true, table_provider).await; + assert!(plan.display_indent().to_string().contains("FusedSeries")); +} + +/// Metric logical scans retain their table-id restriction on the direct fused scan input. +#[tokio::test] +async fn native_fused_metric_engine_table_is_accepted() { + let native = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + true, + build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await, + ) + .await; + let legacy = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + false, + build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await, + ) + .await; + assert!( + find_fused_series(&native).is_some(), + "{}", + native.display_indent() + ); + assert!( + legacy + .display_indent() + .to_string() + .contains("PromRangeManipulate") + ); + let scan = find_metric_physical_scan(&native).expect("fused route keeps physical scan"); + let source_schema = scan.source.schema(); + let projected_names = scan + .projection + .as_ref() + .unwrap() + .iter() + .map(|index| source_schema.field(*index).name()) + .collect::>(); + assert_eq!( + projected_names + .iter() + .filter(|name| name.starts_with("field_")) + .copied() + .collect::>(), + vec!["field_0"] + ); + assert!(scan.filters.iter().any(|filter| { + *filter + == datafusion::logical_expr::Expr::Column(datafusion_common::Column::new( + Some(datafusion_common::TableReference::bare("phy")), + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), + )) + .eq(datafusion::logical_expr::lit(1024u32)) + })); + assert!( + scan.filters + .iter() + .any(|filter| format!("{filter:?}").contains("timestamp")) + ); +} + +/// Metric fields follow logical metadata order even when physical columns use another order. +#[tokio::test] +async fn native_fused_metric_fields_follow_logical_order() { + let native = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + true, + build_test_table_provider_with_tsid_tag_fields_and_field_order( + &[( + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + 1, + 2, + )], + Some(&[1, 0]), + ) + .await, + ) + .await; + assert!( + find_fused_series(&native).is_some(), + "{}", + native.display_indent() + ); + let scan = find_metric_physical_scan(&native).expect("fused route keeps physical scan"); + let source_schema = scan.source.schema(); + let projected_names = scan + .projection + .as_ref() + .unwrap() + .iter() + .map(|index| source_schema.field(*index).name()) + .collect::>(); + assert_eq!(projected_names[..2], ["field_1", "field_0"]); + assert!(scan.filters.iter().any(|filter| { + *filter + == datafusion::logical_expr::Expr::Column(datafusion_common::Column::new( + Some(datafusion_common::TableReference::bare("phy")), + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), + )) + .eq(datafusion::logical_expr::lit(1024u32)) + })); +} + +/// A metric table without logical-table metadata declines to the legacy plan. +#[tokio::test] +async fn native_fused_metric_table_without_logical_metadata_falls_back() { + let native = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + true, + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + METRIC_ENGINE, + ) + .await, + ) + .await; + let legacy = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + false, + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + METRIC_ENGINE, + ) + .await, + ) + .await; + assert!(find_fused_series(&native).is_none()); + assert_eq!( + native.display_indent().to_string(), + legacy.display_indent().to_string() + ); +} + +/// A table whose engine has no raw-series scanner keeps the existing plan. +/// +/// A file-engine table answers an ordinary read with a `SinglePartitionScanner`, whose +/// `raw_series_schema` is the trait's default `None`. Offering the native route to such a table +/// would build the node and then fail in physical planning, because the logical plan is fixed by +/// then; declining here is what keeps the query answering. +#[tokio::test] +async fn native_fused_engine_without_raw_series_falls_back() { + let native = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + true, + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + FILE_ENGINE, + ) + .await, + ) + .await; + let without_flag = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + false, + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + FILE_ENGINE, + ) + .await, + ) + .await; + assert!( + find_fused_series(&native).is_none(), + "a file-engine table must not reach the native route" + ); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string(), + "a file-engine table must keep the existing plan" + ); +} + +/// A value field the raw-series readers cannot both serve keeps the existing plan. +/// +/// The region's preflight serves numeric columns and native histograms, while the columnar +/// series-index samples the scan may be read through store `Float64` fields only, and the logical +/// plan cannot tell which reader the region prepares. So only `Float64` is admitted; a bare +/// selector over any other field type answered normally before this route existed. +#[tokio::test] +async fn native_fused_unsupported_field_type_falls_back() { + for field_type in [ + ConcreteDataType::string_datatype(), + ConcreteDataType::boolean_datatype(), + // Arrow's `is_numeric` includes the decimals while the region's preflight does not, so + // asking Arrow would let a decimal field through to a planning error. + ConcreteDataType::decimal128_datatype(10, 2), + // Numeric, but not the `Float64` the columnar samples store. + ConcreteDataType::int64_datatype(), + ConcreteDataType::float32_datatype(), + ] { + let native = build_native_fused_plan_with_provider( + "some_metric", + true, + build_test_table_provider_with_field_type(field_type.clone()).await, + ) + .await; + let without_flag = build_native_fused_plan_with_provider( + "some_metric", + false, + build_test_table_provider_with_field_type(field_type.clone()).await, + ) + .await; + assert!( + find_fused_series(&native).is_none(), + "{field_type} must not reach the native route" + ); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string(), + "{field_type} must keep the existing plan" + ); + } +} + +/// The `Float64` field every raw reader serves controls the field-type decline above. +#[tokio::test] +async fn native_fused_float64_field_is_accepted() { + let plan = build_native_fused_plan_with_provider( + "some_metric", + true, + build_test_table_provider_with_field_type(ConcreteDataType::float64_datatype()).await, + ) + .await; + + assert!( + find_fused_series(&plan).is_some(), + "a `Float64` field must still fold: {}", + plan.display_indent() + ); +} + +/// Alternative (`or`) matchers keep the existing plan. +/// +/// `a{x="1" or x="2"}` names alternative matchers, which the one scan below the node cannot +/// represent. These queries are only asserted to fall back: what the existing plan answers for +/// them is its own question, not one this route claims to have fixed. +#[tokio::test] +async fn native_fused_or_matchers_fall_back() { + for query in [ + r#"some_metric{tag_0="a" or tag_0="b"}"#, + r#"rate(some_metric{tag_0="a" or tag_0="b"}[5m])"#, + ] { + let native = build_native_fused_plan(query, true).await; + let without_flag = build_native_fused_plan(query, false).await; + assert!( + find_fused_series(&native).is_none(), + "`{query}` names alternative matchers: {}", + native.display_indent() + ); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string(), + "`{query}` must keep the existing plan" + ); + } +} + +/// A `rate`/`increase` over a delta-temporality table keeps the existing plan, which sums each +/// series' increments. +/// +/// The samples of a series tagged `delta` are increments, which the existing planner sums over the +/// window for `increase` and divides by the range for `rate`, while a cumulative series is rated +/// from its counter; it picks that per series. A compiled tree rates every series the same way, so +/// these queries fall back. The same queries over a table without the marker still fold, which is +/// what makes the decline the marker's doing rather than the shape's. +#[tokio::test] +async fn native_fused_delta_temporality_rates_fall_back() { + for query in [ + "rate(some_metric[5m])", + "increase(some_metric[5m])", + // Below another call and inside parentheses, which are the same decline one level down. + "abs(rate(some_metric[5m]))", + "abs(increase(some_metric[5m]))", + "(rate(some_metric[5m]))", + ] { + let marked = build_native_fused_plan_with_provider( + query, + true, + native_engine_temporality_marker_provider().await, + ) + .await; + let marked_without_flag = build_native_fused_plan_with_provider( + query, + false, + native_engine_temporality_marker_provider().await, + ) + .await; + assert!( + find_fused_series(&marked).is_none(), + "`{query}` needs the per-series delta math: {}", + marked.display_indent() + ); + assert_eq!( + marked.display_indent().to_string(), + marked_without_flag.display_indent().to_string(), + "`{query}` must keep the existing plan, not the counter form for every series" + ); + + let unmarked = build_native_fused_plan_with_provider( + query, + true, + native_engine_provider(vec![]).await, + ) + .await; + assert!( + find_fused_series(&unmarked).is_some(), + "`{query}` over a table without the marker must still fold: {}", + unmarked.display_indent() + ); + } +} + +/// The marker alone does not decline a query: only `rate`/`increase` are refused. +#[tokio::test] +async fn native_fused_delta_temporality_non_counter_math_still_folds() { + for query in ["some_metric", "abs(some_metric)"] { + let plan = build_native_fused_plan_with_provider( + query, + true, + native_engine_temporality_marker_provider().await, + ) + .await; + assert!( + find_fused_series(&plan).is_some(), + "`{query}` needs no per-series delta math and must still fold: {}", + plan.display_indent() + ); + } +} + +/// A whole-vector operation is not a per-series tree, so it stays on the existing plan. +#[tokio::test] +async fn native_fused_whole_vector_operation_falls_back() { + for query in ["sort(some_metric)", "absent(some_metric)"] { + let native = build_native_fused_plan(query, true).await; + let without_flag = build_native_fused_plan(query, false).await; + assert!( + find_fused_series(&native).is_none(), + "`{query}` reads the whole vector and must not be placed region-side" + ); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string(), + "`{query}` must keep the existing plan, not an approximation of it" + ); + } +} + +#[tokio::test] +async fn native_fused_unsupported_shape_falls_back() { + for query in [ + "sum(rate(some_metric[5m]))", + "some_metric + some_metric", + "sum_over_time(rate(some_metric[5m])[1h:])", + "histogram_count(some_metric)", + "ceil(some_metric)", + ] { + let plan = build_native_fused_plan(query, true).await; + let display = plan.display_indent().to_string(); + assert!( + !display.contains("FusedSeries"), + "`{query}` must not enter the limited native POC: {display}" + ); + // Falling back is the existing plan, not an approximation of it. + let without_flag = build_native_fused_plan(query, false).await; + assert_eq!( + display, + without_flag.display_indent().to_string(), + "`{query}` must keep the existing plan" + ); + } +} + +/// A label rewrite stays on the existing plan, which is the one that projects its new label. +/// +/// The native node's output schema is the table's own time index, fields and tags, so a tree that +/// rewrote labels would report a label set the existing plan's projection emits and the node cannot. +/// Each query below writes `newlabel`, which the fixture table does not have (its only tag is +/// `tag_0`), so a fused plan would visibly lose it. +#[tokio::test] +async fn native_fused_label_rewrites_keep_legacy_output_labels() { + for query in [ + // The rewrite at the statement root... + r#"label_replace(some_metric, "newlabel", "$1", "tag_0", "(.*)")"#, + r#"label_join(some_metric, "newlabel", "-", "tag_0", "tag_0")"#, + // ... and below another call, which is the same shape one level deeper. + r#"abs(label_replace(some_metric, "newlabel", "$1", "tag_0", "(.*)"))"#, + r#"abs(label_join(some_metric, "newlabel", "-", "tag_0", "tag_0"))"#, + ] { + let native = build_native_fused_plan(query, true).await; + let display = native.display_indent().to_string(); + assert!( + find_fused_series(&native).is_none(), + "`{query}` rewrites labels the node's schema cannot declare: {display}" + ); + assert!( + display.contains("newlabel"), + "`{query}` must keep the projection that emits its new label: {display}" + ); + let without_flag = build_native_fused_plan(query, false).await; + assert_eq!( + display, + without_flag.display_indent().to_string(), + "`{query}` must keep the existing plan, not an approximation of it" + ); + } +} + +/// The flag off: no statement is offered to the native route at all. +#[tokio::test] +async fn native_fused_flag_off_keeps_the_existing_plan() { + let plan = build_native_fused_plan("rate(some_metric[5m])", false).await; + + let display = plan.display_indent().to_string(); + assert!(!display.contains("FusedSeries"), "{display}"); + assert!(display.contains("PromRangeManipulate"), "{display}"); +} + +/// The native selector scan is bounded to the range selector window, not the table retention. +#[tokio::test] +async fn native_fused_rate_scan_has_bounded_time_filter() { + let plan = build_native_fused_plan("rate(some_metric[5m])", true).await; + let display = plan.display_indent().to_string(); + + assert!( + display.contains("timestamp >= TimestampMillisecond(-299999, None)") + && display.contains("timestamp <= TimestampMillisecond(100000000, None)"), + "native scan should include the bounded five-minute statement window: {display}" + ); +} + +/// An anchored native range scan uses the anchor, while preserving the outer statement window. +#[tokio::test] +async fn native_fused_offset_rate_scan_uses_offset_window_bounds() { + let plan = build_native_fused_plan("rate(some_metric[5m] offset 10s)", true).await; + let display = plan.display_indent().to_string(); + + assert!( + display.contains("timestamp >= TimestampMillisecond(-309999, None)") + && display.contains("timestamp <= TimestampMillisecond(99990000, None)"), + "native scan should account for the selector offset: {display}" + ); +} + +#[tokio::test] +async fn native_fused_anchored_rate_scan_uses_anchor_bounds() { + let plan = build_native_fused_plan("rate(some_metric[5m] @ 300)", true).await; + let display = plan.display_indent().to_string(); + + assert!( + display.contains("timestamp >= TimestampMillisecond(1, None)") + && display.contains("timestamp <= TimestampMillisecond(300000, None)"), + "native scan should cover the five-minute window ending at the anchor: {display}" + ); +} + #[tokio::test] async fn single_abs() { do_single_instant_function_call("abs", "abs").await; @@ -3328,6 +4194,7 @@ async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() { table_provider, ctx: PromPlannerContext::from_eval_stmt(&eval_stmt).unwrap(), promql_annotations: None, + native_fused: false, }; let plan = planner .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state()) @@ -4463,6 +5330,7 @@ async fn mixed_histogram_helpers_report_annotations() { &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), &state, Some(collector.clone()), + false, ) .await .unwrap(); @@ -4623,6 +5491,7 @@ async fn mixed_or_topk_bottomk_ignore_native_histograms() { &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")), &state, Some(collector.clone()), + false, ) .await .unwrap(); @@ -4960,6 +5829,7 @@ async fn mixed_native_histogram_rate_executes_real_ranges() { ..Default::default() }, promql_annotations: Some(collector.clone()), + native_fused: false, }; let input = LogicalPlan::Extension(Extension { node: Arc::new( @@ -6289,6 +7159,7 @@ async fn native_scan_bounds_preserve_zero_lookback_and_overflow() { table_provider, ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")).unwrap(), promql_annotations: None, + native_fused: false, }; planner.ctx.time_index_column = Some("timestamp".to_string()); planner.ctx.start = 1_000; @@ -7129,6 +8000,7 @@ async fn test_direct_or_numeric_value_types() { table_provider: provider, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let left_context = direct_or_context("lhs", &["job"], "v"); let right_context = direct_or_context("rhs", &["job"], "v"); @@ -7616,6 +8488,7 @@ async fn malformed_classic_bucket_does_not_drop_native_histogram() { &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"), &state, Some(collector.clone()), + false, ) .await .unwrap(); @@ -7699,6 +8572,7 @@ async fn test_mixed_binary_operator_reports_only_dropped_samples() { &operator_eval_stmt(query), &state, Some(annotations.clone()), + false, ) .await .unwrap(); @@ -8360,6 +9234,7 @@ async fn test_mixed_or_value_aliases_do_not_replace_labels() { table_provider, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let left = LogicalPlanBuilder::from(scan(&left)) .project(vec![ diff --git a/src/query/src/promql/planner/test/delta.rs b/src/query/src/promql/planner/test/delta.rs index 9ca8a8b646c..be0f5b3b28f 100644 --- a/src/query/src/promql/planner/test/delta.rs +++ b/src/query/src/promql/planner/test/delta.rs @@ -411,6 +411,7 @@ async fn delta_mixed_ranges_drop_and_float_ranges_sum() { .await, ctx: PromPlannerContext::default(), promql_annotations: Some(collector.clone()), + native_fused: false, }; let left_context = direct_or_context("lhs", &["job", marker], "v"); let right_context = direct_or_context("rhs", &["job", marker], "v"); @@ -712,6 +713,7 @@ async fn binary_joins_align_only_the_temporality_marker() { .await, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let joined = planner .join_on_non_field_columns( @@ -751,6 +753,7 @@ async fn binary_joins_align_only_the_temporality_marker() { .await, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let set = planner .set_op_on_non_field_columns( @@ -806,6 +809,7 @@ async fn binary_joins_align_only_the_temporality_marker() { .await, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let set = planner .set_op_on_non_field_columns( @@ -918,6 +922,7 @@ async fn binary_joins_align_dictionary_temporality_marker_with_tagless_vector() .await, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let scan = |name, table: Arc| { LogicalPlanBuilder::scan(name, provider_as_source(table), None) diff --git a/src/query/src/query_engine.rs b/src/query/src/query_engine.rs index c512b9c1e5e..1c773ad099e 100644 --- a/src/query/src/query_engine.rs +++ b/src/query/src/query_engine.rs @@ -189,7 +189,7 @@ impl QueryEngineFactory { plugins: Plugins, options: QueryOptions, ) -> datafusion::error::Result { - let state = Arc::new(QueryEngineState::try_new( + Self::try_new_with_plugins_and_region_transport( catalog_manager, partition_rule_manager, region_query_handler, @@ -197,6 +197,40 @@ impl QueryEngineFactory { procedure_service_handler, flow_service_handler, with_dist_planner, + false, + plugins, + options, + ) + } + + /// [`QueryEngineFactory::try_new_with_plugins`], additionally stating how the engine reaches + /// the regions it queries. + /// + /// `in_process_region_transport` reaches `DistPlannerAnalyzer`: when it is `true`, a region + /// input may travel as a `LogicalPlan` instead of Substrait bytes, which only a region server in + /// the same process can accept. + #[allow(clippy::too_many_arguments)] + pub fn try_new_with_plugins_and_region_transport( + catalog_manager: CatalogManagerRef, + partition_rule_manager: Option, + region_query_handler: Option, + table_mutation_handler: Option, + procedure_service_handler: Option, + flow_service_handler: Option, + with_dist_planner: bool, + in_process_region_transport: bool, + plugins: Plugins, + options: QueryOptions, + ) -> datafusion::error::Result { + let state = Arc::new(QueryEngineState::try_new_with_region_transport( + catalog_manager, + partition_rule_manager, + region_query_handler, + table_mutation_handler, + procedure_service_handler, + flow_service_handler, + with_dist_planner, + in_process_region_transport, plugins.clone(), options, )?); diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index e2095c33530..0b37a9a4324 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -99,6 +99,10 @@ pub struct QueryEngineState { table_functions: Arc>>>, extension_rules: Vec>, plugins: Plugins, + experimental_native_promql: bool, + /// Whether this engine's regions are reached in process, which is what lets a region take a + /// `LogicalPlan` instead of Substrait bytes. + in_process_region_transport: bool, } impl fmt::Debug for QueryEngineState { @@ -147,6 +151,39 @@ impl QueryEngineState { with_dist_planner: bool, plugins: Plugins, options: QueryOptionsNew, + ) -> DfResult { + Self::try_new_with_region_transport( + catalog_list, + partition_rule_manager, + region_query_handler, + table_mutation_handler, + procedure_service_handler, + flow_service_handler, + with_dist_planner, + false, + plugins, + options, + ) + } + + /// [`QueryEngineState::try_new`], additionally stating how the engine reaches its regions. + /// + /// `in_process_region_transport` reaches [`DistPlannerAnalyzer`]: when it is `true` the + /// distributed planner may hand a region a `LogicalPlan` instead of Substrait bytes, which only + /// the in-process read path accepts. Every other caller builds an engine whose regions are + /// reached over the wire, which is what [`QueryEngineState::try_new`] states. + #[allow(clippy::too_many_arguments)] + pub(crate) fn try_new_with_region_transport( + catalog_list: CatalogManagerRef, + partition_rule_manager: Option, + region_query_handler: Option, + table_mutation_handler: Option, + procedure_service_handler: Option, + flow_service_handler: Option, + with_dist_planner: bool, + in_process_region_transport: bool, + plugins: Plugins, + options: QueryOptionsNew, ) -> DfResult { let total_memory = get_total_memory_bytes().max(0) as u64; let memory_pool_size = options.memory_pool_size.resolve(total_memory) as usize; @@ -209,7 +246,9 @@ impl QueryEngineState { )), ); if with_dist_planner { - analyzer.rules.push(Arc::new(DistPlannerAnalyzer)); + analyzer.rules.push(Arc::new(DistPlannerAnalyzer::new( + in_process_region_transport, + ))); analyzer.rules.push(Arc::new(JsonSchemaConcretizeRule)); } analyzer.rules.push(Arc::new(FixStateUdafOrderingAnalyzer)); @@ -297,6 +336,8 @@ impl QueryEngineState { table_functions: Arc::new(RwLock::new(HashMap::new())), extension_rules, plugins, + experimental_native_promql: options.experimental_native_promql, + in_process_region_transport, scalar_functions: Arc::new(RwLock::new(HashMap::new())), }) } @@ -445,6 +486,15 @@ impl QueryEngineState { self.df_context.register_udwf(func); } + pub(crate) fn experimental_native_promql(&self) -> bool { + self.experimental_native_promql + } + + /// Whether a region read is served in this process, so a region can take a `LogicalPlan`. + pub(crate) fn in_process_region_transport(&self) -> bool { + self.in_process_region_transport + } + pub fn catalog_manager(&self) -> &CatalogManagerRef { &self.catalog_manager } diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 486b1e8d6da..6e69daa00a2 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; +use std::collections::{BinaryHeap, HashSet}; use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; @@ -109,6 +109,66 @@ impl std::fmt::Debug for RegionScanExec { } } +/// Distribute [`PartitionRange`]s to each partition. +/// +/// Currently we assign ranges to partitions according to their rows so each partition +/// has similar number of rows. This method always return `expected_partition_num` partitions. +pub fn assign_partition_ranges( + mut ranges: Vec, + expected_partition_num: usize, +) -> Vec> { + if ranges.is_empty() { + // Returns a single partition with no range. + return vec![vec![]; expected_partition_num]; + } + + if ranges.len() == 1 { + let mut vec = vec![vec![]; expected_partition_num]; + vec[0] = ranges; + return vec; + } + + // Sort ranges by number of rows in descending order. + ranges.sort_by_key(|b| std::cmp::Reverse(b.num_rows)); + let mut partition_ranges = vec![vec![]; expected_partition_num]; + + #[derive(Eq, PartialEq)] + struct HeapNode { + num_rows: usize, + partition_idx: usize, + } + + impl Ord for HeapNode { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + // Reverse for min-heap. + self.num_rows.cmp(&other.num_rows).reverse() + } + } + + impl PartialOrd for HeapNode { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + + let mut part_heap = + BinaryHeap::from_iter((0..expected_partition_num).map(|partition_idx| HeapNode { + num_rows: 0, + partition_idx, + })); + + // Assigns the range to the partition with the smallest number of rows. + for range in ranges { + let mut node = part_heap.pop().unwrap(); + let partition_idx = node.partition_idx; + node.num_rows += range.num_rows; + partition_ranges[partition_idx].push(range); + part_heap.push(node); + } + + partition_ranges +} + impl RegionScanExec { pub fn new( scanner: RegionScannerRef, @@ -434,6 +494,7 @@ impl RegionScanExec { /// Resolves authorized logical names against this prepared scanner's physical metadata, /// then selects a native raw-series leaf and returns the binding it was validated against. + /// Prepared scans may contain multiple partitions; the producer keeps each series in one partition. /// /// Approval must be based on the returned [`RawSeriesRequest`], because that binding, not the /// requested names, describes the approved output. See [`Self::try_to_named_series_scan`] @@ -474,10 +535,6 @@ impl RegionScanExec { } let scanner = self.scanner.lock().unwrap(); let metadata = scanner.metadata(); - if scanner.properties().num_partitions().max(1) != 1 { - return None; - } - let primary_key_ids = metadata .primary_key_columns() .map(|column| column.column_id) @@ -1302,6 +1359,83 @@ mod test { ) } + #[test] + fn assign_partition_ranges_balances_rows_deterministically() { + let ranges: Vec<_> = [100, 200, 150, 250] + .into_iter() + .enumerate() + .map(|(identifier, num_rows)| PartitionRange { + start: common_time::Timestamp::new( + identifier as i64, + common_time::timestamp::TimeUnit::Second, + ), + end: common_time::Timestamp::new( + identifier as i64 + 1, + common_time::timestamp::TimeUnit::Second, + ), + num_rows, + identifier: identifier + 1, + }) + .collect(); + let assigned = assign_partition_ranges(ranges.clone(), 2); + assert_eq!( + assigned + .iter() + .map(|partition| { + partition + .iter() + .map(|range| range.identifier) + .collect::>() + }) + .collect::>(), + vec![vec![4, 1], vec![2, 3]] + ); + let mut assigned_identifiers = assigned + .iter() + .flatten() + .map(|range| range.identifier) + .collect::>(); + assigned_identifiers.sort_unstable(); + assert_eq!(assigned_identifiers, vec![1, 2, 3, 4]); + + let assigned = assign_partition_ranges(ranges.clone(), 5); + assert_eq!( + assigned + .iter() + .map(|partition| { + partition + .iter() + .map(|range| range.identifier) + .collect::>() + }) + .collect::>(), + vec![vec![4], vec![1], vec![2], vec![], vec![3]] + ); + let mut assigned_identifiers = assigned + .iter() + .flatten() + .map(|range| range.identifier) + .collect::>(); + assigned_identifiers.sort_unstable(); + assert_eq!(assigned_identifiers, vec![1, 2, 3, 4]); + + assert_eq!(assign_partition_ranges(vec![], 5), vec![vec![]; 5]); + } + + #[test] + fn named_series_binding_accepts_multiple_prepared_partitions() { + let (_, plan) = dynamic_filter_fixture(5687); + let prepared = plan + .with_new_partitions(vec![vec![], vec![]], 2) + .expect("prepare two scanner partitions"); + assert!(prepared.is_partition_set()); + assert!( + prepared + .resolve_named_series_binding("cpu", &[], &["a".to_string()]) + .is_some() + ); + } + #[test] fn named_series_scan_requires_prepared_single_partition() { let (_, plan) = dynamic_filter_fixture(5686); diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 074be282543..80fb6ed36ba 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -91,6 +91,7 @@ pub struct GreptimeDbStandaloneBuilder { logical_batcher: Option, table_batcher: BatcherOptions, mito_config: Option, + query_options: query::options::QueryOptions, } impl GreptimeDbStandaloneBuilder { @@ -112,9 +113,17 @@ impl GreptimeDbStandaloneBuilder { logical_batcher: None, table_batcher: BatcherOptions::default(), mito_config: None, + query_options: query::options::QueryOptions::default(), } } + /// Overrides query options for this standalone test instance. + #[must_use] + pub fn with_query_options(mut self, options: query::options::QueryOptions) -> Self { + self.query_options = options; + self + } + /// Overrides the Mito configuration for this test instance. #[must_use] pub fn with_mito_config(mut self, config: mito2::config::MitoConfig) -> Self { @@ -316,6 +325,7 @@ impl GreptimeDbStandaloneBuilder { procedure_executor.clone(), Arc::new(ProcessManager::new(server_addr, None)), ) + .with_in_process_region_transport() .with_local_file_access(local_file_access) .with_plugin(plugins.clone()) .try_build() @@ -389,6 +399,7 @@ impl GreptimeDbStandaloneBuilder { slow_query: self.slow_query_options.clone(), event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, + query: self.query_options.clone(), pending_rows_batcher: PendingRowsBatcherOptions { logical_table: self.logical_batcher.clone(), table: self.table_batcher.clone(), diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 7796be4d15b..65a49b83a28 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -12,32 +12,47 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use api::v1::value::ValueData; use api::v1::{ ColumnDataType, ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, }; +use common_base::Plugins; +use common_query::prelude::{GREPTIME_TEMPORALITY_DELTA, OTLP_AGGREGATION_TEMPORALITY_LABEL}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use common_query::{Output, OutputData}; use common_recordbatch::util::collect_batches; +use datafusion::physical_plan::ExecutionPlanProperties; +use datafusion::physical_plan::coop::CooperativeExec; use datatypes::arrow::array::{ Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, }; use datatypes::arrow::compute::cast; use datatypes::arrow::datatypes::DataType; use frontend::instance::Instance; -use query::datafusion::QUERY_PARALLELISM_HINT; +use query::datafusion::{QUERY_FALLBACK_HINT, QUERY_PARALLELISM_HINT}; +use query::options::QueryOptions; use query::parser::{PromQuery, QueryLanguageParser, QueryStatement}; +use query::physical_wrapper::{PhysicalPlanWrapper, PhysicalPlanWrapperRef}; +use query::promql::planner::NATIVE_PROMQL_EXTENSION_KEY; use rstest::rstest; use rstest_reuse::apply; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; +use sql::ast::{Ident, ObjectName, ObjectNamePart}; +use sql::statements::create::CreateView; +use sql::statements::statement::Statement as SqlStatement; +use sql::statements::tql::{Tql, TqlEval}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::region_engine::RawLabelSource; +use table::table::scan::SeriesScanExec; use super::test_util::{ both_instances_cases, check_unordered_output_stream, distributed, standalone, }; +use crate::standalone::GreptimeDbStandaloneBuilder; use crate::tests::test_util::MockInstance; #[allow(clippy::too_many_arguments)] @@ -1752,3 +1767,1114 @@ async fn f06_scalar_metric_table_lag_order_by_ts_desc(instance: Arc>>); + +impl NativePlanObserver { + fn clear(&self) { + *self.0.lock().unwrap() = None; + } + + fn take(&self) -> Arc { + self.0 + .lock() + .unwrap() + .take() + .expect("native fused plan was captured") + } + + /// The fused node the last wrapped plan contained, if any; `take` requires one instead. + fn fused(&self) -> Option> { + self.0.lock().unwrap().clone() + } +} + +impl PhysicalPlanWrapper for NativePlanObserver { + fn wrap( + &self, + origin: Arc, + _ctx: session::context::QueryContextRef, + ) -> Arc { + let mut stack = vec![origin.clone()]; + while let Some(node) = stack.pop() { + if node.name() == "FusedSeriesExec" { + *self.0.lock().unwrap() = Some(node); + break; + } + stack.extend(node.children().into_iter().cloned()); + } + origin + } +} + +fn native_query_context(enabled: bool, parallelism: usize) -> Arc { + let mut ctx = QueryContext::with_db_name(None); + ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string()); + ctx.set_extension(NATIVE_PROMQL_EXTENSION_KEY, enabled.to_string()); + Arc::new(ctx) +} + +async fn native_tql_output(instance: &Arc, sql: &str, enabled: bool) -> Output { + instance + .do_query(sql, native_query_context(enabled, 4)) + .await + .remove(0) + .unwrap() +} + +async fn native_output_text(output: Output) -> String { + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + _ => unreachable!("query returns record batches"), + }; + batches.pretty_print().unwrap() +} + +async fn native_result_rows(output: Output) -> Vec<(i64, String, f64)> { + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + _ => unreachable!("TQL EVAL returns record batches"), + }; + let mut rows = batches + .iter() + .flat_map(|batch| { + let timestamps = batch + .column_by_name("ts") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let hosts = cast( + batch.column_by_name("host").unwrap().as_ref(), + &DataType::Utf8, + ) + .unwrap(); + let hosts = hosts.as_any().downcast_ref::().unwrap(); + let value_columns = batch + .columns() + .iter() + .filter(|column| column.data_type() == &DataType::Float64) + .collect::>(); + assert_eq!(value_columns.len(), 1, "expected one Float64 value column"); + let values = value_columns[0] + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(timestamps.null_count(), 0); + assert_eq!(hosts.null_count(), 0); + assert_eq!(values.null_count(), 0); + (0..batch.num_rows()) + .map(|row| { + ( + timestamps.value(row), + hosts.value(row).to_string(), + values.value(row), + ) + }) + .collect::>() + }) + .collect::>(); + rows.sort_by(|left, right| (left.0, &left.1).cmp(&(right.0, &right.1))); + rows +} + +/// [`native_result_rows`] plus each row's OTLP temporality marker, sorted by label then time. +async fn native_marker_rows(output: Output) -> Vec<(i64, String, String, f64)> { + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + _ => unreachable!("TQL EVAL returns record batches"), + }; + let mut rows = Vec::new(); + for batch in batches.iter() { + let timestamps = batch + .column_by_name("ts") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let hosts = cast( + batch.column_by_name("host").unwrap().as_ref(), + &DataType::Utf8, + ) + .unwrap(); + let hosts = hosts.as_any().downcast_ref::().unwrap(); + let markers = cast( + batch + .column_by_name(OTLP_AGGREGATION_TEMPORALITY_LABEL) + .unwrap() + .as_ref(), + &DataType::Utf8, + ) + .unwrap(); + let markers = markers.as_any().downcast_ref::().unwrap(); + let values = batch + .columns() + .iter() + .filter(|column| column.data_type() == &DataType::Float64) + .collect::>(); + assert_eq!(values.len(), 1, "expected one Float64 value column"); + let values = values[0].as_any().downcast_ref::().unwrap(); + assert_eq!( + markers.null_count(), + 0, + "every row must keep its temporality label" + ); + assert_eq!(values.null_count(), 0); + for row in 0..batch.num_rows() { + rows.push(( + timestamps.value(row), + hosts.value(row).to_string(), + markers.value(row).to_string(), + values.value(row), + )); + } + } + rows.sort_by(|left, right| (left.1.as_str(), left.0).cmp(&(right.1.as_str(), right.0))); + rows +} + +/// Compares the whole answer against `expected`, in the same label order, key by key and value by +/// value; a duplicated or missing series therefore fails on the keys rather than on the count. +fn assert_delta_rows( + rows: &[(i64, String, String, f64)], + expected: &[(i64, &str, &str, f64)], + context: &str, +) { + let actual = rows + .iter() + .map(|(ts, host, marker, value)| (*ts, host.as_str(), marker.as_str(), *value)) + .collect::>(); + assert_eq!(actual.len(), expected.len(), "{context}: rows {actual:?}"); + for (actual, expected) in actual.iter().zip(expected) { + assert_eq!( + (actual.0, actual.1, actual.2), + (expected.0, expected.1, expected.2), + "{context}: row {actual:?}" + ); + assert!( + (actual.3 - expected.3).abs() <= 1e-9 * expected.3.abs().max(1.0), + "{context}: series {:?} must report {}, found {}", + actual.1, + expected.3, + actual.3 + ); + } +} + +fn assert_four_native_partitions(plan: &Arc) { + use datafusion::physical_plan::metrics::MetricValue; + + assert_eq!(plan.name(), "FusedSeriesExec"); + assert_eq!(plan.output_partitioning().partition_count(), 4); + assert_eq!( + plan.children()[0].output_partitioning().partition_count(), + 4, + "fused scan must have four partitions" + ); + let mut values_by_partition = [0_u64; 4]; + for metric in plan.metrics().expect("fused plan metrics").iter() { + if metric.value().name() == "native.root.output_values" + && let MetricValue::Count { count, .. } = metric.value() + && let Some(partition) = metric.partition() + && let Some(value) = values_by_partition.get_mut(partition) + { + *value += count.value() as u64; + } + } + assert!( + values_by_partition.iter().all(|value| *value > 0), + "every native execution partition must produce values: {values_by_partition:?}" + ); +} + +/// A standalone instance with the native route opted in, and an observer of the plans it runs. +async fn native_test_instance( + builder: GreptimeDbStandaloneBuilder, +) -> ( + crate::standalone::GreptimeDbStandalone, + Arc, +) { + native_test_instance_with_options( + builder, + QueryOptions { + experimental_native_promql: true, + ..Default::default() + }, + ) + .await +} + +/// [`native_test_instance`] with the engine-wide query options spelled out by the caller. +async fn native_test_instance_with_options( + builder: GreptimeDbStandaloneBuilder, + options: QueryOptions, +) -> ( + crate::standalone::GreptimeDbStandalone, + Arc, +) { + let observer = Arc::new(NativePlanObserver::default()); + let plugins = Plugins::new(); + let wrapper: PhysicalPlanWrapperRef = observer.clone(); + plugins.insert(wrapper); + let standalone = builder + .with_plugin(plugins) + .with_query_options(options) + .build() + .await; + (standalone, observer) +} + +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_executes_fused_node_and_matches_legacy() { + common_telemetry::init_default_ut_logging(); + // The fixture opts in so its single bulk insert is distributed across real TSID scan + // partitions; legacy may send its one batch only to partition 0. + let mito_config = mito2::config::MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }; + let builder = + GreptimeDbStandaloneBuilder::new("native_promql_fused_e2e").with_mito_config(mito_config); + let (standalone, observer) = native_test_instance(builder).await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + "CREATE TABLE native_e2e_phy (ts TIMESTAMP TIME INDEX, val DOUBLE) ENGINE=metric WITH (\"physical_metric_table\" = \"true\"); CREATE TABLE native_e2e (host STRING, ts TIMESTAMP TIME INDEX, val DOUBLE, PRIMARY KEY(host)) ENGINE=metric WITH (\"on_physical_table\" = 'native_e2e_phy')", + QueryContext::arc(), + ) + .await; + execute_all(&instance, "INSERT INTO native_e2e (host, ts, val) VALUES ('a', 0, 2), ('b', 0, 4), ('c', 0, 6), ('d', 0, 8), ('e', 0, 10), ('f', 0, 12), ('g', 0, 14), ('h', 0, 16), ('i', 0, 18), ('j', 0, 20), ('k', 0, 22), ('l', 0, 24), ('m', 0, 26), ('n', 0, 28), ('o', 0, 30), ('p', 0, 32)", QueryContext::arc()).await; + let query = PromQuery { + query: "native_e2e".to_string(), + start: "0".to_string(), + end: "0".to_string(), + step: "1s".to_string(), + ..PromQuery::default() + }; + observer.clear(); + let native = instance + .do_promql_query(&query, native_query_context(true, 4)) + .await + .remove(0) + .unwrap(); + let native_rows = native_result_rows(native).await; + let plan = observer.take(); + assert_four_native_partitions(&plan); + let expected = (b'a'..=b'p') + .zip((1..=16).map(|index| index as f64 * 2.0)) + .map(|(host, value)| (0, (host as char).to_string(), value)) + .collect::>(); + let legacy_rows = native_result_rows( + instance + .do_promql_query(&query, native_query_context(false, 4)) + .await + .remove(0) + .unwrap(), + ) + .await; + assert_eq!(native_rows, expected); + assert_eq!(native_rows, legacy_rows); +} + +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_rate_and_abs_rate_match_legacy_across_storage_encodings() { + common_telemetry::init_default_ut_logging(); + let mito_config = mito2::config::MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }; + let (standalone, observer) = native_test_instance( + GreptimeDbStandaloneBuilder::new("native_promql_rate_storage_encodings") + .with_mito_config(mito_config), + ) + .await; + let instance = standalone.fe_instance().clone(); + + for (table, ddl, flush_table, encoding, sparse) in [ + ( + "native_rate_dense", + "CREATE TABLE native_rate_dense (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) ENGINE=mito WITH ('primary_key_encoding' = 'dense')", + "native_rate_dense", + PrimaryKeyEncoding::Dense, + false, + ), + ( + "native_rate_sparse", + "CREATE TABLE native_rate_sparse_phy (ts TIMESTAMP(3) TIME INDEX, val DOUBLE) ENGINE=metric WITH ('physical_metric_table' = 'true', 'primary_key_encoding' = 'sparse'); CREATE TABLE native_rate_sparse (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) ENGINE=metric WITH ('on_physical_table' = 'native_rate_sparse_phy')", + "native_rate_sparse_phy", + PrimaryKeyEncoding::Sparse, + true, + ), + ] { + execute_all(&instance, ddl, QueryContext::arc()).await; + for (start, end) in [(0, 60), (120, 300)] { + let values = (b'a'..=b'p') + .flat_map(|host| { + (start..=end).step_by(60).map(move |seconds| { + format!( + "('{}', {}, {})", + host as char, + seconds * 1000, + seconds as f64 * (host - b'a' + 1) as f64 + ) + }) + }) + .collect::>() + .join(", "); + execute_all( + &instance, + &format!("INSERT INTO {table} (host, ts, val) VALUES {values}"), + QueryContext::arc(), + ) + .await; + if start == 0 { + execute_all( + &instance, + &format!("ADMIN FLUSH_TABLE('{flush_table}')"), + QueryContext::arc(), + ) + .await; + } + } + + let expected = (b'a'..=b'p') + .map(|host| { + ( + 300_000, + (host as char).to_string(), + (host - b'a' + 1) as f64, + ) + }) + .collect::>(); + for expression in [ + format!("rate({table}[5m])"), + format!("abs(rate({table}[5m]))"), + ] { + let sql = format!("TQL EVAL (300, 300, '1m') {expression}"); + let legacy_rows = + native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!( + legacy_rows, expected, + "native-off `{expression}` on {table}" + ); + + observer.clear(); + let native_rows = + native_result_rows(native_tql_output(&instance, &sql, true).await).await; + assert_eq!(native_rows, expected, "native-on `{expression}` on {table}"); + assert_eq!(native_rows, legacy_rows, "`{expression}` on {table}"); + let plan = observer.take(); + assert_eq!(plan.name(), "FusedSeriesExec"); + assert_eq!( + plan.children()[0].output_partitioning().partition_count(), + plan.output_partitioning().partition_count(), + "fused child and output partitions agree for {table}" + ); + assert!( + plan.output_partitioning().partition_count() > 1, + "fused execution must retain partitioning for {table}" + ); + if sparse { + assert_four_native_partitions(&plan); + } + let child = plan.children()[0]; + let child = child + .downcast_ref::() + .map_or(child, |cooperative| cooperative.input()); + let series_scan = child + .downcast_ref::() + .expect("fused node reads a prepared series scan"); + let RawLabelSource::PrimaryKey { + physical_metadata, .. + } = series_scan.label_source() + else { + panic!("fused scan must materialize labels from the primary key"); + }; + assert_eq!(physical_metadata.primary_key_encoding, encoding); + } + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_explain_and_analyze_execute_fused_plan() { + common_telemetry::init_default_ut_logging(); + let (standalone, observer) = native_test_instance(GreptimeDbStandaloneBuilder::new( + "native_promql_explain_analyze_e2e", + )) + .await; + let instance = standalone.fe_instance().clone(); + execute_all(&instance, "CREATE TABLE native_explain (ts TIMESTAMP TIME INDEX, host STRING, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_explain VALUES (0, 'a', 2)", QueryContext::arc()).await; + for (statement, enabled) in [ + ("TQL EXPLAIN (0, 0, '1s') native_explain", true), + ("TQL EXPLAIN VERBOSE (0, 0, '1s') native_explain", true), + ("TQL ANALYZE (0, 0, '1s') native_explain", true), + ("TQL ANALYZE VERBOSE (0, 0, '1s') native_explain", true), + ("TQL EXPLAIN (0, 0, '1s') native_explain", false), + ] { + observer.clear(); + let output = native_tql_output(&instance, statement, enabled).await; + let text = native_output_text(output).await; + if statement.contains("ANALYZE") && enabled { + assert!(text.contains("FusedSeriesExec")); + assert!(text.contains("output_values")); + let plan = observer.take(); + let output_values = plan + .metrics() + .expect("fused plan metrics") + .iter() + .filter_map(|metric| match metric.value() { + datafusion::physical_plan::metrics::MetricValue::Count { name, count } + if name.as_ref() == "native.root.output_values" => + { + Some(count.value()) + } + _ => None, + }) + .sum::(); + assert_eq!(output_values, 1, "ANALYZE executes the fused series"); + } else if statement.contains("EXPLAIN") && enabled { + assert!(text.contains("FusedSeries: native_explain")); + assert!(text.contains("MergeScanExec")); + } else if !enabled { + assert!(!text.contains("FusedSeries:")); + } + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_view_fallback_preserves_native_hint() { + common_telemetry::init_default_ut_logging(); + let (standalone, observer) = native_test_instance(GreptimeDbStandaloneBuilder::new( + "native_promql_view_fallback_e2e", + )) + .await; + let instance = standalone.fe_instance().clone(); + execute_all(&instance, "CREATE TABLE native_view_source (ts TIMESTAMP TIME INDEX, host STRING, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_view_source VALUES (0, 'a', 2)", QueryContext::arc()).await; + let ctx = native_query_context(true, 4); + let expected_extensions = ctx.extensions(); + // Build the TQL view programmatically; the SQL syntax and SHOW CREATE round trip belong to the + // independent view-persistence PR. + let create_view = SqlStatement::CreateView(CreateView { + name: ObjectName(vec![ObjectNamePart::Identifier(Ident::new("native_view"))]), + columns: vec![], + query: Box::new(SqlStatement::Tql(Tql::Eval(TqlEval { + start: "0".to_string(), + end: "0".to_string(), + step: "1s".to_string(), + lookback: Some("5m".to_string()), + query: "native_view_source".to_string(), + alias: None, + }))), + or_replace: false, + if_not_exists: false, + }); + instance + .statement_executor() + .execute_stmt(QueryStatement::Sql(create_view), ctx.clone()) + .await + .unwrap(); + // The view's stored plan must not depend on the caller's native hint, so reading it runs the + // plan the view was created with rather than one that needs the hint to be readable. + observer.clear(); + let view = instance + .do_query("SELECT * FROM native_view", ctx.clone()) + .await + .remove(0) + .unwrap(); + let view_rows = native_result_rows(view).await; + assert_eq!(ctx.extensions(), expected_extensions); + assert_eq!(view_rows, vec![(0, "a".to_string(), 2.0)]); + assert!( + observer.fused().is_none(), + "a view plan must be readable without the native hint" + ); + // The hint is still the caller's: the same context keeps folding a direct query. + observer.clear(); + let direct = instance + .do_query("TQL EVAL (0, 0, '1s') native_view_source", ctx) + .await + .remove(0) + .unwrap(); + let direct_rows = native_result_rows(direct).await; + assert_eq!(direct_rows, vec![(0, "a".to_string(), 2.0)]); + assert_eq!(observer.take().name(), "FusedSeriesExec"); +} + +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_metric_logical_tables_are_isolated() { + common_telemetry::init_default_ut_logging(); + let (standalone, observer) = native_test_instance(GreptimeDbStandaloneBuilder::new( + "native_promql_metric_isolation_e2e", + )) + .await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + r#"CREATE TABLE native_metric_phy ( + ts TIMESTAMP TIME INDEX, + val DOUBLE + ) ENGINE=metric WITH ("physical_metric_table" = "true"); + CREATE TABLE native_metric_a ( + host STRING, + ts TIMESTAMP TIME INDEX, + val DOUBLE, + PRIMARY KEY (host) + ) ENGINE=metric WITH ("on_physical_table" = 'native_metric_phy'); + CREATE TABLE native_metric_b ( + host STRING, + ts TIMESTAMP TIME INDEX, + val DOUBLE, + PRIMARY KEY (host) + ) ENGINE=metric WITH ("on_physical_table" = 'native_metric_phy'); + INSERT INTO native_metric_a (host, ts, val) VALUES ('only_a', 0, 101), ('only_a', 1000, 102); + INSERT INTO native_metric_b (host, ts, val) VALUES ('only_b', 0, 202), ('only_b', 1000, 204);"#, + QueryContext::arc(), + ) + .await; + for (table, host, value) in [ + ("native_metric_a", "only_a", 1.0), + ("native_metric_b", "only_b", 2.0), + ] { + let sql = format!("TQL EVAL (1, 1, '1s') rate({table}[2s])"); + observer.clear(); + let native = native_tql_output(&instance, &sql, true).await; + let native_rows = native_result_rows(native).await; + assert_eq!(observer.take().name(), "FusedSeriesExec"); + let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!(native_rows, vec![(1000, host.to_string(), value)]); + assert_eq!(native_rows, legacy_rows); + } +} + +/// OTLP delta temporality on the public TQL route: a `delta`-marked series answers the raw sample +/// sum (10 + 10 = 20) while a cumulative series with the same samples answers 0. The fused node +/// computes one cumulative window for every series, so it must decline a marker-bearing table +/// instead of reporting 0 for the delta series. +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_delta_temporality_rate_and_increase_match_legacy() { + common_telemetry::init_default_ut_logging(); + const TABLE: &str = "native_delta_temporality"; + const CONTROL: &str = "native_delta_control"; + // Samples at 60s and 120s, evaluated at 120s over a 2m window: the `delta` series' samples are + // raw deltas, so its increase is their sum 10 + 10; the `cumulative` one did not rise, so 0. + const DELTA_INCREASE: f64 = 20.0; + const CUMULATIVE_INCREASE: f64 = 0.0; + const DELTA_RATE: f64 = 20.0 / 120.0; + const CUMULATIVE_RATE: f64 = 0.0; + + let (standalone, observer) = native_test_instance(GreptimeDbStandaloneBuilder::new( + "native_promql_delta_temporality", + )) + .await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + &format!( + r#"CREATE TABLE {TABLE} ( + host STRING, + ts TIMESTAMP(3) TIME INDEX, + val DOUBLE, + otlp_aggregation_temporality STRING, + PRIMARY KEY (host, otlp_aggregation_temporality) + ) ENGINE=mito; + INSERT INTO {TABLE} (host, ts, val, otlp_aggregation_temporality) VALUES + ('delta', 60000, 10, '{GREPTIME_TEMPORALITY_DELTA}'), + ('delta', 120000, 10, '{GREPTIME_TEMPORALITY_DELTA}'), + ('cumulative', 60000, 10, 'cumulative'), + ('cumulative', 120000, 10, 'cumulative'); + CREATE TABLE {CONTROL} (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) ENGINE=mito; + INSERT INTO {CONTROL} (host, ts, val) VALUES ('a', 60000, 10), ('a', 120000, 30);"# + ), + QueryContext::arc(), + ) + .await; + + for (query, delta, cumulative) in [ + ( + format!("increase({TABLE}[2m])"), + DELTA_INCREASE, + CUMULATIVE_INCREASE, + ), + (format!("rate({TABLE}[2m])"), DELTA_RATE, CUMULATIVE_RATE), + ( + format!("abs(increase({TABLE}[2m]))"), + DELTA_INCREASE, + CUMULATIVE_INCREASE, + ), + ( + format!("abs(rate({TABLE}[2m]))"), + DELTA_RATE, + CUMULATIVE_RATE, + ), + ] { + let sql = format!("TQL EVAL (120, 120, '1m') {query}"); + // Both series are in the window, in label order, with their own marker and value. + let expected = [ + (120_000, "cumulative", "cumulative", cumulative), + (120_000, "delta", GREPTIME_TEMPORALITY_DELTA, delta), + ]; + // Values before routes: a wrong number is the regression, a declined plan is not. + let legacy_rows = native_marker_rows(native_tql_output(&instance, &sql, false).await).await; + assert_delta_rows(&legacy_rows, &expected, &format!("native-off `{query}`")); + observer.clear(); + let native_rows = native_marker_rows(native_tql_output(&instance, &sql, true).await).await; + assert_delta_rows(&native_rows, &expected, &format!("native-on `{query}`")); + assert_eq!( + native_rows, legacy_rows, + "`{query}`: same rows on both routes" + ); + assert!( + observer.fused().is_none(), + "`{query}`: the native route must decline a table that carries the temporality marker" + ); + } + + // The control: a marker-free table (samples 10 and 30 here, unlike the table above) still fuses + // and agrees with the existing plan, so the decline is about the marker, not the table shape. + let control_sql = format!("TQL EVAL (120, 120, '1m') sum_over_time({CONTROL}[2m])"); + observer.clear(); + let control_native = + native_result_rows(native_tql_output(&instance, &control_sql, true).await).await; + assert_eq!( + control_native, + vec![(120_000, "a".to_string(), 40.0)], + "the fused control reports the sampled sum" + ); + assert_eq!(observer.take().name(), "FusedSeriesExec"); + let control_legacy = + native_result_rows(native_tql_output(&instance, &control_sql, false).await).await; + assert_eq!( + control_native, control_legacy, + "the control agrees with the existing plan" + ); +} + +/// Enabling query fallback — by the engine option or by the per-request `query_fallback` hint — must +/// keep an otherwise fusible query on the existing plan without changing its answer. +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_fallback_options_keep_expected_answer() { + common_telemetry::init_default_ut_logging(); + const TABLE: &str = "native_fallback_control"; + // One eligible selector, so only the fallback settings decide the route. `sum_over_time` is + // exact over the two samples in the window: 10 + 30 for `a`, 10 + 70 for `b`. + let query = format!("sum_over_time({TABLE}[2m])"); + let expected = vec![ + (120_000, "a".to_string(), 40.0), + (120_000, "b".to_string(), 80.0), + ]; + for (tag, allow_query_fallback) in [("engine_off", false), ("engine_on", true)] { + let (standalone, observer) = native_test_instance_with_options( + GreptimeDbStandaloneBuilder::new(&format!("native_promql_fallback_{tag}")), + QueryOptions { + experimental_native_promql: true, + allow_query_fallback, + ..Default::default() + }, + ) + .await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + &format!( + "CREATE TABLE {TABLE} (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) ENGINE=mito; \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 60000, 10), ('a', 120000, 30), \ + ('b', 60000, 10), ('b', 120000, 70);" + ), + QueryContext::arc(), + ) + .await; + let mut hinted = native_query_context(true, 4).as_ref().clone(); + hinted.set_extension(QUERY_FALLBACK_HINT, "true"); + for (route, ctx) in [ + ("plain", native_query_context(true, 4)), + ("fallback-hint", Arc::new(hinted)), + ] { + let sql = format!("TQL EVAL (120, 120, '1m') {query}"); + observer.clear(); + let rows = + native_result_rows(instance.do_query(&sql, ctx).await.remove(0).unwrap()).await; + assert_eq!( + rows, expected, + "{tag} {route}: the answer must not depend on the fallback settings" + ); + if allow_query_fallback || route == "fallback-hint" { + // A fallback rewrites every table scan into `MergeScan`, which the fused node + // cannot read, so a query allowed to fall back keeps the existing plan. + assert!( + observer.fused().is_none(), + "{tag} {route}: enabling the fallback must keep the existing plan" + ); + } else { + assert_eq!( + observer.take().name(), + "FusedSeriesExec", + "{tag} {route}: with no fallback enabled this query must be fused" + ); + } + } + } +} + +/// A store configured for the columnar series format must not change the answer of a query over a +/// field that format cannot serve: the columnar samples carry `Float64` series only, so an integer +/// field has no raw frame to fold and the route has to decline it while the existing plan answers. +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_columnar_config_falls_back_for_integer_field() { + common_telemetry::init_default_ut_logging(); + const TABLE: &str = "native_columnar_int"; + // Both functions are exact over the two samples inside the window: 10 + 30 and their mean. + const SUM: f64 = 40.0; + const AVG: f64 = 20.0; + let mito_config = mito2::config::MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_format: mito2::config::SeriesIndexFormat::Columnar, + ..Default::default() + }; + let (standalone, observer) = native_test_instance( + GreptimeDbStandaloneBuilder::new("native_promql_columnar_fallback") + .with_mito_config(mito_config), + ) + .await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + &format!( + "CREATE TABLE {TABLE} (host STRING, ts TIMESTAMP(3) TIME INDEX, val BIGINT, PRIMARY KEY (host)) ENGINE=mito; \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 60000, 10), ('a', 120000, 30);" + ), + QueryContext::arc(), + ) + .await; + for (query, expected) in [ + (format!("sum_over_time({TABLE}[2m])"), SUM), + (format!("avg_over_time({TABLE}[2m])"), AVG), + ] { + let sql = format!("TQL EVAL (120, 120, '1m') {query}"); + observer.clear(); + let native_rows = native_result_rows(native_tql_output(&instance, &sql, true).await).await; + assert_eq!( + native_rows.len(), + 1, + "native-on `{query}`: one row for the one series" + ); + assert_eq!(native_rows[0].0, 120_000, "native-on `{query}`: ts"); + assert_eq!(native_rows[0].1, "a", "native-on `{query}`: label"); + assert!( + (native_rows[0].2 - expected).abs() <= 1e-9 * expected.abs().max(1.0), + "native-on `{query}`: expected {expected}, found {}", + native_rows[0].2 + ); + let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!( + native_rows, legacy_rows, + "`{query}`: same rows on both routes" + ); + assert!( + observer.fused().is_none(), + "`{query}`: the native route must decline a field the columnar samples cannot serve" + ); + } +} + +/// The positive counterpart of the fallback above: a sparse physical metric table whose series +/// index is published in the columnar format fuses a `Float64` selector and matches the old plan. +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_columnar_fused_series_matches_legacy() { + common_telemetry::init_default_ut_logging(); + const PHYSICAL: &str = "native_columnar_fused_phy"; + const TABLE: &str = "native_columnar_fused"; + let mito_config = mito2::config::MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_format: mito2::config::SeriesIndexFormat::Columnar, + // Slow the periodic pass, so the manual build below is what publishes the catalog. + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + ..Default::default() + }; + let builder = GreptimeDbStandaloneBuilder::new("native_promql_columnar_fused") + .with_mito_config(mito_config); + let (standalone, observer) = native_test_instance(builder).await; + let instance = standalone.fe_instance().clone(); + // One flushed SST is enough for a manual build, which skips the periodic four-file trigger. + execute_all( + &instance, + &format!( + "CREATE TABLE {PHYSICAL} (ts TIMESTAMP(3) TIME INDEX, val DOUBLE) ENGINE=metric \ + WITH ('physical_metric_table' = 'true', 'primary_key_encoding' = 'sparse'); \ + CREATE TABLE {TABLE} (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) \ + ENGINE=metric WITH ('on_physical_table' = '{PHYSICAL}'); \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 60000, 10), ('a', 120000, 30), \ + ('b', 60000, 20), ('b', 120000, 60); \ + ADMIN FLUSH_TABLE('{PHYSICAL}'); \ + ADMIN BUILD_SERIES_INDEX('{PHYSICAL}');" + ), + QueryContext::arc(), + ) + .await; + // Both samples of each series fall in the 2m window: 10 + 30 and 20 + 60. + let expected = vec![ + (120_000, "a".to_string(), 40.0), + (120_000, "b".to_string(), 80.0), + ]; + let sql = format!("TQL EVAL (120, 120, '1m') sum_over_time({TABLE}[2m])"); + observer.clear(); + let native_rows = native_result_rows(native_tql_output(&instance, &sql, true).await).await; + // The answer is drained first, so the counters below describe execution, not planning. + let plan = observer.take(); + assert_eq!(native_rows, expected, "the fused route reports the sums"); + // `EnsureCooperative` wraps a leaf that does not ask for cooperative scheduling, and the fused + // node reads the same series scan through it. + let child = plan.children()[0]; + let child = child + .downcast_ref::() + .map_or(child, |cooperative| cooperative.input()); + let series_scan = child + .downcast_ref::() + .expect("the fused node reads a prepared series scan"); + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags, + } = series_scan.label_source() + else { + panic!("the fused scan must materialize labels from the primary key"); + }; + assert_eq!(metric_name.as_ref(), TABLE); + assert_eq!( + logical_tags + .iter() + .map(|tag| tag.name.as_ref()) + .collect::>(), + vec!["host"], + "the logical tags are the primary-key columns of the scanned region" + ); + for tag in logical_tags.iter() { + let column = physical_metadata.column_by_id(tag.column_id).unwrap(); + assert_eq!(column.column_schema.name, "host"); + assert_eq!(column.semantic_type, SemanticType::Tag); + } + assert_eq!( + physical_metadata.primary_key_encoding, + PrimaryKeyEncoding::Sparse + ); + let series_schema = series_scan.series_schema(); + assert_eq!(series_schema.fields.len(), 1); + assert_eq!(series_schema.fields[0].name.as_ref(), "val"); + // Only the columnar writer publishes these formats, so the scanned region's catalog rules out + // a configuration-only run. + let samples = std::path::Path::new(&standalone.opts.storage.data_home) + .join("series_index") + .join(physical_metadata.region_id.as_u64().to_string()) + .join("samples.json"); + let catalog: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(samples).unwrap()).unwrap(); + let entries = catalog["entries"].as_array().unwrap(); + let series = catalog["series"].as_array().unwrap(); + assert!(!entries.is_empty() && !series.is_empty(), "published"); + assert!( + entries + .iter() + .all(|entry| entry["artifact"]["format"] == "columnar-samples") + ); + assert!( + series + .iter() + .all(|entry| entry["artifact"]["format"] == "columnar-series") + ); + // One value per series: the node ran, rather than being merely planned. + let output_values = plan + .metrics() + .expect("fused plan metrics") + .iter() + .filter_map(|metric| match metric.value() { + datafusion::physical_plan::metrics::MetricValue::Count { name, count } + if name.as_ref() == "native.root.output_values" => + { + Some(count.value()) + } + _ => None, + }) + .sum::(); + assert_eq!(output_values, 2, "both series must report one value"); + // Native off: the same rows must come from the existing plan. + observer.clear(); + let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!(native_rows, legacy_rows, "same rows on both routes"); + assert!(observer.fused().is_none(), "the legacy route must not fuse"); +} + +/// No published series index: the columnar `Float64` selector must answer from the raw memtable +/// and from an unindexed SST instead of failing to plan. +#[rstest] +#[case::memtable("memtable", false)] +#[case::unindexed_sst("unindexed_sst", true)] +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_columnar_unindexed_float64_matches_legacy( + #[case] stage: &str, + #[case] flush: bool, +) { + common_telemetry::init_default_ut_logging(); + const PHYSICAL: &str = "native_columnar_unindexed_phy"; + const TABLE: &str = "native_columnar_unindexed"; + let mito_config = mito2::config::MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_format: mito2::config::SeriesIndexFormat::Columnar, + // No periodic build, so neither case ever gets an index. + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + ..Default::default() + }; + let (standalone, observer) = native_test_instance( + GreptimeDbStandaloneBuilder::new(&format!("native_promql_columnar_{stage}")) + .with_mito_config(mito_config), + ) + .await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + &format!( + "CREATE TABLE {PHYSICAL} (ts TIMESTAMP(3) TIME INDEX, val DOUBLE) ENGINE=metric \ + WITH ('physical_metric_table' = 'true', 'primary_key_encoding' = 'sparse'); \ + CREATE TABLE {TABLE} (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) \ + ENGINE=metric WITH ('on_physical_table' = '{PHYSICAL}'); \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 60000, 10), ('a', 120000, 30), \ + ('b', 60000, 20), ('b', 120000, 60);" + ), + QueryContext::arc(), + ) + .await; + if flush { + execute_all( + &instance, + &format!("ADMIN FLUSH_TABLE('{PHYSICAL}');"), + QueryContext::arc(), + ) + .await; + } + // Both samples of each series fall in the 2m window: 10 + 30 and 20 + 60. + let expected = vec![ + (120_000, "a".to_string(), 40.0), + (120_000, "b".to_string(), 80.0), + ]; + let sql = format!("TQL EVAL (120, 120, '1m') sum_over_time({TABLE}[2m])"); + // Legacy off first, so a red native query is about the fused route. + observer.clear(); + let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!(legacy_rows, expected, "{stage}: existing plan must answer"); + assert!(observer.fused().is_none(), "{stage}: legacy must not fuse"); + observer.clear(); + let native_rows = native_result_rows(native_tql_output(&instance, &sql, true).await).await; + // Drained first, so the counter below describes execution, not planning; `take` requires the + // fused node to have run. + let plan = observer.take(); + assert_eq!(native_rows, expected, "{stage}: fused route must answer"); + let output_values = plan + .metrics() + .expect("fused plan metrics") + .iter() + .filter_map(|metric| match metric.value() { + datafusion::physical_plan::metrics::MetricValue::Count { name, count } + if name.as_ref() == "native.root.output_values" => + { + Some(count.value()) + } + _ => None, + }) + .sum::(); + assert_eq!(output_values, 2, "{stage}: two values, one per series"); + assert_eq!( + native_rows, legacy_rows, + "{stage}: same rows on both routes" + ); +} + +/// An index over the first SST only: the newer rows still in the memtable, then flushed into a +/// second unindexed SST, must answer the same sums. +#[rstest] +#[case::partial_memtable("partial_memtable", false)] +#[case::partial_sst("partial_sst", true)] +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_columnar_partial_index_float64_matches_legacy( + #[case] stage: &str, + #[case] flush: bool, +) { + common_telemetry::init_default_ut_logging(); + const PHYSICAL: &str = "native_columnar_partial_phy"; + const TABLE: &str = "native_columnar_partial"; + let mito_config = mito2::config::MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_format: mito2::config::SeriesIndexFormat::Columnar, + // Slow the periodic pass, so only the manual build below publishes an index. + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + ..Default::default() + }; + let (standalone, observer) = native_test_instance( + GreptimeDbStandaloneBuilder::new(&format!("native_promql_columnar_{stage}")) + .with_mito_config(mito_config), + ) + .await; + let instance = standalone.fe_instance().clone(); + // The first samples are flushed and indexed before the second ones are inserted, so the index + // covers one SST while the later rows start out unflushed. + execute_all( + &instance, + &format!( + "CREATE TABLE {PHYSICAL} (ts TIMESTAMP(3) TIME INDEX, val DOUBLE) ENGINE=metric \ + WITH ('physical_metric_table' = 'true', 'primary_key_encoding' = 'sparse'); \ + CREATE TABLE {TABLE} (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) \ + ENGINE=metric WITH ('on_physical_table' = '{PHYSICAL}'); \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 60000, 10), ('b', 60000, 20); \ + ADMIN FLUSH_TABLE('{PHYSICAL}'); \ + ADMIN BUILD_SERIES_INDEX('{PHYSICAL}'); \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 120000, 30), ('b', 120000, 60);" + ), + QueryContext::arc(), + ) + .await; + if flush { + // The new SST stays unindexed: the build above is not repeated. + execute_all( + &instance, + &format!("ADMIN FLUSH_TABLE('{PHYSICAL}');"), + QueryContext::arc(), + ) + .await; + } + // Both samples of each series fall in the 2m window: 10 + 30 and 20 + 60. + let expected = vec![ + (120_000, "a".to_string(), 40.0), + (120_000, "b".to_string(), 80.0), + ]; + let sql = format!("TQL EVAL (120, 120, '1m') sum_over_time({TABLE}[2m])"); + // Legacy off first, so a red native query is about the fused route. + observer.clear(); + let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!(legacy_rows, expected, "{stage}: existing plan must answer"); + assert!(observer.fused().is_none(), "{stage}: legacy must not fuse"); + observer.clear(); + let native_rows = native_result_rows(native_tql_output(&instance, &sql, true).await).await; + // Drained first, so the counter below describes execution, not planning; `take` requires the + // fused node to have run. + let plan = observer.take(); + assert_eq!(native_rows, expected, "{stage}: fused route must answer"); + let output_values = plan + .metrics() + .expect("fused plan metrics") + .iter() + .filter_map(|metric| match metric.value() { + datafusion::physical_plan::metrics::MetricValue::Count { name, count } + if name.as_ref() == "native.root.output_values" => + { + Some(count.value()) + } + _ => None, + }) + .sum::(); + assert_eq!(output_values, 2, "{stage}: two values, one per series"); + assert_eq!( + native_rows, legacy_rows, + "{stage}: same rows on both routes" + ); +}