This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 06826f10 fix(datafusion): reduce scan planning peak memory (#666)
06826f10 is described below
commit 06826f107c9e3c70c3f2086066a28da57a410583
Author: XiaoHongbo <[email protected]>
AuthorDate: Tue Aug 4 11:32:04 2026 +0800
fix(datafusion): reduce scan planning peak memory (#666)
---
.../datafusion/src/full_text_search.rs | 2 +-
.../integrations/datafusion/src/hybrid_search.rs | 2 +-
crates/integrations/datafusion/src/table/mod.rs | 6 +-
.../datafusion/src/variant_pushdown.rs | 2 +-
crates/paimon/Cargo.toml | 2 +-
crates/paimon/src/table/source.rs | 81 ++++++++--
crates/paimon/src/table/table_scan.rs | 175 ++++++++++-----------
7 files changed, 162 insertions(+), 108 deletions(-)
diff --git a/crates/integrations/datafusion/src/full_text_search.rs
b/crates/integrations/datafusion/src/full_text_search.rs
index 66a9c773..9173e9fb 100644
--- a/crates/integrations/datafusion/src/full_text_search.rs
+++ b/crates/integrations/datafusion/src/full_text_search.rs
@@ -185,7 +185,7 @@ impl TableProvider for FullTextSearchTableProvider {
PaimonScanBuilder {
table,
schema: &self.schema(),
- plan: &plan,
+ plan,
scan_trace: None,
projection,
pushed_predicate: None,
diff --git a/crates/integrations/datafusion/src/hybrid_search.rs
b/crates/integrations/datafusion/src/hybrid_search.rs
index b37e4caa..67b2d032 100644
--- a/crates/integrations/datafusion/src/hybrid_search.rs
+++ b/crates/integrations/datafusion/src/hybrid_search.rs
@@ -287,7 +287,7 @@ impl TableProvider for HybridSearchTableProvider {
let input = PaimonScanBuilder {
table,
schema: &input_schema,
- plan: &plan,
+ plan,
scan_trace: None,
projection: Some(&input_projection),
pushed_predicate: None,
diff --git a/crates/integrations/datafusion/src/table/mod.rs
b/crates/integrations/datafusion/src/table/mod.rs
index 113c3bb8..9e3007e2 100644
--- a/crates/integrations/datafusion/src/table/mod.rs
+++ b/crates/integrations/datafusion/src/table/mod.rs
@@ -323,7 +323,7 @@ pub(crate) fn bucket_round_robin<T>(items: Vec<T>,
num_buckets: usize) -> Vec<Ve
pub(crate) struct PaimonScanBuilder<'a> {
pub(crate) table: &'a Table,
pub(crate) schema: &'a ArrowSchemaRef,
- pub(crate) plan: &'a paimon::table::Plan,
+ pub(crate) plan: paimon::table::Plan,
pub(crate) scan_trace: Option<paimon::table::ScanTrace>,
pub(crate) projection: Option<&'a Vec<usize>>,
pub(crate) pushed_predicate: Option<paimon::spec::Predicate>,
@@ -361,7 +361,7 @@ impl PaimonScanBuilder<'_> {
(self.schema.clone(), read_fields)
};
- let splits = self.plan.splits().to_vec();
+ let splits = self.plan.into_splits();
let planned_partitions: Vec<Arc<[_]>> = if splits.is_empty() {
vec![Arc::from(Vec::new())]
} else {
@@ -458,7 +458,7 @@ impl TableProvider for PaimonTableProvider {
PaimonScanBuilder {
table: &self.table,
schema: &self.schema,
- plan: &plan,
+ plan,
scan_trace: Some(scan_trace),
projection,
pushed_predicate: filter_analysis.pushed_predicate,
diff --git a/crates/integrations/datafusion/src/variant_pushdown.rs
b/crates/integrations/datafusion/src/variant_pushdown.rs
index 641c6a5a..ba3c7cb1 100644
--- a/crates/integrations/datafusion/src/variant_pushdown.rs
+++ b/crates/integrations/datafusion/src/variant_pushdown.rs
@@ -231,7 +231,7 @@ impl ExtensionPlanner for VariantExtractionExtensionPlanner
{
.await
.map_err(to_datafusion_error)?;
- let splits = plan.splits().to_vec();
+ let splits = plan.into_splits();
let target =
session_state.config_options().execution.target_partitions;
let planned_partitions: Vec<Arc<[_]>> = if splits.is_empty() {
vec![Arc::from(Vec::new())]
diff --git a/crates/paimon/Cargo.toml b/crates/paimon/Cargo.toml
index 22fb72d1..5922270d 100644
--- a/crates/paimon/Cargo.toml
+++ b/crates/paimon/Cargo.toml
@@ -62,7 +62,7 @@ bytes = "1.7.1"
bitflags = "2.6.0"
tokio = { version = "1.39.2", features = ["fs", "io-util", "macros", "sync",
"time"] }
chrono = { version = "0.4.38", features = ["serde"] }
-serde = { version = "1", features = ["derive"] }
+serde = { version = "1", features = ["derive", "rc"] }
serde_bytes = "0.11.15"
serde_json = "1.0.120"
serde_with = "3.9.0"
diff --git a/crates/paimon/src/table/source.rs
b/crates/paimon/src/table/source.rs
index a47dc582..2c1e39bc 100644
--- a/crates/paimon/src/table/source.rs
+++ b/crates/paimon/src/table/source.rs
@@ -22,6 +22,7 @@
use crate::spec::{BinaryRow, DataFileMeta};
use crate::table::stats_filter::group_by_overlapping_row_id;
use serde::{Deserialize, Serialize};
+use std::sync::Arc;
fn is_vector_store_file_name(file_name: &str) -> bool {
file_name.to_ascii_lowercase().contains(".vector.")
@@ -478,19 +479,22 @@ impl PartitionBucket {
/// Input split for reading: partition + bucket + list of data files and
optional deletion files.
///
+/// The metadata collections use shared storage so cloning a planned split for
an
+/// asynchronous reader does not deep-copy every [`DataFileMeta`].
+///
/// Reference:
[org.apache.paimon.table.source.DataSplit](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/table/source/DataSplit.java)
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct DataSplit {
snapshot_id: i64,
- partition: BinaryRow,
+ partition: Arc<BinaryRow>,
bucket: i32,
- bucket_path: String,
+ bucket_path: Arc<str>,
total_buckets: i32,
- data_files: Vec<DataFileMeta>,
+ data_files: Arc<[DataFileMeta]>,
/// Deletion file for each data file, same order as `data_files`.
/// `None` at index `i` means no deletion file for `data_files[i]`
(matches Java getDeletionFiles() / List<DeletionFile> with null elements).
- data_deletion_files: Option<Vec<Option<DeletionFile>>>,
- row_ranges: Option<Vec<RowRange>>,
+ data_deletion_files: Option<Arc<[Option<DeletionFile>]>>,
+ row_ranges: Option<Arc<[RowRange]>>,
/// Whether the split can be read raw, without the merge reader: its
/// physical rows are exactly its logical rows (modulo deletion files).
/// Mirrors Java `DataSplit#rawConvertible`.
@@ -696,7 +700,7 @@ impl DataSplit {
out.extend_from_slice(&0i32.to_be_bytes()); // deprecated beforeFiles
count
out.push(0); // beforeDeletionFiles = null list
out.extend_from_slice(&(self.data_files.len() as i32).to_be_bytes());
- for f in &self.data_files {
+ for f in self.data_files.iter() {
let d = f.to_serialized_row_data()?;
out.extend_from_slice(&(d.len() as i32).to_be_bytes());
out.extend_from_slice(&d);
@@ -847,7 +851,7 @@ impl DataSplit {
out.extend_from_slice(&INDEXED_SPLIT_VERSION.to_be_bytes());
out.extend_from_slice(&self.serialize()?);
out.extend_from_slice(&(ranges.len() as i32).to_be_bytes());
- for r in ranges {
+ for r in ranges.iter() {
out.extend_from_slice(&r.from().to_be_bytes());
out.extend_from_slice(&r.to().to_be_bytes());
}
@@ -921,7 +925,7 @@ impl DataSplit {
});
}
let mut body = body;
- body.row_ranges = Some(ranges);
+ body.row_ranges = Some(ranges.into());
body
}
other => {
@@ -1255,13 +1259,13 @@ impl DataSplitBuilder {
}
Ok(DataSplit {
snapshot_id: self.snapshot_id,
- partition,
+ partition: Arc::new(partition),
bucket: self.bucket,
- bucket_path,
+ bucket_path: bucket_path.into(),
total_buckets: self.total_buckets,
- data_files,
- data_deletion_files: self.data_deletion_files,
- row_ranges: self.row_ranges,
+ data_files: data_files.into(),
+ data_deletion_files: self.data_deletion_files.map(Into::into),
+ row_ranges: self.row_ranges.map(Into::into),
raw_convertible: self.raw_convertible,
})
}
@@ -1290,6 +1294,12 @@ impl Plan {
pub fn splits(&self) -> &[DataSplit] {
&self.splits
}
+
+ /// Consume this plan and return its splits without cloning their file
metadata.
+ #[must_use = "consuming a plan without using its splits drops the planned
work"]
+ pub fn into_splits(self) -> Vec<DataSplit> {
+ self.splits
+ }
}
#[cfg(test)]
@@ -1335,6 +1345,51 @@ mod tests {
.unwrap()
}
+ #[test]
+ fn data_split_clone_shares_planned_metadata() {
+ let split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("file:/tmp/bucket-0".to_string())
+ .with_total_buckets(1)
+ .with_data_files(vec![file("a.parquet", 10, Some(0))])
+ .with_data_deletion_files(vec![Some(DeletionFile::new(
+ "file:/tmp/a.dv".to_string(),
+ 0,
+ 64,
+ Some(1),
+ ))])
+ .with_row_ranges(vec![RowRange::new(0, 9)])
+ .build()
+ .unwrap();
+
+ let cloned = split.clone();
+
+ assert!(Arc::ptr_eq(&split.partition, &cloned.partition));
+ assert!(Arc::ptr_eq(&split.bucket_path, &cloned.bucket_path));
+ assert!(Arc::ptr_eq(&split.data_files, &cloned.data_files));
+ assert!(Arc::ptr_eq(
+ split.data_deletion_files.as_ref().unwrap(),
+ cloned.data_deletion_files.as_ref().unwrap()
+ ));
+ assert!(Arc::ptr_eq(
+ split.row_ranges.as_ref().unwrap(),
+ cloned.row_ranges.as_ref().unwrap()
+ ));
+ }
+
+ #[test]
+ fn plan_into_splits_preserves_metadata_ownership() {
+ let split = split(vec![file("a.parquet", 10, Some(0))], true);
+ let data_files = Arc::clone(&split.data_files);
+
+ let splits = Plan::new(vec![split]).into_splits();
+
+ assert_eq!(splits.len(), 1);
+ assert!(Arc::ptr_eq(&data_files, &splits[0].data_files));
+ }
+
#[test]
fn data_split_serde_json_round_trip() {
let split = DataSplit::builder()
diff --git a/crates/paimon/src/table/table_scan.rs
b/crates/paimon/src/table/table_scan.rs
index ded0cab0..8ce0c4a6 100644
--- a/crates/paimon/src/table/table_scan.rs
+++ b/crates/paimon/src/table/table_scan.rs
@@ -175,98 +175,99 @@ async fn read_all_manifest_entries(
let manifest_path_prefix = format!("{}/{}",
table_path.trim_end_matches('/'), MANIFEST_DIR);
let shared_cache = SharedSchemaCache::new();
- let manifest_results: Vec<(Vec<ManifestEntry>, ManifestReadCounters)> =
- futures::stream::iter(manifest_files)
- .map(|meta| {
- let path = format!("{}/{}", manifest_path_prefix,
meta.file_name());
- let cache = shared_cache.clone();
- async move {
- let input_file = file_io.new_input(&path)?;
- let content = input_file.read().await?;
-
- // Per-task bucket cache (few distinct total_buckets
values per manifest).
- let mut bucket_cache: HashMap<i32, Option<HashSet<i32>>> =
HashMap::new();
- let mut counters = ManifestReadCounters::default();
-
- let entries =
crate::spec::avro::from_manifest_bytes_filtered_shared(
- &content,
- &cache,
- &mut |_kind, partition_bytes, bucket, total_buckets| {
- counters.entries_read += 1;
- // Bucket filter (negative bucket = unassigned)
- if has_primary_keys && !scan_all_files && bucket <
0 {
- counters.pruned_by_bucket += 1;
- return false;
- }
- if let Some(pred) = bucket_predicate {
- let targets =
-
bucket_cache.entry(total_buckets).or_insert_with(|| {
- compute_target_buckets(
- pred,
- bucket_key_fields,
- bucket_function_type,
- total_buckets,
- )
- });
- if let Some(targets) = targets {
- if !targets.contains(&bucket) {
- counters.pruned_by_bucket += 1;
- return false;
- }
+ let (all_entries, mut counters) = futures::stream::iter(manifest_files)
+ .map(|meta| {
+ let path = format!("{}/{}", manifest_path_prefix,
meta.file_name());
+ let cache = shared_cache.clone();
+ async move {
+ let input_file = file_io.new_input(&path)?;
+ let content = input_file.read().await?;
+
+ // Per-task bucket cache (few distinct total_buckets values
per manifest).
+ let mut bucket_cache: HashMap<i32, Option<HashSet<i32>>> =
HashMap::new();
+ let mut counters = ManifestReadCounters::default();
+
+ let entries =
crate::spec::avro::from_manifest_bytes_filtered_shared(
+ &content,
+ &cache,
+ &mut |_kind, partition_bytes, bucket, total_buckets| {
+ counters.entries_read += 1;
+ // Bucket filter (negative bucket = unassigned)
+ if has_primary_keys && !scan_all_files && bucket < 0 {
+ counters.pruned_by_bucket += 1;
+ return false;
+ }
+ if let Some(pred) = bucket_predicate {
+ let targets =
bucket_cache.entry(total_buckets).or_insert_with(|| {
+ compute_target_buckets(
+ pred,
+ bucket_key_fields,
+ bucket_function_type,
+ total_buckets,
+ )
+ });
+ if let Some(targets) = targets {
+ if !targets.contains(&bucket) {
+ counters.pruned_by_bucket += 1;
+ return false;
}
}
+ }
- // Partition filter
- if let Some(pf) = partition_filter {
- match pf.matches_entry(partition_bytes) {
- Ok(false) => {
- counters.pruned_by_partition += 1;
- return false;
- }
- Ok(true) => {}
- Err(_) => {}
+ // Partition filter
+ if let Some(pf) = partition_filter {
+ match pf.matches_entry(partition_bytes) {
+ Ok(false) => {
+ counters.pruned_by_partition += 1;
+ return false;
}
+ Ok(true) => {}
+ Err(_) => {}
}
-
- true
- },
- )?;
- counters.after_entry_pruning = entries.len();
-
- // Post-filter: level-0 and data predicates (need
DataFileMeta)
- let mut filtered = Vec::with_capacity(entries.len());
- for entry in entries {
- if skip_level_zero && has_primary_keys &&
entry.file().level == 0 {
- counters.pruned_by_level += 1;
- continue;
- }
- if !data_predicates.is_empty()
- && !data_file_matches_predicates(
- entry.file(),
- data_predicates,
- current_schema_id,
- schema_fields,
- )
- {
- counters.pruned_by_data_stats += 1;
- continue;
}
- filtered.push(entry);
+
+ true
+ },
+ )?;
+ counters.after_entry_pruning = entries.len();
+
+ // Post-filter: level-0 and data predicates (need DataFileMeta)
+ let mut filtered = Vec::with_capacity(entries.len());
+ for entry in entries {
+ if skip_level_zero && has_primary_keys &&
entry.file().level == 0 {
+ counters.pruned_by_level += 1;
+ continue;
}
- counters.after_manifest_filters = filtered.len();
- Ok::<_, crate::Error>((filtered, counters))
+ if !data_predicates.is_empty()
+ && !data_file_matches_predicates(
+ entry.file(),
+ data_predicates,
+ current_schema_id,
+ schema_fields,
+ )
+ {
+ counters.pruned_by_data_stats += 1;
+ continue;
+ }
+ filtered.push(entry);
}
- })
- .buffered(64)
- .try_collect::<Vec<_>>()
- .await?;
-
- let mut counters = ManifestReadCounters::default();
- let mut all_entries = Vec::new();
- for (entries, manifest_counters) in manifest_results {
- counters.merge(manifest_counters);
- all_entries.extend(entries);
- }
+ counters.after_manifest_filters = filtered.len();
+ Ok::<_, crate::Error>((filtered, counters))
+ }
+ })
+ // Keep manifest read concurrency bounded. `try_fold` releases each
+ // yielded result after merging it, so peak retained results are the
+ // accumulator plus at most this bounded set of in-flight reads.
+ .buffered(64)
+ .try_fold(
+ (Vec::new(), ManifestReadCounters::default()),
+ |(mut all_entries, mut counters), (entries, manifest_counters)|
async move {
+ counters.merge(manifest_counters);
+ all_entries.extend(entries);
+ Ok((all_entries, counters))
+ },
+ )
+ .await?;
let mut all_entries = merge_manifest_entries(all_entries);
let manifest_entries_after_merge = all_entries.len();
if let Some(index) = row_range_index {
@@ -458,17 +459,15 @@ fn build_deletion_files_map(
///
/// Collecting deletes first (rather than insert/remove while iterating) makes
/// the result independent of ADD/DELETE ordering, matching the Java scan path.
-fn merge_manifest_entries(entries: Vec<ManifestEntry>) -> Vec<ManifestEntry> {
+fn merge_manifest_entries(mut entries: Vec<ManifestEntry>) ->
Vec<ManifestEntry> {
use crate::spec::Identifier;
let deleted: HashSet<Identifier> = entries
.iter()
.filter(|e| *e.kind() == FileKind::Delete)
.map(|e| e.identifier())
.collect();
+ entries.retain(|e| *e.kind() == FileKind::Add &&
!deleted.contains(&e.identifier()));
entries
- .into_iter()
- .filter(|e| *e.kind() == FileKind::Add &&
!deleted.contains(&e.identifier()))
- .collect()
}
/// Whether scan-owned pruning still preserves `merged_row_count()` as a safe
