From 1edac3835c18797f8e17a078a82fb76fa68b314d Mon Sep 17 00:00:00 2001 From: JeonDaehong Date: Sun, 18 Oct 2026 22:24:31 +0900 Subject: [PATCH 1/6] [core] Read single-run primary-key sections without sort merge A section with one sorted run of files that hold no retract rows, under the deduplicate or first-row engine, has no key to merge: the writer keeps keys unique within a file and the run's files don't overlap. Planning already reads such a split raw when it has one file. KeyValueFileReader now streams these sections directly instead of running the per-row merge, so small non-overlapping files packed into one split no longer pay for it. --- crates/paimon/src/table/kv_file_reader.rs | 243 ++++++++++++++++-- .../paimon/tests/pk_single_run_read_test.rs | 129 ++++++++++ 2 files changed, 349 insertions(+), 23 deletions(-) create mode 100644 crates/paimon/tests/pk_single_run_read_test.rs diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index e5486e9e1..2df4cbff3 100644 --- a/crates/paimon/src/table/kv_file_reader.rs +++ b/crates/paimon/src/table/kv_file_reader.rs @@ -39,14 +39,14 @@ use crate::file_index::file_index_result::FileIndexResult; use crate::io::FileIO; use crate::spec::{ is_variant_extraction_row, BigIntType, CoreOptions, DataField, DataFileMeta, - DataType as PaimonDataType, MergeEngine, PartialUpdateConfig, Predicate, TinyIntType, + DataType as PaimonDataType, MergeEngine, PartialUpdateConfig, Predicate, RowKind, TinyIntType, KEY_FIELD_ID_START, SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, VALUE_KIND_FIELD_ID, VALUE_KIND_FIELD_NAME, }; use crate::table::schema_manager::SchemaManager; use crate::table::ArrowRecordBatchStream; use crate::{DataSplit, DeletionFile, Error}; -use arrow_array::{RecordBatch, RecordBatchOptions}; +use arrow_array::{Array, BooleanArray, Int8Array, RecordBatch, RecordBatchOptions}; use async_stream::try_stream; use futures::StreamExt; @@ -184,6 +184,43 @@ fn ensure_merge_input_limit(input_stream_count: usize, limit: Option) -> Ok(()) } +/// Handles retract rows of an unmerged single run the way the merge path +/// handles a lone version of a key: first-row without `ignore-delete` rejects +/// them, otherwise they are omitted. +fn drop_lone_retracts( + batch: RecordBatch, + value_kind_index: usize, + reject: bool, +) -> crate::Result { + // The merge cursor reads a missing or NULL `_VALUE_KIND` as INSERT. + let Some(kinds) = batch + .column(value_kind_index) + .as_any() + .downcast_ref::() + else { + return Ok(batch); + }; + if kinds.null_count() == 0 && kinds.values().iter().all(|&kind| kind == 0 || kind == 2) { + return Ok(batch); + } + let mut keep = Vec::with_capacity(kinds.len()); + for kind in kinds.iter() { + let is_add = RowKind::from_value(kind.unwrap_or(0))?.is_add(); + if !is_add && reject { + return Err(Error::Unsupported { + message: "merge-engine=first-row does not support DELETE or UPDATE_BEFORE rows; set ignore-delete=true to ignore them".to_string(), + }); + } + keep.push(is_add); + } + arrow_select::filter::filter_record_batch(&batch, &BooleanArray::from(keep)).map_err(|e| { + Error::DataInvalid { + message: format!("Failed to drop retract rows: {e}"), + source: Some(Box::new(e)), + } + }) +} + /// Java's `KeyValueFieldsExtractor` builds the physical file layout from the /// schema that wrote the file, including its historical key names and types. pub(super) fn key_value_fields_extractor( @@ -397,15 +434,28 @@ impl KeyValueFileReader { } pub fn read(self, data_splits: &[DataSplit]) -> crate::Result { - self.read_with_merge_function(data_splits, Self::new_merge_function) + self.read_merged(data_splits, Self::new_merge_function, true) } + /// Always merges, so a custom merge function (e.g. the audit log, which + /// keeps retract rows) sees every row. pub(super) fn read_with_merge_function( self, data_splits: &[DataSplit], merge_function: impl Fn(&KeyValueReadConfig, &[DataField]) -> crate::Result> + Send + 'static, + ) -> crate::Result { + self.read_merged(data_splits, merge_function, false) + } + + fn read_merged( + self, + data_splits: &[DataSplit], + merge_function: impl Fn(&KeyValueReadConfig, &[DataField]) -> crate::Result> + + Send + + 'static, + skip_single_runs: bool, ) -> crate::Result { // A projected `_ROW_ID` is synthesized as all-nulls here, so the residual // would silently drop every row rather than hit its missing-column guard. @@ -711,6 +761,18 @@ impl KeyValueFileReader { .map(|budget| Arc::new(budget.without_prefetch())) }; let mut file_streams: Vec = Vec::new(); + // One sorted run never repeats a key, so under these engines + // every key is a lone version and needs no merge. Retract + // rows are still dropped per batch below: older Rust writers + // recorded `delete_row_count = 0` for files holding them. + let skip_merge = skip_single_runs + && input_stream_count == 1 + && matches!( + config.merge_engine, + MergeEngine::Deduplicate | MergeEngine::FirstRow + ); + let reject_retracts = config.merge_engine == MergeEngine::FirstRow + && !CoreOptions::new(&config.table_options).ignore_delete(); for MergeRun { files } in merge_group { let reader = DataFileReader::new( @@ -814,25 +876,43 @@ impl KeyValueFileReader { file_streams.push(run_stream); } - // Always go through sort-merge even for a single file: files - // written before the writer merged key groups at flush may - // still contain duplicate keys. - let mut merge_stream = SortMergeReaderBuilder::new( - file_streams, - internal_schema.clone(), - key_indices.clone(), - seq_index, - value_kind_index, - user_sequence_indices.clone(), - value_indices.clone(), - merge_output_schema.clone(), - merge_function(&config, &merge_output_fields)?, - ) - .with_user_sequence_descending( - !CoreOptions::new(&config.table_options) - .sequence_field_sort_order_is_ascending(), - ) - .build()?; + let mut merge_stream = if skip_merge { + let input = file_streams.pop().expect("one input stream"); + let columns: Vec = + key_indices.iter().chain(&value_indices).copied().collect(); + let schema = merge_output_schema.clone(); + input + .map(move |batch| { + let batch = + drop_lone_retracts(batch?, value_kind_index, reject_retracts)?; + let columns = + columns.iter().map(|&i| batch.column(i).clone()).collect(); + RecordBatch::try_new(schema.clone(), columns).map_err(|e| { + Error::UnexpectedError { + message: format!("Failed to project unmerged batch: {e}"), + source: Some(Box::new(e)), + } + }) + }) + .boxed() + } else { + SortMergeReaderBuilder::new( + file_streams, + internal_schema.clone(), + key_indices.clone(), + seq_index, + value_kind_index, + user_sequence_indices.clone(), + value_indices.clone(), + merge_output_schema.clone(), + merge_function(&config, &merge_output_fields)?, + ) + .with_user_sequence_descending( + !CoreOptions::new(&config.table_options) + .sequence_field_sort_order_is_ascending(), + ) + .build()? + }; while let Some(batch) = merge_stream.next().await { let batch = batch?; @@ -2011,7 +2091,9 @@ mod tests { .await .expect("the first output batch") .expect("a late deletion vector must be loaded lazily"); - assert_eq!(first.num_rows(), 128); + // These delete-free files form one run, which streams without the + // merge, so the first batch is one input batch. + assert!(first.num_rows() > 0); let err = stream.try_collect::>().await.unwrap_err(); assert!( @@ -2693,6 +2775,121 @@ mod tests { ); } + /// Writes one KV file with the given keys and row kinds, recorded with + /// `delete_row_count = Some(0)` as Rust writers did before #318. + async fn write_legacy_kv_file( + file_io: &FileIO, + table_path: &str, + file_name: &str, + rows: &[(i32, i8)], + ) -> DataFileMeta { + let schema = crate::arrow::build_target_arrow_schema(&[ + DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::new()), + ), + DataField::new( + VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME.to_string(), + DataType::TinyInt(TinyIntType::new()), + ), + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new(1, "value".to_string(), DataType::Int(IntType::new())), + ]) + .unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from_value(0, rows.len())), + Arc::new(Int8Array::from_iter_values(rows.iter().map(|row| row.1))), + Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.0))), + Arc::new(Int32Array::from_value(0, rows.len())), + ], + ) + .unwrap(); + let mut bytes = Vec::new(); + { + let mut writer = AsyncArrowWriter::try_new(&mut bytes, schema, None).unwrap(); + writer.write(&batch).await.unwrap(); + writer.close().await.unwrap(); + } + let bucket_path = format!("{table_path}/bucket-0"); + file_io.mkdirs(&format!("{bucket_path}/")).await.unwrap(); + file_io + .new_output(&format!("{bucket_path}/{file_name}")) + .unwrap() + .write(bytes::Bytes::from(bytes.clone())) + .await + .unwrap(); + + let mut file = dummy_data_file(file_name.to_string()); + file.file_size = bytes.len() as i64; + file.row_count = rows.len() as i64; + file.min_key = int_key(rows[0].0); + file.max_key = int_key(rows[rows.len() - 1].0); + file + } + + #[tokio::test] + async fn single_run_read_drops_retracts_hidden_by_legacy_metadata() { + let file_io = test_file_io(); + let adds = [(1, 0), (2, 0)]; + let retract = [(3, 3)]; + for (options, expected) in [ + (vec![], Some(vec![1, 2])), + (vec![("merge-engine", "first-row")], None), + ( + vec![("merge-engine", "first-row"), ("ignore-delete", "true")], + Some(vec![1, 2]), + ), + ] { + let table_path = "memory:/kv_legacy_retract"; + let table = pk_table(&file_io, table_path, &options); + let a = write_legacy_kv_file(&file_io, table_path, "a.parquet", &adds).await; + let b = write_legacy_kv_file(&file_io, table_path, "b.parquet", &retract).await; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![a, b]) + .build() + .unwrap(); + let core_options = table.schema().core_options(); + let result = 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(), + table_primary_keys: table.schema().primary_keys().to_vec(), + merge_engine: core_options.merge_engine().unwrap(), + sequence_fields: Vec::new(), + read_batch_size: core_options.read_batch_size().unwrap(), + merge_splits: false, + max_merge_input_streams: None, + parquet_read_budget: None, + mosaic_prefetch: MosaicPrefetchOptions::default(), + }, + ) + .read(&[split]) + .unwrap() + .try_collect::>() + .await; + match expected { + Some(ids) => assert_eq!(int_column(&result.unwrap(), "id"), ids, "{options:?}"), + None => assert!(result.is_err(), "{options:?}"), + } + } + } + #[tokio::test] async fn sorted_run_read_matches_per_file_fan_out() { let file_io = test_file_io(); diff --git a/crates/paimon/tests/pk_single_run_read_test.rs b/crates/paimon/tests/pk_single_run_read_test.rs new file mode 100644 index 000000000..2dca3d568 --- /dev/null +++ b/crates/paimon/tests/pk_single_run_read_test.rs @@ -0,0 +1,129 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Primary-key reads of a split whose sections each hold one sorted run. + +mod common; + +use std::sync::Arc; + +use arrow_array::{ArrayRef, Int32Array, RecordBatch}; +use common::incremental_helpers::{memory_table, persist_table_schema, setup_dirs, write_batch}; +use futures::TryStreamExt; +use paimon::spec::{DataType, Datum, IntType, PredicateBuilder, Schema, TableSchema}; +use paimon::table::Table; + +async fn pk_table() -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("v", DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "1") + .build() + .unwrap(); + let path = "memory:/pk_single_run_read"; + let (io, table) = memory_table(path, TableSchema::new(0, &schema)); + setup_dirs(&io, path).await; + persist_table_schema(&io, path, table.schema()).await; + table +} + +async fn write(table: &Table, ids: std::ops::Range, offset: i32) { + let ids: Vec = ids.collect(); + let values: Vec = ids.iter().map(|id| id + offset).collect(); + write_batch( + table, + &RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(ids)) as ArrayRef), + ("v", Arc::new(Int32Array::from(values)) as ArrayRef), + ]) + .unwrap(), + ) + .await; +} + +async fn read( + table: &Table, + projection: &[&str], + filter: Option, +) -> Vec> { + let mut builder = table.new_read_builder(); + builder.with_projection(projection).unwrap(); + if let Some(filter) = filter { + builder.with_filter(filter); + } + let plan = builder.new_scan().plan().await.unwrap(); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut rows: Vec> = batches + .iter() + .flat_map(|batch| { + (0..batch.num_rows()).map(move |row| { + (0..batch.num_columns()) + .map(|col| { + batch + .column(col) + .as_any() + .downcast_ref::() + .unwrap() + .value(row) + }) + .collect() + }) + }) + .collect(); + rows.sort_unstable(); + rows +} + +#[tokio::test] +async fn disjoint_files_read_without_merge_and_overlaps_still_merge() { + let table = pk_table().await; + // Three files with disjoint keys land in one split, one run per section. + write(&table, 0..100, 1000).await; + write(&table, 100..200, 1000).await; + write(&table, 200..300, 1000).await; + let plan = table.new_read_builder().new_scan().plan().await.unwrap(); + assert_eq!(plan.splits().len(), 1); + + let expected: Vec> = (0..300).map(|id| vec![id, id + 1000]).collect(); + assert_eq!(read(&table, &["id", "v"], None).await, expected); + // Value-only projection and a residual filter on the unmerged path. + let pb = PredicateBuilder::new(table.schema().fields()); + assert_eq!( + read( + &table, + &["v"], + Some(pb.greater_or_equal("v", Datum::Int(1250)).unwrap()) + ) + .await, + (1250..1300).map(|v| vec![v]).collect::>() + ); + + // Overlapping keys make a two-run section, which still merges. + write(&table, 150..160, 2000).await; + let expected: Vec> = (0..300) + .map(|id| vec![id, id + if (150..160).contains(&id) { 2000 } else { 1000 }]) + .collect(); + assert_eq!(read(&table, &["id", "v"], None).await, expected); +} From 6274e9c7599a601a02dfe14b8fafac17ec448b45 Mon Sep 17 00:00:00 2001 From: JeonDaehong Date: Thu, 22 Oct 2026 21:10:20 +0900 Subject: [PATCH 2/6] [core] Keep merging files written under another merge engine A file written under aggregation could hold several versions of a key (before #839), and ALTER can switch the table to deduplicate afterwards. Skip the sort merge only when every file was written under deduplicate or first-row. --- crates/paimon/src/table/kv_file_reader.rs | 147 +++++++++++++++++++++- 1 file changed, 146 insertions(+), 1 deletion(-) diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index 2df4cbff3..f742f4b99 100644 --- a/crates/paimon/src/table/kv_file_reader.rs +++ b/crates/paimon/src/table/kv_file_reader.rs @@ -765,12 +765,36 @@ impl KeyValueFileReader { // every key is a lone version and needs no merge. Retract // rows are still dropped per batch below: older Rust writers // recorded `delete_row_count = 0` for files holding them. - let skip_merge = skip_single_runs + let mut skip_merge = skip_single_runs && input_stream_count == 1 && matches!( config.merge_engine, MergeEngine::Deduplicate | MergeEngine::FirstRow ); + // A file written under another engine may hold several + // versions of a key (the aggregation writer kept them before + // #839), and ALTER can switch the engine afterwards. + if skip_merge { + 'files: for run in &merge_group { + for file in &run.files { + let schema_id = file.file.schema_id; + if schema_id == table_schema_id { + continue; + } + let schema = + config.schema_manager.schema(schema_id).await?; + let engine = + CoreOptions::new(schema.options()).merge_engine()?; + if !matches!( + engine, + MergeEngine::Deduplicate | MergeEngine::FirstRow + ) { + skip_merge = false; + break 'files; + } + } + } + } let reject_retracts = config.merge_engine == MergeEngine::FirstRow && !CoreOptions::new(&config.table_options).ignore_delete(); @@ -2890,6 +2914,127 @@ mod tests { } } + /// Writes one KV file with `(id, sequence, value)` rows, under schema 0. + async fn write_kv_versions( + file_io: &FileIO, + table_path: &str, + file_name: &str, + rows: &[(i32, i64, i32)], + ) -> DataFileMeta { + let schema = crate::arrow::build_target_arrow_schema(&[ + DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::new()), + ), + DataField::new( + VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME.to_string(), + DataType::TinyInt(TinyIntType::new()), + ), + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new(1, "value".to_string(), DataType::Int(IntType::new())), + ]) + .unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from_iter_values(rows.iter().map(|row| row.1))), + Arc::new(Int8Array::from_value(0, rows.len())), + Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.0))), + Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.2))), + ], + ) + .unwrap(); + let mut bytes = Vec::new(); + { + let mut writer = AsyncArrowWriter::try_new(&mut bytes, schema, None).unwrap(); + writer.write(&batch).await.unwrap(); + writer.close().await.unwrap(); + } + let bucket_path = format!("{table_path}/bucket-0"); + file_io.mkdirs(&format!("{bucket_path}/")).await.unwrap(); + file_io + .new_output(&format!("{bucket_path}/{file_name}")) + .unwrap() + .write(bytes::Bytes::from(bytes.clone())) + .await + .unwrap(); + + let mut file = dummy_data_file(file_name.to_string()); + file.file_size = bytes.len() as i64; + file.row_count = rows.len() as i64; + file.min_key = int_key(rows[0].0); + file.max_key = int_key(rows[rows.len() - 1].0); + file.min_sequence_number = rows.iter().map(|row| row.1).min().unwrap(); + file.max_sequence_number = rows.iter().map(|row| row.1).max().unwrap(); + file + } + + #[tokio::test] + async fn single_run_read_merges_files_written_under_aggregation() { + // Before #839 the aggregation writer kept several versions of a key in + // one file. After ALTER to deduplicate, such a file still needs merging. + let file_io = test_file_io(); + let table_path = "memory:/kv_engine_change"; + let old = pk_table( + &file_io, + table_path, + &[ + ("merge-engine", "aggregation"), + ("fields.value.aggregate-function", "sum"), + ], + ); + write_schema_file(&old, old.schema()).await; + let mut json = serde_json::to_value(old.schema()).unwrap(); + json["id"] = 1.into(); + json["options"]["merge-engine"] = "deduplicate".into(); + let schema: TableSchema = serde_json::from_value(json).unwrap(); + write_schema_file(&old, &schema).await; + let table = old.copy_with_resolved_schema(schema, "main").unwrap(); + let a = + write_kv_versions(&file_io, table_path, "a.parquet", &[(1, 0, 10), (1, 1, 11)]).await; + let b = write_kv_versions(&file_io, table_path, "b.parquet", &[(2, 2, 20)]).await; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![a, b]) + .build() + .unwrap(); + let core_options = table.schema().core_options(); + let batches = 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(), + table_primary_keys: table.schema().primary_keys().to_vec(), + merge_engine: core_options.merge_engine().unwrap(), + sequence_fields: Vec::new(), + read_batch_size: core_options.read_batch_size().unwrap(), + merge_splits: false, + max_merge_input_streams: None, + parquet_read_budget: None, + mosaic_prefetch: MosaicPrefetchOptions::default(), + }, + ) + .read(&[split]) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(int_column(&batches, "id"), vec![1, 2]); + assert_eq!(int_column(&batches, "value"), vec![11, 20]); + } + #[tokio::test] async fn sorted_run_read_matches_per_file_fan_out() { let file_io = test_file_io(); From 9de47132e1c5d5288208fa00968b6e989c2c72c2 Mon Sep 17 00:00:00 2001 From: JeonDaehong Date: Sat, 24 Oct 2026 11:17:18 +0900 Subject: [PATCH 3/6] [core] Check the stored engine even when the schema ID matches Table::copy_with_options keeps the schema ID while overriding merge-engine, so files written under aggregation share the read schema's ID. --- .../paimon/src/table/audit_log_table/read.rs | 1 + crates/paimon/src/table/kv_file_reader.rs | 119 +++++++++++++----- crates/paimon/src/table/mod.rs | 6 + crates/paimon/src/table/table_read.rs | 2 + 4 files changed, 98 insertions(+), 30 deletions(-) diff --git a/crates/paimon/src/table/audit_log_table/read.rs b/crates/paimon/src/table/audit_log_table/read.rs index 4719b39b3..2abba6ab1 100644 --- a/crates/paimon/src/table/audit_log_table/read.rs +++ b/crates/paimon/src/table/audit_log_table/read.rs @@ -125,6 +125,7 @@ impl<'a> AuditLogRead<'a> { table_options: self.read.table.schema().options().clone(), schema_manager: self.read.table.schema_manager().clone(), table_schema_id: self.read.table.schema().id(), + merge_engine_overridden: self.read.table.merge_engine_overridden(), table_fields: self.read.table.schema.fields().to_vec(), read_type, predicates: self.read.data_predicates.clone(), diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index f742f4b99..286ca092b 100644 --- a/crates/paimon/src/table/kv_file_reader.rs +++ b/crates/paimon/src/table/kv_file_reader.rs @@ -81,6 +81,9 @@ pub(crate) struct KeyValueReadConfig { /// Full table keys also protect partition-PK fields from aggregation. pub table_primary_keys: Vec, pub merge_engine: MergeEngine, + /// `merge-engine` was set by dynamic options, so files written under + /// `table_schema_id` may still hold another engine's output. + pub merge_engine_overridden: bool, pub sequence_fields: Vec, pub read_batch_size: usize, /// Merge files from all supplied splits into one globally key-sorted stream. @@ -777,12 +780,15 @@ impl KeyValueFileReader { if skip_merge { 'files: for run in &merge_group { for file in &run.files { - let schema_id = file.file.schema_id; - if schema_id == table_schema_id { + // Dynamic options can override the engine but keep + // the schema ID, so only then check the stored one. + if file.file.schema_id == table_schema_id + && !config.merge_engine_overridden + { continue; } let schema = - config.schema_manager.schema(schema_id).await?; + config.schema_manager.schema(file.file.schema_id).await?; let engine = CoreOptions::new(schema.options()).merge_engine()?; if !matches!( @@ -2207,6 +2213,7 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), + merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -2328,6 +2335,7 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), + merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -2533,6 +2541,7 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), + merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -2622,6 +2631,7 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), + merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -2889,6 +2899,7 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), + merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -2971,6 +2982,37 @@ mod tests { file } + async fn read_kv_split(table: &Table, split: DataSplit) -> Vec { + let core_options = table.schema().core_options(); + 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(), + merge_engine_overridden: table.merge_engine_overridden(), + table_fields: table.schema().fields().to_vec(), + read_type: table.schema().fields().to_vec(), + predicates: Vec::new(), + primary_keys: table.schema().trimmed_primary_keys(), + table_primary_keys: table.schema().primary_keys().to_vec(), + merge_engine: core_options.merge_engine().unwrap(), + sequence_fields: Vec::new(), + read_batch_size: core_options.read_batch_size().unwrap(), + merge_splits: false, + max_merge_input_streams: None, + parquet_read_budget: None, + mosaic_prefetch: MosaicPrefetchOptions::default(), + }, + ) + .read(&[split]) + .unwrap() + .try_collect::>() + .await + .unwrap() + } + #[tokio::test] async fn single_run_read_merges_files_written_under_aggregation() { // Before #839 the aggregation writer kept several versions of a key in @@ -3004,33 +3046,48 @@ mod tests { .with_data_files(vec![a, b]) .build() .unwrap(); - let core_options = table.schema().core_options(); - let batches = 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(), - table_primary_keys: table.schema().primary_keys().to_vec(), - merge_engine: core_options.merge_engine().unwrap(), - sequence_fields: Vec::new(), - read_batch_size: core_options.read_batch_size().unwrap(), - merge_splits: false, - max_merge_input_streams: None, - parquet_read_budget: None, - mosaic_prefetch: MosaicPrefetchOptions::default(), - }, - ) - .read(&[split]) - .unwrap() - .try_collect::>() - .await - .unwrap(); + let batches = read_kv_split(&table, split).await; + assert_eq!(int_column(&batches, "id"), vec![1, 2]); + assert_eq!(int_column(&batches, "value"), vec![11, 20]); + } + + #[tokio::test] + async fn single_run_read_merges_when_options_override_the_engine() { + // copy_with_options keeps the schema ID, so the files and the read + // schema share it while the stored engine is still aggregation. + let file_io = test_file_io(); + let table_path = "memory:/kv_engine_override"; + let old = pk_table( + &file_io, + table_path, + &[ + ("merge-engine", "aggregation"), + ("fields.value.aggregate-function", "sum"), + ], + ); + write_schema_file(&old, old.schema()).await; + let table = old.copy_with_options(HashMap::from([( + "merge-engine".to_string(), + "deduplicate".to_string(), + )])); + assert_eq!(table.schema().id(), old.schema().id()); + assert_eq!( + table.schema().core_options().merge_engine().unwrap(), + MergeEngine::Deduplicate + ); + let a = + write_kv_versions(&file_io, table_path, "a.parquet", &[(1, 0, 10), (1, 1, 11)]).await; + let b = write_kv_versions(&file_io, table_path, "b.parquet", &[(2, 2, 20)]).await; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![a, b]) + .build() + .unwrap(); + let batches = read_kv_split(&table, split).await; assert_eq!(int_column(&batches, "id"), vec![1, 2]); assert_eq!(int_column(&batches, "value"), vec![11, 20]); } @@ -3067,6 +3124,7 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), + merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -3136,6 +3194,7 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), + merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 86b8f54de..c763ed250 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -718,6 +718,12 @@ impl Table { PostponeFixedBucketWriteBuilder::new(self) } + /// Whether dynamic options set `merge-engine`. They keep the schema ID, so + /// the stored schema of that ID may name another engine. + pub(crate) fn merge_engine_overridden(&self) -> bool { + self.applied_dynamic_option_keys.contains("merge-engine") + } + /// Create a copy of this table with extra options merged into the schema. /// /// This never switches the schema version; it corresponds to Java diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index 47d38e5e1..cdc749a13 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -1320,6 +1320,7 @@ impl<'a> PaimonTableRead<'a> { table_options: self.table.schema().options().clone(), schema_manager: self.table.schema_manager().clone(), table_schema_id: self.table.schema().id(), + merge_engine_overridden: self.table.merge_engine_overridden(), table_fields: self.table.schema.fields().to_vec(), read_type: read_type.to_vec(), predicates: self.data_predicates.clone(), @@ -2007,6 +2008,7 @@ impl<'a> PaimonTableRead<'a> { table_options: self.table.schema().options().clone(), schema_manager: self.table.schema_manager().clone(), table_schema_id: self.table.schema().id(), + merge_engine_overridden: self.table.merge_engine_overridden(), table_fields: self.table.schema.fields().to_vec(), read_type: self.read_type().to_vec(), predicates: self.data_predicates.clone(), From 0dcaaa8dfe8d457d09ede7c38fcd4369330105c1 Mon Sep 17 00:00:00 2001 From: JeonDaehong Date: Sun, 25 Oct 2026 23:55:02 +0900 Subject: [PATCH 4/6] [core] Always check the writer engine in the stored schema Dynamic options and time travel can both give the read schema another merge engine under the files' schema ID, so the ID match is not enough to trust the read schema's engine. Drop the override flag and look up the stored schema for every file; the schema manager caches it. --- .../paimon/src/table/audit_log_table/read.rs | 1 - crates/paimon/src/table/kv_file_reader.rs | 142 ++++++++++++++---- crates/paimon/src/table/mod.rs | 6 - crates/paimon/src/table/table_read.rs | 2 - 4 files changed, 112 insertions(+), 39 deletions(-) diff --git a/crates/paimon/src/table/audit_log_table/read.rs b/crates/paimon/src/table/audit_log_table/read.rs index 2abba6ab1..4719b39b3 100644 --- a/crates/paimon/src/table/audit_log_table/read.rs +++ b/crates/paimon/src/table/audit_log_table/read.rs @@ -125,7 +125,6 @@ impl<'a> AuditLogRead<'a> { table_options: self.read.table.schema().options().clone(), schema_manager: self.read.table.schema_manager().clone(), table_schema_id: self.read.table.schema().id(), - merge_engine_overridden: self.read.table.merge_engine_overridden(), table_fields: self.read.table.schema.fields().to_vec(), read_type, predicates: self.read.data_predicates.clone(), diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index 286ca092b..c93222d15 100644 --- a/crates/paimon/src/table/kv_file_reader.rs +++ b/crates/paimon/src/table/kv_file_reader.rs @@ -81,9 +81,6 @@ pub(crate) struct KeyValueReadConfig { /// Full table keys also protect partition-PK fields from aggregation. pub table_primary_keys: Vec, pub merge_engine: MergeEngine, - /// `merge-engine` was set by dynamic options, so files written under - /// `table_schema_id` may still hold another engine's output. - pub merge_engine_overridden: bool, pub sequence_fields: Vec, pub read_batch_size: usize, /// Merge files from all supplied splits into one globally key-sorted stream. @@ -780,15 +777,26 @@ impl KeyValueFileReader { if skip_merge { 'files: for run in &merge_group { for file in &run.files { - // Dynamic options can override the engine but keep - // the schema ID, so only then check the stored one. - if file.file.schema_id == table_schema_id - && !config.merge_engine_overridden + // Always check the stored schema: dynamic options and + // time travel can change the read schema's engine + // without changing its ID. Merging is always correct, + // so a missing schema file under the read schema's ID + // only costs the skip. + let schema = match config + .schema_manager + .schema(file.file.schema_id) + .await { - continue; - } - let schema = - config.schema_manager.schema(file.file.schema_id).await?; + Ok(schema) => schema, + Err(crate::Error::IoUnexpected { source, .. }) + if source.kind() == opendal::ErrorKind::NotFound + && file.file.schema_id == table_schema_id => + { + skip_merge = false; + break 'files; + } + Err(err) => Err(err)?, + }; let engine = CoreOptions::new(schema.options()).merge_engine()?; if !matches!( @@ -2076,6 +2084,7 @@ mod tests { ("deletion-vectors.merge-on-read", "true"), ], ); + write_schema_file(&table, table.schema()).await; let mut files = Vec::new(); for i in 0..10 { @@ -2213,7 +2222,6 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), - merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -2335,7 +2343,6 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), - merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -2541,7 +2548,6 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), - merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -2631,7 +2637,6 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), - merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -2691,6 +2696,7 @@ mod tests { ("read.parquet.row-group.parallelism", "1"), ], ); + write_schema_file(&table, table.schema()).await; let low = write_multi_row_group_kv_file(&file_io, table_path, "low.parquet", 0, 0, 10).await; let high = @@ -2747,6 +2753,7 @@ mod tests { ("read.parquet.row-group.parallelism", "1"), ], ); + write_schema_file(&table, table.schema()).await; let first_low = write_multi_row_group_kv_file(&file_io, table_path, "first-low.parquet", 0, 0, 10) .await; @@ -2880,6 +2887,7 @@ mod tests { ] { let table_path = "memory:/kv_legacy_retract"; let table = pk_table(&file_io, table_path, &options); + write_schema_file(&table, table.schema()).await; let a = write_legacy_kv_file(&file_io, table_path, "a.parquet", &adds).await; let b = write_legacy_kv_file(&file_io, table_path, "b.parquet", &retract).await; let split = DataSplitBuilder::new() @@ -2899,7 +2907,6 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), - merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -2991,7 +2998,6 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), - merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -3052,11 +3058,11 @@ mod tests { } #[tokio::test] - async fn single_run_read_merges_when_options_override_the_engine() { - // copy_with_options keeps the schema ID, so the files and the read - // schema share it while the stored engine is still aggregation. + async fn single_run_read_merges_after_time_travel_past_an_engine_alter() { + // aggregation (schema 0) -> ALTER to deduplicate (schema 1). Reading + // snapshot 1 resolves schema 0 but keeps the current engine option. let file_io = test_file_io(); - let table_path = "memory:/kv_engine_override"; + let table_path = "memory:/kv_engine_time_travel"; let old = pk_table( &file_io, table_path, @@ -3066,15 +3072,45 @@ mod tests { ], ); write_schema_file(&old, old.schema()).await; - let table = old.copy_with_options(HashMap::from([( - "merge-engine".to_string(), - "deduplicate".to_string(), - )])); - assert_eq!(table.schema().id(), old.schema().id()); + let mut json = serde_json::to_value(old.schema()).unwrap(); + json["id"] = 1.into(); + json["options"]["merge-engine"] = "deduplicate".into(); + let schema: TableSchema = serde_json::from_value(json).unwrap(); + write_schema_file(&old, &schema).await; + let table = old.copy_with_resolved_schema(schema, "main").unwrap(); + file_io + .mkdirs(&format!("{table_path}/snapshot/")) + .await + .unwrap(); + let snapshot = crate::spec::Snapshot::builder() + .version(3) + .id(1) + .schema_id(0) + .base_manifest_list("base-list".to_string()) + .delta_manifest_list("delta-list".to_string()) + .commit_user("test-user".to_string()) + .commit_identifier(0) + .commit_kind(crate::spec::CommitKind::APPEND) + .time_millis(1000) + .build(); + assert!(table + .snapshot_manager() + .commit_snapshot(&snapshot) + .await + .unwrap()); + let historical = table + .copy_with_time_travel_strict(HashMap::from([( + "scan.version".to_string(), + "1".to_string(), + )])) + .await + .unwrap(); + assert_eq!(historical.schema().id(), 0); assert_eq!( - table.schema().core_options().merge_engine().unwrap(), + historical.schema().core_options().merge_engine().unwrap(), MergeEngine::Deduplicate ); + let a = write_kv_versions(&file_io, table_path, "a.parquet", &[(1, 0, 10), (1, 1, 11)]).await; let b = write_kv_versions(&file_io, table_path, "b.parquet", &[(2, 2, 20)]).await; @@ -3087,16 +3123,64 @@ mod tests { .with_data_files(vec![a, b]) .build() .unwrap(); - let batches = read_kv_split(&table, split).await; + let batches = read_kv_split(&historical, split).await; assert_eq!(int_column(&batches, "id"), vec![1, 2]); assert_eq!(int_column(&batches, "value"), vec![11, 20]); } + #[tokio::test] + async fn single_run_read_merges_when_the_read_schema_names_another_engine() { + // Dynamic options and time travel can give the read schema another + // engine under the files' schema ID; the stored engine is aggregation. + let file_io = test_file_io(); + let table_path = "memory:/kv_engine_override"; + let old = pk_table( + &file_io, + table_path, + &[ + ("merge-engine", "aggregation"), + ("fields.value.aggregate-function", "sum"), + ], + ); + write_schema_file(&old, old.schema()).await; + let overridden = old.copy_with_options(HashMap::from([( + "merge-engine".to_string(), + "deduplicate".to_string(), + )])); + // Same read schema as a historical read after an ALTER to deduplicate. + let rewrapped = old + .copy_with_resolved_schema(overridden.schema().clone(), "main") + .unwrap(); + let a = + write_kv_versions(&file_io, table_path, "a.parquet", &[(1, 0, 10), (1, 1, 11)]).await; + let b = write_kv_versions(&file_io, table_path, "b.parquet", &[(2, 2, 20)]).await; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![a, b]) + .build() + .unwrap(); + for table in [&overridden, &rewrapped] { + assert_eq!(table.schema().id(), old.schema().id()); + assert_eq!( + table.schema().core_options().merge_engine().unwrap(), + MergeEngine::Deduplicate + ); + let batches = read_kv_split(table, split.clone()).await; + assert_eq!(int_column(&batches, "id"), vec![1, 2]); + assert_eq!(int_column(&batches, "value"), vec![11, 20]); + } + } + #[tokio::test] async fn sorted_run_read_matches_per_file_fan_out() { let file_io = test_file_io(); let table_path = "memory:/kv_sorted_run"; let table = pk_table(&file_io, table_path, &[]); + write_schema_file(&table, table.schema()).await; let low = write_multi_row_group_kv_file(&file_io, table_path, "low.parquet", 0, 0, 10).await; let high = @@ -3124,7 +3208,6 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), - merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), @@ -3194,7 +3277,6 @@ mod tests { table_options: table.schema().options().clone(), schema_manager: table.schema_manager().clone(), table_schema_id: table.schema().id(), - merge_engine_overridden: table.merge_engine_overridden(), table_fields: table.schema().fields().to_vec(), read_type: table.schema().fields().to_vec(), predicates: Vec::new(), diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index c763ed250..86b8f54de 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -718,12 +718,6 @@ impl Table { PostponeFixedBucketWriteBuilder::new(self) } - /// Whether dynamic options set `merge-engine`. They keep the schema ID, so - /// the stored schema of that ID may name another engine. - pub(crate) fn merge_engine_overridden(&self) -> bool { - self.applied_dynamic_option_keys.contains("merge-engine") - } - /// Create a copy of this table with extra options merged into the schema. /// /// This never switches the schema version; it corresponds to Java diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index cdc749a13..47d38e5e1 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -1320,7 +1320,6 @@ impl<'a> PaimonTableRead<'a> { table_options: self.table.schema().options().clone(), schema_manager: self.table.schema_manager().clone(), table_schema_id: self.table.schema().id(), - merge_engine_overridden: self.table.merge_engine_overridden(), table_fields: self.table.schema.fields().to_vec(), read_type: read_type.to_vec(), predicates: self.data_predicates.clone(), @@ -2008,7 +2007,6 @@ impl<'a> PaimonTableRead<'a> { table_options: self.table.schema().options().clone(), schema_manager: self.table.schema_manager().clone(), table_schema_id: self.table.schema().id(), - merge_engine_overridden: self.table.merge_engine_overridden(), table_fields: self.table.schema.fields().to_vec(), read_type: self.read_type().to_vec(), predicates: self.data_predicates.clone(), From af7ae3ce69cb677dadb4c3dca799e4fa8954b4ee Mon Sep 17 00:00:00 2001 From: JeonDaehong Date: Tue, 27 Oct 2026 20:46:33 +0900 Subject: [PATCH 5/6] [core] Merge a single run from the first repeated key instead of trusting schemas Neither the stored nor the read schema proves which engine wrote a file: before #839 copy_with_options could switch the writer to aggregation under the same schema ID, and that writer kept every version of a key. Drop the stored-schema lookup. The unmerged path now holds one batch back and checks adjacent keys; at the first repeated key it hands the held batch and the rest of the run to the sort merge. --- crates/paimon/src/table/kv_file_reader.rs | 226 ++++++++++++++++------ crates/paimon/src/table/sort_merge.rs | 6 + 2 files changed, 175 insertions(+), 57 deletions(-) diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index c93222d15..68eaf3bbb 100644 --- a/crates/paimon/src/table/kv_file_reader.rs +++ b/crates/paimon/src/table/kv_file_reader.rs @@ -221,6 +221,48 @@ fn drop_lone_retracts( }) } +/// Whether `batch` repeats a primary key, either in two adjacent rows or in its +/// first row and the last row of `previous`. A key type the comparison kernel +/// does not support counts as repeated, so the caller merges. +fn repeats_a_key( + previous: Option<&RecordBatch>, + batch: &RecordBatch, + key_indices: &[usize], +) -> bool { + let rows = batch.num_rows(); + if rows == 0 { + return false; + } + let previous = previous.filter(|previous| previous.num_rows() > 0); + let mut at_boundary = previous.is_some(); + let mut adjacent: Option = None; + for &index in key_indices { + let column = batch.column(index); + if let Some(previous) = previous.filter(|_| at_boundary) { + let last = previous.column(index).slice(previous.num_rows() - 1, 1); + match arrow_ord::cmp::eq(&last, &column.slice(0, 1)) { + Ok(equal) => at_boundary = equal.true_count() == 1, + Err(_) => return true, + } + } + if rows > 1 { + let Ok(equal) = + arrow_ord::cmp::eq(&column.slice(0, rows - 1), &column.slice(1, rows - 1)) + else { + return true; + }; + adjacent = Some(match adjacent { + None => equal, + Some(so_far) => match arrow_arith::boolean::and(&so_far, &equal) { + Ok(both) => both, + Err(_) => return true, + }, + }); + } + } + at_boundary || adjacent.is_some_and(|equal| equal.true_count() > 0) +} + /// Java's `KeyValueFieldsExtractor` builds the physical file layout from the /// schema that wrote the file, including its historical key names and types. pub(super) fn key_value_fields_extractor( @@ -761,54 +803,17 @@ impl KeyValueFileReader { .map(|budget| Arc::new(budget.without_prefetch())) }; let mut file_streams: Vec = Vec::new(); - // One sorted run never repeats a key, so under these engines - // every key is a lone version and needs no merge. Retract - // rows are still dropped per batch below: older Rust writers - // recorded `delete_row_count = 0` for files holding them. - let mut skip_merge = skip_single_runs + // One sorted run should not repeat a key, so under these + // engines every key is a lone version and needs no merge. + // Retract rows are still dropped per batch below: older Rust + // writers recorded `delete_row_count = 0` for files holding + // them. + let skip_merge = skip_single_runs && input_stream_count == 1 && matches!( config.merge_engine, MergeEngine::Deduplicate | MergeEngine::FirstRow ); - // A file written under another engine may hold several - // versions of a key (the aggregation writer kept them before - // #839), and ALTER can switch the engine afterwards. - if skip_merge { - 'files: for run in &merge_group { - for file in &run.files { - // Always check the stored schema: dynamic options and - // time travel can change the read schema's engine - // without changing its ID. Merging is always correct, - // so a missing schema file under the read schema's ID - // only costs the skip. - let schema = match config - .schema_manager - .schema(file.file.schema_id) - .await - { - Ok(schema) => schema, - Err(crate::Error::IoUnexpected { source, .. }) - if source.kind() == opendal::ErrorKind::NotFound - && file.file.schema_id == table_schema_id => - { - skip_merge = false; - break 'files; - } - Err(err) => Err(err)?, - }; - let engine = - CoreOptions::new(schema.options()).merge_engine()?; - if !matches!( - engine, - MergeEngine::Deduplicate | MergeEngine::FirstRow - ) { - skip_merge = false; - break 'files; - } - } - } - } let reject_retracts = config.merge_engine == MergeEngine::FirstRow && !CoreOptions::new(&config.table_options).ignore_delete(); @@ -915,24 +920,76 @@ impl KeyValueFileReader { } let mut merge_stream = if skip_merge { - let input = file_streams.pop().expect("one input stream"); + let mut input = file_streams.pop().expect("one input stream"); let columns: Vec = key_indices.iter().chain(&value_indices).copied().collect(); let schema = merge_output_schema.clone(); - input - .map(move |batch| { - let batch = - drop_lone_retracts(batch?, value_kind_index, reject_retracts)?; - let columns = - columns.iter().map(|&i| batch.column(i).clone()).collect(); - RecordBatch::try_new(schema.clone(), columns).map_err(|e| { - Error::UnexpectedError { - message: format!("Failed to project unmerged batch: {e}"), - source: Some(Box::new(e)), - } - }) + let project = move |batch: RecordBatch| -> crate::Result { + let batch = + drop_lone_retracts(batch, value_kind_index, reject_retracts)?; + let columns = + columns.iter().map(|&i| batch.column(i).clone()).collect(); + RecordBatch::try_new(schema.clone(), columns).map_err(|e| { + Error::UnexpectedError { + message: format!("Failed to project unmerged batch: {e}"), + source: Some(Box::new(e)), + } }) - .boxed() + }; + // Files written before #839 can repeat a key inside one + // run, whatever engine the schemas name now. Hold one + // batch back, and once a key repeats, merge from there. + let fallback = SortMergeReaderBuilder::new( + Vec::new(), + internal_schema.clone(), + key_indices.clone(), + seq_index, + value_kind_index, + user_sequence_indices.clone(), + value_indices.clone(), + merge_output_schema.clone(), + merge_function(&config, &merge_output_fields)?, + ) + .with_user_sequence_descending( + !CoreOptions::new(&config.table_options) + .sequence_field_sort_order_is_ascending(), + ); + let key_indices = key_indices.clone(); + let stream: ArrowRecordBatchStream = try_stream! { + let mut held: Option = None; + let mut rest: Option = None; + while let Some(batch) = input.next().await { + let batch = batch?; + if repeats_a_key(held.as_ref(), &batch, &key_indices) { + let pending: Vec> = held + .take() + .into_iter() + .chain(std::iter::once(batch)) + .map(Ok) + .collect(); + rest = Some(futures::stream::iter(pending).chain(input).boxed()); + break; + } + if let Some(previous) = held.replace(batch) { + yield project(previous)?; + } + } + match rest { + Some(rest) => { + let mut merged = fallback.with_streams(vec![rest]).build()?; + while let Some(batch) = merged.next().await { + yield batch?; + } + } + None => { + if let Some(previous) = held { + yield project(previous)?; + } + } + } + } + .boxed(); + stream } else { SortMergeReaderBuilder::new( file_streams, @@ -3057,6 +3114,61 @@ mod tests { assert_eq!(int_column(&batches, "value"), vec![11, 20]); } + #[tokio::test] + async fn single_run_read_merges_repeated_keys_under_the_stored_engine() { + // Before #839, copy_with_options could switch the writer to aggregation + // under the same schema ID, so a file can repeat a key even though every + // schema says deduplicate. + let file_io = test_file_io(); + let table_path = "memory:/kv_repeated_keys"; + let table = pk_table(&file_io, table_path, &[]); + write_schema_file(&table, table.schema()).await; + let a = + write_kv_versions(&file_io, table_path, "a.parquet", &[(1, 0, 10), (1, 1, 11)]).await; + let b = write_kv_versions(&file_io, table_path, "b.parquet", &[(2, 2, 20)]).await; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![a, b]) + .build() + .unwrap(); + let batches = read_kv_split(&table, split).await; + assert_eq!(int_column(&batches, "id"), vec![1, 2]); + assert_eq!(int_column(&batches, "value"), vec![11, 20]); + } + + #[tokio::test] + async fn single_run_read_merges_a_key_repeated_across_batches() { + // One-row batches: key 1 is emitted before key 2 repeats across a batch + // boundary, and key 3 comes after the switch to merging. + let file_io = test_file_io(); + let table_path = "memory:/kv_repeated_across_batches"; + let table = pk_table(&file_io, table_path, &[("read.batch-size", "1")]); + write_schema_file(&table, table.schema()).await; + let a = write_kv_versions( + &file_io, + table_path, + "a.parquet", + &[(1, 0, 10), (2, 1, 20), (2, 2, 21), (3, 3, 30)], + ) + .await; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![a]) + .build() + .unwrap(); + let batches = read_kv_split(&table, split).await; + assert_eq!(int_column(&batches, "id"), vec![1, 2, 3]); + assert_eq!(int_column(&batches, "value"), vec![10, 21, 30]); + } + #[tokio::test] async fn single_run_read_merges_after_time_travel_past_an_engine_alter() { // aggregation (schema 0) -> ALTER to deduplicate (schema 1). Reading diff --git a/crates/paimon/src/table/sort_merge.rs b/crates/paimon/src/table/sort_merge.rs index 579ca39b3..f9967a5d6 100644 --- a/crates/paimon/src/table/sort_merge.rs +++ b/crates/paimon/src/table/sort_merge.rs @@ -1258,6 +1258,12 @@ impl SortMergeReaderBuilder { self } + /// Replace the input streams, for a builder prepared before they exist. + pub(crate) fn with_streams(mut self, streams: Vec) -> Self { + self.streams = streams; + self + } + pub(crate) fn with_user_sequence_descending(mut self, descending: bool) -> Self { self.user_sequence_descending = descending; self From 7270d1e16acf28ba49fd1390ce9ad2f11cabc28f Mon Sep 17 00:00:00 2001 From: JeonDaehong Date: Mon, 2 Nov 2026 10:52:37 +0000 Subject: [PATCH 6/6] [core] Compare nullable keys null-safely when checking a single run for repeats --- crates/paimon/src/table/kv_file_reader.rs | 120 ++++++++++++++++++++-- 1 file changed, 113 insertions(+), 7 deletions(-) diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index f5038aa7f..b0652dd3e 100644 --- a/crates/paimon/src/table/kv_file_reader.rs +++ b/crates/paimon/src/table/kv_file_reader.rs @@ -243,15 +243,16 @@ fn repeats_a_key( let column = batch.column(index); if let Some(previous) = previous.filter(|_| at_boundary) { let last = previous.column(index).slice(previous.num_rows() - 1, 1); - match arrow_ord::cmp::eq(&last, &column.slice(0, 1)) { + match arrow_ord::cmp::not_distinct(&last, &column.slice(0, 1)) { Ok(equal) => at_boundary = equal.true_count() == 1, Err(_) => return true, } } if rows > 1 { - let Ok(equal) = - arrow_ord::cmp::eq(&column.slice(0, rows - 1), &column.slice(1, rows - 1)) - else { + let Ok(equal) = arrow_ord::cmp::not_distinct( + &column.slice(0, rows - 1), + &column.slice(1, rows - 1), + ) else { return true; }; adjacent = Some(match adjacent { @@ -2322,6 +2323,17 @@ mod tests { builder.build_serialized() } + fn nullable_int_key(value: Option) -> Vec { + match value { + Some(value) => int_key(value), + None => { + let mut builder = crate::spec::BinaryRowBuilder::new(1); + builder.set_null_at(0); + builder.build_serialized() + } + } + } + async fn write_multi_row_group_kv_file( file_io: &FileIO, table_path: &str, @@ -3117,6 +3129,20 @@ mod tests { table_path: &str, file_name: &str, rows: &[(i32, i64, i32)], + ) -> DataFileMeta { + let rows: Vec<_> = rows + .iter() + .map(|&(id, seq, value)| (Some(id), seq, value)) + .collect(); + write_kv_nullable_versions(file_io, table_path, file_name, &rows).await + } + + /// Like `write_kv_versions`, with a nullable `id`. + async fn write_kv_nullable_versions( + file_io: &FileIO, + table_path: &str, + file_name: &str, + rows: &[(Option, i64, i32)], ) -> DataFileMeta { let schema = crate::arrow::build_target_arrow_schema(&[ DataField::new( @@ -3138,7 +3164,7 @@ mod tests { vec![ Arc::new(Int64Array::from_iter_values(rows.iter().map(|row| row.1))), Arc::new(Int8Array::from_value(0, rows.len())), - Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.0))), + Arc::new(Int32Array::from_iter(rows.iter().map(|row| row.0))), Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.2))), ], ) @@ -3161,8 +3187,8 @@ mod tests { let mut file = dummy_data_file(file_name.to_string()); file.file_size = bytes.len() as i64; file.row_count = rows.len() as i64; - file.min_key = int_key(rows[0].0); - file.max_key = int_key(rows[rows.len() - 1].0); + file.min_key = nullable_int_key(rows[0].0); + file.max_key = nullable_int_key(rows[rows.len() - 1].0); file.min_sequence_number = rows.iter().map(|row| row.1).min().unwrap(); file.max_sequence_number = rows.iter().map(|row| row.1).max().unwrap(); file @@ -3291,6 +3317,86 @@ mod tests { assert_eq!(int_column(&batches, "value"), vec![10, 21, 30]); } + #[tokio::test] + async fn single_run_read_merges_repeated_keys_under_first_row() { + // First-row reads carry the file level next to the merge input, so the + // switch to merging has to use that schema too. The first version wins. + let file_io = test_file_io(); + let table_path = "memory:/kv_repeated_keys_first_row"; + let table = pk_table(&file_io, table_path, &[("merge-engine", "first-row")]); + write_schema_file(&table, table.schema()).await; + let a = write_kv_versions( + &file_io, + table_path, + "a.parquet", + &[(1, 0, 10), (2, 1, 20), (2, 2, 21), (3, 3, 30)], + ) + .await; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![a]) + .build() + .unwrap(); + let batches = read_kv_split(&table, split).await; + assert_eq!(int_column(&batches, "id"), vec![1, 2, 3]); + assert_eq!(int_column(&batches, "value"), vec![10, 20, 30]); + } + + #[tokio::test] + async fn single_run_read_merges_repeated_null_keys() { + // NULL == NULL is NULL under plain equality, so a repeated NULL key needs + // the null-safe comparison. Batch size 64 repeats it inside one batch, 1 + // repeats it across a batch boundary. + for batch_size in ["64", "1"] { + let file_io = test_file_io(); + let table_path = format!("memory:/kv_repeated_null_keys_{batch_size}"); + let table = pk_table( + &file_io, + &table_path, + &[ + ("primary-key.nullable", "true"), + ("read.batch-size", batch_size), + ], + ); + write_schema_file(&table, table.schema()).await; + let a = write_kv_nullable_versions( + &file_io, + &table_path, + "a.parquet", + &[(None, 0, 10), (None, 1, 11), (Some(2), 2, 20)], + ) + .await; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![a]) + .build() + .unwrap(); + let batches = read_kv_split(&table, split).await; + let ids: Vec> = batches + .iter() + .flat_map(|batch| { + let column = batch.column(batch.schema().index_of("id").unwrap()); + let ids = column.as_any().downcast_ref::().unwrap(); + ids.iter().collect::>() + }) + .collect(); + assert_eq!(ids, vec![None, Some(2)], "batch size {batch_size}"); + assert_eq!( + int_column(&batches, "value"), + vec![11, 20], + "batch size {batch_size}" + ); + } + } + #[tokio::test] async fn single_run_read_merges_after_time_travel_past_an_engine_alter() { // aggregation (schema 0) -> ALTER to deduplicate (schema 1). Reading