From 3a951f1398cd6a817bee4fdb28fcb59ca572f8ea Mon Sep 17 00:00:00 2001 From: Huaijin Date: Mon, 1 Jun 2026 14:51:41 +0800 Subject: [PATCH] perf: impl new tantivy index generate for vortex format (#12354) related to - https://github.com/openobserve/openobserve/pull/12331 - https://github.com/openobserve/openobserve/pull/12292 --- src/config/src/utils/inverted_index.rs | 34 +- src/config/src/utils/parquet.rs | 68 +- src/handler/grpc/request/event.rs | 14 +- src/service/compact/bloom_build.rs | 4 +- src/service/compact/deleted.rs | 5 +- src/service/compact/merge.rs | 1275 +------------------ src/service/search/grpc/storage.rs | 19 +- src/service/tantivy/bloom_builder.rs | 22 +- src/service/tantivy/mod.rs | 40 +- src/service/tantivy/parallel.rs | 471 +++++-- src/service/tantivy/puffin/mod.rs | 24 - src/service/tantivy/puffin_directory/mod.rs | 32 - src/service/tantivy/reader/mod.rs | 68 + src/service/tantivy/reader/parquet.rs | 85 ++ src/service/tantivy/reader/vortex.rs | 122 ++ src/service/tantivy/sequential.rs | 56 +- 16 files changed, 744 insertions(+), 1595 deletions(-) delete mode 100644 src/service/tantivy/puffin/mod.rs delete mode 100644 src/service/tantivy/puffin_directory/mod.rs create mode 100644 src/service/tantivy/reader/mod.rs create mode 100644 src/service/tantivy/reader/parquet.rs create mode 100644 src/service/tantivy/reader/vortex.rs diff --git a/src/config/src/utils/inverted_index.rs b/src/config/src/utils/inverted_index.rs index 4371df0861..8af4a3d890 100644 --- a/src/config/src/utils/inverted_index.rs +++ b/src/config/src/utils/inverted_index.rs @@ -17,12 +17,12 @@ use std::borrow::Cow; use crate::{FILE_EXT_PARQUET, FILE_EXT_TANTIVY, FILE_EXT_VORTEX, meta::stream::StreamType}; -/// tantivy inverted index solution has a 1:1 mapping between parquet and idx files. -/// This is a helper function to convert the parquet file name to tantivy file name. +/// Converts a parquet or vortex file name to the corresponding tantivy index file name. +/// Returns `None` for unrecognized stream types, unsupported extensions, or short paths. /// e.g. /// from: files/default/logs/quickstart1/2024/02/16/16/7164299619311026293.parquet /// to: files/default/index/quickstart1_logs/2024/02/16/16/7164299619311026293.ttv -pub fn convert_parquet_file_name_to_tantivy_file(from: &str) -> Option { +pub fn to_tantivy_name(from: &str) -> Option { let mut parts: Vec> = from.split('/').map(Cow::Borrowed).collect(); if parts.len() < 4 { @@ -64,7 +64,7 @@ mod tests { use super::*; #[test] - fn test_convert_parquet_file_name_to_tantivy_file() { + fn test_to_tantivy_name() { let test_cases = vec![ ( "files/default/logs/quickstart1/2024/02/16/16/7164299619311026293.parquet", @@ -101,36 +101,30 @@ mod tests { ]; for (input, expected) in test_cases { - assert_eq!(convert_parquet_file_name_to_tantivy_file(input), expected); + assert_eq!(to_tantivy_name(input), expected); } } #[test] - fn test_convert_vortex_file_name_to_tantivy_file() { + fn test_to_tantivy_name_vortex() { let input = "files/default/logs/quickstart1/2024/02/16/16/7164299619311026293.vortex"; let expected = "files/default/index/quickstart1_logs/2024/02/16/16/7164299619311026293.ttv"; - assert_eq!( - convert_parquet_file_name_to_tantivy_file(input), - Some(expected.to_string()) - ); + assert_eq!(to_tantivy_name(input), Some(expected.to_string())); } #[test] fn test_convert_returns_none_for_unsupported_extension() { // Files that are neither .parquet nor .vortex should return None. let input = "files/default/logs/quickstart1/2024/02/16/16/7164299619311026293.json"; - assert_eq!(convert_parquet_file_name_to_tantivy_file(input), None); + assert_eq!(to_tantivy_name(input), None); } #[test] fn test_convert_returns_none_for_short_path() { // Paths shorter than 4 segments are rejected — there is no stream_type slot. - assert_eq!( - convert_parquet_file_name_to_tantivy_file("files/default/logs"), - None - ); - assert_eq!(convert_parquet_file_name_to_tantivy_file(""), None); - assert_eq!(convert_parquet_file_name_to_tantivy_file("a/b/c"), None); + assert_eq!(to_tantivy_name("files/default/logs"), None); + assert_eq!(to_tantivy_name(""), None); + assert_eq!(to_tantivy_name("a/b/c"), None); } #[test] @@ -138,16 +132,16 @@ mod tests { // service_graph / enrichment_tables / file_list are not handled by this // converter — only logs/metrics/traces/metadata map to a tantivy index. let input = "files/default/service_graph/x/2024/02/16/16/1.parquet"; - assert_eq!(convert_parquet_file_name_to_tantivy_file(input), None); + assert_eq!(to_tantivy_name(input), None); let input = "files/default/enrichment_tables/x/2024/02/16/16/1.parquet"; - assert_eq!(convert_parquet_file_name_to_tantivy_file(input), None); + assert_eq!(to_tantivy_name(input), None); } #[test] fn test_convert_path_without_extension_returns_none() { // Final segment without a recognized extension is rejected. let input = "files/default/logs/quickstart1/2024/02/16/16/no_extension"; - assert_eq!(convert_parquet_file_name_to_tantivy_file(input), None); + assert_eq!(to_tantivy_name(input), None); } } diff --git a/src/config/src/utils/parquet.rs b/src/config/src/utils/parquet.rs index 06eec66bb0..4d2a334d98 100644 --- a/src/config/src/utils/parquet.rs +++ b/src/config/src/utils/parquet.rs @@ -23,18 +23,18 @@ use arrow_schema::Schema; use futures::StreamExt; use futures::{Stream, TryStreamExt}; use parquet::{ - arrow::{ - AsyncArrowWriter, ParquetRecordBatchStreamBuilder, - arrow_reader::ArrowReaderMetadata, - async_reader::{AsyncFileReader, ParquetRecordBatchStream}, - }, + arrow::{AsyncArrowWriter, ParquetRecordBatchStreamBuilder, arrow_reader::ArrowReaderMetadata}, basic::{Compression, Encoding}, file::{metadata::KeyValue, properties::WriterProperties}, }; #[cfg(feature = "vortex")] use vortex::{ - VortexSessionDefault, array::arrow::IntoArrowArray, buffer::Buffer, - file::OpenOptionsSessionExt, io::session::RuntimeSessionExt, session::VortexSession, + VortexSessionDefault, + array::{ArrayRef, arrow::IntoArrowArray}, + buffer::Buffer, + file::OpenOptionsSessionExt, + io::session::RuntimeSessionExt, + session::VortexSession, }; use crate::{FileFormat, config::*, ider, meta::stream::FileMeta}; @@ -113,18 +113,22 @@ pub fn parse_file_key_columns(key: &str) -> Result<(String, String, String), any /// Unified stream type that supports both Vortex and Parquet formats pub type RecordBatchStream = Pin> + Send>>; -/// A generic wrapper for getting a record batch stream from a reader. -/// It can be used for both bytes and file. -async fn get_recordbatch_reader( - reader: T, -) -> Result<(Arc, ParquetRecordBatchStream), anyhow::Error> -where - T: AsyncFileReader + Send + 'static, -{ - let arrow_reader = ParquetRecordBatchStreamBuilder::new(reader).await?; - let schema = arrow_reader.schema().clone(); - let reader = arrow_reader.with_batch_size(get_batch_size()).build()?; - Ok((schema, reader)) +/// Convert a single vortex [`ArrayRef`] to an Arrow [`RecordBatch`]. +#[cfg(feature = "vortex")] +pub fn vortex_array_to_record_batch( + array: ArrayRef, + data_type: &DataType, +) -> Result { + let array = array + .into_arrow(data_type) + .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + let struct_array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + ArrowError::InvalidArgumentError("Expected struct array from vortex".to_string()) + })?; + Ok(RecordBatch::from(struct_array)) } pub async fn get_recordbatch_reader_from_bytes( @@ -133,8 +137,9 @@ pub async fn get_recordbatch_reader_from_bytes( ) -> Result<(Arc, RecordBatchStream), anyhow::Error> { match file_format { FileFormat::Parquet => { - let schema_reader = Cursor::new(data); - let (schema, reader) = get_recordbatch_reader(schema_reader).await?; + let arrow_reader = ParquetRecordBatchStreamBuilder::new(Cursor::new(data)).await?; + let schema = arrow_reader.schema().clone(); + let reader = arrow_reader.with_batch_size(get_batch_size()).build()?; let stream: RecordBatchStream = Box::pin(reader.map_err(ArrowError::from)); Ok((schema, stream)) } @@ -145,8 +150,6 @@ pub async fn get_recordbatch_reader_from_bytes( let buf = Buffer::from(data.to_vec()); let vxf = session.open_options().open_buffer(buf)?; let schema = Arc::new(vxf.dtype().to_arrow_schema()?); - - // Create a stream that converts vortex arrays to record batches let arrow_data_type = DataType::Struct(schema.fields().clone()); let vortex_stream = vxf.scan()?.into_array_stream()?; @@ -154,24 +157,7 @@ pub async fn get_recordbatch_reader_from_bytes( let arrow_data_type = arrow_data_type.clone(); async move { match result { - Ok(vortex_array) => { - // Convert vortex array to arrow array - let arrow_array = vortex_array - .into_arrow(&arrow_data_type) - .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; - - // Convert to struct array and then to record batch - let struct_array = arrow_array - .as_any() - .downcast_ref::() - .ok_or_else(|| { - ArrowError::InvalidArgumentError( - "Expected struct array from vortex".to_string(), - ) - })?; - - Ok(RecordBatch::from(struct_array)) - } + Ok(array) => vortex_array_to_record_batch(array, &arrow_data_type), Err(e) => Err(ArrowError::ExternalError(Box::new(e))), } } diff --git a/src/handler/grpc/request/event.rs b/src/handler/grpc/request/event.rs index 7b86e628f3..774097f070 100644 --- a/src/handler/grpc/request/event.rs +++ b/src/handler/grpc/request/event.rs @@ -18,7 +18,7 @@ use std::ops::Range; use anyhow::Result; use config::{ cluster::LOCAL_NODE, get_config, meta::stream::FileKey, metrics, - utils::inverted_index::convert_parquet_file_name_to_tantivy_file, + utils::inverted_index::to_tantivy_name, }; use infra::cache::file_data::{CacheType, TRACE_ID_FOR_CACHE_LATEST_FILE, disk}; use opentelemetry::global; @@ -78,7 +78,7 @@ impl Event for Eventer { // cache index for the parquet if cfg.cache_latest_files.cache_index && item.meta.index_size > 0 - && let Some(ttv_file) = convert_parquet_file_name_to_tantivy_file(&item.key) + && let Some(ttv_file) = to_tantivy_name(&item.key) { files_to_download.push(( item.id, @@ -157,9 +157,7 @@ impl Event for Eventer { .filter_map(|v| { if v.deleted { match v.meta.as_ref() { - Some(m) if m.index_size > 0 => { - convert_parquet_file_name_to_tantivy_file(&v.key) - } + Some(m) if m.index_size > 0 => to_tantivy_name(&v.key), _ => None, } } else { @@ -455,7 +453,7 @@ mod tests { // Test parquet to tantivy filename conversion let parquet_file = "files/default/logs/quickstart1/2024/02/16/16/7164299619311026293.parquet"; - let tantivy_file = convert_parquet_file_name_to_tantivy_file(parquet_file); + let tantivy_file = to_tantivy_name(parquet_file); // The conversion should return Some for valid parquet files assert!(tantivy_file.is_some()); @@ -466,12 +464,12 @@ mod tests { // Test with non-parquet file let non_parquet_file = "test/file.txt"; - let tantivy_result = convert_parquet_file_name_to_tantivy_file(non_parquet_file); + let tantivy_result = to_tantivy_name(non_parquet_file); assert!(tantivy_result.is_none()); // Test with invalid path format let invalid_path = "test/file.parquet"; - let invalid_result = convert_parquet_file_name_to_tantivy_file(invalid_path); + let invalid_result = to_tantivy_name(invalid_path); assert!(invalid_result.is_none()); } diff --git a/src/service/compact/bloom_build.rs b/src/service/compact/bloom_build.rs index e2368fbff9..bbebe6fa0f 100644 --- a/src/service/compact/bloom_build.rs +++ b/src/service/compact/bloom_build.rs @@ -41,7 +41,7 @@ use bytes::Bytes; use config::{ get_config, meta::stream::{FileKey, FileListDeleted, StreamType}, - utils::{inverted_index::convert_parquet_file_name_to_tantivy_file, time::now_micros}, + utils::{inverted_index::to_tantivy_name, time::now_micros}, }; use infra::{ bloom::{BloomWriter, FieldBloom, path::bloom_path}, @@ -232,7 +232,7 @@ async fn build_for_file( target_fields: &[String], num_blocks: u32, ) -> anyhow::Result> { - let Some(ttv_file_name) = convert_parquet_file_name_to_tantivy_file(&file.key) else { + let Some(ttv_file_name) = to_tantivy_name(&file.key) else { return Ok(Vec::new()); // not an indexable file }; if file.meta.index_size == 0 { diff --git a/src/service/compact/deleted.rs b/src/service/compact/deleted.rs index d5d85a0987..69457aa3d8 100644 --- a/src/service/compact/deleted.rs +++ b/src/service/compact/deleted.rs @@ -15,7 +15,7 @@ use config::{ meta::stream::{FileKey, FileMeta}, - utils::inverted_index::convert_parquet_file_name_to_tantivy_file, + utils::inverted_index::to_tantivy_name, }; use infra::{file_list as infra_file_list, storage}; @@ -56,8 +56,7 @@ pub async fn delete(org_id: &str, time_max: i64) -> Result { .iter() .filter_map(|file| { if file.index_file { - convert_parquet_file_name_to_tantivy_file(&file.file) - .map(|f| (file.account.to_string(), f)) + to_tantivy_name(&file.file).map(|f| (file.account.to_string(), f)) } else { None } diff --git a/src/service/compact/merge.rs b/src/service/compact/merge.rs index 7e06f16e69..6a30a2a3ab 100644 --- a/src/service/compact/merge.rs +++ b/src/service/compact/merge.rs @@ -16,11 +16,10 @@ use std::sync::Arc; use ::datafusion::{arrow::datatypes::Schema, error::DataFusionError}; -use arrow::array::RecordBatch; use bytes::Bytes; use chrono::{DateTime, Datelike, Duration, TimeZone, Utc}; use config::{ - FileFormat, TIMESTAMP_COL_NAME, + FileFormat, cluster::LOCAL_NODE, get_config, ider, is_local_disk_storage, meta::stream::{ @@ -30,7 +29,6 @@ use config::{ metrics, utils::{ parquet::read_schema_from_bytes, - record_batch_ext::concat_batches, schema_ext::SchemaExt, time::{day_micros, hour_micros}, }, @@ -1023,7 +1021,7 @@ pub async fn merge_files( async fn generate_inverted_index( org_id: &str, new_file_key: &str, - full_text_search_fields: &[String], + fts_fields: &[String], index_fields: &[String], retain_file_list: &[FileKey], new_file_meta: &mut FileMeta, @@ -1034,7 +1032,7 @@ async fn generate_inverted_index( "COMPACTOR", org_id, new_file_key, - full_text_search_fields, + fts_fields, index_fields, latest_schema, // Use stream schema to include all configured fields buf, @@ -1123,86 +1121,6 @@ async fn write_file_list( Ok(()) } -pub fn generate_inverted_idx_recordbatch( - schema: Arc, - batches: &[RecordBatch], - stream_type: StreamType, - full_text_search_fields: &[String], - index_fields: &[String], -) -> Result, anyhow::Error> { - let cfg = get_config(); - if !cfg.common.inverted_index_enabled || batches.is_empty() || !stream_type.support_index() { - return Ok(None); - } - - let schema_fields = schema - .fields() - .iter() - .map(|f| f.name()) - .collect::>(); - - let mut inverted_idx_columns: Vec = if !full_text_search_fields.is_empty() { - full_text_search_fields.to_vec() - } else { - config::SQL_FULL_TEXT_SEARCH_FIELDS.to_vec() - }; - inverted_idx_columns.extend(index_fields.to_vec()); - inverted_idx_columns.sort(); - inverted_idx_columns.dedup(); - inverted_idx_columns.retain(|f| schema_fields.contains(f)); - if inverted_idx_columns.is_empty() { - return Ok(None); - } - // add _timestamp column to columns_to_index - if !inverted_idx_columns.contains(&TIMESTAMP_COL_NAME.to_string()) { - inverted_idx_columns.push(TIMESTAMP_COL_NAME.to_string()); - } - - let mut inverted_idx_batches = Vec::with_capacity(batches.len()); - for batch in batches { - let selected_column_indices: Vec = inverted_idx_columns - .iter() - .filter_map(|name| batch.schema().index_of(name).ok()) - .collect(); - - // Use the found indices to select the columns - let selected_columns = selected_column_indices - .iter() - .map(|&i| batch.column(i).clone()) - .collect(); - - // Create a new schema for the new RecordBatch based on the selected columns - let selected_fields: Vec = selected_column_indices - .iter() - .map(|&i| batch.schema().field(i).clone()) - .collect(); - let new_schema = Arc::new(Schema::new(selected_fields)); - - // Create a new RecordBatch with the selected columns - inverted_idx_batches.push(RecordBatch::try_new(new_schema, selected_columns).unwrap()); - } - - if inverted_idx_batches.is_empty() { - Ok(None) - } else { - let new_batch = if inverted_idx_batches.len() == 1 { - inverted_idx_batches.remove(0) - } else { - let new_schema = inverted_idx_batches.first().unwrap().schema(); - concat_batches(new_schema, inverted_idx_batches).map_err(anyhow::Error::from)? - }; - - if matches!( - new_batch.schema().fields().len(), - 0 | 1 if new_batch.schema().field(0).name() == TIMESTAMP_COL_NAME - ) { - Ok(None) - } else { - Ok(Some(new_batch)) - } - } -} - async fn cache_remote_files(files: &[FileKey]) -> Result, anyhow::Error> { let cfg = get_config(); let scan_size = files.iter().map(|f| f.meta.compressed_size).sum::(); @@ -1235,16 +1153,6 @@ async fn cache_remote_files(files: &[FileKey]) -> Result, anyhow::Er log::warn!( "[COMPACT] download file {file_name} found size mismatch, expected: {file_size}, actual: {data_len}, will skip it", ); - // update database - // if let Err(e) = - // file_list::update_compressed_size(&file_name, data_len as i64).await - // { - // log::error!( - // "[COMPACT] update file size for file {} err: {}", - // file_name, - // e - // ); - // } // skip this file for compact Some(file_name) } else { @@ -1325,9 +1233,6 @@ fn sort_by_time_range(mut file_list: Vec) -> Vec { #[cfg(test)] mod tests { - use std::sync::Arc; - - use arrow_schema::{DataType, Field, Schema}; use config::meta::stream::{FileKey, FileMeta}; use super::*; @@ -1585,358 +1490,6 @@ mod tests { assert!(keys.contains(&&"file3.parquet".to_string())); } - // Test cases for generate_inverted_idx_recordbatch function - #[test] - fn test_generate_inverted_idx_recordbatch_empty_batches() { - let schema = Arc::new(Schema::new(vec![ - Field::new("field1", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batches: Vec = vec![]; - let full_text_search_fields = vec!["field1".to_string()]; - let index_fields = vec![]; - - let result = generate_inverted_idx_recordbatch( - schema, - &batches, - StreamType::Logs, - &full_text_search_fields, - &index_fields, - ); - - assert!(result.is_ok()); - assert!(result.unwrap().is_none()); - } - - #[test] - fn test_generate_inverted_idx_recordbatch_with_fts_fields() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("message", DataType::Utf8, true), - Field::new("level", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let message_array = StringArray::from(vec!["log message 1", "log message 2"]); - let level_array = StringArray::from(vec!["info", "error"]); - let timestamp_array = Int64Array::from(vec![1000, 2000]); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(message_array), - Arc::new(level_array), - Arc::new(timestamp_array), - ], - ) - .unwrap(); - - let batches = vec![batch]; - let full_text_search_fields = vec!["message".to_string()]; - let index_fields = vec![]; - - let result = generate_inverted_idx_recordbatch( - schema, - &batches, - StreamType::Logs, - &full_text_search_fields, - &index_fields, - ); - - assert!(result.is_ok()); - let batch_result = result.unwrap(); - assert!(batch_result.is_some()); - - let inverted_batch = batch_result.unwrap(); - assert_eq!(inverted_batch.num_rows(), 2); - assert_eq!(inverted_batch.num_columns(), 2); // message + _timestamp - } - - #[test] - fn test_generate_inverted_idx_recordbatch_with_index_fields() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("user_id", DataType::Utf8, true), - Field::new("session_id", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let user_id_array = StringArray::from(vec!["user1", "user2"]); - let session_id_array = StringArray::from(vec!["session1", "session2"]); - let timestamp_array = Int64Array::from(vec![1000, 2000]); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(user_id_array), - Arc::new(session_id_array), - Arc::new(timestamp_array), - ], - ) - .unwrap(); - - let batches = vec![batch]; - let full_text_search_fields = vec![]; - let index_fields = vec!["user_id".to_string(), "session_id".to_string()]; - - let result = generate_inverted_idx_recordbatch( - schema, - &batches, - StreamType::Logs, - &full_text_search_fields, - &index_fields, - ); - - assert!(result.is_ok()); - let batch_result = result.unwrap(); - assert!(batch_result.is_some()); - - let inverted_batch = batch_result.unwrap(); - assert_eq!(inverted_batch.num_rows(), 2); - assert_eq!(inverted_batch.num_columns(), 3); // user_id + session_id + _timestamp - } - - #[test] - fn test_generate_inverted_idx_recordbatch_with_multiple_batches() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("message", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch1 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["message 1", "message 2"])), - Arc::new(Int64Array::from(vec![1000, 2000])), - ], - ) - .unwrap(); - - let batch2 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["message 3", "message 4"])), - Arc::new(Int64Array::from(vec![3000, 4000])), - ], - ) - .unwrap(); - - let batches = vec![batch1, batch2]; - let full_text_search_fields = vec!["message".to_string()]; - let index_fields = vec![]; - - let result = generate_inverted_idx_recordbatch( - schema, - &batches, - StreamType::Logs, - &full_text_search_fields, - &index_fields, - ); - - assert!(result.is_ok()); - let batch_result = result.unwrap(); - assert!(batch_result.is_some()); - - let inverted_batch = batch_result.unwrap(); - assert_eq!(inverted_batch.num_rows(), 4); // Combined rows from both batches - } - - #[test] - fn test_generate_inverted_idx_recordbatch_no_matching_fields() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("field1", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["value1", "value2"])), - Arc::new(Int64Array::from(vec![1000, 2000])), - ], - ) - .unwrap(); - - let batches = vec![batch]; - // Specify fields that don't exist in schema - let full_text_search_fields = vec!["nonexistent_field".to_string()]; - let index_fields = vec!["another_nonexistent".to_string()]; - - let result = generate_inverted_idx_recordbatch( - schema, - &batches, - StreamType::Logs, - &full_text_search_fields, - &index_fields, - ); - - assert!(result.is_ok()); - let batch_result = result.unwrap(); - // Should return None because only _timestamp would be present - assert!(batch_result.is_none()); - } - - #[test] - fn test_generate_inverted_idx_recordbatch_only_timestamp() { - use arrow::array::Int64Array; - - let schema = Arc::new(Schema::new(vec![Field::new( - TIMESTAMP_COL_NAME, - DataType::Int64, - false, - )])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int64Array::from(vec![1000, 2000]))], - ) - .unwrap(); - - let batches = vec![batch]; - let full_text_search_fields = vec![]; - let index_fields = vec![]; - - let result = generate_inverted_idx_recordbatch( - schema, - &batches, - StreamType::Logs, - &full_text_search_fields, - &index_fields, - ); - - assert!(result.is_ok()); - let batch_result = result.unwrap(); - // Should return None because only timestamp field exists - assert!(batch_result.is_none()); - } - - #[test] - fn test_generate_inverted_idx_recordbatch_with_combined_fts_and_index() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("message", DataType::Utf8, true), - Field::new("user_id", DataType::Utf8, true), - Field::new("level", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["log1", "log2"])), - Arc::new(StringArray::from(vec!["user1", "user2"])), - Arc::new(StringArray::from(vec!["info", "error"])), - Arc::new(Int64Array::from(vec![1000, 2000])), - ], - ) - .unwrap(); - - let batches = vec![batch]; - let full_text_search_fields = vec!["message".to_string()]; - let index_fields = vec!["user_id".to_string(), "level".to_string()]; - - let result = generate_inverted_idx_recordbatch( - schema, - &batches, - StreamType::Logs, - &full_text_search_fields, - &index_fields, - ); - - assert!(result.is_ok()); - let batch_result = result.unwrap(); - assert!(batch_result.is_some()); - - let inverted_batch = batch_result.unwrap(); - assert_eq!(inverted_batch.num_rows(), 2); - // Should have: level, message, user_id (sorted), + _timestamp = 4 columns - assert_eq!(inverted_batch.num_columns(), 4); - } - - #[test] - fn test_generate_inverted_idx_recordbatch_duplicate_fields() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("field1", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["value1", "value2"])), - Arc::new(Int64Array::from(vec![1000, 2000])), - ], - ) - .unwrap(); - - let batches = vec![batch]; - // Same field in both FTS and index fields (should be deduped) - let full_text_search_fields = vec!["field1".to_string()]; - let index_fields = vec!["field1".to_string()]; - - let result = generate_inverted_idx_recordbatch( - schema, - &batches, - StreamType::Logs, - &full_text_search_fields, - &index_fields, - ); - - assert!(result.is_ok()); - let batch_result = result.unwrap(); - assert!(batch_result.is_some()); - - let inverted_batch = batch_result.unwrap(); - assert_eq!(inverted_batch.num_rows(), 2); - assert_eq!(inverted_batch.num_columns(), 2); // field1 + _timestamp (deduplicated) - } - - #[test] - fn test_generate_inverted_idx_recordbatch_unsupported_stream_type() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("message", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["msg1", "msg2"])), - Arc::new(Int64Array::from(vec![1000, 2000])), - ], - ) - .unwrap(); - - let batches = vec![batch]; - let full_text_search_fields = vec!["message".to_string()]; - let index_fields = vec![]; - - // Test with a stream type that might not support indexing - // Note: This depends on the support_index() implementation - let result = generate_inverted_idx_recordbatch( - schema, - &batches, - StreamType::Metrics, // Metrics may not support indexing - &full_text_search_fields, - &index_fields, - ); - - assert!(result.is_ok()); - // Result depends on whether Metrics support_index() returns true - } - // ── Additional sort_by_time_range edge-case tests ──────────────────────── /// Two files share the same min_ts but have different max_ts values. @@ -2057,434 +1610,6 @@ mod tests { assert_eq!(result[0].meta.max_ts, 42); } - // ── Additional generate_inverted_idx_recordbatch edge-case tests ───────── - - /// Stream types that return false from `support_index()` must always yield None. - #[test] - fn test_generate_inverted_idx_recordbatch_non_index_stream_type() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("message", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["hello"])), - Arc::new(Int64Array::from(vec![1000i64])), - ], - ) - .unwrap(); - - let batches = vec![batch]; - let full_text_search_fields = vec!["message".to_string()]; - let index_fields: Vec = vec![]; - - // EnrichmentTables does not support indexing. - let result = generate_inverted_idx_recordbatch( - schema.clone(), - &batches, - StreamType::EnrichmentTables, - &full_text_search_fields, - &index_fields, - ); - assert!(result.is_ok()); - assert!( - result.unwrap().is_none(), - "EnrichmentTables should return None (no index support)" - ); - - // ServiceGraph does not support indexing. - let batch2 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["hello"])), - Arc::new(Int64Array::from(vec![1000i64])), - ], - ) - .unwrap(); - let result2 = generate_inverted_idx_recordbatch( - schema.clone(), - &[batch2], - StreamType::ServiceGraph, - &full_text_search_fields, - &index_fields, - ); - assert!(result2.is_ok()); - assert!( - result2.unwrap().is_none(), - "ServiceGraph should return None (no index support)" - ); - - // Filelist does not support indexing. - let batch3 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["hello"])), - Arc::new(Int64Array::from(vec![1000i64])), - ], - ) - .unwrap(); - let result3 = generate_inverted_idx_recordbatch( - schema, - &[batch3], - StreamType::Filelist, - &full_text_search_fields, - &index_fields, - ); - assert!(result3.is_ok()); - assert!( - result3.unwrap().is_none(), - "Filelist should return None (no index support)" - ); - } - - /// When `full_text_search_fields` is empty and `index_fields` is empty the - /// function falls back to `SQL_FULL_TEXT_SEARCH_FIELDS`. If none of those - /// default names are present in the schema either, the result is None. - #[test] - fn test_generate_inverted_idx_recordbatch_default_fts_fallback_no_match() { - use arrow::array::{Int64Array, StringArray}; - - // Schema has only a custom column and _timestamp – none of the default - // FTS fields (log, message, msg, content, data, body, json, error). - let schema = Arc::new(Schema::new(vec![ - Field::new("custom_col", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["v1", "v2"])), - Arc::new(Int64Array::from(vec![1000i64, 2000i64])), - ], - ) - .unwrap(); - - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Logs, - &[], // empty → falls back to SQL_FULL_TEXT_SEARCH_FIELDS - &[], // empty index_fields too - ); - - assert!(result.is_ok()); - // None of the default FTS fields match, so inverted_idx_columns becomes - // empty after retain, and the function should return None. - assert!( - result.unwrap().is_none(), - "No default FTS field match should produce None" - ); - } - - /// When `full_text_search_fields` is empty but a default FTS field IS present - /// in the schema, the function should produce a valid RecordBatch. - #[test] - fn test_generate_inverted_idx_recordbatch_default_fts_fallback_with_match() { - use arrow::array::{Int64Array, StringArray}; - - // "log" is one of the default SQL_FULL_TEXT_SEARCH_FIELDS. - let schema = Arc::new(Schema::new(vec![ - Field::new("log", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["line1", "line2"])), - Arc::new(Int64Array::from(vec![1000i64, 2000i64])), - ], - ) - .unwrap(); - - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Logs, - &[], // empty → falls back to SQL_FULL_TEXT_SEARCH_FIELDS - &[], - ); - - assert!(result.is_ok()); - let rb = result.unwrap(); - assert!(rb.is_some(), "Default FTS field 'log' should produce Some"); - let rb = rb.unwrap(); - assert_eq!(rb.num_rows(), 2); - // Columns: "log" + "_timestamp" - assert_eq!(rb.num_columns(), 2); - } - - /// When `_timestamp` is explicitly listed in `full_text_search_fields`, it - /// must not be added a second time (dedup). The final batch should still - /// contain only ONE `_timestamp` column. - #[test] - fn test_generate_inverted_idx_recordbatch_timestamp_in_fts_fields_dedup() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("message", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["hello"])), - Arc::new(Int64Array::from(vec![5000i64])), - ], - ) - .unwrap(); - - // Explicitly include _timestamp in fts_fields to exercise the dedup path. - let full_text_search_fields = vec![ - "message".to_string(), - TIMESTAMP_COL_NAME.to_string(), // duplicate that should be deduped - ]; - - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Logs, - &full_text_search_fields, - &[], - ); - - assert!(result.is_ok()); - let rb = result.unwrap(); - assert!(rb.is_some()); - let rb = rb.unwrap(); - // Columns: message + _timestamp (no duplicate _timestamp column). - assert_eq!(rb.num_columns(), 2); - // Verify _timestamp column appears exactly once. - let ts_count = rb - .schema() - .fields() - .iter() - .filter(|f| f.name() == TIMESTAMP_COL_NAME) - .count(); - assert_eq!(ts_count, 1, "_timestamp must appear exactly once"); - } - - /// A single-element batch vector exercises the `len() == 1` fast-path that - /// avoids `concat_batches` and returns the batch directly. - #[test] - fn test_generate_inverted_idx_recordbatch_single_batch_fast_path() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("body", DataType::Utf8, true), - Field::new("level", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["body text"])), - Arc::new(StringArray::from(vec!["warn"])), - Arc::new(Int64Array::from(vec![9999i64])), - ], - ) - .unwrap(); - - // "body" is a default SQL_FULL_TEXT_SEARCH_FIELDS entry. - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], // exactly one batch - StreamType::Logs, - &["body".to_string()], - &["level".to_string()], - ); - - assert!(result.is_ok()); - let rb = result.unwrap(); - assert!(rb.is_some()); - let rb = rb.unwrap(); - assert_eq!(rb.num_rows(), 1); - // Columns: body + level + _timestamp (sorted, deduped) - assert_eq!(rb.num_columns(), 3); - } - - /// When the batch has columns but the selected index columns are missing - /// (because `inverted_idx_columns` only matches `_timestamp`), the result - /// must be None rather than a single-column batch. - #[test] - fn test_generate_inverted_idx_recordbatch_only_timestamp_after_filter() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("irrelevant", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["x"])), - Arc::new(Int64Array::from(vec![100i64])), - ], - ) - .unwrap(); - - // Ask for a field that IS in the schema but is the timestamp itself, - // plus another that does not exist. - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Logs, - &[TIMESTAMP_COL_NAME.to_string()], // only _timestamp passes retain - &[], - ); - - assert!(result.is_ok()); - // inverted_idx_columns after retain would contain _timestamp, then the - // "push _timestamp if absent" step is skipped (it's already there), so - // the final batch has exactly 1 column (_timestamp only) → returns None. - assert!( - result.unwrap().is_none(), - "Batch with only _timestamp column should return None" - ); - } - - /// Traces stream type supports indexing, so the function proceeds normally. - #[test] - fn test_generate_inverted_idx_recordbatch_traces_stream_type() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("msg", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["span-data"])), - Arc::new(Int64Array::from(vec![777i64])), - ], - ) - .unwrap(); - - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Traces, // supports indexing - &["msg".to_string()], - &[], - ); - - assert!(result.is_ok()); - let rb = result.unwrap(); - assert!( - rb.is_some(), - "Traces should produce Some when field matches" - ); - let rb = rb.unwrap(); - assert_eq!(rb.num_rows(), 1); - assert_eq!(rb.num_columns(), 2); // msg + _timestamp - } - - /// Verify the column order in the output: `inverted_idx_columns` is sorted - /// and deduped, so the schema fields of the produced batch must respect that - /// alphabetical ordering (excluding the appended `_timestamp`). - #[test] - fn test_generate_inverted_idx_recordbatch_column_sort_order() { - use arrow::array::{Int64Array, StringArray}; - - // Provide fields in reverse alphabetical order to confirm output is sorted. - let schema = Arc::new(Schema::new(vec![ - Field::new("zebra", DataType::Utf8, true), - Field::new("apple", DataType::Utf8, true), - Field::new("mango", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["z"])), - Arc::new(StringArray::from(vec!["a"])), - Arc::new(StringArray::from(vec!["m"])), - Arc::new(Int64Array::from(vec![1i64])), - ], - ) - .unwrap(); - - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Logs, - &[ - "zebra".to_string(), - "apple".to_string(), - "mango".to_string(), - ], - &[], - ); - - assert!(result.is_ok()); - let rb = result.unwrap().expect("should produce Some"); - let rb_schema = rb.schema(); - let field_names: Vec<&str> = rb_schema - .fields() - .iter() - .map(|f| f.name().as_str()) - .collect(); - // Expected sorted order: apple, mango, zebra, then _timestamp appended at end. - assert_eq!(field_names[0], "apple"); - assert_eq!(field_names[1], "mango"); - assert_eq!(field_names[2], "zebra"); - assert_eq!(field_names[3], TIMESTAMP_COL_NAME); - } - - /// Three-batch concat path: ensure all rows survive concatenation. - #[test] - fn test_generate_inverted_idx_recordbatch_three_batches_concat() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("error", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let make_batch = |msg: &str, ts: i64| { - RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec![msg])) as _, - Arc::new(Int64Array::from(vec![ts])) as _, - ], - ) - .unwrap() - }; - - let batches = vec![ - make_batch("err1", 100), - make_batch("err2", 200), - make_batch("err3", 300), - ]; - - let result = generate_inverted_idx_recordbatch( - schema, - &batches, - StreamType::Logs, - &["error".to_string()], - &[], - ); - - assert!(result.is_ok()); - let rb = result.unwrap().expect("should produce Some"); - // "error" is also a default FTS field, but because we supply fts_fields - // explicitly the result should include it plus _timestamp. - assert_eq!(rb.num_rows(), 3); // one row per batch, concatenated - assert_eq!(rb.num_columns(), 2); // error + _timestamp - } - // ── sort_by_time_range: additional branch-coverage tests ───────────────── /// Two files that don't overlap and have a gap between them. The second @@ -2591,400 +1716,6 @@ mod tests { assert!(keys.contains(&"E.parquet")); } - // ── generate_inverted_idx_recordbatch: additional branch-coverage tests ─── - - /// StreamType::Metadata supports indexing (same as Logs/Traces). - /// A valid FTS field in the schema must produce Some. - #[test] - fn test_generate_inverted_idx_recordbatch_metadata_stream_type() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("message", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["meta-event"])), - Arc::new(Int64Array::from(vec![1234i64])), - ], - ) - .unwrap(); - - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Metadata, // supports indexing - &["message".to_string()], - &[], - ); - - assert!(result.is_ok()); - let rb = result.unwrap(); - assert!( - rb.is_some(), - "Metadata stream should produce Some when field matches" - ); - let rb = rb.unwrap(); - assert_eq!(rb.num_rows(), 1); - assert_eq!(rb.num_columns(), 2); // message + _timestamp - } - - /// StreamType::Metrics supports indexing. Providing explicit FTS fields - /// that exist in the schema must produce a valid RecordBatch (not None). - #[test] - fn test_generate_inverted_idx_recordbatch_metrics_stream_produces_some() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("label", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["cpu_usage"])), - Arc::new(Int64Array::from(vec![9000i64])), - ], - ) - .unwrap(); - - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Metrics, // also supports indexing - &["label".to_string()], - &[], - ); - - assert!(result.is_ok()); - let rb = result.unwrap(); - assert!( - rb.is_some(), - "Metrics supports indexing, so a matching field should produce Some" - ); - let rb = rb.unwrap(); - assert_eq!(rb.num_rows(), 1); - assert_eq!(rb.num_columns(), 2); // label + _timestamp - } - - /// A batch where the batch's schema has the requested field but the column - /// name has different casing — `index_of` is case-sensitive, so it will - /// fail to find the column and produce an empty selection → result is None. - #[test] - fn test_generate_inverted_idx_recordbatch_case_sensitive_field_miss() { - use arrow::array::{Int64Array, StringArray}; - - // Schema has "Message" (capital M), but we request "message" (lowercase). - let schema = Arc::new(Schema::new(vec![ - Field::new("Message", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["hello"])), - Arc::new(Int64Array::from(vec![100i64])), - ], - ) - .unwrap(); - - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Logs, - &["message".to_string()], // lowercase — won't match "Message" - &[], - ); - - assert!(result.is_ok()); - // "message" is not found in schema fields (case-sensitive) → inverted_idx_columns - // becomes empty after retain → returns None early. - assert!( - result.unwrap().is_none(), - "Case mismatch should prevent field from being retained" - ); - } - - /// A batch with zero rows (but correct schema) must still produce a valid - /// (albeit empty) RecordBatch when the fields match. - #[test] - fn test_generate_inverted_idx_recordbatch_zero_row_batch() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("content", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - // Explicitly build a zero-row batch. - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(Vec::<&str>::new())), - Arc::new(Int64Array::from(Vec::::new())), - ], - ) - .unwrap(); - assert_eq!(batch.num_rows(), 0); - - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Logs, - &["content".to_string()], - &[], - ); - - assert!(result.is_ok()); - let rb = result.unwrap(); - // The function does not check num_rows, only column count. - // "content" + "_timestamp" gives 2 columns → Some returned. - assert!( - rb.is_some(), - "Zero-row batch with valid fields should produce Some" - ); - let rb = rb.unwrap(); - assert_eq!(rb.num_rows(), 0); - assert_eq!(rb.num_columns(), 2); // content + _timestamp - } - - /// Providing only `index_fields` (no FTS fields) for a Logs stream should - /// work: the columns are selected from the batch schema and returned. - #[test] - fn test_generate_inverted_idx_recordbatch_index_fields_only_logs() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("request_id", DataType::Utf8, true), - Field::new("status_code", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["req-1", "req-2"])), - Arc::new(StringArray::from(vec!["200", "500"])), - Arc::new(Int64Array::from(vec![1000i64, 2000i64])), - ], - ) - .unwrap(); - - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Logs, - &[], // no FTS fields → falls back to SQL defaults; none match - &["request_id".to_string(), "status_code".to_string()], - ); - - assert!(result.is_ok()); - let rb = result.unwrap(); - assert!(rb.is_some(), "index_fields present should produce Some"); - let rb = rb.unwrap(); - assert_eq!(rb.num_rows(), 2); - // request_id + status_code + _timestamp = 3 columns - assert_eq!(rb.num_columns(), 3); - } - - /// Confirm that when `inverted_idx_columns` would be empty after the retain - /// step (no schema field matches at all), the function returns None before - /// entering the per-batch loop — exercising the early-return branch. - #[test] - fn test_generate_inverted_idx_recordbatch_empty_columns_early_return() { - use arrow::array::{Int64Array, StringArray}; - - // Schema has only fields that are NOT in SQL_FULL_TEXT_SEARCH_FIELDS and - // are not supplied as FTS or index fields. - let schema = Arc::new(Schema::new(vec![ - Field::new("alpha", DataType::Utf8, true), - Field::new("beta", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["a"])), - Arc::new(StringArray::from(vec!["b"])), - Arc::new(Int64Array::from(vec![1i64])), - ], - ) - .unwrap(); - - // Request fields that don't exist in the schema at all. - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Logs, - &["nonexistent_fts".to_string()], - &["nonexistent_idx".to_string()], - ); - - assert!(result.is_ok()); - assert!( - result.unwrap().is_none(), - "No schema field match should yield None (early return)" - ); - } - - /// When the same field name appears in both `full_text_search_fields` AND - /// `index_fields`, `dedup` must collapse it so the output batch column - /// appears exactly once. - #[test] - fn test_generate_inverted_idx_recordbatch_fts_and_index_overlap_dedup() { - use arrow::array::{Int64Array, StringArray}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("tag", DataType::Utf8, true), - Field::new("host", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["web", "db"])), - Arc::new(StringArray::from(vec!["host-1", "host-2"])), - Arc::new(Int64Array::from(vec![100i64, 200i64])), - ], - ) - .unwrap(); - - // "tag" appears in both lists; "host" only in index_fields. - let result = generate_inverted_idx_recordbatch( - schema, - &[batch], - StreamType::Logs, - &["tag".to_string()], - &["tag".to_string(), "host".to_string()], - ); - - assert!(result.is_ok()); - let rb = result.unwrap().expect("should produce Some"); - // After dedup: [host, tag] (sorted) + _timestamp = 3 columns. - assert_eq!(rb.num_columns(), 3); - // "tag" must appear exactly once. - let tag_count = rb - .schema() - .fields() - .iter() - .filter(|f| f.name() == "tag") - .count(); - assert_eq!(tag_count, 1, "tag must be deduped to exactly one column"); - assert_eq!(rb.num_rows(), 2); - } - - /// When the batch's own schema is a subset of the inverted_idx_columns list - /// (some requested columns are absent from the batch), the missing columns - /// are silently skipped via `filter_map`. The result must only contain - /// the fields that the batch actually has. - #[test] - fn test_generate_inverted_idx_recordbatch_partial_batch_schema_match() { - use arrow::array::{Int64Array, StringArray}; - - // Stream schema knows about "present" and "absent". - let stream_schema = Arc::new(Schema::new(vec![ - Field::new("present", DataType::Utf8, true), - Field::new("absent", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - // But the actual batch only has "present" and _timestamp. - let batch_schema = Arc::new(Schema::new(vec![ - Field::new("present", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - let batch = RecordBatch::try_new( - batch_schema, - vec![ - Arc::new(StringArray::from(vec!["here"])), - Arc::new(Int64Array::from(vec![42i64])), - ], - ) - .unwrap(); - - // Pass the stream schema (which includes "absent") but a batch that - // doesn't have it — exercises the filter_map miss path. - let result = generate_inverted_idx_recordbatch( - stream_schema, - &[batch], - StreamType::Logs, - &["present".to_string(), "absent".to_string()], - &[], - ); - - assert!(result.is_ok()); - let rb = result - .unwrap() - .expect("should produce Some because 'present' matched"); - // Only "present" + _timestamp survive (absent was not in the batch schema). - assert_eq!(rb.num_columns(), 2); - let rb_schema = rb.schema(); - let col_names: Vec<&str> = rb_schema - .fields() - .iter() - .map(|f| f.name().as_str()) - .collect(); - assert!(col_names.contains(&"present")); - assert!(col_names.contains(&TIMESTAMP_COL_NAME)); - assert!(!col_names.contains(&"absent")); - } - - /// Two batches with different schemas (first has "msg", second has "body"). - /// Because `inverted_idx_columns` is built from the *stream* schema and both - /// fields are present there, but the per-batch `filter_map` only picks what - /// each batch actually has. After concat the rows from each batch survive. - #[test] - fn test_generate_inverted_idx_recordbatch_two_batches_different_schemas() { - use arrow::array::{Int64Array, StringArray}; - - // Stream schema knows about both columns. - let stream_schema = Arc::new(Schema::new(vec![ - Field::new("msg", DataType::Utf8, true), - Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), - ])); - - // Both batches use the same schema here for simplicity (the column- - // miss path is already tested above); this exercises the two-batch concat. - let batch1 = RecordBatch::try_new( - stream_schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["first"])), - Arc::new(Int64Array::from(vec![10i64])), - ], - ) - .unwrap(); - - let batch2 = RecordBatch::try_new( - stream_schema.clone(), - vec![ - Arc::new(StringArray::from(vec!["second", "third"])), - Arc::new(Int64Array::from(vec![20i64, 30i64])), - ], - ) - .unwrap(); - - let result = generate_inverted_idx_recordbatch( - stream_schema, - &[batch1, batch2], - StreamType::Logs, - &["msg".to_string()], - &[], - ); - - assert!(result.is_ok()); - let rb = result.unwrap().expect("should produce Some"); - assert_eq!(rb.num_rows(), 3); // 1 + 2 - assert_eq!(rb.num_columns(), 2); // msg + _timestamp - } - /// Verify `sort_by_time_range` with exactly two files where the second /// file's min_ts equals the first file's max_ts — the boundary condition /// `min_ts >= max_ts` in the group-fit predicate. diff --git a/src/service/search/grpc/storage.rs b/src/service/search/grpc/storage.rs index f09b2d987b..0903c6bae8 100644 --- a/src/service/search/grpc/storage.rs +++ b/src/service/search/grpc/storage.rs @@ -28,7 +28,7 @@ use config::{ }, metrics::{self, QUERY_PARQUET_CACHE_RATIO_NODE}, utils::{ - inverted_index::convert_parquet_file_name_to_tantivy_file, + inverted_index::to_tantivy_name, size::bytes_to_human_readable, tantivy::tokenizer::{CollectType, O2_TOKENIZER, o2_tokenizer_build}, time::BASE_TIME, @@ -47,6 +47,12 @@ use tantivy::{ Directory, Term, query::{BooleanQuery, Occur, Query, RangeQuery}, }; +use tantivy_utils::puffin_directory::{ + PROP_ROW_GROUP_SIZE, + caching_directory::CachingDirectory, + footer_cache::FooterCache, + reader::{PuffinDirReader, warm_up_terms}, +}; use tokio::sync::Semaphore; use tokio_stream::StreamExt as _; use tracing::Instrument; @@ -62,12 +68,6 @@ use crate::service::{ index::IndexCondition, inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, }, - tantivy::puffin_directory::{ - PROP_ROW_GROUP_SIZE, - caching_directory::CachingDirectory, - footer_cache::FooterCache, - reader::{PuffinDirReader, warm_up_terms}, - }, }; /// search in remote object storage @@ -518,8 +518,7 @@ pub async fn tantivy_search( .filter_map(|(_, f)| { scan_stats.compressed_size += f.meta.index_size; if f.meta.index_size > 0 { - convert_parquet_file_name_to_tantivy_file(&f.key) - .map(|ttv_file| (ttv_file, f.clone())) + to_tantivy_name(&f.key).map(|ttv_file| (ttv_file, f.clone())) } else { None } @@ -808,7 +807,7 @@ async fn search_tantivy_index( parquet_file: &FileKey, ) -> anyhow::Result<(String, TantivyResult, bool)> { let file_account = parquet_file.account.clone(); - let Some(ttv_file_name) = convert_parquet_file_name_to_tantivy_file(&parquet_file.key) else { + let Some(ttv_file_name) = to_tantivy_name(&parquet_file.key) else { return Err(anyhow::anyhow!( "[trace_id {trace_id}] search->storage: Unable to find tantivy index files for parquet file {}", parquet_file.key.clone() diff --git a/src/service/tantivy/bloom_builder.rs b/src/service/tantivy/bloom_builder.rs index f9f1b3fee1..8a334ff6c4 100644 --- a/src/service/tantivy/bloom_builder.rs +++ b/src/service/tantivy/bloom_builder.rs @@ -128,9 +128,10 @@ mod tests { use std::sync::Arc; use arrow::{ - array::{RecordBatch, StringArray}, + array::{Int64Array, RecordBatch, StringArray}, datatypes::{DataType, Field, Schema}, }; + use config::TIMESTAMP_COL_NAME; use tantivy::directory::RamDirectory; use super::*; @@ -144,6 +145,7 @@ mod tests { #[tokio::test] async fn test_build_blooms_round_trip() { let schema = Arc::new(Schema::new(vec![ + Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), Field::new("trace_id", DataType::Utf8, false), Field::new("user_id", DataType::Utf8, false), Field::new("level", DataType::Utf8, false), @@ -155,9 +157,11 @@ mod tests { let levels = (0..100) .map(|i| if i % 2 == 0 { "info" } else { "error" }) .collect::>(); + let timestamps: Vec = (0..100).map(|i| 1000 + i).collect(); let batch = RecordBatch::try_new( schema.clone(), vec![ + Arc::new(Int64Array::from(timestamps)), Arc::new(StringArray::from(trace_ids.clone())), Arc::new(StringArray::from(user_ids.clone())), Arc::new(StringArray::from(levels.clone())), @@ -254,10 +258,18 @@ mod tests { async fn test_empty_field_list_returns_empty() { // No fields requested → no work, no error. let dir = RamDirectory::create(); - let schema = Arc::new(Schema::new(vec![Field::new("f", DataType::Utf8, false)])); - let batch = - RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(vec!["x"]))]) - .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), + Field::new("f", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1000i64])), + Arc::new(StringArray::from(vec!["x"])), + ], + ) + .unwrap(); let buf = create_test_parquet_bytes(vec![batch.clone()]).await; let index = build_index( dir, diff --git a/src/service/tantivy/mod.rs b/src/service/tantivy/mod.rs index 9141384dc2..75085927ab 100644 --- a/src/service/tantivy/mod.rs +++ b/src/service/tantivy/mod.rs @@ -15,8 +15,7 @@ pub mod bloom_builder; mod parallel; -pub mod puffin; -pub mod puffin_directory; +mod reader; mod sequential; use std::{collections::HashMap, sync::Arc}; @@ -29,14 +28,11 @@ use bytes::Bytes; use 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, tantivy::tokenizer::O2_TOKENIZER, - }, + utils::{inverted_index::to_tantivy_name, tantivy::tokenizer::O2_TOKENIZER}, }; use hashbrown::HashSet; use infra::storage; -use puffin_directory::writer::PuffinDirWriter; -use tantivy_utils::puffin_directory::PROP_ROW_GROUP_SIZE; +use tantivy_utils::puffin_directory::{PROP_ROW_GROUP_SIZE, writer::PuffinDirWriter}; pub(crate) async fn create_tantivy_index( caller: &str, @@ -58,8 +54,8 @@ pub(crate) async fn create_tantivy_index( }; let dir = PuffinDirWriter::new(); - let index = if thread_num > 1 && matches!(file_format, FileFormat::Parquet) { - parallel::build_index(dir.clone(), buf, index_schema, thread_num).await? + let index = if thread_num > 1 { + parallel::build_index(dir.clone(), file_format, buf, index_schema, thread_num).await? } else { sequential::build_index(dir.clone(), file_format, buf, index_schema).await? }; @@ -74,7 +70,7 @@ pub(crate) async fn create_tantivy_index( let index_size = puffin_bytes.len(); // write fst bytes into disk - let Some(idx_file_name) = convert_parquet_file_name_to_tantivy_file(parquet_file_name) else { + let Some(idx_file_name) = to_tantivy_name(parquet_file_name) else { return Ok(0); }; @@ -139,10 +135,9 @@ fn build_tantivy_schema( let fts_fields_filtered = fts_fields .iter() .filter(|f| { - schema_fields - .get(f) - .map(|v| v.data_type() == &DataType::Utf8 || v.data_type() == &DataType::LargeUtf8) - .is_some() + schema_fields.get(f).is_some_and(|v| { + v.data_type() == &DataType::Utf8 || v.data_type() == &DataType::LargeUtf8 + }) }) .map(String::from) .collect::>(); @@ -151,10 +146,7 @@ fn build_tantivy_schema( .filter(|f| schema_fields.contains_key(f)) .map(String::from) .collect::>(); - let tantivy_fields = fts_fields_filtered - .union(&index_fields_filtered) - .cloned() - .collect::>(); + let tantivy_fields: HashSet<_> = &fts_fields_filtered | &index_fields_filtered; if tantivy_fields.is_empty() { return None; @@ -264,14 +256,10 @@ pub(super) fn convert_batch_to_docs_sync( } } - // _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_data: &Int64Array = batch + .column_by_name(TIMESTAMP_COL_NAME) + .and_then(|c| c.as_any().downcast_ref::()) + .expect("batch must contain a valid Int64 _timestamp column"); 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)); diff --git a/src/service/tantivy/parallel.rs b/src/service/tantivy/parallel.rs index 6a8236a468..b04a5c40ef 100644 --- a/src/service/tantivy/parallel.rs +++ b/src/service/tantivy/parallel.rs @@ -13,141 +13,204 @@ // 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::{ - TIMESTAMP_COL_NAME, + FileFormat, TIMESTAMP_COL_NAME, utils::tantivy::tokenizer::{CollectType, O2_TOKENIZER, o2_tokenizer_build}, }; -use parquet::arrow::{ProjectionMask, arrow_reader::ParquetRecordBatchReaderBuilder}; -use tantivy::directory::MmapDirectory; +use futures::future::join_all; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use tantivy::{directory::MmapDirectory, indexer::merge_indices}; use tokio::task::JoinHandle; +#[cfg(all(feature = "vortex", feature = "enterprise"))] +use { + config::PARQUET_MAX_ROW_GROUP_SIZE, + vortex::{ + VortexSessionDefault, + file::OpenOptionsSessionExt, + io::{ + runtime::{BlockingRuntime, single::SingleThreadRuntime}, + session::RuntimeSessionExt, + }, + session::VortexSession, + }, +}; use super::{TantivyIndexSchema, convert_batch_to_docs_sync}; +use crate::service::tantivy::reader::{ChunkSelector, chunk_iter}; -type TaskOutput = (usize, tantivy::Index, usize); +pub(super) async fn build_index( + tantivy_dir: D, + file_format: FileFormat, + buf: Bytes, + index_schema: TantivyIndexSchema, + thread_num: usize, +) -> Result, Error> { + match file_format { + FileFormat::Parquet => { + let buf_meta = buf.clone(); + let num_row_groups = tokio::task::spawn_blocking(move || -> Result { + ParquetRecordBatchReaderBuilder::try_new(buf_meta) + .context("failed to open parquet for metadata") + .map(|b| b.metadata().num_row_groups()) + }) + .await??; -/// Row-group-direct parallel tantivy index builder. + build_parallel_index( + tantivy_dir, + buf, + index_schema, + thread_num, + num_row_groups, + ChunkSelector::Parquet, + ) + .await + } + #[cfg(all(feature = "vortex", feature = "enterprise"))] + FileFormat::Vortex => { + let buf_meta = buf.clone(); + let row_count: u64 = tokio::task::spawn_blocking(move || -> Result { + let runtime = SingleThreadRuntime::default(); + let session = VortexSession::default().with_handle(runtime.handle()); + session + .open_options() + .open_buffer(buf_meta) + .context("failed to open vortex file for row_count") + .map(|vxf| vxf.row_count()) + }) + .await??; + + let chunk_size = PARQUET_MAX_ROW_GROUP_SIZE as u64; + let num_chunks = row_count.div_ceil(chunk_size) as usize; + build_parallel_index( + tantivy_dir, + buf, + index_schema, + thread_num, + num_chunks, + move |chunk_idx| { + let chunk_start = chunk_idx as u64 * chunk_size; + let chunk_end = std::cmp::min(chunk_start + chunk_size, row_count); + ChunkSelector::Vortex(chunk_start..chunk_end) + }, + ) + .await + } + #[cfg(not(all(feature = "vortex", feature = "enterprise")))] + FileFormat::Vortex => Err(anyhow::anyhow!( + "Vortex file format requires the vortex feature" + )), + } +} + +struct SegmentOutput { + chunk_idx: usize, + index: tantivy::Index, + row_count: usize, +} + +/// Chunk-parallel tantivy index builder shared by the Parquet and Vortex paths. /// /// ```text -/// parquet (Bytes) +/// Bytes (reference-counted — each worker gets one cheap clone) /// │ -/// ├─ open ParquetRecordBatchReaderBuilder ── read metadata ── num_row_groups +/// │ build_index dispatches by format: +/// │ ┌─ Parquet ── read metadata ─── num_row_groups → chunks +/// │ └─ Vortex ── read row_count ── ÷ PARQUET_MAX_ROW_GROUP_SIZE → chunks /// │ /// ▼ spawn_blocking pool (concurrency = workers) /// ┌──────────────────────┬──────────────────────┬─────────────────────┐ -/// │ rg_0 │ rg_1 │ ... │ -/// │ builder │ builder │ │ -/// │ .with_row_groups │ .with_row_groups │ │ -/// │ ([0]).build() │ ([1]).build() │ │ +/// │ chunk_0 │ chunk_1 │ ... │ +/// │ make_selector(0) │ make_selector(1) │ │ +/// │ → ChunkSelector │ → ChunkSelector │ │ +/// │ → chunk_iter │ → chunk_iter │ │ /// │ → SingleSegmentIdx │ → SingleSegmentIdx │ │ /// │ → finalize → Index │ → finalize → Index │ │ /// └──────────────────────┴──────────────────────┴─────────────────────┘ /// │ -/// ▼ sort by rg_idx ASC +/// ▼ sort by chunk_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( +async fn build_parallel_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 { + num_chunks: usize, + make_selector: F, +) -> Result, Error> +where + D: tantivy::Directory + Send + Sync + 'static, + F: Fn(usize) -> ChunkSelector, +{ + if num_chunks == 0 { return Ok(None); } - // single row_group fast path - if num_row_groups == 1 { + // single chunk fast path — write directly into the caller-supplied dir, no merge + if num_chunks == 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) + let selector = make_selector(0); + let SegmentOutput { + index, row_count, .. + } = tokio::task::spawn_blocking(move || { + build_segment(0, buf, selector, index_schema, tantivy_dir) }) .await??; log::info!( - "parallel::build_index: single row_group fast path, rows={num_rows}, elapsed={} ms", + "parallel::build_index: single chunk, rows={row_count}, elapsed={} ms", start.elapsed().as_millis() ); return Ok(Some(index)); } - let thread_num = thread_num.max(1).min(num_row_groups); + let thread_num = thread_num.max(1).min(num_chunks); 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 mut tasks: Vec>> = Vec::with_capacity(num_chunks); + for chunk_idx in 0..num_chunks { let permit = semaphore.clone().acquire_owned().await?; let buf_c = buf.clone(); let index_schema_c = index_schema.clone(); + let selector = make_selector(chunk_idx); 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) + build_segment(chunk_idx, buf_c, selector, 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(); + // await all tasks concurrently, then sort to restore chunk order. + let mut indexed: Vec = join_all(tasks) + .await + .into_iter() + .map(|r| -> Result { r? }) + .collect::>()?; + indexed.sort_by_key(|s| s.chunk_idx); + let total_rows: usize = indexed.iter().map(|s| s.row_count).sum(); + let indices: Vec = indexed.into_iter().map(|s| s.index).collect(); log::info!( - "parallel::build_index: built {num_row_groups} row_groups, total_rows={total_num_rows}, thread_num={thread_num}, elapsed={} ms", + "parallel::build_index: built {num_chunks} chunks, rows={total_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) + 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", + "parallel::build_index: merged {num_chunks} segments, elapsed={} ms", start.elapsed().as_millis() ); @@ -159,44 +222,29 @@ pub(super) async fn build_index( )); } let merged_rows = merged_segs[0].meta().max_doc() as usize; - if merged_rows != total_num_rows { + if merged_rows != total_rows { return Err(anyhow::anyhow!( - "merged tantivy index row count mismatch: tantivy={merged_rows}, parquet={total_num_rows}", + "merged tantivy index row count mismatch: tantivy={merged_rows}, chunks={total_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, +/// Synchronously builds one tantivy segment from the rows described by +/// `selector`. Used by both the parquet and vortex parallel paths. +fn build_segment( + chunk_idx: usize, buf: Bytes, + selector: ChunkSelector, index_schema: TantivyIndexSchema, dir: D, -) -> Result<(usize, tantivy::Index, usize), Error> { - let builder = ParquetRecordBatchReaderBuilder::try_new(buf)?; +) -> Result { + let mut projection: Vec = index_schema.fields.iter().cloned().collect(); + projection.push(TIMESTAMP_COL_NAME.to_string()); - let arrow_schema = builder.schema(); - let projection_indices: Vec = arrow_schema - .fields() - .iter() - .enumerate() - .filter_map(|(i, f)| { - (f.name() == TIMESTAMP_COL_NAME || index_schema.fields.contains(f.name())).then_some(i) - }) - .collect(); - let projection_mask = ProjectionMask::roots( - builder.metadata().file_metadata().schema_descr(), - projection_indices, - ); - - let reader = builder - .with_projection(projection_mask) - .with_row_groups(vec![rg_idx]) - .build()?; + let iter = chunk_iter(selector, buf, Some(&projection)) + .with_context(|| format!("chunk {chunk_idx}: reader build failed"))?; let tokenizer_manager = tantivy::tokenizer::TokenizerManager::default(); tokenizer_manager.register(O2_TOKENIZER, o2_tokenizer_build(CollectType::Ingest)); @@ -205,13 +253,13 @@ fn build_row_group_segment( .tokenizers(tokenizer_manager) .single_segment_index_writer::(dir, 50_000_000)?; - let mut rg_rows: usize = 0; - for batch_res in reader { + let mut row_count: usize = 0; + for batch_res in iter { let batch = batch_res?; if batch.num_rows() == 0 { continue; } - rg_rows += batch.num_rows(); + row_count += batch.num_rows(); let docs = convert_batch_to_docs_sync(&batch, &index_schema); for doc in docs { writer.add_document(doc)?; @@ -220,9 +268,13 @@ fn build_row_group_segment( let index = writer .finalize() - .map_err(|e| anyhow::anyhow!("worker {rg_idx}: finalize failed: {e}"))?; + .map_err(|e| anyhow::anyhow!("chunk {chunk_idx}: finalize failed: {e}"))?; - Ok((rg_idx, index, rg_rows)) + Ok(SegmentOutput { + chunk_idx, + index, + row_count, + }) } #[cfg(test)] @@ -235,10 +287,12 @@ mod tests { record_batch::RecordBatch, }; use config::{FileFormat, TIMESTAMP_COL_NAME}; - use parquet::{arrow::AsyncArrowWriter, file::properties::WriterProperties}; + use parquet::{ + arrow::{AsyncArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder}, + file::properties::WriterProperties, + }; use tantivy::directory::RamDirectory; - use super::*; use crate::service::tantivy::{ sequential, tests::{create_test_parquet_bytes, make_index_schema}, @@ -332,11 +386,11 @@ mod tests { Field::new("marker", DataType::Utf8, false), ])); - let index = build_index( + let index = super::build_index( RamDirectory::create(), + FileFormat::Parquet, buf, make_index_schema(&[], &["marker".to_string()], &stream_schema), - // workers 3, // force out-of-order completion ) .await @@ -391,8 +445,9 @@ mod tests { Field::new("marker", DataType::Utf8, false), ])); - let index = build_index( + let index = super::build_index( RamDirectory::create(), + FileFormat::Parquet, buf, make_index_schema(&[], &["marker".to_string()], &stream_schema), 4, @@ -444,8 +499,9 @@ mod tests { Field::new("marker", DataType::Utf8, false), ])); - let par = build_index( + let par = super::build_index( RamDirectory::create(), + FileFormat::Parquet, buf_par, make_index_schema(&[], &["marker".to_string()], &schema), 3, @@ -509,4 +565,207 @@ mod tests { assert_eq!(addr_seq.doc_id as usize, row); } } + + #[cfg(all(feature = "vortex", feature = "enterprise"))] + mod vortex_tests { + use std::sync::Arc; + + use arrow::{ + array::{Int64Array, StringArray}, + datatypes::{DataType, Field, Schema}, + record_batch::RecordBatch, + }; + use bytes::Bytes; + use config::{FileFormat, PARQUET_MAX_ROW_GROUP_SIZE, TIMESTAMP_COL_NAME}; + use tantivy::directory::RamDirectory; + use vortex::{ + VortexSessionDefault, + array::{ArrayRef, arrow::FromArrowArray}, + dtype::{DType, arrow::FromArrowType}, + file::{VortexWriteOptions, WriteStrategyBuilder}, + io::session::RuntimeSessionExt, + session::VortexSession, + }; + + const CHUNK_ROWS: u64 = PARQUET_MAX_ROW_GROUP_SIZE as u64; + + use crate::service::tantivy::tests::make_index_schema; + + /// One row per marker, marker text is globally unique → search-by-marker + /// recovers the row's doc_id so we can assert it equals the 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 into an in-memory vortex file using the default write + /// strategy. Test helper — production writes go through `write_vortex` in + /// the enterprise crate. + async fn create_test_vortex_bytes(batches: Vec) -> Bytes { + let schema = batches[0].schema(); + let session = VortexSession::default().with_tokio(); + let dtype = DType::from_arrow(schema.as_ref()); + let write_options = VortexWriteOptions::new(session) + .with_strategy(WriteStrategyBuilder::default().build()); + let mut buf = Vec::new(); + let mut writer = write_options.writer(&mut buf, dtype); + for batch in batches { + let array: ArrayRef = ArrayRef::from_arrow(batch, false).unwrap(); + writer.push(array).await.unwrap(); + } + writer.finish().await.unwrap(); + Bytes::from(buf) + } + + /// Core invariant test: parallel chunks (possibly out of order) must yield + /// `doc_id == global_row_index` after merge. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_vortex_parallel_preserves_row_order() { + // Total rows that span ≥ 2 chunks so the merge path runs. CHUNK_ROWS + // is PARQUET_MAX_ROW_GROUP_SIZE (128k) — using a small multiple keeps + // the test fast while still exercising multi-chunk merge. + let total_rows = (CHUNK_ROWS as usize) * 2 + 137; + let buf = create_test_vortex_bytes(vec![create_marker_batch(0, total_rows)]).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 = super::super::build_index( + RamDirectory::create(), + FileFormat::Vortex, + buf, + make_index_schema(&[], &["marker".to_string()], &stream_schema), + 3, // force out-of-order chunk completion + ) + .await + .expect("vortex parallel build must succeed") + .expect("vortex 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(); + + // Spot-check evenly along the range — full scan of 256k+ rows is too + // slow for a unit test. + let probes = [ + 0, + 1, + CHUNK_ROWS as usize - 1, + CHUNK_ROWS as usize, + CHUNK_ROWS as usize * 2 - 1, + CHUNK_ROWS as usize * 2, + total_rows - 1, + ]; + for row in probes { + 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}" + ); + } + } + + /// Total rows < CHUNK_ROWS ⇒ single-chunk fast path that bypasses + /// `merge_indices` and writes the segment straight into the caller-supplied + /// directory. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_vortex_parallel_single_chunk_fast_path() { + let total_rows = 1234usize; + let buf = create_test_vortex_bytes(vec![create_marker_batch(0, total_rows)]).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 = super::super::build_index( + RamDirectory::create(), + FileFormat::Vortex, + buf, + make_index_schema(&[], &["marker".to_string()], &stream_schema), + 4, + ) + .await + .expect("single chunk build must succeed") + .expect("single chunk 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); + } + } + + /// Zero-row file ⇒ Ok(None), no panic. + #[tokio::test] + async fn test_vortex_parallel_empty_file() { + let empty = create_marker_batch(0, 0); + let buf = create_test_vortex_bytes(vec![empty.clone()]).await; + let stream_schema = Arc::new(Schema::new(vec![ + Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), + Field::new("marker", DataType::Utf8, false), + ])); + let result = super::super::build_index( + RamDirectory::create(), + FileFormat::Vortex, + buf, + make_index_schema(&[], &["marker".to_string()], &stream_schema), + 2, + ) + .await + .expect("empty build must succeed"); + assert!(result.is_none()); + } + } } diff --git a/src/service/tantivy/puffin/mod.rs b/src/service/tantivy/puffin/mod.rs deleted file mode 100644 index 9f059f51e7..0000000000 --- a/src/service/tantivy/puffin/mod.rs +++ /dev/null @@ -1,24 +0,0 @@ -// 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 . - -pub use tantivy_utils::puffin::*; - -pub mod reader { - pub use tantivy_utils::puffin::reader::*; -} - -pub mod writer { - pub use tantivy_utils::puffin::writer::*; -} diff --git a/src/service/tantivy/puffin_directory/mod.rs b/src/service/tantivy/puffin_directory/mod.rs deleted file mode 100644 index 6aab35c7a1..0000000000 --- a/src/service/tantivy/puffin_directory/mod.rs +++ /dev/null @@ -1,32 +0,0 @@ -// 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 . - -pub use tantivy_utils::puffin_directory::*; - -pub mod caching_directory { - pub use tantivy_utils::puffin_directory::caching_directory::*; -} - -pub mod footer_cache { - pub use tantivy_utils::puffin_directory::footer_cache::*; -} - -pub mod reader { - pub use tantivy_utils::puffin_directory::reader::*; -} - -pub mod writer { - pub use tantivy_utils::puffin_directory::writer::*; -} diff --git a/src/service/tantivy/reader/mod.rs b/src/service/tantivy/reader/mod.rs new file mode 100644 index 0000000000..e5270c5ad7 --- /dev/null +++ b/src/service/tantivy/reader/mod.rs @@ -0,0 +1,68 @@ +// 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 . + +mod parquet; +#[cfg(all(feature = "vortex", feature = "enterprise"))] +mod vortex; + +#[cfg(all(feature = "vortex", feature = "enterprise"))] +use std::ops::Range; + +use arrow::{error::ArrowError, record_batch::RecordBatch}; +use bytes::Bytes; +use config::{FileFormat, utils::parquet::RecordBatchStream}; + +pub(super) type RecordBatchIter = Box>>; + +pub(super) enum ChunkSelector { + Parquet(usize), // row_group id + #[cfg(all(feature = "vortex", feature = "enterprise"))] + Vortex(Range), // row range +} + +/// Read a whole file as a record batch stream restricted to `projection` +/// columns. `None` reads every column. +pub(super) async fn file_stream( + file_format: FileFormat, + data: Bytes, + projection: Option<&[String]>, +) -> Result { + match file_format { + FileFormat::Parquet => parquet::scan_parquet_async(data, projection).await, + #[cfg(all(feature = "vortex", feature = "enterprise"))] + FileFormat::Vortex => vortex::scan_vortex_async(data, projection).await, + #[cfg(not(all(feature = "vortex", feature = "enterprise")))] + FileFormat::Vortex => Err(anyhow::anyhow!( + "Vortex file format requires the vortex feature" + )), + } +} + +/// Sync chunk reader. Returns a sync `Iterator`. +pub(super) fn chunk_iter( + selector: ChunkSelector, + data: Bytes, + projection: Option<&[String]>, +) -> Result { + match selector { + ChunkSelector::Parquet(row_group_id) => { + parquet::scan_parquet_row_group(data, projection, row_group_id) + } + #[cfg(all(feature = "vortex", feature = "enterprise"))] + ChunkSelector::Vortex(row_range) => { + vortex::scan_vortex_row_range(data, projection, row_range) + } + } +} diff --git a/src/service/tantivy/reader/parquet.rs b/src/service/tantivy/reader/parquet.rs new file mode 100644 index 0000000000..12fb24ca62 --- /dev/null +++ b/src/service/tantivy/reader/parquet.rs @@ -0,0 +1,85 @@ +// 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 std::{io::Cursor, sync::Arc}; + +use arrow::error::ArrowError; +use arrow_schema::Schema; +use bytes::Bytes; +use config::{get_batch_size, utils::parquet::RecordBatchStream}; +use futures::TryStreamExt; +use parquet::arrow::{ + ParquetRecordBatchStreamBuilder, ProjectionMask, arrow_reader::ParquetRecordBatchReaderBuilder, +}; + +use super::RecordBatchIter; + +/// Open a parquet file and return an async record batch stream over all rows. +pub(super) async fn scan_parquet_async( + data: Bytes, + projection: Option<&[String]>, +) -> Result { + let builder = ParquetRecordBatchStreamBuilder::new(Cursor::new(data)).await?; + let mut builder = builder.with_batch_size(get_batch_size()); + if let Some(mask) = parquet_projection(builder.schema(), builder.metadata(), projection)? { + builder = builder.with_projection(mask); + } + let reader = builder.build()?; + let stream: RecordBatchStream = Box::pin(reader.map_err(ArrowError::from)); + Ok(stream) +} + +/// Build a sync iterator over one parquet row group. +pub(super) fn scan_parquet_row_group( + data: Bytes, + projection: Option<&[String]>, + row_group_id: usize, +) -> Result { + let builder = ParquetRecordBatchReaderBuilder::try_new(data)?; + let mut builder = builder.with_batch_size(get_batch_size()); + if let Some(mask) = parquet_projection(builder.schema(), builder.metadata(), projection)? { + builder = builder.with_projection(mask); + } + let reader = builder.with_row_groups(vec![row_group_id]).build()?; + Ok(Box::new(reader)) +} + +/// Compute the parquet `ProjectionMask` for a given projection column list. +/// Returns `None` when all columns should be read. +pub(super) fn parquet_projection( + full_schema: &Arc, + metadata: &parquet::file::metadata::ParquetMetaData, + projection: Option<&[String]>, +) -> Result, anyhow::Error> { + let Some(cols) = projection else { + return Ok(None); + }; + + let kept: Vec = full_schema + .fields() + .iter() + .enumerate() + .filter_map(|(i, f)| cols.iter().any(|c| c == f.name()).then_some(i)) + .collect(); + + if kept.is_empty() { + return Ok(None); + } + + Ok(Some(ProjectionMask::roots( + metadata.file_metadata().schema_descr(), + kept, + ))) +} diff --git a/src/service/tantivy/reader/vortex.rs b/src/service/tantivy/reader/vortex.rs new file mode 100644 index 0000000000..eb699dc689 --- /dev/null +++ b/src/service/tantivy/reader/vortex.rs @@ -0,0 +1,122 @@ +// 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 +// 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 std::{ops::Range, sync::Arc}; + +use arrow::{datatypes::DataType, error::ArrowError, record_batch::RecordBatch}; +use arrow_schema::Schema; +use bytes::Bytes; +use config::utils::parquet::{RecordBatchStream, vortex_array_to_record_batch}; +use futures::StreamExt; +use vortex::{ + VortexSessionDefault, + array::ArrayRef, + expr::{root, select}, + file::{OpenOptionsSessionExt, VortexFile}, + io::{ + runtime::{BlockingRuntime, single::SingleThreadRuntime}, + session::RuntimeSessionExt, + }, + layout::scan::scan_builder::ScanBuilder, + session::VortexSession, +}; + +use super::RecordBatchIter; + +/// Open a vortex file, apply an optional projection, and return an async +/// record batch stream over all rows. +pub(super) async fn scan_vortex_async( + data: Bytes, + projection: Option<&[String]>, +) -> Result { + let session = VortexSession::default().with_tokio(); + let vxf = session.open_options().open_buffer(data)?; + let scan = open_projected_scan(&vxf, projection)?; + + let stream_schema: Arc = Arc::new(scan.dtype()?.to_arrow_schema()?); + let data_type = DataType::Struct(stream_schema.fields().clone()); + + let stream = scan.into_array_stream()?.then(move |result| { + let data_type = data_type.clone(); + async move { + match result { + Ok(array) => vortex_array_to_record_batch(array, &data_type), + Err(e) => Err(ArrowError::ExternalError(Box::new(e))), + } + } + }); + + let stream: RecordBatchStream = Box::pin(stream); + Ok(stream) +} + +/// Build a lazy sync iterator for one vortex row range. +pub(super) fn scan_vortex_row_range( + data: Bytes, + projection: Option<&[String]>, + row_range: Range, +) -> Result { + let runtime = SingleThreadRuntime::default(); + let session = VortexSession::default().with_handle(runtime.handle()); + let vxf = session.open_options().open_buffer(data)?; + let scan = open_projected_scan(&vxf, projection)?.with_row_range(row_range); + + let stream_schema: Arc = Arc::new(scan.dtype()?.to_arrow_schema()?); + let data_type = DataType::Struct(stream_schema.fields().clone()); + + let iter: Box> + 'static> = Box::new( + scan.into_array_iter(&runtime)? + .map(move |result| match result { + Ok(array) => vortex_array_to_record_batch(array, &data_type), + Err(e) => Err(ArrowError::ExternalError(Box::new(e))), + }), + ); + + Ok(Box::new(VortexRowRangeIter { + iter, + _runtime: runtime, + })) +} + +struct VortexRowRangeIter { + iter: Box> + 'static>, + _runtime: SingleThreadRuntime, // dropped last — keeps Arc alive +} + +impl Iterator for VortexRowRangeIter { + type Item = Result; + + fn next(&mut self) -> Option { + self.iter.next() + } +} + +fn open_projected_scan( + vxf: &VortexFile, + projection: Option<&[String]>, +) -> Result, anyhow::Error> { + let full_schema = vxf.dtype().to_arrow_schema()?; + let mut scan = vxf.scan()?; + if let Some(cols) = projection { + let kept: Vec<&str> = cols + .iter() + .filter(|c| full_schema.field_with_name(c).is_ok()) + .map(String::as_str) + .collect(); + if !kept.is_empty() { + scan = scan.with_projection(select(kept, root())); + } + } + Ok(scan) +} diff --git a/src/service/tantivy/sequential.rs b/src/service/tantivy/sequential.rs index 31bed9dce8..765a799432 100644 --- a/src/service/tantivy/sequential.rs +++ b/src/service/tantivy/sequential.rs @@ -13,23 +13,18 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::io::Cursor; - use anyhow::Context; -use arrow::{array::RecordBatch, error::ArrowError}; +use arrow::array::RecordBatch; use bytes::Bytes; use config::{ - FileFormat, TIMESTAMP_COL_NAME, get_batch_size, - utils::{ - parquet::{RecordBatchStream, get_recordbatch_reader_from_bytes}, - tantivy::tokenizer::{CollectType, O2_TOKENIZER, o2_tokenizer_build}, - }, + FileFormat, TIMESTAMP_COL_NAME, + utils::tantivy::tokenizer::{CollectType, O2_TOKENIZER, o2_tokenizer_build}, }; use futures::TryStreamExt; -use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; use tokio::task::JoinHandle; use super::{TantivyIndexSchema, convert_batch_to_docs_sync}; +use crate::service::tantivy::reader::file_stream; /// Create a tantivy index in the given directory for the input file bytes. pub(super) async fn build_index( @@ -39,7 +34,9 @@ pub(super) async fn build_index( index_schema: TantivyIndexSchema, ) -> Result, anyhow::Error> { let start = std::time::Instant::now(); - let reader = open_reader(file_format, buf, &index_schema).await?; + let mut projection: Vec = index_schema.fields.iter().cloned().collect(); + projection.push(TIMESTAMP_COL_NAME.to_string()); + let reader = file_stream(file_format, buf, Some(&projection)).await?; let tokenizer_manager = tantivy::tokenizer::TokenizerManager::default(); tokenizer_manager.register(O2_TOKENIZER, o2_tokenizer_build(CollectType::Ingest)); let index_writer = tantivy::IndexBuilder::new() @@ -48,7 +45,7 @@ pub(super) async fn build_index( .single_segment_index_writer(tantivy_dir, 50_000_000) .context("failed to create index builder")?; - let (tx, rx) = tokio::sync::mpsc::channel::(2); + let (tx, mut rx) = tokio::sync::mpsc::channel::(2); let producer: JoinHandle> = tokio::task::spawn(async move { let mut total_num_rows = 0usize; @@ -67,7 +64,6 @@ pub(super) async fn build_index( Ok(total_num_rows) }); - let mut rx = rx; let writer_task: JoinHandle> = tokio::task::spawn_blocking(move || { let mut writer = index_writer; @@ -82,7 +78,7 @@ pub(super) async fn build_index( Ok(writer) }); - let _total_num_rows = producer.await??; + producer.await??; let index_writer = writer_task.await??; log::info!( @@ -107,38 +103,6 @@ pub(super) async fn build_index( Ok(Some(index)) } -/// Open a `RecordBatch` stream that only decodes the columns we will index. -async fn open_reader( - file_format: FileFormat, - buf: Bytes, - index_schema: &TantivyIndexSchema, -) -> Result { - if !matches!(file_format, FileFormat::Parquet) { - let (_, reader) = get_recordbatch_reader_from_bytes(file_format, buf).await?; - return Ok(reader); - } - - let builder = ParquetRecordBatchStreamBuilder::new(Cursor::new(buf)).await?; - let projection_indices: Vec = builder - .schema() - .fields() - .iter() - .enumerate() - .filter_map(|(i, f)| { - (f.name() == TIMESTAMP_COL_NAME || index_schema.fields.contains(f.name())).then_some(i) - }) - .collect(); - let projection_mask = ProjectionMask::roots( - builder.metadata().file_metadata().schema_descr(), - projection_indices, - ); - let stream = builder - .with_batch_size(get_batch_size()) - .with_projection(projection_mask) - .build()?; - Ok(Box::pin(stream.map_err(ArrowError::from))) -} - #[cfg(test)] mod tests { use std::sync::Arc; @@ -322,7 +286,7 @@ mod tests { #[tokio::test] async fn test_generate_tantivy_index_without_timestamp() { let dir = RamDirectory::create(); - let batch = create_test_batch(10, false, true, true); + let batch = create_test_batch(10, true, true, true); let buf = create_test_parquet_bytes(vec![batch.clone()]).await; let result = build_index(