From 57ee4c72f995eebe8f01c51b33eeb60d405b8437 Mon Sep 17 00:00:00 2001 From: Hengfei Yang Date: Tue, 19 May 2026 21:54:24 +0800 Subject: [PATCH] perf: store row_group_size in tantivy puffin footer and shrink to 128k (#11965) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Drop `PARQUET_MAX_ROW_GROUP_SIZE` from 1M to 128k. The constant was previously load-bearing for row_id → row_group mapping, which is why the old comment said it couldn't change. - Write the row-group-size value in effect at index-build time into each `.ttv` puffin file's footer (`row_group_size` property) so readers can recover it later. - Generate parquet access plans from tantivy hits using the **per-file** value (carried through `TantivyResult::RowIdsBitVec` → `CacheEntry` → `FileKey.row_group_size` → `SegmentInfo` in the SEGMENTS cache). - Falls back to `1024 * 1024` for legacy `.ttv` files written before the property existed, so old data keeps working unchanged. ## Why this matters With the constant in the path, the assumed row-group size at read time had to match the one at write time. Now each tantivy index is self-describing, so the writer's row-group-size limit can be tuned freely going forward without breaking row_id lookups against historical files. ## Notes for reviewers - No DB schema change. No parquet footer reads at query time. No tantivy/puffin binary format break — the property is additive via `PuffinMeta.properties` (already `#[serde(default)]`). - `FileKey.row_group_size` is set locally on the querier after tantivy search; it's not serialized over gRPC (parallels `segment_ids`). - `helpers.rs::generate_parquet_access_plan` now uses the per-file size; legacy `.ttv` falls back to `LEGACY_ROW_GROUP_SIZE = 1024 * 1024`. ## Test plan - [x] `cargo check --all-targets --workspace` clean - [x] `cargo test --package tantivy_utils` — all 153 tests pass - [x] `cargo test --package config --lib test_file_key` — pass - [ ] Ingest with the new 128k size, then search — confirm row_ids resolve correctly - [ ] Search against a pre-existing `.ttv` (no `row_group_size` property) — confirm fallback path still finds rows 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- src/cli/basic/test.rs | 1 + src/config/src/config.rs | 5 ++- src/config/src/meta/stream.rs | 13 ++++++- src/infra/src/file_list/mod.rs | 1 + src/infra/src/file_list/postgres.rs | 1 + src/infra/src/file_list/sqlite.rs | 1 + src/infra/src/storage/mod.rs | 18 ++++++++- src/service/compact/dump.rs | 2 + src/service/compact/merge.rs | 2 + src/service/db/file_list/broadcast.rs | 4 ++ src/service/file_list.rs | 1 + src/service/search/datafusion/exec.rs | 2 + .../datafusion/plan/tantivy_optimize_exec.rs | 10 +++++ .../search/datafusion/storage/file_list.rs | 29 ++++++++++---- .../datafusion/table_provider/helpers.rs | 30 +++++++++----- src/service/search/grpc/flight.rs | 3 +- src/service/search/grpc/storage.rs | 39 +++++++++++++------ src/service/search/grpc/tantivy.rs | 11 +++--- .../search/grpc/tantivy_result_cache.rs | 29 +++++++------- src/service/tantivy/mod.rs | 8 +++- src/tantivy_utils/src/puffin/writer.rs | 4 ++ src/tantivy_utils/src/puffin_directory/mod.rs | 5 +++ .../src/puffin_directory/reader.rs | 15 +++++++ .../src/puffin_directory/writer.rs | 16 +++++++- 24 files changed, 195 insertions(+), 55 deletions(-) diff --git a/src/cli/basic/test.rs b/src/cli/basic/test.rs index 61f657c7b2..59f6fae8eb 100644 --- a/src/cli/basic/test.rs +++ b/src/cli/basic/test.rs @@ -147,6 +147,7 @@ mod tests { }, deleted: false, segment_ids: None, + row_group_size: None, } } diff --git a/src/config/src/config.rs b/src/config/src/config.rs index e28a74fd84..1a20558cbe 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -67,7 +67,10 @@ pub const GEO_IP_ASN_ENRICHMENT_TABLE: &str = "maxmind_asn"; pub const SIZE_IN_MB: f64 = 1024.0 * 1024.0; pub const SIZE_IN_GB: f64 = 1024.0 * 1024.0 * 1024.0; -pub const PARQUET_MAX_ROW_GROUP_SIZE: usize = 1024 * 1024; // this can't be change, it will cause segment matching error +// The current value is recorded in each tantivy index file (puffin `row_group_size` +// property) so it can be changed safely without breaking row_id → row_group mapping +// for older files. +pub const PARQUET_MAX_ROW_GROUP_SIZE: usize = 128 * 1024; pub const PARQUET_FILE_CHUNK_SIZE: usize = 100 * 1024; // 100k, num_rows pub const DEFAULT_BLOOM_FILTER_FPP: f64 = 0.01; pub const SOURCEMAP_ZIP_MAX_SIZE: usize = 1024 * 1024 * 100; // 100 MB diff --git a/src/config/src/meta/stream.rs b/src/config/src/meta/stream.rs index ee5cec26f6..06c058f759 100644 --- a/src/config/src/meta/stream.rs +++ b/src/config/src/meta/stream.rs @@ -281,6 +281,10 @@ pub struct FileKey { pub meta: FileMeta, pub deleted: bool, pub segment_ids: Option>, + /// Parquet row group size that was in effect when the tantivy index for + /// this file was built. Set together with `segment_ids` after tantivy + /// search runs locally; not serialized over gRPC. + pub row_group_size: Option, } impl FileKey { @@ -292,6 +296,7 @@ impl FileKey { meta, deleted, segment_ids: None, + row_group_size: None, } } @@ -303,11 +308,13 @@ impl FileKey { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, } } - pub fn with_segment_ids(&mut self, segment_ids: BitVec) { + pub fn with_segment_ids(&mut self, segment_ids: BitVec, row_group_size: Option) { self.segment_ids = Some(Arc::new(segment_ids)); + self.row_group_size = row_group_size; } } @@ -669,6 +676,7 @@ impl From<&cluster_rpc::FileKey> for FileKey { meta: FileMeta::from(req.meta.as_ref().unwrap()), deleted: req.deleted, segment_ids: None, + row_group_size: None, } } } @@ -1890,8 +1898,9 @@ mod tests { let mut key = FileKey::from_file_name("files/k.parquet"); assert!(key.segment_ids.is_none()); let bv = BitVec::new(); - key.with_segment_ids(bv); + key.with_segment_ids(bv, Some(1024)); assert!(key.segment_ids.is_some()); + assert_eq!(key.row_group_size, Some(1024)); } #[test] diff --git a/src/infra/src/file_list/mod.rs b/src/infra/src/file_list/mod.rs index 3320e51a05..1723a4ef48 100644 --- a/src/infra/src/file_list/mod.rs +++ b/src/infra/src/file_list/mod.rs @@ -682,6 +682,7 @@ impl From<&FileRecord> for FileKey { meta: r.into(), deleted: r.deleted, segment_ids: None, + row_group_size: None, } } } diff --git a/src/infra/src/file_list/postgres.rs b/src/infra/src/file_list/postgres.rs index c97b309efb..fac6f09fba 100644 --- a/src/infra/src/file_list/postgres.rs +++ b/src/infra/src/file_list/postgres.rs @@ -3384,6 +3384,7 @@ mod tests { deleted, id: 0, segment_ids: None, + row_group_size: None, } } diff --git a/src/infra/src/file_list/sqlite.rs b/src/infra/src/file_list/sqlite.rs index 45e4bab0ad..c03d998044 100644 --- a/src/infra/src/file_list/sqlite.rs +++ b/src/infra/src/file_list/sqlite.rs @@ -2026,6 +2026,7 @@ mod tests { deleted, id: 0, segment_ids: None, + row_group_size: None, } } diff --git a/src/infra/src/storage/mod.rs b/src/infra/src/storage/mod.rs index a7c22f6505..9f6c848fb3 100644 --- a/src/infra/src/storage/mod.rs +++ b/src/infra/src/storage/mod.rs @@ -21,7 +21,11 @@ use std::{ use async_trait::async_trait; use bytes::{Bytes, buf::Buf}; -use config::{get_config, is_local_disk_storage, meta::stream::FileMeta, metrics}; +use config::{ + get_config, is_local_disk_storage, + meta::stream::{FileKey, FileMeta}, + metrics, +}; use datafusion::parquet::{data_type::AsBytes, file::metadata::ParquetMetaData}; use futures::{StreamExt, TryStreamExt, stream::BoxStream}; use hashbrown::HashMap; @@ -318,6 +322,18 @@ pub async fn del(files: Vec<(&str, &str)>) -> Result<()> { Ok(()) } +pub async fn get_row_group_size(file: &FileKey) -> u32 { + if let Some(row_group_size) = file.row_group_size { + return row_group_size; + } + get_parquet_metadata(&file.account, &file.key) + .await + .ok() + .and_then(|(_, meta)| meta.row_groups().iter().map(|rg| rg.num_rows()).max()) + .unwrap_or_default() + .max(config::PARQUET_MAX_ROW_GROUP_SIZE as i64) as u32 +} + pub async fn get_file_meta(account: &str, file: &str) -> Result { let mut file_meta = FileMeta::default(); let (file_size, parquet_meta) = get_parquet_metadata(account, file).await?; diff --git a/src/service/compact/dump.rs b/src/service/compact/dump.rs index f0567e0113..63b0137a7b 100644 --- a/src/service/compact/dump.rs +++ b/src/service/compact/dump.rs @@ -463,6 +463,7 @@ pub async fn delete_by_time_range( .map(|mut f| { f.deleted = true; f.segment_ids = None; + f.row_group_size = None; f }) .collect(); @@ -602,6 +603,7 @@ async fn generate_dump( meta, deleted: false, segment_ids: None, + row_group_size: None, }; Ok(Some(dump_file)) diff --git a/src/service/compact/merge.rs b/src/service/compact/merge.rs index 91e012580d..a57df92f3d 100644 --- a/src/service/compact/merge.rs +++ b/src/service/compact/merge.rs @@ -602,6 +602,7 @@ pub async fn merge_by_stream( events.push(FileKey { deleted: true, segment_ids: None, + row_group_size: None, ..file.clone() }); } @@ -1331,6 +1332,7 @@ mod tests { }, deleted: false, segment_ids: None, + row_group_size: None, } } diff --git a/src/service/db/file_list/broadcast.rs b/src/service/db/file_list/broadcast.rs index cf58044f30..054afa5d97 100644 --- a/src/service/db/file_list/broadcast.rs +++ b/src/service/db/file_list/broadcast.rs @@ -318,6 +318,7 @@ mod tests { meta: config::meta::stream::FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }, FileKey { id: 2, @@ -326,6 +327,7 @@ mod tests { meta: config::meta::stream::FileMeta::default(), deleted: true, segment_ids: None, + row_group_size: None, }, ]; @@ -355,6 +357,7 @@ mod tests { meta: config::meta::stream::FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }); assert_eq!(queue.len(), initial_len + 1); @@ -379,6 +382,7 @@ mod tests { meta: config::meta::stream::FileMeta::default(), deleted, segment_ids: None, + row_group_size: None, } } diff --git a/src/service/file_list.rs b/src/service/file_list.rs index 555125717a..1e267ebcc9 100644 --- a/src/service/file_list.rs +++ b/src/service/file_list.rs @@ -353,6 +353,7 @@ mod tests { }, deleted: false, segment_ids: None, + row_group_size: None, } } diff --git a/src/service/search/datafusion/exec.rs b/src/service/search/datafusion/exec.rs index e15f642737..d4c2ab63c6 100644 --- a/src/service/search/datafusion/exec.rs +++ b/src/service/search/datafusion/exec.rs @@ -817,6 +817,7 @@ mod tests { account: "test_account".to_string(), id: 1, segment_ids: None, + row_group_size: None, }]; let result = register_metrics_table(&session, schema, "test_table", files).await; @@ -856,6 +857,7 @@ mod tests { account: "test_account".to_string(), id: 1, segment_ids: None, + row_group_size: None, }]; let builder = TableBuilder::new().sorted_by_time(true); diff --git a/src/service/search/datafusion/plan/tantivy_optimize_exec.rs b/src/service/search/datafusion/plan/tantivy_optimize_exec.rs index bd8855ed2d..fd2abf7690 100644 --- a/src/service/search/datafusion/plan/tantivy_optimize_exec.rs +++ b/src/service/search/datafusion/plan/tantivy_optimize_exec.rs @@ -985,6 +985,7 @@ mod tests { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }]; let index_condition = None; let index_optimize_mode = IndexOptimizeMode::SimpleCount; @@ -1027,6 +1028,7 @@ mod tests { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }, FileKey { id: 2, @@ -1035,6 +1037,7 @@ mod tests { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }, FileKey { id: 3, @@ -1043,6 +1046,7 @@ mod tests { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }, FileKey { id: 4, @@ -1051,6 +1055,7 @@ mod tests { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }, FileKey { id: 5, @@ -1059,6 +1064,7 @@ mod tests { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }, FileKey { id: 6, @@ -1067,6 +1073,7 @@ mod tests { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }, ]; let index_condition = None; @@ -1115,6 +1122,7 @@ mod tests { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }]; let index_condition = None; let index_optimize_mode = IndexOptimizeMode::SimpleCount; @@ -1156,6 +1164,7 @@ mod tests { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }]; let index_condition = None; let index_optimize_mode = IndexOptimizeMode::SimpleCount; @@ -1198,6 +1207,7 @@ mod tests { meta: FileMeta::default(), deleted: false, segment_ids: None, + row_group_size: None, }]; let index_condition = None; let index_optimize_mode = IndexOptimizeMode::SimpleCount; diff --git a/src/service/search/datafusion/storage/file_list.rs b/src/service/search/datafusion/storage/file_list.rs index 6064f7d26a..a25b9a664a 100644 --- a/src/service/search/datafusion/storage/file_list.rs +++ b/src/service/search/datafusion/storage/file_list.rs @@ -23,7 +23,16 @@ use parking_lot::RwLock; use super::{ACCOUNT_SEPARATOR, TRACE_ID_SEPARATOR}; -type SegmentData = HashMap>; +#[derive(Clone)] +pub struct SegmentInfo { + pub segment_ids: Arc, + /// Parquet row group size that was in effect when the tantivy index for + /// this file was built. `None` for files written before this property + /// existed; callers should fall back to the legacy value. + pub row_group_size: Option, +} + +type SegmentData = HashMap; static FILES: Lazy>>> = Lazy::new(Default::default); static SEGMENTS: Lazy>> = Lazy::new(Default::default); @@ -58,7 +67,13 @@ pub async fn set(trace_id: &str, schema_key: &str, format: &str, files: Vec Option> { +pub fn get_segment_info(file_key: &str) -> Option { let (trace_id, filename) = file_key.split_once("/$$/")?; let r = SEGMENTS.read(); let data = r.get(trace_id)?; @@ -115,14 +130,14 @@ mod tests { } #[test] - fn test_get_segment_ids_nonexistent_returns_none() { - let result = get_segment_ids("nonexistent_trace_file_list/$$/filename.parquet"); + fn test_get_segment_info_nonexistent_returns_none() { + let result = get_segment_info("nonexistent_trace_file_list/$$/filename.parquet"); assert!(result.is_none()); } #[test] - fn test_get_segment_ids_no_separator_returns_none() { - let result = get_segment_ids("no-separator-here"); + fn test_get_segment_info_no_separator_returns_none() { + let result = get_segment_info("no-separator-here"); assert!(result.is_none()); } } diff --git a/src/service/search/datafusion/table_provider/helpers.rs b/src/service/search/datafusion/table_provider/helpers.rs index 296bcd1f17..38331821cb 100644 --- a/src/service/search/datafusion/table_provider/helpers.rs +++ b/src/service/search/datafusion/table_provider/helpers.rs @@ -16,7 +16,7 @@ use std::sync::Arc; use arrow_schema::{DataType, SchemaRef}; -use config::{FileFormat, PARQUET_MAX_ROW_GROUP_SIZE, TIMESTAMP_COL_NAME, meta::bitvec::BitVec}; +use config::{FileFormat, TIMESTAMP_COL_NAME, meta::bitvec::BitVec}; use datafusion::{ common::{DataFusionError, Result, project_schema, stats::Precision}, datasource::{listing::PartitionedFile, physical_plan::parquet::ParquetAccessPlan}, @@ -37,15 +37,21 @@ use o2_enterprise::enterprise::search::vortex::generate_vortex_access_plan; use crate::service::search::{datafusion::storage, index::IndexCondition}; +/// Row group size used by writers before the `row_group_size` puffin property +/// existed. Any .ttv file without the property was produced with this value. +const LEGACY_ROW_GROUP_SIZE: usize = 1024 * 1024; + pub fn generate_access_plan( file: &PartitionedFile, ) -> Option> { - let segment_ids = storage::file_list::get_segment_ids(file.path().as_ref())?; + let info = storage::file_list::get_segment_info(file.path().as_ref())?; let file_format = FileFormat::from_extension(file.path().as_ref())?; match file_format { - FileFormat::Parquet => generate_parquet_access_plan(file, segment_ids), + FileFormat::Parquet => { + generate_parquet_access_plan(file, info.segment_ids, info.row_group_size) + } #[cfg(all(feature = "enterprise", feature = "vortex"))] - FileFormat::Vortex => generate_vortex_access_plan(segment_ids), + FileFormat::Vortex => generate_vortex_access_plan(info.segment_ids), #[cfg(not(all(feature = "enterprise", feature = "vortex")))] FileFormat::Vortex => None, } @@ -54,12 +60,19 @@ pub fn generate_access_plan( fn generate_parquet_access_plan( file: &PartitionedFile, segment_ids: Arc, + row_group_size: Option, ) -> Option> { let stats = file.statistics.as_ref()?; let Precision::Exact(num_rows) = stats.num_rows else { return None; }; - let row_group_count = num_rows.div_ceil(PARQUET_MAX_ROW_GROUP_SIZE); + // Use the row group size recorded in the tantivy index when it was built. + // Fall back to PARQUET_MAX_ROW_GROUP_SIZE for legacy .ttv files that + // predate this property (those were always written with the 1M default). + let row_group_size = row_group_size + .map(|v| v as usize) + .unwrap_or(LEGACY_ROW_GROUP_SIZE); + let row_group_count = num_rows.div_ceil(row_group_size); // Determine sampling mode based on BitVec size: // - If BitVec size == row_group_count: row-group-level sampling (enterprise feature) @@ -77,15 +90,12 @@ fn generate_parquet_access_plan( } let mut access_plan = ParquetAccessPlan::new_none(row_group_count); - for (row_group_id, chunk) in segment_ids.chunks(PARQUET_MAX_ROW_GROUP_SIZE).enumerate() { + for (row_group_id, chunk) in segment_ids.chunks(row_group_size).enumerate() { let mut selection = Vec::new(); let mut current_count = 0; let mut current_select = false; - for val in chunk - .iter() - .take(num_rows - row_group_id * PARQUET_MAX_ROW_GROUP_SIZE) - { + for val in chunk.iter().take(num_rows - row_group_id * row_group_size) { if *val == current_select { current_count += 1; } else { diff --git a/src/service/search/grpc/flight.rs b/src/service/search/grpc/flight.rs index 2f3e3de37b..ddf2584211 100644 --- a/src/service/search/grpc/flight.rs +++ b/src/service/search/grpc/flight.rs @@ -282,7 +282,8 @@ pub async fn search( Some(query_params.time_range), req.search_info.histogram_interval, &trace_id, - ); + ) + .await; } let storage_search_start = std::time::Instant::now(); diff --git a/src/service/search/grpc/storage.rs b/src/service/search/grpc/storage.rs index ee70af1002..5433f5dc97 100644 --- a/src/service/search/grpc/storage.rs +++ b/src/service/search/grpc/storage.rs @@ -62,6 +62,7 @@ use crate::service::{ inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, }, tantivy::puffin_directory::{ + PROP_ROW_GROUP_SIZE, caching_directory::CachingDirectory, footer_cache::FooterCache, reader::{PuffinDirReader, warm_up_terms}, @@ -620,7 +621,7 @@ pub async fn tantivy_search( continue; } match result { - TantivyResult::RowIdsBitVec(num_rows, bitvec) => { + TantivyResult::RowIdsBitVec(bitvec, num_rows, row_group_size) => { if num_rows == 0 { // if the bitmap is empty then we remove the file from the list file_list_map.remove(&file_name); @@ -628,7 +629,7 @@ pub async fn tantivy_search( // Replace the segment IDs in the existing `FileKey` with the found tantivy_result_builder.add_row_nums(num_rows as u64); let file = file_list_map.get_mut(&file_name).unwrap(); - file.with_segment_ids(bitvec); + file.with_segment_ids(bitvec, row_group_size); } } TantivyResult::Count(count) => { @@ -763,6 +764,12 @@ async fn search_tantivy_index( ) .await?, ); + // Read the row group size that the writer used when this tantivy index was + // built. Old .ttv files predate this property — None falls back to the + // legacy assumption in the access-plan code. + let row_group_size = puffin_dir + .get_property(PROP_ROW_GROUP_SIZE) + .and_then(|s| s.parse::().ok()); let footer_cache = FooterCache::from_directory(puffin_dir.clone(), &ttv_file_name).await?; let cache_dir = CachingDirectory::new_with_cacher(puffin_dir, Arc::new(footer_cache)); let reader_directory: Box = Box::new(cache_dir); @@ -898,7 +905,11 @@ async fn search_tantivy_index( TantivyResult::Distinct(distinct) => TantivyResult::Distinct(distinct), TantivyResult::RowIds(row_ids) => { if row_ids.is_empty() || parquet_file.meta.records == 0 { - return Ok((key, TantivyResult::RowIdsBitVec(0, BitVec::EMPTY), false)); + return Ok(( + key, + TantivyResult::RowIdsBitVec(BitVec::EMPTY, 0, row_group_size), + false, + )); } // return early if the number of matched docs is too large let skip_threshold = cfg.limit.inverted_index_skip_threshold; @@ -911,7 +922,11 @@ async fn search_tantivy_index( ); return Ok(( "".to_string(), - TantivyResult::RowIdsBitVec(row_ids_percent as usize, BitVec::EMPTY), + TantivyResult::RowIdsBitVec( + BitVec::EMPTY, + row_ids_percent as usize, + row_group_size, + ), true, )); } @@ -929,7 +944,7 @@ async fn search_tantivy_index( for id in row_ids { res.set(id as usize, true); } - TantivyResult::RowIdsBitVec(num_rows, res) + TantivyResult::RowIdsBitVec(res, num_rows, row_group_size) } TantivyResult::RowIdsBitVec(..) => { unreachable!("unsupported tantivy search result in search_tantivy_index") @@ -1075,7 +1090,7 @@ fn group_files_by_time_range(mut files: Vec, partition_num: usize) -> V fn get_cache_entry(tantivy_result: TantivyResult, percent: f64, parquet_rows: usize) -> CacheEntry { match tantivy_result { - TantivyResult::RowIdsBitVec(num_rows, bitvec) => { + TantivyResult::RowIdsBitVec(bitvec, num_rows, row_group_size) => { // if the percent is less than 1.0, we use roaring bitmap to store the row ids // otherwise, we use bitvec to store the row ids. // because the bitvec is not efficient for small percent, and the roaring bitmap is not @@ -1087,9 +1102,9 @@ fn get_cache_entry(tantivy_result: TantivyResult, percent: f64, parquet_rows: us roaring.insert(i as u32); } } - CacheEntry::RowIdsRoaring(num_rows, roaring, parquet_rows) + CacheEntry::RowIdsRoaring(roaring, num_rows, parquet_rows, row_group_size) } else { - CacheEntry::RowIdsBitVec(num_rows, bitvec) + CacheEntry::RowIdsBitVec(bitvec, num_rows, row_group_size) } } TantivyResult::Count(count) => CacheEntry::Count(count), @@ -1513,13 +1528,13 @@ mod tests { let mut bitvec = BitVec::repeat(false, 4); bitvec.set(0, true); bitvec.set(2, true); - let result = TantivyResult::RowIdsBitVec(2, bitvec); + let result = TantivyResult::RowIdsBitVec(bitvec, 2, None); let percent = 0.5; // Less than 1.0, should use roaring bitmap let parquet_rows = 4; let entry = get_cache_entry(result, percent, parquet_rows); match entry { - tantivy_result_cache::CacheEntry::RowIdsRoaring(num_rows, roaring, rows) => { + tantivy_result_cache::CacheEntry::RowIdsRoaring(roaring, num_rows, rows, _) => { assert_eq!(num_rows, 2); assert_eq!(rows, 4); assert_eq!(roaring.len(), 2); // Should contain positions 0 and 2 @@ -1536,13 +1551,13 @@ mod tests { bitvec.set(0, true); bitvec.set(1, true); bitvec.set(3, true); - let result = TantivyResult::RowIdsBitVec(3, bitvec); + let result = TantivyResult::RowIdsBitVec(bitvec, 3, None); let percent = 2.0; // Greater than 1.0, should use bitvec let parquet_rows = 4; let entry = get_cache_entry(result, percent, parquet_rows); match entry { - tantivy_result_cache::CacheEntry::RowIdsBitVec(num_rows, returned_bitvec) => { + tantivy_result_cache::CacheEntry::RowIdsBitVec(returned_bitvec, num_rows, _) => { assert_eq!(num_rows, 3); assert_eq!(returned_bitvec.len(), 4); assert_eq!(returned_bitvec.get(0).unwrap(), true); diff --git a/src/service/search/grpc/tantivy.rs b/src/service/search/grpc/tantivy.rs index 1b1101a921..0ab283c523 100644 --- a/src/service/search/grpc/tantivy.rs +++ b/src/service/search/grpc/tantivy.rs @@ -39,7 +39,8 @@ use crate::service::search::index::IndexCondition; #[derive(Debug, Clone)] pub enum TantivyResult { RowIds(HashSet), - RowIdsBitVec(usize, BitVec), + /// (row_id_bitvec, matched_row_count, row_group_size_from_index_file) + RowIdsBitVec(BitVec, usize, Option), Count(usize), // simple count optimization Histogram(Vec), // simple histogram optimization TopN(Vec<(String, u64)>), // simple top n optimization @@ -50,7 +51,7 @@ impl TantivyResult { // used for skip tantivy search pub fn percent(&self) -> usize { match self { - Self::RowIdsBitVec(percent, _) => *percent, + Self::RowIdsBitVec(_, percent, _) => *percent, _ => 0, } } @@ -61,7 +62,7 @@ impl TantivyResult { row_ids.capacity() * std::mem::size_of::() + std::mem::size_of::>() } - Self::RowIdsBitVec(_, bitvec) => { + Self::RowIdsBitVec(bitvec, ..) => { bitvec.capacity().div_ceil(8) + std::mem::size_of::() } Self::Count(_) => std::mem::size_of::(), @@ -389,7 +390,7 @@ mod tests { #[test] fn test_tantivy_result_percent() { - let result = TantivyResult::RowIdsBitVec(75, BitVec::repeat(false, 100)); + let result = TantivyResult::RowIdsBitVec(BitVec::repeat(false, 100), 75, None); assert_eq!(result.percent(), 75); let result = TantivyResult::RowIds(HashSet::new()); @@ -417,7 +418,7 @@ mod tests { #[test] fn test_tantivy_result_get_memory_size_bitvec() { let bitvec = BitVec::repeat(false, 1000); - let result = TantivyResult::RowIdsBitVec(50, bitvec); + let result = TantivyResult::RowIdsBitVec(bitvec, 50, None); let memory_size = result.get_memory_size(); // Should include BitVec overhead + bit capacity / 8 diff --git a/src/service/search/grpc/tantivy_result_cache.rs b/src/service/search/grpc/tantivy_result_cache.rs index 46f2da51ed..b9e86d189b 100644 --- a/src/service/search/grpc/tantivy_result_cache.rs +++ b/src/service/search/grpc/tantivy_result_cache.rs @@ -29,9 +29,10 @@ pub static GLOBAL_CACHE: Lazy> = #[derive(Debug, Clone)] pub enum CacheEntry { - RowIdsBitVec(usize, BitVec), - // true number in bitmap, bitmap, parquet row numbers - RowIdsRoaring(usize, RoaringBitmap, usize), + /// (row_id_bitvec, matched_row_count, row_group_size_from_index_file) + RowIdsBitVec(BitVec, usize, Option), + /// (sparse_row_ids, matched_row_count, parquet_row_count, row_group_size_from_index_file) + RowIdsRoaring(RoaringBitmap, usize, usize, Option), Count(usize), // simple count optimization Histogram(Vec), // simple histogram optimization TopN(Vec<(String, u64)>), // simple top n optimization @@ -41,15 +42,15 @@ pub enum CacheEntry { impl From for TantivyResult { fn from(entry: CacheEntry) -> Self { match entry { - CacheEntry::RowIdsBitVec(num_rows, bitvec) => { - TantivyResult::RowIdsBitVec(num_rows, bitvec) + CacheEntry::RowIdsBitVec(bitvec, num_rows, row_group_size) => { + TantivyResult::RowIdsBitVec(bitvec, num_rows, row_group_size) } - CacheEntry::RowIdsRoaring(num_rows, roaring, parquet_rows) => { + CacheEntry::RowIdsRoaring(roaring, num_rows, parquet_rows, row_group_size) => { let mut bitvec = BitVec::repeat(false, parquet_rows); for i in roaring.into_iter() { bitvec.set(i as usize, true); } - TantivyResult::RowIdsBitVec(num_rows, bitvec) + TantivyResult::RowIdsBitVec(bitvec, num_rows, row_group_size) } CacheEntry::Count(count) => TantivyResult::Count(count), CacheEntry::Histogram(histogram) => TantivyResult::Histogram(histogram), @@ -62,10 +63,10 @@ impl From for TantivyResult { impl CacheEntry { pub fn get_memory_size(&self) -> usize { match self { - CacheEntry::RowIdsBitVec(_, bitvec) => { + CacheEntry::RowIdsBitVec(bitvec, ..) => { bitvec.capacity().div_ceil(8) + std::mem::size_of::() } - CacheEntry::RowIdsRoaring(_, roaring, _) => { + CacheEntry::RowIdsRoaring(roaring, ..) => { roaring.serialized_size() + std::mem::size_of::() + std::mem::size_of::() * 2 @@ -177,7 +178,7 @@ mod tests { bitvec.set(10, true); bitvec.set(20, true); bitvec.set(30, true); - CacheEntry::RowIdsBitVec(3, bitvec) + CacheEntry::RowIdsBitVec(bitvec, 3, None) } fn create_test_count_result() -> CacheEntry { @@ -209,7 +210,7 @@ mod tests { bitvec.set(10, true); bitvec.set(20, true); bitvec.set(30, true); - CacheEntry::RowIdsBitVec(3, bitvec) + CacheEntry::RowIdsBitVec(bitvec, 3, None) } #[test] @@ -230,7 +231,7 @@ mod tests { assert!(retrieved.is_some()); assert!(matches!( retrieved.unwrap(), - TantivyResult::RowIdsBitVec(_, _) + TantivyResult::RowIdsBitVec(_, _, _) )); } @@ -453,11 +454,11 @@ mod tests { let mut bitvec = BitVec::repeat(false, 100); bitvec.set(10, true); bitvec.set(20, true); - let bitvec_result = CacheEntry::RowIdsBitVec(25, bitvec); + let bitvec_result = CacheEntry::RowIdsBitVec(bitvec, 25, None); cache.put("percent_key".to_string(), bitvec_result); - if let Some(TantivyResult::RowIdsBitVec(percent, _)) = cache.get("percent_key") { + if let Some(TantivyResult::RowIdsBitVec(_, percent, ..)) = cache.get("percent_key") { assert_eq!(percent, 25); } else { panic!("Expected RowIdsBitVec result"); diff --git a/src/service/tantivy/mod.rs b/src/service/tantivy/mod.rs index 454a276f99..54b226c8f4 100644 --- a/src/service/tantivy/mod.rs +++ b/src/service/tantivy/mod.rs @@ -26,7 +26,7 @@ use arrow::array::{ use arrow_schema::{DataType, Schema}; use bytes::Bytes; use config::{ - INDEX_FIELD_NAME_FOR_ALL, TIMESTAMP_COL_NAME, get_config, + INDEX_FIELD_NAME_FOR_ALL, PARQUET_MAX_ROW_GROUP_SIZE, TIMESTAMP_COL_NAME, get_config, utils::{ inverted_index::convert_parquet_file_name_to_tantivy_file, parquet::RecordBatchStream, @@ -98,6 +98,12 @@ pub(crate) async fn create_tantivy_index( if index.is_none() { return Ok(0); } + // Record the parquet row group size in effect at index build time so the + // reader can map doc_ids back to row groups even if the constant changes. + dir.set_property( + puffin_directory::PROP_ROW_GROUP_SIZE, + PARQUET_MAX_ROW_GROUP_SIZE.to_string(), + ); let puffin_bytes = dir.to_puffin_bytes()?; let index_size = puffin_bytes.len(); diff --git a/src/tantivy_utils/src/puffin/writer.rs b/src/tantivy_utils/src/puffin/writer.rs index 182f00cb6b..49ebcf7edc 100644 --- a/src/tantivy_utils/src/puffin/writer.rs +++ b/src/tantivy_utils/src/puffin/writer.rs @@ -43,6 +43,10 @@ impl PuffinBytesWriter { } } + pub fn set_property(&mut self, key: impl Into, value: impl Into) { + self.properties.insert(key.into(), value.into()); + } + fn build_blob_metadata( &self, blob_type: BlobTypes, diff --git a/src/tantivy_utils/src/puffin_directory/mod.rs b/src/tantivy_utils/src/puffin_directory/mod.rs index 37d4c2b4b2..0d59d2c198 100644 --- a/src/tantivy_utils/src/puffin_directory/mod.rs +++ b/src/tantivy_utils/src/puffin_directory/mod.rs @@ -39,6 +39,11 @@ const EMPTY_FILE_EXT: &[&str] = &["fieldnorm", "store"]; const META_JSON: &str = "meta.json"; const FOOTER_CACHE: &str = "footer_cache"; +/// Puffin file-level property name carrying the parquet row group size in +/// effect when the index was built. Lets readers map tantivy doc_ids back to +/// parquet row groups even if `PARQUET_MAX_ROW_GROUP_SIZE` changes later. +pub const PROP_ROW_GROUP_SIZE: &str = "row_group_size"; + // Lazy loaded global instance of RAM directory which will contain // all the files of an empty tantivy index. This instance will be used to fill the missing files // from the `.ttv` file, as tantivy needs them regardless of the configuration of a field. diff --git a/src/tantivy_utils/src/puffin_directory/reader.rs b/src/tantivy_utils/src/puffin_directory/reader.rs index a641859653..efbc4c7ea4 100644 --- a/src/tantivy_utils/src/puffin_directory/reader.rs +++ b/src/tantivy_utils/src/puffin_directory/reader.rs @@ -39,6 +39,7 @@ use crate::{ pub struct PuffinDirReader { source: Arc, blobs_metadata: Arc>>, + properties: Arc>, } impl PuffinDirReader { @@ -60,11 +61,18 @@ impl PuffinDirReader { } } + let properties = metadata.properties.into_iter().collect::>(); + Ok(Self { source: Arc::new(source), blobs_metadata: Arc::new(blobs_metadata), + properties: Arc::new(properties), }) } + + pub fn get_property(&self, key: &str) -> Option<&str> { + self.properties.get(key).map(|s| s.as_str()) + } } impl Clone for PuffinDirReader { @@ -72,6 +80,7 @@ impl Clone for PuffinDirReader { PuffinDirReader { source: self.source.clone(), blobs_metadata: self.blobs_metadata.clone(), + properties: self.properties.clone(), } } } @@ -448,6 +457,7 @@ mod tests { let reader = PuffinDirReader { source: Arc::new(mock_reader), blobs_metadata: Arc::new(blobs_metadata), + properties: Arc::new(HashbrownHashMap::new()), }; let result = reader.get_file_handle(&path); @@ -465,6 +475,7 @@ mod tests { let reader = PuffinDirReader { source: Arc::new(mock_reader), blobs_metadata: Arc::new(blobs_metadata), + properties: Arc::new(HashbrownHashMap::new()), }; let path = PathBuf::from("nonexistent_file.terms"); @@ -494,6 +505,7 @@ mod tests { let reader = PuffinDirReader { source: Arc::new(mock_reader), blobs_metadata: Arc::new(blobs_metadata), + properties: Arc::new(HashbrownHashMap::new()), }; let path = PathBuf::from("file_without_extension"); @@ -519,6 +531,7 @@ mod tests { let reader = PuffinDirReader { source: Arc::new(mock_reader), blobs_metadata: Arc::new(blobs_metadata), + properties: Arc::new(HashbrownHashMap::new()), }; let result = reader.exists(&path); @@ -536,6 +549,7 @@ mod tests { let reader = PuffinDirReader { source: Arc::new(mock_reader), blobs_metadata: Arc::new(blobs_metadata), + properties: Arc::new(HashbrownHashMap::new()), }; let path = PathBuf::from("nonexistent_file.unknown"); @@ -556,6 +570,7 @@ mod tests { let reader = PuffinDirReader { source: Arc::new(mock_reader), blobs_metadata: Arc::new(blobs_metadata), + properties: Arc::new(HashbrownHashMap::new()), }; let path = PathBuf::from("test_file.txt"); diff --git a/src/tantivy_utils/src/puffin_directory/writer.rs b/src/tantivy_utils/src/puffin_directory/writer.rs index ca32ac3a5c..21241c82ec 100644 --- a/src/tantivy_utils/src/puffin_directory/writer.rs +++ b/src/tantivy_utils/src/puffin_directory/writer.rs @@ -20,7 +20,7 @@ use std::{ }; use anyhow::{Context, Result}; -use hashbrown::HashSet; +use hashbrown::{HashMap, HashSet}; use tantivy::{ HasLen, directory::{Directory, MmapDirectory, WatchCallback, WatchHandle, error::OpenReadError}, @@ -38,6 +38,8 @@ pub struct PuffinDirWriter { mmap_directory: Arc, /// record all the files paths in the puffin file file_paths: Arc>>, + /// file-level puffin properties (written to PuffinMeta.properties) + properties: Arc>>, } impl Default for PuffinDirWriter { @@ -51,6 +53,7 @@ impl Clone for PuffinDirWriter { PuffinDirWriter { mmap_directory: self.mmap_directory.clone(), file_paths: self.file_paths.clone(), + properties: self.properties.clone(), } } } @@ -63,6 +66,7 @@ impl PuffinDirWriter { .expect("failed to create temporary mmap directory"), ), file_paths: Arc::new(RwLock::new(HashSet::default())), + properties: Arc::new(RwLock::new(HashMap::default())), } } @@ -75,10 +79,20 @@ impl PuffinDirWriter { .collect() } + pub fn set_property(&self, key: impl Into, value: impl Into) { + self.properties + .write() + .expect("poisoned lock") + .insert(key.into(), value.into()); + } + // This function will serialize the directory into a single puffin file pub fn to_puffin_bytes(&self) -> Result> { let mut puffin_buf: Vec = Vec::new(); let mut puffin_writer = PuffinBytesWriter::new(&mut puffin_buf); + for (k, v) in self.properties.read().expect("poisoned lock").iter() { + puffin_writer.set_property(k.clone(), v.clone()); + } let mut segment_id = String::new(); let file_paths = self.file_paths.read().expect("poisoned lock");