mirror of
https://github.com/openobserve/openobserve.git
synced 2026-11-06 00:09:22 +00:00
perf: store row_group_size in tantivy puffin footer and shrink to 128k (#11965)
## 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) <noreply@anthropic.com>
This commit is contained in:
Hengfei Yang
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
f877299d31
commit
57ee4c72f9
@@ -147,6 +147,7 @@ mod tests {
|
||||
},
|
||||
deleted: false,
|
||||
segment_ids: None,
|
||||
row_group_size: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -281,6 +281,10 @@ pub struct FileKey {
|
||||
pub meta: FileMeta,
|
||||
pub deleted: bool,
|
||||
pub segment_ids: Option<Arc<BitVec>>,
|
||||
/// 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<u32>,
|
||||
}
|
||||
|
||||
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<u32>) {
|
||||
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]
|
||||
|
||||
@@ -682,6 +682,7 @@ impl From<&FileRecord> for FileKey {
|
||||
meta: r.into(),
|
||||
deleted: r.deleted,
|
||||
segment_ids: None,
|
||||
row_group_size: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3384,6 +3384,7 @@ mod tests {
|
||||
deleted,
|
||||
id: 0,
|
||||
segment_ids: None,
|
||||
row_group_size: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2026,6 +2026,7 @@ mod tests {
|
||||
deleted,
|
||||
id: 0,
|
||||
segment_ids: None,
|
||||
row_group_size: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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<FileMeta, anyhow::Error> {
|
||||
let mut file_meta = FileMeta::default();
|
||||
let (file_size, parquet_meta) = get_parquet_metadata(account, file).await?;
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -353,6 +353,7 @@ mod tests {
|
||||
},
|
||||
deleted: false,
|
||||
segment_ids: None,
|
||||
row_group_size: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -23,7 +23,16 @@ use parking_lot::RwLock;
|
||||
|
||||
use super::{ACCOUNT_SEPARATOR, TRACE_ID_SEPARATOR};
|
||||
|
||||
type SegmentData = HashMap<String, Arc<BitVec>>;
|
||||
#[derive(Clone)]
|
||||
pub struct SegmentInfo {
|
||||
pub segment_ids: Arc<BitVec>,
|
||||
/// 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<u32>,
|
||||
}
|
||||
|
||||
type SegmentData = HashMap<String, SegmentInfo>;
|
||||
|
||||
static FILES: Lazy<RwLock<HashMap<String, Vec<ObjectMeta>>>> = Lazy::new(Default::default);
|
||||
static SEGMENTS: Lazy<RwLock<HashMap<String, SegmentData>>> = Lazy::new(Default::default);
|
||||
@@ -58,7 +67,13 @@ pub async fn set(trace_id: &str, schema_key: &str, format: &str, files: Vec<File
|
||||
version: None,
|
||||
});
|
||||
if let Some(bin_data) = file.segment_ids {
|
||||
segment_data.insert(file.key, bin_data);
|
||||
segment_data.insert(
|
||||
file.key,
|
||||
SegmentInfo {
|
||||
segment_ids: bin_data,
|
||||
row_group_size: file.row_group_size,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
FILES.write().insert(key.clone(), values);
|
||||
@@ -91,7 +106,7 @@ pub fn clear(trace_id: &str) {
|
||||
drop(w);
|
||||
}
|
||||
|
||||
pub fn get_segment_ids(file_key: &str) -> Option<Arc<BitVec>> {
|
||||
pub fn get_segment_info(file_key: &str) -> Option<SegmentInfo> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<Arc<dyn std::any::Any + Send + Sync>> {
|
||||
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<BitVec>,
|
||||
row_group_size: Option<u32>,
|
||||
) -> Option<Arc<dyn std::any::Any + Send + Sync>> {
|
||||
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 {
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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::<u32>().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<dyn Directory> = 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<FileKey>, 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);
|
||||
|
||||
@@ -39,7 +39,8 @@ use crate::service::search::index::IndexCondition;
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum TantivyResult {
|
||||
RowIds(HashSet<u32>),
|
||||
RowIdsBitVec(usize, BitVec),
|
||||
/// (row_id_bitvec, matched_row_count, row_group_size_from_index_file)
|
||||
RowIdsBitVec(BitVec, usize, Option<u32>),
|
||||
Count(usize), // simple count optimization
|
||||
Histogram(Vec<u64>), // 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::<u32>()
|
||||
+ std::mem::size_of::<HashSet<u32>>()
|
||||
}
|
||||
Self::RowIdsBitVec(_, bitvec) => {
|
||||
Self::RowIdsBitVec(bitvec, ..) => {
|
||||
bitvec.capacity().div_ceil(8) + std::mem::size_of::<BitVec>()
|
||||
}
|
||||
Self::Count(_) => std::mem::size_of::<usize>(),
|
||||
@@ -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
|
||||
|
||||
@@ -29,9 +29,10 @@ pub static GLOBAL_CACHE: Lazy<Arc<TantivyResultCache>> =
|
||||
|
||||
#[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<u32>),
|
||||
/// (sparse_row_ids, matched_row_count, parquet_row_count, row_group_size_from_index_file)
|
||||
RowIdsRoaring(RoaringBitmap, usize, usize, Option<u32>),
|
||||
Count(usize), // simple count optimization
|
||||
Histogram(Vec<u64>), // simple histogram optimization
|
||||
TopN(Vec<(String, u64)>), // simple top n optimization
|
||||
@@ -41,15 +42,15 @@ pub enum CacheEntry {
|
||||
impl From<CacheEntry> 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<CacheEntry> 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::<BitVec>()
|
||||
}
|
||||
CacheEntry::RowIdsRoaring(_, roaring, _) => {
|
||||
CacheEntry::RowIdsRoaring(roaring, ..) => {
|
||||
roaring.serialized_size()
|
||||
+ std::mem::size_of::<RoaringBitmap>()
|
||||
+ std::mem::size_of::<usize>() * 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");
|
||||
|
||||
@@ -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();
|
||||
|
||||
|
||||
@@ -43,6 +43,10 @@ impl<W> PuffinBytesWriter<W> {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_property(&mut self, key: impl Into<String>, value: impl Into<String>) {
|
||||
self.properties.insert(key.into(), value.into());
|
||||
}
|
||||
|
||||
fn build_blob_metadata(
|
||||
&self,
|
||||
blob_type: BlobTypes,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -39,6 +39,7 @@ use crate::{
|
||||
pub struct PuffinDirReader {
|
||||
source: Arc<PuffinBytesReader>,
|
||||
blobs_metadata: Arc<HashMap<PathBuf, Arc<BlobMetadata>>>,
|
||||
properties: Arc<HashMap<String, String>>,
|
||||
}
|
||||
|
||||
impl PuffinDirReader {
|
||||
@@ -60,11 +61,18 @@ impl PuffinDirReader {
|
||||
}
|
||||
}
|
||||
|
||||
let properties = metadata.properties.into_iter().collect::<HashMap<_, _>>();
|
||||
|
||||
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");
|
||||
|
||||
@@ -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<MmapDirectory>,
|
||||
/// record all the files paths in the puffin file
|
||||
file_paths: Arc<RwLock<HashSet<PathBuf>>>,
|
||||
/// file-level puffin properties (written to PuffinMeta.properties)
|
||||
properties: Arc<RwLock<HashMap<String, String>>>,
|
||||
}
|
||||
|
||||
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<String>, value: impl Into<String>) {
|
||||
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<Vec<u8>> {
|
||||
let mut puffin_buf: Vec<u8> = 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");
|
||||
|
||||
Reference in New Issue
Block a user