diff --git a/src/cmd/src/datanode/objbench.rs b/src/cmd/src/datanode/objbench.rs index 248acc09bf6..aa26491c351 100644 --- a/src/cmd/src/datanode/objbench.rs +++ b/src/cmd/src/datanode/objbench.rs @@ -203,6 +203,7 @@ impl ObjbenchCommand { max_sequence: None, sst_write_format: FormatType::PrimaryKey, index_options: Default::default(), + index_build_inputs: None, index_config: mito_engine_config.index.clone(), inverted_index_config: MitoConfig::default().inverted_index, fulltext_index_config, diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 55e564fb4e7..3a572b395bd 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -44,6 +44,7 @@ use crate::metrics::{COMPACTION_STAGE_ELAPSED, FLUSH_ELAPSED}; use crate::read::FlatSource; use crate::region::options::IndexOptions; use crate::sst::file::{FileHandle, RegionFileId, RegionIndexId}; +use crate::sst::index::build::IndexBuildInputs; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::puffin_manager::{PuffinManagerFactory, SstPuffinManager}; use crate::sst::index::registry::IndexRegistry; @@ -441,6 +442,7 @@ impl AccessLayer { self.puffin_manager_factory.clone() }; let indexer_builder = IndexerBuilderImpl { + index_build_inputs: request.index_build_inputs, build_type: request.op_type.into(), metadata: request.metadata.clone(), puffin_manager: puffin_factory.build(store, path_provider.clone()), @@ -671,6 +673,8 @@ pub enum OperationType { /// Contents to build a SST. pub struct SstWriteRequest { + /// Optional preplanned index work and physical input files. + pub index_build_inputs: Option, pub op_type: OperationType, pub metadata: RegionMetadataRef, pub source: FlatSource, @@ -707,6 +711,7 @@ impl EncodedSstIndexBuildParams { intermediate_manager: IntermediateManager, ) -> IndexerBuilderImpl { IndexerBuilderImpl { + index_build_inputs: None, build_type: OperationType::Flush.into(), metadata: self.metadata, puffin_manager, @@ -1026,6 +1031,7 @@ mod tests { access_layer .write_sst( SstWriteRequest { + index_build_inputs: None, op_type, metadata, source: new_flat_source_from_record_batches(vec![ diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index 534f0a8c076..37a4964a20e 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -301,6 +301,7 @@ impl WriteCache { let store = self.file_cache.local_store(); let path_provider = WriteCachePathProvider::new(self.file_cache.clone()); let indexer = IndexerBuilderImpl { + index_build_inputs: write_request.index_build_inputs, build_type: write_request.op_type.into(), metadata: write_request.metadata.clone(), puffin_manager: self @@ -746,6 +747,7 @@ mod tests { ]); let write_request = SstWriteRequest { + index_build_inputs: None, op_type, metadata, source, @@ -863,6 +865,7 @@ mod tests { // Write to local cache and upload sst to mock remote store let write_request = SstWriteRequest { + index_build_inputs: None, op_type: OperationType::Flush, metadata, source, @@ -955,6 +958,7 @@ mod tests { // Write to local cache and upload sst to mock remote store let write_request = SstWriteRequest { + index_build_inputs: None, op_type: OperationType::Flush, metadata, source, diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index 6cd51dae9a1..5074dcedadb 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -653,6 +653,7 @@ impl SstMerger for DefaultSstMerger { .access_layer .write_sst( SstWriteRequest { + index_build_inputs: None, op_type: OperationType::Compact, metadata: region_metadata.clone(), source, diff --git a/src/mito2/src/compaction/multi_region/tests.rs b/src/mito2/src/compaction/multi_region/tests.rs index 11829dcca2e..3aadc435cb5 100644 --- a/src/mito2/src/compaction/multi_region/tests.rs +++ b/src/mito2/src/compaction/multi_region/tests.rs @@ -447,6 +447,7 @@ fn write_request( format: FormatType, ) -> SstWriteRequest { SstWriteRequest { + index_build_inputs: None, op_type: OperationType::Flush, metadata, source, diff --git a/src/mito2/src/compaction/multi_region/writer.rs b/src/mito2/src/compaction/multi_region/writer.rs index 3c266584167..737fae51e44 100644 --- a/src/mito2/src/compaction/multi_region/writer.rs +++ b/src/mito2/src/compaction/multi_region/writer.rs @@ -105,6 +105,7 @@ impl<'a> TargetWriter<'a> { bloom.mem_threshold_on_create = self.index_threshold(bloom.mem_threshold_on_create()); let source = output_projection(self.target, source, self.context.cancellation.clone())?; Ok(SstWriteRequest { + index_build_inputs: None, op_type: OperationType::Compact, metadata: self.target.metadata.clone(), source, diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index d0c19c13ec8..1857997fd64 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -851,6 +851,7 @@ impl RegionFlushTask { .map(|f| f == FormatType::Flat) .unwrap_or(self.engine_config.default_flat_format); SstWriteRequest { + index_build_inputs: None, op_type: OperationType::Flush, metadata: version.metadata.clone(), source, diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index 18903abc661..1a445b56faa 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -38,6 +38,11 @@ pub const STAGING_TYPE: &str = "index_staging"; /// Recycle bin type label. pub const RECYCLE_TYPE: &str = "recycle_bin"; +/// Discard reason when an operation cannot admit extension index work. +pub(crate) const INDEX_DISCARD_ADMISSION_REJECTED: &str = "admission_rejected"; +/// Discard reason when a creator exceeds its retained-memory allowance. +pub(crate) const INDEX_DISCARD_MEMORY_BUDGET_EXCEEDED: &str = "memory_budget_exceeded"; + // Write metrics. lazy_static! { /// Global write buffer size in bytes. @@ -505,7 +510,15 @@ lazy_static! { &[TYPE_LABEL], ) .unwrap(); - /// Gauge of index create memory usage. + /// Optional index kinds discarded by a resource policy, not a build error. + pub static ref INDEX_CREATE_DISCARDED_TOTAL: IntCounterVec = register_int_counter_vec!( + "greptime_index_create_discarded_total", + "index kinds omitted or discarded by resource policy", + &[TYPE_LABEL, "reason"], + ).unwrap(); + /// Gauge of index create memory usage. The `extension` type aggregates retained + /// memory across optional creators, not individual algorithms: one creator may + /// share state across several algorithms. pub static ref INDEX_CREATE_MEMORY_USAGE: IntGaugeVec = register_int_gauge_vec!( "greptime_index_create_memory_usage", "index create memory usage", diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index ba29db1736f..a29cc33af0c 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -147,7 +147,7 @@ fn deserialize_level<'de, D: serde::Deserializer<'de>>(deserializer: D) -> Resul } /// Type to store index types for a column. -pub type IndexTypes = SmallVec<[IndexType; 4]>; +pub type IndexTypes = SmallVec<[IndexType; 6]>; /// Cross-region file id. /// diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 635d21f1cb4..47f7834235f 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -13,18 +13,19 @@ // limitations under the License. pub(crate) mod bloom_filter; +pub mod build; mod column; #[cfg(test)] mod column_test; pub(crate) mod fulltext_index; -pub(crate) mod handler; +pub mod handler; mod indexer; pub mod intermediate; pub(crate) mod inverted_index; -pub(crate) mod kind; +pub mod kind; mod primary_key; pub mod puffin_manager; -pub(crate) mod registry; +pub mod registry; #[cfg(test)] mod sparse_test; mod statistics; @@ -36,6 +37,7 @@ use std::ops::ControlFlow; use std::sync::Arc; use bytes::Bytes; +use common_error::ext::BoxedError; use common_telemetry::{debug, error, info, warn}; use datatypes::arrow::record_batch::RecordBatch; use index::target::IndexTarget; @@ -44,7 +46,7 @@ use object_store::ObjectStore; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use puffin_manager::SstPuffinManager; use smallvec::SmallVec; -use snafu::ResultExt; +use snafu::{IntoError, ResultExt}; #[cfg(feature = "enterprise")] pub use statistics::{Statistics, TimerGuard}; use store_api::codec::PrimaryKeyEncoding; @@ -61,11 +63,12 @@ use crate::cache::{CacheManagerRef, CacheStrategy}; #[cfg(test)] use crate::config::{BloomFilterConfig, InvertedIndexConfig}; use crate::error::{ - BuildIndexAsyncSnafu, Error, InvalidConfigSnafu, JoinSnafu, RegionClosedSnafu, + BuildIndexAsyncSnafu, Error, ExternalSnafu, InvalidConfigSnafu, JoinSnafu, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, Result, }; use crate::metrics::{ - INDEX_ARTIFACT_CLEANUP_FAILURE_TOTAL, INDEX_CREATE_MEMORY_USAGE, INDEX_PUBLICATION_STALE_TOTAL, + INDEX_ARTIFACT_CLEANUP_FAILURE_TOTAL, INDEX_CREATE_DISCARDED_TOTAL, INDEX_CREATE_MEMORY_USAGE, + INDEX_DISCARD_MEMORY_BUDGET_EXCEEDED, INDEX_PUBLICATION_STALE_TOTAL, }; use crate::read::Batch; use crate::region::options::IndexOptions; @@ -78,7 +81,9 @@ use crate::request::{ WorkerRequestWithTime, }; use crate::schedule::scheduler::{Job, SchedulerRef}; -use crate::sst::file::{ColumnIndexMetadata, FileHandle, IndexType, RegionFileId, RegionIndexId}; +use crate::sst::file::{ + ColumnIndexMetadata, FileHandle, IndexType, IndexTypes, RegionFileId, RegionIndexId, +}; use crate::sst::file_purger::FilePurgerRef; #[cfg(feature = "testing")] use crate::sst::index::bloom_filter::creator::BloomFilterIndexer; @@ -90,6 +95,7 @@ use crate::sst::index::inverted_index::creator::InvertedIndexer; use crate::sst::parquet::SstInfo; use crate::worker::WorkerListener; +pub(crate) const TYPE_EXTENSION_INDEX: &str = "extension"; pub(crate) const TYPE_INVERTED_INDEX: &str = kind::IndexKind::Inverted.as_str(); pub(crate) const TYPE_FULLTEXT_INDEX: &str = kind::IndexKind::Fulltext.as_str(); pub(crate) const TYPE_BLOOM_FILTER_INDEX: &str = kind::IndexKind::Bloom.as_str(); @@ -161,13 +167,14 @@ pub(crate) fn trigger_index_background_download( /// One index successfully written into the enclosing Puffin file. #[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct CompletedEntry { +pub struct CompletedEntry { /// Indexed column. `None` is reserved for file-level built-in indexes, such as /// JSON-path-only indexes; these appear only in `available_indexes`. pub column_id: Option, /// Registered index identity. pub kind: crate::sst::index::kind::IndexKind, - /// Fulltext bloom version published in column metadata. + /// Fulltext bloom version published in column metadata. Extension format + /// versions belong to their Puffin blob properties and are not published here. pub format_version: u32, } @@ -182,6 +189,12 @@ impl CompletedEntry { } fn validate(&self) -> Result<()> { + if self.column_id.is_none() && !self.kind.is_builtin() { + return InvalidConfigSnafu { + reason: "extension index completion requires a column id".to_string(), + } + .fail(); + } if self.kind.index_type() == IndexType::FulltextIndex { self.fulltext_bloom_version()?; } @@ -197,24 +210,19 @@ pub struct IndexOutput { /// Index version. pub version: u64, /// Successfully written column indexes. - pub(crate) entries: Vec, + pub entries: Vec, } impl IndexOutput { - pub fn build_available_indexes(&self) -> SmallVec<[IndexType; 4]> { - // Preserve the established order independently of handler registration. - [ - IndexType::InvertedIndex, - IndexType::FulltextIndex, - IndexType::BloomFilterIndex, - ] - .into_iter() - .filter(|kind| { - self.entries - .iter() - .any(|entry| entry.kind.index_type() == *kind) - }) - .collect() + pub fn build_available_indexes(&self) -> IndexTypes { + let mut kinds = SmallVec::<[kind::IndexKind; 6]>::new(); + for entry in &self.entries { + if !kinds.contains(&entry.kind) { + kinds.push(entry.kind); + } + } + kinds.sort_unstable(); + kinds.into_iter().map(kind::IndexKind::index_type).collect() } /// Builds column metadata, rejecting values that cannot be represented on disk. @@ -251,11 +259,34 @@ pub struct Indexer { puffin_manager: Option, write_cache_enabled: bool, creators: Vec, + /// Shared retained-memory ceiling for all injected creators. + extension_memory_limit: Option, dense_pk_decoder: Option, intermediate_manager: Option, + /// First failure retained even if best-effort execution discards a creator. + failure: Option>, + /// Optional index kinds dropped by the operation's resource policy. + discarded: Vec, } impl Indexer { + /// Reports failures even when surviving creators produced a partial result. + pub fn check_success(&self) -> Result<(), Arc> { + self.failure.clone().map_or(Ok(()), Err) + } + + /// Index kinds intentionally discarded because their memory budget was exceeded. + pub fn discarded(&self) -> &[kind::IndexKind] { + &self.discarded + } + + /// Whether any active creator requires decoded rows. + pub fn requires_row_scan(&self) -> bool { + self.creators + .iter() + .any(|active| active.creator.requires_row_scan()) + } + /// Returns true if no index creator is active. pub(crate) fn is_empty(&self) -> bool { self.creators.is_empty() @@ -263,17 +294,23 @@ impl Indexer { #[cfg(test)] fn has_kind(&self, kind: kind::IndexKind) -> bool { - self.creators.iter().any(|active| active.kind == kind) + self.creators + .iter() + .any(|active| active.kinds.contains(&kind)) } - fn push_creator(&mut self, kind: kind::IndexKind, creator: Box) { + fn push_creator( + &mut self, + kinds: Vec, + creator: Box, + memory_limit: Option, + label: &'static str, + ) { self.creators.push(ActiveIndex { - kind, + kinds, creator, - memory: IndexMemoryUsage { - label: kind.as_str(), - bytes: 0, - }, + memory_limit, + memory: IndexMemoryUsage { label, bytes: 0 }, }); } @@ -289,10 +326,20 @@ impl Indexer { ..Default::default() }; if let Some(creator) = inverted_indexer { - indexer.push_creator(kind::IndexKind::Inverted, Box::new(creator)); + indexer.push_creator( + vec![kind::IndexKind::Inverted], + Box::new(creator), + None, + TYPE_INVERTED_INDEX, + ); } if let Some(creator) = bloom_filter_indexer { - indexer.push_creator(kind::IndexKind::Bloom, Box::new(creator)); + indexer.push_creator( + vec![kind::IndexKind::Bloom], + Box::new(creator), + None, + TYPE_BLOOM_FILTER_INDEX, + ); } indexer.prepare_dense_pk_decoder(metadata); indexer @@ -326,15 +373,17 @@ impl Indexer { Ok(()) } - /// Updates all active creators using the borrowed batch. + /// Updates creators that require rows using the borrowed batch. pub async fn update(&mut self, batch: &mut Batch) { self.do_update(batch).await; + self.enforce_memory_limits().await; self.flush_mem_metrics(); } - /// Updates all active creators using the borrowed flat batch. + /// Updates creators that require rows using the borrowed flat batch. pub async fn update_flat(&mut self, batch: &RecordBatch) { self.do_update_flat(batch).await; + self.enforce_memory_limits().await; self.flush_mem_metrics(); } @@ -351,6 +400,57 @@ impl Indexer { self.flush_mem_metrics(); } + fn record_failure(&mut self, source: BoxedError) { + self.failure.get_or_insert_with(|| { + Arc::new( + ExternalSnafu { + context: "index creation".to_string(), + } + .into_error(source), + ) + }); + } + + async fn enforce_memory_limits(&mut self) { + let mut position = 0; + let mut extension_bytes = 0usize; + while position < self.creators.len() { + let active = &self.creators[position]; + if active.memory_limit.is_none() + && (active.is_builtin() || self.extension_memory_limit.is_none()) + { + position += 1; + continue; + } + let bytes = active.creator.memory_usage(); + let retained = extension_bytes.saturating_add(bytes); + let exceeds_shared_limit = !active.is_builtin() + && self + .extension_memory_limit + .is_some_and(|limit| retained > limit); + if exceeds_shared_limit || active.memory_limit.is_some_and(|limit| bytes > limit) { + warn!( + "Discarding index exceeding memory budget, region_id: {}, file_id: {}, kinds: {:?}", + self.region_id, self.file_id, active.kinds + ); + for kind in &active.kinds { + INDEX_CREATE_DISCARDED_TOTAL + .with_label_values(&[kind.as_str(), INDEX_DISCARD_MEMORY_BUDGET_EXCEEDED]) + .inc(); + if !self.discarded.contains(kind) { + self.discarded.push(*kind); + } + } + self.abort_creator(position).await; + } else { + if !active.is_builtin() { + extension_bytes = retained; + } + position += 1; + } + } + } + fn flush_mem_metrics(&mut self) { for active in &mut self.creators { let bytes = active.creator.memory_usage(); @@ -364,11 +464,18 @@ impl Indexer { /// One handler's active creator and its shared memory accounting. struct ActiveIndex { - kind: kind::IndexKind, + kinds: Vec, creator: Box, + memory_limit: Option, memory: IndexMemoryUsage, } +impl ActiveIndex { + fn is_builtin(&self) -> bool { + self.kinds.iter().all(|kind| kind.is_builtin()) + } +} + /// Releases each creator's gauge contribution on completion or cancellation. struct IndexMemoryUsage { label: &'static str, @@ -393,9 +500,9 @@ pub trait IndexerBuilder { row_group_size: Option, ) -> Indexer; } - #[derive(Clone)] pub(crate) struct IndexerBuilderImpl { + pub(crate) index_build_inputs: Option, pub(crate) registry: registry::IndexRegistryRef, pub(crate) build_type: IndexBuildType, pub(crate) metadata: RegionMetadataRef, @@ -423,29 +530,98 @@ impl IndexerBuilder for IndexerBuilderImpl { ..Default::default() }; - let context = handler::PlanContext { - metadata: &self.metadata, - index_options: &self.index_options, - build_type: self.build_type.clone(), + let plan = if let Some(inputs) = &self.index_build_inputs { + inputs.plan.clone() + } else { + let context = handler::PlanContext { + metadata: &self.metadata, + index_options: &self.index_options, + build_type: self.build_type.clone(), + existing_indexes: &[], + }; + Arc::new(self.registry.plan(&context)) }; - let plan = self.registry.plan(&context); - let context = handler::BuildContext { + indexer.failure = plan.check_success().err(); + let mut context = handler::BuildContext { metadata: &self.metadata, output: RegionIndexId::new(region_file_id, index_version), + puffin_manager: &self.puffin_manager, intermediate_manager: &self.intermediate_manager, + inputs: self + .index_build_inputs + .as_ref() + .map_or(&[], |inputs| &inputs.inputs), + input_puffin_manager: self + .index_build_inputs + .as_ref() + .map_or(&self.puffin_manager, |inputs| &inputs.input_puffin_manager), + can_merge_inputs: self + .index_build_inputs + .as_ref() + .is_some_and(|inputs| inputs.can_merge_inputs), + memory_limit: self + .index_build_inputs + .as_ref() + .and_then(|inputs| inputs.memory_limit), row_group_size, + output_rows: self + .index_build_inputs + .as_ref() + .and_then(|inputs| inputs.output_rows), }; + indexer.extension_memory_limit = context.memory_limit; for planned in plan.planned() { + let Some(first_kind) = planned.kinds().first() else { + let err = BoxedError::new( + InvalidConfigSnafu { + reason: "index plan must own at least one kind", + } + .build(), + ); + warn!(err; "Failed to create index, region_id: {}", self.metadata.region_id); + indexer.record_failure(err); + continue; + }; + let builtin = planned.kinds().iter().all(|kind| kind.is_builtin()); + // Optional creators share the operation allocation. Pass only the + // remaining retained budget before allowing another creator to build. + context.memory_limit = if builtin { + indexer.extension_memory_limit + } else { + let remaining = indexer.extension_memory_limit.map(|limit| { + let retained = indexer + .creators + .iter() + .filter(|active| !active.is_builtin()) + .fold(0usize, |bytes, active| { + bytes.saturating_add(active.creator.memory_usage()) + }); + limit.saturating_sub(retained) + }); + remaining.into_iter().chain(planned.memory_limit()).min() + }; match planned.build(&context).await { Ok(Some(creator)) => { - indexer.push_creator(planned.kind(), creator); + let label = if builtin { + first_kind.as_str() + } else { + TYPE_EXTENSION_INDEX + }; + let limit = if builtin { None } else { context.memory_limit }; + indexer.push_creator(planned.kinds().to_vec(), creator, limit, label); + if !builtin { + indexer.enforce_memory_limits().await; + } } Ok(None) => {} Err(err) => { - if cfg!(any(test, feature = "test")) { + if planned.kinds().iter().all(|kind| kind.is_builtin()) + && cfg!(any(test, feature = "test")) + { panic!("Failed to create index: {err:?}"); } - warn!(err; "Failed to create index, region_id: {}, kind: {:?}", self.metadata.region_id, planned.kind()); + warn!(err; "Failed to create index, region_id: {}, kinds: {:?}", self.metadata.region_id, planned.kinds()); + indexer.record_failure(err); } } } @@ -1649,6 +1825,7 @@ mod tests { let mut index_config = MitoConfig::default().index; index_config.build_mode = build_mode; let write_request = SstWriteRequest { + index_build_inputs: None, op_type: OperationType::Flush, metadata: metadata.clone(), source, @@ -1701,6 +1878,7 @@ mod tests { ), ); Arc::new(IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata, puffin_manager, @@ -1715,6 +1893,786 @@ mod tests { }) } + #[tokio::test] + async fn test_builtin_output_with_missing_or_failing_provider() { + use datatypes::arrow::array::{ + Float64Array, Int64Array, StringArray, TimestampMillisecondArray, UInt8Array, + UInt64Array, + }; + + for failure in [ + None, + Some(ProviderFailure::Plan), + Some(ProviderFailure::Finish), + ] { + let (dir, factory) = + PuffinManagerFactory::new_for_test_async("builtin_provider_failure").await; + let mut metadata = mock_region_metadata(MetaConfig { + with_inverted: true, + with_fulltext: true, + with_skipping_bloom: true, + }); + Arc::make_mut(&mut metadata).column_metadatas[0] + .column_schema + .mut_metadata() + .insert( + datatypes::schema::HLL_KEY.into(), + "{\"precision\":10}".into(), + ); + let registry = + registry::IndexRegistry::new( + Default::default(), + Default::default(), + Default::default(), + ) + .with_provider(failure.map(|failure| { + Arc::new(FailingProvider(failure)) as handler::IndexKindProviderRef + })) + .unwrap(); + let builder = IndexerBuilderImpl { + index_build_inputs: None, + build_type: IndexBuildType::Compact, + metadata, + puffin_manager: factory.build( + mock_object_store(), + RegionFilePathFactory::new( + "builtin-provider-test/".into(), + store_api::region_request::PathType::Bare, + ), + ), + write_cache_enabled: false, + intermediate_manager: mock_intm_mgr(dir.path().to_string_lossy()).await, + index_options: IndexOptions::default(), + registry: Arc::new(registry), + }; + let mut indexer = builder.build(random_region_file_id(), 0, Some(1024)).await; + let mut batch = crate::read::BatchBuilder::new(Vec::new()); + batch + .push_field_array(1, Arc::new(Int64Array::from(vec![42]))) + .unwrap(); + batch + .push_field_array(2, Arc::new(Float64Array::from(vec![1.0]))) + .unwrap(); + batch + .push_field_array(4, Arc::new(StringArray::from(vec!["hello world"]))) + .unwrap(); + batch + .push_field_array(5, Arc::new(StringArray::from(vec!["hello"]))) + .unwrap(); + batch + .timestamps_array(Arc::new(TimestampMillisecondArray::from(vec![0]))) + .unwrap(); + batch + .sequences_array(Arc::new(UInt64Array::from(vec![0]))) + .unwrap(); + batch + .op_types_array(Arc::new(UInt8Array::from(vec![0]))) + .unwrap(); + indexer.update(&mut batch.build().unwrap()).await; + let output = indexer.finish().await; + assert_eq!( + indexer.check_success().is_err(), + failure.is_some(), + "{failure:?}" + ); + if failure == Some(ProviderFailure::Finish) { + assert_eq!(output.file_size, 0); + assert!(output.entries.is_empty()); + } else { + assert!(output.file_size > 0, "{failure:?}"); + assert_eq!( + output.build_available_indexes().as_slice(), + &[ + crate::sst::file::IndexType::InvertedIndex, + crate::sst::file::IndexType::FulltextIndex, + crate::sst::file::IndexType::BloomFilterIndex, + ], + "{failure:?}" + ); + assert!(!output.build_indexes().unwrap().is_empty()); + } + } + } + + #[rstest::rstest] + #[case::batch(false)] + #[case::flat(true)] + #[tokio::test] + async fn test_mixed_reused_and_scanning_creators(#[case] flat: bool) { + use datatypes::arrow::array::{ArrayRef, Int64Array}; + use datatypes::vectors::{TimestampMillisecondVector, UInt8Vector, UInt64Vector}; + use index::inverted_index::format::reader::InvertedIndexBlobReader; + + let (dir, factory) = PuffinManagerFactory::new_for_test_async("mixed_reuse").await; + let manager = factory.build( + mock_object_store(), + RegionFilePathFactory::new( + "mixed-reuse/".into(), + store_api::region_request::PathType::Bare, + ), + ); + let builder = IndexerBuilderImpl { + index_build_inputs: None, + build_type: IndexBuildType::Compact, + metadata: mock_region_metadata(MetaConfig { + with_inverted: true, + with_fulltext: false, + with_skipping_bloom: false, + }), + puffin_manager: manager.clone(), + write_cache_enabled: false, + intermediate_manager: mock_intm_mgr(dir.path().to_string_lossy()).await, + index_options: IndexOptions::default(), + registry: Arc::new( + registry::IndexRegistry::new( + Default::default(), + Default::default(), + Default::default(), + ) + .with_provider(Some(Arc::new(FailingProvider(ProviderFailure::Reused)))) + .unwrap(), + ), + }; + let file = random_region_file_id(); + let mut indexer = builder.build(file, 0, Some(1024)).await; + assert!(indexer.requires_row_scan()); + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3])); + if flat { + let batch = RecordBatch::try_from_iter([("a", values)]).unwrap(); + indexer.update_flat(&batch).await; + } else { + let mut batch = crate::read::BatchBuilder::with_required_columns( + Vec::new(), + Arc::new(TimestampMillisecondVector::from_vec(vec![0, 1, 2])), + Arc::new(UInt64Vector::from_vec(vec![0; 3])), + Arc::new(UInt8Vector::from_vec(vec![0; 3])), + ); + batch.push_field_array(1, values).unwrap(); + indexer.update(&mut batch.build().unwrap()).await; + } + indexer.check_success().unwrap(); + let output = indexer.finish().await; + indexer.check_success().unwrap(); + assert_eq!( + output.build_available_indexes().as_slice(), + &[IndexType::InvertedIndex, IndexType::HllIndex] + ); + let reader = manager.reader(&RegionIndexId::new(file, 0)).await.unwrap(); + reader.blob("test-sketch").await.unwrap(); + let blob = reader.blob("greptime-inverted-index-v1").await.unwrap(); + let inverted = InvertedIndexBlobReader::new(blob.reader().await.unwrap()); + let metadata = inverted.metadata(None).await.unwrap(); + assert_eq!(metadata.total_row_count, 3); + } + + #[derive(Debug, Default, Clone)] + struct TestProvider { + build_calls: Arc, + kind: Option, + received_limits: Arc>>>, + hard_limit: Option, + } + + impl handler::IndexKindProvider for TestProvider { + fn handlers(&self) -> Vec> { + vec![Arc::new(self.clone())] + } + } + + impl handler::IndexKindHandler for TestProvider { + fn kinds(&self) -> &[kind::IndexKind] { + match self.kind { + Some(kind::ExtensionIndexKind::Kll) => { + &[kind::IndexKind::Extension(kind::ExtensionIndexKind::Kll)] + } + _ => &[kind::IndexKind::Extension(kind::ExtensionIndexKind::Hll)], + } + } + fn plan( + &self, + _context: &handler::PlanContext<'_>, + ) -> std::result::Result< + Option>, + common_error::ext::BoxedError, + > { + Ok(Some(Arc::new(self.clone()))) + } + } + + #[async_trait::async_trait] + impl handler::PlannedIndex for TestProvider { + fn kinds(&self) -> &[kind::IndexKind] { + handler::IndexKindHandler::kinds(self) + } + fn entries(&self) -> &[handler::PlanEntry] { + match self.kind { + Some(kind::ExtensionIndexKind::Kll) => &[handler::PlanEntry { + column_id: 0, + kind: kind::IndexKind::Extension(kind::ExtensionIndexKind::Kll), + format_version: 2, + }], + _ => &[handler::PlanEntry { + column_id: 0, + kind: kind::IndexKind::Extension(kind::ExtensionIndexKind::Hll), + format_version: 2, + }], + } + } + + fn memory_limit(&self) -> Option { + self.hard_limit + } + + async fn build( + &self, + context: &handler::BuildContext<'_>, + ) -> std::result::Result< + Option>, + common_error::ext::BoxedError, + > { + self.build_calls + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.received_limits + .lock() + .unwrap() + .push(context.memory_limit); + Ok(Some(Box::new(TestCreator))) + } + } + + #[derive(Debug)] + struct TestCreator; + + #[async_trait::async_trait] + impl handler::IndexCreator for TestCreator { + async fn update( + &mut self, + _batch: &mut Batch, + ) -> std::result::Result<(), common_error::ext::BoxedError> { + Ok(()) + } + + async fn update_flat( + &mut self, + _batch: &RecordBatch, + ) -> std::result::Result<(), common_error::ext::BoxedError> { + Ok(()) + } + + async fn finish( + &mut self, + _writer: &mut puffin_manager::SstPuffinWriter, + ) -> std::result::Result, common_error::ext::BoxedError> { + Ok(Vec::new()) + } + + async fn abort(&mut self) -> std::result::Result<(), common_error::ext::BoxedError> { + Ok(()) + } + + fn memory_usage(&self) -> usize { + 17 + } + } + + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + enum ProviderFailure { + None, + Plan, + EmptyKinds, + Reused, + Budget, + Create, + Update, + UpdateFlat, + Finish, + Close, + } + + #[derive(Debug, Clone)] + struct FailingProvider(ProviderFailure); + + fn provider_test_error() -> common_error::ext::BoxedError { + common_error::ext::BoxedError::new(common_error::ext::PlainError::new( + "injected provider failure".into(), + common_error::status_code::StatusCode::Internal, + )) + } + + impl handler::IndexKindProvider for FailingProvider { + fn handlers(&self) -> Vec> { + vec![Arc::new(self.clone())] + } + } + + impl handler::IndexKindHandler for FailingProvider { + fn kinds(&self) -> &[kind::IndexKind] { + &[kind::IndexKind::Extension(kind::ExtensionIndexKind::Hll)] + } + fn plan( + &self, + context: &handler::PlanContext<'_>, + ) -> std::result::Result< + Option>, + common_error::ext::BoxedError, + > { + if self.0 == ProviderFailure::Plan { + return Err(provider_test_error()); + } + if context.build_type == IndexBuildType::Flush { + return Ok(None); + } + Ok(Some(Arc::new(self.clone()))) + } + } + + #[async_trait::async_trait] + impl handler::PlannedIndex for FailingProvider { + fn kinds(&self) -> &[kind::IndexKind] { + if self.0 == ProviderFailure::EmptyKinds { + return &[]; + } + handler::IndexKindHandler::kinds(self) + } + fn entries(&self) -> &[handler::PlanEntry] { + &[handler::PlanEntry { + column_id: 0, + kind: kind::IndexKind::Extension(kind::ExtensionIndexKind::Hll), + format_version: 2, + }] + } + + fn memory_limit(&self) -> Option { + (self.0 == ProviderFailure::Budget).then_some(16) + } + + async fn build( + &self, + _context: &handler::BuildContext<'_>, + ) -> std::result::Result< + Option>, + common_error::ext::BoxedError, + > { + if self.0 == ProviderFailure::Create { + return Err(provider_test_error()); + } + Ok(Some(Box::new(FailingCreator(self.0)))) + } + } + + #[derive(Debug)] + struct FailingCreator(ProviderFailure); + + #[async_trait::async_trait] + impl handler::IndexCreator for FailingCreator { + fn requires_row_scan(&self) -> bool { + self.0 != ProviderFailure::Reused + } + + async fn update( + &mut self, + _batch: &mut Batch, + ) -> std::result::Result<(), common_error::ext::BoxedError> { + if matches!(self.0, ProviderFailure::Update | ProviderFailure::Reused) { + Err(provider_test_error()) + } else { + Ok(()) + } + } + + async fn update_flat( + &mut self, + _batch: &RecordBatch, + ) -> std::result::Result<(), common_error::ext::BoxedError> { + if matches!( + self.0, + ProviderFailure::UpdateFlat | ProviderFailure::Reused + ) { + Err(provider_test_error()) + } else { + Ok(()) + } + } + + async fn finish( + &mut self, + writer: &mut puffin_manager::SstPuffinWriter, + ) -> std::result::Result, common_error::ext::BoxedError> { + use puffin::puffin_manager::PuffinWriter; + + writer + .put_blob( + "test-sketch", + futures::io::Cursor::new(vec![1_u8, 2, 3]), + puffin::puffin_manager::PutOptions::default(), + Default::default(), + ) + .await + .map_err(common_error::ext::BoxedError::new)?; + if self.0 == ProviderFailure::Finish { + // Fail after writing to exercise partially written artifact handling. + return Err(provider_test_error()); + } + Ok(vec![CompletedEntry { + column_id: Some(0), + kind: kind::IndexKind::Extension(kind::ExtensionIndexKind::Hll), + format_version: 2, + }]) + } + + async fn abort(&mut self) -> std::result::Result<(), common_error::ext::BoxedError> { + Ok(()) + } + + fn memory_usage(&self) -> usize { + 17 + } + } + + struct FailingIndexClose { + inner: object_store::layers::mock::Writer, + attempts: Arc, + } + + impl object_store::layers::mock::Write for FailingIndexClose { + async fn write( + &mut self, + bytes: object_store::layers::mock::Buffer, + ) -> object_store::layers::mock::Result<()> { + self.inner.write(bytes).await + } + + async fn close( + &mut self, + ) -> object_store::layers::mock::Result { + self.attempts + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + Err(object_store::layers::mock::Error::new( + object_store::layers::mock::ErrorKind::Unexpected, + "injected Puffin close failure", + )) + } + + async fn abort(&mut self) -> object_store::layers::mock::Result<()> { + self.inner.abort().await + } + } + + #[tokio::test] + async fn test_provider_failures_do_not_publish_sketch_descriptors() { + for failure in [ + ProviderFailure::None, + ProviderFailure::Plan, + ProviderFailure::EmptyKinds, + ProviderFailure::Create, + ProviderFailure::Update, + ProviderFailure::UpdateFlat, + ProviderFailure::Finish, + ProviderFailure::Close, + ] { + let (dir, factory) = PuffinManagerFactory::new_for_test_async("provider_failure").await; + let close_attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let store = if failure == ProviderFailure::Close { + let attempts = close_attempts.clone(); + let layer = object_store::layers::mock::MockLayerBuilder::default() + .writer_factory(Arc::new(move |_, _, inner| { + Box::new(FailingIndexClose { + inner, + attempts: attempts.clone(), + }) + })) + .build() + .unwrap(); + mock_object_store().layer(layer) + } else { + mock_object_store() + }; + let manager = factory.build( + store, + crate::access_layer::RegionFilePathFactory::new( + "provider-test/".into(), + store_api::region_request::PathType::Bare, + ), + ); + let builder = IndexerBuilderImpl { + index_build_inputs: None, + build_type: IndexBuildType::Compact, + metadata: mock_region_metadata(MetaConfig { + with_inverted: false, + with_fulltext: false, + with_skipping_bloom: false, + }), + puffin_manager: manager, + write_cache_enabled: false, + intermediate_manager: mock_intm_mgr(dir.path().to_string_lossy()).await, + index_options: IndexOptions::default(), + registry: Arc::new( + registry::IndexRegistry::new( + InvertedIndexConfig::default(), + FulltextIndexConfig::default(), + BloomFilterConfig::default(), + ) + .with_provider(Some(Arc::new(FailingProvider(failure)))) + .unwrap(), + ), + }; + let mut indexer = builder.build(random_region_file_id(), 0, Some(1024)).await; + if failure == ProviderFailure::Update { + let mut batch = crate::read::BatchBuilder::with_required_columns( + Vec::new(), + Arc::new(datatypes::vectors::TimestampMillisecondVector::from_vec( + vec![0], + )), + Arc::new(datatypes::vectors::UInt64Vector::from_vec(vec![0])), + Arc::new(datatypes::vectors::UInt8Vector::from_vec(vec![0])), + ) + .build() + .unwrap(); + indexer.update(&mut batch).await; + } else { + indexer + .update_flat(&new_record_batch_by_range(&["a", "b"], 0, 1)) + .await; + } + let output = indexer.finish().await; + if failure == ProviderFailure::None { + assert_eq!( + output + .entries + .iter() + .filter(|entry| !entry.kind.is_builtin()) + .count(), + 1 + ); + assert!(output.file_size > 0); + assert!(output.build_indexes().unwrap().iter().any(|index| { + index.created_indexes.iter().any(|kind| { + matches!( + kind, + crate::sst::file::IndexType::HllIndex + | crate::sst::file::IndexType::KllIndex + | crate::sst::file::IndexType::MgIndex + ) + }) + })); + } else { + assert!( + output.entries.iter().all(|entry| entry.kind.is_builtin()), + "{failure:?}" + ); + assert!(output.build_indexes().unwrap().is_empty(), "{failure:?}"); + assert_eq!(output.file_size, 0, "{failure:?}"); + } + assert!(indexer.creators.is_empty()); + assert_eq!( + indexer.check_success().is_err(), + failure != ProviderFailure::None + ); + assert_eq!( + close_attempts.load(std::sync::atomic::Ordering::Relaxed), + usize::from(failure == ProviderFailure::Close) + ); + } + } + + #[tokio::test] + async fn test_provider_compaction_only_and_per_output() { + let (dir, factory) = PuffinManagerFactory::new_for_test_async("provider_scope").await; + let provider = Arc::new(TestProvider { + hard_limit: Some(17), + ..Default::default() + }); + let mut builder = IndexerBuilderImpl { + index_build_inputs: None, + build_type: IndexBuildType::Flush, + metadata: mock_region_metadata(MetaConfig { + with_inverted: false, + with_fulltext: false, + with_skipping_bloom: false, + }), + puffin_manager: factory.build(mock_object_store(), NoopPathProvider), + write_cache_enabled: false, + intermediate_manager: mock_intm_mgr(dir.path().to_string_lossy()).await, + index_options: IndexOptions::default(), + registry: Arc::new( + registry::IndexRegistry::new( + InvertedIndexConfig::default(), + FulltextIndexConfig::default(), + BloomFilterConfig::default(), + ) + .with_provider(Some(provider.clone())) + .unwrap(), + ), + }; + let flush = builder.build(random_region_file_id(), 0, Some(1024)).await; + assert!(flush.creators.is_empty()); + assert_eq!( + provider + .build_calls + .load(std::sync::atomic::Ordering::Relaxed), + 0 + ); + + builder.build_type = IndexBuildType::Compact; + let mut first = builder.build(random_region_file_id(), 0, Some(1024)).await; + let mut second = builder.build(random_region_file_id(), 0, Some(1024)).await; + assert_eq!( + provider + .build_calls + .load(std::sync::atomic::Ordering::Relaxed), + 2 + ); + assert!(!first.is_empty()); + assert!(!second.is_empty()); + assert_eq!(first.creators[0].creator.memory_usage(), 17); + first.abort().await; + assert!(first.creators.is_empty()); + assert!(!second.creators.is_empty()); + second.abort().await; + + // Reject a creator at preparation when its retained state already exceeds + // the caller's allocation; equality is allowed. + let plan = builder.registry.plan(&handler::PlanContext { + metadata: &builder.metadata, + index_options: &builder.index_options, + build_type: builder.build_type.clone(), + existing_indexes: &[], + }); + builder.index_build_inputs = Some(build::IndexBuildInputs { + output_rows: None, + plan: Arc::new(plan), + inputs: Arc::from([]), + input_puffin_manager: builder.puffin_manager.clone(), + can_merge_inputs: false, + memory_limit: Some(16), + }); + let rejected = builder.build(random_region_file_id(), 0, Some(1024)).await; + assert!(rejected.is_empty()); + builder.index_build_inputs.as_mut().unwrap().memory_limit = Some(17); + let mut bounded = builder.build(random_region_file_id(), 0, Some(1024)).await; + assert!(!bounded.creators.is_empty()); + bounded.creators[0].memory_limit = Some(16); + bounded.enforce_memory_limits().await; + bounded.flush_mem_metrics(); + assert!(bounded.creators.is_empty()); + + // A larger operation allocation must not relax the handler's ceiling. + builder.index_build_inputs.as_mut().unwrap().memory_limit = Some(34); + let mut bounded = builder.build(random_region_file_id(), 0, Some(1024)).await; + assert_eq!(bounded.creators[0].memory_limit, Some(17)); + bounded.abort().await; + } + + #[tokio::test] + async fn test_builder_passes_remaining_extension_budget() { + let (dir, factory) = PuffinManagerFactory::new_for_test_async("build_budget").await; + for (operation_limit, ceiling, expected, retained) in [ + (Some(30), None, vec![Some(30), Some(13)], 1), + (Some(17), None, vec![Some(17), Some(0)], 1), + (Some(30), Some(20), vec![Some(20), Some(13)], 1), + (Some(40), Some(20), vec![Some(20), Some(20)], 2), + (Some(16), None, vec![Some(16), Some(16)], 0), + (None, None, vec![None, None], 2), + (None, Some(17), vec![Some(17), Some(17)], 2), + ] { + let observed = Arc::new(std::sync::Mutex::new(Vec::new())); + let mut registry = registry::IndexRegistry::new( + InvertedIndexConfig::default(), + FulltextIndexConfig::default(), + BloomFilterConfig::default(), + ); + for kind in [kind::ExtensionIndexKind::Hll, kind::ExtensionIndexKind::Kll] { + registry + .register(Arc::new(TestProvider { + kind: Some(kind), + received_limits: observed.clone(), + hard_limit: ceiling, + ..Default::default() + })) + .unwrap(); + } + let mut builder = IndexerBuilderImpl { + index_build_inputs: None, + build_type: IndexBuildType::Compact, + metadata: mock_region_metadata(MetaConfig { + with_inverted: false, + with_fulltext: false, + with_skipping_bloom: false, + }), + puffin_manager: factory.build(mock_object_store(), NoopPathProvider), + write_cache_enabled: false, + intermediate_manager: mock_intm_mgr(dir.path().to_string_lossy()).await, + index_options: IndexOptions::default(), + registry: Arc::new(registry), + }; + let plan = builder.registry.plan(&handler::PlanContext { + metadata: &builder.metadata, + index_options: &builder.index_options, + build_type: builder.build_type.clone(), + existing_indexes: &[], + }); + builder.index_build_inputs = Some(build::IndexBuildInputs { + output_rows: None, + plan: Arc::new(plan), + inputs: Arc::from([]), + input_puffin_manager: builder.puffin_manager.clone(), + can_merge_inputs: false, + memory_limit: operation_limit, + }); + let mut indexer = builder.build(random_region_file_id(), 0, Some(1024)).await; + assert_eq!(*observed.lock().unwrap(), expected); + assert_eq!(indexer.creators.len(), retained); + indexer.abort().await; + } + } + + #[tokio::test] + async fn test_extension_creators_share_memory_budget() { + for (limit, expected_extensions) in [(34, 2), (33, 1), (16, 0)] { + let mut indexer = Indexer { + extension_memory_limit: Some(limit), + ..Default::default() + }; + // Built-in state uses its own spill budget and does not consume the + // optional extension allocation. Each test creator retains 17 bytes. + indexer.push_creator( + vec![kind::IndexKind::Inverted], + Box::new(TestCreator), + None, + TYPE_INVERTED_INDEX, + ); + for kind in [kind::ExtensionIndexKind::Hll, kind::ExtensionIndexKind::Kll] { + indexer.push_creator( + vec![kind::IndexKind::Extension(kind)], + Box::new(TestCreator), + Some(34), + TYPE_EXTENSION_INDEX, + ); + } + let before = INDEX_CREATE_DISCARDED_TOTAL + .with_label_values(&["kll_index", "memory_budget_exceeded"]) + .get(); + indexer.enforce_memory_limits().await; + assert_eq!(indexer.creators.len(), 1 + expected_extensions); + assert!(indexer.check_success().is_ok()); + assert_eq!(indexer.discarded().len(), 2 - expected_extensions); + let kinds = [ + kind::IndexKind::Extension(kind::ExtensionIndexKind::Hll), + kind::IndexKind::Extension(kind::ExtensionIndexKind::Kll), + ]; + assert_eq!(indexer.discarded(), &kinds[expected_extensions..]); + if expected_extensions < 2 { + assert!( + INDEX_CREATE_DISCARDED_TOTAL + .with_label_values(&["kll_index", "memory_budget_exceeded"]) + .get() + > before + ); + } + assert!(indexer.creators[0].is_builtin()); + } + } + #[test] fn test_file_level_completion_preserves_json_only_index() { let output = IndexOutput { @@ -1732,6 +2690,50 @@ mod tests { assert!(output.build_indexes().unwrap().is_empty()); } + #[test] + fn test_completed_entries_preserve_all_kinds_in_order() { + use crate::sst::index::kind::{ExtensionIndexKind, IndexKind}; + + let output = IndexOutput { + entries: [ + IndexKind::Extension(ExtensionIndexKind::Mg), + IndexKind::Fulltext, + IndexKind::Extension(ExtensionIndexKind::Kll), + IndexKind::Bloom, + IndexKind::Extension(ExtensionIndexKind::Hll), + IndexKind::Inverted, + IndexKind::Extension(ExtensionIndexKind::Hll), + ] + .into_iter() + .enumerate() + .map(|(position, kind)| CompletedEntry { + column_id: Some(if position == 6 { 2 } else { 1 }), + kind, + format_version: 0, + }) + .collect(), + ..Default::default() + }; + let expected = [ + IndexType::InvertedIndex, + IndexType::FulltextIndex, + IndexType::BloomFilterIndex, + IndexType::HllIndex, + IndexType::KllIndex, + IndexType::MgIndex, + ]; + let available = output.build_available_indexes(); + assert_eq!(available.as_slice(), &expected); + assert!(!available.spilled()); + let columns = output.build_indexes().unwrap(); + assert_eq!(columns.len(), 2); + assert_eq!(columns[0].created_indexes.as_slice(), &expected); + assert_eq!( + columns[1].created_indexes.as_slice(), + &[IndexType::HllIndex] + ); + } + #[test] fn test_completed_entries_preserve_builtin_metadata() { use crate::sst::index::kind::IndexKind; @@ -1822,6 +2824,10 @@ mod tests { assert!(entry.validate().is_ok()); entry.format_version += 1; assert!(entry.validate().is_err()); + entry.kind = kind::IndexKind::Extension(kind::ExtensionIndexKind::Hll); + assert!(entry.validate().is_ok()); // Extension versions belong to Puffin. + entry.column_id = None; + assert!(entry.validate().is_err()); } #[tokio::test] @@ -1836,6 +2842,7 @@ mod tests { with_skipping_bloom: true, }); let indexer = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata, puffin_manager: factory.build(mock_object_store(), NoopPathProvider), @@ -1868,6 +2875,7 @@ mod tests { with_skipping_bloom: true, }); let indexer = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata: metadata.clone(), puffin_manager: factory.build(mock_object_store(), NoopPathProvider), @@ -1891,6 +2899,7 @@ mod tests { assert!(indexer.has_kind(kind::IndexKind::Bloom)); let indexer = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Compact, metadata: metadata.clone(), puffin_manager: factory.build(mock_object_store(), NoopPathProvider), @@ -1914,6 +2923,7 @@ mod tests { assert!(indexer.has_kind(kind::IndexKind::Bloom)); let indexer = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Compact, metadata, puffin_manager: factory.build(mock_object_store(), NoopPathProvider), @@ -1949,6 +2959,7 @@ mod tests { with_skipping_bloom: true, }); let indexer = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata: metadata.clone(), puffin_manager: factory.build(mock_object_store(), NoopPathProvider), @@ -1974,6 +2985,7 @@ mod tests { with_skipping_bloom: true, }); let indexer = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata: metadata.clone(), puffin_manager: factory.build(mock_object_store(), NoopPathProvider), @@ -1999,6 +3011,7 @@ mod tests { with_skipping_bloom: false, }); let indexer = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata: metadata.clone(), puffin_manager: factory.build(mock_object_store(), NoopPathProvider), @@ -2031,6 +3044,7 @@ mod tests { with_skipping_bloom: true, }); let indexer = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata, puffin_manager: factory.build(mock_object_store(), NoopPathProvider), @@ -2433,6 +3447,7 @@ mod tests { ); // Indexer builder built from write cache. let indexer_builder = Arc::new(IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata: metadata.clone(), puffin_manager: write_cache.build_puffin_manager().clone(), diff --git a/src/mito2/src/sst/index/build.rs b/src/mito2/src/sst/index/build.rs new file mode 100644 index 00000000000..863aa84b7a4 --- /dev/null +++ b/src/mito2/src/sst/index/build.rs @@ -0,0 +1,46 @@ +// 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. + +//! Shared index plans and physical inputs carried into SST writers. + +use std::sync::Arc; + +use crate::sst::file::FileHandle; +use crate::sst::index::puffin_manager::SstPuffinManager; +use crate::sst::index::registry::IndexBuildPlan; + +/// An immutable plan with operation-scoped input files and output constraints. +/// Currently compaction and explicit rebuild supply these inputs; flush does not +/// inject extension handlers and builds only the existing built-in indexes. +#[derive(Clone)] +pub struct IndexBuildInputs { + /// Typed plans shared by output files of this operation. + pub plan: Arc, + /// Pinned physical input files; flushes have no inputs. + pub inputs: Arc<[FileHandle]>, + /// Reads persisted inputs even when output uses a local write cache. + /// Input availability is determined by `inputs.is_empty()`, not this manager. + pub input_puffin_manager: SstPuffinManager, + /// Allows coarse input reuse only for the complete input scope without slicing. + /// The caller that owns the output plan validates this scope. + /// This is not a promise that input rows equal live output rows: deduplication + /// or deletion may remove rows. Handlers must record input-derived coverage + /// separately from actual output counts when reusing approximate summaries. + pub can_merge_inputs: bool, + /// Existing allocation available to optional index creators; `None` sets no limit here. + /// Single-region compaction and rebuild currently provide no shared total limit. + pub memory_limit: Option, + /// Exact row count of an unchanged SST during rebuild, otherwise unknown. + pub output_rows: Option, +} diff --git a/src/mito2/src/sst/index/column_test.rs b/src/mito2/src/sst/index/column_test.rs index db6e3f58dd4..a9d3ef9cf1b 100644 --- a/src/mito2/src/sst/index/column_test.rs +++ b/src/mito2/src/sst/index/column_test.rs @@ -188,6 +188,7 @@ async fn materialized_runs_match_legacy_index_bytes_and_bloom_lookups() { let mut options = IndexOptions::default(); options.inverted_index.segment_row_count = 3; let builder = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata: metadata.clone(), puffin_manager: puffin.clone(), diff --git a/src/mito2/src/sst/index/handler.rs b/src/mito2/src/sst/index/handler.rs index 31ff82d3bb5..dba75c143ce 100644 --- a/src/mito2/src/sst/index/handler.rs +++ b/src/mito2/src/sst/index/handler.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Planning and creation contracts for built-in indexes. +//! Planning and creation contracts for built-in and injected indexes. pub(crate) mod builtin; @@ -24,63 +24,128 @@ use common_error::ext::BoxedError; use datatypes::arrow::record_batch::RecordBatch; use index::target::IndexTarget; use store_api::metadata::RegionMetadataRef; +use store_api::storage::ColumnId; use crate::read::Batch; use crate::region::options::IndexOptions; -use crate::sst::file::RegionIndexId; +use crate::sst::file::{ColumnIndexMetadata, FileHandle, RegionIndexId}; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::kind::IndexKind; -use crate::sst::index::puffin_manager::SstPuffinWriter; +use crate::sst::index::puffin_manager::{SstPuffinManager, SstPuffinWriter}; use crate::sst::index::{CompletedEntry, IndexBuildType}; +/// A requested column format; typed implementation parameters stay in its plan. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PlanEntry { + /// Column to index. + pub column_id: ColumnId, + /// Implementation owning this entry. + pub kind: IndexKind, + /// Requested format version. The manifest records extension types, not versions. + /// A format upgrade requires a forced rebuild; the provider validates Puffin + /// properties and payload compatibility before reusing input state. + pub format_version: u32, +} + /// Schema and options used to prepare work without reading input files. -pub(crate) struct PlanContext<'a> { +pub struct PlanContext<'a> { /// Shared region schema. pub metadata: &'a RegionMetadataRef, /// Region index options. pub index_options: &'a IndexOptions, /// Lifecycle event requesting the build. pub build_type: IndexBuildType, + /// Existing column descriptors of the target SST during rebuild. + /// Empty for new outputs, including compaction; input descriptors are available + /// through [`BuildContext::inputs`] when constructing a creator. + pub existing_indexes: &'a [ColumnIndexMetadata], } /// Inputs for constructing all of one handler's indexes in a single output file. -pub(crate) struct BuildContext<'a> { +pub struct BuildContext<'a> { /// Shared output schema. pub metadata: &'a RegionMetadataRef, /// Physical output index identity. pub output: RegionIndexId, + /// Output manager, possibly backed by a local write cache. + pub puffin_manager: &'a SstPuffinManager, /// Existing intermediate-file facility shared by built-in creators. pub intermediate_manager: &'a IntermediateManager, + /// Pinned physical input files. + pub inputs: &'a [FileHandle], + /// Manager for reading persisted input indexes. With no inputs this falls back + /// to the output manager; determine input availability from `inputs.is_empty()`. + pub input_puffin_manager: &'a SstPuffinManager, + /// Permits coarse input-summary reuse only for a complete single output. + /// The caller that owns the output plan verifies this scope before construction. + pub can_merge_inputs: bool, + /// Retained-memory allowance for an injected creator, capped by the remaining + /// operation budget and its plan ceiling. `None` applies no limit at this layer. + /// Built-in creators use spill budgets. + pub memory_limit: Option, /// Actual output row-group size used to align inverted-index segments. pub row_group_size: Option, + /// Exact row count when rebuilding an existing SST; unknown for compaction outputs. + pub output_rows: Option, } -/// Prepares one built-in index implementation. -pub(crate) trait IndexKindHandler: Send + Sync + Debug { - /// Built-in implementation owned by this handler. - fn kind(&self) -> IndexKind; +/// Shared plugin installed by the embedding application. +pub type IndexKindProviderRef = Arc; - /// Prepares typed parameters. `None` means no work. +/// Supplies extension handlers while the engine retains the complete lifecycle. +pub trait IndexKindProvider: Send + Sync + Debug { + /// Related algorithms may share one handler to reuse row traversal. + fn handlers(&self) -> Vec>; +} + +/// Owns one or more related kinds that can share preparation and row traversal. +pub trait IndexKindHandler: Send + Sync + Debug { + /// Nonempty, unique kinds exclusively owned by this handler. A group must be + /// entirely built-in or entirely extension kinds. + fn kinds(&self) -> &[IndexKind]; + + /// Prepares typed parameters. `None` means no work; an empty entry list may + /// still request removal of obsolete descriptors during maintenance. fn plan(&self, context: &PlanContext<'_>) -> Result>, BoxedError>; } /// An immutable typed plan shared by output files of an operation. #[async_trait] -pub(crate) trait PlannedIndex: Send + Sync + Debug { - /// Built-in implementation owned by this plan. - fn kind(&self) -> IndexKind; +pub trait PlannedIndex: Send + Sync + Debug { + /// Nonempty, unique subset of the handler's kinds, including removal plans. + fn kinds(&self) -> &[IndexKind]; - /// Creates a per-file index creator before any output batch is consumed. + /// Requested column formats. Every entry kind must belong to this plan. + fn entries(&self) -> &[PlanEntry]; + + /// Prepares input reuse before any output batch is consumed. async fn build( &self, context: &BuildContext<'_>, ) -> Result>, BoxedError>; + + /// Whether building this plan can reuse persisted input summaries. + /// Row-only plans do not require checking input object existence. + fn needs_input_reuse(&self) -> bool { + false + } + + /// Optional hard ceiling; built-in creators use their existing spill budgets. + fn memory_limit(&self) -> Option { + None + } } /// Accumulates rows and writes blobs into the engine-owned Puffin lifecycle. /// Implementations release retained state on drop, including cancellation. #[async_trait] -pub(crate) trait IndexCreator: Send { +pub trait IndexCreator: Send { + /// Whether rebuilding this index requires scanning SST rows after preparation. + /// A creator reusing persisted state must account for output rows without updates. + fn requires_row_scan(&self) -> bool { + true + } + /// Borrows a batch and its reusable primary-key decode cache. async fn update(&mut self, batch: &mut Batch) -> Result<(), BoxedError>; diff --git a/src/mito2/src/sst/index/handler/builtin.rs b/src/mito2/src/sst/index/handler/builtin.rs index ec21eaf00f2..c814cb4b43b 100644 --- a/src/mito2/src/sst/index/handler/builtin.rs +++ b/src/mito2/src/sst/index/handler/builtin.rs @@ -27,7 +27,7 @@ use crate::config::{BloomFilterConfig, FulltextIndexConfig, InvertedIndexConfig} use crate::sst::index::handler::builtin::bloom::BloomPlan; use crate::sst::index::handler::builtin::fulltext::FulltextPlan; use crate::sst::index::handler::builtin::inverted::InvertedPlan; -use crate::sst::index::handler::{IndexKindHandler, PlanContext, PlannedIndex}; +use crate::sst::index::handler::{IndexKindHandler, PlanContext, PlanEntry, PlannedIndex}; use crate::sst::index::kind::IndexKind; /// Configuration for one of the existing file-level built-in creators. @@ -39,11 +39,11 @@ pub(crate) enum BuiltinIndexHandler { } impl IndexKindHandler for BuiltinIndexHandler { - fn kind(&self) -> IndexKind { + fn kinds(&self) -> &[IndexKind] { match self { - Self::Inverted(_) => IndexKind::Inverted, - Self::Fulltext(_) => IndexKind::Fulltext, - Self::Bloom(_) => IndexKind::Bloom, + Self::Inverted(_) => &[IndexKind::Inverted], + Self::Fulltext(_) => &[IndexKind::Fulltext], + Self::Bloom(_) => &[IndexKind::Bloom], } } @@ -57,6 +57,13 @@ impl IndexKindHandler for BuiltinIndexHandler { } } +fn target_column_id(target: &IndexTarget) -> store_api::storage::ColumnId { + match target { + IndexTarget::ColumnId(id) => *id, + IndexTarget::JsonPath(path) => path.column_id(), + } +} + fn column_target_id(target: &IndexTarget) -> Option { match target { IndexTarget::ColumnId(id) => Some(*id), @@ -64,6 +71,18 @@ fn column_target_id(target: &IndexTarget) -> Option) { + entries.sort_by_key(|entry| (entry.column_id, entry.format_version)); + entries.dedup_by(|a, b| { + if a.column_id == b.column_id { + b.format_version = b.format_version.max(a.format_version); + true + } else { + false + } + }); +} + #[cfg(test)] mod tests { use super::*; @@ -103,6 +122,7 @@ mod tests { metadata: &metadata, index_options: &options, build_type: IndexBuildType::Flush, + existing_indexes: &[], }; assert!(handler.plan(&context).is_err()); } diff --git a/src/mito2/src/sst/index/handler/builtin/bloom.rs b/src/mito2/src/sst/index/handler/builtin/bloom.rs index 1b011b79086..aff50955bb8 100644 --- a/src/mito2/src/sst/index/handler/builtin/bloom.rs +++ b/src/mito2/src/sst/index/handler/builtin/bloom.rs @@ -25,8 +25,10 @@ use index::target::IndexTarget; use crate::config::BloomFilterConfig; use crate::read::Batch; use crate::sst::index::bloom_filter::creator::BloomFilterIndexer; -use crate::sst::index::handler::builtin::column_target_id; -use crate::sst::index::handler::{BuildContext, IndexCreator, PlanContext, PlannedIndex}; +use crate::sst::index::handler::builtin::{column_target_id, normalize_entries, target_column_id}; +use crate::sst::index::handler::{ + BuildContext, IndexCreator, PlanContext, PlanEntry, PlannedIndex, +}; use crate::sst::index::kind::IndexKind; use crate::sst::index::puffin_manager::SstPuffinWriter; use crate::sst::index::{CompletedEntry, IndexBuildType}; @@ -35,6 +37,7 @@ use crate::sst::index::{CompletedEntry, IndexBuildType}; #[derive(Debug)] pub(crate) struct BloomPlan { config: BloomFilterConfig, + entries: Vec, targets: Vec<(IndexTarget, SkippingIndexOptions)>, } @@ -53,18 +56,26 @@ impl BloomPlan { return Ok(None); } let mut targets = Vec::new(); + let mut entries = Vec::new(); for column in &context.metadata.column_metadatas { for (target, options) in BloomFilterIndexer::targets(column, &context.build_type).map_err(BoxedError::new)? { + entries.push(PlanEntry { + column_id: target_column_id(&target), + kind: IndexKind::Bloom, + format_version: 0, + }); targets.push((target, options)); } } - if targets.is_empty() { + normalize_entries(&mut entries); + if entries.is_empty() { return Ok(None); } Ok(Some(Arc::new(Self { config: config.clone(), + entries, targets, }))) } @@ -72,8 +83,12 @@ impl BloomPlan { #[async_trait] impl PlannedIndex for BloomPlan { - fn kind(&self) -> IndexKind { - IndexKind::Bloom + fn kinds(&self) -> &[IndexKind] { + &[IndexKind::Bloom] + } + + fn entries(&self) -> &[PlanEntry] { + &self.entries } async fn build( diff --git a/src/mito2/src/sst/index/handler/builtin/fulltext.rs b/src/mito2/src/sst/index/handler/builtin/fulltext.rs index 160b7923f49..4a287ef4017 100644 --- a/src/mito2/src/sst/index/handler/builtin/fulltext.rs +++ b/src/mito2/src/sst/index/handler/builtin/fulltext.rs @@ -19,14 +19,17 @@ use std::sync::Arc; use async_trait::async_trait; use common_error::ext::BoxedError; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::schema::FulltextBackend; use index::target::IndexTarget; use crate::config::FulltextIndexConfig; use crate::read::Batch; use crate::sst::index::fulltext_index; use crate::sst::index::fulltext_index::creator::{FulltextIndexer, FulltextTarget}; -use crate::sst::index::handler::builtin::column_target_id; -use crate::sst::index::handler::{BuildContext, IndexCreator, PlanContext, PlannedIndex}; +use crate::sst::index::handler::builtin::{column_target_id, normalize_entries, target_column_id}; +use crate::sst::index::handler::{ + BuildContext, IndexCreator, PlanContext, PlanEntry, PlannedIndex, +}; use crate::sst::index::kind::IndexKind; use crate::sst::index::puffin_manager::SstPuffinWriter; use crate::sst::index::{CompletedEntry, IndexBuildType}; @@ -35,6 +38,7 @@ use crate::sst::index::{CompletedEntry, IndexBuildType}; #[derive(Debug)] pub(crate) struct FulltextPlan { config: FulltextIndexConfig, + entries: Vec, targets: Vec, } @@ -53,10 +57,20 @@ impl FulltextPlan { return Ok(None); } let mut targets = Vec::new(); + let mut entries = Vec::new(); for column in &context.metadata.column_metadatas { for (target, options) in FulltextIndexer::targets(column, &context.build_type).map_err(BoxedError::new)? { + entries.push(PlanEntry { + column_id: target_column_id(&target), + kind: IndexKind::Fulltext, + format_version: if options.backend == FulltextBackend::Bloom { + u32::from(fulltext_index::FULLTEXT_BLOOM_VERSION) + } else { + 0 + }, + }); targets.push(( column.column_id, column.column_schema.name.clone(), @@ -65,11 +79,13 @@ impl FulltextPlan { )); } } - if targets.is_empty() { + normalize_entries(&mut entries); + if entries.is_empty() { return Ok(None); } Ok(Some(Arc::new(Self { config: config.clone(), + entries, targets, }))) } @@ -77,8 +93,12 @@ impl FulltextPlan { #[async_trait] impl PlannedIndex for FulltextPlan { - fn kind(&self) -> IndexKind { - IndexKind::Fulltext + fn kinds(&self) -> &[IndexKind] { + &[IndexKind::Fulltext] + } + + fn entries(&self) -> &[PlanEntry] { + &self.entries } async fn build( diff --git a/src/mito2/src/sst/index/handler/builtin/inverted.rs b/src/mito2/src/sst/index/handler/builtin/inverted.rs index baf87f93258..353afd5cff9 100644 --- a/src/mito2/src/sst/index/handler/builtin/inverted.rs +++ b/src/mito2/src/sst/index/handler/builtin/inverted.rs @@ -27,8 +27,10 @@ use index::target::IndexTarget; use crate::config::InvertedIndexConfig; use crate::read::Batch; -use crate::sst::index::handler::builtin::column_target_id; -use crate::sst::index::handler::{BuildContext, IndexCreator, PlanContext, PlannedIndex}; +use crate::sst::index::handler::builtin::{column_target_id, normalize_entries, target_column_id}; +use crate::sst::index::handler::{ + BuildContext, IndexCreator, PlanContext, PlanEntry, PlannedIndex, +}; use crate::sst::index::inverted_index::creator::InvertedIndexer; use crate::sst::index::kind::IndexKind; use crate::sst::index::puffin_manager::SstPuffinWriter; @@ -38,6 +40,7 @@ use crate::sst::index::{CompletedEntry, IndexBuildType}; #[derive(Debug)] pub(crate) struct InvertedPlan { config: InvertedIndexConfig, + entries: Vec, targets: Vec, segment_row_count: NonZeroUsize, } @@ -78,11 +81,21 @@ impl InvertedPlan { ); return Ok(None); }; - if targets.is_empty() { + let mut entries: Vec<_> = targets + .iter() + .map(|target| PlanEntry { + column_id: target_column_id(target), + kind: IndexKind::Inverted, + format_version: 0, + }) + .collect(); + normalize_entries(&mut entries); + if entries.is_empty() { return Ok(None); } Ok(Some(Arc::new(Self { config: config.clone(), + entries, targets, segment_row_count, }))) @@ -91,8 +104,12 @@ impl InvertedPlan { #[async_trait] impl PlannedIndex for InvertedPlan { - fn kind(&self) -> IndexKind { - IndexKind::Inverted + fn kinds(&self) -> &[IndexKind] { + &[IndexKind::Inverted] + } + + fn entries(&self) -> &[PlanEntry] { + &self.entries } async fn build( @@ -250,6 +267,7 @@ mod tests { metadata: &metadata, index_options: &options, build_type: build_type.clone(), + existing_indexes: &[], }; let plan = InvertedPlan::plan(&config, &context).unwrap().unwrap(); let columns: Vec<_> = plan diff --git a/src/mito2/src/sst/index/indexer/abort.rs b/src/mito2/src/sst/index/indexer/abort.rs index e34b8ba4e7f..0bc659666b1 100644 --- a/src/mito2/src/sst/index/indexer/abort.rs +++ b/src/mito2/src/sst/index/indexer/abort.rs @@ -22,10 +22,10 @@ impl Indexer { pub(crate) async fn abort_creator(&mut self, position: usize) { let mut active = self.creators.remove(position); if let Err(err) = active.creator.abort().await { - if cfg!(any(test, feature = "test")) { + if active.is_builtin() && cfg!(any(test, feature = "test")) { panic!("Failed to abort index: {err:?}"); } - warn!(err; "Failed to abort index, region_id: {}, file_id: {}, kind: {:?}", self.region_id, self.file_id, active.kind); + warn!(err; "Failed to abort index, region_id: {}, file_id: {}, kinds: {:?}", self.region_id, self.file_id, active.kinds); } } diff --git a/src/mito2/src/sst/index/indexer/finish.rs b/src/mito2/src/sst/index/indexer/finish.rs index 366bde9a196..30c3064c9fe 100644 --- a/src/mito2/src/sst/index/indexer/finish.rs +++ b/src/mito2/src/sst/index/indexer/finish.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + +use common_error::ext::BoxedError; use common_telemetry::warn; use puffin::puffin_manager::{PuffinManager, PuffinWriter}; @@ -37,6 +40,7 @@ impl Indexer { Ok(entries) => { if let Err(err) = entries.iter().try_for_each(|entry| entry.validate()) { warn!(err; "Invalid index completion, region_id: {}, file_id: {}", self.region_id, self.file_id); + self.failure.get_or_insert_with(|| Arc::new(err)); drop(active); self.do_abort().await; return IndexOutput::default(); @@ -44,10 +48,11 @@ impl Indexer { output.entries.extend(entries); } Err(err) => { - if cfg!(any(test, feature = "test")) { + if active.is_builtin() && cfg!(any(test, feature = "test")) { panic!("Failed to finish index: {err:?}"); } - warn!(err; "Failed to finish index, region_id: {}, file_id: {}, kind: {:?}", self.region_id, self.file_id, active.kind); + warn!(err; "Failed to finish index, region_id: {}, file_id: {}, kinds: {:?}", self.region_id, self.file_id, active.kinds); + self.record_failure(err); drop(active); self.do_abort().await; return IndexOutput::default(); @@ -89,6 +94,7 @@ impl Indexer { ); } + self.record_failure(BoxedError::new(err)); None } @@ -105,6 +111,7 @@ impl Indexer { self.region_id, self.file_id, ); + self.record_failure(BoxedError::new(err)); 0 } diff --git a/src/mito2/src/sst/index/indexer/update.rs b/src/mito2/src/sst/index/indexer/update.rs index 54cdf254586..98fcbd3949b 100644 --- a/src/mito2/src/sst/index/indexer/update.rs +++ b/src/mito2/src/sst/index/indexer/update.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use common_telemetry::warn; use datatypes::arrow::record_batch::RecordBatch; @@ -28,20 +30,30 @@ impl Indexer { panic!("Failed to decode primary key for indexes: {err:?}"); } warn!(err; "Failed to decode primary key for indexes, region_id: {}, file_id: {}", self.region_id, self.file_id); + self.failure.get_or_insert_with(|| Arc::new(err)); self.do_abort().await; return; } let mut position = 0; while position < self.creators.len() { + if !self.creators[position].creator.requires_row_scan() { + position += 1; + continue; + } match self.creators[position].creator.update(batch).await { Ok(()) => position += 1, Err(err) => { - if cfg!(any(test, feature = "test")) { + let builtin = self.creators[position].is_builtin(); + if builtin && cfg!(any(test, feature = "test")) { panic!("Failed to update index: {err:?}"); } - warn!(err; "Failed to update index, region_id: {}, file_id: {}, kind: {:?}", self.region_id, self.file_id, self.creators[position].kind); - self.do_abort().await; - return; + warn!(err; "Failed to update index, region_id: {}, file_id: {}, kinds: {:?}", self.region_id, self.file_id, self.creators[position].kinds); + self.record_failure(err); + if builtin { + self.do_abort().await; + return; + } + self.abort_creator(position).await; } } } @@ -53,15 +65,24 @@ impl Indexer { } let mut position = 0; while position < self.creators.len() { + if !self.creators[position].creator.requires_row_scan() { + position += 1; + continue; + } match self.creators[position].creator.update_flat(batch).await { Ok(()) => position += 1, Err(err) => { - if cfg!(any(test, feature = "test")) { + let builtin = self.creators[position].is_builtin(); + if builtin && cfg!(any(test, feature = "test")) { panic!("Failed to update flat index: {err:?}"); } - warn!(err; "Failed to update flat index, region_id: {}, file_id: {}, kind: {:?}", self.region_id, self.file_id, self.creators[position].kind); - self.do_abort().await; - return; + warn!(err; "Failed to update flat index, region_id: {}, file_id: {}, kinds: {:?}", self.region_id, self.file_id, self.creators[position].kinds); + self.record_failure(err); + if builtin { + self.do_abort().await; + return; + } + self.abort_creator(position).await; } } } diff --git a/src/mito2/src/sst/index/kind.rs b/src/mito2/src/sst/index/kind.rs index e55c837bc55..a997bd7ed8f 100644 --- a/src/mito2/src/sst/index/kind.rs +++ b/src/mito2/src/sst/index/kind.rs @@ -16,30 +16,78 @@ use crate::sst::file::IndexType; -/// Index implementations supplied by the storage engine. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub(crate) enum IndexKind { +/// An index implementation known to the storage engine or supplied by a plugin. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum IndexKind { + /// Inverted index. Inverted, + /// Full-text index. Fulltext, + /// Bloom filter index. Bloom, + /// An index whose implementation is supplied by the embedding application. + Extension(ExtensionIndexKind), +} + +/// Extension algorithm identities; implementations remain outside mito2. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum ExtensionIndexKind { + /// HyperLogLog distinct-value estimate. + Hll, + /// KLL rank estimate. + Kll, + /// Misra–Gries frequency estimate. + Mg, +} + +impl ExtensionIndexKind { + /// Known extension algorithms checked by operation-level provider validation. + pub const ALL: [Self; 3] = [Self::Hll, Self::Kll, Self::Mg]; + + /// Column metadata key that requires this algorithm's provider. + pub const fn required_metadata_key(self) -> &'static str { + match self { + Self::Hll => datatypes::schema::HLL_KEY, + Self::Kll => datatypes::schema::KLL_KEY, + Self::Mg => datatypes::schema::MG_KEY, + } + } + + /// Stable label used by index diagnostics and metrics. + pub const fn as_str(self) -> &'static str { + match self { + Self::Hll => "hll_index", + Self::Kll => "kll_index", + Self::Mg => "mg_index", + } + } } impl IndexKind { /// Stable label used by the existing index metrics. - pub(crate) const fn as_str(self) -> &'static str { + pub const fn as_str(self) -> &'static str { match self { Self::Inverted => "inverted_index", Self::Fulltext => "fulltext_index", Self::Bloom => "bloom_filter_index", + Self::Extension(kind) => kind.as_str(), } } /// Returns the identity recorded in file and column index metadata. - pub(crate) fn index_type(self) -> IndexType { + pub fn index_type(self) -> IndexType { match self { Self::Inverted => IndexType::InvertedIndex, Self::Fulltext => IndexType::FulltextIndex, Self::Bloom => IndexType::BloomFilterIndex, + Self::Extension(ExtensionIndexKind::Hll) => IndexType::HllIndex, + Self::Extension(ExtensionIndexKind::Kll) => IndexType::KllIndex, + Self::Extension(ExtensionIndexKind::Mg) => IndexType::MgIndex, } } + + /// Whether this kind is implemented by the storage engine. + pub fn is_builtin(self) -> bool { + !matches!(self, Self::Extension(_)) + } } diff --git a/src/mito2/src/sst/index/registry.rs b/src/mito2/src/sst/index/registry.rs index 7f7c43f3671..7b9da2e51dd 100644 --- a/src/mito2/src/sst/index/registry.rs +++ b/src/mito2/src/sst/index/registry.rs @@ -12,29 +12,40 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Ordered planning of built-in index implementations. +//! Ordered registration and planning of index implementations. +use std::collections::HashSet; use std::sync::Arc; use common_error::ext::BoxedError; use common_telemetry::warn; +use snafu::IntoError; +use store_api::metadata::RegionMetadataRef; +use store_api::storage::RegionId; use crate::config::{BloomFilterConfig, FulltextIndexConfig, InvertedIndexConfig}; +use crate::error::{Error, ExternalSnafu, InvalidConfigSnafu}; +use crate::metrics::{INDEX_CREATE_DISCARDED_TOTAL, INDEX_DISCARD_ADMISSION_REJECTED}; +use crate::sst::file::{FileMeta, IndexType}; +use crate::sst::index::IndexBuildType; use crate::sst::index::handler::builtin::BuiltinIndexHandler; -use crate::sst::index::handler::{IndexKindHandler, PlanContext, PlannedIndex}; +use crate::sst::index::handler::{ + IndexKindHandler, IndexKindProviderRef, PlanContext, PlanEntry, PlannedIndex, +}; +use crate::sst::index::kind::{ExtensionIndexKind, IndexKind}; -/// Shared registry installed by the engine. -pub(crate) type IndexRegistryRef = Arc; +/// Shared registry installed by the engine and its embedding application. +pub type IndexRegistryRef = Arc; -/// Built-in handlers in index creation order. +/// One handler per owner; related algorithms may share one creator and row traversal. #[derive(Debug, Clone)] -pub(crate) struct IndexRegistry { +pub struct IndexRegistry { handlers: Vec>, } impl IndexRegistry { /// Registers the existing built-in creators in their established order. - pub(crate) fn new( + pub fn new( inverted: InvertedIndexConfig, fulltext: FulltextIndexConfig, bloom: BloomFilterConfig, @@ -48,46 +59,452 @@ impl IndexRegistry { } } - /// Prepares each handler once, retaining the existing test panic and - /// production log-and-skip policy for construction failures. - pub(crate) fn plan(&self, context: &PlanContext<'_>) -> IndexBuildPlan { + /// Adds the embedding application's handlers after the built-in configuration. + pub fn with_provider( + mut self, + provider: Option, + ) -> Result { + if let Some(provider) = provider { + for handler in provider.handlers() { + self.register(handler)?; + } + } + Ok(self) + } + + /// Installs a handler, rejecting empty or overlapping kind ownership. + pub fn register(&mut self, handler: Arc) -> Result<(), BoxedError> { + self.validate_registration(handler.kinds())?; + self.handlers.push(handler); + self.handlers + .sort_by_key(|handler| handler.kinds().iter().min().copied()); + Ok(()) + } + + /// Rejects required extension indexes whose implementation is unavailable. + /// Checks current schema options and both file-level and column-level metadata. + pub fn validate_required_kinds<'a>( + &self, + metadata: &RegionMetadataRef, + files: impl IntoIterator, + ) -> Result<(), BoxedError> { + self.validate_schema_kinds(metadata)?; + for file in files { + for kind in file.available_indexes.iter().chain( + file.indexes + .iter() + .flat_map(|column| &column.created_indexes), + ) { + self.validate_required_kind(*kind)?; + } + } + Ok(()) + } + + /// Prepares each handler once. Construction failures retain the existing + /// built-in test panic and production skip policy. Other failures are retained + /// for the caller to check before publishing rebuilt indexes. Required-provider + /// validation belongs to the operation entry point, before calling this method. + pub fn plan(&self, context: &PlanContext<'_>) -> IndexBuildPlan { let mut planned = Vec::new(); + let mut failure = None; for handler in &self.handlers { - match handler.plan(context) { + if context.build_type == IndexBuildType::Flush + && !handler.kinds().iter().all(|kind| kind.is_builtin()) + { + continue; + } + let result = handler.plan(context).and_then(|plan| { + if let Some(plan) = &plan { + Self::validate_plan(handler.as_ref(), plan.as_ref())?; + } + Ok(plan) + }); + match result { Ok(Some(plan)) => planned.push(plan), Ok(None) => {} Err(err) => { - Self::record_planning_error(handler.as_ref(), context, err); + Self::record_planning_error(handler.as_ref(), context, err, &mut failure); } } } - IndexBuildPlan { planned } + IndexBuildPlan { planned, failure } + } + + fn validate_registration(&self, kinds: &[IndexKind]) -> Result<(), BoxedError> { + if kinds.is_empty() + || kinds + .iter() + .enumerate() + .any(|(i, kind)| kinds[..i].contains(kind)) + { + return Err(BoxedError::new( + InvalidConfigSnafu { + reason: "index handler kinds must be nonempty and unique", + } + .build(), + )); + } + if kinds.iter().any(|kind| kind.is_builtin()) && kinds.iter().any(|kind| !kind.is_builtin()) + { + return Err(BoxedError::new( + InvalidConfigSnafu { + reason: "index handler cannot mix built-in and extension kinds", + } + .build(), + )); + } + if kinds.iter().any(|kind| self.has_handler(kind.index_type())) { + return Err(BoxedError::new( + InvalidConfigSnafu { + reason: "index handler kind is already registered", + } + .build(), + )); + } + Ok(()) + } + + fn has_handler(&self, kind: IndexType) -> bool { + self.handlers.iter().any(|handler| { + handler + .kinds() + .iter() + .any(|owned| owned.index_type() == kind) + }) + } + + fn validate_plan( + handler: &dyn IndexKindHandler, + plan: &dyn PlannedIndex, + ) -> Result<(), BoxedError> { + let kinds = plan.kinds(); + if kinds.is_empty() + || kinds + .iter() + .enumerate() + .any(|(i, kind)| kinds[..i].contains(kind) || !handler.kinds().contains(kind)) + || plan + .entries() + .iter() + .any(|entry| !kinds.contains(&entry.kind)) + { + return Err(BoxedError::new( + InvalidConfigSnafu { + reason: "index plan kinds must be nonempty, unique, owned by the handler, and include every entry kind", + } + .build(), + )); + } + Ok(()) } fn record_planning_error( handler: &dyn IndexKindHandler, context: &PlanContext<'_>, err: BoxedError, + failure: &mut Option>, ) { - if cfg!(any(test, feature = "test")) { + if handler.kinds().iter().all(|kind| kind.is_builtin()) && cfg!(any(test, feature = "test")) + { panic!( "Failed to prepare built-in index, region_id: {}, err: {err}", context.metadata.region_id ); } - warn!(err; "Failed to prepare index, region_id: {}, kind: {:?}", context.metadata.region_id, handler.kind()); + warn!(err; "Failed to prepare index, region_id: {}, kinds: {:?}", context.metadata.region_id, handler.kinds()); + failure.get_or_insert_with(|| { + Arc::new( + ExternalSnafu { + context: "index planning".to_string(), + } + .into_error(err), + ) + }); + } + + fn validate_schema_kinds(&self, metadata: &RegionMetadataRef) -> Result<(), BoxedError> { + for column in &metadata.column_metadatas { + let schema = &column.column_schema; + for kind in ExtensionIndexKind::ALL { + if schema.metadata().contains_key(kind.required_metadata_key()) { + self.validate_required_kind(IndexKind::Extension(kind).index_type())?; + } + } + } + Ok(()) + } + + fn validate_required_kind(&self, kind: IndexType) -> Result<(), BoxedError> { + if ExtensionIndexKind::ALL + .iter() + .any(|extension| IndexKind::Extension(*extension).index_type() == kind) + && !self.has_handler(kind) + { + return Err(BoxedError::new( + InvalidConfigSnafu { + reason: format!("required index {kind:?} has no registered handler"), + } + .build(), + )); + } + Ok(()) } } /// Typed plans retained across output files of one operation. #[derive(Debug, Clone, Default)] -pub(crate) struct IndexBuildPlan { +pub struct IndexBuildPlan { planned: Vec>, + failure: Option>, } impl IndexBuildPlan { - /// Plans in index creation order. - pub(crate) fn planned(&self) -> &[Arc] { + /// Reports preparation failure separately from a successful empty plan. + pub fn check_success(&self) -> Result<(), Arc> { + self.failure.clone().map_or(Ok(()), Err) + } + + /// Plans in execution order, grouped by owner rather than by column. + pub fn planned(&self) -> &[Arc] { &self.planned } + + /// All requested column formats. + pub fn entries(&self) -> impl Iterator { + self.planned.iter().flat_map(|plan| plan.entries()) + } + + /// Whether any planned implementation can reuse persisted input indexes. + pub fn needs_input_reuse(&self) -> bool { + self.planned.iter().any(|plan| plan.needs_input_reuse()) + } + + /// Omits optional extension work when the caller cannot admit its memory budget. + /// Records each requested kind once per target plan, rather than once per column. + pub fn omit_extensions_for_budget(&mut self, region_id: RegionId) { + let omitted: HashSet<_> = self + .entries() + .filter(|entry| !entry.kind.is_builtin()) + .map(|entry| entry.kind) + .collect(); + for kind in omitted { + warn!( + "Omitting index before construction, region_id: {}, kind: {}, reason: admission_rejected", + region_id, + kind.as_str() + ); + INDEX_CREATE_DISCARDED_TOTAL + .with_label_values(&[kind.as_str(), INDEX_DISCARD_ADMISSION_REJECTED]) + .inc(); + } + self.planned + .retain(|plan| plan.kinds().iter().all(|kind| kind.is_builtin())); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sst::index::kind::ExtensionIndexKind; + + #[derive(Debug)] + struct TestHandler(Vec); + + impl IndexKindHandler for TestHandler { + fn kinds(&self) -> &[IndexKind] { + &self.0 + } + fn plan(&self, _: &PlanContext<'_>) -> Result>, BoxedError> { + Ok(None) + } + } + + #[derive(Debug)] + struct TestPlan { + kinds: Vec, + entries: Vec, + } + + #[async_trait::async_trait] + impl PlannedIndex for TestPlan { + fn kinds(&self) -> &[IndexKind] { + &self.kinds + } + fn entries(&self) -> &[PlanEntry] { + &self.entries + } + async fn build( + &self, + _: &crate::sst::index::handler::BuildContext<'_>, + ) -> Result>, BoxedError> { + unreachable!("admission must not construct creators") + } + } + + #[test] + fn test_admission_omission_records_each_requested_kind_once() { + let builtin = IndexKind::Bloom; + let enabled = IndexKind::Extension(ExtensionIndexKind::Mg); + let disabled = IndexKind::Extension(ExtensionIndexKind::Kll); + let entry = |kind, column_id| PlanEntry { + kind, + column_id, + format_version: 1, + }; + let mut plan = IndexBuildPlan { + planned: vec![ + Arc::new(TestPlan { + kinds: vec![builtin], + entries: vec![entry(builtin, 0)], + }), + Arc::new(TestPlan { + kinds: vec![enabled, disabled], + entries: vec![entry(enabled, 0), entry(enabled, 1)], + }), + ], + failure: None, + }; + let counter = INDEX_CREATE_DISCARDED_TOTAL + .with_label_values(&[enabled.as_str(), "admission_rejected"]); + let disabled_counter = INDEX_CREATE_DISCARDED_TOTAL + .with_label_values(&[disabled.as_str(), "admission_rejected"]); + let before = counter.get(); + let disabled_before = disabled_counter.get(); + plan.omit_extensions_for_budget(RegionId::new(1, 1)); + assert_eq!( + plan.entries().cloned().collect::>(), + vec![entry(builtin, 0)] + ); + assert_eq!(counter.get(), before + 1); + assert_eq!(disabled_counter.get(), disabled_before); + assert!(plan.check_success().is_ok()); + plan.omit_extensions_for_budget(RegionId::new(1, 1)); + assert_eq!(counter.get(), before + 1); + } + + #[test] + fn test_registration_preserves_grouped_owner_and_rejects_overlap() { + let mut registry = + IndexRegistry::new(Default::default(), Default::default(), Default::default()); + let kinds = vec![ + IndexKind::Extension(ExtensionIndexKind::Hll), + IndexKind::Extension(ExtensionIndexKind::Kll), + IndexKind::Extension(ExtensionIndexKind::Mg), + ]; + registry + .register(Arc::new(TestHandler(kinds.clone()))) + .unwrap(); + assert_eq!(registry.handlers.len(), 4); + assert_eq!(registry.handlers[3].kinds(), kinds); + for invalid in [vec![], vec![kinds[0]], vec![IndexKind::Fulltext]] { + assert!(registry.register(Arc::new(TestHandler(invalid))).is_err()); + assert_eq!(registry.handlers.len(), 4); + } + } + + #[test] + fn test_registration_rejects_duplicates_within_one_owner() { + let mut registry = + IndexRegistry::new(Default::default(), Default::default(), Default::default()); + let kind = IndexKind::Extension(ExtensionIndexKind::Hll); + assert!( + registry + .register(Arc::new(TestHandler(vec![kind, kind]))) + .is_err() + ); + assert_eq!(registry.handlers.len(), 3); + } + + #[derive(Debug)] + struct PlannedHandler { + owned: Vec, + plan: Arc, + } + + impl IndexKindHandler for PlannedHandler { + fn kinds(&self) -> &[IndexKind] { + &self.owned + } + + fn plan(&self, _: &PlanContext<'_>) -> Result>, BoxedError> { + Ok(Some(self.plan.clone())) + } + } + + #[test] + fn test_plans_validate_ownership_without_removing_valid_plans() { + let hll = IndexKind::Extension(ExtensionIndexKind::Hll); + let kll = IndexKind::Extension(ExtensionIndexKind::Kll); + let metadata = Arc::new(crate::test_util::sst_util::sst_region_metadata()); + let options = Default::default(); + let context = PlanContext { + metadata: &metadata, + index_options: &options, + build_type: crate::sst::index::IndexBuildType::Compact, + existing_indexes: &[], + }; + for (kinds, entry_kind, invalid) in [ + (vec![], hll, true), + (vec![hll, hll], hll, true), + (vec![kll], kll, true), + (vec![hll], kll, true), + (vec![hll], hll, false), + ] { + let mut registry = + IndexRegistry::new(Default::default(), Default::default(), Default::default()); + registry + .register(Arc::new(PlannedHandler { + owned: vec![hll, IndexKind::Extension(ExtensionIndexKind::Mg)], + plan: Arc::new(TestPlan { + kinds, + entries: vec![PlanEntry { + column_id: 0, + kind: entry_kind, + format_version: 1, + }], + }), + })) + .unwrap(); + registry + .register(Arc::new(PlannedHandler { + owned: vec![kll], + plan: Arc::new(TestPlan { + kinds: vec![kll], + entries: vec![PlanEntry { + column_id: 0, + kind: kll, + format_version: 1, + }], + }), + })) + .unwrap(); + let plan = registry.plan(&context); + assert_eq!(plan.check_success().is_err(), invalid); + assert_eq!( + plan.entries() + .filter(|entry| !entry.kind.is_builtin()) + .map(|entry| entry.kind) + .collect::>(), + if invalid { vec![kll] } else { vec![hll, kll] } + ); + assert_eq!(registry.handlers.len(), 5); + } + } + + #[test] + fn test_registration_rejects_mixed_builtin_and_extension_owner() { + let mut registry = IndexRegistry { + handlers: Vec::new(), + }; + assert!( + registry + .register(Arc::new(TestHandler(vec![ + IndexKind::Inverted, + IndexKind::Extension(ExtensionIndexKind::Hll), + ]))) + .is_err() + ); + assert!(registry.handlers.is_empty()); + } } diff --git a/src/mito2/src/sst/index/sparse_test.rs b/src/mito2/src/sst/index/sparse_test.rs index e12d77f7903..1c2ba260ba9 100644 --- a/src/mito2/src/sst/index/sparse_test.rs +++ b/src/mito2/src/sst/index/sparse_test.rs @@ -130,6 +130,7 @@ async fn sparse_and_materialized_tags_produce_identical_indexes() { let mut options = IndexOptions::default(); options.inverted_index.segment_row_count = 3; let builder = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata: metadata.clone(), puffin_manager: puffin.clone(), diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index b8a8395088b..911e02cbf76 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -1560,6 +1560,7 @@ mod tests { let intermediate_manager = env.get_intermediate_manager(); let indexer_builder = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata: metadata.clone(), puffin_manager, @@ -2106,6 +2107,7 @@ mod tests { let intermediate_manager = env.get_intermediate_manager(); IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata, puffin_manager, @@ -2227,6 +2229,7 @@ mod tests { let intermediate_manager = env.get_intermediate_manager(); let indexer_builder = IndexerBuilderImpl { + index_build_inputs: None, build_type: IndexBuildType::Flush, metadata: metadata.clone(), puffin_manager, diff --git a/src/mito2/src/worker/handle_rebuild_index.rs b/src/mito2/src/worker/handle_rebuild_index.rs index 917a58d0c68..c17008bab27 100644 --- a/src/mito2/src/worker/handle_rebuild_index.rs +++ b/src/mito2/src/worker/handle_rebuild_index.rs @@ -76,6 +76,7 @@ impl RegionWorkerLoop { let indexer_builder_ref = Arc::new(IndexerBuilderImpl { build_type: build_type.clone(), metadata: snapshot.metadata.clone(), + index_build_inputs: None, registry: Arc::new(crate::sst::index::registry::IndexRegistry::new( self.config.inverted_index.clone(), self.config.fulltext_index.clone(),