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

refactor: merge result / access-plan / deletion helpers groundwork for TSID-major metrics (#13861)

Browse Source 
Behavior-neutral refactors split out of #13839 so that PR only carries
the TSID-major feature. Nothing here changes what is written, read or
deleted; every commit is independently reviewable.

## Commits

- `cba75d244b` **promql**: `apply_matchers` split into
`matcher_predicates(schema, matchers) -> Vec<Expr>` + `apply_matchers`
(unchanged behavior; lets other inputs evaluate the same matcher
semantics).
- `b5ff6fed21` **parquet**: `new_parquet_writer` / `append_metadata`
generic over `W: AsyncFileWriter` instead of `&mut Vec<u8>`; existing
callers unchanged.
- `b07d84216a` **merge**: `MergeParquetResult::{Single, Multiple}` →
`MergeParquetResult { files: Vec<MergedFile { buf, meta }>, file_format
}` + `into_single()`. The compactor's two copies of the upload / index /
push sequence become one loop (`compressed_size` is set from the buffer
for every file, as the `Multiple` arm already did); downsampling maps
its `(bufs, metas)` into it; ingester movers use `into_single()` instead
of panicking/erroring on `Multiple`.
- `014b6c5721` **compactor**: `deleted.rs` `derived_files()` /
`delete_from_storage()` replace three copies of collect / `storage::del`
/ ignore-not-found (data, `.ttv`, flattened).
- `7163786362` **search**: `generate_parquet_access_plan_from_ranges`
extracted from the bitmap path (`set_slices` → ranges); an out-of-range,
empty or unsorted range is rejected with `assert!` (`56de386522`; a
producer that violates the sorted/non-overlapping contract is a bug, not
something to degrade around) instead of the previous `unreachable!()` on
the bitmap path. Adds `FileSelection::RowRanges(Arc<Vec<Range<usize>>>)`
as the matching selection kind (no producer yet; Parquet only, Vortex
full scan like `RowGroups`).
- `754014ccf3` **test**: `ListingTableAdapter` plan-shape test for
`FileSortOrder::HashTimestampAsc` (from #13853): three overlapping
hash-sorted files → `SortPreservingMergeExec` over three single-file
groups, no `SortExec`, globally ordered rows. `tempfile` as `search`
dev-dependency.

## Validation

- `cargo fmt --check`
- `cargo clippy -p config -p promql -p search -p compaction -p
openobserve-jobs -p search_service -p promql-service --tests` clean
- `cargo test -p promql --lib -- utils` 10, `-p search --lib -- merge::
helpers listing_adapter` 14, `-p compaction --lib -- merge deleted` 31,
`-p config --lib -- test_file_key` 4 passed
- `merge/downsampling.rs` is enterprise-only and was updated by hand;
needs the enterprise CI compile.
This commit is contained in:
Huaijin
2026-08-18 10:54:23 +00:00
committed by GitHub
parent a6c910f015
commit 4e805038fe
13 changed files with 446 additions and 285 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1
View File
@@ -10961,6 +10961,7 @@ dependencies = [
"sqlparser",
"tantivy",
"tantivy_utils",
"tempfile",
"thiserror 2.0.19",
"tokio",
"tokio-stream",
src/compaction/src/deleted.rs
+62 -75
View File
@@ -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<String>,
) -> 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::<Vec<_>>(),
)
.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<i64, anyhow::Error> {
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<i64, anyhow::Error> {
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::<Vec<_>>(),
.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::<Vec<_>>();
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::<Vec<_>>(),
)
.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::<Vec<_>>();
if !flattened_files.is_empty()
&& let Err(e) = storage::del(
flattened_files
.iter()
.map(|file| (file.0.as_str(), file.1.as_str()))
.collect::<Vec<_>>(),
)
.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(
src/compaction/src/merge.rs
+62 -117
View File
@@ -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::<Vec<_>>(),
new_files.iter().map(|f| f.meta.original_size).sum::<i64>(),
new_files
.iter()
.map(|f| f.meta.compressed_size)
.sum::<i64>(),
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::<Vec<_>>(),
new_files.iter().map(|f| f.meta.original_size).sum::<i64>(),
new_files
.iter()
.map(|f| f.meta.compressed_size)
.sum::<i64>(),
start.elapsed().as_millis(),
);
Ok((new_files, retain_file_list))
}
src/config/src/meta/stream.rs
+13 -1
View File
@@ -13,7 +13,7 @@
// 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::{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<BooleanBuffer>),
/// Sorted, non-overlapping physical row ranges selected by a compact
/// secondary index. This avoids materializing one bitmap bit per row.
RowRanges(Arc<Vec<Range<usize>>>),
/// Row group ids selected by row-group-level sampling.
RowGroups(Arc<Vec<u32>>),
}
@@ -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);
src/config/src/utils/parquet.rs
+10 -7
View File
@@ -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<u8> {
.expect("file meta is always serializable")
}
pub fn new_parquet_writer<'a>(
buf: &'a mut Vec<u8>,
schema: &'a Arc<Schema>,
bloom_filter_fields: &'a [String],
metadata: &'a FileMeta,
pub fn new_parquet_writer<W: AsyncFileWriter>(
buf: W,
schema: &Arc<Schema>,
bloom_filter_fields: &[String],
metadata: &FileMeta,
write_metadata: bool,
compression: Option<&str>,
) -> AsyncArrowWriter<&'a mut Vec<u8>> {
) -> AsyncArrowWriter<W> {
let cfg = get_config();
let compression = compression.unwrap_or(&cfg.common.parquet_compression);
let mut writer_props = WriterProperties::builder()
src/jobs/src/job/files/pack.rs
+4 -13
View File
@@ -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 {
src/jobs/src/job/files/parquet.rs
+4 -16
View File
@@ -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!(
src/promql/src/utils.rs
+25 -20
View File
@@ -30,19 +30,20 @@ use datafusion::{
use hashbrown::HashSet;
use promql_parser::label::{MatchOp, Matchers};
pub fn apply_matchers(df: DataFrame, matchers: &Matchers) -> Result<DataFrame> {
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<Expr> {
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<DataFrame> {
_ => 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<DataFrame> {
} 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<DataFrame> {
} 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<DataFrame> {
let predicates = matcher_predicates(df.schema().as_arrow(), matchers);
let mut df = df;
for predicate in predicates {
df = df.filter(predicate)?;
}
Ok(df)
}
src/search/Cargo.toml
+3
View File
@@ -65,3 +65,6 @@ utoipa.workspace = true
vortex = { workspace = true }
vortex-btrblocks = { workspace = true }
vortex-datafusion = { workspace = true }
[dev-dependencies]
tempfile.workspace = true
src/search/src/datafusion/merge/downsampling.rs
+7 -6
View File
@@ -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(
src/search/src/datafusion/merge/mod.rs
+40 -22
View File
@@ -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<u8>,
file_meta: FileMeta,
file_format: FileFormat,
},
#[allow(unused)]
Multiple {
bufs: Vec<Vec<u8>>,
file_metas: Vec<FileMeta>,
file_format: FileFormat,
},
/// One file written by [`merge_parquet_files`].
pub struct MergedFile {
pub buf: Vec<u8>,
pub meta: FileMeta,
}
pub struct MergeParquetResult {
pub files: Vec<MergedFile>,
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<u8>>,
pub fn append_metadata<W: AsyncFileWriter>(
writer: &mut AsyncArrowWriter<W>,
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()
src/search/src/datafusion/table_provider/helpers.rs
+74 -8
View File
@@ -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<u32>,
) -> Option<ParquetAccessPlan> {
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<Item = Range<usize>>,
row_group_size: Option<u32>,
) -> Option<ParquetAccessPlan> {
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);
src/search/src/datafusion/table_provider/listing_adapter.rs
+141
View File
@@ -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<Schema> {
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::<UInt64Array>()
.unwrap();
let ts = batch
.column(1)
.as_any()
.downcast_ref::<Int64Array>()
.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);
}
}