diff --git a/src/config/src/config.rs b/src/config/src/config.rs index 19896ccdf9..e555d87f5e 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -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)] diff --git a/src/job/files/parquet.rs b/src/job/files/parquet.rs index b53c9755e9..fc7c659622 100644 --- a/src/job/files/parquet.rs +++ b/src/job/files/parquet.rs @@ -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}"))?; diff --git a/src/service/compact/merge.rs b/src/service/compact/merge.rs index 15c51fe427..b602b0b737 100644 --- a/src/service/compact/merge.rs +++ b/src/service/compact/merge.rs @@ -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, 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| { diff --git a/src/service/tantivy/mod.rs b/src/service/tantivy/mod.rs index 54b226c8f4..23c2117bc4 100644 --- a/src/service/tantivy/mod.rs +++ b/src/service/tantivy/mod.rs @@ -13,97 +13,62 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . +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, - reader: RecordBatchStream, + buf: Bytes, ) -> Result { 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( - 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, + /// Field handle for `INDEX_FIELD_NAME_FOR_ALL`, present when there is at + /// least one full-text-search field. + pub(super) fts_field: Option, +} + +/// 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, -) -> Result, anyhow::Error> { + arrow_schema: &Schema, +) -> Option { 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::>(); - // 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( }) .map(String::from) .collect::>(); - let index_fields = index_fields + let index_fields_filtered = index_fields .iter() .filter(|f| schema_fields.contains_key(f)) .map(String::from) .collect::>(); - let tantivy_fields = fts_fields - .union(&index_fields) + let tantivy_fields = fts_fields_filtered + .union(&index_fields_filtered) .cloned() .collect::>(); 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( .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::>(2); - let task: JoinHandle> = 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::() { - 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 { + 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::() { + 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> = Vec::new(); - // Add timestamp field if requested if include_timestamp { fields.push(Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false)); let timestamps: Vec = (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 = (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 = (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) -> 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, + ) -> 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 = (0..10).map(|i| 1000 + i as i64).collect(); - let names: Vec = (0..10).map(|i| format!("Name {i}")).collect(); - let flag_urls: Vec = (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. } diff --git a/src/service/tantivy/parallel.rs b/src/service/tantivy/parallel.rs new file mode 100644 index 0000000000..74d95b5002 --- /dev/null +++ b/src/service/tantivy/parallel.rs @@ -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 . + +//! 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( + tantivy_dir: D, + buf: Bytes, + index_schema: TantivyIndexSchema, + thread_num: usize, +) -> Result, 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 { + 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>> = 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 = 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 = 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( + 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::(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 = (0..num_rows).map(|i| (start + i) as i64).collect(); + let markers: Vec = (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, + 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); + } + } +} diff --git a/src/service/tantivy/sequential.rs b/src/service/tantivy/sequential.rs new file mode 100644 index 0000000000..a12c550d2d --- /dev/null +++ b/src/service/tantivy/sequential.rs @@ -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 . + +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( + tantivy_dir: D, + file_format: FileFormat, + buf: Bytes, + index_schema: TantivyIndexSchema, +) -> Result, 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::(2); + + let producer: JoinHandle> = 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> = + 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 = (0..10).map(|i| 1000 + i as i64).collect(); + let names: Vec = (0..10).map(|i| format!("Name {i}")).collect(); + let flag_urls: Vec = (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); + } +}