starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-05 08:09:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf: reduce memory usage and time when generate tantivy index (#12292)

Browse Source 
- close https://github.com/openobserve/o2-enterprise/issues/1658

## Summary

Speed up tantivy index generation during compaction by indexing parquet
row_groups in parallel.

- New dispatcher in `src/service/tantivy/mod.rs` picks one of two
builders:
- **sequential** (`sequential.rs`) — legacy producer/consumer path, used
for non-parquet inputs or when `thread_num <= 1`.
- **parallel** (`parallel.rs`) — one `SingleSegmentIndexWriter` per
parquet row_group built under `spawn_blocking` +
`Semaphore(thread_num)`, then stitched via
`tantivy::indexer::merge_indices`.
- `create_tantivy_index` now takes raw parquet `Bytes` so each
per-row_group worker gets a cheap `Bytes::clone()` instead of re-reading
the file. Call sites in `src/job/files/parquet.rs` and
`src/service/compact/merge.rs` are updated.

## Config

- `ZO_COMPACT_TANTIVY_BUILDER_THREAD_NUM` (default `2`) — per-file
concurrent row_group workers. `<= 1` keeps the legacy single-threaded
path.

## Benchmark

Input: 21 row_groups, 2,672,970 rows, ~1.38 GB .ttv output.

| variant | cpu | total tantivy | peak mem | vs main (time) | vs main
(mem) |
|---|---:|---:|---:|---:|---:|
| **main (baseline)** | 2 | 108,673 ms | **10.62 GB** | — | — |
| PR, workers=2 (default) | 2 | 109,060 ms | 2.92 GB | ≈ same | **−72%**
|
| PR, workers=4 | 4 | 92,232 ms | 3.99 GB | **−14%** | **−62%** |
| PR, workers=10 | 10 | 76,167 ms | 7.75 GB | **−27%** | **−27%** |

- **Memory is the big win.** Even sequential mode (workers=1) cuts peak
from 10.6 GB → 3.0 GB (~3.6×).
- **At the default (workers=2), wall time matches main exactly while
peak memory drops 72%.** Safe upgrade.
- **Doubling workers buys time at the cost of memory** (per-row_group
segments buffer before merge). 4 workers is the best speed/memory
trade-off; ~10 is the cap (single-threaded `merge_indices` floor is ~60
s).

## Why it's correct

OpenObserve relies on `tantivy doc_id == global parquet row_index` to
map search hits back to row offsets. The parallel path preserves it
because:

1. Each per-row_group writer writes rows in stored order → local
`doc_id` == offset inside that row_group.
2. Worker results are sorted by `rg_idx` before `merge_indices`.
3. `merge_indices` stacks segments in input order (no `sort_by_field`
path in current tantivy).
4. Parquet stores row_groups in row order, so the stitched layout
matches the original global row order.

Two guards backstop this: the merged index must be a single segment, and
its `max_doc` must equal the sum of per-row_group row counts.

## Tests

- `test_row_group_parallel_preserves_row_order` — 4 row_groups, 3
workers (forces out-of-order completion), asserts every marker's `doc_id
== row_index`.
- `test_row_group_parallel_matches_sequential` — same data through both
paths yields identical `(marker → doc_id)` mappings.
- `test_row_group_parallel_single_row_group_fast_path` —
single-row_group input skips `merge_indices` and still preserves row
order.
This commit is contained in:
Huaijin
2026-05-27 11:37:57 +00:00
committed by GitHub
parent e79cb30b30
commit e2d0a460bc
6 changed files with 1158 additions and 603 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config/src/config.rs
+6
View File
@@ -1956,6 +1956,12 @@ pub struct Compact {
help = "Comma-separated list of hours (0-23) when retention can run. Empty means run at all hours. Example: 5,6,8"
)]
pub retention_allowed_hours: String,
#[env_config(
name = "ZO_COMPACT_TANTIVY_BUILDER_THREAD_NUM",
default = 2,
help = "Per-file concurrent row_group workers for tantivy index generation during compaction. less than or equal to 1 disables (single-threaded)"
)]
pub tantivy_builder_thread_num: usize,
}
#[derive(Serialize, EnvConfig, Default)]
src/job/files/parquet.rs
+4 -7
View File
@@ -22,6 +22,8 @@ use std::{
use arrow_schema::Schema;
use bytes::Bytes;
use chrono::{Duration, Utc};
#[cfg(feature = "enterprise")]
use config::utils::parquet::get_recordbatch_reader_from_bytes;
use config::{
FxIndexMap, cluster, get_config,
meta::{
@@ -32,9 +34,7 @@ use config::{
utils::{
async_file::{get_file_meta, get_file_size},
file::scan_files_with_channel,
parquet::{
get_recordbatch_reader_from_bytes, read_metadata_from_file, read_schema_from_file,
},
parquet::{read_metadata_from_file, read_schema_from_file},
schema_ext::SchemaExt,
},
};
@@ -928,9 +928,6 @@ async fn merge_files(
return Ok((account, new_file_key, new_file_meta, retain_file_list));
}
// generate tantivy inverted index and write to storage
let file_format = config::get_config().common.file_format;
let (_, reader) = get_recordbatch_reader_from_bytes(file_format, buf).await?;
let index_size = create_tantivy_index(
"INGESTER",
&org_id,
@@ -938,7 +935,7 @@ async fn merge_files(
&full_text_search_fields,
&index_fields,
latest_schema.clone(), // Use stream schema to include all configured fields
reader,
buf,
)
.await
.map_err(|e| anyhow::anyhow!("generate_tantivy_index_on_ingester error: {e}"))?;
src/service/compact/merge.rs
+2 -4
View File
@@ -29,7 +29,7 @@ use config::{
},
metrics,
utils::{
parquet::{get_recordbatch_reader_from_bytes, read_schema_from_bytes},
parquet::read_schema_from_bytes,
record_batch_ext::concat_batches,
schema_ext::SchemaExt,
time::{day_micros, hour_micros},
@@ -996,8 +996,6 @@ async fn generate_inverted_index(
latest_schema: Arc<Schema>,
buf: Bytes,
) -> Result<(), anyhow::Error> {
let file_format = get_config().common.file_format;
let (_, reader) = get_recordbatch_reader_from_bytes(file_format, buf).await?;
let index_size = create_tantivy_index(
"COMPACTOR",
org_id,
@@ -1005,7 +1003,7 @@ async fn generate_inverted_index(
full_text_search_fields,
index_fields,
latest_schema, // Use stream schema to include all configured fields
reader,
buf,
)
.await
.map_err(|e| {
src/service/tantivy/mod.rs
+172 -592
View File
@@ -13,97 +13,62 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
mod parallel;
pub mod puffin;
pub mod puffin_directory;
mod sequential;
use std::{collections::HashMap, sync::Arc};
use anyhow::Context;
use arrow::array::{
Array, BooleanArray, Float64Array, Int64Array, LargeStringArray, StringArray, StringViewArray,
UInt64Array,
Array, BooleanArray, Float64Array, Int64Array, LargeStringArray, RecordBatch, StringArray,
StringViewArray, UInt64Array,
};
use arrow_schema::{DataType, Schema};
use bytes::Bytes;
use config::{
INDEX_FIELD_NAME_FOR_ALL, PARQUET_MAX_ROW_GROUP_SIZE, TIMESTAMP_COL_NAME, get_config,
FileFormat, 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,
tantivy::tokenizer::{CollectType, O2_TOKENIZER, o2_tokenizer_build},
inverted_index::convert_parquet_file_name_to_tantivy_file, tantivy::tokenizer::O2_TOKENIZER,
},
};
use futures::TryStreamExt;
use hashbrown::HashSet;
use infra::storage;
use puffin_directory::writer::PuffinDirWriter;
use tokio::task::JoinHandle;
// macro to reduce duplication in array processing
macro_rules! process_string_array {
($data:expr, $array_type:ty, $docs:expr, $field:expr) => {
if let Some(array) = $data.as_any().downcast_ref::<$array_type>() {
for (i, doc) in $docs.iter_mut().enumerate() {
if array.is_null(i) {
doc.add_text($field, "");
} else {
doc.add_text($field, array.value(i));
}
tokio::task::coop::consume_budget().await;
}
continue;
}
};
}
// macro to reduce duplication in array processing
macro_rules! process_numeric_array {
($data:expr, $array_type:ty, $docs:expr, $field:expr) => {
if let Some(array) = $data.as_any().downcast_ref::<$array_type>() {
for (i, doc) in $docs.iter_mut().enumerate() {
if array.is_null(i) {
doc.add_text($field, "");
} else {
let text = array.value(i).to_string();
doc.add_text($field, &text);
}
tokio::task::coop::consume_budget().await;
}
continue;
}
};
}
use tantivy_utils::puffin_directory::PROP_ROW_GROUP_SIZE;
pub(crate) async fn create_tantivy_index(
caller: &str,
org_id: &str,
parquet_file_name: &str,
full_text_search_fields: &[String],
fts_fields: &[String],
index_fields: &[String],
schema: Arc<Schema>,
reader: RecordBatchStream,
buf: Bytes,
) -> Result<usize, anyhow::Error> {
let start = std::time::Instant::now();
let cfg = get_config();
let caller = format!("[{caller}:JOB]");
let file_format = FileFormat::from_extension(parquet_file_name).unwrap_or_default();
let thread_num = cfg.compact.tantivy_builder_thread_num;
let Some(index_schema) = build_tantivy_schema(fts_fields, index_fields, &schema) else {
return Ok(0);
};
let dir = PuffinDirWriter::new();
let index = generate_tantivy_index(
dir.clone(),
reader,
full_text_search_fields,
index_fields,
schema,
)
.await?;
let index = if thread_num > 1 && matches!(file_format, FileFormat::Parquet) {
parallel::build_index(dir.clone(), buf, index_schema, thread_num).await?
} else {
sequential::build_index(dir.clone(), file_format, buf, index_schema).await?
};
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(),
);
dir.set_property(PROP_ROW_GROUP_SIZE, PARQUET_MAX_ROW_GROUP_SIZE.to_string());
let puffin_bytes = dir.to_puffin_bytes()?;
let index_size = puffin_bytes.len();
@@ -113,7 +78,6 @@ pub(crate) async fn create_tantivy_index(
};
let buf = Bytes::from(puffin_bytes);
let cfg = get_config();
if cfg.cache_latest_files.enabled
&& cfg.cache_latest_files.cache_index
&& cfg.cache_latest_files.download_from_node
@@ -139,23 +103,39 @@ pub(crate) async fn create_tantivy_index(
Ok(index_size)
}
/// Create a tantivy index in the given directory for the record batch
pub(crate) async fn generate_tantivy_index<D: tantivy::Directory>(
tantivy_dir: D,
reader: RecordBatchStream,
full_text_search_fields: &[String],
/// Bundle of the tantivy schema and the per-row metadata both build paths need.
///
/// Built once by [`build_tantivy_schema`] and shared by the
/// sequential and parallel builders so they don't each rebuild the schema.
#[derive(Clone)]
pub(super) struct TantivyIndexSchema {
/// The built tantivy schema.
pub(super) schema: tantivy::schema::Schema,
/// Set of column names we must populate per row — union of fts fields and
/// secondary index fields, filtered by data type.
pub(super) fields: HashSet<String>,
/// Field handle for `INDEX_FIELD_NAME_FOR_ALL`, present when there is at
/// least one full-text-search field.
pub(super) fts_field: Option<tantivy::schema::Field>,
}
/// Builds the tantivy `Schema` shared by both index-build paths.
///
/// Returns `None` when there are no fields to index (configured fields don't
/// intersect with the stream schema).
fn build_tantivy_schema(
fts_fields: &[String],
index_fields: &[String],
schema: Arc<Schema>,
) -> Result<Option<tantivy::Index>, anyhow::Error> {
arrow_schema: &Schema,
) -> Option<TantivyIndexSchema> {
let mut tantivy_schema_builder = tantivy::schema::SchemaBuilder::new();
let schema_fields = schema
let schema_fields = arrow_schema
.fields()
.iter()
.map(|f| (f.name(), f))
.collect::<HashMap<_, _>>();
// filter out fields that are not in schema & not of type Utf8
let fts_fields = full_text_search_fields
let fts_fields_filtered = fts_fields
.iter()
.filter(|f| {
schema_fields
@@ -165,22 +145,21 @@ pub(crate) async fn generate_tantivy_index<D: tantivy::Directory>(
})
.map(String::from)
.collect::<HashSet<_>>();
let index_fields = index_fields
let index_fields_filtered = index_fields
.iter()
.filter(|f| schema_fields.contains_key(f))
.map(String::from)
.collect::<HashSet<_>>();
let tantivy_fields = fts_fields
.union(&index_fields)
let tantivy_fields = fts_fields_filtered
.union(&index_fields_filtered)
.cloned()
.collect::<HashSet<_>>();
if tantivy_fields.is_empty() {
return Ok(None);
return None;
}
// add fields to tantivy schema
if !full_text_search_fields.is_empty() {
if !fts_fields_filtered.is_empty() {
let fts_opts = tantivy::schema::TextOptions::default().set_indexing_options(
tantivy::schema::TextFieldIndexing::default()
.set_index_option(tantivy::schema::IndexRecordOption::Basic)
@@ -198,142 +177,110 @@ pub(crate) async fn generate_tantivy_index<D: tantivy::Directory>(
.set_fieldnorms(false),
)
.set_fast(None);
for field in index_fields.iter() {
for field in index_fields_filtered.iter() {
if field == TIMESTAMP_COL_NAME {
continue;
}
tantivy_schema_builder.add_text_field(field, index_opts.clone());
}
// add _timestamp field to tantivy schema
tantivy_schema_builder.add_i64_field(TIMESTAMP_COL_NAME, tantivy::schema::FAST);
let tantivy_schema = tantivy_schema_builder.build();
let fts_field = tantivy_schema.get_field(INDEX_FIELD_NAME_FOR_ALL).ok();
let tokenizer_manager = tantivy::tokenizer::TokenizerManager::default();
tokenizer_manager.register(O2_TOKENIZER, o2_tokenizer_build(CollectType::Ingest));
let mut index_writer = tantivy::IndexBuilder::new()
.schema(tantivy_schema.clone())
.tokenizers(tokenizer_manager)
.single_segment_index_writer(tantivy_dir, 50_000_000)
.context("failed to create index builder")?;
Some(TantivyIndexSchema {
schema: tantivy_schema,
fields: tantivy_fields,
fts_field,
})
}
// docs per row to be added in the tantivy index
log::debug!("start write documents to tantivy index");
let (tx, mut rx) = tokio::sync::mpsc::channel::<Vec<tantivy::TantivyDocument>>(2);
let task: JoinHandle<Result<usize, anyhow::Error>> = tokio::task::spawn(async move {
let mut total_num_rows = 0;
let mut reader = reader;
loop {
let batch = reader.try_next().await?;
let Some(inverted_idx_batch) = batch else {
break;
};
let num_rows = inverted_idx_batch.num_rows();
if num_rows == 0 {
continue;
}
// update total_num_rows
total_num_rows += num_rows;
// process full text search fields
let mut docs = vec![tantivy::doc!(); num_rows];
for column_name in tantivy_fields.iter() {
// get field
let field = match tantivy_schema.get_field(column_name) {
Ok(f) => f,
Err(_) => fts_field.unwrap(),
};
// get column data and convert to strings for indexing
if let Some(data) = inverted_idx_batch.column_by_name(column_name) {
// handle string types directly
process_string_array!(data, StringViewArray, docs, field);
process_string_array!(data, StringArray, docs, field);
process_string_array!(data, LargeStringArray, docs, field);
// handle numeric and boolean types with to_string conversion
process_numeric_array!(data, Int64Array, docs, field);
process_numeric_array!(data, UInt64Array, docs, field);
process_numeric_array!(data, Float64Array, docs, field);
process_numeric_array!(data, BooleanArray, docs, field);
// unsupported type, add empty string
for doc in docs.iter_mut() {
doc.add_text(field, "");
tokio::task::coop::consume_budget().await;
}
macro_rules! process_string_array_sync {
($data:expr, $array_type:ty, $docs:expr, $field:expr) => {
if let Some(array) = $data.as_any().downcast_ref::<$array_type>() {
for (i, doc) in $docs.iter_mut().enumerate() {
if array.is_null(i) {
doc.add_text($field, "");
} else {
// column not found, add empty string
for doc in docs.iter_mut() {
doc.add_text(field, "");
tokio::task::coop::consume_budget().await;
}
doc.add_text($field, array.value(i));
}
}
// process _timestamp field
let column_data = match inverted_idx_batch.column_by_name(TIMESTAMP_COL_NAME) {
Some(column_data) => match column_data.as_any().downcast_ref::<Int64Array>() {
Some(column_data) => column_data,
None => {
// generate empty array to ensure the tantivy and parquet have same rows
&Int64Array::from(vec![0; num_rows])
}
},
None => {
// generate empty array to ensure the tantivy and parquet have same rows
&Int64Array::from(vec![0; num_rows])
}
};
let ts_field = tantivy_schema.get_field(TIMESTAMP_COL_NAME).unwrap(); // unwrap directly since added above
const YIELD_THRESHOLD: usize = 100;
let mut batch_size = 0;
for (i, doc) in docs.iter_mut().enumerate() {
doc.add_i64(ts_field, column_data.value(i));
batch_size += 1;
if batch_size >= YIELD_THRESHOLD {
tokio::task::coop::consume_budget().await;
batch_size = 0;
}
}
tx.send(docs).await?;
continue;
}
Ok(total_num_rows)
});
};
}
while let Some(docs) = rx.recv().await {
for doc in docs {
if let Err(e) = index_writer.add_document(doc) {
log::error!("generate_tantivy_index: Failed to add document to index: {e}");
return Err(anyhow::anyhow!("Failed to add document to index: {}", e));
macro_rules! process_numeric_array_sync {
($data:expr, $array_type:ty, $docs:expr, $field:expr) => {
if let Some(array) = $data.as_any().downcast_ref::<$array_type>() {
for (i, doc) in $docs.iter_mut().enumerate() {
if array.is_null(i) {
doc.add_text($field, "");
} else {
let text = array.value(i).to_string();
doc.add_text($field, &text);
}
}
continue;
}
};
}
// Callers run this inside `spawn_blocking` so the per-row CPU work stays off the async runtime.
pub(super) fn convert_batch_to_docs_sync(
batch: &RecordBatch,
index_schema: &TantivyIndexSchema,
) -> Vec<tantivy::TantivyDocument> {
let tantivy_schema = &index_schema.schema;
let tantivy_fields = &index_schema.fields;
let fts_field = index_schema.fts_field;
let num_rows = batch.num_rows();
let mut docs = vec![tantivy::doc!(); num_rows];
for column_name in tantivy_fields.iter() {
let field = match tantivy_schema.get_field(column_name) {
Ok(f) => f,
Err(_) => fts_field.expect("fts field must exist when index field is missing"),
};
if let Some(data) = batch.column_by_name(column_name) {
process_string_array_sync!(data, StringViewArray, docs, field);
process_string_array_sync!(data, StringArray, docs, field);
process_string_array_sync!(data, LargeStringArray, docs, field);
process_numeric_array_sync!(data, Int64Array, docs, field);
process_numeric_array_sync!(data, UInt64Array, docs, field);
process_numeric_array_sync!(data, Float64Array, docs, field);
process_numeric_array_sync!(data, BooleanArray, docs, field);
for doc in docs.iter_mut() {
doc.add_text(field, "");
}
} else {
for doc in docs.iter_mut() {
doc.add_text(field, "");
}
tokio::task::coop::consume_budget().await;
}
}
let total_num_rows = task.await??;
// Create index even with 0 rows since we have valid configured fields in stream schema
// (empty index acts as a marker to prevent expensive DataFusion scans)
log::debug!(
"write documents to tantivy index success (rows: {}, empty_index: {})",
total_num_rows,
total_num_rows == 0
);
let index = tokio::task::spawn_blocking(move || {
index_writer.finalize().map_err(|e| {
log::error!("generate_tantivy_index: Failed to finalize the index writer: {e}");
anyhow::anyhow!("Failed to finalize the index writer: {}", e)
})
})
.await??;
// _timestamp field, should we report error when _timestamp not found
let ts_data: &Int64Array = match batch.column_by_name(TIMESTAMP_COL_NAME) {
Some(column_data) => match column_data.as_any().downcast_ref::<Int64Array>() {
Some(ts) => ts,
None => &Int64Array::from(vec![0; num_rows]),
},
None => &Int64Array::from(vec![0; num_rows]),
};
let ts_field = tantivy_schema.get_field(TIMESTAMP_COL_NAME).unwrap();
for (i, doc) in docs.iter_mut().enumerate() {
doc.add_i64(ts_field, ts_data.value(i));
}
Ok(Some(index))
docs
}
#[cfg(test)]
mod tests {
pub(super) mod tests {
use std::sync::Arc;
use arrow::{
@@ -341,13 +288,13 @@ mod tests {
datatypes::{DataType, Field, Schema},
record_batch::RecordBatch,
};
use config::{FileFormat, INDEX_FIELD_NAME_FOR_ALL, TIMESTAMP_COL_NAME};
use tantivy::directory::RamDirectory;
use config::TIMESTAMP_COL_NAME;
use super::*;
// Helper function to create test record batches
fn create_test_batch(
/// Helper function to create test record batches.
/// Shared with [`super::sequential`] and [`super::parallel`] test modules.
pub(in crate::service::tantivy) fn create_test_batch(
num_rows: usize,
include_timestamp: bool,
include_fts_field: bool,
@@ -356,14 +303,12 @@ mod tests {
let mut fields = Vec::new();
let mut columns: Vec<Arc<dyn arrow::array::Array>> = Vec::new();
// Add timestamp field if requested
if include_timestamp {
fields.push(Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false));
let timestamps: Vec<i64> = (0..num_rows).map(|i| 1000 + i as i64).collect();
columns.push(Arc::new(Int64Array::from(timestamps)));
}
// Add full-text search field if requested
if include_fts_field {
fields.push(Field::new("content", DataType::Utf8, false));
let content: Vec<String> = (0..num_rows)
@@ -372,7 +317,6 @@ mod tests {
columns.push(Arc::new(StringArray::from(content)));
}
// Add index field if requested
if include_index_field {
fields.push(Field::new("status", DataType::Utf8, false));
let status: Vec<String> = (0..num_rows)
@@ -391,13 +335,14 @@ mod tests {
RecordBatch::try_new(schema, columns).unwrap()
}
// Helper function to create a mock RecordBatchStream
async fn create_test_stream(batches: Vec<RecordBatch>) -> RecordBatchStream {
// Create a simple parquet file from the batches
/// Helper: serializes `batches` to an in-memory parquet file and returns
/// the raw bytes. Shared by both the stream-based and the new
/// row_group-direct parallel tests.
pub(in crate::service::tantivy) async fn create_test_parquet_bytes(
batches: Vec<RecordBatch>,
) -> bytes::Bytes {
let schema = batches[0].schema();
let mut buffer = Vec::new();
// Use the parquet writer utilities from config
let file_meta = config::meta::stream::FileMeta {
min_ts: 1000,
max_ts: 2000,
@@ -405,7 +350,6 @@ mod tests {
original_size: 1000,
..Default::default()
};
let mut writer = config::utils::parquet::new_parquet_writer(
&mut buffer,
&schema,
@@ -414,321 +358,46 @@ mod tests {
false,
None,
);
for batch in batches {
writer.write(&batch).await.unwrap();
}
writer.close().await.unwrap();
bytes::Bytes::from(buffer)
}
// Create stream from the buffer using the new API
let bytes = bytes::Bytes::from(buffer);
let (_schema, stream) =
config::utils::parquet::get_recordbatch_reader_from_bytes(FileFormat::Parquet, bytes)
.await
.unwrap();
stream
/// Test helper: build the per-index schema struct or panic on `None`.
pub(in crate::service::tantivy) fn make_index_schema(
fts: &[String],
index: &[String],
arrow_schema: &Schema,
) -> TantivyIndexSchema {
build_tantivy_schema(fts, index, arrow_schema).expect("schema helper returned None")
}
#[tokio::test]
async fn test_generate_tantivy_index_empty_data() {
let dir = RamDirectory::create();
let empty_batch = create_test_batch(0, true, true, true);
let stream = create_test_stream(vec![empty_batch.clone()]).await;
let result = generate_tantivy_index(
dir,
stream,
&["content".to_string()],
&["status".to_string()],
empty_batch.schema(),
)
.await;
assert!(result.is_ok());
// Should create empty index when fields are configured (even with 0 rows)
assert!(result.unwrap().is_some());
}
#[tokio::test]
async fn test_generate_tantivy_index_no_fields() {
let dir = RamDirectory::create();
async fn test_build_tantivy_schema_no_fields() {
let batch = create_test_batch(10, true, true, true);
let stream = create_test_stream(vec![batch.clone()]).await;
let result = generate_tantivy_index(
dir,
stream,
&[], // No full-text search fields
&[], // No index fields
batch.schema(),
)
.await;
assert!(result.is_ok());
assert!(result.unwrap().is_none()); // Should return None when no fields to index
// No fields to index → helper returns None and the orchestrator short-circuits.
assert!(build_tantivy_schema(&[], &[], &batch.schema()).is_none());
}
#[tokio::test]
async fn test_generate_tantivy_index_with_fts_fields() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, true, true, false);
let stream = create_test_stream(vec![batch.clone()]).await;
let result =
generate_tantivy_index(dir, stream, &["content".to_string()], &[], batch.schema())
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
// Verify the index has the expected schema
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_with_index_fields() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, true, false, true);
let stream = create_test_stream(vec![batch.clone()]).await;
let result =
generate_tantivy_index(dir, stream, &[], &["status".to_string()], batch.schema()).await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
// Verify the index has the expected schema
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field("status").is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_with_both_field_types() {
let dir = RamDirectory::create();
async fn test_build_tantivy_schema_missing_fields_in_schema() {
let batch = create_test_batch(10, true, true, true);
let stream = create_test_stream(vec![batch.clone()]).await;
let result = generate_tantivy_index(
dir,
stream,
&["content".to_string()],
&["status".to_string()],
batch.schema(),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
// Verify the index has the expected schema
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field("status").is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_missing_fields_in_schema() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, true, true, true);
let stream = create_test_stream(vec![batch.clone()]).await;
// Request fields that don't exist in the schema at all
let result = generate_tantivy_index(
dir,
stream,
// Configured fields don't exist in the stream schema → helper returns None
// and the orchestrator returns 0 without creating an index.
let result = build_tantivy_schema(
&["nonexistent_field".to_string()],
&["another_nonexistent_field".to_string()],
batch.schema(),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
// Should NOT create index when configured fields don't exist in stream schema
// (this indicates a configuration error, not just missing data)
assert!(index.is_none());
}
#[tokio::test]
async fn test_generate_tantivy_index_mixed_existing_and_missing_fields() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, true, true, true);
let stream = create_test_stream(vec![batch.clone()]).await;
// Mix of existing and non-existing fields
let result = generate_tantivy_index(
dir,
stream,
&["content".to_string(), "nonexistent_field".to_string()],
&[
"status".to_string(),
"another_nonexistent_field".to_string(),
],
batch.schema(),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
// Verify the index has only the existing fields
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field("status").is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_multiple_batches() {
let dir = RamDirectory::create();
let batch1 = create_test_batch(5, true, true, true);
let batch2 = create_test_batch(5, true, true, true);
let stream = create_test_stream(vec![batch1.clone(), batch2.clone()]).await;
let result = generate_tantivy_index(
dir,
stream,
&["content".to_string()],
&["status".to_string()],
batch1.schema(),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
// Verify the index was created successfully
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field("status").is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_without_timestamp() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, false, true, true);
let stream = create_test_stream(vec![batch.clone()]).await;
let result = generate_tantivy_index(
dir,
stream,
&["content".to_string()],
&["status".to_string()],
batch.schema(),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
// Verify the index has the expected schema (including timestamp field added by the
// function)
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field("status").is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_non_string_fields() {
// Create a batch with non-string fields in the schema
let fields = vec![
Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false),
Field::new("content", DataType::Utf8, false),
Field::new("number_field", DataType::Int32, false), // Non-string field
];
let schema = Arc::new(Schema::new(fields));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int64Array::from(vec![1000, 1001, 1002])),
Arc::new(StringArray::from(vec!["content1", "content2", "content3"])),
Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3])),
],
)
.unwrap();
let dir = RamDirectory::create();
let stream = create_test_stream(vec![batch.clone()]).await;
// Try to index the non-string field
let result = generate_tantivy_index(
dir,
stream,
&["content".to_string()],
&["number_field".to_string()], // This field is not Utf8
batch.schema(),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
// Verify that both fields are indexed (index fields are not filtered by data type)
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field("number_field").is_ok()); // Non-string fields are still indexed
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_timestamp_field_handling() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, true, true, true);
let stream = create_test_stream(vec![batch.clone()]).await;
// Try to index the timestamp field as a regular index field
let result = generate_tantivy_index(
dir,
stream,
&["content".to_string()],
&[TIMESTAMP_COL_NAME.to_string()], // This should be ignored
batch.schema(),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
// Verify that timestamp field is handled specially
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
// Verify that the timestamp field is indexed as Int64, not as text
let ts_field = schema.get_field(TIMESTAMP_COL_NAME).unwrap();
assert!(matches!(
schema.get_field_entry(ts_field).field_type(),
tantivy::schema::FieldType::I64(_)
));
&batch.schema(),
);
assert!(result.is_none());
}
#[tokio::test]
async fn test_create_tantivy_index_with_empty_data() {
let empty_batch = create_test_batch(0, true, true, true);
let stream = create_test_stream(vec![empty_batch.clone()]).await;
let buf = create_test_parquet_bytes(vec![empty_batch.clone()]).await;
// Test that create_tantivy_index returns 0 for empty data
let result = create_tantivy_index(
@@ -738,7 +407,7 @@ mod tests {
&["content".to_string()],
&["status".to_string()],
empty_batch.schema(),
stream,
buf,
)
.await;
@@ -749,7 +418,7 @@ mod tests {
#[tokio::test]
async fn test_create_tantivy_index_with_no_fields() {
let batch = create_test_batch(10, true, true, true);
let stream = create_test_stream(vec![batch.clone()]).await;
let buf = create_test_parquet_bytes(vec![batch.clone()]).await;
// Test that create_tantivy_index returns 0 when no fields to index
let result = create_tantivy_index(
@@ -759,7 +428,7 @@ mod tests {
&[], // No FTS fields
&[], // No index fields
batch.schema(),
stream,
buf,
)
.await;
@@ -770,7 +439,7 @@ mod tests {
#[tokio::test]
async fn test_create_tantivy_index_with_invalid_parquet_filename() {
let batch = create_test_batch(10, true, true, true);
let stream = create_test_stream(vec![batch.clone()]).await;
let buf = create_test_parquet_bytes(vec![batch.clone()]).await;
// Test with an invalid parquet filename that won't convert to tantivy filename
let result = create_tantivy_index(
@@ -780,7 +449,7 @@ mod tests {
&["content".to_string()],
&["status".to_string()],
batch.schema(),
stream,
buf,
)
.await;
@@ -788,97 +457,8 @@ mod tests {
assert_eq!(result.unwrap(), 0); // Should return 0 when filename conversion fails
}
#[tokio::test]
async fn test_generate_tantivy_index_stream_schema_vs_parquet_schema() {
// This test validates the critical fix where we pass stream schema (with all
// configured fields) instead of parquet schema (only fields in actual data).
//
// Scenario: Stream settings configure `continent`, `name`, `flag_url` as secondary
// index fields, but the parquet data only contains `name` and `flag_url`.
// The missing field `continent` should still be included in the .ttv schema.
let dir = RamDirectory::create();
// Create parquet data with only `name` and `flag_url` (no `continent`)
let parquet_fields = vec![
Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
Field::new("flag_url", DataType::Utf8, false),
];
let parquet_schema = Arc::new(Schema::new(parquet_fields));
let timestamps: Vec<i64> = (0..10).map(|i| 1000 + i as i64).collect();
let names: Vec<String> = (0..10).map(|i| format!("Name {i}")).collect();
let flag_urls: Vec<String> = (0..10)
.map(|i| format!("https://example.com/{i}"))
.collect();
let parquet_batch = RecordBatch::try_new(
parquet_schema.clone(),
vec![
Arc::new(Int64Array::from(timestamps)),
Arc::new(StringArray::from(names)),
Arc::new(StringArray::from(flag_urls)),
],
)
.unwrap();
// Create stream schema with all configured fields including `continent`
let stream_fields = vec![
Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false),
Field::new("continent", DataType::Utf8, false), // Configured but not in data
Field::new("name", DataType::Utf8, false),
Field::new("flag_url", DataType::Utf8, false),
];
let stream_schema = Arc::new(Schema::new(stream_fields));
let stream = create_test_stream(vec![parquet_batch]).await;
// Generate index with stream schema (all configured fields)
let result = generate_tantivy_index(
dir,
stream,
&[], // No FTS fields
&[
"continent".to_string(),
"name".to_string(),
"flag_url".to_string(),
],
stream_schema.clone(), // Pass stream schema, not parquet schema
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
let index = index.unwrap();
let tantivy_schema = index.schema();
// Verify that ALL configured fields are in the tantivy schema,
// including `continent` which is missing from the parquet data
assert!(
tantivy_schema.get_field("continent").is_ok(),
"continent field should be in tantivy schema even though it's not in parquet data"
);
assert!(
tantivy_schema.get_field("name").is_ok(),
"name field should be in tantivy schema"
);
assert!(
tantivy_schema.get_field("flag_url").is_ok(),
"flag_url field should be in tantivy schema"
);
assert!(tantivy_schema.get_field(TIMESTAMP_COL_NAME).is_ok());
// Verify we have 10 documents
let reader = index.reader().unwrap();
let searcher = reader.searcher();
assert_eq!(searcher.num_docs(), 10);
}
// Note: Full testing of create_tantivy_index with storage operations would require
// mocking the storage layer, which is complex for unit tests. The main logic is
// tested in generate_tantivy_index. Integration tests would be more appropriate
// tested in sequential::build_index. Integration tests would be more appropriate
// for testing the full create_tantivy_index function with actual storage operations.
}
src/service/tantivy/parallel.rs
+491
View File
@@ -0,0 +1,491 @@
// Copyright 2026 OpenObserve Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
//! Row-group-direct parallel tantivy index builder.
//!
//! Activated when `compact.tantivy_parallel_build_workers > 1` and the input
//! file format is parquet. Each parquet row_group becomes one independent
//! chunk segment built off the runtime in `tokio::task::spawn_blocking`; the
//! resulting per-row_group segments are stitched together by
//! `tantivy::indexer::merge_indices` in row_group order.
//!
//! Correctness of `doc_id == global_row_index` is preserved because:
//! 1. each per-row_group `SingleSegmentIndexWriter` writes rows in stored order, so its local
//! doc_ids equal the row's offset inside that row_group;
//! 2. `IndexMerger` stacks input segments in input order (current tantivy has no `sort_by_field`
//! path), so each segment's doc_ids land at `i * PARQUET_MAX_ROW_GROUP_SIZE`. For example,
//! pretending the fixed row_group size were 100, segment 0 would occupy doc_ids `[0, 100)`,
//! segment 1 `[100, 200)`, segment 2 `[200, 300)`, etc. (the final row_group may be shorter,
//! but it is always the last segment);
//! 3. parquet stores row_groups in row order, so the global row index of row_group `i`'s row `k`
//! is `i * PARQUET_MAX_ROW_GROUP_SIZE + k` (with size 100, row 5 of row_group 1 maps to global
//! row index `100 + 5 = 105`).
use std::sync::Arc;
use anyhow::{Context, Error};
use bytes::Bytes;
use config::utils::tantivy::tokenizer::{CollectType, O2_TOKENIZER, o2_tokenizer_build};
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use tantivy::directory::MmapDirectory;
use tokio::task::JoinHandle;
use super::{TantivyIndexSchema, convert_batch_to_docs_sync};
type TaskOutput = (usize, tantivy::Index, usize);
/// Row-group-direct parallel tantivy index builder.
///
/// ```text
/// parquet (Bytes)
/// │
/// ├─ open ParquetRecordBatchReaderBuilder ── read metadata ── num_row_groups
/// │
/// ▼ spawn_blocking pool (concurrency = workers)
/// ┌──────────────────────┬──────────────────────┬─────────────────────┐
/// │ rg_0 │ rg_1 │ ... │
/// │ builder │ builder │ │
/// │ .with_row_groups │ .with_row_groups │ │
/// │ ([0]).build() │ ([1]).build() │ │
/// │ → SingleSegmentIdx │ → SingleSegmentIdx │ │
/// │ → finalize → Index │ → finalize → Index │ │
/// └──────────────────────┴──────────────────────┴─────────────────────┘
/// │
/// ▼ sort by rg_idx ASC
/// merge_indices([Index_0, Index_1, ...], PuffinDirWriter)
/// │
/// ▼
/// single segment, doc_id == global_row_index
/// ```
///
/// `Bytes` is reference-counted, so handing one clone per worker is cheap.
pub(super) async fn build_index<D: tantivy::Directory + Send + Sync + 'static>(
tantivy_dir: D,
buf: Bytes,
index_schema: TantivyIndexSchema,
thread_num: usize,
) -> Result<Option<tantivy::Index>, Error> {
// Parse parquet metadata once on a blocking thread (footer decode is sync I/O).
let num_row_groups = {
let buf = buf.clone();
tokio::task::spawn_blocking(move || -> Result<usize, Error> {
let builder = ParquetRecordBatchReaderBuilder::try_new(buf)
.context("failed to open parquet for metadata")?;
Ok(builder.metadata().num_row_groups())
})
.await??
};
if num_row_groups == 0 {
return Ok(None);
}
// single row_group fast path
if num_row_groups == 1 {
let start = std::time::Instant::now();
let (_, index, num_rows) = tokio::task::spawn_blocking(move || {
build_row_group_segment(0, buf, index_schema, tantivy_dir)
})
.await??;
log::info!(
"parallel::build_index: single row_group fast path, rows={num_rows}, elapsed={} ms",
start.elapsed().as_millis()
);
return Ok(Some(index));
}
let thread_num = thread_num.max(1).min(num_row_groups);
let semaphore = Arc::new(tokio::sync::Semaphore::new(thread_num));
let start = std::time::Instant::now();
let mut tasks: Vec<JoinHandle<Result<TaskOutput, Error>>> = Vec::with_capacity(num_row_groups);
for rg_idx in 0..num_row_groups {
let permit = semaphore.clone().acquire_owned().await?;
let buf_c = buf.clone();
let index_schema_c = index_schema.clone();
tasks.push(tokio::task::spawn_blocking(move || {
let _permit = permit;
let dir = MmapDirectory::create_from_tempdir()?;
build_row_group_segment(rg_idx, buf_c, index_schema_c, dir)
}));
}
let mut indexed: Vec<TaskOutput> = Vec::with_capacity(tasks.len());
for t in tasks {
indexed.push(t.await??);
}
// Sort by row_group index to preserve global row order in the merged index.
indexed.sort_by_key(|(rg, ..)| *rg);
let total_num_rows: usize = indexed.iter().map(|(_, _, n)| n).sum();
let indices: Vec<tantivy::Index> = indexed.into_iter().map(|(_, i, _)| i).collect();
log::info!(
"parallel::build_index: built {num_row_groups} row_groups, total_rows={total_num_rows}, thread_num={thread_num}, elapsed={} ms",
start.elapsed().as_millis()
);
// merge all row_group indices into a single index.
let start = std::time::Instant::now();
let merged = tokio::task::spawn_blocking(move || {
tantivy::indexer::merge_indices(&indices, tantivy_dir)
.map_err(|e| anyhow::anyhow!("merge_indices failed: {e}"))
})
.await??;
log::info!(
"parallel::build_index: merged row_group indices into single tantivy index, elapsed={} ms",
start.elapsed().as_millis()
);
let merged_segs = merged.searchable_segments()?;
if merged_segs.len() != 1 {
return Err(anyhow::anyhow!(
"merged tantivy index produced {} segments (expected 1)",
merged_segs.len()
));
}
let merged_rows = merged_segs[0].meta().max_doc() as usize;
if merged_rows != total_num_rows {
return Err(anyhow::anyhow!(
"merged tantivy index row count mismatch: tantivy={merged_rows}, parquet={total_num_rows}",
));
}
Ok(Some(merged))
}
/// Synchronously builds the tantivy segment for exactly one parquet row_group
/// into the caller-supplied directory.
/// Return: (row_group_index, per-row_group Index, row_count)
fn build_row_group_segment<D: tantivy::Directory>(
rg_idx: usize,
buf: Bytes,
index_schema: TantivyIndexSchema,
dir: D,
) -> Result<(usize, tantivy::Index, usize), Error> {
let builder = ParquetRecordBatchReaderBuilder::try_new(buf)?;
let reader = builder.with_row_groups(vec![rg_idx]).build()?;
let tokenizer_manager = tantivy::tokenizer::TokenizerManager::default();
tokenizer_manager.register(O2_TOKENIZER, o2_tokenizer_build(CollectType::Ingest));
let mut writer = tantivy::IndexBuilder::new()
.schema(index_schema.schema.clone())
.tokenizers(tokenizer_manager)
.single_segment_index_writer::<tantivy::TantivyDocument>(dir, 50_000_000)?;
let mut rg_rows: usize = 0;
for batch_res in reader {
let batch = batch_res?;
if batch.num_rows() == 0 {
continue;
}
rg_rows += batch.num_rows();
let docs = convert_batch_to_docs_sync(&batch, &index_schema);
for doc in docs {
writer.add_document(doc)?;
}
}
let index = writer
.finalize()
.map_err(|e| anyhow::anyhow!("worker {rg_idx}: finalize failed: {e}"))?;
Ok((rg_idx, index, rg_rows))
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow::{
array::{Int64Array, StringArray},
datatypes::{DataType, Field, Schema},
record_batch::RecordBatch,
};
use config::{FileFormat, TIMESTAMP_COL_NAME};
use parquet::{arrow::AsyncArrowWriter, file::properties::WriterProperties};
use tantivy::directory::RamDirectory;
use super::*;
use crate::service::tantivy::{
sequential,
tests::{create_test_parquet_bytes, make_index_schema},
};
/// Build a record batch where every row has a globally-unique marker token
/// in the `marker` column. The marker is what we search to recover the
/// doc_id, so we can assert `doc_id == global_row_index`.
fn create_marker_batch(start: usize, num_rows: usize) -> RecordBatch {
let fields = vec![
Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false),
Field::new("marker", DataType::Utf8, false),
];
let timestamps: Vec<i64> = (0..num_rows).map(|i| (start + i) as i64).collect();
let markers: Vec<String> = (0..num_rows)
.map(|i| format!("row{:09}", start + i))
.collect();
let schema = Arc::new(Schema::new(fields));
RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(timestamps)),
Arc::new(StringArray::from(markers)),
],
)
.unwrap()
}
/// Writes batches to an in-memory parquet file with a caller-specified
/// fixed row_group row count. We deliberately bypass
/// `config::utils::parquet::new_parquet_writer` because it hard-codes the
/// row_group cap to `PARQUET_MAX_ROW_GROUP_SIZE` (128k). The parallel
/// builder's `doc_id == i * row_group_size + k` invariant only holds when
/// every non-tail row_group has the same row count, so the test needs to
/// drive that exact splitting boundary at a size small enough to keep the
/// test fast.
async fn create_test_parquet_bytes_fixed_rg(
batches: Vec<RecordBatch>,
row_group_row_count: usize,
) -> bytes::Bytes {
let schema = batches[0].schema();
let mut buffer = Vec::new();
let writer_props = WriterProperties::builder()
.set_max_row_group_row_count(Some(row_group_row_count))
.build();
let mut writer =
AsyncArrowWriter::try_new(&mut buffer, schema, Some(writer_props)).unwrap();
for batch in batches {
writer.write(&batch).await.unwrap();
}
writer.close().await.unwrap();
bytes::Bytes::from(buffer)
}
/// The core invariant test: with multiple row_groups parsed by
/// per-row_group workers (potentially out of order), every marker's
/// `doc_id` must equal its original global row index after merge.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_row_group_parallel_preserves_row_order() {
// Fixed row_group size of 500 + 2137 total rows ⇒ 5 row_groups of
// exactly 500 rows and a 137-row partial tail. This matches the
// production layout where every non-tail row_group has the same size.
let row_group_size = 500usize;
let total_rows = 2137usize;
let buf = create_test_parquet_bytes_fixed_rg(
vec![create_marker_batch(0, total_rows)],
row_group_size,
)
.await;
// Verify the parquet split into the expected fixed-size row_groups.
let expected_num_rg = total_rows.div_ceil(row_group_size);
let metadata_num_rg = {
let buf = buf.clone();
tokio::task::spawn_blocking(move || {
ParquetRecordBatchReaderBuilder::try_new(buf)
.unwrap()
.metadata()
.num_row_groups()
})
.await
.unwrap()
};
assert_eq!(
metadata_num_rg, expected_num_rg,
"test setup must yield {expected_num_rg} fixed-size row_groups"
);
let stream_schema = Arc::new(Schema::new(vec![
Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false),
Field::new("marker", DataType::Utf8, false),
]));
let index = build_index(
RamDirectory::create(),
buf,
make_index_schema(&[], &["marker".to_string()], &stream_schema),
// workers
3, // force out-of-order completion
)
.await
.expect("parallel build must succeed")
.expect("parallel build must return Some(index)");
let segs = index.searchable_segments().unwrap();
assert_eq!(segs.len(), 1);
assert_eq!(segs[0].meta().max_doc() as usize, total_rows);
let reader = index.reader().unwrap();
let searcher = reader.searcher();
let marker_field = index.schema().get_field("marker").unwrap();
for row in 0..total_rows {
let m = format!("row{:09}", row);
let q = tantivy::query::TermQuery::new(
tantivy::Term::from_field_text(marker_field, &m),
tantivy::schema::IndexRecordOption::Basic,
);
let hits = searcher
.search(&q, &tantivy::collector::DocSetCollector)
.unwrap();
assert_eq!(
hits.len(),
1,
"marker for row {row} must match exactly 1 doc"
);
let addr = *hits.iter().next().unwrap();
assert_eq!(addr.segment_ord, 0);
assert_eq!(
addr.doc_id as usize, row,
"INVARIANT VIOLATION: doc_id != row_index at row {row}"
);
}
}
/// Single row_group hits the fast path that bypasses `merge_indices` and
/// writes the segment straight into the caller-supplied directory.
/// Validates that the resulting index lands in `tantivy_dir` and that
/// `doc_id == row_index` is preserved.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_row_group_parallel_single_row_group_fast_path() {
// 321 rows with a 500-row max row_group ⇒ exactly 1 row_group, so the
// builder takes the single-row_group fast path.
let total_rows = 321usize;
let buf =
create_test_parquet_bytes_fixed_rg(vec![create_marker_batch(0, total_rows)], 500).await;
let stream_schema = Arc::new(Schema::new(vec![
Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false),
Field::new("marker", DataType::Utf8, false),
]));
let index = build_index(
RamDirectory::create(),
buf,
make_index_schema(&[], &["marker".to_string()], &stream_schema),
4,
)
.await
.expect("single row_group build must succeed")
.expect("single row_group build must return Some(index)");
let segs = index.searchable_segments().unwrap();
assert_eq!(segs.len(), 1);
assert_eq!(segs[0].meta().max_doc() as usize, total_rows);
let reader = index.reader().unwrap();
let searcher = reader.searcher();
let marker_field = index.schema().get_field("marker").unwrap();
for row in 0..total_rows {
let m = format!("row{:09}", row);
let q = tantivy::query::TermQuery::new(
tantivy::Term::from_field_text(marker_field, &m),
tantivy::schema::IndexRecordOption::Basic,
);
let hits = searcher
.search(&q, &tantivy::collector::DocSetCollector)
.unwrap();
assert_eq!(hits.len(), 1);
let addr = *hits.iter().next().unwrap();
assert_eq!(addr.segment_ord, 0);
assert_eq!(addr.doc_id as usize, row);
}
}
/// Same data through parallel-row_group path and sequential path must
/// yield the same `(marker → doc_id)` mapping for every row. Verifies
/// behavioural equivalence end-to-end.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_row_group_parallel_matches_sequential() {
// 750 rows, fixed row_group size 250 ⇒ 3 full row_groups of 250 each.
let total = 750usize;
let buf_par =
create_test_parquet_bytes_fixed_rg(vec![create_marker_batch(0, total)], 250).await;
let buf_seq = create_test_parquet_bytes(vec![
create_marker_batch(0, 250),
create_marker_batch(250, 250),
create_marker_batch(500, 250),
])
.await;
let schema = Arc::new(Schema::new(vec![
Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false),
Field::new("marker", DataType::Utf8, false),
]));
let par = build_index(
RamDirectory::create(),
buf_par,
make_index_schema(&[], &["marker".to_string()], &schema),
3,
)
.await
.unwrap()
.unwrap();
let seq = sequential::build_index(
RamDirectory::create(),
FileFormat::Parquet,
buf_seq,
make_index_schema(&[], &["marker".to_string()], &schema),
)
.await
.unwrap()
.unwrap();
let par_total: usize = par
.searchable_segments()
.unwrap()
.iter()
.map(|s| s.meta().max_doc() as usize)
.sum();
let seq_total: usize = seq
.searchable_segments()
.unwrap()
.iter()
.map(|s| s.meta().max_doc() as usize)
.sum();
assert_eq!(par_total, total);
assert_eq!(seq_total, total);
let par_reader = par.reader().unwrap();
let par_searcher = par_reader.searcher();
let par_marker = par.schema().get_field("marker").unwrap();
let seq_reader = seq.reader().unwrap();
let seq_searcher = seq_reader.searcher();
let seq_marker = seq.schema().get_field("marker").unwrap();
for row in 0..total {
let m = format!("row{:09}", row);
let q_par = tantivy::query::TermQuery::new(
tantivy::Term::from_field_text(par_marker, &m),
tantivy::schema::IndexRecordOption::Basic,
);
let q_seq = tantivy::query::TermQuery::new(
tantivy::Term::from_field_text(seq_marker, &m),
tantivy::schema::IndexRecordOption::Basic,
);
let hits_par = par_searcher
.search(&q_par, &tantivy::collector::DocSetCollector)
.unwrap();
let hits_seq = seq_searcher
.search(&q_seq, &tantivy::collector::DocSetCollector)
.unwrap();
assert_eq!(hits_par.len(), 1);
assert_eq!(hits_seq.len(), 1);
let addr_par = *hits_par.iter().next().unwrap();
let addr_seq = *hits_seq.iter().next().unwrap();
assert_eq!(addr_par.doc_id as usize, row);
assert_eq!(addr_seq.doc_id as usize, row);
}
}
}
src/service/tantivy/sequential.rs
+483
View File
@@ -0,0 +1,483 @@
// Copyright 2026 OpenObserve Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use anyhow::Context;
use arrow::array::RecordBatch;
use bytes::Bytes;
use config::{
FileFormat,
utils::{
parquet::get_recordbatch_reader_from_bytes,
tantivy::tokenizer::{CollectType, O2_TOKENIZER, o2_tokenizer_build},
},
};
use futures::TryStreamExt;
use tokio::task::JoinHandle;
use super::{TantivyIndexSchema, convert_batch_to_docs_sync};
/// Create a tantivy index in the given directory for the input file bytes.
pub(super) async fn build_index<D: tantivy::Directory>(
tantivy_dir: D,
file_format: FileFormat,
buf: Bytes,
index_schema: TantivyIndexSchema,
) -> Result<Option<tantivy::Index>, anyhow::Error> {
let start = std::time::Instant::now();
let (_, reader) = get_recordbatch_reader_from_bytes(file_format, buf).await?;
let tokenizer_manager = tantivy::tokenizer::TokenizerManager::default();
tokenizer_manager.register(O2_TOKENIZER, o2_tokenizer_build(CollectType::Ingest));
let index_writer = tantivy::IndexBuilder::new()
.schema(index_schema.schema.clone())
.tokenizers(tokenizer_manager)
.single_segment_index_writer(tantivy_dir, 50_000_000)
.context("failed to create index builder")?;
let (tx, rx) = tokio::sync::mpsc::channel::<RecordBatch>(2);
let producer: JoinHandle<Result<usize, anyhow::Error>> = tokio::task::spawn(async move {
let mut total_num_rows = 0usize;
let mut reader = reader;
while let Some(batch) = reader.try_next().await? {
let num_rows = batch.num_rows();
if num_rows == 0 {
continue;
}
total_num_rows += num_rows;
if tx.send(batch).await.is_err() {
// Worker exited (likely an error on its side); stop reading.
break;
}
}
Ok(total_num_rows)
});
let mut rx = rx;
let writer_task: JoinHandle<Result<_, anyhow::Error>> =
tokio::task::spawn_blocking(move || {
let mut writer = index_writer;
while let Some(batch) = rx.blocking_recv() {
let docs = convert_batch_to_docs_sync(&batch, &index_schema);
for doc in docs {
writer
.add_document(doc)
.map_err(|e| anyhow::anyhow!("Failed to add document to index: {e}"))?;
}
}
Ok(writer)
});
let _total_num_rows = producer.await??;
let index_writer = writer_task.await??;
log::info!(
"sequential::build_index: index building completed in {} ms",
start.elapsed().as_millis()
);
let start = std::time::Instant::now();
let index = tokio::task::spawn_blocking(move || {
index_writer.finalize().map_err(|e| {
log::error!("sequential::build_index: Failed to finalize the index writer: {e}");
anyhow::anyhow!("Failed to finalize the index writer: {e}")
})
})
.await??;
log::info!(
"sequential::build_index: finalizing index in {}",
start.elapsed().as_millis()
);
Ok(Some(index))
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow::{
array::{Int64Array, StringArray},
datatypes::{DataType, Field, Schema},
record_batch::RecordBatch,
};
use config::{FileFormat, INDEX_FIELD_NAME_FOR_ALL, TIMESTAMP_COL_NAME};
use tantivy::directory::RamDirectory;
use super::*;
use crate::service::tantivy::tests::{
create_test_batch, create_test_parquet_bytes, make_index_schema,
};
#[tokio::test]
async fn test_generate_tantivy_index_empty_data() {
let dir = RamDirectory::create();
let empty_batch = create_test_batch(0, true, true, true);
let buf = create_test_parquet_bytes(vec![empty_batch.clone()]).await;
let result = build_index(
dir,
FileFormat::Parquet,
buf,
make_index_schema(
&["content".to_string()],
&["status".to_string()],
&empty_batch.schema(),
),
)
.await;
assert!(result.is_ok());
// Should create empty index when fields are configured (even with 0 rows)
assert!(result.unwrap().is_some());
}
#[tokio::test]
async fn test_generate_tantivy_index_with_fts_fields() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, true, true, false);
let buf = create_test_parquet_bytes(vec![batch.clone()]).await;
let result = build_index(
dir,
FileFormat::Parquet,
buf,
make_index_schema(&["content".to_string()], &[], &batch.schema()),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_with_index_fields() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, true, false, true);
let buf = create_test_parquet_bytes(vec![batch.clone()]).await;
let result = build_index(
dir,
FileFormat::Parquet,
buf,
make_index_schema(&[], &["status".to_string()], &batch.schema()),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field("status").is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_with_both_field_types() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, true, true, true);
let buf = create_test_parquet_bytes(vec![batch.clone()]).await;
let result = build_index(
dir,
FileFormat::Parquet,
buf,
make_index_schema(
&["content".to_string()],
&["status".to_string()],
&batch.schema(),
),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field("status").is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_mixed_existing_and_missing_fields() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, true, true, true);
let buf = create_test_parquet_bytes(vec![batch.clone()]).await;
// Mix of existing and non-existing fields
let result = build_index(
dir,
FileFormat::Parquet,
buf,
make_index_schema(
&["content".to_string(), "nonexistent_field".to_string()],
&[
"status".to_string(),
"another_nonexistent_field".to_string(),
],
&batch.schema(),
),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field("status").is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_multiple_batches() {
let dir = RamDirectory::create();
let batch1 = create_test_batch(5, true, true, true);
let batch2 = create_test_batch(5, true, true, true);
let buf = create_test_parquet_bytes(vec![batch1.clone(), batch2.clone()]).await;
let result = build_index(
dir,
FileFormat::Parquet,
buf,
make_index_schema(
&["content".to_string()],
&["status".to_string()],
&batch1.schema(),
),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field("status").is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_without_timestamp() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, false, true, true);
let buf = create_test_parquet_bytes(vec![batch.clone()]).await;
let result = build_index(
dir,
FileFormat::Parquet,
buf,
make_index_schema(
&["content".to_string()],
&["status".to_string()],
&batch.schema(),
),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field("status").is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_non_string_fields() {
// Create a batch with non-string fields in the schema
let fields = vec![
Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false),
Field::new("content", DataType::Utf8, false),
Field::new("number_field", DataType::Int32, false), // Non-string field
];
let schema = Arc::new(Schema::new(fields));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int64Array::from(vec![1000, 1001, 1002])),
Arc::new(StringArray::from(vec!["content1", "content2", "content3"])),
Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3])),
],
)
.unwrap();
let dir = RamDirectory::create();
let buf = create_test_parquet_bytes(vec![batch.clone()]).await;
let result = build_index(
dir,
FileFormat::Parquet,
buf,
make_index_schema(
&["content".to_string()],
&["number_field".to_string()], // This field is not Utf8
&batch.schema(),
),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field("number_field").is_ok()); // Non-string fields are still indexed
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
}
#[tokio::test]
async fn test_generate_tantivy_index_timestamp_field_handling() {
let dir = RamDirectory::create();
let batch = create_test_batch(10, true, true, true);
let buf = create_test_parquet_bytes(vec![batch.clone()]).await;
// Try to index the timestamp field as a regular index field
let result = build_index(
dir,
FileFormat::Parquet,
buf,
make_index_schema(
&["content".to_string()],
&[TIMESTAMP_COL_NAME.to_string()], // This should be ignored
&batch.schema(),
),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
let index = index.unwrap();
let schema = index.schema();
assert!(schema.get_field(INDEX_FIELD_NAME_FOR_ALL).is_ok());
assert!(schema.get_field(TIMESTAMP_COL_NAME).is_ok());
// Verify that the timestamp field is indexed as Int64, not as text
let ts_field = schema.get_field(TIMESTAMP_COL_NAME).unwrap();
assert!(matches!(
schema.get_field_entry(ts_field).field_type(),
tantivy::schema::FieldType::I64(_)
));
}
#[tokio::test]
async fn test_generate_tantivy_index_stream_schema_vs_parquet_schema() {
// This test validates the critical fix where we pass stream schema (with all
// configured fields) instead of parquet schema (only fields in actual data).
//
// Scenario: Stream settings configure `continent`, `name`, `flag_url` as secondary
// index fields, but the parquet data only contains `name` and `flag_url`.
// The missing field `continent` should still be included in the .ttv schema.
let dir = RamDirectory::create();
let parquet_fields = vec![
Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
Field::new("flag_url", DataType::Utf8, false),
];
let parquet_schema = Arc::new(Schema::new(parquet_fields));
let timestamps: Vec<i64> = (0..10).map(|i| 1000 + i as i64).collect();
let names: Vec<String> = (0..10).map(|i| format!("Name {i}")).collect();
let flag_urls: Vec<String> = (0..10)
.map(|i| format!("https://example.com/{i}"))
.collect();
let parquet_batch = RecordBatch::try_new(
parquet_schema.clone(),
vec![
Arc::new(Int64Array::from(timestamps)),
Arc::new(StringArray::from(names)),
Arc::new(StringArray::from(flag_urls)),
],
)
.unwrap();
let stream_fields = vec![
Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false),
Field::new("continent", DataType::Utf8, false), // Configured but not in data
Field::new("name", DataType::Utf8, false),
Field::new("flag_url", DataType::Utf8, false),
];
let stream_schema = Arc::new(Schema::new(stream_fields));
let buf = create_test_parquet_bytes(vec![parquet_batch]).await;
let result = build_index(
dir,
FileFormat::Parquet,
buf,
make_index_schema(
&[], // No FTS fields
&[
"continent".to_string(),
"name".to_string(),
"flag_url".to_string(),
],
&stream_schema, // Pass stream schema, not parquet schema
),
)
.await;
assert!(result.is_ok());
let index = result.unwrap();
assert!(index.is_some());
let index = index.unwrap();
let tantivy_schema = index.schema();
assert!(
tantivy_schema.get_field("continent").is_ok(),
"continent field should be in tantivy schema even though it's not in parquet data"
);
assert!(
tantivy_schema.get_field("name").is_ok(),
"name field should be in tantivy schema"
);
assert!(
tantivy_schema.get_field("flag_url").is_ok(),
"flag_url field should be in tantivy schema"
);
assert!(tantivy_schema.get_field(TIMESTAMP_COL_NAME).is_ok());
let reader = index.reader().unwrap();
let searcher = reader.searcher();
assert_eq!(searcher.num_docs(), 10);
}
}