diff --git a/Cargo.toml b/Cargo.toml index 398888eb0c..c481cf9b36 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,7 @@ default = [] enterprise = [] vectorscan = [] cloud = [] +vortex = ["config/vortex"] mimalloc = ["dep:mimalloc"] jemalloc = ["dep:tikv-jemallocator"] profiling = [ diff --git a/src/config/Cargo.toml b/src/config/Cargo.toml index b7fc4fa8ec..8ab903e41d 100644 --- a/src/config/Cargo.toml +++ b/src/config/Cargo.toml @@ -8,6 +8,7 @@ license.workspace = true default = ["gxhash"] enterprise = [] gxhash = ["dep:gxhash"] +vortex = ["dep:vortex"] [dependencies] aes-siv.workspace = true @@ -93,7 +94,7 @@ tonic-prost.workspace = true tracing-opentelemetry.workspace = true vortex = { git = "https://github.com/openobserve/vortex", rev = "7f3491d604b1e298b9162e174a20a15cf863ff83", features = [ "tokio", -] } +], optional = true } [dev-dependencies] expect-test.workspace = true diff --git a/src/config/src/utils/parquet.rs b/src/config/src/utils/parquet.rs index e9088d4172..da74347ce5 100644 --- a/src/config/src/utils/parquet.rs +++ b/src/config/src/utils/parquet.rs @@ -21,9 +21,13 @@ use std::{ sync::Arc, }; -use arrow::{array::StructArray, error::ArrowError, record_batch::RecordBatch}; -use arrow_schema::{DataType, Schema}; -use futures::{Stream, StreamExt, TryStreamExt}; +#[cfg(feature = "vortex")] +use arrow::{array::StructArray, datatypes::DataType}; +use arrow::{error::ArrowError, record_batch::RecordBatch}; +use arrow_schema::Schema; +#[cfg(feature = "vortex")] +use futures::StreamExt; +use futures::{Stream, TryStreamExt}; use parquet::{ arrow::{ AsyncArrowWriter, ParquetRecordBatchStreamBuilder, @@ -33,6 +37,7 @@ use parquet::{ 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, @@ -161,6 +166,7 @@ pub async fn get_recordbatch_reader_from_bytes( let stream: RecordBatchStream = Box::pin(reader.map_err(ArrowError::from)); Ok((schema, stream)) } + #[cfg(feature = "vortex")] FileFormat::Vortex => { // Read vortex file from bytes and convert to record batches let session = VortexSession::default().with_tokio(); @@ -202,6 +208,10 @@ pub async fn get_recordbatch_reader_from_bytes( let stream: RecordBatchStream = Box::pin(stream); Ok((schema, stream)) } + #[cfg(not(feature = "vortex"))] + FileFormat::Vortex => Err(anyhow::anyhow!( + "Vortex file format requires the vortex feature" + )), } } @@ -221,6 +231,7 @@ pub async fn read_schema_from_file(path: &PathBuf) -> Result, anyhow let format = FileFormat::from_extension(path_str); match format { + #[cfg(feature = "vortex")] Some(FileFormat::Vortex) => { // Read vortex file let session = VortexSession::default().with_tokio(); @@ -248,6 +259,7 @@ pub async fn read_schema_from_bytes( let arrow_reader = ParquetRecordBatchStreamBuilder::new(schema_reader).await?; Ok(arrow_reader.schema().clone()) } + #[cfg(feature = "vortex")] FileFormat::Vortex => { let session = VortexSession::default().with_tokio(); let buf = Buffer::from(data.to_vec()); @@ -255,6 +267,10 @@ pub async fn read_schema_from_bytes( let schema = Arc::new(vxf.dtype().to_arrow_schema()?); Ok(schema) } + #[cfg(not(feature = "vortex"))] + FileFormat::Vortex => { + anyhow::bail!("Vortex file format requires the vortex feature to be enabled") + } } } diff --git a/src/service/search/datafusion/exec.rs b/src/service/search/datafusion/exec.rs index d1afeb921f..35059e8b1c 100644 --- a/src/service/search/datafusion/exec.rs +++ b/src/service/search/datafusion/exec.rs @@ -48,8 +48,9 @@ use datafusion::{ prelude::{SessionContext, col}, }; #[cfg(feature = "enterprise")] +use o2_enterprise::enterprise::{common::config::get_config as get_o2_config, search::WorkGroup}; +#[cfg(all(feature = "enterprise", feature = "vortex"))] use { - o2_enterprise::enterprise::{common::config::get_config as get_o2_config, search::WorkGroup}, vortex::{VortexSessionDefault, io::session::RuntimeSessionExt, session::VortexSession}, vortex_datafusion::VortexFormat, }; @@ -412,25 +413,25 @@ impl TableBuilder { // Group files by format let mut parquet_files = Vec::new(); - #[cfg(feature = "enterprise")] + #[cfg(all(feature = "enterprise", feature = "vortex"))] let mut vortex_files = Vec::new(); for file in files { match FileFormat::from_extension(&file.key) { - #[cfg(feature = "enterprise")] + #[cfg(all(feature = "enterprise", feature = "vortex"))] Some(FileFormat::Vortex) => vortex_files.push(file), _ => parquet_files.push(file), // Default to parquet } } - #[cfg(feature = "enterprise")] + #[cfg(all(feature = "enterprise", feature = "vortex"))] log::info!( "[trace_id: {}] parquet_files numbers: {}, vortex_files numbers: {}", session.id, parquet_files.len(), vortex_files.len() ); - #[cfg(not(feature = "enterprise"))] + #[cfg(not(all(feature = "enterprise", feature = "vortex")))] log::info!( "[trace_id: {}] parquet_files numbers: {}", session.id, @@ -453,7 +454,7 @@ impl TableBuilder { tables.push(table); } - #[cfg(feature = "enterprise")] + #[cfg(all(feature = "enterprise", feature = "vortex"))] if !vortex_files.is_empty() { let table = self .build_table_for_format( @@ -481,15 +482,15 @@ impl TableBuilder { // Configure listing options with the appropriate file format let file_format: Arc = match format { FileFormat::Parquet => Arc::new(ParquetFormat::default()), - #[cfg(feature = "enterprise")] + #[cfg(all(feature = "enterprise", feature = "vortex"))] FileFormat::Vortex => { let vortex_session = VortexSession::default().with_tokio(); Arc::new(VortexFormat::new(vortex_session)) } - #[cfg(not(feature = "enterprise"))] + #[cfg(not(all(feature = "enterprise", feature = "vortex")))] FileFormat::Vortex => { return Err(DataFusionError::Execution( - "Vortex file format requires enterprise feature".to_string(), + "Vortex file format requires enterprise and vortex features".to_string(), )); } }; diff --git a/src/service/search/datafusion/merge/downsampling.rs b/src/service/search/datafusion/merge/downsampling.rs index 0b6ae82d3f..7162cbf456 100644 --- a/src/service/search/datafusion/merge/downsampling.rs +++ b/src/service/search/datafusion/merge/downsampling.rs @@ -31,7 +31,9 @@ use datafusion::{ physical_plan::execute_stream, }; use futures::TryStreamExt; +#[cfg(all(feature = "enterprise", feature = "vortex"))] use o2_enterprise::enterprise::search::vortex::{Utf8Compressor, VORTEX_RUNTIME}; +#[cfg(all(feature = "enterprise", feature = "vortex"))] use vortex::{ VortexSessionDefault, array::{ArrayRef, arrow::FromArrowArray}, @@ -114,9 +116,16 @@ pub async fn merge_parquet_files_with_downsampling( FileFormat::Parquet => { write_downsampled_parquet(rx, &schema, bloom_filter_fields, &metadata, &cfg).await? } + #[cfg(all(feature = "enterprise", feature = "vortex"))] FileFormat::Vortex => { write_downsampled_vortex(rx, schema.clone(), cfg.compact.max_file_size as i64).await? } + #[cfg(not(all(feature = "enterprise", feature = "vortex")))] + FileFormat::Vortex => { + return Err(DataFusionError::Execution( + "Vortex file format requires enterprise and vortex features".to_string(), + )); + } }; // Wait for read task to complete @@ -197,6 +206,7 @@ async fn write_downsampled_parquet( Ok((bufs, file_metas)) } +#[cfg(all(feature = "enterprise", feature = "vortex"))] async fn write_downsampled_vortex( mut rx: tokio::sync::mpsc::Receiver, schema: Arc, diff --git a/src/service/search/datafusion/merge/mod.rs b/src/service/search/datafusion/merge/mod.rs index 8a789694ce..14441c4988 100644 --- a/src/service/search/datafusion/merge/mod.rs +++ b/src/service/search/datafusion/merge/mod.rs @@ -220,15 +220,15 @@ async fn write_vortex( rx: tokio::sync::mpsc::Receiver, read_task: tokio::task::JoinHandle>, ) -> Result> { - #[cfg(feature = "enterprise")] + #[cfg(all(feature = "enterprise", feature = "vortex"))] { o2_enterprise::enterprise::search::vortex::write_vortex(schema, rx, read_task).await } - #[cfg(not(feature = "enterprise"))] + #[cfg(not(all(feature = "enterprise", feature = "vortex")))] { let _ = (schema, rx, read_task); Err(DataFusionError::Execution( - "Vortex file format requires enterprise feature".to_string(), + "Vortex file format requires enterprise and vortex features".to_string(), )) } } diff --git a/src/service/search/datafusion/table_provider/helpers.rs b/src/service/search/datafusion/table_provider/helpers.rs index 25fc764b1d..28d3135a9d 100644 --- a/src/service/search/datafusion/table_provider/helpers.rs +++ b/src/service/search/datafusion/table_provider/helpers.rs @@ -32,7 +32,7 @@ use datafusion::{ scalar::ScalarValue, }; use hashbrown::HashMap; -#[cfg(feature = "enterprise")] +#[cfg(all(feature = "enterprise", feature = "vortex"))] use o2_enterprise::enterprise::search::vortex::generate_vortex_access_plan; use crate::service::search::{datafusion::storage, index::IndexCondition}; @@ -44,9 +44,9 @@ pub fn generate_access_plan( let file_format = FileFormat::from_extension(file.path().as_ref())?; match file_format { FileFormat::Parquet => generate_parquet_access_plan(file, segment_ids), - #[cfg(feature = "enterprise")] + #[cfg(all(feature = "enterprise", feature = "vortex"))] FileFormat::Vortex => generate_vortex_access_plan(segment_ids), - #[cfg(not(feature = "enterprise"))] + #[cfg(not(all(feature = "enterprise", feature = "vortex")))] FileFormat::Vortex => None, } }