fix: correct batch incremental diff reads

Author: Pandas886

crates/paimon/src/table/audit_log_table.rs

@@ -81,6 +81,7 @@ impl AuditLogTable {
     }
 
     pub fn to_arrow(&self, plan: &IncrementalPlan) -> crate::Result<ArrowRecordBatchStream> {
+        plan.validate()?;
         let read = self.wrapped.new_read_builder().new_read()?;
         read.to_audit_log_arrow(plan)
     }

crates/paimon/src/table/incremental_scan.rs

@@ -47,7 +47,7 @@ pub enum IncrementalScanMode {
 #[derive(Debug, Clone)]
 pub enum IncrementalSplit {
     Data(DataSplit),
-    /// Per-(partition, bucket) diff pair. Memory bounded by one bucket's data.
+    /// Per-(partition, bucket) diff pair.
     DiffPair {
         before: Vec<DataSplit>,
         after: Vec<DataSplit>,
@@ -66,6 +66,80 @@ impl IncrementalPlan {
         Self { mode, splits }
     }
 
+    pub fn try_new(
+        mode: IncrementalScanMode,
+        splits: Vec<IncrementalSplit>,
+    ) -> crate::Result<Self> {
+        let plan = Self::new(mode, splits);
+        plan.validate()?;
+        Ok(plan)
+    }
+
+    /// Validate the plan at every point it crosses into a reader.
+    ///
+    /// `new` is retained for source compatibility, so callers can still build
+    /// an invalid plan. Readers must call this method instead of assuming a
+    /// plan came from the scanner.
+    pub fn validate(&self) -> crate::Result<()> {
+        if self.mode == IncrementalScanMode::Auto {
+            return Err(crate::Error::DataInvalid {
+                message: "Incremental plan mode Auto must be resolved before consumption"
+                    .to_string(),
+                source: None,
+            });
+        }
+        if self.mode == IncrementalScanMode::Diff {
+            let mut before_snapshot_id = None;
+            let mut after_snapshot_id = None;
+            for split in &self.splits {
+                let IncrementalSplit::DiffPair { before, after } = split else {
+                    return Err(crate::Error::DataInvalid {
+                        message: "Diff incremental plan contains a Data split".to_string(),
+                        source: None,
+                    });
+                };
+                validate_diff_pair(before, after)?;
+                if let Some(snapshot_id) = before.first().map(DataSplit::snapshot_id) {
+                    if before_snapshot_id.is_some_and(|expected| expected != snapshot_id) {
+                        return Err(crate::Error::DataInvalid {
+                            message: "Diff plan contains different before snapshots".to_string(),
+                            source: None,
+                        });
+                    }
+                    before_snapshot_id = Some(snapshot_id);
+                }
+                if let Some(snapshot_id) = after.first().map(DataSplit::snapshot_id) {
+                    if after_snapshot_id.is_some_and(|expected| expected != snapshot_id) {
+                        return Err(crate::Error::DataInvalid {
+                            message: "Diff plan contains different after snapshots".to_string(),
+                            source: None,
+                        });
+                    }
+                    after_snapshot_id = Some(snapshot_id);
+                }
+            }
+            if let (Some(before), Some(after)) = (before_snapshot_id, after_snapshot_id) {
+                if before >= after {
+                    return Err(crate::Error::DataInvalid {
+                        message: "Diff plan before snapshot must be earlier than after snapshot"
+                            .to_string(),
+                        source: None,
+                    });
+                }
+            }
+        } else if self
+            .splits
+            .iter()
+            .any(|split| matches!(split, IncrementalSplit::DiffPair { .. }))
+        {
+            return Err(crate::Error::DataInvalid {
+                message: "Non-Diff incremental plan contains a DiffPair".to_string(),
+                source: None,
+            });
+        }
+        Ok(())
+    }
+
     /// Resolved mode (`Auto` already collapsed to `Delta` / `Changelog`).
     pub fn mode(&self) -> IncrementalScanMode {
         self.mode
@@ -86,6 +160,49 @@ impl IncrementalPlan {
     }
 }
 
+pub(crate) fn validate_diff_pair(before: &[DataSplit], after: &[DataSplit]) -> crate::Result<()> {
+    if before
+        .iter()
+        .chain(after)
+        .any(|split| split.row_ranges().is_some())
+    {
+        return Err(crate::Error::DataInvalid {
+            message: "Diff pair must not contain physical row ranges".to_string(),
+            source: None,
+        });
+    }
+    let first = before.first().or(after.first());
+    let Some(first) = first else {
+        return Ok(());
+    };
+    for side in [before, after] {
+        if let Some(first_in_side) = side.first() {
+            if side
+                .iter()
+                .any(|split| split.snapshot_id() != first_in_side.snapshot_id())
+            {
+                return Err(crate::Error::DataInvalid {
+                    message: "Diff pair side contains splits from different snapshots".to_string(),
+                    source: None,
+                });
+            }
+        }
+    }
+    for split in before.iter().chain(after) {
+        if split.partition() != first.partition()
+            || split.bucket() != first.bucket()
+            || split.bucket_path() != first.bucket_path()
+            || split.total_buckets() != first.total_buckets()
+        {
+            return Err(crate::Error::DataInvalid {
+                message: "Diff pair contains splits from different partition buckets".to_string(),
+                source: None,
+            });
+        }
+    }
+    Ok(())
+}
+
 /// Batch incremental scan over a snapshot id range.
 pub struct IncrementalScan<'a> {
     table: &'a Table,
@@ -202,7 +319,7 @@ impl<'a> IncrementalScan<'a> {
             let plan = self.scan.plan_snapshot_delta(&snapshot).await?;
             splits.extend(plan.splits().iter().cloned().map(IncrementalSplit::Data));
         }
-        Ok(IncrementalPlan::new(mode, splits))
+        IncrementalPlan::try_new(mode, splits)
     }
 
     async fn plan_changelog(&self, mode: IncrementalScanMode) -> crate::Result<IncrementalPlan> {
@@ -220,10 +337,15 @@ impl<'a> IncrementalScan<'a> {
             let plan = self.scan.plan_snapshot_changelog(&snapshot).await?;
             splits.extend(plan.splits().iter().cloned().map(IncrementalSplit::Data));
         }
-        Ok(IncrementalPlan::new(mode, splits))
+        IncrementalPlan::try_new(mode, splits)
     }
 
     async fn plan_diff(&self, mode: IncrementalScanMode) -> crate::Result<IncrementalPlan> {
+        if self.table.schema().primary_keys().is_empty() {
+            return Err(crate::Error::Unsupported {
+                message: "Batch incremental Diff requires a table with primary keys".to_string(),
+            });
+        }
         let core_options = CoreOptions::new(self.table.schema().options());
         if core_options.merge_engine()? != crate::spec::MergeEngine::Deduplicate {
             return Err(crate::Error::Unsupported {
@@ -269,6 +391,6 @@ impl<'a> IncrementalScan<'a> {
             splits.push(IncrementalSplit::DiffPair { before, after });
         }
 
-        Ok(IncrementalPlan::new(mode, splits))
+        IncrementalPlan::try_new(mode, splits)
     }
 }

crates/paimon/src/table/kv_file_reader.rs

@@ -70,6 +70,8 @@ pub(crate) struct KeyValueReadConfig {
     pub primary_keys: Vec<String>,
     pub merge_engine: MergeEngine,
     pub sequence_fields: Vec<String>,
+    /// Merge files from all supplied splits into one globally key-sorted stream.
+    pub merge_splits: bool,
 }
 
 /// Keep only the conjuncts of `predicates` that reference primary-key columns,
@@ -301,7 +303,15 @@ impl KeyValueFileReader {
             }
         }
 
-        let splits: Vec<DataSplit> = data_splits.to_vec();
+        let split_groups: Vec<Vec<DataSplit>> = if self.config.merge_splits {
+            vec![data_splits.to_vec()]
+        } else {
+            data_splits
+                .iter()
+                .cloned()
+                .map(|split| vec![split])
+                .collect()
+        };
         let file_io = self.file_io;
         let merge_engine = self.config.merge_engine;
         let schema_manager = self.config.schema_manager;
@@ -321,21 +331,23 @@ impl KeyValueFileReader {
         let merge_output_schema = build_target_arrow_schema(&merge_output_fields)?;
 
         Ok(try_stream! {
-            for split in &splits {
+            for split_group in &split_groups {
                 // DV mode should not reach KeyValueFileReader.
-                if split
-                    .data_deletion_files()
-                    .is_some_and(|files| files.iter().any(Option::is_some))
-                {
-                    Err(Error::Unsupported {
-                        message: "KeyValueFileReader does not support deletion vectors".to_string(),
-                    })?;
+                for split in split_group {
+                    if split
+                        .data_deletion_files()
+                        .is_some_and(|files| files.iter().any(Option::is_some))
+                    {
+                        Err(Error::Unsupported {
+                            message: "KeyValueFileReader does not support deletion vectors".to_string(),
+                        })?;
+                    }
                 }
-
                 // Create one stream per data file.
                 let mut file_streams: Vec<ArrowRecordBatchStream> = Vec::new();
 
-                for file_meta in split.data_files().to_vec() {
+                for split in split_group {
+                    for file_meta in split.data_files().to_vec() {
                     let data_fields: Option<Vec<DataField>> = if file_meta.schema_id != table_schema_id {
                         let data_schema = schema_manager.schema(file_meta.schema_id).await?;
                         Some(data_schema.fields().to_vec())
@@ -357,9 +369,10 @@ impl KeyValueFileReader {
                         file_meta,
                         data_fields,
                         None,
-                        None,
+                        split.row_ranges().map(|ranges| ranges.to_vec()),
                     )?;
                     file_streams.push(stream);
+                    }
                 }
 
                 if file_streams.is_empty() {