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 81d2986d perf(datafusion): honor Paimon read.batch-size (#547)
81d2986d is described below
commit 81d2986d2bdfb6524ecfd38a233b1ef564a9389b
Author: Jingsong Lee <[email protected]>
AuthorDate: Sun Jul 19 22:49:18 2026 +0800
perf(datafusion): honor Paimon read.batch-size (#547)
---
.../datafusion/src/physical_plan/scan.rs | 69 ++++++++++++
crates/integrations/datafusion/src/sql_context.rs | 58 ++++++++++
crates/paimon/src/spec/core_options.rs | 42 +++++++
crates/paimon/src/spec/schema.rs | 53 ++++++++-
crates/paimon/src/table/data_evolution_reader.rs | 125 ++++++++++++++++++++-
crates/paimon/src/table/data_file_reader.rs | 10 +-
crates/paimon/src/table/format_read_builder.rs | 5 +-
crates/paimon/src/table/format_table_read.rs | 7 +-
crates/paimon/src/table/kv_file_reader.rs | 97 +++++++++++++++-
crates/paimon/src/table/read_builder.rs | 3 +-
crates/paimon/src/table/table_read.rs | 18 +--
11 files changed, 470 insertions(+), 17 deletions(-)
diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs
b/crates/integrations/datafusion/src/physical_plan/scan.rs
index d31891ed..25709034 100644
--- a/crates/integrations/datafusion/src/physical_plan/scan.rs
+++ b/crates/integrations/datafusion/src/physical_plan/scan.rs
@@ -486,4 +486,73 @@ mod tests {
assert_eq!(actual_ids, vec![2, 3, 4]);
}
+
+ #[tokio::test]
+ async fn test_execute_uses_read_batch_size_option() {
+ let tempdir = tempdir().unwrap();
+ let table_path = local_file_path(tempdir.path());
+ let bucket_dir = tempdir.path().join("bucket-0");
+ fs::create_dir_all(&bucket_dir).unwrap();
+
+ write_int_parquet_file(
+ &bucket_dir.join("data.parquet"),
+ vec![("id", vec![1, 2, 3, 4, 5])],
+ None,
+ );
+ let file_size =
fs::metadata(bucket_dir.join("data.parquet")).unwrap().len() as i64;
+
+ let file_io = FileIOBuilder::new("file").build().unwrap();
+ let table_schema = TableSchema::new(
+ 0,
+ &PaimonSchema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .option("read.batch-size", "2")
+ .build()
+ .unwrap(),
+ );
+ let table = Table::new(
+ file_io,
+ Identifier::new("default", "t"),
+ table_path,
+ table_schema,
+ None,
+ );
+ let split = paimon::DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(local_file_path(&bucket_dir))
+ .with_total_buckets(1)
+ .with_data_files(vec![test_data_file("data.parquet", 5,
file_size)])
+ .build()
+ .unwrap();
+ let scan = PaimonTableScan::new(
+ test_schema(),
+ table,
+ test_read_type(),
+ None,
+ vec![Arc::from(vec![split])],
+ None,
+ false,
+ None,
+ None,
+ true,
+ );
+
+ let ctx = SessionContext::new();
+ let batches = scan
+ .execute(0, ctx.task_ctx())
+ .expect("execute should succeed")
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert_eq!(
+ batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>(),
+ vec![2, 2, 1]
+ );
+ }
}
diff --git a/crates/integrations/datafusion/src/sql_context.rs
b/crates/integrations/datafusion/src/sql_context.rs
index dedd7ebc..3190487b 100644
--- a/crates/integrations/datafusion/src/sql_context.rs
+++ b/crates/integrations/datafusion/src/sql_context.rs
@@ -6626,6 +6626,64 @@ mod tests {
(temp_dir, sql_context)
}
+ #[tokio::test]
+ async fn test_dynamic_read_batch_size_overrides_table_option() {
+ let (_tmp, sql_context) = setup_fs_sql_context().await;
+
+ sql_context
+ .sql(
+ "CREATE TABLE paimon.test_db.batch_size_t (id INT) \
+ WITH ('read.batch-size' = '3')",
+ )
+ .await
+ .unwrap();
+ sql_context
+ .sql("INSERT INTO paimon.test_db.batch_size_t VALUES (1), (2),
(3), (4), (5)")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+
+ sql_context
+ .sql("SET 'paimon.read.batch-size' = '2'")
+ .await
+ .unwrap();
+ let batches = sql_context
+ .sql("SELECT id FROM paimon.test_db.batch_size_t")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_eq!(
+ batches
+ .iter()
+ .map(|batch| batch.num_rows())
+ .collect::<Vec<_>>(),
+ vec![2, 2, 1]
+ );
+
+ sql_context
+ .sql("RESET 'paimon.read.batch-size'")
+ .await
+ .unwrap();
+ let batches = sql_context
+ .sql("SELECT id FROM paimon.test_db.batch_size_t")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_eq!(
+ batches
+ .iter()
+ .map(|batch| batch.num_rows())
+ .collect::<Vec<_>>(),
+ vec![3, 2]
+ );
+ }
+
#[tokio::test]
async fn test_truncate_table() {
let (_tmp, sql_context) = setup_fs_sql_context().await;
diff --git a/crates/paimon/src/spec/core_options.rs
b/crates/paimon/src/spec/core_options.rs
index d1865191..e5a61b54 100644
--- a/crates/paimon/src/spec/core_options.rs
+++ b/crates/paimon/src/spec/core_options.rs
@@ -73,6 +73,7 @@ const MANIFEST_TARGET_FILE_SIZE_OPTION: &str =
"manifest.target-file-size";
const MANIFEST_TARGET_SIZE_OPTION: &str = "manifest.target-size";
const MANIFEST_MERGE_MIN_COUNT_OPTION: &str = "manifest.merge-min-count";
const WRITE_PARQUET_BUFFER_SIZE_OPTION: &str = "write.parquet-buffer-size";
+const READ_BATCH_SIZE_OPTION: &str = "read.batch-size";
pub(crate) const SEQUENCE_FIELD_OPTION: &str = "sequence.field";
pub(crate) const DISABLE_EXPLICIT_TYPE_CASTING_OPTION: &str =
"disable-explicit-type-casting";
pub(crate) const DISABLE_ALTER_COLUMN_NULL_TO_NOT_NULL_OPTION: &str =
@@ -109,6 +110,7 @@ const DEFAULT_PARTITION_DEFAULT_NAME: &str =
"__DEFAULT_PARTITION__";
const DEFAULT_CHANGELOG_FILE_PREFIX: &str = "changelog-";
const DEFAULT_TARGET_FILE_SIZE: i64 = 256 * 1024 * 1024;
const DEFAULT_WRITE_PARQUET_BUFFER_SIZE: i64 = 256 * 1024 * 1024;
+const DEFAULT_READ_BATCH_SIZE: usize = 1024;
const DYNAMIC_BUCKET_TARGET_ROW_NUM_OPTION: &str =
"dynamic-bucket.target-row-num";
const DEFAULT_DYNAMIC_BUCKET_TARGET_ROW_NUM: i64 = 200_000;
const DEFAULT_GLOBAL_INDEX_ROW_COUNT_PER_SHARD: i64 = 100_000;
@@ -305,6 +307,32 @@ impl<'a> CoreOptions<'a> {
Self { options }
}
+ /// Preferred number of rows emitted by file-format readers.
+ ///
+ /// Mirrors Java Paimon's `CoreOptions.READ_BATCH_SIZE`.
+ pub fn read_batch_size(&self) -> crate::Result<usize> {
+ let Some(raw) = self.options.get(READ_BATCH_SIZE_OPTION) else {
+ return Ok(DEFAULT_READ_BATCH_SIZE);
+ };
+ let value = raw
+ .parse::<i32>()
+ .map_err(|error| crate::Error::DataInvalid {
+ message: format!(
+ "Option '{READ_BATCH_SIZE_OPTION}' must be a positive
integer, got: {raw}"
+ ),
+ source: Some(Box::new(error)),
+ })?;
+ if value <= 0 {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Option '{READ_BATCH_SIZE_OPTION}' must be greater than 0,
got: {value}"
+ ),
+ source: None,
+ });
+ }
+ Ok(value as usize)
+ }
+
/// Reject scan options whose semantics the Rust core does not yet
implement.
///
/// These are not malformed input — they are unimplemented scan modes — so
@@ -1182,6 +1210,20 @@ fn parse_memory_size(value: &str) -> Option<i64> {
mod tests {
use super::*;
+ #[test]
+ fn test_read_batch_size() {
+ let options = HashMap::new();
+ assert_eq!(CoreOptions::new(&options).read_batch_size().unwrap(),
1024);
+
+ let options = HashMap::from([("read.batch-size".to_string(),
"8192".to_string())]);
+ assert_eq!(CoreOptions::new(&options).read_batch_size().unwrap(),
8192);
+
+ for value in ["0", "-1", "invalid"] {
+ let options = HashMap::from([("read.batch-size".to_string(),
value.to_string())]);
+ assert!(CoreOptions::new(&options).read_batch_size().is_err());
+ }
+ }
+
#[test]
fn test_source_split_defaults() {
let options = HashMap::new();
diff --git a/crates/paimon/src/spec/schema.rs b/crates/paimon/src/spec/schema.rs
index f92e5ef8..413c0d99 100644
--- a/crates/paimon/src/spec/schema.rs
+++ b/crates/paimon/src/spec/schema.rs
@@ -135,11 +135,16 @@ impl TableSchema {
&self.options
}
+ /// Typed view over this schema's table options.
+ pub fn core_options(&self) -> CoreOptions<'_> {
+ CoreOptions::new(&self.options)
+ }
+
/// Create a copy of this schema with extra options merged in.
///
/// A stored `query-auth.enabled = true` can't be turned off by a dynamic
override.
pub fn copy_with_options(&self, mut extra: HashMap<String, String>) ->
Self {
- if CoreOptions::new(&self.options).query_auth_enabled() {
+ if self.core_options().query_auth_enabled() {
extra.insert(QUERY_AUTH_ENABLED_OPTION.to_string(),
"true".to_string());
}
let mut new_schema = self.clone();
@@ -475,6 +480,7 @@ impl TableSchema {
&new_schema.primary_keys,
&new_schema.fields,
)?;
+ Schema::validate_read_batch_size(&new_schema.options)?;
Ok(new_schema)
}
@@ -921,6 +927,7 @@ impl Schema {
AggregationConfig::new(&options).validate_create_mode(&primary_keys,
&fields)?;
Self::validate_first_row_changelog_producer(&options)?;
Self::validate_rowkind_field(&options, &primary_keys, &fields)?;
+ Self::validate_read_batch_size(&options)?;
Ok(Self {
fields,
@@ -1344,6 +1351,13 @@ impl Schema {
}
}
+ fn validate_read_batch_size(options: &HashMap<String, String>) ->
crate::Result<()> {
+ CoreOptions::new(options)
+ .read_batch_size()
+ .map(|_| ())
+ .map_err(Self::options_error_to_config_invalid)
+ }
+
/// Returns top-level Blob field names for create-time Blob contract
checks.
fn top_level_blob_field_names(fields: &[DataField]) -> Vec<&str> {
fields
@@ -1670,6 +1684,43 @@ mod tests {
);
}
+ #[test]
+ fn test_schema_rejects_invalid_read_batch_size() {
+ for value in ["0", "-1", "invalid"] {
+ let err = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .option("read.batch-size", value)
+ .build()
+ .unwrap_err();
+ assert!(
+ matches!(err, crate::Error::ConfigInvalid { ref message }
+ if message.contains("read.batch-size")),
+ "got {err:?} for read.batch-size={value}"
+ );
+ }
+ }
+
+ #[test]
+ fn test_apply_changes_rejects_invalid_read_batch_size() {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .build()
+ .unwrap();
+ let table_schema = TableSchema::new(0, &schema);
+
+ let err = table_schema
+ .apply_changes(vec![crate::spec::SchemaChange::set_option(
+ "read.batch-size".to_string(),
+ "0".to_string(),
+ )])
+ .unwrap_err();
+ assert!(
+ matches!(err, crate::Error::ConfigInvalid { ref message }
+ if message.contains("read.batch-size")),
+ "got {err:?}"
+ );
+ }
+
#[test]
fn test_copy_with_replaced_options() {
let schema = Schema::builder()
diff --git a/crates/paimon/src/table/data_evolution_reader.rs
b/crates/paimon/src/table/data_evolution_reader.rs
index 1d5a1306..e166fea8 100644
--- a/crates/paimon/src/table/data_evolution_reader.rs
+++ b/crates/paimon/src/table/data_evolution_reader.rs
@@ -112,6 +112,7 @@ pub(crate) struct DataEvolutionReader {
blob_view_resolve_enabled: bool,
blob_view_rest_env: Option<RESTEnv>,
blob_read_limiter: BlobReadLimiter,
+ batch_size: Option<usize>,
}
impl DataEvolutionReader {
@@ -172,9 +173,15 @@ impl DataEvolutionReader {
blob_view_resolve_enabled,
blob_view_rest_env,
blob_read_limiter: BlobReadLimiter::new(),
+ batch_size: None,
})
}
+ pub(crate) fn with_batch_size(mut self, batch_size: Option<usize>) -> Self
{
+ self.batch_size = batch_size;
+ self
+ }
+
/// Read data files in data evolution mode.
pub fn read(self, data_splits: &[DataSplit]) ->
crate::Result<ArrowRecordBatchStream> {
let splits: Vec<DataSplit> = data_splits.to_vec();
@@ -198,7 +205,8 @@ impl DataEvolutionReader {
self.table_fields.clone(),
self.wide_file_read_type.clone(),
Vec::new(),
- );
+ )
+ .with_batch_size(self.batch_size);
for split in splits {
let row_ranges = split.row_ranges().map(|r| r.to_vec());
@@ -495,7 +503,8 @@ impl DataEvolutionReader {
HashSet::new(),
false,
None,
- )?;
+ )?
+ .with_batch_size(self.batch_size);
let mut stream = prescan.read(splits)?;
let mut view_structs = HashSet::new();
while let Some(batch) = stream.next().await {
@@ -557,6 +566,7 @@ impl DataEvolutionReader {
let table_fields = self.table_fields.clone();
let blob_descriptor_fields = self.blob_descriptor_fields.clone();
let blob_as_descriptor = self.blob_as_descriptor;
+ let batch_size = self.batch_size;
let anchor_deletion_vector = anchor_deletion_vector.clone();
// Batch size for column-merge output. Matches the default Parquet
reader batch size.
const MERGE_BATCH_SIZE: usize = 1024;
@@ -623,6 +633,7 @@ impl DataEvolutionReader {
schema_manager.clone(),
table_schema_id,
table_fields.clone(),
+ batch_size,
blob_as_descriptor,
anchor_deletion_vector.as_ref(),
)
@@ -1109,6 +1120,7 @@ fn open_source_stream(
schema_manager: SchemaManager,
table_schema_id: i64,
table_fields: Vec<DataField>,
+ batch_size: Option<usize>,
blob_as_descriptor: bool,
anchor_deletion_vector: Option<&DeletionVectorContext>,
) -> crate::Result<ArrowRecordBatchStream> {
@@ -1135,6 +1147,7 @@ fn open_source_stream(
source.read_fields().to_vec(),
Vec::new(),
)
+ .with_batch_size(batch_size)
.with_blob_as_descriptor(blob_as_descriptor);
match source {
@@ -5075,6 +5088,114 @@ mod tests {
}
}
+ #[tokio::test]
+ async fn
test_evolution_input_decode_honors_read_batch_size_on_all_file_paths() {
+ let tempdir = tempdir().unwrap();
+ let table_path = local_file_path(tempdir.path());
+ let bucket_dir = tempdir.path().join("bucket-0");
+ fs::create_dir_all(&bucket_dir).unwrap();
+
+ let raw_path = bucket_dir.join("raw.parquet");
+ write_int_parquet_file(
+ &raw_path,
+ vec![
+ ("id", vec![1, 2, 3, 4, 5]),
+ ("value", vec![10, 20, 30, 40, 50]),
+ ],
+ None,
+ );
+ let id_path = bucket_dir.join("id.parquet");
+ write_int_parquet_file(&id_path, vec![("id", vec![6, 7, 8, 9, 10])],
None);
+ let value_path = bucket_dir.join("value.parquet");
+ write_int_parquet_file(
+ &value_path,
+ vec![("value", vec![60, 70, 80, 90, 100])],
+ None,
+ );
+
+ let file_io = FileIOBuilder::new("file").build().unwrap();
+ let table_schema = TableSchema::new(
+ 0,
+ &Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("value", DataType::Int(IntType::new()))
+ .option("data-evolution.enabled", "true")
+ .option("read.batch-size", "2")
+ .build()
+ .unwrap(),
+ );
+ let table = Table::new(
+ file_io,
+ Identifier::new("default", "evolution_batch_size_t"),
+ table_path,
+ table_schema,
+ None,
+ );
+
+ let raw_split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(local_file_path(&bucket_dir))
+ .with_total_buckets(1)
+ .with_data_files(vec![data_file_meta_with_path(
+ "raw.parquet",
+ 0,
+ 5,
+ 1,
+ raw_path.metadata().unwrap().len() as i64,
+ Some(vec!["id", "value"]),
+ )])
+ .build()
+ .unwrap();
+ let merge_split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(local_file_path(&bucket_dir))
+ .with_total_buckets(1)
+ .with_data_files(vec![
+ data_file_meta_with_path(
+ "id.parquet",
+ 5,
+ 5,
+ 1,
+ id_path.metadata().unwrap().len() as i64,
+ Some(vec!["id"]),
+ ),
+ data_file_meta_with_path(
+ "value.parquet",
+ 5,
+ 5,
+ 2,
+ value_path.metadata().unwrap().len() as i64,
+ Some(vec!["value"]),
+ ),
+ ])
+ .build()
+ .unwrap();
+
+ let read = TableRead::new(&table, table.schema().fields().to_vec(),
Vec::new());
+ let batches = read
+ .to_arrow(&[raw_split, merge_split])
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert_eq!(
+ batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>(),
+ vec![2, 2, 1, 2, 2, 1]
+ );
+ assert_eq!(
+ collect_int_values(&batches, "id"),
+ vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+ );
+ }
+
/// _ROW_ID + predicate, raw branch: surviving rows keep their ORIGINAL row
/// ids (ids are attached before the residual filter), not renumbered ones.
#[tokio::test]
diff --git a/crates/paimon/src/table/data_file_reader.rs
b/crates/paimon/src/table/data_file_reader.rs
index 771f69bc..f5b2647e 100644
--- a/crates/paimon/src/table/data_file_reader.rs
+++ b/crates/paimon/src/table/data_file_reader.rs
@@ -44,6 +44,7 @@ pub(crate) struct DataFileReader {
read_type: Vec<DataField>,
predicates: Vec<Predicate>,
blob_as_descriptor: bool,
+ batch_size: Option<usize>,
}
impl DataFileReader {
@@ -63,6 +64,7 @@ impl DataFileReader {
read_type,
predicates,
blob_as_descriptor: false,
+ batch_size: None,
}
}
@@ -71,6 +73,11 @@ impl DataFileReader {
self
}
+ pub(crate) fn with_batch_size(mut self, batch_size: Option<usize>) -> Self
{
+ self.batch_size = batch_size;
+ self
+ }
+
/// Return a copy with a replaced read-type. Used by
`pk_vector_position_read`
/// to inject the internal `_ROW_ID` column for physical-position recovery.
pub(super) fn with_read_type(mut self, read_type: Vec<DataField>) -> Self {
@@ -247,6 +254,7 @@ impl DataFileReader {
let file_io = self.file_io.clone();
let split = split.clone();
let blob_as_descriptor = self.blob_as_descriptor;
+ let batch_size = self.batch_size;
let target_schema = build_target_arrow_schema(&read_type)?;
let file_fields = data_fields.clone().unwrap_or_else(||
table_fields.clone());
@@ -319,7 +327,7 @@ impl DataFileReader {
file_meta.file_size as u64,
&format_read_fields,
file_predicates.as_ref(),
- None,
+ batch_size,
row_selection,
).await?;
diff --git a/crates/paimon/src/table/format_read_builder.rs
b/crates/paimon/src/table/format_read_builder.rs
index 4fa03b62..f97805c0 100644
--- a/crates/paimon/src/table/format_read_builder.rs
+++ b/crates/paimon/src/table/format_read_builder.rs
@@ -21,7 +21,7 @@ use super::partition_filter::PartitionFilter;
use super::read_builder::split_scan_predicates;
use super::read_builder::{resolve_projected_fields,
validate_projection_possible};
use super::{Table, TableRead, TableScan};
-use crate::spec::{CoreOptions, DataField, Predicate};
+use crate::spec::{DataField, Predicate};
use crate::table::source::RowRange;
use crate::Result;
@@ -114,7 +114,8 @@ impl<'a> FormatReadBuilder<'a> {
}
pub(crate) fn new_read(&self) -> Result<TableRead<'a>> {
-
CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?;
+ let core_options = self.table.schema().core_options();
+ core_options.ensure_read_authorized()?;
let read_type = match self.resolve_read_type()? {
None => self.table.schema().fields().to_vec(),
Some(fields) => fields,
diff --git a/crates/paimon/src/table/format_table_read.rs
b/crates/paimon/src/table/format_table_read.rs
index 398ec912..0d561ab1 100644
--- a/crates/paimon/src/table/format_table_read.rs
+++ b/crates/paimon/src/table/format_table_read.rs
@@ -21,7 +21,7 @@ use super::data_file_reader::DataFileReader;
use super::read_builder::split_scan_predicates;
use super::{ArrowRecordBatchStream, Table};
use crate::arrow::{build_target_arrow_schema, paimon_type_to_arrow};
-use crate::spec::{extract_datum, BinaryRow, CoreOptions, DataField, DataType,
Datum, Predicate};
+use crate::spec::{extract_datum, BinaryRow, DataField, DataType, Datum,
Predicate};
use crate::{DataSplit, Error};
use arrow_array::{
new_null_array, ArrayRef, BinaryArray, BooleanArray, Date32Array,
Float32Array, Float64Array,
@@ -77,7 +77,8 @@ impl<'a> FormatTableRead<'a> {
&self,
data_splits: &[DataSplit],
) -> crate::Result<ArrowRecordBatchStream> {
-
CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?;
+ let core_options = self.table.schema().core_options();
+ core_options.ensure_read_authorized()?;
let read_type = self.read_type.clone();
let output_schema = build_target_arrow_schema(&read_type)?;
let partition_keys = self.table.schema().partition_keys().to_vec();
@@ -93,6 +94,7 @@ impl<'a> FormatTableRead<'a> {
let schema_manager = self.table.schema_manager().clone();
let schema_id = self.table.schema().id();
let mut remaining = self.limit;
+ let batch_size = Some(core_options.read_batch_size()?);
Ok(try_stream! {
for split in splits {
@@ -108,6 +110,7 @@ impl<'a> FormatTableRead<'a> {
data_read_type.clone(),
data_predicates.clone(),
)
+ .with_batch_size(batch_size)
.read(std::slice::from_ref(&split))?;
while let Some(batch) = stream.next().await {
diff --git a/crates/paimon/src/table/kv_file_reader.rs
b/crates/paimon/src/table/kv_file_reader.rs
index 2a6224b8..4450b273 100644
--- a/crates/paimon/src/table/kv_file_reader.rs
+++ b/crates/paimon/src/table/kv_file_reader.rs
@@ -55,6 +55,8 @@ pub(crate) struct KeyValueFileReader {
/// of a key survives); they are enforced by the post-merge residual
/// filter using the full `config.predicates` instead.
pushdown_predicates: Vec<Predicate>,
+ #[cfg(test)]
+ input_batch_sizes: Option<std::sync::Arc<std::sync::Mutex<Vec<usize>>>>,
}
/// Configuration for [`KeyValueFileReader`], grouping table schema and
@@ -70,6 +72,7 @@ pub(crate) struct KeyValueReadConfig {
pub primary_keys: Vec<String>,
pub merge_engine: MergeEngine,
pub sequence_fields: Vec<String>,
+ pub read_batch_size: usize,
}
/// Keep only the conjuncts of `predicates` that reference primary-key columns,
@@ -115,9 +118,20 @@ impl KeyValueFileReader {
file_io,
config,
pushdown_predicates,
+ #[cfg(test)]
+ input_batch_sizes: None,
}
}
+ #[cfg(test)]
+ fn with_input_batch_sizes(
+ mut self,
+ input_batch_sizes: std::sync::Arc<std::sync::Mutex<Vec<usize>>>,
+ ) -> Self {
+ self.input_batch_sizes = Some(input_batch_sizes);
+ self
+ }
+
fn new_merge_function(
merge_engine: MergeEngine,
table_options: &HashMap<String, String>,
@@ -313,6 +327,9 @@ impl KeyValueFileReader {
let residual_predicates = self.config.predicates;
let primary_keys = self.config.primary_keys;
let sequence_fields = self.config.sequence_fields;
+ let read_batch_size = self.config.read_batch_size;
+ #[cfg(test)]
+ let input_batch_sizes = self.input_batch_sizes;
// Build the merge output schema (keys + values, no system columns).
let mut merge_output_fields: Vec<DataField> = Vec::new();
@@ -350,7 +367,8 @@ impl KeyValueFileReader {
table_fields.clone(),
internal_read_type.clone(),
pushdown_predicates.clone(),
- );
+ )
+ .with_batch_size(Some(read_batch_size));
let stream = reader.read_single_file_stream(
split,
@@ -359,6 +377,18 @@ impl KeyValueFileReader {
None,
None,
)?;
+ #[cfg(test)]
+ let stream = if let Some(batch_sizes) =
input_batch_sizes.clone() {
+ stream
+ .inspect(move |batch| {
+ if let Ok(batch) = batch {
+
batch_sizes.lock().unwrap().push(batch.num_rows());
+ }
+ })
+ .boxed()
+ } else {
+ stream
+ };
file_streams.push(stream);
}
@@ -454,6 +484,7 @@ mod tests {
use crate::table::{Table, TableWrite};
use arrow_array::{Array, Int32Array};
use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema
as ArrowSchema};
+ use futures::TryStreamExt;
use std::sync::Arc;
fn test_file_io() -> FileIO {
@@ -646,6 +677,70 @@ mod tests {
assert!(matches!(&kept[0], Predicate::AlwaysTrue));
}
+ #[tokio::test]
+ async fn
kv_input_decode_honors_read_batch_size_without_changing_merge_batching() {
+ let file_io = test_file_io();
+ let table_path = "memory:/kv_read_batch_size";
+ setup_dirs(&file_io, table_path).await;
+ let table = pk_table(&file_io, table_path, &[("read.batch-size",
"2")]);
+
+ write_commit(
+ &table,
+ &int_batch(
+ vec![1, 2, 3, 4, 5],
+ vec![Some(10), Some(20), Some(30), Some(40), Some(50)],
+ ),
+ )
+ .await;
+ write_commit(
+ &table,
+ &int_batch(
+ vec![1, 2, 3, 4, 5],
+ vec![Some(11), Some(21), Some(31), Some(41), Some(51)],
+ ),
+ )
+ .await;
+
+ let read_builder = table.new_read_builder();
+ let plan = read_builder.new_scan().plan().await.unwrap();
+ let core_options = table.schema().core_options();
+ let input_batch_sizes = Arc::new(std::sync::Mutex::new(Vec::new()));
+ let reader = KeyValueFileReader::new(
+ table.file_io().clone(),
+ KeyValueReadConfig {
+ table_name: table.identifier().full_name(),
+ table_options: table.schema().options().clone(),
+ schema_manager: table.schema_manager().clone(),
+ table_schema_id: table.schema().id(),
+ table_fields: table.schema().fields().to_vec(),
+ read_type: table.schema().fields().to_vec(),
+ predicates: Vec::new(),
+ primary_keys: table.schema().trimmed_primary_keys(),
+ merge_engine: core_options.merge_engine().unwrap(),
+ sequence_fields: core_options
+ .sequence_fields()
+ .iter()
+ .map(|field| field.to_string())
+ .collect(),
+ read_batch_size: core_options.read_batch_size().unwrap(),
+ },
+ )
+ .with_input_batch_sizes(input_batch_sizes.clone());
+ let batches = reader
+ .read(plan.splits())
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ let mut decoded_batch_sizes =
input_batch_sizes.lock().unwrap().clone();
+ decoded_batch_sizes.sort_unstable();
+ assert_eq!(decoded_batch_sizes, vec![1, 1, 2, 2, 2, 2]);
+ assert_eq!(batches.len(), 1, "merge output batching stays
independent");
+ assert_eq!(batches[0].num_rows(), 5);
+ assert_eq!(int_column(&batches, "value"), vec![11, 21, 31, 41, 51]);
+ }
+
/// Non-PK equality filter on a dedup PK table read through the sort-merge
/// path must return only matching rows. Before the post-merge residual,
/// the non-PK conjunct was silently dropped and all rows came back.
diff --git a/crates/paimon/src/table/read_builder.rs
b/crates/paimon/src/table/read_builder.rs
index 6046ab0b..a3d390fd 100644
--- a/crates/paimon/src/table/read_builder.rs
+++ b/crates/paimon/src/table/read_builder.rs
@@ -450,7 +450,8 @@ impl<'a> PaimonReadBuilder<'a> {
pub fn new_read(&self) -> Result<TableRead<'a>> {
// Fail closed at read construction so bindings that short-circuit
before
// `to_arrow` (e.g. an empty-splits fast path) can't bypass the guard.
-
CoreOptions::new(self.table.schema.options()).ensure_read_authorized()?;
+ let core_options = self.table.schema.core_options();
+ core_options.ensure_read_authorized()?;
let read_type = match self.resolve_read_type()? {
None => self.table.schema.fields().to_vec(),
Some(fields) => fields,
diff --git a/crates/paimon/src/table/table_read.rs
b/crates/paimon/src/table/table_read.rs
index f3d1f64e..d95c0104 100644
--- a/crates/paimon/src/table/table_read.rs
+++ b/crates/paimon/src/table/table_read.rs
@@ -229,7 +229,7 @@ impl<'a> PaimonTableRead<'a> {
}
// Delta / Changelog rows are read as-is from planned files (no
full-table
// merge against historical base versions).
- self.new_data_file_reader().read(&data_splits)
+ self.new_data_file_reader()?.read(&data_splits)
}
/// Returns an audit-log stream for a planned incremental scan.
@@ -285,7 +285,8 @@ impl<'a> PaimonTableRead<'a> {
self.table.schema.fields().to_vec(),
read_type,
self.data_predicates.clone(),
- );
+ )
+
.with_batch_size(Some(self.table.schema().core_options().read_batch_size()?));
let raw_stream = reader.read(&data_splits)?;
Ok(Box::pin(async_stream::try_stream! {
@@ -340,7 +341,7 @@ impl<'a> PaimonTableRead<'a> {
/// Returns an [`ArrowRecordBatchStream`].
pub fn to_arrow(&self, data_splits: &[DataSplit]) ->
crate::Result<ArrowRecordBatchStream> {
let has_primary_keys = !self.table.schema.primary_keys().is_empty();
- let core_options = CoreOptions::new(self.table.schema.options());
+ let core_options = self.table.schema.core_options();
// Fail closed for a direct `TableRead` (bypassing
`ReadBuilder::new_read`).
core_options.ensure_read_authorized()?;
let merge_engine = core_options.merge_engine()?;
@@ -440,6 +441,7 @@ impl<'a> PaimonTableRead<'a> {
.iter()
.map(|s| s.to_string())
.collect(),
+ read_batch_size: core_options.read_batch_size()?,
},
);
reader.read(splits)
@@ -463,17 +465,18 @@ impl<'a> PaimonTableRead<'a> {
core_options.blob_view_fields(),
core_options.blob_view_resolve_enabled(),
self.table.rest_env().cloned(),
- )?;
+ )?
+ .with_batch_size(Some(core_options.read_batch_size()?));
reader.read(data_splits)
}
/// Read raw data files without dedup or evolution.
fn read_raw(&self, data_splits: &[DataSplit]) ->
crate::Result<ArrowRecordBatchStream> {
- self.new_data_file_reader().read(data_splits)
+ self.new_data_file_reader()?.read(data_splits)
}
- fn new_data_file_reader(&self) -> DataFileReader {
- DataFileReader::new(
+ fn new_data_file_reader(&self) -> crate::Result<DataFileReader> {
+ Ok(DataFileReader::new(
self.table.file_io.clone(),
self.table.schema_manager().clone(),
self.table.schema().id(),
@@ -481,6 +484,7 @@ impl<'a> PaimonTableRead<'a> {
self.read_type().to_vec(),
self.data_predicates.clone(),
)
+
.with_batch_size(Some(self.table.schema().core_options().read_batch_size()?)))
}
}
