From c96cdd0026bfda6bb1c1dc0d71d2ada192275339 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 26 Oct 2026 08:54:02 +0000 Subject: [PATCH] perf: stream and prune SST manifest scans (#9931) * refactor: add reusable metadata pruning and projected SST batches Signed-off-by: WenyXu * perf(mito2): stream and prune manifest SST metadata Signed-off-by: WenyXu * perf(datanode): stream inspection queries through reusable providers Signed-off-by: WenyXu * perf(catalog): push manifest filters through inspection queries Signed-off-by: WenyXu * refactor: align manifest column constants with module conventions Signed-off-by: WenyXu * perf: reduce manifest inspection allocations Signed-off-by: WenyXu * refactor: narrow manifest inspection API visibility Signed-off-by: WenyXu * refactor: retain residual filters for manifest inspection Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/catalog/src/information_extension.rs | 119 ++++++ .../system_schema/information_schema/ssts.rs | 7 + src/common/recordbatch/src/filter.rs | 38 +- src/common/recordbatch/src/filter/tests.rs | 88 ++++ src/datanode/src/region_server/catalog.rs | 18 +- .../src/region_server/catalog/inspect.rs | 311 ++++++++++++++ .../region_server/catalog/tests/manifest.rs | 257 ++++++++++++ src/metric-engine/src/engine/bulk_insert.rs | 11 +- src/mito2/src/engine.rs | 116 +++++- src/mito2/src/engine/basic_test.rs | 244 +++++++++++ src/mito2/src/engine/staging_test.rs | 98 +++++ src/mito2/src/region.rs | 355 ++++++++++++++-- src/store-api/src/sst_entry.rs | 381 ++++++++++++------ .../information_schema/ssts_limit.result | 53 +++ .../information_schema/ssts_limit.sql | 22 + .../common/information_schema/ssts.result | 78 ++++ .../common/information_schema/ssts.sql | 35 ++ 17 files changed, 2054 insertions(+), 177 deletions(-) create mode 100644 src/datanode/src/region_server/catalog/inspect.rs create mode 100644 src/datanode/src/region_server/catalog/tests/manifest.rs diff --git a/src/catalog/src/information_extension.rs b/src/catalog/src/information_extension.rs index 9d55f9f7a1e..f22c169df13 100644 --- a/src/catalog/src/information_extension.rs +++ b/src/catalog/src/information_extension.rs @@ -310,3 +310,122 @@ impl InformationExtension for DistributedInformationExtension { )))) } } + +#[cfg(test)] +mod tests { + use common_meta::node_manager::{DatanodeManager, DatanodeRef}; + use common_meta::peer::Peer; + use datafusion::common::tree_node::TreeNode; + use datafusion::logical_expr::LogicalPlan; + use datafusion::physical_plan::sorts::sort::SortExec; + use meta_client::client::MetaClientBuilder; + use store_api::sst_entry::ManifestSstEntry; + use store_api::storage::ScanRequest; + + use super::*; + use crate::information_schema::DatanodeInspectKind; + + struct UnusedDatanodeManager; + + #[async_trait::async_trait] + impl DatanodeManager for UnusedDatanodeManager { + async fn datanode(&self, _node: &Peer) -> DatanodeRef { + panic!("plan tests must not contact datanodes") + } + } + + #[derive(Debug)] + struct ManifestScanProvider; + + #[async_trait::async_trait] + impl datafusion::catalog::TableProvider for ManifestScanProvider { + fn schema(&self) -> ArrowSchemaRef { + ManifestSstEntry::schema().arrow_schema().clone() + } + + fn table_type(&self) -> datafusion::logical_expr::TableType { + datafusion::logical_expr::TableType::View + } + + fn supports_filters_pushdown( + &self, + filters: &[&datafusion::logical_expr::Expr], + ) -> Result> { + Ok(vec![datafusion::logical_expr::TableProviderFilterPushDown::Inexact; filters.len()]) + } + + async fn scan( + &self, + _state: &dyn datafusion::catalog::Session, + projection: Option<&Vec>, + filters: &[datafusion::logical_expr::Expr], + limit: Option, + ) -> Result> { + let schema = ManifestSstEntry::schema(); + let schema = match projection { + Some(projection) => Arc::new(schema.try_project(projection).unwrap()), + None => schema, + }; + Ok(Arc::new(DistributedInspectExec::new( + DatanodeInspectRequest { + kind: DatanodeInspectKind::SstManifest, + scan: ScanRequest { + projection: projection.cloned(), + filters: filters.to_vec(), + limit, + ..Default::default() + }, + }, + schema, + Arc::new(MetaClientBuilder::new(0, api::v1::meta::Role::Frontend).build()), + Arc::new(UnusedDatanodeManager), + ))) + } + } + + #[tokio::test] + async fn test_sql_manifest_sort_stays_on_frontend() { + let ctx = datafusion::prelude::SessionContext::new(); + ctx.register_table("ssts_manifest", Arc::new(ManifestScanProvider)) + .unwrap(); + for sql in [ + "SELECT region_id, file_id, min_ts, max_ts, file_size FROM ssts_manifest WHERE region_id = 42 ORDER BY min_ts LIMIT 100", + "SELECT file_id FROM ssts_manifest WHERE region_id = 42 ORDER BY max_ts DESC LIMIT 10 OFFSET 3", + "SELECT min_ts AS ts, file_id FROM ssts_manifest WHERE region_id = 42 ORDER BY ts DESC NULLS FIRST LIMIT 10", + "SELECT file_id FROM ssts_manifest WHERE region_id = 42 ORDER BY min_ts, max_ts DESC LIMIT 10", + ] { + let plan = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let mut found_scan = false; + let mut found_sort = false; + plan.apply(|plan| { + found_sort |= plan.is::(); + if let Some(scan) = plan.downcast_ref::() { + found_scan = true; + // All matching rows must reach the frontend global sort. + assert_eq!(scan.request.scan.limit, None); + let remote = scan.request.clone().build_plan().unwrap(); + remote.apply(|node| { + assert!(!matches!( + node, + LogicalPlan::Sort(_) | LogicalPlan::Limit(_) + )); + Ok(TreeNodeRecursion::Continue) + })?; + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + assert!( + found_scan && found_sort, + "{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(true) + ); + } + } +} diff --git a/src/catalog/src/system_schema/information_schema/ssts.rs b/src/catalog/src/system_schema/information_schema/ssts.rs index 9006ab6fed9..ace6a87bce5 100644 --- a/src/catalog/src/system_schema/information_schema/ssts.rs +++ b/src/catalog/src/system_schema/information_schema/ssts.rs @@ -349,6 +349,13 @@ mod tests { .unwrap() .unwrap(); + let filter = + datafusion::logical_expr::col("region_id").eq(datafusion::logical_expr::lit(42_u64)); + assert_eq!( + table.supports_filters_pushdown(&[&filter]).unwrap(), + vec![table::metadata::FilterPushDownType::Inexact], + ); + let request = ScanRequest { projection: Some(vec![0]), ..Default::default() diff --git a/src/common/recordbatch/src/filter.rs b/src/common/recordbatch/src/filter.rs index 60ff399601f..0893d2aa3ad 100644 --- a/src/common/recordbatch/src/filter.rs +++ b/src/common/recordbatch/src/filter.rs @@ -19,10 +19,11 @@ mod predicate; mod tests; mod timestamp; +use std::borrow::Borrow; use std::sync::Arc; use datafusion::error::Result as DfResult; -use datafusion::logical_expr::{Expr, Literal, Operator}; +use datafusion::logical_expr::{Expr, Literal, Operator, utils}; use datafusion::physical_plan::PhysicalExpr; use datafusion_common::arrow::array::{ArrayRef, Datum}; use datafusion_common::arrow::buffer::BooleanBuffer; @@ -41,6 +42,41 @@ pub use self::timestamp::timestamp_scalar_value; use crate::error::{Result, ToArrowScalarSnafu}; use crate::filter::predicate::SimplePredicate; +/// Prunes candidates using known metadata columns. Missing columns and unsupported +/// predicates cannot exclude a candidate. Complete SQL predicates must still run +/// after pruning. +#[derive(Debug, Clone, Default)] +pub struct MetadataFilter { + filters: Arc<[SimpleFilterEvaluator]>, +} + +impl MetadataFilter { + /// Compiles supported conjuncts once for reuse across scan stages. OR is + /// never split: an unknown alternative might match the candidate. + pub fn new(expressions: &[Expr]) -> Self { + Self { + filters: expressions + .iter() + .flat_map(utils::split_conjunction) + .filter_map(SimpleFilterEvaluator::try_new) + .collect(), + } + } + + /// Returns false only when known metadata proves a conjunct cannot match. + /// `None` means unavailable metadata, distinct from a known SQL NULL scalar. + /// Values may be borrowed so repeated predicates can reuse allocated strings. + pub fn matches>( + &self, + mut value: impl FnMut(&str) -> Option, + ) -> bool { + self.filters.iter().all(|filter| { + value(filter.column_name()) + .is_none_or(|value| filter.evaluate_scalar(value.borrow()).unwrap_or(true)) + }) + } +} + /// Evaluates supported single-column predicates with precompiled regular expressions. /// Compound predicates preserve SQL null semantics until conversion to a scan mask. #[derive(Debug, Clone)] diff --git a/src/common/recordbatch/src/filter/tests.rs b/src/common/recordbatch/src/filter/tests.rs index 77abf59b404..f48b4a3c126 100644 --- a/src/common/recordbatch/src/filter/tests.rs +++ b/src/common/recordbatch/src/filter/tests.rs @@ -70,3 +70,91 @@ fn test_nullable_boolean_predicate_becomes_scan_mask() { boolean_array_to_scan_mask(&predicate) ); } + +#[test] +fn test_partial_metadata() { + use datafusion_common::ScalarValue; + + use crate::filter::MetadataFilter; + + let node = col("node_id").eq(lit(7u64)); + let path = col("file_path").eq(lit("table/region/file.parquet")); + let filters = MetadataFilter::new(&[node.clone().and(path.clone())]); + assert!(!filters.matches(|name| (name == "node_id").then_some(ScalarValue::UInt64(Some(8))))); + assert!(filters.matches(|name| (name == "node_id").then_some(ScalarValue::UInt64(Some(7))))); + assert!(!filters.matches(|name| { + (name == "file_path").then(|| ScalarValue::Utf8(Some("other/file.parquet".into()))) + })); + assert!(filters.matches(|name| { + (name == "file_path").then(|| ScalarValue::Utf8(Some("table/region/file.parquet".into()))) + })); + // Unknown alternatives cannot be discarded from OR, nor extracted from NOT. + let disjunction = MetadataFilter::new(&[node.clone().or(path.clone())]); + assert!( + disjunction.matches(|name| (name == "node_id").then_some(ScalarValue::UInt64(Some(8)))) + ); + let negation = MetadataFilter::new(&[datafusion::logical_expr::Expr::Not(Box::new( + node.and(path), + ))]); + assert!(negation.matches(|_| None::)); + // Unsupported LIKE remains a residual even when its column is available. + let residual = MetadataFilter::new(&[col("file_path").like(lit("table/%"))]); + assert!(residual.matches(|_| Some(ScalarValue::Utf8(Some("other/file".into()))))); + // Missing metadata differs from known NULL; type errors conservatively pass. + let null_filter = MetadataFilter::new(&[col("node_id").is_not_null()]); + assert!(null_filter.matches(|_| None::)); + assert!(!null_filter.matches(|_| Some(ScalarValue::UInt64(None)))); + let typed = MetadataFilter::new(&[col("node_id").eq(lit(7u64))]); + assert!(typed.matches(|_| Some(ScalarValue::Utf8(Some("invalid".into()))))); +} + +#[test] +fn test_borrowed_and_owned_values() { + use datafusion_common::ScalarValue; + + use crate::filter::MetadataFilter; + + let pruner = MetadataFilter::new(&[col("path").gt_eq(lit("a")).and(col("path").lt(lit("z")))]); + for (value, expected) in [ + (ScalarValue::Utf8(Some("a".into())), true), + (ScalarValue::Utf8(Some("z".into())), false), + (ScalarValue::Utf8(None), false), + (ScalarValue::UInt64(Some(1)), true), + ] { + assert_eq!(pruner.matches(|_| Some(&value)), expected); + assert_eq!(pruner.matches(|_| Some(value.clone())), expected); + } +} + +#[test] +fn test_borrowed_lazy_metadata() { + use std::cell::OnceCell; + + use datafusion_common::ScalarValue; + + use crate::filter::MetadataFilter; + + let pruner = MetadataFilter::new(&[col("path").gt(lit("a")).and(col("path").lt(lit("z")))]); + let path = OnceCell::new(); + let mut builds = 0; + assert!(pruner.matches(|name| { + (name == "path").then(|| { + path.get_or_init(|| { + builds += 1; + ScalarValue::Utf8(Some("region/file.parquet".into())) + }) + }) + })); + assert_eq!(builds, 1); + assert_eq!( + path.into_inner(), + Some(ScalarValue::Utf8(Some("region/file.parquet".into()))) + ); + + let path = OnceCell::new(); + assert!( + MetadataFilter::default() + .matches(|_| { Some(path.get_or_init(|| ScalarValue::Utf8(Some("unused".into())))) }) + ); + assert!(path.get().is_none()); +} diff --git a/src/datanode/src/region_server/catalog.rs b/src/datanode/src/region_server/catalog.rs index 264408342e9..789aa514785 100644 --- a/src/datanode/src/region_server/catalog.rs +++ b/src/datanode/src/region_server/catalog.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod inspect; + use std::collections::HashMap; use std::sync::Arc; @@ -33,6 +35,7 @@ use store_api::storage::RegionId; use crate::error::{DataFusionSnafu, ListStorageSstsSnafu, Result, UnexpectedSnafu}; use crate::region_server::RegionServer; +use crate::region_server::catalog::inspect::{InspectTable, ManifestSstSource}; /// Reserved internal table kinds used. /// These are recognized by reserved table names and mapped to providers. @@ -75,7 +78,7 @@ impl InternalTableKind { } impl RegionServer { - /// Expose SSTs listed in Manifest as an in-memory table for inspection. + /// Exposes manifest SSTs through a lazy table with metadata predicate pruning. pub async fn inspect_sst_manifest_provider(&self) -> Result> { let mito = { let guard = self.inner.mito_engine.read().unwrap(); @@ -84,14 +87,9 @@ impl RegionServer { })? }; - let entries = mito.all_ssts_from_manifest().await; - let schema = ManifestSstEntry::schema().arrow_schema().clone(); - let batch = ManifestSstEntry::to_record_batch(&entries) - .map_err(DataFusionError::from) - .context(DataFusionSnafu)?; - - let table = MemTable::try_new(schema, vec![vec![batch]]).context(DataFusionSnafu)?; - Ok(Arc::new(table)) + Ok(Arc::new(InspectTable::new(Arc::new( + ManifestSstSource::new(mito), + )))) } /// Expose SSTs found in storage as an in-memory table for inspection. @@ -355,6 +353,8 @@ impl TreeNodeRewriter for NameAwareDataSourceInjector { #[cfg(test)] mod tests { + mod manifest; + use std::sync::Arc; use datafusion::catalog::MemTable as DfMemTable; diff --git a/src/datanode/src/region_server/catalog/inspect.rs b/src/datanode/src/region_server/catalog/inspect.rs new file mode 100644 index 00000000000..307c8b42883 --- /dev/null +++ b/src/datanode/src/region_server/catalog/inspect.rs @@ -0,0 +1,311 @@ +// 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. + +//! Streaming providers for internal inspection tables. + +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::{Session, TableProvider}; +use datafusion::error::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::streaming::{PartitionStream, StreamingTableExec}; +use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream}; +use datafusion_common::DataFusionError; +use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; +use futures::StreamExt; +use mito2::engine::MitoEngine; +use store_api::sst_entry::{ManifestSstBatchBuilder, ManifestSstEntry}; + +/// Supplies batches for an internal inspection table. Implementations may prune +/// conservatively; DataFusion retains complete residual predicates. +pub(crate) trait InspectSource: std::fmt::Debug + Send + Sync { + fn schema(&self) -> SchemaRef; + + /// Creates a fresh scan at execution time, including on repeated execution. + /// Projection includes columns required by residual predicates. + fn scan(&self, request: &InspectScan, ctx: Arc) -> SendableRecordBatchStream; +} + +/// Parameters shared by inspection sources. Limits are enforced by the physical +/// plan, after residual filters, rather than by individual sources. +#[derive(Debug, Clone)] +pub(crate) struct InspectScan { + schema: SchemaRef, + projection: Option>, + filters: Vec, + batch_size: usize, +} + +/// Adapts a lazy inspection source to DataFusion's streaming execution plan. +#[derive(Debug)] +pub(crate) struct InspectTable { + source: Arc, +} + +impl InspectTable { + /// Wraps a source without opening a scan. + pub(crate) fn new(source: Arc) -> Self { + Self { source } + } +} + +#[async_trait::async_trait] +impl TableProvider for InspectTable { + fn schema(&self) -> SchemaRef { + self.source.schema() + } + + fn table_type(&self) -> TableType { + TableType::View + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> Result> { + Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> Result> { + // An inexact predicate must be checked before selecting a limited prefix. + let limit = limit.filter(|_| filters.is_empty()); + let configured_batch_size: usize = state.config_options().execution.batch_size.into(); + let schema = match projection { + Some(projection) => Arc::new(self.schema().project(projection)?), + None => self.schema(), + }; + let partition = InspectPartition { + source: self.source.clone(), + request: InspectScan { + schema: schema.clone(), + projection: projection.cloned(), + filters: filters.to_vec(), + batch_size: limit + .unwrap_or(configured_batch_size) + .min(configured_batch_size) + .max(1), + }, + }; + Ok(Arc::new(StreamingTableExec::try_new( + schema, + vec![Arc::new(partition)], + None, + None, + false, + limit, + )?)) + } +} + +/// Reads manifest snapshots without materializing the full SST inventory. +pub(crate) struct ManifestSstSource { + engine: MitoEngine, +} + +impl ManifestSstSource { + /// Creates a source backed by the engine's open-region manifests. + pub(crate) fn new(engine: MitoEngine) -> Self { + Self { engine } + } +} + +impl std::fmt::Debug for ManifestSstSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ManifestSstSource").finish_non_exhaustive() + } +} + +impl InspectSource for ManifestSstSource { + fn schema(&self) -> SchemaRef { + ManifestSstEntry::schema().arrow_schema().clone() + } + + fn scan(&self, request: &InspectScan, _ctx: Arc) -> SendableRecordBatchStream { + let builder = ManifestSstBatchBuilder::new(request.schema.clone()); + let stream = self + .engine + .scan_ssts_from_manifest_projected( + &request.filters, + request.batch_size, + request.projection.as_deref(), + ) + .map(move |entries| builder.build(&entries).map_err(DataFusionError::from)); + Box::pin(RecordBatchStreamAdapter::new( + request.schema.clone(), + stream, + )) + } +} + +#[derive(Debug)] +struct InspectPartition { + source: Arc, + request: InspectScan, +} + +impl PartitionStream for InspectPartition { + fn schema(&self) -> &SchemaRef { + &self.request.schema + } + + fn execute(&self, ctx: Arc) -> SendableRecordBatchStream { + self.source.scan(&self.request, ctx) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use datafusion::arrow::array::{ArrayRef, Int32Array}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::prelude::SessionContext; + + use super::*; + + #[derive(Debug, Default)] + struct CountingSource { + scans: Arc, + rows: Arc, + drops: Arc, + } + + struct ScanGuard(Arc); + + impl Drop for ScanGuard { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + impl InspectSource for CountingSource { + fn schema(&self) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])) + } + + fn scan(&self, request: &InspectScan, _ctx: Arc) -> SendableRecordBatchStream { + self.scans.fetch_add(1, Ordering::SeqCst); + let guard = ScanGuard(self.drops.clone()); + let rows = self.rows.clone(); + let schema = request.schema.clone(); + let projection = request.projection.clone().unwrap_or_else(|| vec![0, 1]); + // Deliberately ignore source predicates: the provider must preserve + // residual filtering before limit, including hidden predicate columns. + let stream = futures::stream::iter(0..100).map(move |id| { + let _guard = &guard; + rows.fetch_add(1, Ordering::SeqCst); + let columns: Vec = projection + .iter() + .map(|column| { + Arc::new(Int32Array::from(vec![if *column == 0 { + id + } else { + id * 10 + }])) as ArrayRef + }) + .collect(); + RecordBatch::try_new(schema.clone(), columns).map_err(DataFusionError::from) + }); + Box::pin(RecordBatchStreamAdapter::new( + request.schema.clone(), + stream, + )) + } + } + + #[tokio::test] + async fn test_lazy_repeatable_scan_and_limit_cancellation() { + let source = Arc::new(CountingSource::default()); + let ctx = SessionContext::new(); + ctx.register_table("inspection", Arc::new(InspectTable::new(source.clone()))) + .unwrap(); + let frame = ctx + .sql("SELECT value FROM inspection LIMIT 2") + .await + .unwrap(); + let plan = frame.create_physical_plan().await.unwrap(); + assert_eq!(source.scans.load(Ordering::SeqCst), 0); + assert_eq!(source.rows.load(Ordering::SeqCst), 0); + for run in 1..=2 { + let batches = datafusion::physical_plan::collect(plan.clone(), ctx.task_ctx()) + .await + .unwrap(); + let values = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + assert_eq!(values, vec![0, 10]); + assert_eq!(source.scans.load(Ordering::SeqCst), run); + assert_eq!(source.rows.load(Ordering::SeqCst), run * 2); + assert_eq!(source.drops.load(Ordering::SeqCst), run); + } + } + + #[tokio::test] + async fn test_hidden_residual_filter_before_limit() { + let source = Arc::new(CountingSource::default()); + let ctx = SessionContext::new(); + ctx.register_table("inspection", Arc::new(InspectTable::new(source.clone()))) + .unwrap(); + let batches = ctx + .sql("SELECT value FROM inspection WHERE id >= 8 LIMIT 2") + .await + .unwrap() + .collect() + .await + .unwrap(); + let values = batches + .iter() + .flat_map(|batch| { + assert_eq!(batch.num_columns(), 1); + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + // No ORDER BY: repartitioning may choose any two qualifying rows. + assert_eq!(values.len(), 2); + assert!( + values + .iter() + .all(|value| (80..1000).contains(value) && value % 10 == 0) + ); + assert_eq!(source.scans.load(Ordering::SeqCst), 1); + assert_eq!(source.drops.load(Ordering::SeqCst), 1); + } +} diff --git a/src/datanode/src/region_server/catalog/tests/manifest.rs b/src/datanode/src/region_server/catalog/tests/manifest.rs new file mode 100644 index 00000000000..91b9f757a3d --- /dev/null +++ b/src/datanode/src/region_server/catalog/tests/manifest.rs @@ -0,0 +1,257 @@ +// 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. + +use std::sync::Arc; + +use common_time::Timestamp; +use datafusion::arrow::compute::concat_batches; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::catalog::{MemTable, TableProvider}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use mito2::config::MitoConfig; +use mito2::engine::MitoEngine; +use mito2::manifest::action::RegionEdit; +use mito2::sst::file::FileMeta; +use mito2::test_util::{CreateRequestBuilder, TestEnv}; +use store_api::region_engine::RegionEngine; +use store_api::region_request::RegionRequest; +use store_api::sst_entry::ManifestSstEntry; +use store_api::storage::{FileId, RegionId}; + +use crate::region_server::catalog::inspect::{InspectTable, ManifestSstSource}; + +/// Installs deterministic metadata only. No SST object is read by inspection. +async fn manifest_engine(regions: u32, files_per_region: usize) -> (TestEnv, MitoEngine) { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + schedule_compaction_after_edit: false, + ..Default::default() + }) + .await; + for number in 0..regions { + let region_id = RegionId::new(11, number); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + engine + .edit_region( + region_id, + RegionEdit { + files_to_add: (0..files_per_region) + .map(|index| FileMeta { + region_id, + file_id: FileId::from_bytes( + (((number as u128) << 64) | index as u128).to_be_bytes(), + ), + time_range: ( + Timestamp::new_millisecond(index as i64), + Timestamp::new_millisecond(index as i64 + 1), + ), + file_size: 1024 + index as u64, + sequence: None, + ..Default::default() + }) + .collect(), + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + ) + .await + .unwrap(); + } + (env, engine) +} + +/// Baseline preserves eager full-row materialization, then applies SQL. Both +/// variants use the current engine enumerator, so this does not measure removal +/// of the old manifest-map cloning independently. +async fn inspect_query(engine: &MitoEngine, streaming: bool, sql: &str) -> RecordBatch { + let schema = ManifestSstEntry::schema().arrow_schema().clone(); + let table: Arc = if streaming { + Arc::new(InspectTable::new(Arc::new(ManifestSstSource::new( + engine.clone(), + )))) + } else { + let entries = engine.all_ssts_from_manifest().await; + let batch = ManifestSstEntry::to_record_batch(&entries).unwrap(); + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()) + }; + let context = SessionContext::new_with_config( + SessionConfig::new() + .with_batch_size(17) + .with_target_partitions(1), + ); + context.register_table("manifest", table).unwrap(); + let frame = context.sql(sql).await.unwrap(); + let schema = frame.schema().as_arrow().clone(); + let batches = frame.collect().await.unwrap(); + concat_batches(&Arc::new(schema), &batches).unwrap() +} + +#[tokio::test] +async fn test_manifest_provider_matches_eager_scan() { + let (_env, engine) = manifest_engine(3, 64).await; + let region = RegionId::new(11, 1).as_u64(); + for sql in [ + "SELECT file_id, min_ts FROM manifest WHERE table_id = 11 ORDER BY min_ts, file_id LIMIT 7" + .to_string(), + "SELECT file_id FROM manifest WHERE table_id = 12 LIMIT 7".to_string(), + format!( + "SELECT file_id FROM manifest WHERE table_id = 11 AND region_id = {region} ORDER BY min_ts DESC LIMIT 7" + ), + format!("SELECT * FROM manifest WHERE region_id = {region} ORDER BY file_id"), + format!( + "SELECT min_ts, max_ts FROM manifest WHERE region_id = {region} ORDER BY min_ts ASC NULLS LAST, file_id LIMIT 7" + ), + format!( + "SELECT file_size FROM manifest WHERE region_id = {region} ORDER BY max_ts DESC NULLS FIRST, file_id LIMIT 7 OFFSET 3" + ), + // Arithmetic is deliberately left as a residual predicate. Limiting + // metadata enumeration before it would return too few rows. + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} AND file_size % 3 = 0 ORDER BY min_ts, file_id LIMIT 7" + ), + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} AND min_ts >= TIMESTAMP '1970-01-01 00:00:00.020' ORDER BY min_ts, file_id LIMIT 7" + ), + // Variable-sized metadata needed only by a residual filter must still + // be materialized even when it is absent from the SELECT list. + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} AND length(file_path) > 0 ORDER BY min_ts, file_id LIMIT 7" + ), + // Repeated path predicates reuse metadata while both path columns stay + // hidden from the output. The range selects only region 1's 64 files. + "SELECT file_id FROM manifest WHERE table_dir >= 'test' AND table_dir < 'tesu' AND file_path >= 'test/11_0000000001/' AND file_path < 'test/11_0000000002/' ORDER BY file_id LIMIT 7".to_string(), + // Visibility snapshots are needed for residual predicates even when + // the SELECT list does not include visible. + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} AND visible ORDER BY file_id LIMIT 7" + ), + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} AND NOT visible ORDER BY file_id LIMIT 7" + ), + "SELECT file_id FROM manifest WHERE region_id = 0 ORDER BY min_ts LIMIT 1".to_string(), + "SELECT min_ts FROM manifest ORDER BY min_ts LIMIT 0".to_string(), + "SELECT count(*) FROM manifest".to_string(), + // Unsupported compound predicates must not prune matching regions. + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} OR file_size = 1024 ORDER BY file_id" + ), + ] { + let expected = inspect_query(&engine, false, &sql).await; + let actual = inspect_query(&engine, true, &sql).await; + assert_eq!(actual, expected, "{sql}"); + } +} + +#[tokio::test] +async fn test_manifest_referring_region_and_origin_path() { + use datafusion::arrow::array::{StringArray, UInt64Array}; + + let (_env, engine) = manifest_engine(2, 2).await; + let referring = RegionId::new(11, 1); + // The creating region need not be open on this datanode. + let origin = RegionId::new(11, 99); + let file_id = FileId::from_bytes(999u128.to_be_bytes()); + engine + .edit_region( + referring, + RegionEdit { + files_to_add: vec![FileMeta { + region_id: origin, + file_id, + file_size: 4096, + ..Default::default() + }], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + ) + .await + .unwrap(); + let sql = format!( + "SELECT region_id, origin_region_id, file_path FROM manifest WHERE table_id = 11 AND region_id = {} AND origin_region_id = {} AND file_path LIKE '%/11_0000000099/%' LIMIT 1", + referring.as_u64(), + origin.as_u64(), + ); + // Check source pruning without SQL residuals: the origin is not a loaded + // region, but its file must be found in the referring region's manifest. + use datafusion_expr::{col, lit}; + use futures::StreamExt; + for (predicate, expected) in [ + (col("origin_region_id").eq(lit(origin.as_u64())), 1), + ( + col("file_path").eq(lit(format!("test/11_0000000099/{file_id}.parquet"))), + 1, + ), + (col("file_path").eq(lit("missing/file.parquet")), 0), + (col("table_dir").eq(lit("missing/")), 0), + (col("node_id").is_not_null(), 0), + ] { + let batches = engine + .scan_ssts_from_manifest_projected(std::slice::from_ref(&predicate), 2, None) + .collect::>() + .await; + assert_eq!( + batches.iter().map(Vec::len).sum::(), + expected, + "{predicate}" + ); + } + let actual = inspect_query(&engine, true, &sql).await; + assert_eq!(actual, inspect_query(&engine, false, &sql).await); + assert_eq!(actual.num_rows(), 1); + assert_eq!( + actual + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + referring.as_u64() + ); + assert_eq!( + actual + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + origin.as_u64() + ); + assert!( + actual + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + .ends_with(&format!("11_0000000099/{file_id}.parquet")) + ); + let missing = sql.replace("%/11_0000000099/%", "%/11_0000000001/%"); + assert_eq!(inspect_query(&engine, true, &missing).await.num_rows(), 0); +} diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index a9080745e2e..6467c710501 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -768,11 +768,16 @@ mod tests { .unwrap(); let region = env.mito().find_region(data_region_id).unwrap(); - let entry = region - .manifest_sst_entries() + let entry = env + .mito() + .all_ssts_from_manifest() .await .into_iter() - .find(|entry| entry.visible && entry.file_path.ends_with(".parquet")) + .find(|entry| { + entry.region_id == data_region_id + && entry.visible + && entry.file_path.ends_with(".parquet") + }) .unwrap(); let mut cache_metrics = MetadataCacheMetrics::default(); let footer = MetadataLoader::new( diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 715357a39b3..92d848f5677 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -104,10 +104,13 @@ use common_base::Plugins; use common_error::ext::BoxedError; use common_meta::error::UnexpectedSnafu; use common_meta::key::SchemaMetadataManagerRef; +use common_recordbatch::filter::MetadataFilter; use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream}; use common_stat::get_total_memory_bytes; use common_telemetry::{debug, info, tracing, warn}; use common_wal::options::WalOptions; +use datafusion_common::ScalarValue; +use datafusion_expr::Expr; use futures::FutureExt; use futures::future::{self, join_all, try_join_all}; use futures::stream::{self, Stream, StreamExt}; @@ -132,7 +135,7 @@ use store_api::region_info::RegionInfoEntry; use store_api::region_request::{ AffectedRows, RegionCatchupRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, }; -use store_api::sst_entry::{ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry}; +use store_api::sst_entry::{self, ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry}; use store_api::storage::{FileId, FileRefsManifest, RegionId, ScanRequest, SequenceNumber}; use tokio::sync::{Semaphore, oneshot}; @@ -578,19 +581,65 @@ impl MitoEngine { /// Lists all SSTs from the manifest of all regions in the engine. pub async fn all_ssts_from_manifest(&self) -> Vec { - let node_id = self.inner.workers.file_ref_manager().node_id(); - let regions = self.inner.workers.all_regions(); + self.scan_ssts_from_manifest(&[], 1024) + .flat_map(stream::iter) + .collect() + .await + } - let mut results = Vec::new(); - for region in regions { - let mut entries = region.manifest_sst_entries().await; - for e in &mut entries { - e.node_id = node_id; + /// Lazily emits bounded batches of manifest entries, pruning regions before + /// opening their manifest snapshots and files before constructing entries. + /// Filters are conservative; callers must retain their residual predicates. + pub(crate) fn scan_ssts_from_manifest( + &self, + filters: &[Expr], + batch_size: usize, + ) -> futures::stream::BoxStream<'static, Vec> { + self.scan_ssts_from_manifest_projected(filters, batch_size, None) + } + + /// Like [`Self::scan_ssts_from_manifest`], but avoids allocating unrequested + /// variable-sized fields. The projection must include residual filter columns. + pub fn scan_ssts_from_manifest_projected( + &self, + filters: &[Expr], + batch_size: usize, + projection: Option<&[usize]>, + ) -> futures::stream::BoxStream<'static, Vec> { + let engine = self.clone(); + let projection = crate::region::ManifestSstProjection::new(projection); + let filters = MetadataFilter::new(filters); + Box::pin(async_stream::stream! { + let node_id = engine.inner.workers.file_ref_manager().node_id(); + if !filters.matches(|column| match column { + sst_entry::NODE_ID => Some(ScalarValue::UInt64(node_id)), + _ => None, + }) { + return; } - results.extend(entries); - } - - results + let mut batch = Vec::with_capacity(batch_size.max(1)); + for region in engine.inner.workers.all_regions() { + if !manifest_region_matches(region.region_id, node_id, &filters) + || !filters.matches(|column| match column { + sst_entry::TABLE_DIR => Some(ScalarValue::Utf8(Some(region.table_dir().to_string()))), + _ => None, + }) + { + continue; + } + let mut entries = region.manifest_sst_stream(filters.clone(), projection); + while let Some(mut entry) = entries.next().await { + entry.node_id = node_id; + batch.push(entry); + if batch.len() == batch_size.max(1) { + yield std::mem::replace(&mut batch, Vec::with_capacity(batch_size.max(1))); + } + } + } + if !batch.is_empty() { + yield batch; + } + }) } /// Lists metadata about all puffin index targets stored in the engine. @@ -1700,6 +1749,25 @@ impl MitoEngine { } } +/// Evaluates region-constant predicates before visiting any of its files. +fn manifest_region_matches( + region_id: RegionId, + node_id: Option, + filters: &MetadataFilter, +) -> bool { + filters.matches(|column| { + Some(match column { + sst_entry::REGION_ID => ScalarValue::UInt64(Some(region_id.as_u64())), + sst_entry::TABLE_ID => ScalarValue::UInt32(Some(region_id.table_id())), + sst_entry::REGION_NUMBER => ScalarValue::UInt32(Some(region_id.region_number())), + sst_entry::REGION_GROUP => ScalarValue::UInt8(Some(region_id.region_group())), + sst_entry::REGION_SEQUENCE => ScalarValue::UInt32(Some(region_id.region_sequence())), + sst_entry::NODE_ID => ScalarValue::UInt64(node_id), + _ => return None, + }) + }) +} + #[cfg(test)] mod tests { use std::time::Duration; @@ -1707,6 +1775,30 @@ mod tests { use super::*; use crate::sst::file::FileMeta; + #[test] + fn test_manifest_region_predicates() { + use super::manifest_region_matches; + use common_recordbatch::filter::MetadataFilter; + use datafusion_expr::{col, lit}; + let region = store_api::storage::RegionId::new(42, 3); + for (expr, expected) in [ + (col("region_id").eq(lit(region.as_u64())), true), + (col("region_id").eq(lit(region.as_u64() + 1)), false), + (col("table_id").eq(lit(42u32)), true), + (col("node_id").is_null(), true), + (col("file_size").gt(lit(100u64)), true), + // A type mismatch is left to the residual filter, never pruned. + (col("region_id").eq(lit("invalid")), true), + ] { + let filter = MetadataFilter::new(std::slice::from_ref(&expr)); + assert_eq!( + manifest_region_matches(region, None, &filter), + expected, + "{expr}" + ); + } + } + #[test] fn test_is_valid_region_edit() { // Valid: has only "files_to_add" diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 2576158bab6..1b3d27120c8 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -2047,3 +2047,247 @@ async fn test_bulk_insert_encode_failure_keeps_worker_batch() { assert_eq!(3, scan_num_rows(&engine, region_id).await); assert_eq!(2, scan_num_rows(&engine, other_region_id).await); } + +#[tokio::test] +async fn test_manifest_scan_region_filter() { + use datafusion_expr::{col, lit}; + use futures::StreamExt; + + let regions = 3; + let files_per_region = 64; + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + for number in 0..regions { + add_manifest_scan_region(&engine, RegionId::new(11, number), files_per_region).await; + } + + let target = RegionId::new(11, 1); + let batch_size = 17; + let mut reference = None; + for pushed in [false, true] { + let filters = if pushed { + vec![col("region_id").eq(lit(target.as_u64()))] + } else { + vec![] + }; + let mut stream = engine.scan_ssts_from_manifest(&filters, batch_size); + let mut examined_entries = 0; + let mut selected = Vec::new(); + while let Some(batch) = stream.next().await { + assert!(batch.len() <= batch_size); + examined_entries += batch.len(); + selected.extend(batch.into_iter().filter(|entry| entry.region_id == target)); + } + assert_eq!(selected.len(), files_per_region); + assert_eq!( + examined_entries, + if pushed { + files_per_region + } else { + regions as usize * files_per_region + } + ); + selected.sort_unstable_by(|a, b| a.file_id.cmp(&b.file_id)); + if let Some(reference) = &reference { + assert_eq!(&selected, reference); + } else { + reference = Some(selected); + } + } + // Region pruning must happen before acquiring a manifest snapshot. Holding + // another region's write lock would deadlock an unpruned enumeration. + let unrelated = engine.get_region(RegionId::new(11, 0)).unwrap(); + let guard = unrelated.manifest_ctx.manifest_manager.write().await; + // Node and directory predicates also reject candidates before opening any + // manifest; otherwise the locked region would block these empty scans. + for predicate in [ + col("node_id").is_not_null(), + col("table_dir").eq(lit("missing/")), + ] { + let batches = tokio::time::timeout( + Duration::from_secs(5), + engine + .scan_ssts_from_manifest(&[predicate], batch_size) + .collect::>(), + ) + .await + .expect("pruned scan must not acquire a manifest lock"); + assert!(batches.is_empty()); + } + let filtered = + engine.scan_ssts_from_manifest(&[col("region_id").eq(lit(target.as_u64()))], batch_size); + let batches = tokio::time::timeout(Duration::from_secs(5), filtered.collect::>()) + .await + .expect("filtered scan must not acquire an unrelated manifest lock"); + assert_eq!( + batches.iter().map(Vec::len).sum::(), + files_per_region + ); + drop(guard); + + let schema = store_api::sst_entry::ManifestSstEntry::schema(); + let projection = ["file_id", "min_ts"] + .map(|name| { + schema + .column_schemas() + .iter() + .position(|c| c.name == name) + .unwrap() + }) + .to_vec(); + let mut projected = engine.scan_ssts_from_manifest_projected( + &[col("region_id").eq(lit(target.as_u64()))], + batch_size, + Some(&projection), + ); + let mut projected_entries = Vec::new(); + while let Some(batch) = projected.next().await { + assert!( + batch + .iter() + .all(|entry| entry.table_dir.is_empty() && entry.file_path.is_empty()) + ); + projected_entries.extend(batch); + } + projected_entries.sort_unstable_by(|a, b| a.file_id.cmp(&b.file_id)); + let builder = store_api::sst_entry::ManifestSstBatchBuilder::new(Arc::new( + store_api::sst_entry::ManifestSstEntry::schema() + .arrow_schema() + .project(&projection) + .unwrap(), + )); + let expected = builder.build(&reference.unwrap()).unwrap(); + let actual = builder.build(&projected_entries).unwrap(); + assert_eq!(actual, expected); + assert!( + engine + .scan_ssts_from_manifest(&[col("region_id").eq(lit(0u64))], batch_size) + .next() + .await + .is_none() + ); +} + +/// Installs synthetic immutable file metadata without creating physical SSTs. +async fn add_manifest_scan_region( + engine: &MitoEngine, + region_id: RegionId, + files_per_region: usize, +) { + use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; + use crate::sst::file::FileMeta; + use store_api::storage::FileId; + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let edit = RegionEdit { + files_to_add: (0..files_per_region) + .map(|index| FileMeta { + region_id, + file_id: FileId::from_bytes( + (((region_id.region_number() as u128) << 64) | index as u128).to_be_bytes(), + ), + time_range: ( + Timestamp::new_millisecond(index as i64), + Timestamp::new_millisecond(index as i64 + 1), + ), + file_size: 1024 + index as u64, + ..Default::default() + }) + .collect(), + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }; + region + .manifest_ctx + .manifest_manager + .write() + .await + .update( + RegionMetaActionList::new(vec![RegionMetaAction::Edit(edit)]), + false, + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_manifest_scan_table_filter() { + use datafusion_expr::{col, lit}; + use futures::StreamExt; + use store_api::sst_entry::ManifestSstEntry; + + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let target_regions = [RegionId::new(11, 0), RegionId::new(11, 1)]; + let unrelated_id = RegionId::new(12, 0); + for region_id in target_regions.into_iter().chain([unrelated_id]) { + add_manifest_scan_region(&engine, region_id, 3).await; + } + let mut expected = engine.all_ssts_from_manifest().await; + expected.retain(|entry| entry.table_id == 11); + expected.sort_unstable_by_key(|entry| (entry.region_id, entry.file_id.clone())); + assert_eq!(expected.len(), 6); + + // A table predicate must exclude other tables before locking their manifests, + // while keeping every region of the selected table (including the same number). + let unrelated = engine.get_region(unrelated_id).unwrap(); + let guard = unrelated.manifest_ctx.manifest_manager.write().await; + let batches = tokio::time::timeout( + Duration::from_secs(5), + engine + .scan_ssts_from_manifest(&[col("table_id").eq(lit(11u32))], 2) + .collect::>(), + ) + .await + .expect("table filtering must not acquire another table's manifest lock"); + assert!(batches.iter().all(|batch| batch.len() <= 2)); + let mut actual = batches.into_iter().flatten().collect::>(); + actual.sort_unstable_by_key(|entry| (entry.region_id, entry.file_id.clone())); + assert_eq!(actual, expected); + + let schema = ManifestSstEntry::schema(); + let projection = ["region_id", "file_id"] + .map(|name| { + schema + .column_schemas() + .iter() + .position(|column| column.name == name) + .unwrap() + }) + .to_vec(); + let batches = tokio::time::timeout( + Duration::from_secs(5), + engine + .scan_ssts_from_manifest_projected( + &[col("table_id").eq(lit(11u32))], + 2, + Some(&projection), + ) + .collect::>(), + ) + .await + .expect("hidden table predicate must prune unrelated manifests"); + let mut actual = batches.into_iter().flatten().collect::>(); + actual.sort_unstable_by_key(|entry| (entry.region_id, entry.file_id.clone())); + let builder = store_api::sst_entry::ManifestSstBatchBuilder::new(Arc::new( + ManifestSstEntry::schema() + .arrow_schema() + .project(&projection) + .unwrap(), + )); + assert_eq!( + builder.build(&actual).unwrap(), + builder.build(&expected).unwrap() + ); + drop(guard); +} diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index 1b2a7ffee13..98973de4f30 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -843,6 +843,104 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) assert_eq!(sst_entries.iter().filter(|e| e.visible).count(), 1); assert_eq!(sst_entries.iter().filter(|e| !e.visible).count(), 2); + // Streaming inspection must retain staging entries and their visibility. + { + use futures::StreamExt; + + let mut batches = engine.scan_ssts_from_manifest( + &[datafusion_expr::col("region_id").eq(datafusion_expr::lit(region_id.as_u64()))], + 1, + ); + let mut actual = Vec::new(); + while let Some(batch) = batches.next().await { + assert_eq!(batch.len(), 1); + actual.extend(batch); + } + actual.sort_unstable_by(|a, b| a.file_id.cmp(&b.file_id)); + let mut expected = sst_entries; + expected.sort_unstable_by(|a, b| a.file_id.cmp(&b.file_id)); + assert_eq!(actual, expected); + } + + // Staging metadata overrides the same file in the normal manifest. Filter + // evaluation must use that winning metadata, even if it is not projected. + { + use futures::StreamExt; + + let region = engine.get_region(region_id).unwrap(); + let original = region + .manifest_ctx + .manifest_manager + .read() + .await + .manifest() + .files + .values() + .next() + .unwrap() + .clone(); + let mut overridden = original.clone(); + overridden.file_size = u64::MAX - 1; + for metadata in [overridden, original.clone()] { + region + .manifest_ctx + .manifest_manager + .write() + .await + .update( + RegionMetaActionList::new(vec![RegionMetaAction::Edit(RegionEdit { + files_to_add: vec![metadata.clone()], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + })]), + true, + ) + .await + .unwrap(); + let entries = engine.all_ssts_from_manifest().await; + assert_eq!(entries.len(), 3); + let matches = entries + .iter() + .filter(|entry| entry.file_id == original.file_id.to_string()) + .collect::>(); + assert_eq!(matches.len(), 1); + assert_eq!(matches[0].file_size, metadata.file_size); + assert!(matches[0].visible); + + let schema = store_api::sst_entry::ManifestSstEntry::schema(); + let projection = ["file_id", "visible"] + .map(|name| { + schema + .column_schemas() + .iter() + .position(|c| c.name == name) + .unwrap() + }) + .to_vec(); + let mut stream = engine.scan_ssts_from_manifest_projected( + &[datafusion_expr::col("file_size").eq(datafusion_expr::lit(metadata.file_size))], + 1, + Some(&projection), + ); + let mut filtered = Vec::new(); + while let Some(batch) = stream.next().await { + filtered.extend(batch); + } + assert!( + filtered + .iter() + .any(|entry| entry.file_id == original.file_id.to_string()) + ); + if metadata.file_size == u64::MAX - 1 { + assert_eq!(filtered.len(), 1); + } + } + } + // Exit staging mode successfully engine .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader) diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index ba32f7aff56..613b1ce428e 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -20,6 +20,7 @@ pub mod options; pub mod utils; pub(crate) mod version; +use std::cell::OnceCell; use std::collections::hash_map::Entry; use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; @@ -27,8 +28,12 @@ use std::sync::{Arc, Mutex, RwLock}; use common_base::hash::partition_expr_version; use common_recordbatch::adapter::RegionQueryStatCounters; +use common_recordbatch::filter::MetadataFilter; use common_telemetry::{error, info, warn}; +use common_time::timestamp::TimeUnit; use crossbeam_utils::atomic::AtomicCell; +use datafusion_common::ScalarValue; +use futures::stream::{BoxStream, StreamExt}; use object_store::ObjectStore; use partition::expr::PartitionExpr; use snafu::{OptionExt, ResultExt, ensure}; @@ -42,7 +47,7 @@ use store_api::region_engine::{ }; use store_api::region_info::RegionInfoEntry; use store_api::region_request::{PathType, StagingPartitionDirective}; -use store_api::sst_entry::ManifestSstEntry; +use store_api::sst_entry::{self, ManifestSstEntry}; use store_api::storage::{FileId, RegionId, SequenceNumber}; use tokio::sync::RwLockWriteGuard; pub use utils::*; @@ -823,58 +828,85 @@ impl MitoRegion { } /// Returns the SST entries of the region. - pub async fn manifest_sst_entries(&self) -> Vec { - let table_dir = self.table_dir(); - let path_type = self.access_layer.path_type(); - - let visible_ssts = self - .version() - .ssts - .levels() - .iter() - .flat_map(|level| level.files().map(|file| file.file_id().file_id())) - .collect::>(); - - let manifest_files = self.manifest_ctx.manifest().await.files.clone(); - let staging_files = self - .manifest_ctx - .staging_manifest() + pub(crate) async fn manifest_sst_entries(&self) -> Vec { + self.manifest_sst_stream(MetadataFilter::default(), ManifestSstProjection::new(None)) + .collect() .await - .map(|m| m.files.clone()) - .unwrap_or_default(); - let files = manifest_files - .into_iter() - .chain(staging_files) - .collect::>(); + } - files - .values() - .map(|meta| { + /// Traverses manifest snapshots without copying the file maps. Staging metadata + /// overrides the normal manifest for the same file. + pub(crate) fn manifest_sst_stream( + &self, + filters: MetadataFilter, + projection: ManifestSstProjection, + ) -> BoxStream<'_, ManifestSstEntry> { + Box::pin(async_stream::stream! { + let table_dir = self.table_dir(); + let path_type = self.access_layer.path_type(); + let visible_files = self.manifest_visible_file_ids(projection); + let (manifest, staging) = { + let manager = self.manifest_ctx.manifest_manager.read().await; + (manager.manifest(), manager.staging_manifest()) + }; + let normal_files = manifest.files.values().map(|meta| (false, meta)); + let staging_files = staging.iter().flat_map(|s| s.files.values()).map(|meta| (true, meta)); + for (index, (is_staging, meta)) in normal_files.chain(staging_files).enumerate() { + // A selective filter may otherwise traverse millions of files in + // a single poll without yielding any rows. + if index > 0 && index % 1024 == 0 { + tokio::task::yield_now().await; + } + if !is_staging && staging.as_ref().is_some_and(|s| s.files.contains_key(&meta.file_id)) { + continue; + } + if !manifest_file_matches(meta, &filters) { + continue; + } + // Paths describe the creating region, not the referring region. + // Construct them only when a path predicate or projection needs them. + let file_path = OnceCell::new(); + if !filters.matches(|column| match column { + sst_entry::FILE_PATH => Some(file_path.get_or_init(|| { + ScalarValue::Utf8(Some(sst_file_path(table_dir, meta.file_id(), path_type))) + })), + _ => None, + }) { + continue; + } let region_id = self.region_id; let origin_region_id = meta.region_id; let (index_version, index_file_path, index_file_size) = if meta.index_file_size > 0 { - let index_file_path = index_file_path(table_dir, meta.index_id(), path_type); + let index_file_path = projection.index_file_path.then(|| index_file_path(table_dir, meta.index_id(), path_type)); ( meta.index_version, - Some(index_file_path), + index_file_path, Some(meta.index_file_size), ) } else { (0, None, None) }; - let visible = visible_ssts.contains(&meta.file_id); - ManifestSstEntry { - table_dir: table_dir.to_string(), + let visible = visible_files.as_ref().is_some_and(|files| files.contains(&meta.file_id)); + let file_path = if projection.file_path { + match file_path.into_inner() { + Some(ScalarValue::Utf8(Some(path))) => path, + _ => sst_file_path(table_dir, meta.file_id(), path_type), + } + } else { + String::new() + }; + yield ManifestSstEntry { + table_dir: if projection.table_dir { table_dir.to_string() } else { String::new() }, region_id, table_id: region_id.table_id(), region_number: region_id.region_number(), region_group: region_id.region_group(), region_sequence: region_id.region_sequence(), - file_id: meta.file_id.to_string(), + file_id: if projection.file_id { meta.file_id.to_string() } else { String::new() }, index_version, level: meta.level, - file_path: sst_file_path(table_dir, meta.file_id(), path_type), + file_path, file_size: meta.file_size, max_row_group_uncompressed_size: meta.max_row_group_uncompressed_size, index_file_path, @@ -885,15 +917,31 @@ impl MitoRegion { min_ts: meta.time_range.0, max_ts: meta.time_range.1, sequence: meta.sequence.map(|s| s.get()), - partition_expr: meta.partition_expr.as_ref().map(ToString::to_string), + partition_expr: if projection.partition_expr { meta.partition_expr.as_ref().map(ToString::to_string) } else { None }, origin_region_id, node_id: None, visible, - primary_key_min: meta.primary_key_min.clone(), - primary_key_max: meta.primary_key_max.clone(), - } - }) - .collect() + primary_key_min: if projection.primary_key_min { meta.primary_key_min.clone() } else { None }, + primary_key_max: if projection.primary_key_max { meta.primary_key_max.clone() } else { None }, + }; + } + }) + } + + /// Copies only file IDs so a paused inspection cannot delay SST purging. + /// Hidden residual columns are included in the query's projection. + fn manifest_visible_file_ids( + &self, + projection: ManifestSstProjection, + ) -> Option> { + projection.visible.then(|| { + self.version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files().map(|file| file.file_id().file_id())) + .collect() + }) } /// Returns the file metas of the region by file ids. @@ -1993,6 +2041,75 @@ pub fn parse_partition_expr(partition_expr_str: Option<&str>) -> Result) -> Self { + let schema = ManifestSstEntry::schema(); + let included = |name: &str| { + projection.is_none_or(|indices| { + indices.iter().any(|index| { + schema + .column_schemas() + .get(*index) + .is_some_and(|column| column.name == name) + }) + }) + }; + Self { + visible: included(sst_entry::VISIBLE), + table_dir: included(sst_entry::TABLE_DIR), + file_id: included(sst_entry::FILE_ID), + file_path: included(sst_entry::FILE_PATH), + index_file_path: included(sst_entry::INDEX_FILE_PATH), + partition_expr: included(sst_entry::PARTITION_EXPR), + primary_key_min: included(sst_entry::PRIMARY_KEY_MIN), + primary_key_max: included(sst_entry::PRIMARY_KEY_MAX), + } + } +} + +/// Applies supported file predicates before allocating an inspection row. Unknown +/// predicates and evaluation failures are left to the query's residual filter. +fn manifest_file_matches(meta: &FileMeta, filters: &MetadataFilter) -> bool { + filters.matches(|column| { + Some(match column { + sst_entry::FILE_SIZE => ScalarValue::UInt64(Some(meta.file_size)), + sst_entry::LEVEL => ScalarValue::UInt8(Some(meta.level)), + sst_entry::NUM_ROWS => ScalarValue::UInt64(Some(meta.num_rows)), + sst_entry::NUM_ROW_GROUPS => ScalarValue::UInt64(Some(meta.num_row_groups)), + sst_entry::ORIGIN_REGION_ID => ScalarValue::UInt64(Some(meta.region_id.as_u64())), + sst_entry::MIN_TS => ScalarValue::TimestampNanosecond( + meta.time_range + .0 + .convert_to(TimeUnit::Nanosecond) + .map(|ts| ts.value()), + None, + ), + sst_entry::MAX_TS => ScalarValue::TimestampNanosecond( + meta.time_range + .1 + .convert_to(TimeUnit::Nanosecond) + .map(|ts| ts.value()), + None, + ), + _ => return None, + }) + }) +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -2024,6 +2141,66 @@ mod tests { use crate::test_util::version_util::VersionControlBuilder; use crate::time_provider::StdTimeProvider; + #[test] + fn test_manifest_file_predicates() { + use common_recordbatch::filter::MetadataFilter; + use common_time::Timestamp; + use datafusion_common::ScalarValue; + use datafusion_expr::{col, lit}; + let meta = super::FileMeta { + file_size: 100, + time_range: ( + Timestamp::new_millisecond(10), + Timestamp::new_millisecond(20), + ), + ..Default::default() + }; + for (expr, expected) in [ + (col("file_size").lt(lit(101u64)), true), + (col("file_size").lt(lit(100u64)), false), + ( + col("min_ts").eq(lit(ScalarValue::TimestampNanosecond( + Some(10_000_000), + None, + ))), + true, + ), + ( + col("max_ts").lt(lit(ScalarValue::TimestampNanosecond( + Some(20_000_000), + None, + ))), + false, + ), + (col("file_id").eq(lit("unknown")), true), + ] { + let filter = MetadataFilter::new(std::slice::from_ref(&expr)); + assert_eq!( + super::manifest_file_matches(&meta, &filter), + expected, + "{expr}" + ); + } + } + + #[test] + fn test_manifest_timestamp_overflow_is_null() { + use common_recordbatch::filter::MetadataFilter; + use common_time::Timestamp; + use datafusion_expr::col; + let meta = super::FileMeta { + time_range: ( + Timestamp::new_second(i64::MAX), + Timestamp::new_second(i64::MAX), + ), + ..Default::default() + }; + let filter = MetadataFilter::new(&[col("min_ts").is_null()]); + assert!(super::manifest_file_matches(&meta, &filter)); + let filter = MetadataFilter::new(&[col("max_ts").is_not_null()]); + assert!(!super::manifest_file_matches(&meta, &filter)); + } + #[test] fn test_region_state_lock_free() { assert!(AtomicCell::::is_lock_free()); @@ -2092,6 +2269,106 @@ mod tests { } } + #[tokio::test] + async fn test_manifest_stream_does_not_pin_sst_handles() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_recordbatch::filter::MetadataFilter; + use datafusion_expr::{col, lit}; + use futures::StreamExt; + + use super::ManifestSstProjection; + use crate::sst::file::FileMeta; + use crate::sst::file_purger::FilePurger; + + #[derive(Debug, Default)] + struct CountingPurger(AtomicUsize); + + impl FilePurger for CountingPurger { + fn remove_file(&self, _: FileMeta, is_delete: bool, _: bool) { + if is_delete { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + } + + for include_visible in [false, true] { + let env = SchedulerEnv::new().await; + let region = build_test_region(&env).await; + let purger = Arc::new(CountingPurger::default()); + let files = (0..2) + .map(|_| FileMeta { + region_id: region.region_id, + file_id: FileId::random(), + ..Default::default() + }) + .collect::>(); + let mut edit = empty_edit(); + edit.files_to_add = files.clone(); + region + .manifest_ctx + .manifest_manager + .write() + .await + .update( + RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone())), + false, + ) + .await + .unwrap(); + region + .version_control + .apply_edit(Some(edit), &[], purger.clone()); + + let schema = store_api::sst_entry::ManifestSstEntry::schema(); + let columns = if include_visible { + vec!["file_id", "visible"] + } else { + vec!["file_id"] + }; + let projection = columns + .iter() + .map(|name| { + schema + .column_schemas() + .iter() + .position(|column| column.name == *name) + .unwrap() + }) + .collect::>(); + let projection = ManifestSstProjection::new(Some(&projection)); + let snapshot = region.manifest_visible_file_ids(projection); + assert_eq!(snapshot.is_some(), include_visible); + if let Some(snapshot) = snapshot { + assert_eq!(snapshot, files.iter().map(|file| file.file_id).collect()); + } + // A direct caller may omit a residual column. The source must not + // interpret an unmaterialized visible value as grounds for pruning. + let filters = MetadataFilter::new(&[col("visible").eq(lit(true))]); + let mut stream = region.manifest_sst_stream(filters, projection); + let first = stream.next().await.unwrap(); + if include_visible { + assert!(first.visible); + } + let mut removal = empty_edit(); + removal.files_to_remove = files; + region + .version_control + .apply_edit(Some(removal), &[], purger.clone()); + // Both final FileHandles must reach the purger while the stream is + // suspended at its first yield, rather than waiting for cancellation. + assert_eq!(purger.0.load(Ordering::SeqCst), 2); + let second = stream.next().await.unwrap(); + if include_visible { + assert!( + second.visible, + "visibility must retain the original snapshot" + ); + } + assert!(stream.next().await.is_none()); + } + } + fn empty_edit() -> RegionEdit { RegionEdit { files_to_add: Vec::new(), diff --git a/src/store-api/src/sst_entry.rs b/src/store-api/src/sst_entry.rs index c70e909ebd6..6017f26bed3 100644 --- a/src/store-api/src/sst_entry.rs +++ b/src/store-api/src/sst_entry.rs @@ -24,6 +24,7 @@ use datatypes::arrow::array::{ ArrayRef, BinaryArray, BooleanArray, TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt32Array, UInt64Array, }; +use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef; use datatypes::arrow::error::ArrowError; use datatypes::arrow_array::StringArray; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; @@ -31,6 +32,34 @@ use serde::{Deserialize, Serialize}; use crate::storage::{RegionGroup, RegionId, RegionNumber, RegionSeq, ScanRequest, TableId}; +// Manifest SST column names. +pub const TABLE_DIR: &str = "table_dir"; +pub const REGION_ID: &str = "region_id"; +pub const TABLE_ID: &str = "table_id"; +pub const REGION_NUMBER: &str = "region_number"; +pub const REGION_GROUP: &str = "region_group"; +pub const REGION_SEQUENCE: &str = "region_sequence"; +pub const FILE_ID: &str = "file_id"; +pub const INDEX_VERSION: &str = "index_version"; +pub const LEVEL: &str = "level"; +pub const FILE_PATH: &str = "file_path"; +pub const FILE_SIZE: &str = "file_size"; +pub const INDEX_FILE_PATH: &str = "index_file_path"; +pub const INDEX_FILE_SIZE: &str = "index_file_size"; +pub const NUM_ROWS: &str = "num_rows"; +pub const NUM_ROW_GROUPS: &str = "num_row_groups"; +pub const NUM_SERIES: &str = "num_series"; +pub const MIN_TS: &str = "min_ts"; +pub const MAX_TS: &str = "max_ts"; +pub const SEQUENCE: &str = "sequence"; +pub const ORIGIN_REGION_ID: &str = "origin_region_id"; +pub const NODE_ID: &str = "node_id"; +pub const VISIBLE: &str = "visible"; +pub const PRIMARY_KEY_MIN: &str = "primary_key_min"; +pub const PRIMARY_KEY_MAX: &str = "primary_key_max"; +pub const MAX_ROW_GROUP_UNCOMPRESSED_SIZE: &str = "max_row_group_uncompressed_size"; +pub const PARTITION_EXPR: &str = "partition_expr"; + /// Puffin index type for bloom filter indexes. pub const PUFFIN_INDEX_TYPE_BLOOM_FILTER: &str = "bloom_filter"; /// Puffin index type for bloom-backed full-text indexes. @@ -102,110 +131,55 @@ impl ManifestSstEntry { pub fn schema() -> SchemaRef { use datatypes::prelude::ConcreteDataType as Ty; Arc::new(Schema::new(vec![ - ColumnSchema::new("table_dir", Ty::string_datatype(), false), - ColumnSchema::new("region_id", Ty::uint64_datatype(), false), - ColumnSchema::new("table_id", Ty::uint32_datatype(), false), - ColumnSchema::new("region_number", Ty::uint32_datatype(), false), - ColumnSchema::new("region_group", Ty::uint8_datatype(), false), - ColumnSchema::new("region_sequence", Ty::uint32_datatype(), false), - ColumnSchema::new("file_id", Ty::string_datatype(), false), - ColumnSchema::new("index_version", Ty::uint64_datatype(), false), - ColumnSchema::new("level", Ty::uint8_datatype(), false), - ColumnSchema::new("file_path", Ty::string_datatype(), false), - ColumnSchema::new("file_size", Ty::uint64_datatype(), false), - ColumnSchema::new("index_file_path", Ty::string_datatype(), true), - ColumnSchema::new("index_file_size", Ty::uint64_datatype(), true), - ColumnSchema::new("num_rows", Ty::uint64_datatype(), false), - ColumnSchema::new("num_row_groups", Ty::uint64_datatype(), false), - ColumnSchema::new("num_series", Ty::uint64_datatype(), true), - ColumnSchema::new("min_ts", Ty::timestamp_nanosecond_datatype(), true), - ColumnSchema::new("max_ts", Ty::timestamp_nanosecond_datatype(), true), - ColumnSchema::new("sequence", Ty::uint64_datatype(), true), - ColumnSchema::new("origin_region_id", Ty::uint64_datatype(), false), - ColumnSchema::new("node_id", Ty::uint64_datatype(), true), - ColumnSchema::new("visible", Ty::boolean_datatype(), false), - ColumnSchema::new("primary_key_min", Ty::binary_datatype(), true), - ColumnSchema::new("primary_key_max", Ty::binary_datatype(), true), + ColumnSchema::new(TABLE_DIR, Ty::string_datatype(), false), + ColumnSchema::new(REGION_ID, Ty::uint64_datatype(), false), + ColumnSchema::new(TABLE_ID, Ty::uint32_datatype(), false), + ColumnSchema::new(REGION_NUMBER, Ty::uint32_datatype(), false), + ColumnSchema::new(REGION_GROUP, Ty::uint8_datatype(), false), + ColumnSchema::new(REGION_SEQUENCE, Ty::uint32_datatype(), false), + ColumnSchema::new(FILE_ID, Ty::string_datatype(), false), + ColumnSchema::new(INDEX_VERSION, Ty::uint64_datatype(), false), + ColumnSchema::new(LEVEL, Ty::uint8_datatype(), false), + ColumnSchema::new(FILE_PATH, Ty::string_datatype(), false), + ColumnSchema::new(FILE_SIZE, Ty::uint64_datatype(), false), + ColumnSchema::new(INDEX_FILE_PATH, Ty::string_datatype(), true), + ColumnSchema::new(INDEX_FILE_SIZE, Ty::uint64_datatype(), true), + ColumnSchema::new(NUM_ROWS, Ty::uint64_datatype(), false), + ColumnSchema::new(NUM_ROW_GROUPS, Ty::uint64_datatype(), false), + ColumnSchema::new(NUM_SERIES, Ty::uint64_datatype(), true), + ColumnSchema::new(MIN_TS, Ty::timestamp_nanosecond_datatype(), true), + ColumnSchema::new(MAX_TS, Ty::timestamp_nanosecond_datatype(), true), + ColumnSchema::new(SEQUENCE, Ty::uint64_datatype(), true), + ColumnSchema::new(ORIGIN_REGION_ID, Ty::uint64_datatype(), false), + ColumnSchema::new(NODE_ID, Ty::uint64_datatype(), true), + ColumnSchema::new(VISIBLE, Ty::boolean_datatype(), false), + ColumnSchema::new(PRIMARY_KEY_MIN, Ty::binary_datatype(), true), + ColumnSchema::new(PRIMARY_KEY_MAX, Ty::binary_datatype(), true), ColumnSchema::new( - "max_row_group_uncompressed_size", + MAX_ROW_GROUP_UNCOMPRESSED_SIZE, Ty::uint64_datatype(), false, ), - ColumnSchema::new("partition_expr", Ty::string_datatype(), true), + ColumnSchema::new(PARTITION_EXPR, Ty::string_datatype(), true), ])) } /// Converts a list of manifest sst entries to a record batch. pub fn to_record_batch(entries: &[Self]) -> std::result::Result { - let schema = Self::schema(); - let table_dirs = entries.iter().map(|e| e.table_dir.as_str()); - let region_ids = entries.iter().map(|e| e.region_id.as_u64()); - let table_ids = entries.iter().map(|e| e.table_id); - let region_numbers = entries.iter().map(|e| e.region_number); - let region_groups = entries.iter().map(|e| e.region_group); - let region_sequences = entries.iter().map(|e| e.region_sequence); - let file_ids = entries.iter().map(|e| e.file_id.as_str()); - let index_versions = entries.iter().map(|e| e.index_version); - let levels = entries.iter().map(|e| e.level); - let file_paths = entries.iter().map(|e| e.file_path.as_str()); - let file_sizes = entries.iter().map(|e| e.file_size); - let max_row_group_uncompressed_sizes = - entries.iter().map(|e| e.max_row_group_uncompressed_size); - let index_file_paths = entries.iter().map(|e| e.index_file_path.as_ref()); - let index_file_sizes = entries.iter().map(|e| e.index_file_size); - let num_rows = entries.iter().map(|e| e.num_rows); - let num_row_groups = entries.iter().map(|e| e.num_row_groups); - let num_series = entries.iter().map(|e| e.num_series); - let min_ts = entries.iter().map(|e| { - e.min_ts - .convert_to(TimeUnit::Nanosecond) - .map(|ts| ts.value()) - }); - let max_ts = entries.iter().map(|e| { - e.max_ts - .convert_to(TimeUnit::Nanosecond) - .map(|ts| ts.value()) - }); - let sequences = entries.iter().map(|e| e.sequence); - let partition_exprs = entries.iter().map(|e| e.partition_expr.as_ref()); - let origin_region_ids = entries.iter().map(|e| e.origin_region_id.as_u64()); - let node_ids = entries.iter().map(|e| e.node_id); - let visible_flags = entries.iter().map(|e| Some(e.visible)); - let primary_key_min = entries.iter().map(|e| e.primary_key_min.as_deref()); - let primary_key_max = entries.iter().map(|e| e.primary_key_max.as_deref()); + Self::to_projected_record_batch(entries, None) + } - let columns: Vec = vec![ - Arc::new(StringArray::from_iter_values(table_dirs)), - Arc::new(UInt64Array::from_iter_values(region_ids)), - Arc::new(UInt32Array::from_iter_values(table_ids)), - Arc::new(UInt32Array::from_iter_values(region_numbers)), - Arc::new(UInt8Array::from_iter_values(region_groups)), - Arc::new(UInt32Array::from_iter_values(region_sequences)), - Arc::new(StringArray::from_iter_values(file_ids)), - Arc::new(UInt64Array::from_iter(index_versions)), - Arc::new(UInt8Array::from_iter_values(levels)), - Arc::new(StringArray::from_iter_values(file_paths)), - Arc::new(UInt64Array::from_iter_values(file_sizes)), - Arc::new(StringArray::from_iter(index_file_paths)), - Arc::new(UInt64Array::from_iter(index_file_sizes)), - Arc::new(UInt64Array::from_iter_values(num_rows)), - Arc::new(UInt64Array::from_iter_values(num_row_groups)), - Arc::new(UInt64Array::from_iter(num_series)), - Arc::new(TimestampNanosecondArray::from_iter(min_ts)), - Arc::new(TimestampNanosecondArray::from_iter(max_ts)), - Arc::new(UInt64Array::from_iter(sequences)), - Arc::new(UInt64Array::from_iter_values(origin_region_ids)), - Arc::new(UInt64Array::from_iter(node_ids)), - Arc::new(BooleanArray::from_iter(visible_flags)), - Arc::new(BinaryArray::from_iter(primary_key_min)), - Arc::new(BinaryArray::from_iter(primary_key_max)), - Arc::new(UInt64Array::from_iter_values( - max_row_group_uncompressed_sizes, - )), - Arc::new(StringArray::from_iter(partition_exprs)), - ]; - - DfRecordBatch::try_new(schema.arrow_schema().clone(), columns) + /// Builds only the requested columns, including empty and repeated projections. + fn to_projected_record_batch( + entries: &[Self], + projection: Option<&[usize]>, + ) -> std::result::Result { + let schema = Self::schema().arrow_schema().clone(); + let schema = match projection { + Some(projection) => Arc::new(schema.project(projection)?), + None => schema, + }; + ManifestSstBatchBuilder::new(schema).build(entries) } /// Reserved internal inspect table name. @@ -227,6 +201,140 @@ impl ManifestSstEntry { } } +/// Converts manifest SST batches using a schema prepared once per scan. +/// Repeated projected columns share the same Arrow array. +pub struct ManifestSstBatchBuilder { + schema: ArrowSchemaRef, + duplicates: Vec>, +} + +impl ManifestSstBatchBuilder { + /// Creates a builder for a projected manifest SST schema. + pub fn new(schema: ArrowSchemaRef) -> Self { + let duplicates = schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + schema.fields()[..index] + .iter() + .position(|previous| previous == field) + }) + .collect(); + Self { schema, duplicates } + } + + /// Builds the projected columns, copying string contents only into Arrow buffers. + pub fn build( + &self, + entries: &[ManifestSstEntry], + ) -> std::result::Result { + let mut columns: Vec = Vec::with_capacity(self.schema.fields().len()); + for (index, field) in self.schema.fields().iter().enumerate() { + if let Some(previous) = self.duplicates[index] { + columns.push(columns[previous].clone()); + continue; + } + let array: ArrayRef = match field.name().as_str() { + TABLE_DIR => Arc::new(StringArray::from_iter_values( + entries.iter().map(|e| e.table_dir.as_str()), + )), + REGION_ID => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.region_id.as_u64()), + )), + TABLE_ID => Arc::new(UInt32Array::from_iter_values( + entries.iter().map(|e| e.table_id), + )), + REGION_NUMBER => Arc::new(UInt32Array::from_iter_values( + entries.iter().map(|e| e.region_number), + )), + REGION_GROUP => Arc::new(UInt8Array::from_iter_values( + entries.iter().map(|e| e.region_group), + )), + REGION_SEQUENCE => Arc::new(UInt32Array::from_iter_values( + entries.iter().map(|e| e.region_sequence), + )), + FILE_ID => Arc::new(StringArray::from_iter_values( + entries.iter().map(|e| e.file_id.as_str()), + )), + INDEX_VERSION => Arc::new(UInt64Array::from_iter( + entries.iter().map(|e| e.index_version), + )), + LEVEL => Arc::new(UInt8Array::from_iter_values( + entries.iter().map(|e| e.level), + )), + FILE_PATH => Arc::new(StringArray::from_iter_values( + entries.iter().map(|e| e.file_path.as_str()), + )), + FILE_SIZE => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.file_size), + )), + INDEX_FILE_PATH => Arc::new(StringArray::from_iter( + entries.iter().map(|e| e.index_file_path.as_ref()), + )), + INDEX_FILE_SIZE => Arc::new(UInt64Array::from_iter( + entries.iter().map(|e| e.index_file_size), + )), + NUM_ROWS => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.num_rows), + )), + NUM_ROW_GROUPS => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.num_row_groups), + )), + NUM_SERIES => { + Arc::new(UInt64Array::from_iter(entries.iter().map(|e| e.num_series))) + } + MIN_TS => Arc::new(TimestampNanosecondArray::from_iter(entries.iter().map( + |e| { + e.min_ts + .convert_to(TimeUnit::Nanosecond) + .map(|ts| ts.value()) + }, + ))), + MAX_TS => Arc::new(TimestampNanosecondArray::from_iter(entries.iter().map( + |e| { + e.max_ts + .convert_to(TimeUnit::Nanosecond) + .map(|ts| ts.value()) + }, + ))), + SEQUENCE => Arc::new(UInt64Array::from_iter(entries.iter().map(|e| e.sequence))), + ORIGIN_REGION_ID => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.origin_region_id.as_u64()), + )), + NODE_ID => Arc::new(UInt64Array::from_iter(entries.iter().map(|e| e.node_id))), + VISIBLE => Arc::new(BooleanArray::from_iter( + entries.iter().map(|e| Some(e.visible)), + )), + PRIMARY_KEY_MIN => Arc::new(BinaryArray::from_iter( + entries.iter().map(|e| e.primary_key_min.as_deref()), + )), + PRIMARY_KEY_MAX => Arc::new(BinaryArray::from_iter( + entries.iter().map(|e| e.primary_key_max.as_deref()), + )), + MAX_ROW_GROUP_UNCOMPRESSED_SIZE => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.max_row_group_uncompressed_size), + )), + PARTITION_EXPR => Arc::new(StringArray::from_iter( + entries.iter().map(|e| e.partition_expr.as_ref()), + )), + name => { + return Err(ArrowError::SchemaError(format!( + "Unknown manifest SST column: {name}" + ))); + } + }; + columns.push(array); + } + DfRecordBatch::try_new_with_options( + self.schema.clone(), + columns, + &datatypes::arrow::record_batch::RecordBatchOptions::new() + .with_row_count(Some(entries.len())), + ) + } +} + /// An entry describing a SST file listed from storage layer directly. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct StorageSstEntry { @@ -421,13 +529,25 @@ fn build_plan_helper( ) -> Result { let table_source = LogicalTableSource::new(schema.arrow_schema().clone()); + // Filters may reference columns omitted from the requested output. Apply + // them to the full schema and let the optimizer push the minimal projection. let projection = scan_request.projection; - let mut builder = LogicalPlanBuilder::scan(table_name, Arc::new(table_source), projection)?; + let mut builder = LogicalPlanBuilder::scan(table_name, Arc::new(table_source), None)?; for filter in scan_request.filters { builder = builder.filter(filter)?; } + if let Some(projection) = projection { + let projected = schema.arrow_schema().project(&projection)?; + builder = builder.project( + projected + .fields() + .iter() + .map(|field| datafusion_expr::col(field.name())), + )?; + } + if let Some(limit) = scan_request.limit { builder = builder.limit(0, Some(limit))?; } @@ -438,7 +558,7 @@ fn build_plan_helper( #[cfg(test)] mod tests { use datafusion_common::TableReference; - use datafusion_expr::{LogicalPlan, Operator, binary_expr, col, lit}; + use datafusion_expr::{Expr, LogicalPlan, Operator, binary_expr, col, lit}; use datatypes::arrow::array::{ Array, BinaryArray, TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt32Array, UInt64Array, @@ -523,6 +643,29 @@ mod tests { let schema = ManifestSstEntry::schema(); let batch = ManifestSstEntry::to_record_batch(&entries).unwrap(); + for projection in [vec![], vec![16, 6, 10], vec![1, 1], (0..26).rev().collect()] { + let projected = + ManifestSstEntry::to_projected_record_batch(&entries, Some(&projection)).unwrap(); + assert_eq!(projected, batch.project(&projection).unwrap()); + assert_eq!(projected.num_rows(), entries.len()); + } + assert!(ManifestSstEntry::to_projected_record_batch(&entries, Some(&[26])).is_err()); + + let projected_schema = Arc::new(batch.schema().project(&[6, 9, 6, 11, 25, 9]).unwrap()); + let builder = ManifestSstBatchBuilder::new(projected_schema.clone()); + for rows in [entries.as_slice(), &entries[..1], &entries[..0]] { + let projected = builder.build(rows).unwrap(); + assert_eq!( + projected, + ManifestSstEntry::to_record_batch(rows) + .unwrap() + .project(&[6, 9, 6, 11, 25, 9]) + .unwrap() + ); + assert!(Arc::ptr_eq(&projected.schema(), &projected_schema)); + assert!(Arc::ptr_eq(projected.column(0), projected.column(2))); + assert!(Arc::ptr_eq(projected.column(1), projected.column(5))); + } // Schema checks assert_eq!(schema.arrow_schema().fields().len(), batch.num_columns()); @@ -1021,35 +1164,46 @@ mod tests { #[test] fn test_manifest_build_plan() { - // Note: filter must reference a column in the projected schema - let projection = Some(vec![0, 1, 2]); let request = ScanRequest { - projection, - filters: vec![binary_expr(col("table_id"), Operator::Gt, lit(0))], + projection: Some(vec![0, 1]), + filters: vec![col("table_id").gt(lit(0u32))], limit: Some(5), ..Default::default() }; let plan = ManifestSstEntry::build_plan(request).unwrap(); - // Expect plan to be Filter -> Limit -> TableScan or Filter+Limit wrapped. - // We'll pattern match to reach TableScan and verify key fields. - let (scan, has_filter, has_limit) = extract_scan(&plan); - - assert!(has_filter); - assert!(has_limit); + // The hidden table_id predicate must run before projection and limit. + let LogicalPlan::Limit(limit) = &plan else { + panic!("expected limit") + }; + let LogicalPlan::Projection(projection) = limit.input.as_ref() else { + panic!("expected projection") + }; + let LogicalPlan::Filter(filter) = projection.input.as_ref() else { + panic!("expected filter") + }; + let LogicalPlan::TableScan(scan) = filter.input.as_ref() else { + panic!("expected table scan") + }; assert_eq!( scan.table_name, TableReference::bare(ManifestSstEntry::reserved_table_name_for_inspection()) ); - assert_eq!(scan.projection, Some(vec![0, 1, 2])); + assert_eq!(scan.projection, None); + assert_eq!( + filter.predicate, + Expr::Column(datafusion_common::Column::new( + Some(scan.table_name.clone()), + "table_id", + )) + .gt(lit(0u32)) + ); - // projected schema should match projection - let fields = scan.projected_schema.fields(); - assert_eq!(fields.len(), 3); + let fields = plan.schema().fields(); + assert_eq!(fields.len(), 2); assert_eq!(fields[0].name(), "table_dir"); assert_eq!(fields[1].name(), "region_id"); - assert_eq!(fields[2].name(), "table_id"); } #[test] @@ -1070,9 +1224,9 @@ mod tests { scan.table_name, TableReference::bare(StorageSstEntry::reserved_table_name_for_inspection()) ); - assert_eq!(scan.projection, Some(vec![0, 2])); + assert_eq!(scan.projection, None); - let fields = scan.projected_schema.fields(); + let fields = plan.schema().fields(); assert_eq!(fields.len(), 2); assert_eq!(fields[0].name(), "file_path"); assert_eq!(fields[1].name(), "last_modified_ms"); @@ -1091,6 +1245,7 @@ mod tests { let (scan, has_filter, _) = extract_scan(input); (scan, has_filter, true) } + LogicalPlan::Projection(projection) => extract_scan(&projection.input), LogicalPlan::TableScan(scan) => (scan, false, false), other => panic!("unexpected plan: {other:?}"), } diff --git a/tests/cases/distributed/information_schema/ssts_limit.result b/tests/cases/distributed/information_schema/ssts_limit.result index 1b21a84c7af..67c6abf3cb1 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.result +++ b/tests/cases/distributed/information_schema/ssts_limit.result @@ -95,6 +95,59 @@ FROM ( | | | +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +-- The frontend orders results from all datanodes before applying the global offset. +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY min_ts DESC +LIMIT 2 OFFSET 1; + ++-------------------------+-------------------------+ +| min_ts | max_ts | ++-------------------------+-------------------------+ +| 1970-01-01T00:00:00.002 | 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.001 | 1970-01-01T00:00:00.001 | ++-------------------------+-------------------------+ + +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +EXPLAIN SELECT min_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY min_ts DESC +LIMIT 1; + ++---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Limit: skip=0, fetch=1 | +| | Sort: information_schema.ssts_manifest.min_ts DESC NULLS FIRST | +| | Projection: information_schema.ssts_manifest.min_ts | +| | Filter: information_schema.ssts_manifest.region_id > UInt64(0) | +| | TableScan: information_schema.ssts_manifest, partial_filters=[information_schema.ssts_manifest.region_id > UInt64(0)] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [min_ts@0 DESC], fetch=1 | +| | SortExec: TopK(fetch=1), expr=[min_ts@0 DESC], preserve_partitioning=[true] | +| | FilterExec: region_id@0 > 0, projection=[min_ts@1] | +| | RepartitionExec: REDACTED +| | DistributedInspectExec: kind=SstManifest, scan=ScanRequest { projection: [1, 16], filters: [region_id > UInt64(0)] }, schema=Schema { fields: [Field { name: "region_id", data_type: UInt64 }, Field { name: "min_ts", data_type: Timestamp(Nanosecond, None), nullable: true }], metadata: {"greptime:version": "0"} } | +| | | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +-- Timestamp filtering happens before global ordering and offset. The sort key +-- must remain available even when it is absent from the output projection. +SELECT max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts DESC +LIMIT 1 OFFSET 1; + ++-------------------------+ +| max_ts | ++-------------------------+ +| 1970-01-01T00:00:00.002 | ++-------------------------+ + DROP TABLE ssts_limit_case; Affected Rows: 0 diff --git a/tests/cases/distributed/information_schema/ssts_limit.sql b/tests/cases/distributed/information_schema/ssts_limit.sql index d548bd283a5..e634f7a2378 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.sql +++ b/tests/cases/distributed/information_schema/ssts_limit.sql @@ -44,4 +44,26 @@ FROM ( LIMIT 1 ); +-- The frontend orders results from all datanodes before applying the global offset. +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY min_ts DESC +LIMIT 2 OFFSET 1; + +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +EXPLAIN SELECT min_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY min_ts DESC +LIMIT 1; + +-- Timestamp filtering happens before global ordering and offset. The sort key +-- must remain available even when it is absent from the output projection. +SELECT max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts DESC +LIMIT 1 OFFSET 1; + DROP TABLE ssts_limit_case; diff --git a/tests/cases/standalone/common/information_schema/ssts.result b/tests/cases/standalone/common/information_schema/ssts.result index d7c7814a995..d8149d27769 100644 --- a/tests/cases/standalone/common/information_schema/ssts.result +++ b/tests/cases/standalone/common/information_schema/ssts.result @@ -243,6 +243,84 @@ FROM ( | 1 | +-----------------------+ +-- Timestamp filters must run before ordering and limiting. Sort columns must +-- remain available even when they are absent from the output projection. +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts ASC NULLS LAST, max_ts ASC +LIMIT 2; + ++-------------------------+-------------------------+ +| min_ts | max_ts | ++-------------------------+-------------------------+ +| 1970-01-01T00:00:00.002 | 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.003 | 1970-01-01T00:00:00.003 | ++-------------------------+-------------------------+ + +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY max_ts DESC NULLS FIRST, min_ts DESC +LIMIT 2 OFFSET 1; + ++-------------------------+-------------------------+ +| min_ts | max_ts | ++-------------------------+-------------------------+ +| 1970-01-01T00:00:00.100 | 1970-01-01T00:00:00.200 | +| 1970-01-01T00:00:00.003 | 1970-01-01T00:00:00.003 | ++-------------------------+-------------------------+ + +SELECT visible +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts DESC +LIMIT 2; + ++---------+ +| visible | ++---------+ +| true | +| true | ++---------+ + +-- Table predicates share region metadata pruning, including when table_id +-- is not part of the output projection. +SELECT min_ts, max_ts FROM information_schema.ssts_manifest +WHERE table_id > 0 ORDER BY min_ts, max_ts LIMIT 2; + ++-------------------------+-------------------------+ +| min_ts | max_ts | ++-------------------------+-------------------------+ +| 1970-01-01T00:00:00.001 | 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | 1970-01-01T00:00:00.002 | ++-------------------------+-------------------------+ + +SELECT min_ts FROM information_schema.ssts_manifest +WHERE table_id = 0 ORDER BY min_ts LIMIT 1; + ++--------+ +| min_ts | ++--------+ ++--------+ + +-- A false region predicate and a zero limit must produce empty results. +SELECT min_ts FROM information_schema.ssts_manifest +WHERE region_id = 0 ORDER BY min_ts LIMIT 1; + ++--------+ +| min_ts | ++--------+ ++--------+ + +SELECT min_ts FROM information_schema.ssts_manifest +ORDER BY min_ts LIMIT 0; + ++--------+ +| min_ts | ++--------+ ++--------+ + DROP TABLE sst_case; Affected Rows: 0 diff --git a/tests/cases/standalone/common/information_schema/ssts.sql b/tests/cases/standalone/common/information_schema/ssts.sql index 7b40b0056bb..e630ab9a0ba 100644 --- a/tests/cases/standalone/common/information_schema/ssts.sql +++ b/tests/cases/standalone/common/information_schema/ssts.sql @@ -80,4 +80,39 @@ FROM ( LIMIT 1 ); +-- Timestamp filters must run before ordering and limiting. Sort columns must +-- remain available even when they are absent from the output projection. +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts ASC NULLS LAST, max_ts ASC +LIMIT 2; + +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY max_ts DESC NULLS FIRST, min_ts DESC +LIMIT 2 OFFSET 1; + +SELECT visible +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts DESC +LIMIT 2; + +-- Table predicates share region metadata pruning, including when table_id +-- is not part of the output projection. +SELECT min_ts, max_ts FROM information_schema.ssts_manifest +WHERE table_id > 0 ORDER BY min_ts, max_ts LIMIT 2; + +SELECT min_ts FROM information_schema.ssts_manifest +WHERE table_id = 0 ORDER BY min_ts LIMIT 1; + +-- A false region predicate and a zero limit must produce empty results. +SELECT min_ts FROM information_schema.ssts_manifest +WHERE region_id = 0 ORDER BY min_ts LIMIT 1; + +SELECT min_ts FROM information_schema.ssts_manifest +ORDER BY min_ts LIMIT 0; + DROP TABLE sst_case;