perf: stream and prune SST manifest scans (#9931)

* refactor: add reusable metadata pruning and projected SST batches

Signed-off-by: WenyXu <wenymedia@gmail.com>

* perf(mito2): stream and prune manifest SST metadata

Signed-off-by: WenyXu <wenymedia@gmail.com>

* perf(datanode): stream inspection queries through reusable providers

Signed-off-by: WenyXu <wenymedia@gmail.com>

* perf(catalog): push manifest filters through inspection queries

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor: align manifest column constants with module conventions

Signed-off-by: WenyXu <wenymedia@gmail.com>

* perf: reduce manifest inspection allocations

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor: narrow manifest inspection API visibility

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor: retain residual filters for manifest inspection

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-10-26 08:54:02 +00:00
committed by GitHub
parent 513b9cee58
commit c96cdd0026
17 changed files with 2054 additions and 177 deletions
+119
View File
@@ -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<Vec<datafusion::logical_expr::TableProviderFilterPushDown>> {
Ok(vec![datafusion::logical_expr::TableProviderFilterPushDown::Inexact; filters.len()])
}
async fn scan(
&self,
_state: &dyn datafusion::catalog::Session,
projection: Option<&Vec<usize>>,
filters: &[datafusion::logical_expr::Expr],
limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
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::<SortExec>();
if let Some(scan) = plan.downcast_ref::<DistributedInspectExec>() {
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)
);
}
}
}
@@ -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()
+37 -1
View File
@@ -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<T: Borrow<ScalarValue>>(
&self,
mut value: impl FnMut(&str) -> Option<T>,
) -> 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)]
@@ -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::<ScalarValue>));
// 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::<ScalarValue>));
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());
}
+9 -9
View File
@@ -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<Arc<dyn TableProvider>> {
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;
@@ -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<TaskContext>) -> 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<Vec<usize>>,
filters: Vec<Expr>,
batch_size: usize,
}
/// Adapts a lazy inspection source to DataFusion's streaming execution plan.
#[derive(Debug)]
pub(crate) struct InspectTable {
source: Arc<dyn InspectSource>,
}
impl InspectTable {
/// Wraps a source without opening a scan.
pub(crate) fn new(source: Arc<dyn InspectSource>) -> 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<Vec<TableProviderFilterPushDown>> {
Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
}
async fn scan(
&self,
state: &dyn Session,
projection: Option<&Vec<usize>>,
filters: &[Expr],
limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
// 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<TaskContext>) -> 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<dyn InspectSource>,
request: InspectScan,
}
impl PartitionStream for InspectPartition {
fn schema(&self) -> &SchemaRef {
&self.request.schema
}
fn execute(&self, ctx: Arc<TaskContext>) -> 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<AtomicUsize>,
rows: Arc<AtomicUsize>,
drops: Arc<AtomicUsize>,
}
struct ScanGuard(Arc<AtomicUsize>);
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<TaskContext>) -> 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<ArrayRef> = 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::<Int32Array>()
.unwrap()
.values()
.to_vec()
})
.collect::<Vec<_>>();
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::<Int32Array>()
.unwrap()
.values()
.to_vec()
})
.collect::<Vec<_>>();
// 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);
}
}
@@ -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<dyn TableProvider> = 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::<Vec<_>>()
.await;
assert_eq!(
batches.iter().map(Vec::len).sum::<usize>(),
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::<UInt64Array>()
.unwrap()
.value(0),
referring.as_u64()
);
assert_eq!(
actual
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap()
.value(0),
origin.as_u64()
);
assert!(
actual
.column(2)
.as_any()
.downcast_ref::<StringArray>()
.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);
}
+8 -3
View File
@@ -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(
+104 -12
View File
@@ -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<ManifestSstEntry> {
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<ManifestSstEntry>> {
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<ManifestSstEntry>> {
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<u64>,
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"
+244
View File
@@ -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::<Vec<_>>(),
)
.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::<Vec<_>>())
.await
.expect("filtered scan must not acquire an unrelated manifest lock");
assert_eq!(
batches.iter().map(Vec::len).sum::<usize>(),
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::<Vec<_>>(),
)
.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::<Vec<_>>();
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::<Vec<_>>(),
)
.await
.expect("hidden table predicate must prune unrelated manifests");
let mut actual = batches.into_iter().flatten().collect::<Vec<_>>();
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);
}
+98
View File
@@ -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::<Vec<_>>();
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)
+316 -39
View File
@@ -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<ManifestSstEntry> {
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::<HashSet<_>>();
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<ManifestSstEntry> {
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::<HashMap<_, _>>();
}
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<HashSet<FileId>> {
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<Option<P
}
}
/// Columns requiring allocations or snapshots during an inspection scan.
#[derive(Clone, Copy)]
pub(crate) struct ManifestSstProjection {
visible: bool,
table_dir: bool,
file_id: bool,
file_path: bool,
index_file_path: bool,
partition_expr: bool,
primary_key_min: bool,
primary_key_max: bool,
}
impl ManifestSstProjection {
/// Selects columns by index in [`ManifestSstEntry::schema`]; `None` selects all columns.
pub(crate) fn new(projection: Option<&[usize]>) -> 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::<RegionRoleState>::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::<Vec<_>>();
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::<Vec<_>>();
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(),
+268 -113
View File
@@ -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<DfRecordBatch, ArrowError> {
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<ArrayRef> = 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<DfRecordBatch, ArrowError> {
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<Option<usize>>,
}
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<DfRecordBatch, ArrowError> {
let mut columns: Vec<ArrayRef> = 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<LogicalPlan, DataFusionError> {
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:?}"),
}
@@ -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
@@ -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;
@@ -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
@@ -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;