feat(table): extend incremental primary-key compaction

crates/paimon/src/table/compact.rs

@@ -115,20 +115,22 @@ impl<'a> CompactBuilder<'a> {
         ensure_supported(self.table)?;
         let options = self.table.schema().core_options();
         let lookup = options.try_changelog_producer()? == ChangelogProducer::Lookup
+            || (!self.table.schema().primary_keys().is_empty()
+                && options.merge_engine()? == MergeEngine::FirstRow)
             || (options.deletion_vectors_enabled()
                 && !self.table.schema().primary_keys().is_empty());
         let local_cache = self.table.file_io().local_cache();
         if lookup && local_cache.is_none() {
             return Err(Error::ConfigInvalid {
                 message:
-                    "Lookup or deletion-vector compaction requires local-cache.enabled=true on the catalog/FileIO"
+                    "Lookup, first-row or deletion-vector compaction requires local-cache.enabled=true on the catalog/FileIO"
                         .into(),
             });
         }
         Ok(CompactExecutor {
             table: self.table.clone(),
             resources: self.resources.clone(),
-            lookup_cache: if lookup && options.merge_engine()? == MergeEngine::Deduplicate {
+            lookup_cache: if lookup {
                 Some(Arc::new(primary_key::lookup_file::LookupFileCache::new(
                     local_cache.expect("lookup cache validated above"),
                     self.table.schema().options(),
@@ -820,13 +822,13 @@ pub(super) fn prepare_results(table: &Table, results: &[CompactResult]) -> Resul
         }
         let options = task.context.schema.core_options();
         let dv = task.is_primary_key() && options.deletion_vectors_enabled();
-        let ignores_retracts = dv
+        // First-row can leave every accepted key in retained history, even
+        // with ignore-delete=false, and therefore produce no data files.
+        let can_omit_keys = dv
             || (task.is_primary_key()
-                && options.ignore_delete()
-                && matches!(
-                    options.merge_engine()?,
-                    MergeEngine::FirstRow | MergeEngine::PartialUpdate
-                ));
+                && (options.merge_engine()? == MergeEngine::FirstRow
+                    || (options.ignore_delete()
+                        && options.merge_engine()? == MergeEngine::PartialUpdate)));
         if dv
             && result.files.iter().any(|file| {
                 file.delete_row_count != Some(0) || file.row_count > i64::from(u32::MAX) + 1
@@ -853,11 +855,10 @@ pub(super) fn prepare_results(table: &Table, results: &[CompactResult]) -> Resul
             collected.extend(updates.iter().cloned());
         }
         if task.is_primary_key()
-            && (output_rows > input_rows
-                || (input_rows > 0 && output_rows == 0 && !ignores_retracts))
+            && (output_rows > input_rows || (input_rows > 0 && output_rows == 0 && !can_omit_keys))
         {
             return Err(invalid(
-                "PK compaction must retain each key's merged state unless retracts are ignored",
+                "PK compaction must retain merged keys unless its merge mode can omit them",
             ));
         }
         let lookup = options.try_changelog_producer()? == ChangelogProducer::Lookup;

crates/paimon/src/table/compact/primary_key.rs

@@ -182,9 +182,38 @@ impl MergeFunction for CompactReducer {
     }
 }
 
+struct CompactFirstRow {
+    inner: Box<dyn MergeFunction>,
+    output_level: i32,
+}
+
+impl MergeFunction for CompactFirstRow {
+    fn merge(
+        &self,
+        rows: &[MergeRow],
+        batches: &[BufferedBatch],
+        columns: &[usize],
+        schema: &SchemaRef,
+    ) -> Result<MergeResult> {
+        // Validate all selected records before checking retained history, as
+        // Java FirstRowMergeFunctionWrapper does when feeding its reducer.
+        let result = self.inner.merge(rows, batches, columns, schema)?;
+        for row in rows {
+            // Only a lookup hit can be above the output level. Its accepted
+            // record stays in retained history; selected positive rows must
+            // still be rewritten into this compaction's output.
+            if matches!(row.value_kind, 0 | 2) && row.level(batches)? > self.output_level {
+                return Ok(MergeResult::Omit);
+            }
+        }
+        Ok(result)
+    }
+}
+
 fn new_merge_function(
     config: &KeyValueReadConfig,
     fields: &[DataField],
+    output_level: i32,
 ) -> Result<Box<dyn MergeFunction>> {
     if config.merge_engine == MergeEngine::Deduplicate {
         return Ok(Box::new(CompactDeduplicate));
@@ -220,7 +249,10 @@ fn new_merge_function(
     };
     Ok(if config.merge_engine == MergeEngine::FirstRow {
         // First-row must still reject retracts unless ignore-delete is enabled.
-        inner
+        Box::new(CompactFirstRow {
+            inner,
+            output_level,
+        })
     } else {
         Box::new(CompactReducer {
             inner,
@@ -314,14 +346,15 @@ pub(super) async fn execute(
     let schema = table.schema();
     let options = schema.core_options();
     let lookup_changelog = options.try_changelog_producer()? == ChangelogProducer::Lookup;
-    let drop_delete = options.deletion_vectors_enabled() && task.output_level() > 0;
+    let output_level = task.output_level();
+    let drop_delete = options.deletion_vectors_enabled() && output_level > 0;
     let reader = new_reader(table, resources.clone())?.with_lookup(lookup);
     let mut batches =
         reader.read_with_merge_function(std::slice::from_ref(split), move |config, fields| {
             if lookup_changelog {
-                lookup::new_merge_function(config, fields)
+                lookup::new_merge_function(config, fields, output_level)
             } else {
-                new_merge_function(config, fields)
+                new_merge_function(config, fields, output_level)
             }
         })?;
     let keys = schema.trimmed_primary_keys();

crates/paimon/src/table/compact/primary_key/lookup.rs

@@ -26,9 +26,7 @@ use arrow_schema::{DataType as ArrowType, SchemaRef};
 
 use crate::spec::{DataField, MergeEngine, RowKind, VALUE_KIND_FIELD_NAME};
 use crate::table::kv_file_reader::KeyValueReadConfig;
-use crate::table::sort_merge::{
-    lookup_rows, BufferedBatch, MergeFunction, MergeResult, MergeRow, FILE_LEVEL_METADATA_KEY,
-};
+use crate::table::sort_merge::{lookup_rows, BufferedBatch, MergeFunction, MergeResult, MergeRow};
 use crate::{Error, Result};
 
 struct ChangelogOptions {
@@ -81,10 +79,11 @@ struct LookupChangelogMergeFunction {
 pub(super) fn new_merge_function(
     config: &KeyValueReadConfig,
     fields: &[DataField],
+    output_level: i32,
 ) -> Result<Box<dyn MergeFunction>> {
     let options = ChangelogOptions::parse(&config.table_options)?;
     Ok(Box::new(LookupChangelogMergeFunction {
-        inner: super::new_merge_function(config, fields)?,
+        inner: super::new_merge_function(config, fields, output_level)?,
         first_row: config.merge_engine == MergeEngine::FirstRow,
         value_columns: fields
             .iter()
@@ -317,13 +316,13 @@ impl MergeFunction for LookupChangelogMergeFunction {
             return Ok((result, None));
         }
         if self.first_row {
-            let contains_history = rows.iter().any(|row| {
-                matches!(row.value_kind, 0 | 2)
-                    && matches!(&batches[row.batch_idx], BufferedBatch::Source(batch)
-                        if batch.schema().metadata().get(FILE_LEVEL_METADATA_KEY)
-                            .and_then(|level| level.parse::<i32>().ok())
-                            .is_some_and(|level| level > 0))
-            });
+            let mut contains_history = false;
+            for row in rows {
+                if matches!(row.value_kind, 0 | 2) && row.level(batches)? > 0 {
+                    contains_history = true;
+                    break;
+                }
+            }
             // FirstRowMergeFunctionWrapper emits the first add unchanged,
             // including UPDATE_AFTER; an existing accepted history emits none.
             let changelog = if contains_history {

crates/paimon/src/table/compact/primary_key/lookup_store.rs

@@ -20,7 +20,7 @@ use crate::io::FileRead;
 use crate::resource::ResourceContext;
 use crate::spec::{BinaryRowBuilder, DataField, DataFileMeta};
 use crate::table::merge_tree_split_generator::KeyComparator;
-use crate::table::sort_merge::{LookupProvider, FILE_LEVEL_METADATA_KEY};
+use crate::table::sort_merge::{with_level, LookupProvider};
 use crate::table::{DataSplit, Table};
 use crate::{Error, Result};
 use arrow_array::RecordBatch;
@@ -331,15 +331,7 @@ impl LookupProvider for LevelLookup {
                             .or_default()
                             .insert(position);
                     }
-                    let mut metadata = batch.schema().metadata().clone();
-                    metadata.insert(FILE_LEVEL_METADATA_KEY.into(), file.meta.level.to_string());
-                    found = Some(
-                        RecordBatch::try_new(
-                            Arc::new(batch.schema().as_ref().clone().with_metadata(metadata)),
-                            batch.columns().to_vec(),
-                        )
-                        .map_err(invalid)?,
-                    );
+                    found = Some(with_level(batch, file.meta.level)?);
                 }
             }
             if found.is_some() {

crates/paimon/src/table/compact/primary_key/planner.rs

@@ -19,17 +19,125 @@
 //! before forcing L0 progress, as ForceUpLevel0Compaction does, without a
 //! full-bucket override.
 
-use crate::spec::{ChangelogProducer, DataFileMeta, MergeEngine, TableSchema};
+use crate::spec::{
+    AggregationConfig, ChangelogProducer, DataField, DataFileMeta, DataType, MergeEngine,
+    PartialUpdateConfig, TableSchema,
+};
 use crate::{Error, Result};
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashSet};
 
 pub(in crate::table::compact) fn enabled(schema: &TableSchema) -> Result<bool> {
     let options = schema.core_options();
-    Ok(options.merge_engine()? == MergeEngine::Deduplicate
-        && matches!(
-            options.try_changelog_producer()?,
+    let producer = options.try_changelog_producer()?;
+    Ok(match options.merge_engine()? {
+        MergeEngine::Deduplicate => matches!(
+            producer,
             ChangelogProducer::None | ChangelogProducer::Input | ChangelogProducer::Lookup
-        ))
+        ),
+        MergeEngine::FirstRow => {
+            !options.deletion_vectors_enabled()
+                && matches!(
+                    producer,
+                    ChangelogProducer::None | ChangelogProducer::Lookup
+                )
+        }
+        MergeEngine::PartialUpdate | MergeEngine::Aggregation => {
+            !options.deletion_vectors_enabled()
+                && match producer {
+                    ChangelogProducer::Lookup => true,
+                    ChangelogProducer::None | ChangelogProducer::Input => {
+                        plain_state_mergeable(schema)?
+                    }
+                    _ => false,
+                }
+        }
+    })
+}
+
+// Plain queries combine all positive levels. Only states that can be folded
+// again without depending on the original row order may leave history behind.
+fn plain_state_mergeable(schema: &TableSchema) -> Result<bool> {
+    let options = schema.core_options();
+    let keys = schema.primary_keys();
+    match options.merge_engine()? {
+        MergeEngine::Aggregation => {
+            let config = AggregationConfig::new(schema.options());
+            // Global sequence fields use last_value, whose retract clears the
+            // sequence tuple. A materialized subset can then sort before older
+            // retained states and change the visible sequence fields.
+            if config.remove_record_on_delete() || !options.sequence_fields().is_empty() {
+                return Ok(false);
+            }
+            Ok(schema
+                .fields()
+                .iter()
+                .filter(|field| !keys.iter().any(|key| key == field.name()))
+                .all(|field| {
+                    mergeable_field(
+                        schema,
+                        field,
+                        config
+                            .agg_function_for_field(field.name())
+                            .or_else(|| config.default_agg_function())
+                            .unwrap_or("last_non_null_value"),
+                    )
+                }))
+        }
+        MergeEngine::PartialUpdate => {
+            if !options.ignore_delete() || !options.sequence_fields().is_empty() {
+                return Ok(false);
+            }
+            let config = PartialUpdateConfig::new(schema.options());
+            let functions = config.validated_aggregate_functions(schema.fields(), keys)?;
+            let groups = config.validated_sequence_groups(schema.fields(), keys)?;
+            let ordering = groups
+                .iter()
+                .flat_map(|group| group.sequence_fields.iter().cloned())
+                .collect::<HashSet<_>>();
+            // Another group can overwrite an ordering field with NULL, making
+            // its aggregate state impossible to replay into retained history.
+            if groups.iter().any(|group| {
+                group
+                    .protected_fields
+                    .iter()
+                    .any(|field| ordering.contains(field))
+            }) {
+                return Ok(false);
+            }
+            Ok(schema
+                .fields()
+                .iter()
+                .filter(|field| !keys.iter().any(|key| key == field.name()))
+                .all(|field| {
+                    ordering.contains(field.name())
+                        || functions
+                            .get(field.name())
+                            .is_some_and(|function| mergeable_field(schema, field, function))
+                }))
+        }
+        _ => Ok(false),
+    }
+}
+
+fn mergeable_field(schema: &TableSchema, field: &DataField, function: &str) -> bool {
+    match crate::spec::canonical_aggregator_name(function) {
+        "min" | "max" | "bool_and" | "bool_or" => true,
+        // Floating sums can round differently when regrouped. Checked integer
+        // sums can fail in a subset even if the complete merge does not overflow.
+        "sum" => {
+            matches!(
+                field.data_type(),
+                DataType::TinyInt(_)
+                    | DataType::SmallInt(_)
+                    | DataType::Int(_)
+                    | DataType::BigInt(_)
+            ) && !schema
+                .options()
+                .get(&format!("fields.{}.sum.fail-on-overflow", field.name()))
+                .is_some_and(|value| value.eq_ignore_ascii_case("true"))
+        }
+        _ => false,
+    }
 }
 
 fn number(

crates/paimon/src/table/kv_file_reader.rs

@@ -27,9 +27,9 @@
 
 use super::data_file_reader::{file_index_selection_to_local_ranges, DataFileReader};
 use super::sort_merge::{
-    AggregateMergeFunction, DeduplicateMergeFunction, FirstRowMergeFunction, LookupMergeFunction,
-    MergeFunction, PartialUpdateMergeFunction, QueryReducerMergeFunction, SortMergeReaderBuilder,
-    CHANGELOG_BATCH_METADATA_KEY, FILE_LEVEL_METADATA_KEY,
+    schema_with_level, with_level, AggregateMergeFunction, DeduplicateMergeFunction,
+    FirstRowMergeFunction, LookupMergeFunction, MergeFunction, PartialUpdateMergeFunction,
+    QueryReducerMergeFunction, SortMergeReaderBuilder, CHANGELOG_BATCH_METADATA_KEY,
 };
 use crate::arrow::format::MosaicPrefetchOptions;
 use crate::arrow::shredding::variant::assemble_variant_extraction_array;
@@ -353,10 +353,11 @@ impl KeyValueFileReader {
             &config.table_fields,
             &config.primary_keys,
         );
-        let include_file_levels = config
-            .table_options
-            .get("changelog-producer")
-            .is_some_and(|value| value.eq_ignore_ascii_case("lookup"));
+        let include_file_levels = config.merge_engine == MergeEngine::FirstRow
+            || config
+                .table_options
+                .get("changelog-producer")
+                .is_some_and(|value| value.eq_ignore_ascii_case("lookup"));
         Self {
             file_io,
             config,
@@ -642,6 +643,11 @@ impl KeyValueFileReader {
         internal_read_type.extend(user_fields.clone());
 
         let internal_schema = crate::arrow::build_read_arrow_schema(&internal_read_type)?;
+        let merge_input_schema = if self.include_file_levels {
+            schema_with_level(internal_schema.as_ref())?
+        } else {
+            internal_schema
+        };
 
         // Output schema: user's read_type order
         let output_schema = crate::arrow::build_read_arrow_schema(&self.config.read_type)?;
@@ -882,11 +888,7 @@ impl KeyValueFileReader {
                                 while let Some(batch) = file_stream.next().await {
                                     let batch = batch?;
                                     if include_file_levels {
-                                        let mut metadata = batch.schema().metadata().clone();
-                                        metadata.insert(FILE_LEVEL_METADATA_KEY.into(), file_level.to_string());
-                                        yield RecordBatch::try_new(
-                                            Arc::new(batch.schema().as_ref().clone().with_metadata(metadata)), batch.columns().to_vec(),
-                                        ).map_err(|error| Error::DataInvalid { message: format!("Invalid lookup input batch: {error}"), source: Some(Box::new(error)) })?;
+                                        yield with_level(batch, file_level)?;
                                     } else { yield batch; }
                                 }
                             }
@@ -911,7 +913,7 @@ impl KeyValueFileReader {
                     // still contain duplicate keys.
                     let mut merge_stream = SortMergeReaderBuilder::new(
                         file_streams,
-                        internal_schema.clone(),
+                        merge_input_schema.clone(),
                         key_indices.clone(),
                         seq_index,
                         value_kind_index,