diff --git a/licenserc-enterprise.toml b/licenserc-enterprise.toml index c700e696bf1..939346f5cba 100644 --- a/licenserc-enterprise.toml +++ b/licenserc-enterprise.toml @@ -15,6 +15,7 @@ headerPath = "licenses/enterprise-header.txt" # with the "# enterprise" block in `licenserc.toml`'s `excludes` — both are # enforced by `make check-enterprise-license`. includes = [ + "src/datatypes/src/schema/sketch.rs", "src/common/meta/src/rpc/ddl/trigger.rs", "src/operator/src/expr_helper/trigger.rs", "src/sql/src/statements/alter/trigger.rs", diff --git a/licenserc.toml b/licenserc.toml index 8810015c752..e601ca6f7c1 100644 --- a/licenserc.toml +++ b/licenserc.toml @@ -31,6 +31,7 @@ excludes = [ # Keep this list in sync with that file's `includes`; enforced by # `make check-enterprise-license`. # enterprise:start + "src/datatypes/src/schema/sketch.rs", "src/common/meta/src/rpc/ddl/trigger.rs", "src/operator/src/expr_helper/trigger.rs", "src/sql/src/statements/alter/trigger.rs", diff --git a/src/datatypes/Cargo.toml b/src/datatypes/Cargo.toml index e836a6abb02..d684415400b 100644 --- a/src/datatypes/Cargo.toml +++ b/src/datatypes/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true [features] default = [] +enterprise = [] test = [] [lints] diff --git a/src/datatypes/src/error.rs b/src/datatypes/src/error.rs index 745268f6e35..b41e199e3a0 100644 --- a/src/datatypes/src/error.rs +++ b/src/datatypes/src/error.rs @@ -278,6 +278,14 @@ pub enum Error { location: Location, }, + #[cfg(feature = "enterprise")] + #[snafu(display("Invalid sketch option: {}", msg))] + InvalidSketchOption { + msg: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid skipping index option: {}", msg))] InvalidSkippingIndexOption { msg: String, @@ -370,6 +378,9 @@ impl ErrorExt for Error { | InvalidSkippingIndexOption { .. } | InvalidCompressionOption { .. } => StatusCode::InvalidArguments, + #[cfg(feature = "enterprise")] + InvalidSketchOption { .. } => StatusCode::InvalidArguments, + ValueExceedsPrecision { .. } | CastType { .. } | CastTimeType { .. } diff --git a/src/datatypes/src/schema.rs b/src/datatypes/src/schema.rs index 4652b0104d1..cae7b9aa2bc 100644 --- a/src/datatypes/src/schema.rs +++ b/src/datatypes/src/schema.rs @@ -15,6 +15,8 @@ mod column_schema; pub mod constraint; pub mod ext; +#[cfg(feature = "enterprise")] +mod sketch; use std::collections::HashMap; use std::sync::Arc; @@ -36,10 +38,12 @@ pub use crate::schema::column_schema::{ COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY, COMPRESSION_KEY, ColumnEncoding, ColumnExtType, ColumnSchema, CompressionOptions, CompressionType, FULLTEXT_KEY, - FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, - SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, TIME_INDEX_KEY, + FulltextAnalyzer, FulltextBackend, FulltextOptions, HLL_KEY, INVERTED_INDEX_KEY, KLL_KEY, + MG_KEY, Metadata, SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, TIME_INDEX_KEY, }; pub use crate::schema::constraint::ColumnDefaultConstraint; +#[cfg(feature = "enterprise")] +pub use crate::schema::sketch::{HllOptions, KllOptions, MgOptions}; /// Key used to store version number of the schema in metadata. pub const VERSION_KEY: &str = "greptime:version"; diff --git a/src/datatypes/src/schema/column_schema.rs b/src/datatypes/src/schema/column_schema.rs index 9dac9258116..d43f95b4c44 100644 --- a/src/datatypes/src/schema/column_schema.rs +++ b/src/datatypes/src/schema/column_schema.rs @@ -32,6 +32,8 @@ use crate::error::{ use crate::extension::json::Json2ExtensionType; use crate::schema::TYPE_KEY; use crate::schema::constraint::ColumnDefaultConstraint; +#[cfg(feature = "enterprise")] +use crate::schema::sketch::{HllOptions, KllOptions, MgOptions}; use crate::value::Value; use crate::vectors::json::builder::JsonVectorBuilder; use crate::vectors::{MutableVector, VectorRef}; @@ -51,6 +53,12 @@ pub const INVERTED_INDEX_KEY: &str = "greptime:inverted_index"; pub const SKIPPING_INDEX_KEY: &str = "greptime:skipping_index"; /// Key used to store SST compression options in arrow field's metadata. pub const COMPRESSION_KEY: &str = "greptime:compression"; +/// Key used to store HLL options in arrow field's metadata. +pub const HLL_KEY: &str = "greptime:hll"; +/// Key used to store KLL options in arrow field's metadata. +pub const KLL_KEY: &str = "greptime:kll"; +/// Key used to store MG options in arrow field's metadata. +pub const MG_KEY: &str = "greptime:mg"; /// Keys used in fulltext options pub const COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE: &str = "enable"; @@ -436,6 +444,99 @@ impl ColumnSchema { Ok(()) } + /// Retrieves independently enabled HLL options, validating stored parameters. + #[cfg(feature = "enterprise")] + pub fn hll_options(&self) -> Result> { + self.metadata + .get(HLL_KEY) + .map(|json| serde_json::from_str(json).context(error::DeserializeSnafu { json })) + .transpose() + } + + /// Enables HLL with validated options. + #[cfg(feature = "enterprise")] + pub fn with_hll_options(mut self, options: HllOptions) -> Result { + self.set_hll_options(&options)?; + Ok(self) + } + + /// Sets HLL options without modifying other algorithm metadata. + #[cfg(feature = "enterprise")] + pub fn set_hll_options(&mut self, options: &HllOptions) -> Result<()> { + options.validate()?; + let json = serde_json::to_string(options).context(error::SerializeSnafu)?; + self.metadata.insert(HLL_KEY.to_string(), json); + Ok(()) + } + + /// Removes only HLL options. + #[cfg(feature = "enterprise")] + pub fn unset_hll_options(&mut self) { + self.metadata.remove(HLL_KEY); + } + + /// Retrieves independently enabled KLL options, validating stored parameters. + #[cfg(feature = "enterprise")] + pub fn kll_options(&self) -> Result> { + self.metadata + .get(KLL_KEY) + .map(|json| serde_json::from_str(json).context(error::DeserializeSnafu { json })) + .transpose() + } + + /// Enables KLL with validated options. + #[cfg(feature = "enterprise")] + pub fn with_kll_options(mut self, options: KllOptions) -> Result { + self.set_kll_options(&options)?; + Ok(self) + } + + /// Sets KLL options without modifying other algorithm metadata. + #[cfg(feature = "enterprise")] + pub fn set_kll_options(&mut self, options: &KllOptions) -> Result<()> { + options.validate()?; + let json = serde_json::to_string(options).context(error::SerializeSnafu)?; + self.metadata.insert(KLL_KEY.to_string(), json); + Ok(()) + } + + /// Removes only KLL options. + #[cfg(feature = "enterprise")] + pub fn unset_kll_options(&mut self) { + self.metadata.remove(KLL_KEY); + } + + /// Retrieves independently enabled MG options, validating stored parameters. + #[cfg(feature = "enterprise")] + pub fn mg_options(&self) -> Result> { + self.metadata + .get(MG_KEY) + .map(|json| serde_json::from_str(json).context(error::DeserializeSnafu { json })) + .transpose() + } + + /// Enables MG with validated options. + #[cfg(feature = "enterprise")] + pub fn with_mg_options(mut self, options: MgOptions) -> Result { + self.set_mg_options(&options)?; + Ok(self) + } + + /// Sets MG options without modifying other algorithm metadata. + #[cfg(feature = "enterprise")] + pub fn set_mg_options(&mut self, options: &MgOptions) -> Result<()> { + options.validate()?; + let json = serde_json::to_string(options).context(error::SerializeSnafu)?; + self.metadata.insert(MG_KEY.to_string(), json); + Ok(()) + } + + /// Removes only MG options. + #[cfg(feature = "enterprise")] + pub fn unset_mg_options(&mut self) { + self.metadata.remove(MG_KEY); + } + /// Retrieves the SST compression options for the column. pub fn compression_options(&self) -> Result> { match self.metadata.get(COMPRESSION_KEY) { @@ -1350,6 +1451,30 @@ mod tests { assert_eq!(column_schema, new_column_schema); } + #[test] + fn test_column_schema_preserves_sketch_metadata() { + let metadata = Metadata::from([ + ( + "greptime:hll".to_string(), + r#"{"precision":10}"#.to_string(), + ), + ("greptime:kll".to_string(), r#"{"capacity":64}"#.to_string()), + ( + "greptime:mg".to_string(), + r#"{"max-counters":64}"#.to_string(), + ), + ]); + let column = ColumnSchema::new("v", ConcreteDataType::int64_datatype(), true) + .with_metadata(metadata.clone()); + let field = Field::try_from(&column).unwrap(); + assert_eq!(&metadata, field.metadata()); + assert_eq!(column, ColumnSchema::try_from(&field).unwrap()); + assert_eq!( + column, + serde_json::from_str::(&serde_json::to_string(&column).unwrap()).unwrap() + ); + } + #[test] fn test_column_schema_with_duplicate_metadata() { let metadata = Metadata::from([(DEFAULT_CONSTRAINT_KEY.to_string(), "v1".to_string())]); diff --git a/src/datatypes/src/schema/sketch.rs b/src/datatypes/src/schema/sketch.rs new file mode 100644 index 00000000000..7017ad8b920 --- /dev/null +++ b/src/datatypes/src/schema/sketch.rs @@ -0,0 +1,316 @@ +// Copyright 2023-2026 GrepTime Inc. +// +// This file is part of the GreptimeDB Enterprise Edition and is licensed under +// the GreptimeDB Enterprise License. You may not use this file except in +// compliance with that license. A copy of the license is available at the root +// of this repository in the file LICENSE-ENTERPRISE. +// +// Unless required by applicable law or agreed to in writing, this software is +// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +// either express or implied. + +use serde::{Deserialize, Deserializer, Serialize}; +use snafu::ensure; +use sqlparser_derive::{Visit, VisitMut}; + +use crate::error::{InvalidSketchOptionSnafu, Result}; + +/// HLL options for a column. Metadata presence enables the algorithm. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Visit, VisitMut)] +#[serde(default, rename_all = "kebab-case")] +pub struct HllOptions { + /// Base-two logarithm of the register count, in 4..=18; defaults to 10. + #[serde(deserialize_with = "deserialize_hll_precision")] + pub precision: u8, +} + +impl HllOptions { + /// Validates the supported algorithm parameter range. + pub fn validate(&self) -> Result<()> { + ensure!( + (4..=18).contains(&self.precision), + InvalidSketchOptionSnafu { + msg: format!("HLL precision must be in 4..=18, got {}", self.precision) + } + ); + Ok(()) + } +} + +impl Default for HllOptions { + fn default() -> Self { + Self { precision: 10 } + } +} + +/// KLL options for a column. Metadata presence enables the algorithm. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Visit, VisitMut)] +#[serde(default, rename_all = "kebab-case")] +pub struct KllOptions { + /// Per-level capacity parameter, at least 8; defaults to 64. + #[serde(deserialize_with = "deserialize_kll_capacity")] + pub capacity: u16, +} + +impl KllOptions { + /// Validates the supported algorithm parameter range. + pub fn validate(&self) -> Result<()> { + ensure!( + self.capacity >= 8, + InvalidSketchOptionSnafu { + msg: format!("KLL capacity must be at least 8, got {}", self.capacity) + } + ); + Ok(()) + } +} + +impl Default for KllOptions { + fn default() -> Self { + Self { capacity: 64 } + } +} + +/// MG options for a column. Metadata presence enables the algorithm. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Visit, VisitMut)] +#[serde(default, rename_all = "kebab-case")] +pub struct MgOptions { + /// Maximum number of counters, greater than zero; defaults to 64. + #[serde(deserialize_with = "deserialize_mg_max_counters")] + pub max_counters: u32, +} + +impl MgOptions { + /// Validates the supported algorithm parameter range. + pub fn validate(&self) -> Result<()> { + ensure!( + self.max_counters > 0, + InvalidSketchOptionSnafu { + msg: format!( + "MG max_counters must be positive, got {}", + self.max_counters + ) + } + ); + Ok(()) + } +} + +impl Default for MgOptions { + fn default() -> Self { + Self { max_counters: 64 } + } +} + +fn deserialize_hll_precision<'de, D: Deserializer<'de>>( + deserializer: D, +) -> std::result::Result { + let precision = u8::deserialize(deserializer)?; + HllOptions { precision } + .validate() + .map_err(serde::de::Error::custom)?; + Ok(precision) +} + +fn deserialize_kll_capacity<'de, D: Deserializer<'de>>( + deserializer: D, +) -> std::result::Result { + let capacity = u16::deserialize(deserializer)?; + KllOptions { capacity } + .validate() + .map_err(serde::de::Error::custom)?; + Ok(capacity) +} + +fn deserialize_mg_max_counters<'de, D: Deserializer<'de>>( + deserializer: D, +) -> std::result::Result { + let max_counters = u32::deserialize(deserializer)?; + MgOptions { max_counters } + .validate() + .map_err(serde::de::Error::custom)?; + Ok(max_counters) +} + +#[cfg(test)] +mod tests { + use arrow::datatypes::Field; + + use super::*; + use crate::data_type::ConcreteDataType; + use crate::schema::{ColumnSchema, HLL_KEY, KLL_KEY, MG_KEY}; + + #[test] + fn test_defaults_and_parameter_validation() { + assert_eq!( + HllOptions::default(), + serde_json::from_str::("{}").unwrap() + ); + assert_eq!( + KllOptions::default(), + serde_json::from_str::("{}").unwrap() + ); + assert_eq!( + MgOptions::default(), + serde_json::from_str::("{}").unwrap() + ); + for precision in [4, 10, 18] { + let options = HllOptions { precision }; + assert_eq!( + options, + serde_json::from_str::(&serde_json::to_string(&options).unwrap()) + .unwrap() + ); + } + for precision in [0, 3, 19, 255] { + assert!( + serde_json::from_str::(&format!(r#"{{"precision":{precision}}}"#)) + .is_err() + ); + } + for capacity in [0, 7] { + assert!( + serde_json::from_str::(&format!(r#"{{"capacity":{capacity}}}"#)) + .is_err() + ); + } + assert!(serde_json::from_str::(r#"{"max-counters":0}"#).is_err()); + assert!(serde_json::from_str::(r#"{"max-counters":4294967296}"#).is_err()); + for capacity in [8, u16::MAX] { + let options = KllOptions { capacity }; + assert_eq!( + options, + serde_json::from_str::(&serde_json::to_string(&options).unwrap()) + .unwrap() + ); + } + for capacity in [1, u32::MAX] { + let options = MgOptions { + max_counters: capacity, + }; + assert_eq!( + options, + serde_json::from_str::(&serde_json::to_string(&options).unwrap()) + .unwrap() + ); + } + } + + #[test] + fn test_independent_metadata_and_roundtrip() { + let mut column = ColumnSchema::new("v", ConcreteDataType::int64_datatype(), true); + assert_eq!(None, column.hll_options().unwrap()); + assert_eq!(None, column.kll_options().unwrap()); + assert_eq!(None, column.mg_options().unwrap()); + column = column + .with_hll_options(HllOptions::default()) + .unwrap() + .with_kll_options(KllOptions::default()) + .unwrap() + .with_mg_options(MgOptions::default()) + .unwrap(); + let field = Field::try_from(&column).unwrap(); + assert_eq!(field.metadata()[HLL_KEY], r#"{"precision":10}"#); + assert_eq!(field.metadata()[KLL_KEY], r#"{"capacity":64}"#); + assert_eq!(field.metadata()[MG_KEY], r#"{"max-counters":64}"#); + assert_eq!(column, ColumnSchema::try_from(&field).unwrap()); + assert_eq!( + column, + serde_json::from_str::(&serde_json::to_string(&column).unwrap()).unwrap() + ); + let before = column.clone(); + assert!( + column + .set_hll_options(&HllOptions { precision: 3 }) + .is_err() + ); + assert!(column.set_kll_options(&KllOptions { capacity: 7 }).is_err()); + assert!( + column + .set_mg_options(&MgOptions { max_counters: 0 }) + .is_err() + ); + assert_eq!(before, column); + column + .set_hll_options(&HllOptions { precision: 12 }) + .unwrap(); + column + .set_kll_options(&KllOptions { capacity: 128 }) + .unwrap(); + column + .set_mg_options(&MgOptions { max_counters: 256 }) + .unwrap(); + assert_eq!( + Some(HllOptions { precision: 12 }), + column.hll_options().unwrap() + ); + assert_eq!( + Some(KllOptions { capacity: 128 }), + column.kll_options().unwrap() + ); + assert_eq!( + Some(MgOptions { max_counters: 256 }), + column.mg_options().unwrap() + ); + column.unset_hll_options(); + assert_eq!(None, column.hll_options().unwrap()); + assert_eq!( + Some(KllOptions { capacity: 128 }), + column.kll_options().unwrap() + ); + assert_eq!( + Some(MgOptions { max_counters: 256 }), + column.mg_options().unwrap() + ); + column.unset_kll_options(); + assert_eq!(None, column.kll_options().unwrap()); + assert_eq!( + Some(MgOptions { max_counters: 256 }), + column.mg_options().unwrap() + ); + column.unset_mg_options(); + assert_eq!(None, column.mg_options().unwrap()); + } + + #[test] + fn test_invalid_persisted_options_are_rejected_by_getters() { + let column = ColumnSchema::new("v", ConcreteDataType::int64_datatype(), true); + let field = Field::try_from(&column).unwrap().with_metadata( + [ + (HLL_KEY.to_string(), r#"{"precision":19}"#.to_string()), + (KLL_KEY.to_string(), r#"{"capacity":0}"#.to_string()), + (MG_KEY.to_string(), "malformed json".to_string()), + ] + .into(), + ); + let column = ColumnSchema::try_from(&field).unwrap(); + assert!(column.hll_options().is_err()); + assert!(column.kll_options().is_err()); + assert!(column.mg_options().is_err()); + } + + #[test] + fn test_mg_max_counters_serialization() { + let expected = MgOptions { max_counters: 128 }; + assert_eq!( + expected, + serde_json::from_str::(r#"{"max-counters":128}"#).unwrap() + ); + assert_eq!( + r#"{"max-counters":128}"#, + serde_json::to_string(&expected).unwrap() + ); + assert!(serde_json::from_str::(r#"{"max-counters":0}"#).is_err()); + } + + #[test] + fn test_validation_errors_include_values() { + for (error, value) in [ + (HllOptions { precision: 19 }.validate().unwrap_err(), "19"), + (KllOptions { capacity: 7 }.validate().unwrap_err(), "7"), + (MgOptions { max_counters: 0 }.validate().unwrap_err(), "0"), + ] { + assert!(error.to_string().contains(value), "{error}"); + } + } +} diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index f9a8cea42b1..5265554728d 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -9,7 +9,7 @@ default = [] test = ["common-test-util", "rstest", "rstest_reuse", "rskafka"] testing = ["test"] test-shared-fs-region-migration = [] -enterprise = ["tokio-util/rt"] +enterprise = ["datatypes/enterprise", "tokio-util/rt"] hdfs-object-store = ["object-store/hdfs-object-store"] [lints] diff --git a/src/mito2/src/engine/edit_region_test.rs b/src/mito2/src/engine/edit_region_test.rs index 4e44d9fdd8d..68db3def4cf 100644 --- a/src/mito2/src/engine/edit_region_test.rs +++ b/src/mito2/src/engine/edit_region_test.rs @@ -33,13 +33,15 @@ use store_api::region_request::{ use store_api::storage::{FileId, RegionId}; use tokio::sync::{Barrier, mpsc, oneshot}; +use crate::cache::file_cache::{FileType, IndexKey}; use crate::config::MitoConfig; use crate::engine::MitoEngine; use crate::engine::listener::EventListener; use crate::manifest::action::RegionEdit; use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState}; use crate::request::{RegionEditRequest, WorkerRequest}; -use crate::sst::file::FileMeta; +use crate::sst::file::{ColumnIndexMetadata, FileMeta, IndexType}; +use crate::sst::location; use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, rows_schema}; use crate::time_provider::mock::MockTimeProvider; @@ -171,13 +173,15 @@ async fn test_edit_region_schedule_compaction_with_format( } } +#[rstest::rstest] +#[case::without_index(None)] +#[case::known_index_size(Some(7))] +#[case::unknown_index_size(Some(0))] #[tokio::test] -async fn test_edit_region_fill_cache() { - test_edit_region_fill_cache_with_format(false).await; - test_edit_region_fill_cache_with_format(true).await; -} - -async fn test_edit_region_fill_cache_with_format(flat_format: bool) { +async fn test_edit_region_fill_cache( + #[values(false, true)] flat_format: bool, + #[case] index_file_size: Option, +) { let mut env = TestEnv::new().await; struct EditRegionListener { @@ -229,13 +233,33 @@ async fn test_edit_region_fill_cache_with_format(flat_format: bool) { .await .unwrap(); - let edit = RegionEdit { - files_to_add: vec![FileMeta { - region_id: region.region_id, - file_id, - level: 0, + let mut meta = FileMeta { + region_id, + file_id, + index_version: 3, + ..Default::default() + }; + if let Some(size) = index_file_size { + meta.index_file_size = size; + meta.indexes.push(ColumnIndexMetadata { + column_id: 0, + created_indexes: smallvec::smallvec![IndexType::HllIndex], ..Default::default() - }], + }); + } + let index_path = location::index_file_path( + region.table_dir(), + meta.index_id(), + region.access_layer.path_type(), + ); + region + .access_layer + .object_store() + .write(&index_path, b"payload".as_slice()) + .await + .unwrap(); + let edit = RegionEdit { + files_to_add: vec![meta], files_to_remove: vec![], timestamp_ms: None, compaction_time_window: None, @@ -251,6 +275,21 @@ async fn test_edit_region_fill_cache_with_format(flat_format: bool) { .unwrap() .unwrap(); assert_eq!(file_id, actual); + // The listener runs after both download attempts, including the skipped-index path. + let cache_manager = engine.cache_manager(); + let file_cache = cache_manager.write_cache().unwrap().file_cache(); + let index_key = IndexKey::new(region_id, file_id, FileType::Puffin(3)); + let expected_cached = index_file_size.is_some_and(|size| size > 0); + assert_eq!(file_cache.contains_key(&index_key), expected_cached); + let path = file_cache.cache_file_path(index_key); + if expected_cached { + assert_eq!( + file_cache.local_store().read(&path).await.unwrap().to_vec(), + b"payload" + ); + } else { + assert!(!file_cache.local_store().exists(&path).await.unwrap()); + } } #[tokio::test] diff --git a/src/mito2/src/gc/worker_test.rs b/src/mito2/src/gc/worker_test.rs index 0283de625b7..55a04dbad5e 100644 --- a/src/mito2/src/gc/worker_test.rs +++ b/src/mito2/src/gc/worker_test.rs @@ -39,7 +39,7 @@ use crate::gc::{GcConfig, LocalGcWorker, should_delete_file}; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList, RemovedFile}; use crate::metrics::GC_SKIPPED_UNPARSABLE_FILES; use crate::region::{MitoRegionRef, RegionLeaderState}; -use crate::sst::file::{FileHandle, FileMeta, IndexType, RegionIndexId}; +use crate::sst::file::{ColumnIndexMetadata, FileHandle, FileMeta, IndexType, RegionIndexId}; use crate::sst::location; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, rows_schema, @@ -280,8 +280,17 @@ async fn test_gc_reference_protects_parquet_and_only_its_index( async fn test_gc_reclaims_obsolete_index_after_readers_and_linger_release( #[values(false, true)] full_file_listing: bool, #[values(false, true)] linger_expired: bool, + #[values(false, true)] column_only: bool, ) { - let env = IndexGcTestEnv::new(full_file_listing).await; + let mut env = IndexGcTestEnv::new(full_file_listing).await; + if column_only { + env.meta.available_indexes.clear(); + env.meta.indexes = vec![ColumnIndexMetadata { + column_id: 0, + created_indexes: smallvec::smallvec![IndexType::HllIndex], + ..Default::default() + }]; + } let region_id = env.meta.region_id; env.engine .set_region_role(region_id, RegionRole::Leader) diff --git a/src/mito2/src/manifest/action.rs b/src/mito2/src/manifest/action.rs index 0fccd8382c8..5d5620bdab8 100644 --- a/src/mito2/src/manifest/action.rs +++ b/src/mito2/src/manifest/action.rs @@ -663,6 +663,51 @@ mod tests { // These tests are used to ensure backward compatibility of manifest files. // DO NOT modify the serialized string when they fail, check if your // modification to manifest-related structs is compatible with older manifests. + #[test] + fn test_sketch_index_manifest_fixture() { + use crate::sst::file::IndexType; + + // Frozen persisted names. Keep this fixture unchanged when extending the format. + let fixture = r#"{"actions":[{"Edit":{"files_to_add":[{ + "region_id":0,"file_id":"bc5896ec-e4d8-4017-a80d-f2de73188d55", + "time_range":[{"value":0,"unit":"Millisecond"},{"value":1,"unit":"Millisecond"}], + "level":0,"index_version":7,"index_file_size":128, + "available_indexes":["InvertedIndex","HllIndex","KllIndex","MgIndex"], + "indexes":[{"column_id":3,"created_indexes":["InvertedIndex","HllIndex","KllIndex","MgIndex"]}] + }],"files_to_remove":[],"compaction_time_window":null}}]}"#; + let actions = RegionMetaActionList::decode(fixture.as_bytes()).unwrap(); + let RegionMetaAction::Edit(edit) = &actions.actions[0] else { + panic!("expected a file edit"); + }; + assert_eq!(actions.actions.len(), 1); + assert_eq!(edit.files_to_add.len(), 1); + let file = &edit.files_to_add[0]; + let expected = [ + IndexType::InvertedIndex, + IndexType::HllIndex, + IndexType::KllIndex, + IndexType::MgIndex, + ]; + assert_eq!(file.available_indexes.as_slice(), expected); + assert_eq!(file.indexes.len(), 1); + assert_eq!(file.indexes[0].column_id, 3); + assert_eq!(file.indexes[0].created_indexes.as_slice(), expected); + assert_eq!(file.index_version(), Some(7)); + assert_eq!(file.index_file_size, 128); + let encoded = actions.encode().unwrap(); + assert_eq!(RegionMetaActionList::decode(&encoded).unwrap(), actions); + let encoded: serde_json::Value = serde_json::from_slice(&encoded).unwrap(); + let names = serde_json::json!(["InvertedIndex", "HllIndex", "KllIndex", "MgIndex"]); + assert_eq!( + encoded["actions"][0]["Edit"]["files_to_add"][0]["available_indexes"], + names + ); + assert_eq!( + encoded["actions"][0]["Edit"]["files_to_add"][0]["indexes"][0]["created_indexes"], + names + ); + } + #[test] fn test_region_action_compatibility() { let region_edit = r#"{ diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index ffefd1243fb..85881028ab9 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -1115,7 +1115,7 @@ impl RegionLoadCacheTask { for level in version.ssts.levels() { for file_handle in level.files.values() { let file_meta = file_handle.meta_ref(); - if file_meta.exists_index() { + if file_meta.has_downloadable_index() { let puffin_key = IndexKey::new( file_meta.region_id, file_meta.file_id, @@ -1747,6 +1747,95 @@ mod tests { parquet_bytes } + #[rstest::rstest] + #[case::known_size(7, true)] + #[case::unknown_size(0, false)] + #[tokio::test] + async fn test_preload_column_only_index( + #[case] index_file_size: u64, + #[case] expected_cached: bool, + ) { + use store_api::region_engine::RegionEngine; + use store_api::region_request::RegionRequest; + + use crate::manifest::action::RegionEdit; + use crate::sst::file::{ColumnIndexMetadata, IndexType}; + use crate::test_util::CreateRequestBuilder; + + let mut env = TestEnv::new().await; + let engine = env.create_engine(Default::default()).await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let write_cache = env + .create_write_cache( + ObjectStore::new(Memory::default()).unwrap(), + ReadableSize::mb(1), + ) + .await; + let meta = FileMeta { + region_id, + file_id: FileId::random(), + index_file_size, + index_version: 3, + indexes: vec![ColumnIndexMetadata { + column_id: 0, + created_indexes: smallvec::smallvec![IndexType::HllIndex], + ..Default::default() + }], + ..Default::default() + }; + assert!(meta.available_indexes.is_empty()); + assert!(meta.exists_index()); + assert_eq!(meta.index_version(), Some(3)); + let key = IndexKey::new(region_id, meta.file_id, FileType::Puffin(3)); + let path = crate::sst::location::index_file_path( + region.access_layer.table_dir(), + meta.index_id(), + region.access_layer.path_type(), + ); + region + .access_layer + .object_store() + .write(&path, b"payload".as_slice()) + .await + .unwrap(); + region.version_control.apply_edit( + Some(RegionEdit { + files_to_add: vec![meta], + ..Default::default() + }), + &[], + region.file_purger.clone(), + ); + + super::RegionLoadCacheTask::new(region) + .fill_cache(&write_cache.file_cache()) + .await; + let file_cache = write_cache.file_cache(); + assert_eq!(file_cache.contains_key(&key), expected_cached); + let cached_path = file_cache.cache_file_path(key); + if expected_cached { + assert_eq!( + file_cache + .local_store() + .read(&cached_path) + .await + .unwrap() + .to_vec(), + b"payload" + ); + } else { + assert!(!file_cache.local_store().exists(&cached_path).await.unwrap()); + } + } + #[tokio::test] async fn test_preload_parquet_meta_cache_uses_file_cache() { let env = TestEnv::new().await; diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index c5217b494bd..ba29db1736f 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -366,7 +366,7 @@ impl Debug for FileMeta { "max_row_group_uncompressed_size", &ReadableSize(self.max_row_group_uncompressed_size), ); - if !self.available_indexes.is_empty() { + if self.exists_index() { debug_struct .field("available_indexes", &self.available_indexes) .field("indexes", &self.indexes) @@ -414,6 +414,12 @@ pub enum IndexType { BloomFilterIndex, /// Retained for manifests written with the removed experimental vector index. VectorIndex, + /// HyperLogLog distinct-value estimate. + HllIndex, + /// KLL rank estimate. + KllIndex, + /// Misra–Gries frequency estimate. + MgIndex, } /// Metadata of indexes created for a specific column in an SST file. @@ -466,6 +472,16 @@ impl FileMeta { pub fn exists_index(&self) -> bool { !self.available_indexes.is_empty() + || self + .indexes + .iter() + .any(|index| !index.created_indexes.is_empty()) + } + + /// Returns whether the index has a known readable length for cache downloads. + /// A zero size does not remove the descriptor or its cleanup obligations. + pub fn has_downloadable_index(&self) -> bool { + self.exists_index() && self.index_file_size > 0 } pub fn index_version(&self) -> Option { @@ -1208,6 +1224,50 @@ mod tests { } } + #[test] + fn test_sketch_manifest_compatibility() { + let legacy = r#"{ + "region_id": 0, + "file_id": "bc5896ec-e4d8-4017-a80d-f2de73188d55", + "time_range": [ + {"value": 0, "unit": "Millisecond"}, + {"value": 0, "unit": "Millisecond"} + ], + "available_indexes": [], + "indexes": [], + "level": 0 + }"#; + let legacy: FileMeta = serde_json::from_str(legacy).unwrap(); + assert!(legacy.indexes.is_empty()); + assert!(!legacy.exists_index()); + let old_column = r#"{"column_id":3,"created_indexes":["InvertedIndex"]}"#; + let column: ColumnIndexMetadata = serde_json::from_str(old_column).unwrap(); + assert_eq!( + serde_json::to_value(&column).unwrap(), + serde_json::from_str::(old_column).unwrap() + ); + let mut metadata = legacy.clone(); + metadata.index_version = 7; + metadata.indexes = vec![ColumnIndexMetadata { + column_id: 3, + created_indexes: SmallVec::from_iter([ + IndexType::HllIndex, + IndexType::KllIndex, + IndexType::MgIndex, + ]), + ..Default::default() + }]; + let encoded = serde_json::to_string(&metadata).unwrap(); + let decoded: FileMeta = serde_json::from_str(&encoded).unwrap(); + assert_eq!(metadata, decoded); + assert!(decoded.exists_index()); + assert_eq!(decoded.index_version(), Some(7)); + assert_eq!( + serde_json::to_value(&decoded.indexes[0].created_indexes).unwrap(), + serde_json::json!(["HllIndex", "KllIndex", "MgIndex"]) + ); + } + #[test] fn test_deserialize_file_meta() { let file_meta = create_file_meta(FileId::random(), 0); diff --git a/src/mito2/src/sst/file_purger.rs b/src/mito2/src/sst/file_purger.rs index da404eb140e..2b9732bda40 100644 --- a/src/mito2/src/sst/file_purger.rs +++ b/src/mito2/src/sst/file_purger.rs @@ -548,7 +548,10 @@ mod tests { #[rstest::rstest] #[tokio::test] - async fn test_file_purge_with_index(#[values(0, LEVEL_SHARED)] level: u8) { + async fn test_file_purge_with_index( + #[values(0, LEVEL_SHARED)] level: u8, + #[values(false, true)] sketch_only: bool, + ) { let env = PurgeTestEnv::new().await; let object_store = env.layer.object_store(); let sst_file_id = env.file_id; @@ -575,10 +578,18 @@ mod tests { level, file_size: 4096, max_row_group_uncompressed_size: 4096, - available_indexes: SmallVec::from_iter([IndexType::InvertedIndex]), + available_indexes: if sketch_only { + SmallVec::new() + } else { + SmallVec::from_iter([IndexType::InvertedIndex]) + }, indexes: vec![ColumnIndexMetadata { column_id: 0, - created_indexes: SmallVec::from_iter([IndexType::InvertedIndex]), + created_indexes: SmallVec::from_iter([if sketch_only { + IndexType::HllIndex + } else { + IndexType::InvertedIndex + }]), ..Default::default() }], index_file_size: 4096, @@ -592,6 +603,14 @@ mod tests { }, file_purger, ); + assert!(handle.meta_ref().exists_index()); + assert_eq!(handle.meta_ref().available_indexes.is_empty(), sketch_only); + assert_eq!( + handle.meta_ref().indexes[0] + .created_indexes + .contains(&IndexType::HllIndex), + sketch_only + ); // Shared SSTs and their indexes require global GC even after local release. handle.mark_deleted(); } diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index b5cd6d300f8..4fb0f08e217 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -45,6 +45,8 @@ use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use puffin_manager::SstPuffinManager; use smallvec::SmallVec; use snafu::ResultExt; +#[cfg(feature = "enterprise")] +pub use statistics::{Statistics, TimerGuard}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::storage::{ColumnId, FileId, RegionId}; diff --git a/src/mito2/src/sst/index/statistics.rs b/src/mito2/src/sst/index/statistics.rs index 698ac29c94e..5d63a0e826b 100644 --- a/src/mito2/src/sst/index/statistics.rs +++ b/src/mito2/src/sst/index/statistics.rs @@ -16,22 +16,28 @@ use std::time::{Duration, Instant}; use crate::metrics::{INDEX_CREATE_BYTES_TOTAL, INDEX_CREATE_ELAPSED, INDEX_CREATE_ROWS_TOTAL}; -pub(crate) type ByteCount = u64; -pub(crate) type RowCount = usize; +pub type ByteCount = u64; +pub type RowCount = usize; /// Stage of the index creation process. enum Stage { Update, + #[cfg(feature = "enterprise")] + Merge, Finish, Cleanup, } /// Statistics for index creation. Flush metrics when dropped. -pub(crate) struct Statistics { +pub struct Statistics { /// Index type. index_type: &'static str, /// Accumulated elapsed time for the index update stage. update_elapsed: Duration, + /// Whether update and total elapsed time can be attributed to this index. + update_timing_available: bool, + /// Accumulated time spent merging existing index states. + merge_elapsed: Duration, /// Accumulated elapsed time for the index finish stage. finish_elapsed: Duration, /// Accumulated elapsed time for the cleanup stage. @@ -43,10 +49,13 @@ pub(crate) struct Statistics { } impl Statistics { + /// Creates index statistics whose accumulated metrics are published on drop. pub fn new(index_type: &'static str) -> Self { Self { index_type, update_elapsed: Duration::default(), + update_timing_available: true, + merge_elapsed: Duration::default(), finish_elapsed: Duration::default(), cleanup_eplased: Duration::default(), row_count: 0, @@ -54,12 +63,38 @@ impl Statistics { } } + /// Creates counters for interleaved work whose update time cannot be attributed + /// to this index. Update and total histograms remain absent until `record_update`. + #[cfg(feature = "enterprise")] + pub fn new_without_update_timing(index_type: &'static str) -> Self { + let mut stats = Self::new(index_type); + stats.update_timing_available = false; + stats + } + + /// Counts rows without taking a clock sample or claiming an update duration. + #[cfg(feature = "enterprise")] + pub fn inc_row_count(&mut self, rows: RowCount) { + self.row_count += rows; + } + /// Starts timing the update stage, returning a `TimerGuard` to automatically record duration. #[must_use] pub fn record_update(&mut self) -> TimerGuard<'_> { + self.update_timing_available = true; TimerGuard::new(self, Stage::Update) } + /// Starts timing merges of existing index states. + /// + /// Merged input rows are represented rows, not newly scanned output rows. Callers + /// count actual output rows separately through the update or finish guard. + #[must_use] + #[cfg(feature = "enterprise")] + pub fn record_merge(&mut self) -> TimerGuard<'_> { + TimerGuard::new(self, Stage::Merge) + } + /// Starts timing the finish stage, returning a `TimerGuard` to automatically record duration. #[must_use] pub fn record_finish(&mut self) -> TimerGuard<'_> { @@ -85,20 +120,33 @@ impl Statistics { impl Drop for Statistics { fn drop(&mut self) { - INDEX_CREATE_ELAPSED - .with_label_values(&["update", self.index_type]) - .observe(self.update_elapsed.as_secs_f64()); + if self.update_timing_available { + INDEX_CREATE_ELAPSED + .with_label_values(&["update", self.index_type]) + .observe(self.update_elapsed.as_secs_f64()); + } + if !self.merge_elapsed.is_zero() { + INDEX_CREATE_ELAPSED + .with_label_values(&["merge", self.index_type]) + .observe(self.merge_elapsed.as_secs_f64()); + } INDEX_CREATE_ELAPSED .with_label_values(&["finish", self.index_type]) .observe(self.finish_elapsed.as_secs_f64()); INDEX_CREATE_ELAPSED .with_label_values(&["cleanup", self.index_type]) .observe(self.cleanup_eplased.as_secs_f64()); - INDEX_CREATE_ELAPSED - .with_label_values(&["total", self.index_type]) - .observe( - (self.update_elapsed + self.finish_elapsed + self.cleanup_eplased).as_secs_f64(), - ); + if self.update_timing_available { + INDEX_CREATE_ELAPSED + .with_label_values(&["total", self.index_type]) + .observe( + (self.update_elapsed + + self.merge_elapsed + + self.finish_elapsed + + self.cleanup_eplased) + .as_secs_f64(), + ); + } INDEX_CREATE_ROWS_TOTAL .with_label_values(&[self.index_type]) @@ -111,7 +159,7 @@ impl Drop for Statistics { /// `TimerGuard` is a RAII struct that ensures elapsed time /// is recorded when it goes out of scope. -pub(crate) struct TimerGuard<'a> { +pub struct TimerGuard<'a> { stats: &'a mut Statistics, stage: Stage, timer: Instant, @@ -144,6 +192,10 @@ impl Drop for TimerGuard<'_> { Stage::Update => { self.stats.update_elapsed += self.timer.elapsed(); } + #[cfg(feature = "enterprise")] + Stage::Merge => { + self.stats.merge_elapsed += self.timer.elapsed(); + } Stage::Finish => { self.stats.finish_elapsed += self.timer.elapsed(); } @@ -158,6 +210,46 @@ impl Drop for TimerGuard<'_> { mod tests { use super::*; + #[cfg(feature = "enterprise")] + #[test] + fn test_unmeasured_updates_do_not_publish_zero_durations() { + let name = "test_sketch_fused_counter_only"; + let update = INDEX_CREATE_ELAPSED.with_label_values(&["update", name]); + let total = INDEX_CREATE_ELAPSED.with_label_values(&["total", name]); + let update_before = update.get_sample_count(); + let total_before = total.get_sample_count(); + let mut stats = Statistics::new_without_update_timing(name); + stats.inc_row_count(7); + assert_eq!(stats.row_count(), 7); + drop(stats); + assert_eq!(update.get_sample_count(), update_before); + assert_eq!(total.get_sample_count(), total_before); + let mut measured = Statistics::new_without_update_timing(name); + { + let _timer = measured.record_update(); + } + drop(measured); + assert_eq!(update.get_sample_count(), update_before + 1); + assert_eq!(total.get_sample_count(), total_before + 1); + } + + #[cfg(feature = "enterprise")] + #[test] + fn test_merge_does_not_count_input_rows_as_output_rows() { + let mut stats = Statistics::new("test_merge"); + { + let _guard = stats.record_merge(); + } + assert_eq!(stats.row_count(), 0); + { + let mut guard = stats.record_finish(); + guard.inc_row_count(7); + guard.inc_byte_count(32); + } + assert_eq!(stats.row_count(), 7); + assert_eq!(stats.byte_count(), 32); + } + #[test] fn test_statistics_basic() { let mut stats = Statistics::new("test"); diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index 1cbc47d8975..d8f4ecf8c16 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -906,7 +906,7 @@ async fn edit_region( let remote_path = location::sst_file_path(layer.table_dir(), file_meta.file_id(), layer.path_type()); - let is_index_exist = file_meta.exists_index(); + let has_downloadable_index = file_meta.has_downloadable_index(); let index_file_size = file_meta.index_file_size(); let index_file_index_key = IndexKey::new( @@ -940,12 +940,8 @@ async fn edit_region( PageIndexPolicy::Optional, ) .await; - - if matches!(parquet_cached, Ok(true)) { - listener.on_file_cache_filled(index_key.file_id); - } } - if is_index_exist { + if has_downloadable_index { // also download puffin file if let Err(err) = write_cache .download( @@ -963,6 +959,9 @@ async fn edit_region( } WRITE_CACHE_INFLIGHT_DOWNLOAD.sub(1); + if matches!(parquet_cached, Ok(true)) { + listener.on_file_cache_filled(index_key.file_id); + } }); } } diff --git a/tests/compatibility/README.md b/tests/compatibility/README.md index 31a447dbae3..b84b55c92e5 100644 --- a/tests/compatibility/README.md +++ b/tests/compatibility/README.md @@ -35,6 +35,19 @@ Frozen objects of the object store WAL are checked byte for byte by the format t mito2's Rust tests. They cover ordinary file-edit compatibility, supported SST levels, and explicit rejection of unpublished prototype task-state records. +The inline sketch-index fixtures in `src/mito2/src/manifest/action.rs` cover both +file-level and column-level descriptors and mixed builtin/sketch indexes. They run without the enterprise feature: + +```shell +cargo nextest run -p mito2 -E 'test(test_sketch_index_manifest_fixture)' +``` + +These fixtures verify supported-reader decoding, not downgrade support for binaries +that predate the new index types. Upgrade all manifest readers before enabling a +writer that emits these descriptors; downgrade only to readers that recognize them. +The OSS build does not generate sketch payloads, so SQL restart cases cannot create +these descriptors through an OSS writer. + ## Quick Start ```shell