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

perf: impl new tantivy index generate for vortex format (#12354)

Browse Source 
related to 
- https://github.com/openobserve/openobserve/pull/12331
- https://github.com/openobserve/openobserve/pull/12292
This commit is contained in:
Huaijin
2026-06-01 06:51:41 +00:00
committed by GitHub
parent 9a1ed606aa
commit 3a951f1398
16 changed files with 744 additions and 1595 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config/src/utils/inverted_index.rs
+14 -20
View File
@@ -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<String> {
pub fn to_tantivy_name(from: &str) -> Option<String> {
let mut parts: Vec<Cow<str>> = 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);
}
}
src/config/src/utils/parquet.rs
+27 -41
View File
@@ -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<Box<dyn Stream<Item = Result<RecordBatch, ArrowError>> + 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<T>(
reader: T,
) -> Result<(Arc<Schema>, ParquetRecordBatchStream<T>), 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<RecordBatch, ArrowError> {
let array = array
.into_arrow(data_type)
.map_err(|e| ArrowError::ExternalError(Box::new(e)))?;
let struct_array = array
.as_any()
.downcast_ref::<StructArray>()
.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<Schema>, 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::<StructArray>()
.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))),
}
}
src/handler/grpc/request/event.rs
+6 -8
View File
@@ -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());
}
src/service/compact/bloom_build.rs
+2 -2
View File
@@ -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<Vec<FieldBloom>> {
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 {
src/service/compact/deleted.rs
+2 -3
View File
@@ -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<i64, anyhow::Error> {
.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
}
src/service/compact/merge.rs
+3 -1272
View File
File diff suppressed because it is too large Load Diff
src/service/search/grpc/storage.rs
+9 -10
View File
@@ -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()
src/service/tantivy/bloom_builder.rs
+17 -5
View File
@@ -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::<Vec<_>>();
let timestamps: Vec<i64> = (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,
src/service/tantivy/mod.rs
+14 -26
View File
@@ -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::<HashSet<_>>();
@@ -151,10 +146,7 @@ fn build_tantivy_schema(
.filter(|f| schema_fields.contains_key(f))
.map(String::from)
.collect::<HashSet<_>>();
let tantivy_fields = fts_fields_filtered
.union(&index_fields_filtered)
.cloned()
.collect::<HashSet<_>>();
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::<Int64Array>() {
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::<Int64Array>())
.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));
src/service/tantivy/parallel.rs
+365 -106
View File
@@ -13,141 +13,204 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
//! Row-group-direct parallel tantivy index builder.
//!
//! Activated when `compact.tantivy_parallel_build_workers > 1` and the input
//! file format is parquet. Each parquet row_group becomes one independent
//! chunk segment built off the runtime in `tokio::task::spawn_blocking`; the
//! resulting per-row_group segments are stitched together by
//! `tantivy::indexer::merge_indices` in row_group order.
//!
//! Correctness of `doc_id == global_row_index` is preserved because:
//! 1. each per-row_group `SingleSegmentIndexWriter` writes rows in stored order, so its local
//! doc_ids equal the row's offset inside that row_group;
//! 2. `IndexMerger` stacks input segments in input order (current tantivy has no `sort_by_field`
//! path), so each segment's doc_ids land at `i * PARQUET_MAX_ROW_GROUP_SIZE`. For example,
//! pretending the fixed row_group size were 100, segment 0 would occupy doc_ids `[0, 100)`,
//! segment 1 `[100, 200)`, segment 2 `[200, 300)`, etc. (the final row_group may be shorter,
//! but it is always the last segment);
//! 3. parquet stores row_groups in row order, so the global row index of row_group `i`'s row `k`
//! is `i * PARQUET_MAX_ROW_GROUP_SIZE + k` (with size 100, row 5 of row_group 1 maps to global
//! row index `100 + 5 = 105`).
use std::sync::Arc;
use anyhow::{Context, Error};
use bytes::Bytes;
use config::{
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<D: tantivy::Directory + Send + Sync + 'static>(
tantivy_dir: D,
file_format: FileFormat,
buf: Bytes,
index_schema: TantivyIndexSchema,
thread_num: usize,
) -> Result<Option<tantivy::Index>, Error> {
match file_format {
FileFormat::Parquet => {
let buf_meta = buf.clone();
let num_row_groups = tokio::task::spawn_blocking(move || -> Result<usize, Error> {
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<u64, Error> {
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<D: tantivy::Directory + Send + Sync + 'static>(
async fn build_parallel_index<D, F>(
tantivy_dir: D,
buf: Bytes,
index_schema: TantivyIndexSchema,
thread_num: usize,
) -> Result<Option<tantivy::Index>, Error> {
// Parse parquet metadata once on a blocking thread (footer decode is sync I/O).
let num_row_groups = {
let buf = buf.clone();
tokio::task::spawn_blocking(move || -> Result<usize, Error> {
let builder = ParquetRecordBatchReaderBuilder::try_new(buf)
.context("failed to open parquet for metadata")?;
Ok(builder.metadata().num_row_groups())
})
.await??
};
if num_row_groups == 0 {
num_chunks: usize,
make_selector: F,
) -> Result<Option<tantivy::Index>, 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<JoinHandle<Result<TaskOutput, Error>>> = Vec::with_capacity(num_row_groups);
for rg_idx in 0..num_row_groups {
let mut tasks: Vec<JoinHandle<Result<SegmentOutput, Error>>> = 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<TaskOutput> = Vec::with_capacity(tasks.len());
for t in tasks {
indexed.push(t.await??);
}
// Sort by row_group index to preserve global row order in the merged index.
indexed.sort_by_key(|(rg, ..)| *rg);
let total_num_rows: usize = indexed.iter().map(|(_, _, n)| n).sum();
let indices: Vec<tantivy::Index> = indexed.into_iter().map(|(_, i, _)| i).collect();
// await all tasks concurrently, then sort to restore chunk order.
let mut indexed: Vec<SegmentOutput> = join_all(tasks)
.await
.into_iter()
.map(|r| -> Result<SegmentOutput, Error> { r? })
.collect::<Result<_, _>>()?;
indexed.sort_by_key(|s| s.chunk_idx);
let total_rows: usize = indexed.iter().map(|s| s.row_count).sum();
let indices: Vec<tantivy::Index> = 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<D: tantivy::Directory + Send + Sync + 'static>(
));
}
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<D: tantivy::Directory>(
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<D: tantivy::Directory>(
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<SegmentOutput, Error> {
let mut projection: Vec<String> = index_schema.fields.iter().cloned().collect();
projection.push(TIMESTAMP_COL_NAME.to_string());
let arrow_schema = builder.schema();
let projection_indices: Vec<usize> = 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<D: tantivy::Directory>(
.tokenizers(tokenizer_manager)
.single_segment_index_writer::<tantivy::TantivyDocument>(dir, 50_000_000)?;
let mut rg_rows: usize = 0;
for batch_res in reader {
let 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<D: tantivy::Directory>(
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<i64> = (0..num_rows).map(|i| (start + i) as i64).collect();
let markers: Vec<String> = (0..num_rows)
.map(|i| format!("row{:09}", start + i))
.collect();
let schema = Arc::new(Schema::new(fields));
RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(timestamps)),
Arc::new(StringArray::from(markers)),
],
)
.unwrap()
}
/// Writes batches 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<RecordBatch>) -> 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());
}
}
}
src/service/tantivy/puffin/mod.rs
-24
View File
@@ -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 <http://www.gnu.org/licenses/>.
pub use tantivy_utils::puffin::*;
pub mod reader {
pub use tantivy_utils::puffin::reader::*;
}
pub mod writer {
pub use tantivy_utils::puffin::writer::*;
}
src/service/tantivy/puffin_directory/mod.rs
-32
View File
@@ -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 <http://www.gnu.org/licenses/>.
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::*;
}
src/service/tantivy/reader/mod.rs
+68
View File
@@ -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 <http://www.gnu.org/licenses/>.
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<dyn Iterator<Item = Result<RecordBatch, ArrowError>>>;
pub(super) enum ChunkSelector {
Parquet(usize), // row_group id
#[cfg(all(feature = "vortex", feature = "enterprise"))]
Vortex(Range<u64>), // 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<RecordBatchStream, anyhow::Error> {
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<RecordBatchIter, anyhow::Error> {
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)
}
}
}
src/service/tantivy/reader/parquet.rs
+85
View File
@@ -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 <http://www.gnu.org/licenses/>.
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<RecordBatchStream, anyhow::Error> {
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<RecordBatchIter, anyhow::Error> {
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<Schema>,
metadata: &parquet::file::metadata::ParquetMetaData,
projection: Option<&[String]>,
) -> Result<Option<ProjectionMask>, anyhow::Error> {
let Some(cols) = projection else {
return Ok(None);
};
let kept: Vec<usize> = 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,
)))
}
src/service/tantivy/reader/vortex.rs
+122
View File
@@ -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 <http://www.gnu.org/licenses/>.
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<RecordBatchStream, anyhow::Error> {
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<Schema> = 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<u64>,
) -> Result<RecordBatchIter, anyhow::Error> {
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<Schema> = Arc::new(scan.dtype()?.to_arrow_schema()?);
let data_type = DataType::Struct(stream_schema.fields().clone());
let iter: Box<dyn Iterator<Item = Result<RecordBatch, ArrowError>> + '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<dyn Iterator<Item = Result<RecordBatch, ArrowError>> + 'static>,
_runtime: SingleThreadRuntime, // dropped last — keeps Arc<Sender> alive
}
impl Iterator for VortexRowRangeIter {
type Item = Result<RecordBatch, ArrowError>;
fn next(&mut self) -> Option<Self::Item> {
self.iter.next()
}
}
fn open_projected_scan(
vxf: &VortexFile,
projection: Option<&[String]>,
) -> Result<ScanBuilder<ArrayRef>, 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)
}
src/service/tantivy/sequential.rs
+10 -46
View File
@@ -13,23 +13,18 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use 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<D: tantivy::Directory>(
@@ -39,7 +34,9 @@ pub(super) async fn build_index<D: tantivy::Directory>(
index_schema: TantivyIndexSchema,
) -> Result<Option<tantivy::Index>, anyhow::Error> {
let start = std::time::Instant::now();
let reader = open_reader(file_format, buf, &index_schema).await?;
let mut projection: Vec<String> = 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<D: tantivy::Directory>(
.single_segment_index_writer(tantivy_dir, 50_000_000)
.context("failed to create index builder")?;
let (tx, rx) = tokio::sync::mpsc::channel::<RecordBatch>(2);
let (tx, mut rx) = tokio::sync::mpsc::channel::<RecordBatch>(2);
let producer: JoinHandle<Result<usize, anyhow::Error>> = tokio::task::spawn(async move {
let mut total_num_rows = 0usize;
@@ -67,7 +64,6 @@ pub(super) async fn build_index<D: tantivy::Directory>(
Ok(total_num_rows)
});
let mut rx = rx;
let writer_task: JoinHandle<Result<_, anyhow::Error>> =
tokio::task::spawn_blocking(move || {
let mut writer = index_writer;
@@ -82,7 +78,7 @@ pub(super) async fn build_index<D: tantivy::Directory>(
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<D: tantivy::Directory>(
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<RecordBatchStream, anyhow::Error> {
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<usize> = 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(