diff --git a/Cargo.lock b/Cargo.lock index 5a417b2011..e1ebdb2b16 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10961,6 +10961,7 @@ dependencies = [ "sqlparser", "tantivy", "tantivy_utils", + "tempfile", "thiserror 2.0.19", "tokio", "tokio-stream", diff --git a/src/compaction/src/deleted.rs b/src/compaction/src/deleted.rs index 6860210af8..a3f816d73e 100644 --- a/src/compaction/src/deleted.rs +++ b/src/compaction/src/deleted.rs @@ -16,7 +16,7 @@ //! Delayed deletion of compacted files. use config::{ - meta::stream::{FileKey, FileMeta}, + meta::stream::{FileKey, FileListDeleted, FileMeta}, utils::inverted_index::to_tantivy_name, }; use infra::{file_list as infra_file_list, storage}; @@ -24,6 +24,42 @@ use infra::{file_list as infra_file_list, storage}; // Batch size for deleting files from file_list_deleted table const BATCH_SIZE: i64 = 10000; +/// `(account, derived object key)` pairs of the files for which `derive` +/// returns a key. +fn derived_files( + files: &[FileListDeleted], + derive: impl Fn(&FileListDeleted) -> Option, +) -> Vec<(String, String)> { + files + .iter() + .filter_map(|file| derive(file).map(|key| (file.account.clone(), key))) + .collect() +} + +/// Delete objects from storage, ignoring `not found` (already deleted, or a +/// derived object that was never written). +async fn delete_from_storage( + kind: &str, + files: Vec<(String, String)>, +) -> Result<(), anyhow::Error> { + if files.is_empty() { + return Ok(()); + } + if let Err(e) = storage::del( + files + .iter() + .map(|(account, key)| (account.as_str(), key.as_str())) + .collect::>(), + ) + .await + && !e.to_string().to_lowercase().contains("not found") + { + log::error!("[COMPACTOR] delete {kind} files from storage failed: {e}"); + return Err(e.into()); + } + Ok(()) +} + pub async fn delete(org_id: &str, time_max: i64) -> Result { let files = infra_file_list::query_deleted(org_id, time_max, BATCH_SIZE).await?; if files.is_empty() { @@ -32,84 +68,35 @@ pub async fn delete(org_id: &str, time_max: i64) -> Result { let files_num = files.len() as i64; // delete files from storage - if let Err(e) = storage::del( + delete_from_storage( + "data", files .iter() - .filter_map(|file| { - if !ingester::is_wal_file(&file.file) { - Some((file.account.as_str(), file.file.as_str())) - } else { - None - } - }) - .collect::>(), + .filter(|file| !ingester::is_wal_file(&file.file)) + .map(|file| (file.account.clone(), file.file.clone())) + .collect(), ) - .await - { - // maybe the file already deleted, so we just skip the `not found` error - if !e.to_string().to_lowercase().contains("not found") { - log::error!("[COMPACTOR] delete files from storage failed: {e}"); - return Err(e.into()); - } - } + .await?; - // delete related inverted index puffin files - let inverted_index_files = files - .iter() - .filter_map(|file| { - if file.index_file { - to_tantivy_name(&file.file).map(|f| (file.account.to_string(), f)) - } else { - None - } - }) - .collect::>(); - if !inverted_index_files.is_empty() - && let Err(e) = storage::del( - inverted_index_files - .iter() - .map(|file| (file.0.as_str(), file.1.as_str())) - .collect::>(), - ) - .await - { - // maybe the file already deleted or there's not related index files, - // so we just skip the `not found` error - if !e.to_string().to_lowercase().contains("not found") { - log::error!("[COMPACTOR] delete files from storage failed: {e}"); - return Err(e.into()); - } - } - - // delete flattened files from storage - let flattened_files = files - .iter() - .filter_map(|file| { - if file.flattened { - Some(( - file.account.to_string(), - super::flatten::generate_flatten_file_key(&file.file), - )) - } else { - None - } - }) - .collect::>(); - if !flattened_files.is_empty() - && let Err(e) = storage::del( - flattened_files - .iter() - .map(|file| (file.0.as_str(), file.1.as_str())) - .collect::>(), - ) - .await - { - // maybe the file already deleted, so we just skip the `not found` error - if !e.to_string().to_lowercase().contains("not found") { - log::error!("[COMPACTOR] delete files from storage failed: {e}"); - return Err(e.into()); - } - } + // derived objects are not tracked in file_list: delete them together with + // their parent data file + delete_from_storage( + "inverted index", + derived_files(&files, |file| { + file.index_file + .then(|| to_tantivy_name(&file.file)) + .flatten() + }), + ) + .await?; + delete_from_storage( + "flattened", + derived_files(&files, |file| { + file.flattened + .then(|| super::flatten::generate_flatten_file_key(&file.file)) + }), + ) + .await?; // delete files from file_list_deleted table if let Err(e) = infra_file_list::batch_remove_deleted( diff --git a/src/compaction/src/merge.rs b/src/compaction/src/merge.rs index c69c11fad5..8453d834d6 100644 --- a/src/compaction/src/merge.rs +++ b/src/compaction/src/merge.rs @@ -50,7 +50,7 @@ use o2_enterprise::enterprise::common::downsampling::get_largest_downsampling_ru use schema::generate_schema_for_defined_schema_fields; use search::datafusion::{ exec::TableBuilder, - merge::{self, MergeParquetResult}, + merge::{self, MergeParquetResult, MergedFile}, sort_order::FileSortOrder, }; use search_service::file_list; @@ -901,126 +901,71 @@ pub async fn merge_files( log::debug!("skip index generation for stream: {org_id}/{stream_type}/{stream_name}"); } - let mut new_files = Vec::new(); - match buf { - MergeParquetResult::Single { - buf, - file_meta: mut new_file_meta, - file_format, - } => { - if new_file_meta.compressed_size == 0 { - return Err(anyhow::anyhow!( - "merge_parquet_files error: compressed_size is 0" - )); - } - - let id = ider::generate_file_name(); - let new_file_key = format!("{prefix}/{id}{}", file_format.extension()); - log::info!( - "[COMPACTOR:WORKER:{thread_id}] merged {} files into a new file: {new_file_key}, original_size: {}, compressed_size: {}, took: {} ms", - retain_file_list.len(), - new_file_meta.original_size, - new_file_meta.compressed_size, - start.elapsed().as_millis(), - ); - - // upload file to storage - let buf = Bytes::from(buf); - if cfg.cache_latest_files.enabled - && cfg.cache_latest_files.cache_parquet - && cfg.cache_latest_files.download_from_node - { - infra::cache::file_data::disk::set(&new_file_key, buf.clone()).await?; - log::debug!("merge_files {new_file_key} file_data::disk::set success"); - } - - // TODO: check how compliance will interact with org storage - let account = storage::get_account(org_id, &new_file_key).unwrap_or_default(); - if cfg.s3.feature_force_infrequent_access && storage_type.is_compliance() { - storage::put_with_compliance(&account, &new_file_key, buf.clone()).await?; - } else { - storage::put(&account, &new_file_key, buf.clone()).await?; - } - - if cfg.search.inverted_index_enabled && stream_type.support_index() && need_index { - generate_inverted_index( - org_id, - &new_file_key, - &full_text_search_fields, - &index_fields, - &retain_file_list, - &mut new_file_meta, - latest_schema.clone(), - buf, - ) - .await?; - } - new_files.push(FileKey::new(0, account, new_file_key, new_file_meta, false)); + let MergeParquetResult { + files: merged_files, + file_format, + } = buf; + let mut new_files = Vec::with_capacity(merged_files.len()); + for MergedFile { + buf, + meta: mut new_file_meta, + } in merged_files + { + new_file_meta.compressed_size = buf.len() as i64; + if new_file_meta.compressed_size == 0 { + return Err(anyhow::anyhow!( + "merge_parquet_files error: compressed_size is 0" + )); } - MergeParquetResult::Multiple { - bufs, - file_metas, - file_format, - } => { - for (buf, file_meta) in bufs.into_iter().zip(file_metas) { - let mut new_file_meta = file_meta; - new_file_meta.compressed_size = buf.len() as i64; - if new_file_meta.compressed_size == 0 { - return Err(anyhow::anyhow!( - "merge_parquet_files error: compressed_size is 0" - )); - } - let id = ider::generate_file_name(); - let new_file_key = format!("{prefix}/{id}{}", file_format.extension()); + let id = ider::generate_file_name(); + let new_file_key = format!("{prefix}/{id}{}", file_format.extension()); - // upload file to storage - let buf = Bytes::from(buf); - if cfg.cache_latest_files.enabled - && cfg.cache_latest_files.cache_parquet - && cfg.cache_latest_files.download_from_node - { - infra::cache::file_data::disk::set(&new_file_key, buf.clone()).await?; - log::debug!("merge_files {new_file_key} file_data::disk::set success"); - } - - // TODO: check how compliance will interact with org storage - let account = storage::get_account(org_id, &new_file_key).unwrap_or_default(); - if cfg.s3.feature_force_infrequent_access && storage_type.is_compliance() { - storage::put_with_compliance(&account, &new_file_key, buf.clone()).await?; - } else { - storage::put(&account, &new_file_key, buf.clone()).await?; - } - - if cfg.search.inverted_index_enabled && stream_type.support_index() && need_index { - generate_inverted_index( - org_id, - &new_file_key, - &full_text_search_fields, - &index_fields, - &retain_file_list, - &mut new_file_meta, - latest_schema.clone(), - buf, - ) - .await?; - } - - new_files.push(FileKey::new(0, account, new_file_key, new_file_meta, false)); - } - log::info!( - "[COMPACTOR:WORKER:{thread_id}] merged {} files into a new file: {:?}, original_size: {}, compressed_size: {}, took: {} ms", - retain_file_list.len(), - new_files.iter().map(|f| f.key.as_str()).collect::>(), - new_files.iter().map(|f| f.meta.original_size).sum::(), - new_files - .iter() - .map(|f| f.meta.compressed_size) - .sum::(), - start.elapsed().as_millis(), - ); + // upload file to storage + let buf = Bytes::from(buf); + if cfg.cache_latest_files.enabled + && cfg.cache_latest_files.cache_parquet + && cfg.cache_latest_files.download_from_node + { + infra::cache::file_data::disk::set(&new_file_key, buf.clone()).await?; + log::debug!("merge_files {new_file_key} file_data::disk::set success"); } - }; + + // TODO: check how compliance will interact with org storage + let account = storage::get_account(org_id, &new_file_key).unwrap_or_default(); + if cfg.s3.feature_force_infrequent_access && storage_type.is_compliance() { + storage::put_with_compliance(&account, &new_file_key, buf.clone()).await?; + } else { + storage::put(&account, &new_file_key, buf.clone()).await?; + } + + if cfg.search.inverted_index_enabled && stream_type.support_index() && need_index { + generate_inverted_index( + org_id, + &new_file_key, + &full_text_search_fields, + &index_fields, + &retain_file_list, + &mut new_file_meta, + latest_schema.clone(), + buf, + ) + .await?; + } + new_files.push(FileKey::new(0, account, new_file_key, new_file_meta, false)); + } + log::info!( + "[COMPACTOR:WORKER:{thread_id}] merged {} files into {} new file(s): {:?}, original_size: {}, compressed_size: {}, took: {} ms", + retain_file_list.len(), + new_files.len(), + new_files.iter().map(|f| f.key.as_str()).collect::>(), + new_files.iter().map(|f| f.meta.original_size).sum::(), + new_files + .iter() + .map(|f| f.meta.compressed_size) + .sum::(), + start.elapsed().as_millis(), + ); Ok((new_files, retain_file_list)) } diff --git a/src/config/src/meta/stream.rs b/src/config/src/meta/stream.rs index fa2108098a..07f988ff1a 100644 --- a/src/config/src/meta/stream.rs +++ b/src/config/src/meta/stream.rs @@ -13,7 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::{cmp::max, fmt::Display, str::FromStr, sync::Arc}; +use std::{cmp::max, fmt::Display, ops::Range, str::FromStr, sync::Arc}; use arrow::buffer::BooleanBuffer; use chrono::{DateTime, Duration, TimeZone, Utc}; @@ -278,6 +278,9 @@ pub enum FileSelection { /// Row ids matched by the tantivy index, as a per-row bitmap of length /// `num_rows` (one bit per parquet row). Rows(Arc), + /// Sorted, non-overlapping physical row ranges selected by a compact + /// secondary index. This avoids materializing one bitmap bit per row. + RowRanges(Arc>>), /// Row group ids selected by row-group-level sampling. RowGroups(Arc>), } @@ -2027,6 +2030,15 @@ mod tests { assert_eq!(key.row_group_size, Some(1024)); } + #[test] + fn test_file_key_with_row_range_selection() { + let mut key = FileKey::from_file_name("files/k.parquet"); + let selection = FileSelection::RowRanges(Arc::new(vec![1..4, 8..12])); + key.with_selection(selection.clone(), Some(1024)); + assert_eq!(key.selection, Some(selection)); + assert_eq!(key.row_group_size, Some(1024)); + } + #[test] fn test_stream_params_is_valid() { let valid = StreamParams::new("org", "stream", StreamType::Logs); diff --git a/src/config/src/utils/parquet.rs b/src/config/src/utils/parquet.rs index 821b2a9a44..aa78a19a15 100644 --- a/src/config/src/utils/parquet.rs +++ b/src/config/src/utils/parquet.rs @@ -30,7 +30,10 @@ use arrow::{ use arrow_schema::Schema; use futures::{Stream, StreamExt, TryStreamExt}; use parquet::{ - arrow::{AsyncArrowWriter, ParquetRecordBatchStreamBuilder, arrow_reader::ArrowReaderMetadata}, + arrow::{ + AsyncArrowWriter, ParquetRecordBatchStreamBuilder, arrow_reader::ArrowReaderMetadata, + async_writer::AsyncFileWriter, + }, basic::{Compression, Encoding}, file::{metadata::KeyValue, properties::WriterProperties}, }; @@ -71,14 +74,14 @@ pub fn encode_vortex_file_meta(metadata: &FileMeta) -> Vec { .expect("file meta is always serializable") } -pub fn new_parquet_writer<'a>( - buf: &'a mut Vec, - schema: &'a Arc, - bloom_filter_fields: &'a [String], - metadata: &'a FileMeta, +pub fn new_parquet_writer( + buf: W, + schema: &Arc, + bloom_filter_fields: &[String], + metadata: &FileMeta, write_metadata: bool, compression: Option<&str>, -) -> AsyncArrowWriter<&'a mut Vec> { +) -> AsyncArrowWriter { let cfg = get_config(); let compression = compression.unwrap_or(&cfg.common.parquet_compression); let mut writer_props = WriterProperties::builder() diff --git a/src/jobs/src/job/files/pack.rs b/src/jobs/src/job/files/pack.rs index 95ba6a71cf..f1796a5ecf 100644 --- a/src/jobs/src/job/files/pack.rs +++ b/src/jobs/src/job/files/pack.rs @@ -42,7 +42,7 @@ use infra::{ }; use ingester::{PackSegment, PendingStreamStats}; use schema::generate_schema_for_defined_schema_fields; -use search::datafusion::merge::{self, MergeParquetResult}; +use search::datafusion::merge; use tantivy_utils::index_builder::create_tantivy_index; use tokio::sync::{Mutex, RwLock}; @@ -442,18 +442,9 @@ async fn upload_chunk( true, ) .await?; - match merge_result { - MergeParquetResult::Single { - buf, - file_meta, - file_format, - } => (Bytes::from(buf), file_meta, file_format), - MergeParquetResult::Multiple { .. } => { - return Err(anyhow::anyhow!( - "merge_parquet_files error: unexpected multiple files on ingester" - )); - } - } + // the ingester always merges into exactly one file + let (merged, file_format) = merge_result.into_single()?; + (Bytes::from(merged.buf), merged.meta, file_format) }; if new_file_meta.compressed_size == 0 { diff --git a/src/jobs/src/job/files/parquet.rs b/src/jobs/src/job/files/parquet.rs index bacdcfc45b..aad26f64f8 100644 --- a/src/jobs/src/job/files/parquet.rs +++ b/src/jobs/src/job/files/parquet.rs @@ -47,11 +47,7 @@ use infra::{ }; use ingester::WAL_PARQUET_METADATA; use schema::generate_schema_for_defined_schema_fields; -use search::datafusion::{ - exec::TableBuilder, - merge::{self, MergeParquetResult}, - sort_order::FileSortOrder, -}; +use search::datafusion::{exec::TableBuilder, merge, sort_order::FileSortOrder}; use tantivy_utils::index_builder::create_tantivy_index; use tokio::{ fs::remove_file, @@ -808,17 +804,9 @@ async fn merge_files( } }; - let (buf, mut new_file_meta, file_format) = match buf { - MergeParquetResult::Single { - buf, - file_meta, - file_format, - } => (buf, file_meta, file_format), - MergeParquetResult::Multiple { .. } => { - // ingester should not support multiple files, it will be handled in compactor mode - panic!("[INGESTER:JOB] merge_parquet_files error: multiple files"); - } - }; + // the ingester always merges into exactly one file + let (merged, file_format) = buf.into_single()?; + let (buf, mut new_file_meta) = (merged.buf, merged.meta); if new_file_meta.compressed_size == 0 { return Err(anyhow::anyhow!( diff --git a/src/promql/src/utils.rs b/src/promql/src/utils.rs index d4d9ee7aaf..df5ed8a31e 100644 --- a/src/promql/src/utils.rs +++ b/src/promql/src/utils.rs @@ -30,19 +30,20 @@ use datafusion::{ use hashbrown::HashSet; use promql_parser::label::{MatchOp, Matchers}; -pub fn apply_matchers(df: DataFrame, matchers: &Matchers) -> Result { - let mut df = df; +/// Build the DataFusion predicates used for PromQL label matchers. +/// +/// Keeping predicate construction separate lets storage-side secondary +/// indexes evaluate exactly the same matcher semantics as the final scan. +pub fn matcher_predicates(schema: &Schema, matchers: &Matchers) -> Vec { + let mut predicates = Vec::new(); for mat in matchers.matchers.iter() { if mat.name == TIMESTAMP_COL_NAME || mat.name == VALUE_LABEL { continue; } - let field_type = { - let schema = df.schema().as_arrow(); - let Ok(field) = schema.field_with_name(&mat.name) else { - continue; - }; - field.data_type().clone() + let Ok(field) = schema.field_with_name(&mat.name) else { + continue; }; + let field_type = field.data_type().clone(); let literal = |value: String| -> Expr { match &field_type { // Explicitly type equality matcher literals to the label column; @@ -52,13 +53,9 @@ pub fn apply_matchers(df: DataFrame, matchers: &Matchers) -> Result { _ => lit(value), } }; - match &mat.op { - MatchOp::Equal => { - df = df.filter(col(mat.name.clone()).eq(literal(mat.value.clone())))? - } - MatchOp::NotEqual => { - df = df.filter(col(mat.name.clone()).not_eq(literal(mat.value.clone())))? - } + let predicate = match &mat.op { + MatchOp::Equal => col(mat.name.clone()).eq(literal(mat.value.clone())), + MatchOp::NotEqual => col(mat.name.clone()).not_eq(literal(mat.value.clone())), MatchOp::Re(regex) => { let regex = format!("^{}$", regex.as_str()); // DataFusion 54 can lower a regex on Utf8View to a mixed-type @@ -69,8 +66,7 @@ pub fn apply_matchers(df: DataFrame, matchers: &Matchers) -> Result { } else { col(mat.name.clone()) }; - let predicate = regexp_like().call(vec![value, lit(regex)]); - df = df.filter(predicate)? + regexp_like().call(vec![value, lit(regex)]) } MatchOp::NotRe(regex) => { let regex = format!("^{}$", regex.as_str()); @@ -79,10 +75,19 @@ pub fn apply_matchers(df: DataFrame, matchers: &Matchers) -> Result { } else { col(mat.name.clone()) }; - let predicate = regexp_like().call(vec![value, lit(regex)]); - df = df.filter(predicate.not())? + regexp_like().call(vec![value, lit(regex)]).not() } - } + }; + predicates.push(predicate); + } + predicates +} + +pub fn apply_matchers(df: DataFrame, matchers: &Matchers) -> Result { + let predicates = matcher_predicates(df.schema().as_arrow(), matchers); + let mut df = df; + for predicate in predicates { + df = df.filter(predicate)?; } Ok(df) } diff --git a/src/search/Cargo.toml b/src/search/Cargo.toml index 8051dd68cf..78f3024805 100644 --- a/src/search/Cargo.toml +++ b/src/search/Cargo.toml @@ -65,3 +65,6 @@ utoipa.workspace = true vortex = { workspace = true } vortex-btrblocks = { workspace = true } vortex-datafusion = { workspace = true } + +[dev-dependencies] +tempfile.workspace = true diff --git a/src/search/src/datafusion/merge/downsampling.rs b/src/search/src/datafusion/merge/downsampling.rs index 2969e2c072..3163f33c5d 100644 --- a/src/search/src/datafusion/merge/downsampling.rs +++ b/src/search/src/datafusion/merge/downsampling.rs @@ -43,7 +43,7 @@ use vortex::{ use crate::datafusion::{ exec::DataFusionContextBuilder, - merge::{MergeParquetResult, append_metadata}, + merge::{MergeParquetResult, MergedFile, append_metadata}, sort_order::FileSortOrder, table_provider::uniontable::NewUnionTable, vortex::{VORTEX_RUNTIME, vortex_write_strategy}, @@ -127,11 +127,12 @@ pub async fn merge_parquet_files_with_downsampling( start.elapsed().as_millis() ); - Ok(MergeParquetResult::Multiple { - bufs, - file_metas, - file_format, - }) + let files = bufs + .into_iter() + .zip(file_metas) + .map(|(buf, meta)| MergedFile { buf, meta }) + .collect(); + Ok(MergeParquetResult { files, file_format }) } async fn write_downsampled_parquet( diff --git a/src/search/src/datafusion/merge/mod.rs b/src/search/src/datafusion/merge/mod.rs index 4cea126fba..a71e180cb9 100644 --- a/src/search/src/datafusion/merge/mod.rs +++ b/src/search/src/datafusion/merge/mod.rs @@ -31,7 +31,10 @@ use datafusion::{ physical_plan::execute_stream, }; use futures::TryStreamExt; -use parquet::{arrow::AsyncArrowWriter, file::metadata::KeyValue}; +use parquet::{ + arrow::{AsyncArrowWriter, async_writer::AsyncFileWriter}, + file::metadata::KeyValue, +}; use vortex::{ VortexSessionDefault, array::ArrayRef, @@ -57,18 +60,33 @@ use { o2_enterprise::enterprise::common::downsampling::get_largest_downsampling_rule, }; -pub enum MergeParquetResult { - Single { - buf: Vec, - file_meta: FileMeta, - file_format: FileFormat, - }, - #[allow(unused)] - Multiple { - bufs: Vec>, - file_metas: Vec, - file_format: FileFormat, - }, +/// One file written by [`merge_parquet_files`]. +pub struct MergedFile { + pub buf: Vec, + pub meta: FileMeta, +} + +pub struct MergeParquetResult { + pub files: Vec, + pub file_format: FileFormat, +} + +impl MergeParquetResult { + /// The merged file, for callers that always merge into exactly one file + /// (the ingester movers). + pub fn into_single(self) -> Result<(MergedFile, FileFormat)> { + let Self { + mut files, + file_format, + } = self; + if files.len() != 1 { + return Err(DataFusionError::Execution(format!( + "merge_parquet_files produced {} files, expected exactly one", + files.len() + ))); + } + Ok((files.pop().unwrap(), file_format)) + } } pub async fn merge_parquet_files( @@ -185,9 +203,11 @@ pub async fn merge_parquet_files( ); metadata.compressed_size = buf.len() as i64; - Ok(MergeParquetResult::Single { - buf, - file_meta: metadata, + Ok(MergeParquetResult { + files: vec![MergedFile { + buf, + meta: metadata, + }], file_format, }) } @@ -291,8 +311,8 @@ async fn write_vortex( .map_err(|e| DataFusionError::Execution(format!("Failed to write vortex file: {e}"))) } -pub fn append_metadata( - writer: &mut AsyncArrowWriter<&mut Vec>, +pub fn append_metadata( + writer: &mut AsyncArrowWriter, file_meta: &FileMeta, ) -> Result<()> { writer.append_key_value_metadata(KeyValue::new( @@ -425,10 +445,8 @@ mod tests { ) .await .unwrap(); - let MergeParquetResult::Single { buf, .. } = merged else { - panic!("trace time index merge must produce a single parquet file"); - }; - let batches = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buf)) + let (merged, _) = merged.into_single().unwrap(); + let batches = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(merged.buf)) .unwrap() .build() .unwrap() diff --git a/src/search/src/datafusion/table_provider/helpers.rs b/src/search/src/datafusion/table_provider/helpers.rs index 414e6cfe6c..3a2bba6758 100644 --- a/src/search/src/datafusion/table_provider/helpers.rs +++ b/src/search/src/datafusion/table_provider/helpers.rs @@ -68,6 +68,15 @@ pub fn generate_access_plan(file: &mut PartitionedFile) { file.extensions.insert(access_plan); } } + FileSelection::RowRanges(ranges) => { + if let Some(access_plan) = generate_parquet_access_plan_from_ranges( + file, + ranges.iter().cloned(), + row_group_size, + ) { + file.extensions.insert(access_plan); + } + } #[cfg(feature = "enterprise")] FileSelection::RowGroups(row_group_ids) => { if let Some((num_rows, row_group_size)) = parquet_file_layout(file, row_group_size) @@ -89,8 +98,9 @@ pub fn generate_access_plan(file: &mut PartitionedFile) { file.extensions.insert(access_plan); } } - // row-group sampling is parquet only; vortex falls back to a full scan - FileSelection::RowGroups(_) => {} + // compact row ranges and row-group sampling are parquet only; + // vortex falls back to a full scan + FileSelection::RowRanges(_) | FileSelection::RowGroups(_) => {} }, } } @@ -99,6 +109,21 @@ fn generate_parquet_access_plan( file: &PartitionedFile, row_ids: &BooleanBuffer, row_group_size: Option, +) -> Option { + generate_parquet_access_plan_from_ranges( + file, + row_ids.set_slices().map(|(start, end)| start..end), + row_group_size, + ) +} + +/// Build a Parquet access plan directly from sorted, non-overlapping row +/// ranges. This avoids allocating a bitmap proportional to a high-cardinality +/// metrics file when a series index already stores compact contiguous runs. +fn generate_parquet_access_plan_from_ranges( + file: &PartitionedFile, + ranges: impl IntoIterator>, + row_group_size: Option, ) -> Option { let (num_rows, row_group_size) = parquet_file_layout(file, row_group_size)?; let row_group_count = num_rows.div_ceil(row_group_size); @@ -107,12 +132,15 @@ fn generate_parquet_access_plan( let mut row_group_selection = RowGroupSelectionBuilder::new(&mut access_plan, row_group_size, num_rows); - for (start, end) in row_ids.set_slices() { - if end > num_rows { - unreachable!("row_ids set slice end {end} exceeds num_rows {num_rows}"); - } - - row_group_selection.push_selected_range(start, end); + let mut previous_end = 0; + for range in ranges { + assert!( + previous_end <= range.start && range.start < range.end && range.end <= num_rows, + "invalid sorted row range {range:?} for file {} with {num_rows} rows", + file.path().as_ref() + ); + previous_end = range.end; + row_group_selection.push_selected_range(range.start, range.end); } row_group_selection.finish(); @@ -344,6 +372,44 @@ mod tests { assert_eq!(plan, expected); } + #[test] + fn test_generate_parquet_access_plan_from_series_ranges() { + let file = make_partitioned_file(12); + let plan = + generate_parquet_access_plan_from_ranges(&file, [1..3, 5..8, 10..12], Some(4)).unwrap(); + + let mut expected = ParquetAccessPlan::new_none(3); + expected.scan(0); + expected.scan_selection( + 0, + RowSelection::from(vec![ + RowSelector::skip(1), + RowSelector::select(2), + RowSelector::skip(1), + ]), + ); + expected.scan(1); + expected.scan_selection( + 1, + RowSelection::from(vec![RowSelector::skip(1), RowSelector::select(3)]), + ); + expected.scan(2); + expected.scan_selection( + 2, + RowSelection::from(vec![RowSelector::skip(2), RowSelector::select(2)]), + ); + assert_eq!(plan, expected); + } + + #[test] + #[should_panic(expected = "invalid sorted row range")] + fn test_generate_parquet_access_plan_rejects_out_of_bounds_row_ids() { + let file = make_partitioned_file(4); + let row_ids = BooleanBuffer::from_iter((0..5).map(|i| i == 4)); + + generate_parquet_access_plan(&file, &row_ids, Some(4)); + } + #[test] fn test_generate_parquet_access_plan_keeps_disjoint_ranges_in_same_row_group() { let file = make_partitioned_file(8); diff --git a/src/search/src/datafusion/table_provider/listing_adapter.rs b/src/search/src/datafusion/table_provider/listing_adapter.rs index b2498b214a..479c0112e7 100644 --- a/src/search/src/datafusion/table_provider/listing_adapter.rs +++ b/src/search/src/datafusion/table_provider/listing_adapter.rs @@ -285,3 +285,144 @@ fn handler_tantivy_index( } plan } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{Float64Array, Int64Array, RecordBatch, UInt64Array}; + use arrow_schema::{DataType, Field, Schema}; + use config::meta::promql::HASH_LABEL; + use datafusion::{ + datasource::{ + file_format::parquet::ParquetFormat, + listing::{ListingOptions, ListingTableUrl}, + }, + physical_plan::{collect, displayable}, + }; + use parquet::arrow::ArrowWriter; + + use super::*; + use crate::datafusion::exec::DataFusionContextBuilder; + + fn hash_sorted_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new(HASH_LABEL, DataType::UInt64, false), + Field::new(TIMESTAMP_COL_NAME, DataType::Int64, false), + Field::new("value", DataType::Float64, false), + ])) + } + + /// Write one Parquet file whose rows are ordered by (__hash__, _timestamp). + fn write_hash_sorted_file(dir: &std::path::Path, name: &str, rows: &[(u64, i64)]) { + let schema = hash_sorted_schema(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from_iter_values(rows.iter().map(|r| r.0))), + Arc::new(Int64Array::from_iter_values(rows.iter().map(|r| r.1))), + Arc::new(Float64Array::from_iter_values( + rows.iter().map(|r| r.0 as f64), + )), + ], + ) + .unwrap(); + let file = std::fs::File::create(dir.join(name)).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + + /// Hash-sorted files with overlapping hash ranges (every ingester file + /// covers the whole hash space) must still yield an ordered scan: one file + /// per partition, merged by SortPreservingMergeExec without a SortExec. + #[tokio::test] + async fn test_hash_sorted_files_merge_without_sort_exec() { + let dir = tempfile::tempdir().unwrap(); + write_hash_sorted_file( + dir.path(), + "a.parquet", + &[(1, 10), (1, 30), (5, 10), (9, 20)], + ); + write_hash_sorted_file(dir.path(), "b.parquet", &[(1, 20), (2, 10), (9, 10)]); + write_hash_sorted_file(dir.path(), "c.parquet", &[(3, 10), (5, 5), (5, 20)]); + + let sort_order = FileSortOrder::HashTimestampAsc; + let ctx = DataFusionContextBuilder::new() + .trace_id("test_hash_sorted_merge") + .sort_order(sort_order) + .build(2) + .await + .unwrap(); + + let listing_options = ListingOptions::new(Arc::new(ParquetFormat::default())) + .with_target_partitions(2) + .with_collect_stat(true) + .with_file_sort_order(vec![sort_order.logical_sort_exprs()]); + let url = ListingTableUrl::parse(format!("file://{}/", dir.path().display())).unwrap(); + let config = ListingTableConfig::new(url) + .with_listing_options(listing_options) + .with_schema(hash_sorted_schema()); + let table = ListingTableAdapter::try_new( + config, + "test_hash_sorted_merge".to_string(), + sort_order, + None, + vec![], + None, + ) + .unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + + let plan = ctx + .state() + .create_logical_plan(&format!( + "SELECT * FROM t ORDER BY {}", + sort_order.order_by_clause().unwrap() + )) + .await + .unwrap(); + let physical_plan = ctx.state().create_physical_plan(&plan).await.unwrap(); + let display = displayable(physical_plan.as_ref()).indent(true).to_string(); + assert!( + display.contains("SortPreservingMergeExec"), + "expected a merge of pre-sorted partitions, got:\n{display}" + ); + assert!( + !display.contains("SortExec"), + "hash-sorted inputs must not be re-sorted, got:\n{display}" + ); + + let batches = collect(physical_plan, ctx.task_ctx()).await.unwrap(); + let mut rows = Vec::new(); + for batch in &batches { + let hashes = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let ts = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + rows.push((hashes.value(i), ts.value(i))); + } + } + let mut expected = vec![ + (1, 10), + (1, 30), + (5, 10), + (9, 20), + (1, 20), + (2, 10), + (9, 10), + (3, 10), + (5, 5), + (5, 20), + ]; + expected.sort(); + assert_eq!(rows, expected); + } +}