diff --git a/crates/paimon/src/arrow/format/orc.rs b/crates/paimon/src/arrow/format/orc.rs index 861151636..dd5308bf0 100644 --- a/crates/paimon/src/arrow/format/orc.rs +++ b/crates/paimon/src/arrow/format/orc.rs @@ -386,7 +386,11 @@ impl FormatFileReader for OrcFormatReader { let mut builder = builder.with_projection(projection); - if let Some(predicate) = orc_predicate { + // orc-rust 0.8 composes stripe predicates with row selections using + // and_then(), which expects already-filtered coordinates and can panic + // for skipped rows. Positional reads use the selection and the exact + // residual below; ordinary scans retain predicate pruning. + if let Some(predicate) = orc_predicate.filter(|_| row_selection.is_none()) { builder = builder.with_predicate(predicate); } diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index 89071ed7c..052fc6e90 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -92,6 +92,8 @@ const ROW_TRACKING_ENABLED_OPTION: &str = "row-tracking.enabled"; const ROW_TRACKING_PARTITION_GROUP_ON_COMMIT_OPTION: &str = "row-tracking.partition-group-on-commit"; const CLUSTERING_INCREMENTAL_OPTION: &str = "clustering.incremental"; +const COMPACTION_MIN_FILE_NUM_OPTION: &str = "compaction.min.file-num"; +const COMPACTION_SMALL_FILE_RATIO_OPTION: &str = "compaction.small-file-ratio"; pub(crate) const TABLE_TYPE_OPTION: &str = "type"; pub(crate) const PATH_OPTION: &str = "path"; @@ -1637,6 +1639,35 @@ impl<'a> CoreOptions<'a> { .unwrap_or(DEFAULT_TARGET_FILE_SIZE) } + /// Minimum append input files for an underfilled compaction task. Default is 5. + pub fn compaction_min_file_num(&self) -> crate::Result { + self.options + .get(COMPACTION_MIN_FILE_NUM_OPTION) + .map_or(Ok(5), |raw| { + raw.parse::() + .ok() + .filter(|count| *count >= 2) + .map(|count| count as usize) + .ok_or_else(|| crate::Error::ConfigInvalid { + message: format!("compaction.min.file-num must be at least 2, got '{raw}'"), + }) + }) + } + + /// Strict append small-file threshold: target size times the configured ratio. + /// `compaction.small-file-ratio` defaults to 0.7, matching Java. + pub fn compaction_file_size(&self) -> crate::Result { + let ratio = self + .options + .get(COMPACTION_SMALL_FILE_RATIO_OPTION) + .map_or(Some(0.7), |raw| raw.parse::().ok()) + .filter(|ratio| ratio.is_finite() && *ratio >= 0.0) + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "compaction.small-file-ratio must be a finite nonnegative number".into(), + })?; + Ok((self.target_file_size() as f64 * ratio) as i64) + } + /// Maximum automatically inferred postpone bucket count (Java default: 2048). pub fn postpone_batch_write_fixed_bucket_max_parallelism(&self) -> crate::Result { let key = "postpone.batch-write-fixed-bucket.max-parallelism"; diff --git a/crates/paimon/src/spec/stats.rs b/crates/paimon/src/spec/stats.rs index 6e4017bc7..63263ea83 100644 --- a/crates/paimon/src/spec/stats.rs +++ b/crates/paimon/src/spec/stats.rs @@ -149,57 +149,66 @@ impl Display for BinaryTableStats { } } -/// Compute per-column independent min/max/null_count for the specified columns -/// in a RecordBatch. Each entry in `col_indices` is the column index in the batch, -/// and the corresponding entry in `col_types` is its Paimon DataType. -pub fn compute_column_stats( - batch: &RecordBatch, - col_indices: &[usize], - col_types: &[DataType], -) -> crate::Result { - let num_cols = col_indices.len(); - let num_rows = batch.num_rows(); - let mut min_datums: Vec> = vec![None; num_cols]; - let mut max_datums: Vec> = vec![None; num_cols]; - let mut null_counts: Vec> = vec![Some(0); num_cols]; +/// Incremental full column statistics without retaining the input batches. +pub(crate) struct ColumnStatsCollector<'a> { + types: &'a [DataType], + min_datums: Vec>, + max_datums: Vec>, + null_counts: Vec>, +} - for row_idx in 0..num_rows { - for (pos, (&col_idx, data_type)) in col_indices.iter().zip(col_types.iter()).enumerate() { - let datum = extract_datum_from_arrow(batch, row_idx, col_idx, data_type)?; - match datum { - Some(d) => { - if min_datums[pos].as_ref().is_none_or(|m| d < *m) { - min_datums[pos] = Some(d.clone()); +impl<'a> ColumnStatsCollector<'a> { + pub(crate) fn new(types: &'a [DataType]) -> Self { + Self { + types, + min_datums: vec![None; types.len()], + max_datums: vec![None; types.len()], + null_counts: vec![Some(0); types.len()], + } + } + + pub(crate) fn collect(&mut self, batch: &RecordBatch, indices: &[usize]) -> crate::Result<()> { + for row_idx in 0..batch.num_rows() { + for (pos, (&col_idx, data_type)) in indices.iter().zip(self.types.iter()).enumerate() { + let datum = extract_datum_from_arrow(batch, row_idx, col_idx, data_type)?; + match datum { + Some(d) => { + if self.min_datums[pos].as_ref().is_none_or(|m| d < *m) { + self.min_datums[pos] = Some(d.clone()); + } + if self.max_datums[pos].as_ref().is_none_or(|m| d > *m) { + self.max_datums[pos] = Some(d); + } } - if max_datums[pos].as_ref().is_none_or(|m| d > *m) { - max_datums[pos] = Some(d); + None => { + *self.null_counts[pos].as_mut().unwrap() += 1; } } - None => { - *null_counts[pos].as_mut().unwrap() += 1; - } } } + Ok(()) } - let mut min_builder = BinaryRowBuilder::new(num_cols as i32); - let mut max_builder = BinaryRowBuilder::new(num_cols as i32); - for (pos, data_type) in col_types.iter().enumerate() { - match &min_datums[pos] { - Some(d) => min_builder.write_datum(pos, d, data_type), - None => min_builder.set_null_at(pos), - } - match &max_datums[pos] { - Some(d) => max_builder.write_datum(pos, d, data_type), - None => max_builder.set_null_at(pos), + pub(crate) fn finish(&self) -> BinaryTableStats { + let mut min_builder = BinaryRowBuilder::new(self.types.len() as i32); + let mut max_builder = BinaryRowBuilder::new(self.types.len() as i32); + for (pos, data_type) in self.types.iter().enumerate() { + match &self.min_datums[pos] { + Some(d) => min_builder.write_datum(pos, d, data_type), + None => min_builder.set_null_at(pos), + } + match &self.max_datums[pos] { + Some(d) => max_builder.write_datum(pos, d, data_type), + None => max_builder.set_null_at(pos), + } } - } - Ok(BinaryTableStats::new( - min_builder.build_serialized(), - max_builder.build_serialized(), - null_counts, - )) + BinaryTableStats::new( + min_builder.build_serialized(), + max_builder.build_serialized(), + self.null_counts.clone(), + ) + } } #[cfg(test)] @@ -207,6 +216,31 @@ mod tests { use super::*; use crate::spec::BinaryRow; + #[test] + fn streaming_stats_keep_projected_bounds_and_nulls_across_batches() { + use arrow_array::Int32Array; + use std::sync::Arc; + let types = vec![DataType::Int(crate::spec::IntType::new()); 2]; + let mut collector = ColumnStatsCollector::new(&types); + for (x, y) in [ + (vec![Some(5), None], vec![Some(100), Some(20)]), + (vec![Some(-3), Some(9)], vec![None, Some(200)]), + ] { + let batch = RecordBatch::try_from_iter([ + ("x", Arc::new(Int32Array::from(x)) as arrow_array::ArrayRef), + ("y", Arc::new(Int32Array::from(y)) as arrow_array::ArrayRef), + ]) + .unwrap(); + collector.collect(&batch, &[1, 0]).unwrap(); + } + let stats = collector.finish(); + let min = BinaryRow::from_serialized_bytes(stats.min_values()).unwrap(); + let max = BinaryRow::from_serialized_bytes(stats.max_values()).unwrap(); + assert_eq!((min.get_int(0).unwrap(), min.get_int(1).unwrap()), (20, -3)); + assert_eq!((max.get_int(0).unwrap(), max.get_int(1).unwrap()), (200, 9)); + assert_eq!(stats.null_counts(), &[Some(1), Some(1)]); + } + /// Empty stats must produce min/max bytes that the Java side's /// `SerializationUtils.deserializeBinaryRow` accepts: at minimum a 4-byte BE /// arity prefix. A bare `Vec::new()` would trigger `BufferUnderflowException` diff --git a/crates/paimon/src/table/compact.rs b/crates/paimon/src/table/compact.rs new file mode 100644 index 000000000..abfec835b --- /dev/null +++ b/crates/paimon/src/table/compact.rs @@ -0,0 +1,733 @@ +// 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. + +//! Snapshot-bound append compaction and full-bucket primary-key compaction. + +mod append; +mod primary_key; + +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::sync::Arc; + +use futures::TryStreamExt; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use super::data_file_index_writer::FileIndexOptions; +use super::data_file_writer::DataFileWriter; +use super::partition_filter::validate_partition_filter; +use super::snapshot_manifest::read_snapshot_entries; +use super::write_builder::{ensure_table_write_authorized, validate_commit_user}; +use super::{DataSplit, Table, TableCommit}; +use crate::resource::ResourceContext; +use crate::spec::{ + bucket_path_under, BinaryRow, ChangelogProducer, DataFileMeta, FileKind, IndexManifest, + ManifestEntry, MergeEngine, PartitionComputer, Predicate, Snapshot, TableSchema, +}; +use crate::{Error, Result}; + +/// Creates planners, executors and a committer with one shared configuration. +/// Output size and append task selection use the captured table options. +/// Supports ordinary append tables and fixed or dynamic bucket PK tables, +/// without row tracking, deletion vectors, data evolution, query authorization +/// or dedicated blob/vector storage. +pub struct CompactBuilder<'a> { + table: &'a Table, + partition_filter: Option, + resources: Option, + commit_user: String, +} + +impl<'a> CompactBuilder<'a> { + pub(super) fn new(table: &'a Table) -> Self { + Self { + table, + partition_filter: None, + resources: None, + commit_user: uuid::Uuid::new_v4().to_string(), + } + } + + /// Restrict planning to partitions. Repeated calls combine filters with AND. + pub fn with_partition_filter(mut self, filter: Predicate) -> Result { + validate_partition_filter(self.table, &filter)?; + self.partition_filter = Some(match self.partition_filter { + Some(previous) => Predicate::and(vec![previous, filter]), + None => filter, + }); + Ok(self) + } + + pub fn with_resources(mut self, resources: ResourceContext) -> Self { + self.resources = Some(resources); + self + } + + pub fn with_commit_user(mut self, commit_user: impl Into) -> Result { + let commit_user = commit_user.into(); + validate_commit_user(&commit_user)?; + self.commit_user = commit_user; + Ok(self) + } + + pub fn new_planner(&self) -> Result { + ensure_supported(self.table)?; + let options = self.table.schema().core_options(); + if options.target_file_size() <= 0 { + return Err(Error::ConfigInvalid { + message: "Compaction requires a positive target-file-size".into(), + }); + } + let append = if self.table.schema().primary_keys().is_empty() { + Some(append::AppendCompactPlanner::new(&options)?) + } else { + None + }; + Ok(CompactPlanner { + table: self.table.clone(), + append, + partition_filter: self.partition_filter.clone(), + }) + } + + /// Executor resources are local; input files and output settings come from each task. + pub fn new_executor(&self) -> Result { + ensure_supported(self.table)?; + Ok(CompactExecutor { + table: self.table.clone(), + resources: self.resources.clone(), + }) + } + + pub fn new_commit(&self) -> Result { + ensure_supported(self.table)?; + Ok(TableCommit::new( + self.table.clone(), + self.commit_user.clone(), + )) + } +} + +/// Selects append small files or whole PK buckets from each captured snapshot. +pub struct CompactPlanner { + table: Table, + append: Option, + partition_filter: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct CompactContext { + plan_id: String, + location: String, + identifier: String, + branch: String, + schema: TableSchema, + snapshot: Snapshot, +} + +/// An immutable, serializable set of independently executable compaction tasks. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CompactPlan { + snapshot_id: Option, + tasks: Vec, +} + +impl CompactPlan { + pub fn snapshot_id(&self) -> Option { + self.snapshot_id + } + + pub fn tasks(&self) -> &[CompactTask] { + &self.tasks + } +} + +/// Complete metadata for rewriting files in one partition and bucket. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CompactTask { + context: CompactContext, + task_id: usize, + partition: Vec, + bucket: i32, + total_buckets: i32, + files: Vec, +} + +impl CompactTask { + pub fn snapshot_id(&self) -> i64 { + self.context.snapshot.id() + } + + pub fn partition(&self) -> &[u8] { + &self.partition + } + + pub fn bucket(&self) -> i32 { + self.bucket + } + + pub fn files(&self) -> &[DataFileMeta] { + &self.files + } +} + +impl CompactPlanner { + /// Capture a snapshot without writing data or publishing metadata. + pub async fn plan(&self) -> Result { + ensure_supported(&self.table)?; + let mut reader = self.table.new_read_builder(); + if let Some(filter) = &self.partition_filter { + reader.with_filter(filter.clone()); + } + let source = reader.new_scan().plan_all().await?; + let Some(snapshot_id) = source.snapshot_id() else { + return Ok(CompactPlan { + snapshot_id: None, + tasks: vec![], + }); + }; + let snapshot = self + .table + .snapshot_manager() + .get_snapshot(snapshot_id) + .await?; + validate_snapshot_indexes(&self.table, &snapshot).await?; + let context = CompactContext { + plan_id: uuid::Uuid::new_v4().to_string(), + location: self.table.location().to_string(), + identifier: self.table.identifier().full_name(), + branch: self.table.branch().to_string(), + schema: self.table.schema().clone(), + snapshot, + }; + let mut groups = BTreeMap::<_, Vec>::new(); + let primary_key = !self.table.schema().primary_keys().is_empty(); + for split in source.splits() { + let group = groups + .entry(( + split.partition().to_serialized_bytes(), + split.bucket(), + split.total_buckets(), + )) + .or_default(); + for file in split.data_files() { + validate_file(file, primary_key)?; + group.push(file.clone()); + } + } + let mut tasks = Vec::new(); + for ((partition, bucket, total_buckets), files) in groups { + if primary_key { + // L0 requires upgrading; multiple levels require merging. Files + // at one positive level already form a non-overlapping sorted + // run, including when the previous compaction rolled files. + let first_level = files.first().map(|file| file.level); + if files + .iter() + .any(|file| file.level == 0 || Some(file.level) != first_level) + { + tasks.push(CompactTask { + context: context.clone(), + task_id: tasks.len(), + partition, + bucket, + total_buckets, + files, + }); + } + continue; + } + for files in self.append.as_ref().unwrap().plan(files) { + tasks.push(CompactTask { + context: context.clone(), + task_id: tasks.len(), + partition: partition.clone(), + bucket, + total_buckets, + files, + }); + } + } + Ok(CompactPlan { + snapshot_id: Some(snapshot_id), + tasks, + }) + } +} + +/// Streams task inputs into new data files without publishing a snapshot. +pub struct CompactExecutor { + table: Table, + resources: Option, +} + +/// Prepared files owned by the caller until submitted to `TableCommit::commit_compact`. +/// A failed or uncertain commit must retain these files for retry. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CompactResult { + task: CompactTask, + files: Vec, +} + +impl CompactResult { + pub fn files(&self) -> &[DataFileMeta] { + &self.files + } + + pub fn metrics(&self) -> CompactMetrics { + CompactMetrics { + files_removed: self.task.files.len(), + files_added: self.files.len(), + rows_written: self.files.iter().map(|file| file.row_count as u64).sum(), + bytes_written: self.files.iter().map(|file| file.file_size as u64).sum(), + } + } +} + +/// Counts for the submitted compaction results; removed files remain for time travel. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct CompactMetrics { + pub files_removed: usize, + pub files_added: usize, + pub rows_written: u64, + pub bytes_written: u64, +} + +impl CompactExecutor { + pub async fn execute(&self, task: &CompactTask) -> Result { + ensure_supported(&self.table)?; + task.validate_target(&self.table)?; + let manager = self.table.snapshot_manager(); + let snapshot = manager.get_snapshot(task.snapshot_id()).await?; + if !task.context.snapshot.has_same_read_metadata(&snapshot) { + return Err(invalid("Compaction snapshot was replaced")); + } + validate_snapshot_indexes(&self.table, &snapshot).await?; + let active = read_snapshot_entries(&manager, &snapshot, None).await?; + task.validate_inputs(&active)?; + let table = self + .table + .copy_with_resolved_schema(task.context.schema.clone(), &task.context.branch)? + .copy_with_pinned_snapshot(Some(&snapshot)); + let options = table.schema().core_options(); + let partition = BinaryRow::from_serialized_bytes(&task.partition)?; + let computer = PartitionComputer::new( + table.schema().partition_keys(), + table.schema().fields(), + options.partition_default_name(), + options.legacy_partition_name(), + )?; + let partition_path = if table.schema().partition_keys().is_empty() { + String::new() + } else { + computer.generate_partition_path(&partition)? + }; + let bucket_path = + bucket_path_under(&table.data_file_location(), &partition_path, task.bucket); + let split = DataSplit::builder() + .with_snapshot(snapshot.id()) + .with_partition(partition) + .with_bucket(task.bucket) + .with_total_buckets(task.total_buckets) + .with_bucket_path(bucket_path) + .with_data_files(task.files.clone()) + .build()?; + if task.is_primary_key() { + let files = primary_key::execute(&table, task, &split, self.resources.clone()).await?; + return Ok(CompactResult { + task: task.clone(), + files, + }); + } + let mut read = table.new_read_builder().new_read()?.with_ordered(true); + if let Some(resources) = &self.resources { + read = read.with_resources(resources.clone()); + } + let mut batches = read.to_arrow(&[split])?; + let index = FileIndexOptions::parse(table.schema().options(), table.schema().fields())?; + let mut writer = DataFileWriter::new( + table.file_io().clone(), + table.location().to_string(), + partition_path, + task.bucket, + table.schema().id(), + options.target_file_size(), + options.file_compression().to_string(), + options.file_compression_zstd_level(), + options.write_buffer_size()? as i64, + options.file_format(), + table.schema().fields().to_vec(), + table.schema().options().clone(), + Some(1), // FileSource.COMPACT + None, + None, + )? + .with_file_index(index.map(Arc::new)) + .with_target_file_row_num(options.target_file_row_num()?) + .with_resources(self.resources.clone()); + let written = async { + let mut rows = 0u64; + while let Some(batch) = batches.try_next().await? { + rows += batch.num_rows() as u64; + writer.write(&batch).await?; + } + if rows + != task + .files + .iter() + .map(|file| file.row_count as u64) + .sum::() + { + return Err(invalid("Compaction input row count changed")); + } + writer.prepare_commit().await + } + .await; + match written { + Ok(files) => Ok(CompactResult { + task: task.clone(), + files, + }), + Err(error) => { + writer.abort().await; + Err(error) + } + } + } +} + +impl CompactTask { + fn is_primary_key(&self) -> bool { + !self.context.schema.primary_keys().is_empty() + } + + fn output_level(&self) -> i32 { + if self.is_primary_key() { + self.files + .iter() + .map(|file| file.level) + .max() + .unwrap_or(0) + .max(1) + } else { + 0 + } + } + + fn validate_inputs(&self, active: &[ManifestEntry]) -> Result<()> { + let bucket_files = active + .iter() + .filter(|entry| { + entry.partition() == self.partition + && entry.bucket() == self.bucket + && entry.total_buckets() == self.total_buckets + }) + .map(|entry| (entry.file().file_name.as_str(), entry)) + .collect::>(); + if self.is_primary_key() && bucket_files.len() != self.files.len() { + return Err(invalid( + "PK compaction must include the complete captured bucket", + )); + } + for file in &self.files { + if !bucket_files + .get(file.file_name.as_str()) + .is_some_and(|entry| entry.file() == file) + { + return Err(invalid( + "Compaction input does not match the captured snapshot", + )); + } + } + Ok(()) + } + + fn validate_target(&self, table: &Table) -> Result<()> { + if table.location() != self.context.location + || table.identifier().full_name() != self.context.identifier + || table.branch() != self.context.branch + || table.schema() != &self.context.schema + { + return Err(invalid( + "Compaction task belongs to a different table, branch or schema", + )); + } + ensure_snapshot_supported(table, &self.context.snapshot)?; + let minimum_files = if self.is_primary_key() { 1 } else { 2 }; + if table.schema().core_options().target_file_size() <= 0 + || self.files.len() < minimum_files + || self.bucket < 0 + { + return Err(invalid("Invalid compaction task")); + } + let mut names = HashSet::new(); + for file in &self.files { + validate_file(file, self.is_primary_key())?; + if !names.insert(&file.file_name) { + return Err(invalid("Duplicate compaction input file")); + } + } + Ok(()) + } +} + +/// Validated inputs for the existing atomic snapshot publication machinery. +pub(super) struct PreparedCompact { + pub entries: Vec, + pub schema_id: i64, + pub snapshot_id: i64, + pub read_snapshot: Option>, + pub commit_suffix: String, + pub metrics: CompactMetrics, +} + +pub(super) fn prepare_results(table: &Table, results: &[CompactResult]) -> Result { + ensure_supported(table)?; + let first = results + .first() + .ok_or_else(|| invalid("No compaction results"))?; + let mut ids = Vec::new(); + let mut inputs = HashSet::new(); + let mut outputs = HashSet::new(); + let mut entries = Vec::new(); + let mut metrics = CompactMetrics::default(); + for result in results { + let task = &result.task; + task.validate_target(table)?; + if task.context.plan_id != first.task.context.plan_id + || !task + .context + .snapshot + .has_same_read_metadata(&first.task.context.snapshot) + { + return Err(invalid("Compaction results must belong to one plan")); + } + ids.push(task.task_id); + for file in &task.files { + if !inputs.insert((task.partition.clone(), task.bucket, file.file_name.clone())) { + return Err(invalid("Overlapping compaction results")); + } + entries.push(ManifestEntry::new( + FileKind::Delete, + task.partition.clone(), + task.bucket, + task.total_buckets, + file.clone(), + 2, + )); + } + let input_rows = task + .files + .iter() + .map(|file| file.row_count as u64) + .sum::(); + let mut output_rows = 0u64; + for file in &result.files { + validate_file(file, task.is_primary_key())?; + if file.schema_id != task.context.schema.id() + || file.file_source != Some(1) + || file.level != task.output_level() + || !outputs.insert(file.file_name.clone()) + || task.files.iter().any(|old| old.file_name == file.file_name) + { + return Err(invalid("Invalid compaction output file")); + } + output_rows += file.row_count as u64; + entries.push(ManifestEntry::new( + FileKind::Add, + task.partition.clone(), + task.bucket, + task.total_buckets, + file.clone(), + 2, + )); + } + if !task.is_primary_key() && input_rows != output_rows { + return Err(invalid("Compaction must preserve the input row count")); + } + let options = task.context.schema.core_options(); + let ignores_retracts = task.is_primary_key() + && options.ignore_delete() + && matches!( + options.merge_engine()?, + MergeEngine::FirstRow | MergeEngine::PartialUpdate + ); + if task.is_primary_key() + && (output_rows > input_rows + || (input_rows > 0 && output_rows == 0 && !ignores_retracts)) + { + return Err(invalid( + "PK compaction must retain each key's merged state unless retracts are ignored", + )); + } + let part = result.metrics(); + metrics.files_removed += part.files_removed; + metrics.files_added += part.files_added; + metrics.rows_written += part.rows_written; + metrics.bytes_written += part.bytes_written; + } + ids.sort_unstable(); + if ids.windows(2).any(|pair| pair[0] == pair[1]) { + return Err(invalid("Duplicate compaction task")); + } + // Stable across result ordering and task retries; distinct subsets may commit separately. + let mut digest = Sha256::new(); + digest.update(first.task.context.plan_id.as_bytes()); + for id in ids { + digest.update((id as u64).to_le_bytes()); + } + Ok(PreparedCompact { + entries, + schema_id: first.task.context.schema.id(), + snapshot_id: first.task.snapshot_id(), + read_snapshot: first + .task + .is_primary_key() + .then(|| Box::new(first.task.context.snapshot.clone())), + commit_suffix: hex::encode(digest.finalize()), + metrics, + }) +} + +pub(super) fn ensure_supported(table: &Table) -> Result<()> { + ensure_table_write_authorized(table)?; + table.ensure_not_branch_reference_for_write()?; + let options = table.schema().core_options(); + let properties = table.schema().options(); + if table.schema().primary_keys().is_empty() + && options.bucket() > 0 + && !properties + .get("bucket-append-ordered") + .is_some_and(|value| value.eq_ignore_ascii_case("false")) + { + // Ordered append needs contiguous sequence windows and rewrite counters. + // Ordinary Rust append writers do not yet restore those counters. + return Err(Error::Unsupported { + message: "Compaction of ordered bucketed append tables is not supported; bucket-append-ordered=false is required".into(), + }); + } + if properties + .get("pk-clustering-override") + .is_some_and(|value| value.eq_ignore_ascii_case("true")) + || properties.contains_key("record-level.expire-time") + || (table.schema().primary_keys().is_empty() && options.clustering_incremental_enabled()) + { + // Clustered KV files are not necessarily primary-key ordered; record + // expiration must filter versions before merging, as Java does. + return Err(Error::Unsupported { + message: + "Compaction does not support pk-clustering-override, append incremental clustering or record-level.expire-time" + .into(), + }); + } + if table.is_format_table() + || options.row_tracking_enabled() + || options.data_evolution_enabled() + || options.deletion_vectors_enabled() + || options.query_auth_enabled() + || !options.blob_fields().is_empty() + || table + .schema() + .fields() + .iter() + .any(|field| field.data_type().is_blob_file_field()) + || options.vector_file_format().is_some() + || options.rowkind_field().is_some() + || options.bucket() < -1 + { + return Err(Error::Unsupported { + message: "Compaction does not support row tracking, data evolution, deletion vectors, query authorization, rowkind fields, postpone buckets or dedicated blob/vector storage".into(), + }); + } + if !table.schema().primary_keys().is_empty() + && (options.bucket() == 0 + || !matches!( + options.try_changelog_producer()?, + ChangelogProducer::None | ChangelogProducer::Input + ) + || (options.merge_engine()? == MergeEngine::FirstRow + && options.try_changelog_producer()? != ChangelogProducer::None)) + { + return Err(Error::Unsupported { + message: "PK compaction requires fixed or dynamic buckets and changelog-producer=none or input; first-row requires changelog-producer=none".into(), + }); + } + options.validate_data_file_path_directory() +} + +fn ensure_snapshot_supported(table: &Table, snapshot: &Snapshot) -> Result<()> { + if snapshot.index_manifest().is_some() + && (table.schema().primary_keys().is_empty() + || table.schema().core_options().bucket() != -1) + { + return Err(Error::Unsupported { + message: "Compaction only supports dynamic-bucket HASH index manifests".into(), + }); + } + Ok(()) +} + +/// HASH indexes contain key-to-bucket assignments, independent of data files. +/// Preserve their manifest verbatim, including hashes of deleted keys. Other +/// index types may address rewritten files or row positions and cannot be reused. +pub(super) async fn validate_snapshot_indexes( + table: &Table, + snapshot: &Snapshot, +) -> Result> { + ensure_snapshot_supported(table, snapshot)?; + let Some(name) = snapshot.index_manifest() else { + return Ok(None); + }; + let path = table.snapshot_manager().manifest_path(name); + let entries = IndexManifest::read(table.file_io(), &path).await?; + if entries.iter().any(|entry| { + entry.index_file.index_type != "HASH" + || entry.index_file.deletion_vectors_ranges.is_some() + || entry.index_file.global_index_meta.is_some() + }) { + return Err(Error::Unsupported { + message: "Compaction only supports dynamic-bucket HASH index manifests".into(), + }); + } + Ok(Some(name.to_string())) +} + +fn validate_file(file: &DataFileMeta, primary_key: bool) -> Result<()> { + if file.row_count < 0 + || file.file_size < 0 + || file.first_row_id.is_some() + || file.write_cols.is_some() + || file.level < 0 + || file + .delete_row_count + .is_some_and(|count| count < 0 || (!primary_key && count != 0)) + || file.file_name.contains(".vector.") + || super::dedicated_format_file_writer::is_blob_or_video_file_name(&file.file_name) + { + return Err(Error::Unsupported { + message: "Compaction requires complete data files with known row counts".into(), + }); + } + Ok(()) +} + +fn invalid(message: &str) -> Error { + Error::DataInvalid { + message: message.into(), + source: None, + } +} diff --git a/crates/paimon/src/table/compact/append.rs b/crates/paimon/src/table/compact/append.rs new file mode 100644 index 000000000..8e4f1e63f --- /dev/null +++ b/crates/paimon/src/table/compact/append.rs @@ -0,0 +1,274 @@ +// 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. + +use crate::spec::{CoreOptions, DataFileMeta}; +use crate::Result; +use std::collections::VecDeque; + +/// Unaware batch packing and bucketed best-effort windows from Java. +pub(super) struct AppendCompactPlanner { + bucketed: bool, + small_file_size: i64, + target_weight: i128, + open_file_cost: i64, + min_file_num: usize, +} + +impl AppendCompactPlanner { + pub(super) fn new(options: &CoreOptions<'_>) -> Result { + Ok(Self { + bucketed: options.bucket() > 0, + small_file_size: options.compaction_file_size()?, + target_weight: i128::from(options.target_file_size()) * 2, + open_file_cost: options.source_split_open_file_cost(), + min_file_num: options.compaction_min_file_num()?, + }) + } + + pub(super) fn plan(&self, mut files: Vec) -> Vec> { + if self.bucketed { + return self.plan_bucketed(files); + } + files.retain(|file| file.file_size < self.small_file_size); + // Size ordering matches Java; sequence/name break ties deterministically + // so regrouping input read splits cannot change task boundaries. + files.sort_by(|a, b| { + (a.file_size, a.min_sequence_number, &a.file_name).cmp(&( + b.file_size, + b.min_sequence_number, + &b.file_name, + )) + }); + let mut tasks = Vec::new(); + let mut pending = Vec::new(); + let mut weight = 0i128; + for file in files { + weight += i128::from(file.file_size) + i128::from(self.open_file_cost); + pending.push(file); + if pending.len() > 1 && weight >= self.target_weight { + tasks.push(std::mem::take(&mut pending)); + weight = 0; + } + } + if pending.len() >= self.min_file_num { + tasks.push(pending); + } + tasks + } + + /// Java BucketedAppendCompactManager.pickCompactBefore: a sequence-ordered + /// sliding window, without the unaware coordinator's size filter/open cost. + fn plan_bucketed(&self, mut files: Vec) -> Vec> { + files.sort_by(|a, b| { + (a.min_sequence_number, &a.file_name).cmp(&(b.min_sequence_number, &b.file_name)) + }); + let mut tasks = Vec::new(); + let mut pending = VecDeque::new(); + let mut size = 0i128; + for file in files { + size += i128::from(file.file_size); + pending.push_back(file); + if pending.len() >= self.min_file_num { + tasks.push(pending.drain(..).collect()); + size = 0; + } else if size >= self.target_weight { + size -= i128::from(pending.pop_front().unwrap().file_size); + } + } + tasks + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::stats::BinaryTableStats; + use std::collections::HashMap; + + fn files(sizes: &[i64]) -> Vec { + sizes + .iter() + .enumerate() + .map(|(i, &size)| DataFileMeta { + file_name: format!("file-{i}.parquet"), + file_size: size, + row_count: 1, + min_key: vec![], + max_key: vec![], + key_stats: BinaryTableStats::new(vec![], vec![], vec![]), + value_stats: BinaryTableStats::new(vec![], vec![], vec![]), + min_sequence_number: i as i64, + max_sequence_number: i as i64, + schema_id: 0, + level: 0, + extra_files: vec![], + creation_time: None, + delete_row_count: None, + embedded_index: None, + first_row_id: None, + write_cols: None, + external_path: None, + file_source: None, + value_stats_cols: None, + column_max_sequence_numbers: None, + }) + .collect() + } + + fn make_planner(options: &[(&str, &str)]) -> AppendCompactPlanner { + let options = options + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(); + AppendCompactPlanner::new(&CoreOptions::new(&options)).unwrap() + } + + #[test] + fn java_coordinator_packing_examples() { + let options = HashMap::from([("compaction.min.file-num".into(), "3".into())]); + let core = CoreOptions::new(&options); + let target = core.target_file_size(); + let planner = AppendCompactPlanner::new(&core).unwrap(); + for (count, size, tasks) in [ + (200, 0, 2), + (100, i64::MAX, 0), + (100, target / 3 + 1, 17), + (100, target / 10 * 8, 0), + (1000, target / 10, 56), + (1050, target / 5, 105), + (1089, target / 5, 109), + ] { + let groups = planner.plan(files(&vec![size; count])); + assert_eq!(groups.len(), tasks, "count={count}, size={size}"); + } + } + + #[test] + fn strict_small_file_ratio_threshold_and_minimum_tail() { + let planner = make_planner(&[ + ("target-file-size", "100b"), + ("source.split.open-file-cost", "0b"), + ]); + assert!(planner.plan(files(&[10; 4])).is_empty()); + let groups = planner.plan(files(&[10, 10, 10, 10, 10, 70, 80, 100])); + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].len(), 5); + assert!(groups[0].iter().all(|file| file.file_size == 10)); + + let custom = make_planner(&[ + ("target-file-size", "100b"), + ("source.split.open-file-cost", "0b"), + ("compaction.small-file-ratio", "0.5"), + ("compaction.min.file-num", "2"), + ]); + assert_eq!(custom.plan(files(&[49, 49, 50, 70]))[0].len(), 2); + assert!(make_planner(&[("compaction.small-file-ratio", "0")]) + .plan(files(&[0; 5])) + .is_empty()); + } + + #[test] + fn weight_adds_open_cost_and_emits_at_twice_target_with_at_least_two_files() { + let planner = make_planner(&[ + ("target-file-size", "100b"), + ("source.split.open-file-cost", "60b"), + ]); + let groups = planner.plan(files(&[40; 5])); + assert_eq!(groups.iter().map(Vec::len).collect::>(), vec![2, 2]); + assert!(make_planner(&[ + ("target-file-size", "100b"), + ("source.split.open-file-cost", "59b"), + ]) + .plan(files(&[40; 2])) + .is_empty()); + assert!(make_planner(&[ + ("target-file-size", "100b"), + ("source.split.open-file-cost", "200b"), + ]) + .plan(files(&[40])) + .is_empty()); + } + + #[test] + fn compaction_sorts_by_size_while_read_splits_sort_by_sequence() { + let input = files(&[60, 10, 90, 40, 70, 30]); + let groups = make_planner(&[ + ("target-file-size", "100b"), + ("compaction.small-file-ratio", "1"), + ("source.split.open-file-cost", "0b"), + ]) + .plan(input.clone()); + assert_eq!(groups.len(), 1); + assert_eq!( + groups[0] + .iter() + .map(|file| file.file_size) + .collect::>(), + vec![10, 30, 40, 60, 70] + ); + let read_splits = crate::table::bin_pack::split_for_batch(input, 100, 0); + assert_eq!( + read_splits + .iter() + .map(|split| split.iter().map(|file| file.file_size).collect::>()) + .collect::>(), + vec![vec![60, 10], vec![90], vec![40], vec![70, 30]] + ); + } + + #[test] + fn task_weight_does_not_overflow_at_large_target_sizes() { + let options = HashMap::from([ + ("target-file-size".into(), i64::MAX.to_string()), + ("compaction.small-file-ratio".into(), "1".into()), + ("source.split.open-file-cost".into(), "0b".into()), + ]); + let planner = AppendCompactPlanner::new(&CoreOptions::new(&options)).unwrap(); + assert!(planner.plan(files(&[i64::MAX / 2; 3])).is_empty()); + let groups = planner.plan(files(&[i64::MAX / 2; 5])); + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].len(), 5); + } + + #[test] + fn bucketed_windows_match_java_and_do_not_pack_across_large_files() { + let planner = make_planner(&[ + ("bucket", "1"), + ("target-file-size", "100b"), + ("compaction.min.file-num", "3"), + ("source.split.open-file-cost", "1mb"), + ]); + // Open cost cannot force a two-file task; a large file shifts the + // contiguous window rather than being filtered out before packing. + assert!(planner.plan(files(&[10, 10])).is_empty()); + let groups = planner.plan(files(&[10, 300, 10, 10, 10, 10, 10, 10])); + assert_eq!( + groups + .iter() + .map(|group| group + .iter() + .map(|file| file.min_sequence_number) + .collect::>()) + .collect::>(), + vec![vec![2, 3, 4], vec![5, 6, 7]] + ); + // The minimum count check precedes the byte limit, just as in Java. + let groups = planner.plan(files(&[10, 10, 300])); + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].len(), 3); + } +} diff --git a/crates/paimon/src/table/compact/primary_key.rs b/crates/paimon/src/table/compact/primary_key.rs new file mode 100644 index 000000000..e228838a9 --- /dev/null +++ b/crates/paimon/src/table/compact/primary_key.rs @@ -0,0 +1,304 @@ +// 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. + +use std::sync::Arc; + +use arrow_schema::SchemaRef; +use futures::TryStreamExt; + +use super::CompactTask; +use crate::resource::ResourceContext; +use crate::spec::{ + BigIntType, DataField, DataFileMeta, DataType, MergeEngine, RowKind, TinyIntType, + SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME, +}; +use crate::table::data_file_index_writer::FileIndexOptions; +use crate::table::data_file_path_factory::DataFilePathFactory; +use crate::table::kv_data_file_writer::{KeyValueFileOptions, KeyValueRollingWriter}; +use crate::table::kv_file_reader::{KeyValueFileReader, KeyValueReadConfig}; +use crate::table::sort_merge::{ + compare_sequence_order, BufferedBatch, MergeFunction, MergeResult, MergeRow, +}; +use crate::table::table_read::{ + configured_mosaic_prefetch, configured_parquet_read_budget, MAX_MERGE_INPUT_STREAMS, +}; +use crate::table::{DataSplit, Table}; +use crate::Result; + +/// Keep the same winner as a deduplicate query, including DELETE/UPDATE_BEFORE. +/// A later write may carry a lower user sequence, so even full compaction +/// must retain a winning retract rather than resurrect an older value. +struct CompactDeduplicate; + +impl MergeFunction for CompactDeduplicate { + fn merge( + &self, + rows: &[MergeRow], + _batches: &[BufferedBatch], + _columns: &[usize], + _schema: &SchemaRef, + ) -> Result { + let row = rows + .iter() + .max_by(|a, b| compare_sequence_order(a, b)) + .expect("merge called with empty key group"); + RowKind::from_value(row.value_kind)?; + Ok(MergeResult::SourceRow { + batch_idx: row.batch_idx, + row_idx: row.row_idx, + }) + } +} + +/// The query reader omits singleton retracts and materialized DELETE rows. +/// Compaction persists these states so later inserts/retractions merge exactly +/// as they would against the original files. Singletons remain unreduced, as +/// in Java's reducer wrapper; compound records use the existing engine reducer. +struct CompactReducer { + inner: Box, + reject_delete_state: bool, + reject_retract_group: bool, +} + +impl MergeFunction for CompactReducer { + fn merge( + &self, + rows: &[MergeRow], + batches: &[BufferedBatch], + columns: &[usize], + schema: &SchemaRef, + ) -> Result { + if let [row] = rows { + RowKind::from_value(row.value_kind)?; + return Ok(MergeResult::SourceRow { + batch_idx: row.batch_idx, + row_idx: row.row_idx, + }); + } + if self.reject_retract_group && rows.iter().any(|row| matches!(row.value_kind, 1 | 3)) { + // Some aggregators preserve flags or raw DELETE payloads that a + // single physical INSERT/DELETE cannot encode, even if the final + // merged row is an INSERT. Leave this group in its source files. + return Err(crate::Error::Unsupported { + message: "Compaction cannot preserve aggregate state for this retract group".into(), + }); + } + Ok(match self.inner.merge(rows, batches, columns, schema)? { + MergeResult::MaterializedDeleteRow(_) if self.reject_delete_state => { + // A partial-update reducer initializes from the first retract + // and then retracts its payload. Replaying an accumulated + // DELETE can cancel its own state. First-value aggregators + // also carry an initialized flag that DELETE payloads do not + // serialize. Keep the sources until replay is safe. + return Err(crate::Error::Unsupported { + message: + "Compaction cannot preserve DELETE aggregate state for this merge engine" + .into(), + }); + } + MergeResult::MaterializedDeleteRow(batch) => MergeResult::MaterializedRow(batch), + result => result, + }) + } +} + +fn new_merge_function( + config: &KeyValueReadConfig, + fields: &[DataField], +) -> Result> { + if config.merge_engine == MergeEngine::Deduplicate { + return Ok(Box::new(CompactDeduplicate)); + } + let inner = KeyValueFileReader::new_merge_function(config, fields)?; + let aggregate_functions = match config.merge_engine { + MergeEngine::PartialUpdate => crate::spec::PartialUpdateConfig::new(&config.table_options) + .validated_aggregate_functions(&config.table_fields, &config.primary_keys)? + .into_values() + .collect::>(), + MergeEngine::Aggregation => { + let aggregation = crate::spec::AggregationConfig::new(&config.table_options); + config + .table_fields + .iter() + .filter(|field| { + !config + .table_primary_keys + .iter() + .chain(&config.sequence_fields) + .any(|name| name == field.name()) + }) + .map(|field| { + aggregation + .agg_function_for_field(field.name()) + .or_else(|| aggregation.default_agg_function()) + .unwrap_or("last_non_null_value") + .to_string() + }) + .collect() + } + _ => vec![], + }; + Ok(if config.merge_engine == MergeEngine::FirstRow { + // First-row must still reject retracts unless ignore-delete is enabled. + inner + } else { + Box::new(CompactReducer { + inner, + reject_delete_state: config.merge_engine == MergeEngine::PartialUpdate + && !aggregate_functions.is_empty(), + reject_retract_group: !(config.merge_engine == MergeEngine::PartialUpdate + && crate::spec::CoreOptions::new(&config.table_options).ignore_delete()) + && aggregate_functions.iter().any(|function| { + matches!( + crate::spec::canonical_aggregator_name(function), + "first_value" + | "first_non_null_value" + | "collect" + | "nested_update" + | "nested_partial_update" + | "listagg" + ) + }), + }) + }) +} + +pub(super) async fn execute( + table: &Table, + task: &CompactTask, + split: &DataSplit, + resources: Option, +) -> Result> { + let schema = table.schema(); + let options = schema.core_options(); + let merge_engine = options.merge_engine()?; + let mut read_type = vec![ + DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.into(), + DataType::BigInt(BigIntType::new()), + ), + DataField::new( + VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME.into(), + DataType::TinyInt(TinyIntType::new()), + ), + ]; + read_type.extend_from_slice(schema.fields()); + let budget = configured_parquet_read_budget(table)?; + let budget = match &resources { + Some(resources) => Arc::new(budget.with_resources(resources.clone()).without_prefetch()), + None => Arc::new(budget.without_prefetch()), + }; + let reader = KeyValueFileReader::new( + table.file_io().clone(), + KeyValueReadConfig { + table_name: table.identifier().full_name(), + table_options: schema.options().clone(), + schema_manager: table.schema_manager().clone(), + table_schema_id: schema.id(), + table_fields: schema.fields().to_vec(), + read_type, + predicates: vec![], + primary_keys: schema.trimmed_primary_keys(), + table_primary_keys: schema.primary_keys().to_vec(), + merge_engine, + sequence_fields: options + .sequence_fields() + .iter() + .map(|s| s.to_string()) + .collect(), + read_batch_size: options.read_batch_size()?, + merge_splits: true, + max_merge_input_streams: Some(MAX_MERGE_INPUT_STREAMS), + parquet_read_budget: Some(budget), + mosaic_prefetch: configured_mosaic_prefetch(table)?, + }, + ); + let mut batches = + reader.read_with_merge_function(std::slice::from_ref(split), new_merge_function)?; + let keys = schema.trimmed_primary_keys(); + let primary_key_indices = keys + .iter() + .map(|name| { + schema + .fields() + .iter() + .position(|field| field.name() == name) + .unwrap() + }) + .collect::>(); + let primary_key_types = primary_key_indices + .iter() + .map(|&i| schema.fields()[i].data_type().clone()) + .collect::>(); + let partition_path = if schema.partition_keys().is_empty() { + String::new() + } else { + crate::spec::PartitionComputer::new( + schema.partition_keys(), + schema.fields(), + options.partition_default_name(), + options.legacy_partition_name(), + )? + .generate_partition_path(split.partition())? + }; + let index = FileIndexOptions::parse(schema.options(), schema.fields())?; + let paths = DataFilePathFactory::new( + table.location(), + &partition_path, + task.bucket, + schema.options(), + )?; + let file_format = options.file_format(); + let mut writer = KeyValueRollingWriter::new( + KeyValueFileOptions { + file_io: table.file_io(), + paths: &paths, + table_options: schema.options(), + value_fields: schema.fields(), + primary_key_indices: &primary_key_indices, + primary_key_types: &primary_key_types, + schema_id: schema.id(), + file_format: &file_format, + file_prefix: options.data_file_prefix(), + file_compression: options.file_compression(), + file_compression_zstd_level: options.file_compression_zstd_level(), + write_buffer_size: options.write_buffer_size()? as i64, + level: task.output_level(), + file_source: 1, + is_changelog: false, + index_options: index.as_ref(), + managed_blob_enabled: false, + resources: resources.as_ref(), + }, + options.target_file_size(), + options.target_file_row_num()?, + ); + let result = async { + while let Some(batch) = batches.try_next().await? { + writer.write(&batch).await?; + } + writer.prepare_commit().await + } + .await; + if result.is_err() { + writer.abort().await; + } + result +} diff --git a/crates/paimon/src/table/kv_data_file_writer.rs b/crates/paimon/src/table/kv_data_file_writer.rs new file mode 100644 index 000000000..0943e3a7a --- /dev/null +++ b/crates/paimon/src/table/kv_data_file_writer.rs @@ -0,0 +1,717 @@ +// 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. + +//! Physical KV file output shared by write-buffer flushes and compaction. +//! Inputs are already ordered thin batches: [sequence, kind, value columns...]. + +use super::data_file_index_writer::{DataFileIndexWriter, FileIndexOptions}; +use super::data_file_path_factory::{DataFilePath, DataFilePathFactory}; +use super::managed_blob_reference::ManagedBlobReferences; +use crate::arrow::format::{ + create_format_writer_factory, with_write_resources, FormatFileWriter, FormatWriterFactory, +}; +use crate::io::FileIO; +use crate::resource::ResourceContext; +use crate::spec::stats::{BinaryTableStats, ColumnStatsCollector}; +use crate::spec::{ + data_file_to_file_index_file_name, extract_datum_from_arrow, BigIntType, BinaryRowBuilder, + CoreOptions, DataField, DataFileMeta, DataType, RowKind, TinyIntType, KEY_FIELD_ID_START, + SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME, +}; +use crate::{Error, Result}; +use arrow_array::{Array, Int64Array, Int8Array, RecordBatch, UInt64Array}; +use arrow_schema::{ + DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef, +}; +use std::collections::HashMap; +use std::sync::Arc; + +/// Immutable physical-file settings; no sequence allocator or merge-engine state. +#[derive(Clone, Copy)] +pub(super) struct KeyValueFileOptions<'a> { + pub file_io: &'a FileIO, + pub paths: &'a DataFilePathFactory, + pub table_options: &'a HashMap, + pub value_fields: &'a [DataField], + pub primary_key_indices: &'a [usize], + pub primary_key_types: &'a [DataType], + pub schema_id: i64, + pub file_format: &'a str, + pub file_prefix: &'a str, + pub file_compression: &'a str, + pub file_compression_zstd_level: i32, + pub write_buffer_size: i64, + pub level: i32, + pub file_source: i32, + pub is_changelog: bool, + pub index_options: Option<&'a FileIndexOptions>, + pub managed_blob_enabled: bool, + pub resources: Option<&'a ResourceContext>, +} + +struct RowKeyProjection { + schema: Arc, + column_indices: Vec, +} + +impl RowKeyProjection { + fn project(&self, batch: &RecordBatch) -> Result { + let columns = self + .column_indices + .iter() + .map(|&index| Arc::clone(batch.column(index))) + .collect(); + RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(|error| { + crate::Error::DataInvalid { + message: format!("Failed to prepend ROW key columns: {error}"), + source: None, + } + }) + } +} + +impl<'a> KeyValueFileOptions<'a> { + pub(super) fn new_format_writer_factory( + &self, + mut schema: arrow_schema::SchemaRef, + ) -> Result> { + schema = canonical_schema(&schema)?; + let mut physical_fields = build_physical_fields(&schema, self.value_fields)?; + let mut options = self.table_options.clone(); + let mode = CoreOptions::new(&options) + .pk_file_metadata_stats_mode(self.level, self.is_changelog)? + .to_string(); + options.insert("metadata.stats-mode".to_string(), mode); + let file_format = self.file_format; + if file_format.eq_ignore_ascii_case("row") { + let projection = self.row_key_projection(&schema)?; + let key_count = self.primary_key_indices.len(); + physical_fields = projection + .column_indices + .iter() + .enumerate() + .map(|(index, &source)| { + let field = physical_fields[source].clone(); + if index < key_count { + let id = field.id() + KEY_FIELD_ID_START; + field + .with_id(id) + .with_name(projection.schema.field(index).name().clone()) + } else { + field + } + }) + .collect(); + schema = projection.schema; + } + create_format_writer_factory( + file_format, + schema, + self.file_compression_zstd_level, + None, + Some(&physical_fields), + Some(&options), + Some(self.value_fields), + ) + } + + fn row_key_projection(&self, schema: &Arc) -> Result { + let mut fields = Vec::new(); + let mut column_indices = Vec::new(); + // These indices contain only trimmed keys; partition key columns are + // stored in the value portion but do not belong in the key prefix. + for &index in self.primary_key_indices { + let key = &self.value_fields[index]; + let source = + schema + .index_of(key.name()) + .map_err(|error| crate::Error::DataInvalid { + message: format!( + "ROW key '{}' is absent from the physical schema: {error}", + key.name() + ), + source: None, + })?; + let arrow_field = schema.field(source); + fields.push(Arc::new(ArrowField::new( + format!("_KEY_{}", key.name()), + arrow_field.data_type().clone(), + arrow_field.is_nullable(), + ))); + column_indices.push(source); + } + fields.extend(schema.fields().iter().cloned()); + column_indices.extend(0..schema.fields().len()); + Ok(RowKeyProjection { + schema: Arc::new(ArrowSchema::new_with_metadata( + fields, + schema.metadata().clone(), + )), + column_indices, + }) + } + + pub(super) async fn open( + self, + schema: &SchemaRef, + factory: &dyn FormatWriterFactory, + ordinal: usize, + ) -> Result> { + let schema = canonical_schema(schema)?; + let row_projection = self + .file_format + .eq_ignore_ascii_case("row") + .then(|| self.row_key_projection(&schema)) + .transpose()?; + let index = if self.is_changelog { + None + } else { + self.index_options + .map(FileIndexOptions::create_writer) + .transpose()? + }; + let blob_references = ManagedBlobReferences::new( + self.value_fields, + &CoreOptions::new(self.table_options), + &schema, + self.managed_blob_enabled, + self.is_changelog, + )?; + let name = format!( + "{}{}-{}.{}", + self.file_prefix, + uuid::Uuid::new_v4(), + ordinal, + self.file_format + ); + let location = self.paths.new_path(&name)?; + self.file_io + .mkdirs(&format!("{}/", location.parent())) + .await?; + let output = self.file_io.new_output(&location.path)?; + let writer = match factory.create_writer(&output, self.file_compression).await { + Ok(writer) => with_write_resources(writer, self.resources), + Err(error) => { + let _ = self.file_io.delete_file(&location.path).await; + return Err(error); + } + }; + Ok(KeyValueDataFileWriter { + options: self, + schema, + row_projection, + name, + location, + writer: Some(writer), + index, + blob_references, + key_columns: self.primary_key_indices.iter().map(|i| i + 2).collect(), + key_stats: ColumnStatsCollector::new(self.primary_key_types), + rows: 0, + min_key: vec![], + max_key: vec![], + min_sequence: i64::MAX, + max_sequence: i64::MIN, + deletes: 0, + buffered_input_bytes: 0, + needs_memory_roll: false, + }) + } +} + +/// One physical file. Flush releases format buffers; close finishes the file. +pub(super) struct KeyValueDataFileWriter<'a> { + options: KeyValueFileOptions<'a>, + schema: SchemaRef, + row_projection: Option, + name: String, + location: DataFilePath, + writer: Option>, + index: Option, + blob_references: Option, + key_columns: Vec, + key_stats: ColumnStatsCollector<'a>, + rows: usize, + min_key: Vec, + max_key: Vec, + min_sequence: i64, + max_sequence: i64, + deletes: i64, + buffered_input_bytes: usize, + needs_memory_roll: bool, +} + +impl KeyValueDataFileWriter<'_> { + pub(super) async fn write(&mut self, batch: &RecordBatch) -> Result<()> { + if batch.num_rows() == 0 { + return Ok(()); + } + // Normalize Arrow field metadata/nullability without changing column order + // or projecting physical rows back through the logical write-buffer schema. + let batch = RecordBatch::try_new(self.schema.clone(), batch.columns().to_vec()) + .map_err(|e| invalid(format!("Invalid physical KV batch: {e}")))?; + let sequences = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("KV sequence column must be Int64"))?; + let kinds = batch + .column(1) + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("KV kind column must be Int8"))?; + for row in 0..batch.num_rows() { + self.min_sequence = self.min_sequence.min(sequences.value(row)); + self.max_sequence = self.max_sequence.max(sequences.value(row)); + if RowKind::from_value(kinds.value(row))?.is_retract() { + self.deletes += 1; + } + } + if self.rows == 0 { + self.min_key = self.key(&batch, 0)?; + } + self.max_key = self.key(&batch, batch.num_rows() - 1)?; + self.key_stats.collect(&batch, &self.key_columns)?; + ManagedBlobReferences::collect(&mut self.blob_references, &batch)?; + let projected = self + .row_projection + .as_ref() + .map(|p| p.project(&batch)) + .transpose()?; + let writer = self.writer.as_mut().unwrap(); + writer.write(projected.as_ref().unwrap_or(&batch)).await?; + if let Some(index) = self.index.as_mut() { + let columns = (2..batch.num_columns()).collect::>(); + let values = batch + .project(&columns) + .map_err(|e| invalid(format!("Failed to project KV index values: {e}")))?; + index.write(&values)?; + } + self.rows += batch.num_rows(); + self.buffered_input_bytes = self + .buffered_input_bytes + .saturating_add(batch.get_array_memory_size()); + if !retains_input_rows(writer.as_ref()) { + self.buffered_input_bytes = 0; + } + self.needs_memory_roll = false; + if self.buffered_input_bytes as i64 >= self.options.write_buffer_size + || writer.in_progress_size() as i64 >= self.options.write_buffer_size + { + writer.flush().await?; + // Mosaic/Vortex may retain their pending rows after flush. Closing + // is the bounded fallback when a format cannot release its buffer. + self.needs_memory_roll = retains_input_rows(writer.as_ref()); + if !self.needs_memory_roll { + self.buffered_input_bytes = 0; + } + } + Ok(()) + } + + fn key(&self, batch: &RecordBatch, row: usize) -> Result> { + let mut builder = BinaryRowBuilder::new(self.key_columns.len() as i32); + for (pos, (&column, ty)) in self + .key_columns + .iter() + .zip(self.options.primary_key_types) + .enumerate() + { + match extract_datum_from_arrow(batch, row, column, ty)? { + Some(value) => builder.write_datum(pos, &value, ty), + None => builder.set_null_at(pos), + } + } + Ok(builder.build_serialized()) + } + + pub(super) fn rows(&self) -> usize { + self.rows + } + pub(super) fn num_bytes(&self) -> usize { + self.writer.as_ref().unwrap().num_bytes() + } + + pub(super) async fn close(mut self) -> Result { + let result = self.finish().await; + if result.is_err() { + self.abort().await; + } + result + } + + async fn finish(&mut self) -> Result { + if self.rows == 0 { + return Err(invalid("Cannot close an empty key-value file")); + } + let result = self.writer.take().unwrap().close().await?; + let (value_stats, value_stats_cols) = result + .value_stats + .map(|stats| (stats.stats, stats.columns)) + .unwrap_or_else(|| (BinaryTableStats::empty(), Some(vec![]))); + let mut meta = DataFileMeta { + file_name: self.name.clone(), + file_size: result.file_size as i64, + row_count: self.rows as i64, + min_key: std::mem::take(&mut self.min_key), + max_key: std::mem::take(&mut self.max_key), + key_stats: self.key_stats.finish(), + value_stats, + min_sequence_number: self.min_sequence, + max_sequence_number: self.max_sequence, + schema_id: self.options.schema_id, + level: self.options.level, + extra_files: vec![], + creation_time: None, + delete_row_count: Some(self.deletes), + embedded_index: None, + file_source: Some(self.options.file_source), + value_stats_cols, + external_path: self.location.external_path.clone(), + first_row_id: None, + write_cols: None, + column_max_sequence_numbers: None, + }; + if let Some(index) = self.index.take() { + let bytes = index.serialize()?; + if bytes.len() as u64 > self.options.index_options.unwrap().in_manifest_threshold as u64 + { + let name = data_file_to_file_index_file_name(&self.name); + self.options + .file_io + .new_output(&format!("{}/{}", self.location.parent(), name))? + .write(bytes) + .await?; + meta.extra_files.push(name); + } else { + meta.embedded_index = Some(bytes.to_vec()); + } + } + ManagedBlobReferences::finish( + self.blob_references.take(), + self.options.file_io, + &self.location.path, + self.location.parent(), + &mut meta, + ) + .await?; + meta.creation_time = Some(crate::common::date_time::current_local_timestamp()); + Ok(meta) + } + + pub(super) async fn abort(mut self) { + if let Some(writer) = self.writer.take() { + let _ = writer.close().await; + } + let _ = self.options.file_io.delete_file(&self.location.path).await; + if self.options.index_options.is_some() && !self.options.is_changelog { + let path = format!( + "{}/{}", + self.location.parent(), + data_file_to_file_index_file_name(&self.name) + ); + let _ = self.options.file_io.delete_file(&path).await; + } + } +} + +/// Roll an ordered KV stream independently of input batch and buffer boundaries. +pub(super) struct KeyValueRollingWriter<'a> { + options: KeyValueFileOptions<'a>, + target_size: i64, + target_rows: usize, + factory: Option>, + current: Option>, + files: Vec, +} + +impl<'a> KeyValueRollingWriter<'a> { + pub(super) fn new( + options: KeyValueFileOptions<'a>, + target_size: i64, + target_rows: i64, + ) -> Self { + Self { + options, + target_size, + target_rows: target_rows.try_into().unwrap_or(usize::MAX), + factory: None, + current: None, + files: vec![], + } + } + + pub(super) async fn write(&mut self, batch: &RecordBatch) -> Result<()> { + let mut offset = 0; + while offset < batch.num_rows() { + if self.current.is_none() { + if self.factory.is_none() { + self.factory = Some(self.options.new_format_writer_factory(batch.schema())?); + } + self.current = Some( + self.options + .open( + &batch.schema(), + self.factory.as_ref().unwrap().as_ref(), + self.files.len(), + ) + .await?, + ); + } + let file = self.current.as_mut().unwrap(); + let rows = (self.target_rows - file.rows()) + .min(batch.num_rows() - offset) + .min(4096); + if offset == 0 && rows == batch.num_rows() { + file.write(batch).await?; + } else { + // A zero-copy slice retains the whole source allocation and is + // charged as such by ResourceFormatWriter. Materialize bounded + // chunks so each row limit/format write owns only its rows. + let indices = + UInt64Array::from_iter_values((offset..offset + rows).map(|i| i as u64)); + let chunk = arrow_select::take::take_record_batch(batch, &indices) + .map_err(|e| invalid(format!("Failed to materialize KV output chunk: {e}")))?; + file.write(&chunk).await?; + } + offset += rows; + if file.rows() >= self.target_rows + || file.num_bytes() as i64 >= self.target_size + || file.needs_memory_roll + { + self.close_current().await?; + } + } + Ok(()) + } + + async fn close_current(&mut self) -> Result<()> { + if let Some(file) = self.current.take() { + self.files.push(file.close().await?); + } + Ok(()) + } + + pub(super) async fn prepare_commit(&mut self) -> Result> { + self.close_current().await?; + Ok(std::mem::take(&mut self.files)) + } + + pub(super) async fn abort(&mut self) { + if let Some(file) = self.current.take() { + file.abort().await; + } + for file in self.files.drain(..) { + for path in file.collect_files(self.options.paths.bucket_path()) { + let _ = self.options.file_io.delete_file(&path).await; + } + } + } +} + +fn canonical_schema(schema: &ArrowSchema) -> Result { + if schema.fields().len() < 2 + || schema.field(0).name() != SEQUENCE_NUMBER_FIELD_NAME + || schema.field(1).name() != VALUE_KIND_FIELD_NAME + { + return Err(invalid( + "Expected thin KV columns: sequence, kind, then values", + )); + } + let mut fields = schema.fields().to_vec(); + fields[0] = Arc::new(ArrowField::new( + SEQUENCE_NUMBER_FIELD_NAME, + ArrowDataType::Int64, + false, + )); + fields[1] = Arc::new(ArrowField::new( + VALUE_KIND_FIELD_NAME, + ArrowDataType::Int8, + false, + )); + Ok(Arc::new(ArrowSchema::new(fields))) +} + +fn invalid(message: impl Into) -> Error { + Error::DataInvalid { + message: message.into(), + source: None, + } +} + +fn retains_input_rows(writer: &dyn FormatFileWriter) -> bool { + // ROW keeps a nonempty block header after flush, even with no pending rows. + writer + .pending_rows() + .map_or_else(|| writer.retains_batch_data(), |rows| rows != 0) +} + +fn build_physical_fields( + physical_schema: &ArrowSchema, + value_fields: &[DataField], +) -> Result> { + physical_schema + .fields() + .iter() + .map(|arrow_field| match arrow_field.name().as_str() { + SEQUENCE_NUMBER_FIELD_NAME => Ok(DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.into(), + DataType::BigInt(BigIntType::new()), + )), + VALUE_KIND_FIELD_NAME => Ok(DataField::new( + VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME.into(), + DataType::TinyInt(TinyIntType::new()), + )), + name => value_fields + .iter() + .find(|field| field.name() == name) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "Physical file column '{name}' is missing from the table schema" + ), + source: None, + }) + .and_then(|field| { + // The table remains logically NOT NULL, but Java's KV + // value type makes non-key DELETE payloads nullable. + Ok(DataField::new( + field.id(), + field.name().into(), + field + .data_type() + .copy_with_nullable(arrow_field.is_nullable())?, + ) + .with_description(field.description().map(str::to_owned)) + .with_default_value(field.default_value().map(str::to_owned))) + }), + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::io::FileIOBuilder; + use crate::spec::IntType; + use arrow_array::Int32Array; + + async fn write_with_budget( + format: &str, + buffer: &str, + rows: usize, + batches: usize, + limit: usize, + expected_files: usize, + ) { + let io = FileIOBuilder::new("memory").build().unwrap(); + let path = format!("memory:/kv_stream/{}", uuid::Uuid::new_v4()); + let options = HashMap::from([ + ("write-buffer-size".into(), buffer.into()), + ("file.block-size".into(), "64mb".into()), + ("file.compression".into(), "zstd".into()), + ]); + let core = CoreOptions::new(&options); + let fields = vec![ + DataField::new(0, "id".into(), DataType::Int(IntType::new())), + DataField::new(1, "value".into(), DataType::Int(IntType::new())), + ]; + let types = vec![DataType::Int(IntType::new())]; + let paths = DataFilePathFactory::new(&path, "", 0, &options).unwrap(); + let resources = ResourceContext::builder() + .memory_limit(limit) + .build() + .unwrap(); + let mut writer = KeyValueRollingWriter::new( + KeyValueFileOptions { + file_io: &io, + paths: &paths, + table_options: &options, + value_fields: &fields, + primary_key_indices: &[0], + primary_key_types: &types, + schema_id: 0, + file_format: format, + file_prefix: "data-", + file_compression: "zstd", + file_compression_zstd_level: 0, + write_buffer_size: core.write_buffer_size().unwrap() as i64, + level: 1, + file_source: 1, + is_changelog: false, + index_options: None, + managed_blob_enabled: false, + resources: Some(&resources), + }, + 256 * 1024 * 1024, + i64::MAX, + ); + for batch in 0..batches { + let start = batch * rows; + let data = RecordBatch::try_from_iter([ + ( + SEQUENCE_NUMBER_FIELD_NAME, + Arc::new(Int64Array::from_iter_values( + (start..start + rows).map(|i| i as i64), + )) as arrow_array::ArrayRef, + ), + ( + VALUE_KIND_FIELD_NAME, + Arc::new(Int8Array::from(vec![0; rows])) as arrow_array::ArrayRef, + ), + ( + "id", + Arc::new(Int32Array::from_iter_values( + (start..start + rows).map(|i| i as i32), + )) as arrow_array::ArrayRef, + ), + ( + "value", + Arc::new(Int32Array::from(vec![10; rows])) as arrow_array::ArrayRef, + ), + ]) + .unwrap(); + writer.write(&data).await.unwrap(); + } + let files = writer.prepare_commit().await.unwrap(); + assert_eq!(files.len(), expected_files, "{format}"); + assert_eq!( + files.iter().map(|file| file.row_count).sum::(), + (rows * batches) as i64 + ); + assert_eq!(resources.metrics().reserved_memory_bytes, 0); + assert!(resources.metrics().peak_reserved_memory_bytes <= limit); + } + + #[tokio::test] + async fn large_input_batches_charge_only_materialized_output_chunks() { + write_with_budget("parquet", "64mb", 32768, 1, 1024 * 1024, 1).await; + } + + #[tokio::test] + async fn formats_that_retain_rows_after_flush_keep_bounded_output_files() { + #[cfg(feature = "vortex")] + let formats = ["mosaic", "vortex"]; + #[cfg(not(feature = "vortex"))] + let formats = ["mosaic"]; + for format in formats { + write_with_budget(format, "8kb", 1024, 8, 64 * 1024, 8).await; + } + } +} diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index d14b9f120..e97d3a1bf 100644 --- a/crates/paimon/src/table/kv_file_reader.rs +++ b/crates/paimon/src/table/kv_file_reader.rs @@ -368,7 +368,7 @@ impl KeyValueFileReader { self } - fn new_merge_function( + pub(super) fn new_merge_function( config: &KeyValueReadConfig, merge_output_fields: &[DataField], ) -> crate::Result> { diff --git a/crates/paimon/src/table/kv_file_writer.rs b/crates/paimon/src/table/kv_file_writer.rs index 35bcb8e72..f8a566d3c 100644 --- a/crates/paimon/src/table/kv_file_writer.rs +++ b/crates/paimon/src/table/kv_file_writer.rs @@ -26,21 +26,16 @@ //! //! Reference: [org.apache.paimon.io.KeyValueDataFileWriterImpl](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/io/KeyValueDataFileWriterImpl.java) +use super::kv_data_file_writer::KeyValueFileOptions; use crate::arrow::arrow_fields_to_paimon; -use crate::arrow::format::{ - create_format_writer_factory, with_write_resources, FormatWriterFactory, -}; +use crate::arrow::format::FormatWriterFactory; use crate::io::FileIO; use crate::resource::{MemoryReservation, ResourceContext}; -use crate::spec::stats::{compute_column_stats, BinaryTableStats}; use crate::spec::{ - data_file_to_file_index_file_name, extract_datum_from_arrow, AggregationConfig, BigIntType, - BinaryRowBuilder, CoreOptions, DataField, DataFileMeta, DataType, MergeEngine, - PartialUpdateConfig, RowKind, TinyIntType, KEY_FIELD_ID_START, SEQUENCE_NUMBER_FIELD_ID, - SEQUENCE_NUMBER_FIELD_NAME, VALUE_KIND_FIELD_ID, VALUE_KIND_FIELD_NAME, + AggregationConfig, CoreOptions, DataField, DataFileMeta, DataType, MergeEngine, + PartialUpdateConfig, RowKind, SEQUENCE_NUMBER_FIELD_NAME, VALUE_KIND_FIELD_NAME, }; use crate::table::data_file_index_writer::FileIndexOptions; -use crate::table::managed_blob_reference::ManagedBlobReferences; use crate::table::managed_blob_writer::ManagedBlobWriteState; use crate::table::prepared_files::PreparedFiles; use crate::table::sort_merge::{ @@ -116,34 +111,7 @@ pub(crate) struct KeyValueWriteConfig { struct IndexedFileWrite<'a> { format_writer_factory: &'a dyn FormatWriterFactory, is_changelog: bool, - file_prefix: &'a str, file_ordinal: usize, - file_format: &'a str, - file_compression: &'a str, - min_sequence_number: i64, - max_sequence_number: i64, - delete_row_count: i64, -} - -struct RowKeyProjection { - schema: Arc, - column_indices: Vec, -} - -impl RowKeyProjection { - fn project(&self, batch: &RecordBatch) -> Result { - let columns = self - .column_indices - .iter() - .map(|&index| Arc::clone(batch.column(index))) - .collect(); - RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(|error| { - crate::Error::DataInvalid { - message: format!("Failed to prepend ROW key columns: {error}"), - source: None, - } - }) - } } impl KeyValueFileWriter { @@ -210,90 +178,48 @@ impl KeyValueFileWriter { self } - /// Construct one reusable factory for this flush's data or changelog output. - /// Statistics retain logical value fields independently of the physical KV schema. + fn file_options(&self, is_changelog: bool) -> KeyValueFileOptions<'_> { + KeyValueFileOptions { + file_io: &self.file_io, + paths: &self.paths, + table_options: &self.config.table_options, + value_fields: &self.config.value_fields, + primary_key_indices: &self.config.primary_key_indices, + primary_key_types: &self.config.primary_key_types, + schema_id: self.config.schema_id, + file_format: if is_changelog { + &self.config.changelog_file_format + } else { + &self.config.file_format + }, + file_prefix: if is_changelog { + &self.config.changelog_file_prefix + } else { + &self.config.data_file_prefix + }, + file_compression: if is_changelog { + &self.config.changelog_file_compression + } else { + &self.config.file_compression + }, + file_compression_zstd_level: self.config.file_compression_zstd_level, + write_buffer_size: self.config.write_buffer_size, + level: 0, + file_source: 0, + is_changelog, + index_options: self.config.file_index_options.as_deref(), + managed_blob_enabled: self.managed_blob_writer.enabled(), + resources: self.resources.as_ref(), + } + } + fn new_format_writer_factory( &self, - mut schema: arrow_schema::SchemaRef, + schema: arrow_schema::SchemaRef, is_changelog: bool, ) -> Result> { - let mut physical_fields = build_physical_fields(&schema, &self.config.value_fields)?; - let mut options = self.config.table_options.clone(); - let mode = CoreOptions::new(&options) - .pk_file_metadata_stats_mode(0, is_changelog)? - .to_string(); - options.insert("metadata.stats-mode".to_string(), mode); - let file_format = if is_changelog { - &self.config.changelog_file_format - } else { - &self.config.file_format - }; - if file_format.eq_ignore_ascii_case("row") { - let projection = self.row_key_projection(&schema)?; - let key_count = self.config.primary_key_indices.len(); - physical_fields = projection - .column_indices - .iter() - .enumerate() - .map(|(index, &source)| { - let field = physical_fields[source].clone(); - if index < key_count { - let id = field.id() + KEY_FIELD_ID_START; - field - .with_id(id) - .with_name(projection.schema.field(index).name().clone()) - } else { - field - } - }) - .collect(); - schema = projection.schema; - } - create_format_writer_factory( - file_format, - schema, - self.config.file_compression_zstd_level, - None, - Some(&physical_fields), - Some(&options), - Some(&self.config.value_fields), - ) - } - - fn row_key_projection(&self, schema: &Arc) -> Result { - let mut fields = Vec::new(); - let mut column_indices = Vec::new(); - // These indices contain only trimmed keys; partition key columns are - // stored in the value portion but do not belong in the key prefix. - for &index in &self.config.primary_key_indices { - let key = &self.config.value_fields[index]; - let source = - schema - .index_of(key.name()) - .map_err(|error| crate::Error::DataInvalid { - message: format!( - "ROW key '{}' is absent from the physical schema: {error}", - key.name() - ), - source: None, - })?; - let arrow_field = schema.field(source); - fields.push(Arc::new(ArrowField::new( - format!("_KEY_{}", key.name()), - arrow_field.data_type().clone(), - arrow_field.is_nullable(), - ))); - column_indices.push(source); - } - fields.extend(schema.fields().iter().cloned()); - column_indices.extend(0..schema.fields().len()); - Ok(RowKeyProjection { - schema: Arc::new(ArrowSchema::new_with_metadata( - fields, - schema.metadata().clone(), - )), - column_indices, - }) + self.file_options(is_changelog) + .new_format_writer_factory(schema) } pub(crate) fn with_resources(mut self, resources: Option) -> Self { @@ -450,17 +376,9 @@ impl KeyValueFileWriter { // buffer flush. Rolled files share history only within that flush. let data_factory = self.new_format_writer_factory(build_physical_schema(&data_batch.schema()), false)?; - let data_sequences = data_seq.as_any().downcast_ref::().unwrap(); for offset in (0..data_indices.len()).step_by(self.target_file_row_num) { let len = self.target_file_row_num.min(data_indices.len() - offset); let file_indices = data_indices.slice(offset, len); - let (min_sequence_number, max_sequence_number) = file_indices - .values() - .iter() - .map(|&idx| data_sequences.value(idx as usize)) - .fold((i64::MAX, i64::MIN), |(min, max), seq| { - (min.min(seq), max.max(seq)) - }); let file = self .write_indexed_file( &data_batch, @@ -469,16 +387,7 @@ impl KeyValueFileWriter { IndexedFileWrite { format_writer_factory: data_factory.as_ref(), is_changelog: false, - file_prefix: &self.config.data_file_prefix, file_ordinal: self.written_files.len(), - file_format: &self.config.file_format, - file_compression: &self.config.file_compression, - min_sequence_number, - max_sequence_number, - delete_row_count: Self::indexed_delete_row_count( - &data_batch, - &file_indices, - )?, }, ) .await?; @@ -488,17 +397,9 @@ impl KeyValueFileWriter { if self.config.input_changelog { let changelog_factory = self.new_format_writer_factory(build_physical_schema(&combined.schema()), true)?; - let input_sequences = seq_array.as_any().downcast_ref::().unwrap(); for offset in (0..sorted_indices.len()).step_by(self.target_file_row_num) { let len = self.target_file_row_num.min(sorted_indices.len() - offset); let file_indices = sorted_indices.slice(offset, len); - let (min_sequence_number, max_sequence_number) = file_indices - .values() - .iter() - .map(|&idx| input_sequences.value(idx as usize)) - .fold((i64::MAX, i64::MIN), |(min, max), seq| { - (min.min(seq), max.max(seq)) - }); let file = self .write_indexed_file( &combined, @@ -507,16 +408,7 @@ impl KeyValueFileWriter { IndexedFileWrite { format_writer_factory: changelog_factory.as_ref(), is_changelog: true, - file_prefix: &self.config.changelog_file_prefix, file_ordinal: self.written_changelog_files.len(), - file_format: &self.config.changelog_file_format, - file_compression: &self.config.changelog_file_compression, - min_sequence_number, - max_sequence_number, - delete_row_count: Self::indexed_delete_row_count( - &combined, - &file_indices, - )?, }, ) .await?; @@ -564,6 +456,8 @@ impl KeyValueFileWriter { }) } + /// Materialize sorted selections into thin KV chunks. Physical file metadata, + /// format encoding, indexes and cleanup are owned by KeyValueDataFileWriter. async fn write_indexed_file( &self, batch: &RecordBatch, @@ -573,298 +467,89 @@ impl KeyValueFileWriter { ) -> Result { if indices.is_empty() { return Err(crate::Error::DataInvalid { - message: "Cannot write an empty key-value data file".to_string(), + message: "Cannot write an empty key-value data file".into(), source: None, }); } - - let user_schema = batch.schema(); - let first_row = indices.value(0) as usize; - let last_row = indices.value(indices.len() - 1) as usize; - let min_key = self.extract_key_binary_row(batch, first_row)?; - let max_key = self.extract_key_binary_row(batch, last_row)?; - let mut file_index = if write.is_changelog { - None - } else { - self.config - .file_index_options - .as_ref() - .map(|options| options.create_writer()) - .transpose()? - }; - let physical_schema = build_physical_schema(&user_schema); - let row_projection = write - .file_format - .eq_ignore_ascii_case("row") - .then(|| self.row_key_projection(&physical_schema)) - .transpose()?; - let mut blob_references = ManagedBlobReferences::new( - &self.config.value_fields, - &CoreOptions::new(&self.config.table_options), - &physical_schema, - self.managed_blob_writer.enabled(), - write.is_changelog, - )?; - - let file_name = format!( - "{}{}-{}.{}", - write.file_prefix, - uuid::Uuid::new_v4(), - write.file_ordinal, - write.file_format, - ); - let location = self.paths.new_path(&file_name)?; - let bucket_dir = location.parent(); - self.file_io.mkdirs(&format!("{bucket_dir}/")).await?; - let file_path = location.path.clone(); - let output = self.file_io.new_output(&file_path)?; - let writer = write - .format_writer_factory - .create_writer(&output, write.file_compression) + let physical_schema = build_physical_schema(&batch.schema()); + let mut file = self + .file_options(write.is_changelog) + .open( + &physical_schema, + write.format_writer_factory, + write.file_ordinal, + ) .await?; - let mut writer = with_write_resources(writer, self.resources.as_ref()); + let result = async { + let vk_idx = batch + .schema() + .fields() + .iter() + .position(|f| f.name() == crate::spec::VALUE_KIND_FIELD_NAME); - let vk_idx = batch - .schema() - .fields() - .iter() - .position(|f| f.name() == crate::spec::VALUE_KIND_FIELD_NAME); + for chunk_start in (0..indices.len()).step_by(Self::FLUSH_CHUNK_ROWS) { + let chunk_len = Self::FLUSH_CHUNK_ROWS.min(indices.len() - chunk_start); + let chunk_indices = indices.slice(chunk_start, chunk_len); - for chunk_start in (0..indices.len()).step_by(Self::FLUSH_CHUNK_ROWS) { - let chunk_len = Self::FLUSH_CHUNK_ROWS.min(indices.len() - chunk_start); - let chunk_indices = indices.slice(chunk_start, chunk_len); + let mut physical_columns: Vec> = Vec::new(); + physical_columns.push( + arrow_select::take::take(seq_array, &chunk_indices, None).map_err(|e| { + crate::Error::DataInvalid { + message: format!("Failed to reorder sequence numbers: {e}"), + source: None, + } + })?, + ); - let mut physical_columns: Vec> = Vec::new(); - physical_columns.push( - arrow_select::take::take(seq_array, &chunk_indices, None).map_err(|e| { - crate::Error::DataInvalid { - message: format!("Failed to reorder sequence numbers: {e}"), - source: None, + match vk_idx { + Some(vk_idx) => { + physical_columns.push( + arrow_select::take::take( + batch.column(vk_idx).as_ref(), + &chunk_indices, + None, + ) + .map_err(|e| { + crate::Error::DataInvalid { + message: format!("Failed to reorder value kind column: {e}"), + source: None, + } + })?, + ); } - })?, - ); + None => { + physical_columns.push(Arc::new(Int8Array::from(vec![0i8; chunk_len]))); + } + } - match vk_idx { - Some(vk_idx) => { + for idx in 0..batch.num_columns() { + if Some(idx) == vk_idx { + continue; + } physical_columns.push( - arrow_select::take::take( - batch.column(vk_idx).as_ref(), - &chunk_indices, - None, - ) - .map_err(|e| crate::Error::DataInvalid { - message: format!("Failed to reorder value kind column: {e}"), - source: None, - })?, + arrow_select::take::take(batch.column(idx).as_ref(), &chunk_indices, None) + .map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to reorder by sort indices: {e}"), + source: None, + })?, ); } - None => { - physical_columns.push(Arc::new(Int8Array::from(vec![0i8; chunk_len]))); - } - } - - for idx in 0..batch.num_columns() { - if Some(idx) == vk_idx { - continue; - } - physical_columns.push( - arrow_select::take::take(batch.column(idx).as_ref(), &chunk_indices, None) - .map_err(|e| crate::Error::DataInvalid { - message: format!("Failed to reorder by sort indices: {e}"), - source: None, - })?, - ); - } - let chunk_batch = RecordBatch::try_new(physical_schema.clone(), physical_columns) - .map_err(|e| crate::Error::DataInvalid { + let chunk_batch = RecordBatch::try_new(physical_schema.clone(), physical_columns) + .map_err(|e| crate::Error::DataInvalid { message: format!("Failed to create physical batch: {e}"), source: None, })?; - if let Err(error) = ManagedBlobReferences::collect(&mut blob_references, &chunk_batch) { - let _ = writer.close().await; - let _ = self.file_io.delete_file(&file_path).await; - return Err(error); - } - let row_batch = row_projection - .as_ref() - .map(|projection| projection.project(&chunk_batch)) - .transpose()?; - if let Err(error) = writer - .write(row_batch.as_ref().unwrap_or(&chunk_batch)) - .await - { - let _ = writer.close().await; - let _ = self.file_io.delete_file(&file_path).await; - return Err(error); - } - if let Some(index) = file_index.as_mut() { - // The index positions must match the physical file after PK - // sorting and flush-time merging. Its field positions refer - // to the logical value schema, so omit the two KV metadata - // columns from the same output chunk used by the file writer. - let logical_indices = (2..chunk_batch.num_columns()).collect::>(); - let index_result = chunk_batch - .project(&logical_indices) - .map_err(|error| crate::Error::DataInvalid { - message: format!("Failed to project KV index values: {error}"), - source: None, - }) - .and_then(|logical_batch| index.write(&logical_batch)); - if let Err(error) = index_result { - let _ = writer.close().await; - let _ = self.file_io.delete_file(&file_path).await; - return Err(error); - } - } - } - - let write_result = writer.close().await?; - let file_size = write_result.file_size as i64; - let (value_stats, value_stats_cols) = match write_result.value_stats { - Some(stats) => (stats.stats, stats.columns), - None => (BinaryTableStats::empty(), Some(Vec::new())), - }; - - let key_columns: Vec> = self - .config - .primary_key_indices - .iter() - .map(|&idx| { - arrow_select::take::take(batch.column(idx).as_ref(), indices, None).map_err(|e| { - crate::Error::DataInvalid { - message: format!("Failed to take key column for stats: {e}"), - source: None, - } - }) - }) - .collect::>>()?; - let key_batch = RecordBatch::try_new( - Arc::new(ArrowSchema::new( - self.config - .primary_key_indices - .iter() - .map(|&idx| user_schema.field(idx).clone()) - .collect::>(), - )), - key_columns, - ) - .map_err(|e| crate::Error::DataInvalid { - message: format!("Failed to build key batch for stats: {e}"), - source: None, - })?; - let stats_col_indices: Vec = (0..self.config.primary_key_indices.len()).collect(); - let key_stats = compute_column_stats( - &key_batch, - &stats_col_indices, - &self.config.primary_key_types, - )?; - - let mut meta = DataFileMeta { - file_name, - file_size, - row_count: indices.len() as i64, - min_key, - max_key, - key_stats, - value_stats, - min_sequence_number: write.min_sequence_number, - max_sequence_number: write.max_sequence_number, - schema_id: self.config.schema_id, - level: 0, - extra_files: vec![], - creation_time: None, - delete_row_count: Some(write.delete_row_count), - embedded_index: None, - file_source: Some(0), // FileSource.APPEND - value_stats_cols, - external_path: location.external_path.clone(), - first_row_id: None, - write_cols: None, - column_max_sequence_numbers: None, - }; - if let Some(index) = file_index { - let index_result = index.serialize(); - let bytes = match index_result { - Ok(bytes) => bytes, - Err(error) => { - let _ = self.file_io.delete_file(&file_path).await; - return Err(error); - } - }; - let threshold = self - .config - .file_index_options - .as_ref() - .expect("file index writer must have options") - .in_manifest_threshold; - if bytes.len() as u64 > threshold as u64 { - let name = data_file_to_file_index_file_name(&meta.file_name); - let index_path = format!("{bucket_dir}/{name}"); - let output = match self.file_io.new_output(&index_path) { - Ok(output) => output, - Err(error) => { - let _ = self.file_io.delete_file(&file_path).await; - return Err(error); - } - }; - if let Err(error) = output.write(bytes).await { - let _ = self.file_io.delete_file(&index_path).await; - let _ = self.file_io.delete_file(&file_path).await; - return Err(error); - } - meta.extra_files.push(name); - } else { - meta.embedded_index = Some(bytes.to_vec()); + file.write(&chunk_batch).await?; } + Ok(()) } - - ManagedBlobReferences::finish( - blob_references, - &self.file_io, - &file_path, - bucket_dir, - &mut meta, - ) - .await?; - - // Java creates file metadata after the index and Blob references finish. - meta.creation_time = Some(crate::common::date_time::current_local_timestamp()); - Ok(meta) - } - - fn indexed_delete_row_count(batch: &RecordBatch, indices: &UInt32Array) -> Result { - let Some(vk_idx) = batch - .schema() - .fields() - .iter() - .position(|f| f.name() == crate::spec::VALUE_KIND_FIELD_NAME) - else { - return Ok(0); - }; - - let column = batch.column(vk_idx); - let Some(value_kinds) = column.as_any().downcast_ref::() else { - return Err(crate::Error::DataInvalid { - message: "_VALUE_KIND column must be Int8".to_string(), - source: None, - }); - }; - - let mut delete_count = 0; - for idx in 0..indices.len() { - let row = indices.value(idx) as usize; - let value = if column.is_null(row) { - 0 - } else { - value_kinds.value(row) - }; - match RowKind::from_value(value)? { - RowKind::UpdateBefore | RowKind::Delete => delete_count += 1, - RowKind::Insert | RowKind::UpdateAfter => {} - } + .await; + if let Err(error) = result { + file.abort().await; + return Err(error); } - Ok(delete_count) + file.close().await } /// Select output row indices from sorted inputs according to merge engine. @@ -1292,25 +977,6 @@ impl KeyValueFileWriter { changelog_files: std::mem::take(&mut self.written_changelog_files), }) } - - /// Extract primary key columns from a batch at a given row index into a serialized BinaryRow. - fn extract_key_binary_row(&self, batch: &RecordBatch, row_idx: usize) -> Result> { - let num_keys = self.config.primary_key_indices.len(); - let mut builder = BinaryRowBuilder::new(num_keys as i32); - for (pos, (&col_idx, data_type)) in self - .config - .primary_key_indices - .iter() - .zip(self.config.primary_key_types.iter()) - .enumerate() - { - match extract_datum_from_arrow(batch, row_idx, col_idx, data_type)? { - Some(datum) => builder.write_datum(pos, &datum, data_type), - None => builder.set_null_at(pos), - } - } - Ok(builder.build_serialized()) - } } /// Build the physical schema: [_SEQUENCE_NUMBER, _VALUE_KIND, user_cols (excluding _VALUE_KIND)...] @@ -1334,57 +1000,14 @@ pub(crate) fn build_physical_schema(user_schema: &ArrowSchema) -> Arc Result> { - physical_schema - .fields() - .iter() - .map(|arrow_field| match arrow_field.name().as_str() { - SEQUENCE_NUMBER_FIELD_NAME => Ok(DataField::new( - SEQUENCE_NUMBER_FIELD_ID, - SEQUENCE_NUMBER_FIELD_NAME.into(), - DataType::BigInt(BigIntType::new()), - )), - VALUE_KIND_FIELD_NAME => Ok(DataField::new( - VALUE_KIND_FIELD_ID, - VALUE_KIND_FIELD_NAME.into(), - DataType::TinyInt(TinyIntType::new()), - )), - name => value_fields - .iter() - .find(|field| field.name() == name) - .ok_or_else(|| crate::Error::DataInvalid { - message: format!( - "Physical file column '{name}' is missing from the table schema" - ), - source: None, - }) - .and_then(|field| { - // The table remains logically NOT NULL, but Java's KV - // value type makes non-key DELETE payloads nullable. - Ok(DataField::new( - field.id(), - field.name().into(), - field - .data_type() - .copy_with_nullable(arrow_field.is_nullable())?, - ) - .with_description(field.description().map(str::to_owned)) - .with_default_value(field.default_value().map(str::to_owned))) - }), - }) - .collect() -} - #[cfg(test)] mod tests { use super::*; use crate::io::FileIOBuilder; - use crate::spec::IntType; + use crate::spec::{ + BigIntType, IntType, TinyIntType, KEY_FIELD_ID_START, SEQUENCE_NUMBER_FIELD_ID, + VALUE_KIND_FIELD_ID, + }; use arrow_array::{Int32Array, RecordBatchReader, StringArray, UInt32Array}; use std::collections::HashMap; @@ -2749,30 +2372,30 @@ mod tests { ); } - #[test] - fn test_indexed_delete_row_count_rejects_invalid_value_kind() { - let schema = Arc::new(ArrowSchema::new(vec![ - Arc::new(ArrowField::new("id", ArrowDataType::Int32, false)), - Arc::new(ArrowField::new( + #[tokio::test] + async fn test_flush_rejects_invalid_value_kind() { + let batch = RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(vec![1])) as Arc), + ( + "value", + Arc::new(Int32Array::from(vec![10])) as Arc, + ), + ( VALUE_KIND_FIELD_NAME, - ArrowDataType::Int8, - false, - )), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(vec![1])) as Arc, - Arc::new(Int8Array::from(vec![4])) as Arc, - ], + Arc::new(Int8Array::from(vec![4])) as Arc, + ), + ]) + .unwrap(); + let mut writer = KeyValueFileWriter::new( + FileIOBuilder::new("memory").build().unwrap(), + test_write_config(MergeEngine::Deduplicate), + 0, ) .unwrap(); - let indices = UInt32Array::from(vec![0]); - - let err = KeyValueFileWriter::indexed_delete_row_count(&batch, &indices).unwrap_err(); - + writer.write(&batch).await.unwrap(); + let error = writer.flush().await.unwrap_err(); assert!( - matches!(err, crate::Error::DataInvalid { message, .. } if message.contains("Invalid RowKind value")) + matches!(error, crate::Error::DataInvalid { message, .. } if message.contains("Invalid RowKind value")) ); } diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 7b8985d92..8e64944e4 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -44,6 +44,7 @@ mod bucket_router; mod chunk_shuffle; mod commit_callback; mod commit_message; +mod compact; mod consumer_manager; pub(crate) mod cow_writer; mod data_evolution_fields; @@ -91,6 +92,7 @@ mod incremental_scan; pub(crate) mod index_file_path; mod inline_blob; mod instant; +mod kv_data_file_writer; mod kv_file_reader; mod kv_file_writer; mod managed_blob_reader; @@ -204,6 +206,10 @@ pub use branch_manager::BranchManager; pub use bucket_router::BucketRouter; pub use commit_callback::{CommitCallback, CommitCallbackContext}; pub use commit_message::{CommitMessage, COMMIT_MESSAGE_SERIALIZER_VERSION}; +pub use compact::{ + CompactBuilder, CompactExecutor, CompactMetrics, CompactPlan, CompactPlanner, CompactResult, + CompactTask, +}; pub use consumer_manager::ConsumerManager; pub use cow_writer::{CopyOnWriteMergeWriter, FileInfo}; pub use data_evolution_writer::{DataEvolutionDeleteWriter, DataEvolutionWriter}; @@ -701,6 +707,11 @@ impl Table { WriteBuilder::new(self) } + /// Plan and execute append or fixed/dynamic bucket PK compaction. + pub fn new_compact_builder(&self) -> CompactBuilder<'_> { + CompactBuilder::new(self) + } + /// Create a one-shot fixed-bucket builder for a postpone table. pub fn new_postpone_fixed_bucket_write_builder( &self, diff --git a/crates/paimon/src/table/table_commit.rs b/crates/paimon/src/table/table_commit.rs index 1df59e0fc..3592b191d 100644 --- a/crates/paimon/src/table/table_commit.rs +++ b/crates/paimon/src/table/table_commit.rs @@ -418,6 +418,39 @@ impl TableCommit { .await } + /// Publish prepared compaction results in one COMPACT snapshot. + /// Results must belong to one plan; independent subsets may be submitted separately. + /// Retrying the same tasks with this commit user is idempotent, regardless of order. + /// Never abort prepared files merely because publication returned an error. + pub async fn commit_compact( + &self, + results: Vec, + ) -> Result { + super::compact::ensure_supported(&self.table)?; + if results.is_empty() { + return Ok(super::CompactMetrics::default()); + } + let prepared = super::compact::prepare_results(&self.table, &results)?; + let mut commit = self.clone(); + // One stable identity per plan and task subset, allowing repeated plans + // and idempotent retry without consuming the caller's checkpoint sequence. + commit.commit_user = format!("{}-compact-{}", self.commit_user, prepared.commit_suffix); + commit + .try_commit( + CommitEntriesPlan::Compact { + entries: prepared.entries, + schema_id: prepared.schema_id, + read_snapshot: prepared.read_snapshot, + planned_snapshot_id: prepared.snapshot_id, + }, + None, + BATCH_COMMIT_IDENTIFIER, + true, + ) + .await?; + Ok(prepared.metrics) + } + /// Commit new files with a caller-provided commit identifier. /// /// Identifiers must increase monotonically for a given `commit_user`. @@ -1095,8 +1128,19 @@ impl TableCommit { loop { let latest_snapshot = self.snapshot_manager.get_latest_snapshot().await?; if filter_committed { + let recovery_start = match &plan { + CommitEntriesPlan::Compact { + planned_snapshot_id, + .. + } => Some(planned_snapshot_id.saturating_add(1)), + _ => None, + }; if self - .is_committed_identifier(&latest_snapshot, commit_identifier) + .is_committed_identifier_since( + &latest_snapshot, + commit_identifier, + recovery_start, + ) .await? { break; @@ -1205,6 +1249,9 @@ impl TableCommit { new_index_entries, .. } => (new_entries.as_slice(), &[][..], new_index_entries), + CommitEntriesPlan::Compact { entries, .. } => { + return self.check_recovery_entries(entries, &[], &[]).await; + } }; self.check_recovery_entries(entries, changelog_entries, index_entries) .await @@ -1752,10 +1799,23 @@ impl TableCommit { &self, latest_snapshot: &Option, commit_identifier: i64, + ) -> Result { + self.is_committed_identifier_since(latest_snapshot, commit_identifier, None) + .await + } + + async fn is_committed_identifier_since( + &self, + latest_snapshot: &Option, + commit_identifier: i64, + start: Option, ) -> Result { let Some(latest) = latest_snapshot else { return Ok(false); }; + if start.is_some_and(|start| start > latest.id()) { + return Ok(false); + } if latest.commit_user() == self.commit_user { return Ok(commit_identifier <= latest.commit_identifier()); } @@ -1763,7 +1823,8 @@ impl TableCommit { .snapshot_manager .earliest_snapshot_id() .await? - .unwrap_or(latest.id()); + .unwrap_or(latest.id()) + .max(start.unwrap_or(i64::MIN)); for snapshot_id in (earliest_snapshot_id..=latest.id()).rev() { let snapshot = if snapshot_id == latest.id() { latest.clone() @@ -1815,6 +1876,129 @@ impl TableCommit { let manifest_dir = self.snapshot_manager.manifest_dir(); match plan { + CommitEntriesPlan::Compact { + entries, + schema_id, + read_snapshot, + planned_snapshot_id, + } => { + if latest_snapshot + .as_ref() + .is_none_or(|latest| latest.id() < *planned_snapshot_id) + { + return Err(crate::Error::DataInvalid { + message: "Table rolled back after compaction was planned".into(), + source: None, + }); + } + if self.latest_schema_id().await? != *schema_id { + return Err(crate::Error::DataInvalid { + message: "Table schema changed after compaction was planned".into(), + source: None, + }); + } + // Concurrent writes may replace the HASH index after planning. + // Reuse the latest manifest on every retry, never the captured one. + let index_manifest_name = match latest_snapshot { + Some(snapshot) => { + super::compact::validate_snapshot_indexes(&self.table, snapshot).await? + } + None => None, + }; + validate_file_entries(entries.iter())?; + // This runs after duplicate-commit filtering: an uncertain commit + // can be retried even after the captured read snapshot expires. + if let Some(captured) = read_snapshot { + let snapshot = self.snapshot_manager.get_snapshot(captured.id()).await?; + if !captured.has_same_read_metadata(&snapshot) { + return Err(crate::Error::DataInvalid { + message: "Compaction snapshot was replaced".into(), + source: None, + }); + } + let captured_files = super::snapshot_manifest::read_snapshot_entries( + &self.snapshot_manager, + &snapshot, + None, + ) + .await?; + let deletes = entries + .iter() + .filter(|e| *e.kind() == FileKind::Delete) + .collect::>(); + let buckets = deletes + .iter() + .map(|e| (e.partition(), e.bucket())) + .collect::>(); + let captured_files = captured_files + .iter() + .filter(|e| buckets.contains(&(e.partition(), e.bucket()))) + .map(|entry| (commit_file_identifier(entry), entry)) + .collect::>(); + if captured_files.len() != deletes.len() + || deletes.iter().any(|deleted| { + !captured_files + .get(&commit_file_identifier(deleted)) + .is_some_and(|file| { + file.file() == deleted.file() + && file.total_buckets() == deleted.total_buckets() + }) + }) + { + return Err(crate::Error::DataInvalid { + message: "PK compaction must include the complete captured bucket" + .into(), + source: None, + }); + } + } + let base_data_files = self + .detect_commit_conflicts( + latest_snapshot, + retry_state, + entries, + &CommitKind::COMPACT, + None, + ) + .await?; + // File identity alone does not detect metadata-only changes. + let active_files = base_data_files + .as_deref() + .unwrap_or_default() + .iter() + .map(|entry| (commit_file_identifier(entry), entry)) + .collect::>(); + for deleted in entries.iter().filter(|e| *e.kind() == FileKind::Delete) { + if !active_files + .get(&commit_file_identifier(deleted)) + .is_some_and(|entry| { + entry.file() == deleted.file() + && entry.total_buckets() == deleted.total_buckets() + }) + { + return Err(crate::Error::DataInvalid { + message: format!( + "Compaction input '{}' metadata changed", + deleted.file().file_name + ), + source: None, + }); + } + } + self.check_compact_lsm_conflicts( + base_data_files.as_deref().unwrap_or_default(), + entries, + )?; + Ok(ResolvedCommit { + entries: entries.clone(), + changelog_entries: vec![], + kind: CommitKind::COMPACT, + index_manifest_name, + index_manifest_changed: false, + new_index_manifest_entries: None, + base_data_files, + }) + } CommitEntriesPlan::Direct { entries, changelog_entries, @@ -1905,7 +2089,9 @@ impl TableCommit { new_index_entries.clone(), *check_from_snapshot, ), - CommitEntriesPlan::Direct { .. } => unreachable!(), + CommitEntriesPlan::Direct { .. } | CommitEntriesPlan::Compact { .. } => { + unreachable!() + } }; let base_data_files = self .detect_commit_conflicts( @@ -2523,6 +2709,54 @@ impl TableCommit { Ok(Some(base_data_files)) } + /// Java PrimaryKeyConflictDetection: each positive level is a sorted run + /// with strictly disjoint primary-key ranges. Concurrent L0 appends remain + /// valid, but a concurrent rewrite into our output level may conflict. + fn check_compact_lsm_conflicts( + &self, + base: &[ManifestEntry], + delta: &[ManifestEntry], + ) -> Result<()> { + if self.table.schema().primary_keys().is_empty() { + return Ok(()); + } + let comparator = super::merge_tree_split_generator::KeyComparator::from_table_schema( + self.table.schema(), + ); + let active = merge_active_entries(base.iter().chain(delta).cloned().collect()); + let mut levels = HashMap::<_, Vec>::new(); + for entry in active { + if entry.file().level > 0 { + levels + .entry(( + entry.partition().to_vec(), + entry.bucket(), + entry.file().level, + )) + .or_default() + .push(entry.file().clone()); + } + } + for files in levels.into_values().filter(|files| files.len() > 1) { + let overlaps = match &comparator { + Some(comparator) => { + super::merge_tree_split_generator::interval_partition(files, comparator) + .iter() + .any(|section| section.len() > 1) + } + // Empty trimmed keys identify one key in this partition. + None => true, + }; + if overlaps { + return Err(crate::Error::DataInvalid { + message: "LSM conflicts detected: overlapping or unverifiable key ranges in a positive level".into(), + source: None, + }); + } + } + Ok(()) + } + async fn resolve_conflict_base_entries( &self, latest_snapshot: &Option, @@ -3564,6 +3798,13 @@ fn build_partition_stats_row(datums: &[Option], data_types: &[DataType]) /// Plan for resolving commit entries. enum CommitEntriesPlan { + /// Replace unchanged data files without publishing logical row changes. + Compact { + entries: Vec, + schema_id: i64, + read_snapshot: Option>, + planned_snapshot_id: i64, + }, /// Caller-provided entries. May contain `FileKind::Delete` entries from CoW /// rewrites, in which case `resolve_commit` auto-promotes to `CommitKind::OVERWRITE`. Direct { @@ -3588,6 +3829,7 @@ enum CommitEntriesPlan { impl CommitEntriesPlan { fn commit_kind_hint(&self) -> CommitKind { match self { + CommitEntriesPlan::Compact { .. } => CommitKind::COMPACT, CommitEntriesPlan::Direct { entries, new_index_entries, @@ -4328,7 +4570,7 @@ mod tests { delta_probe_count, .. } => (*full_scan_count, *delta_probe_count), - CommitEntriesPlan::Direct { .. } => unreachable!(), + CommitEntriesPlan::Direct { .. } | CommitEntriesPlan::Compact { .. } => unreachable!(), } } diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index 47d38e5e1..f578aff02 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -49,7 +49,7 @@ use std::sync::Arc; mod audit; pub use audit::AuditLogRead; -const MAX_MERGE_INPUT_STREAMS: usize = 256; +pub(super) const MAX_MERGE_INPUT_STREAMS: usize = 256; const MAX_TOTAL_BLOB_READS: usize = 64; /// Table read: reads data from splits (e.g. produced by [TableScan::plan]). diff --git a/crates/paimon/tests/append_compact/planning.rs b/crates/paimon/tests/append_compact/planning.rs new file mode 100644 index 000000000..e67dd7fea --- /dev/null +++ b/crates/paimon/tests/append_compact/planning.rs @@ -0,0 +1,175 @@ +// 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. + +use super::*; + +#[tokio::test] +async fn compaction_recovery_does_not_read_snapshots_older_than_its_plan() { + let table = table(schema(&[])).await; + append_files(&table, 3).await; + let builder = table + .new_compact_builder() + .with_commit_user("bounded-recovery") + .unwrap(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + // Old snapshot contents are irrelevant to a fresh plan's identity. Fail + // reads deterministically instead of relying on remote request timings. + table + .file_io() + .new_output(&format!("{}/snapshot/snapshot-1", table.location())) + .unwrap() + .write(bytes::Bytes::from_static(b"unreadable old snapshot")) + .await + .unwrap(); + let commit = builder.new_commit().unwrap(); + commit.commit_compact(results.clone()).await.unwrap(); + append_files(&table, 1).await; + commit.commit_compact(results).await.unwrap(); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(5) + ); + assert_eq!(rows(&table, None).await.len(), 8); +} + +#[tokio::test] +async fn rollback_invalidates_a_prepared_compaction_plan() { + let table = table(schema(&[])).await; + append_files(&table, 2).await; + table + .new_write_builder() + .new_commit() + .with_ignore_empty_commit(false) + .commit(vec![]) + .await + .unwrap(); + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.snapshot_id(), Some(3)); + let results = execute(&table, &plan).await; + // Equivalent to Java rollback: lower LATEST and remove later snapshots. + table + .file_io() + .delete_file(&table.snapshot_manager().snapshot_path(3)) + .await + .unwrap(); + table.snapshot_manager().write_latest_hint(2).await.unwrap(); + let error = builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap_err(); + assert!(error.to_string().contains("rolled back"), "{error}"); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(2) + ); + assert_eq!(rows(&table, None).await.len(), 4); +} + +#[tokio::test] +async fn append_plan_waits_for_the_default_minimum_and_has_no_file_count_cap() { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .option("bucket", "-1") + .build() + .unwrap(); + let table = table(TableSchema::new(0, &schema)).await; + append_files(&table, 4).await; + let planner = table.new_compact_builder().new_planner().unwrap(); + assert!(planner.plan().await.unwrap().tasks().is_empty()); + append_files(&table, 1).await; + let plan = planner.plan().await.unwrap(); + assert_eq!(plan.tasks().len(), 1); + assert_eq!(plan.tasks()[0].files().len(), 5); + append_files(&table, 35).await; + let plan = planner.plan().await.unwrap(); + assert_eq!(plan.tasks().len(), 1); + assert_eq!(plan.tasks()[0].files().len(), 40); +} + +#[tokio::test] +async fn append_compaction_groups_are_independent_of_read_split_sizes() { + let table = table(schema(&[ + ("target-file-size", "8kb"), + ("source.split.open-file-cost", "8kb"), + ("source.split.target-size", "1b"), + ])) + .await; + append_files(&table, 5).await; + let read_plan = table.new_read_builder().new_scan().plan().await.unwrap(); + assert_eq!(read_plan.splits().len(), 5); + let plan = table + .new_compact_builder() + .new_planner() + .unwrap() + .plan() + .await + .unwrap(); + assert_eq!(plan.tasks().len(), 2); + assert!(plan.tasks().iter().all(|task| task.files().len() == 2)); + let larger_splits = table.copy_with_options(HashMap::from([( + "source.split.target-size".into(), + "1mb".into(), + )])); + let read_plan = larger_splits + .new_read_builder() + .new_scan() + .plan() + .await + .unwrap(); + assert_eq!(read_plan.splits().len(), 1); + let larger_plan = larger_splits + .new_compact_builder() + .new_planner() + .unwrap() + .plan() + .await + .unwrap(); + let groups = |plan: &CompactPlan| { + plan.tasks() + .iter() + .map(|task| { + task.files() + .iter() + .map(|file| file.file_name.clone()) + .collect::>() + }) + .collect::>() + }; + assert_eq!(groups(&plan), groups(&larger_plan)); + let result = execute(&table, &plan).await; + table + .new_compact_builder() + .new_commit() + .unwrap() + .commit_compact(result) + .await + .unwrap(); + assert_eq!(rows(&table, None).await.len(), 10); +} diff --git a/crates/paimon/tests/append_compact_test.rs b/crates/paimon/tests/append_compact_test.rs new file mode 100644 index 000000000..a67da1847 --- /dev/null +++ b/crates/paimon/tests/append_compact_test.rs @@ -0,0 +1,748 @@ +// 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. + +mod common; + +#[path = "append_compact/planning.rs"] +mod planning_tests; + +use arrow_array::{Array, Int32Array, RecordBatch, StringArray}; +use common::incremental_helpers::{ + make_batch, make_partitioned_batch, memory_table, persist_table_schema, pk_schema, setup_dirs, + write_batch, +}; +use futures::TryStreamExt; +use paimon::spec::{ + CommitKind, DataType, Datum, IntType, PredicateBuilder, Schema, TableSchema, VarCharType, +}; +use paimon::table::{CompactPlan, CompactResult, IncrementalScanMode, Table}; +use std::collections::HashMap; + +fn schema(options: &[(&str, &str)]) -> TableSchema { + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .option("bucket", "-1") + .option("file.format", "parquet") + // Most behavioral cases deliberately compact tiny fixtures. + .option("compaction.min.file-num", "2"); + for (key, value) in options { + builder = builder.option(*key, *value); + } + TableSchema::new(0, &builder.build().unwrap()) +} + +async fn table(schema: TableSchema) -> Table { + let path = format!("memory:/compact/{}", uuid::Uuid::new_v4()); + let (io, table) = memory_table(&path, schema); + setup_dirs(&io, &path).await; + persist_table_schema(&io, &path, table.schema()).await; + table +} + +async fn batches(table: &Table, snapshot: Option) -> Vec { + let table = match snapshot { + Some(id) => table.copy_with_pinned_snapshot(Some( + &table.snapshot_manager().get_snapshot(id).await.unwrap(), + )), + None => table.clone(), + }; + let builder = table.new_read_builder(); + let plan = builder.new_scan().plan().await.unwrap(); + builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap() +} + +async fn rows(table: &Table, snapshot: Option) -> Vec<(i32, Option)> { + let mut rows = Vec::new(); + for batch in batches(table, snapshot).await { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + rows.extend( + (0..batch.num_rows()) + .map(|i| (ids.value(i), (!values.is_null(i)).then(|| values.value(i)))), + ); + } + rows.sort_unstable(); + rows +} + +async fn append_files(table: &Table, count: i32) { + for id in 0..count { + write_batch(table, &make_batch(vec![id, id], vec![id * 10, id * 10])).await; + } +} + +async fn execute(table: &Table, plan: &CompactPlan) -> Vec { + let executor = table.new_compact_builder().new_executor().unwrap(); + let mut results = Vec::new(); + for task in plan.tasks() { + results.push(executor.execute(task).await.unwrap()); + } + results +} + +#[tokio::test] +async fn compact_preserves_duplicates_nulls_old_snapshots_and_incremental_events() { + for format in ["parquet", "orc", "avro", "row"] { + let table = table(schema(&[ + ("file.format", format), + ("file.compression", "none"), + ])) + .await; + append_files(&table, 3).await; + let batch = make_batch(vec![9], vec![90]); + let nullable = RecordBatch::try_new( + std::sync::Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", arrow_schema::DataType::Int32, false), + arrow_schema::Field::new("value", arrow_schema::DataType::Int32, true), + ])), + vec![ + batch.column(0).clone(), + std::sync::Arc::new(Int32Array::from(vec![None])), + ], + ) + .unwrap(); + write_batch(&table, &nullable).await; + let expected = rows(&table, None).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.snapshot_id(), Some(4)); + assert_eq!(plan.tasks().len(), 1); + assert_eq!(plan.tasks()[0].files().len(), 4); + + // Both task and result cross a serialization boundary, as on a remote worker. + let plan: CompactPlan = + serde_json::from_slice(&serde_json::to_vec(&plan).unwrap()).unwrap(); + let results = execute(&table, &plan).await; + let results: Vec = + serde_json::from_slice(&serde_json::to_vec(&results).unwrap()).unwrap(); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(4) + ); + let committer = builder.new_commit().unwrap(); + let metrics = committer.commit_compact(results.clone()).await.unwrap(); + assert_eq!(metrics.files_removed, 4, "{format}"); + assert_eq!(metrics.files_added, 1, "{format}"); + assert_eq!(metrics.rows_written, 7, "{format}"); + assert_eq!(rows(&table, None).await, expected, "{format}"); + assert_eq!(rows(&table, Some(4)).await, expected, "{format}"); + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + assert_eq!(snapshot.commit_kind(), &CommitKind::COMPACT); + assert_eq!(snapshot.total_record_count(), Some(7)); + assert_eq!(snapshot.delta_record_count(), Some(0)); + assert_eq!(snapshot.changelog_record_count(), None); + let read = table.new_read_builder(); + let latest = read.new_scan().plan().await.unwrap(); + assert_eq!( + latest.splits().iter().flat_map(|s| s.data_files()).count(), + 1 + ); + assert!(latest + .splits() + .iter() + .flat_map(|s| s.data_files()) + .all(|f| f.file_source == Some(1))); + let incremental = read + .new_incremental_scan(IncrementalScanMode::Delta, 0, 5) + .plan_combined_delta() + .await + .unwrap(); + let events: Vec = read + .new_read() + .unwrap() + .to_arrow(incremental.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(events.iter().map(|b| b.num_rows()).sum::(), 7); + + // External retries reuse the same task identity and publish no second snapshot. + assert_eq!(committer.commit_compact(results).await.unwrap(), metrics); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(5) + ); + } +} + +#[tokio::test] +async fn compact_allows_concurrent_append_but_rejects_stale_inputs() { + let table = table(schema(&[])).await; + append_files(&table, 2).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + let competing = table.new_compact_builder(); + let other = execute( + &table, + &competing.new_planner().unwrap().plan().await.unwrap(), + ) + .await; + write_batch(&table, &make_batch(vec![99], vec![990])).await; + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap(); + let expected = vec![ + (0, Some(0)), + (0, Some(0)), + (1, Some(10)), + (1, Some(10)), + (99, Some(990)), + ]; + assert_eq!(rows(&table, None).await, expected); + let error = competing + .new_commit() + .unwrap() + .commit_compact(other.clone()) + .await + .unwrap_err(); + assert!(error.to_string().contains("conflict"), "{error}"); + assert_eq!(rows(&table, None).await, expected); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(4) + ); + // A rejected commit retains its prepared output; it never deletes a possible published file. + let old = table.copy_with_pinned_snapshot(Some( + &table.snapshot_manager().get_snapshot(2).await.unwrap(), + )); + let split = old.new_read_builder().new_scan().plan().await.unwrap(); + let bucket = split.splits()[0].bucket_path(); + assert!(table + .file_io() + .exists(&format!("{bucket}/{}", other[0].files()[0].file_name)) + .await + .unwrap()); +} + +#[tokio::test] +async fn compact_after_overwrite_does_not_resurrect_old_rows() { + let table = table(schema(&[])).await; + append_files(&table, 2).await; + let compact = table.new_compact_builder(); + let plan = compact.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + let overwrite = table.new_write_builder().with_overwrite(); + let mut writer = overwrite.new_write().unwrap(); + writer + .write_arrow_batch(&make_batch(vec![88], vec![880])) + .await + .unwrap(); + overwrite + .new_commit() + .overwrite(writer.prepare_commit().await.unwrap(), None) + .await + .unwrap(); + assert!(compact + .new_commit() + .unwrap() + .commit_compact(results) + .await + .is_err()); + assert_eq!(rows(&table, None).await, vec![(88, Some(880))]); +} + +#[tokio::test] +async fn compact_skips_empty_single_and_target_sized_files_and_bounds_tasks() { + let table = table(schema(&[ + ("target-file-size", "4mb"), + ("source.split.open-file-cost", "4mb"), + ])) + .await; + let builder = table.new_compact_builder(); + assert!(builder + .new_planner() + .unwrap() + .plan() + .await + .unwrap() + .tasks() + .is_empty()); + builder + .new_commit() + .unwrap() + .commit_compact(vec![]) + .await + .unwrap(); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + None + ); + append_files(&table, 1).await; + assert!(builder + .new_planner() + .unwrap() + .plan() + .await + .unwrap() + .tasks() + .is_empty()); + append_files(&table, 4).await; + let compact = table.new_compact_builder(); + let plan = compact.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.tasks().len(), 2); + assert!(plan.tasks().iter().all(|task| task.files().len() == 2)); + let min_size = plan + .tasks() + .iter() + .flat_map(|t| t.files()) + .map(|f| f.file_size) + .min() + .unwrap(); + let small = table.copy_with_options(HashMap::from([( + "target-file-size".into(), + format!("{min_size}b"), + )])); + assert!(small + .new_compact_builder() + .new_planner() + .unwrap() + .plan() + .await + .unwrap() + .tasks() + .is_empty()); + let mut results = execute(&table, &plan).await; + let committer = compact.new_commit().unwrap(); + committer + .commit_compact(vec![results.remove(0)]) + .await + .unwrap(); + committer.commit_compact(results).await.unwrap(); + assert_eq!(rows(&table, None).await.len(), 10); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(7) + ); +} + +#[tokio::test] +async fn compact_keeps_partitions_and_buckets_separate_and_honors_partition_filter() { + let schema = Schema::builder() + .column("pt", DataType::VarChar(VarCharType::string_type())) + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .partition_keys(["pt"]) + .option("bucket", "2") + .option("bucket-key", "id") + .option("bucket-append-ordered", "false") + .option("compaction.min.file-num", "2") + .build() + .unwrap(); + let table = table(TableSchema::new(0, &schema)).await; + for _ in 0..2 { + for bucket in 0..2 { + for pt in ["a", "b"] { + let write = table.new_write_builder(); + let mut writer = write.new_write().unwrap(); + writer + .write_arrow( + &[make_partitioned_batch( + vec![pt], + vec![bucket], + vec![bucket * 10], + )], + Some(bucket), + ) + .await + .unwrap(); + write + .new_commit() + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + } + } + } + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("pt", Datum::String("a".into())) + .unwrap(); + let builder = table + .new_compact_builder() + .with_partition_filter(predicate) + .unwrap(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.tasks().len(), 2); + assert!(plan.tasks().iter().all(|task| task.files().len() == 2)); + assert_ne!(plan.tasks()[0].bucket(), plan.tasks()[1].bucket()); + let expected = rows(&table, None).await; + let mut results = execute(&table, &plan).await; + results.reverse(); + let committer = builder.new_commit().unwrap(); + committer.commit_compact(results.clone()).await.unwrap(); + results.reverse(); + committer.commit_compact(results).await.unwrap(); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(9) + ); + assert_eq!(rows(&table, None).await, expected); + let mut partition_counts = std::collections::BTreeMap::new(); + for batch in batches(&table, None).await { + let partitions = batch + .column_by_name("pt") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + *partition_counts + .entry(partitions.value(i).to_string()) + .or_insert(0) += 1; + } + } + assert_eq!( + partition_counts, + std::collections::BTreeMap::from([("a".to_string(), 4), ("b".to_string(), 4)]) + ); + let plan = table.new_read_builder().new_scan().plan().await.unwrap(); + assert_eq!(plan.splits().iter().flat_map(|s| s.data_files()).count(), 6); +} + +#[tokio::test] +async fn compact_refuses_unsupported_tables_invalid_options_and_foreign_tasks() { + for options in [ + vec![("row-tracking.enabled", "true")], + vec![ + ("data-evolution.enabled", "true"), + ("row-tracking.enabled", "true"), + ], + vec![("deletion-vectors.enabled", "true")], + vec![("bucket", "1")], + vec![ + ("bucket", "1"), + ("bucket-append-ordered", "false"), + ("clustering.incremental", "true"), + ("clustering.columns", "value"), + ], + vec![ + ("clustering.incremental", "true"), + ("clustering.columns", "value"), + ], + ] { + let table = table(schema(&options)).await; + assert!(table.new_compact_builder().new_planner().is_err()); + assert!(table.new_compact_builder().new_executor().is_err()); + } + let pk = table(pk_schema(&[("changelog-producer", "full-compaction")])).await; + assert!(pk.new_compact_builder().new_planner().is_err()); + let table = table(schema(&[])).await; + let invalid_size = + table.copy_with_options(HashMap::from([("target-file-size".into(), "0b".into())])); + assert!(invalid_size.new_compact_builder().new_planner().is_err()); + for value in ["0", "1", "-1", "invalid", "2147483648"] { + let invalid = table.copy_with_options(HashMap::from([( + "compaction.min.file-num".into(), + value.into(), + )])); + let error = invalid.new_compact_builder().new_planner().err().unwrap(); + assert!(error.to_string().contains("compaction.min.file-num")); + } + for value in ["-0.1", "NaN", "inf", "invalid"] { + let invalid = table.copy_with_options(HashMap::from([( + "compaction.small-file-ratio".into(), + value.into(), + )])); + let error = invalid.new_compact_builder().new_planner().err().unwrap(); + assert!(error.to_string().contains("compaction.small-file-ratio")); + } + append_files(&table, 2).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let foreign = self::table(schema(&[])).await; + assert!(foreign + .new_compact_builder() + .new_executor() + .unwrap() + .execute(&plan.tasks()[0]) + .await + .is_err()); + let results = execute(&table, &plan).await; + assert!(foreign + .new_compact_builder() + .new_commit() + .unwrap() + .commit_compact(results.clone()) + .await + .is_err()); + assert!(builder + .new_commit() + .unwrap() + .commit_compact(vec![results[0].clone(), results[0].clone()]) + .await + .is_err()); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(2) + ); +} + +#[tokio::test] +async fn compact_rebuilds_embedded_and_sidecar_indexes_for_the_merged_row_positions() { + for threshold in ["0 B", "1 MB"] { + let table = table(schema(&[ + ("file-index.bitmap.columns", "id"), + ("file-index.in-manifest-threshold", threshold), + ])) + .await; + append_files(&table, 3).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + assert_eq!(results[0].files().len(), 1); + let file = &results[0].files()[0]; + assert_eq!(file.embedded_index.is_some(), threshold == "1 MB"); + if threshold == "0 B" { + assert!(file.extra_files.iter().any(|name| name.ends_with(".index"))); + } + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap(); + for id in [0, 1, 2, 9] { + let mut read = table.new_read_builder(); + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("id", Datum::Int(id)) + .unwrap(); + read.with_filter(predicate); + let plan = read.new_scan().plan().await.unwrap(); + let batches: Vec = read + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + if id == 9 { 0 } else { 2 } + ); + } + } +} + +#[tokio::test] +async fn compact_rejects_mismatched_input_metadata_and_changed_table_schema() { + let table = table(schema(&[])).await; + append_files(&table, 2).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + // A result arriving from a worker must retain the source metadata exactly. + let mut encoded = serde_json::to_value(results).unwrap(); + encoded[0]["task"]["files"][0]["_MIN_SEQUENCE_NUMBER"] = serde_json::json!(99); + let results: Vec = serde_json::from_value(encoded).unwrap(); + let error = builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap_err(); + assert!(error.to_string().contains("metadata changed"), "{error}"); + assert_eq!(rows(&table, None).await.len(), 4); + + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + let changed = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .column("extra", DataType::Int(IntType::new())) + .option("bucket", "-1") + .build() + .unwrap(); + persist_table_schema( + table.file_io(), + table.location(), + &TableSchema::new(1, &changed), + ) + .await; + let error = builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap_err(); + assert!(error.to_string().contains("schema changed"), "{error}"); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(2) + ); +} + +#[tokio::test] +async fn compact_execution_failure_cleans_only_its_unsubmitted_output() { + let table = table(schema(&[])).await; + append_files(&table, 2).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let files = plan.tasks()[0].files(); + let source = table.new_read_builder().new_scan().plan().await.unwrap(); + let bucket = source.splits()[0].bucket_path(); + table + .file_io() + .delete_file(&format!("{bucket}/{}", files[1].file_name)) + .await + .unwrap(); + let before = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(); + assert!(builder + .new_executor() + .unwrap() + .execute(&plan.tasks()[0]) + .await + .is_err()); + assert_eq!( + table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(), + before + ); + assert!(table + .file_io() + .exists(&format!("{bucket}/{}", files[0].file_name)) + .await + .unwrap()); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(2) + ); +} + +#[tokio::test] +async fn compact_obeys_the_executor_memory_budget_without_changing_the_table() { + let table = table(schema(&[])).await; + append_files(&table, 2).await; + let plan = table + .new_compact_builder() + .new_planner() + .unwrap() + .plan() + .await + .unwrap(); + let resources = paimon::resource::ResourceContext::builder() + .memory_limit(1) + .build() + .unwrap(); + let executor = table + .new_compact_builder() + .with_resources(resources.clone()) + .new_executor() + .unwrap(); + let before = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(); + let error = executor.execute(&plan.tasks()[0]).await.unwrap_err(); + assert!( + matches!(error, paimon::Error::ResourceExhausted { .. }), + "{error}" + ); + assert_eq!(resources.metrics().reserved_memory_bytes, 0); + assert_eq!( + table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(), + before + ); + assert_eq!(rows(&table, None).await.len(), 4); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(2) + ); +} diff --git a/crates/paimon/tests/pk_compact/dynamic_bucket.rs b/crates/paimon/tests/pk_compact/dynamic_bucket.rs new file mode 100644 index 000000000..4a9b69549 --- /dev/null +++ b/crates/paimon/tests/pk_compact/dynamic_bucket.rs @@ -0,0 +1,472 @@ +// 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. + +use super::*; +use paimon::spec::{IndexManifest, IndexManifestEntry, Snapshot}; +use std::collections::BTreeMap; + +async fn latest(table: &Table) -> Snapshot { + table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap() +} + +async fn indexes(table: &Table) -> Vec { + let snapshot = latest(table).await; + let Some(name) = snapshot.index_manifest() else { + return vec![]; + }; + IndexManifest::read( + table.file_io(), + &table.snapshot_manager().manifest_path(name), + ) + .await + .unwrap() +} + +async fn compact(table: &Table) { + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert!(!plan.tasks().is_empty()); + builder + .new_commit() + .unwrap() + .commit_compact(execute(table, &plan).await) + .await + .unwrap(); +} + +// Read each bucket separately so a lost HASH mapping cannot hide behind an +// otherwise correct full-table result before the next writer restarts. +async fn bucket_rows(table: &Table) -> BTreeMap { + let builder = table.new_read_builder(); + let plan = builder.new_scan().plan_all().await.unwrap(); + let read = builder.new_read().unwrap(); + let mut result = BTreeMap::new(); + for split in plan.splits() { + assert_eq!(split.total_buckets(), -1); + let batches: Vec = read + .to_arrow(std::slice::from_ref(split)) + .unwrap() + .try_collect() + .await + .unwrap(); + for batch in batches { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + assert!(result + .insert(ids.value(i), (split.bucket(), values.value(i))) + .is_none()); + } + } + } + result +} + +#[tokio::test] +async fn dynamic_compaction_preserves_hash_assignments_for_all_merge_engines() { + for engine in ["deduplicate", "first-row", "partial-update", "aggregation"] { + for in_data_dir in ["false", "true"] { + let mut options = vec![ + ("bucket", "-1"), + ("dynamic-bucket.target-row-num", "1"), + ("merge-engine", engine), + ("index-file-in-data-file-dir", in_data_dir), + ]; + if engine == "aggregation" { + options.push(("fields.value.aggregate-function", "sum")); + } + let table = table(&options).await; + write_batch(&table, &make_batch(vec![1, 2, 3], vec![10, 20, 30])).await; + write_batch(&table, &make_batch(vec![1, 2], vec![11, 21])).await; + let before = bucket_rows(&table).await; + assert!(before.values().map(|r| r.0).max().unwrap() > 0); + let snapshot = latest(&table).await; + let original_indexes = indexes(&table).await; + assert_eq!(original_indexes.len(), 3); + for entry in &original_indexes { + let dir = if in_data_dir == "true" { + format!("bucket-{}", entry.bucket) + } else { + "index".into() + }; + let path = format!("{}/{dir}/{}", table.location(), entry.index_file.file_name); + assert!(table.file_io().exists(&path).await.unwrap(), "{path}"); + } + compact(&table).await; + assert_eq!(bucket_rows(&table).await, before, "{engine}"); + assert_eq!( + latest(&table).await.index_manifest(), + snapshot.index_manifest() + ); + assert_eq!(indexes(&table).await, original_indexes); + + // A fresh writer restores HASH assignments and capacity after compaction. + write_batch(&table, &make_batch(vec![1, 4], vec![100, 40])).await; + compact(&table).await; + let after = bucket_rows(&table).await; + assert_eq!(after.len(), 4); + let value = match engine { + "first-row" => 10, + "aggregation" => 121, + _ => 100, + }; + assert_eq!(after[&1], (before[&1].0, value)); + assert_eq!(after[&2], before[&2]); + assert_eq!(after[&3], before[&3]); + assert_eq!(after[&4].1, 40); + assert!(after[&4].0 > before.values().map(|r| r.0).max().unwrap()); + } + } +} + +#[tokio::test] +async fn dynamic_compaction_commits_the_latest_hash_index_after_concurrent_append() { + let table = table(&[ + ("bucket", "-1"), + ("dynamic-bucket.target-row-num", "2"), + ("dynamic-bucket.max-buckets", "1"), + ]) + .await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + write_batch(&table, &make_batch(vec![1], vec![11])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + let captured = latest(&table).await; + write_batch(&table, &make_batch(vec![1, 2], vec![12, 20])).await; + let appended = latest(&table).await; + let appended_indexes = indexes(&table).await; + assert_ne!(captured.index_manifest(), appended.index_manifest()); + let commit = builder.new_commit().unwrap(); + commit.commit_compact(results.clone()).await.unwrap(); + assert_eq!( + latest(&table).await.index_manifest(), + appended.index_manifest() + ); + assert_eq!(indexes(&table).await, appended_indexes); + assert_eq!(rows(&table).await, vec![(1, 12), (2, 20)]); + let snapshot_id = latest(&table).await.id(); + commit.commit_compact(results).await.unwrap(); + assert_eq!(latest(&table).await.id(), snapshot_id); + write_batch(&table, &make_batch(vec![2, 3], vec![21, 30])).await; + compact(&table).await; + assert_eq!(rows(&table).await, vec![(1, 12), (2, 21), (3, 30)]); +} + +#[tokio::test] +async fn dynamic_compaction_retains_deleted_hashes_and_bucket_capacity() { + let table = table(&[("bucket", "-1"), ("dynamic-bucket.target-row-num", "1")]).await; + write_batch(&table, &make_batch(vec![1, 2], vec![10, 20])).await; + let before = bucket_rows(&table).await; + write_batch(&table, &make_batch_with_kinds(vec![1], vec![10], vec![3])).await; + let original_indexes = indexes(&table).await; + compact(&table).await; + assert_eq!(rows(&table).await, vec![(2, 20)]); + assert_eq!(indexes(&table).await, original_indexes); + write_batch(&table, &make_batch(vec![1, 3], vec![11, 30])).await; + let after = bucket_rows(&table).await; + assert_eq!(after[&1], (before[&1].0, 11)); + assert_eq!(after[&2], before[&2]); + assert!(after[&3].0 > before.values().map(|r| r.0).max().unwrap()); +} + +async fn partitioned(cross_partition: bool, options: &[(&str, &str)]) -> Table { + let mut schema = Schema::builder() + .column("pt", DataType::VarChar(VarCharType::string_type())) + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .primary_key(if cross_partition { + vec!["id"] + } else { + vec!["pt", "id"] + }) + .partition_keys(["pt"]) + .option("bucket", "-1") + .option("dynamic-bucket.target-row-num", "1"); + for (key, value) in options { + schema = schema.option(*key, *value); + } + table_with_schema(TableSchema::new(0, &schema.build().unwrap())).await +} + +async fn partition_rows(table: &Table) -> Vec<(String, i32, i32)> { + let builder = table.new_read_builder(); + let plan = builder.new_scan().plan_all().await.unwrap(); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut result = vec![]; + for batch in batches { + let pts = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let ids = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + result.extend( + (0..batch.num_rows()) + .map(|i| (pts.value(i).to_string(), ids.value(i), values.value(i))), + ); + } + result.sort(); + result +} + +#[tokio::test] +async fn dynamic_partition_compaction_keeps_unselected_partitions_and_indexes() { + let table = partitioned(false, &[]).await; + for values in [vec![10, 20], vec![11, 21]] { + write_batch( + &table, + &make_partitioned_batch(vec!["a", "b"], vec![1, 2], values), + ) + .await; + } + let before = indexes(&table).await; + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("pt", Datum::String("a".into())) + .unwrap(); + let builder = table + .new_compact_builder() + .with_partition_filter(predicate) + .unwrap(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.tasks().len(), 1); + builder + .new_commit() + .unwrap() + .commit_compact(execute(&table, &plan).await) + .await + .unwrap(); + assert_eq!(indexes(&table).await, before); + assert_eq!( + partition_rows(&table).await, + vec![("a".into(), 1, 11), ("b".into(), 2, 21)] + ); + let plan = table + .new_compact_builder() + .new_planner() + .unwrap() + .plan() + .await + .unwrap(); + assert_eq!(plan.tasks().len(), 1, "only partition b still has L0 files"); + write_batch( + &table, + &make_partitioned_batch(vec!["a", "b"], vec![1, 2], vec![12, 22]), + ) + .await; + compact(&table).await; + assert_eq!( + partition_rows(&table).await, + vec![("a".into(), 1, 12), ("b".into(), 2, 22)] + ); +} + +#[tokio::test] +async fn cross_partition_compaction_preserves_migrations_and_restart_for_all_engines() { + for engine in ["deduplicate", "first-row", "partial-update", "aggregation"] { + let mut options = vec![("merge-engine", engine)]; + if engine == "aggregation" { + options.push(("fields.value.aggregate-function", "sum")); + } + let table = partitioned(true, &options).await; + write_batch( + &table, + &make_partitioned_batch(vec!["a"], vec![1], vec![10]), + ) + .await; + write_batch( + &table, + &make_partitioned_batch(vec!["b"], vec![1], vec![20]), + ) + .await; + let before = partition_rows(&table).await; + compact(&table).await; + assert_eq!(partition_rows(&table).await, before, "{engine}"); + assert!(latest(&table).await.index_manifest().is_none()); + write_batch( + &table, + &make_partitioned_batch(vec!["c"], vec![1], vec![30]), + ) + .await; + if engine == "first-row" { + // The restored global index ignores this existing key altogether. + assert!(table + .new_compact_builder() + .new_planner() + .unwrap() + .plan() + .await + .unwrap() + .tasks() + .is_empty()); + } else { + compact(&table).await; + } + let (partition, value) = match engine { + "deduplicate" => ("c", 30), + "first-row" => ("a", 10), + "partial-update" => ("a", 30), + "aggregation" => ("a", 60), + _ => unreachable!(), + }; + assert_eq!( + partition_rows(&table).await, + vec![(partition.into(), 1, value)], + "{engine}" + ); + } +} + +#[tokio::test] +async fn cross_partition_compaction_preserves_concurrent_migration_deletes() { + let table = partitioned(true, &[]).await; + write_batch( + &table, + &make_partitioned_batch(vec!["a"], vec![1], vec![10]), + ) + .await; + write_batch( + &table, + &make_partitioned_batch(vec!["b"], vec![1], vec![20]), + ) + .await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + write_batch( + &table, + &make_partitioned_batch(vec!["c"], vec![1], vec![30]), + ) + .await; + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap(); + assert_eq!(partition_rows(&table).await, vec![("c".into(), 1, 30)]); + write_batch( + &table, + &make_partitioned_batch(vec!["a"], vec![1], vec![40]), + ) + .await; + compact(&table).await; + assert_eq!(partition_rows(&table).await, vec![("a".into(), 1, 40)]); +} + +#[tokio::test] +async fn dynamic_compaction_rejects_other_index_types_at_plan_and_commit() { + for index_type in [ + "DELETION_VECTORS", + "GLOBAL_INDEX", + "UNKNOWN", + "HASH_WITH_DV_METADATA", + "HASH_WITH_GLOBAL_METADATA", + ] { + let table = table(&[("bucket", "-1")]).await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + write_batch(&table, &make_batch(vec![2], vec![20])).await; + let snapshot = latest(&table).await; + let mut entries = indexes(&table).await; + let mut unsupported = entries[0].clone(); + match index_type { + "HASH_WITH_DV_METADATA" => { + unsupported.index_file.deletion_vectors_ranges = Some(Default::default()); + } + "HASH_WITH_GLOBAL_METADATA" => { + unsupported.index_file.global_index_meta = Some(paimon::spec::GlobalIndexMeta { + row_range_start: 0, + row_range_end: 1, + index_field_id: 0, + extra_field_ids: None, + index_meta: None, + source_meta: None, + }); + } + _ => unsupported.index_file.index_type = index_type.into(), + } + entries.push(unsupported); + let manager = table.snapshot_manager(); + let name = "unsupported-index-manifest"; + IndexManifest::write(table.file_io(), &manager.manifest_path(name), &entries) + .await + .unwrap(); + let mut altered = serde_json::to_value(&snapshot).unwrap(); + altered["indexManifest"] = name.into(); + table + .file_io() + .new_output(&manager.snapshot_path(snapshot.id())) + .unwrap() + .write(serde_json::to_vec(&altered).unwrap().into()) + .await + .unwrap(); + let error = builder.new_planner().unwrap().plan().await.unwrap_err(); + assert!( + error + .to_string() + .contains("only supports dynamic-bucket HASH"), + "{error}" + ); + let error = builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("only supports dynamic-bucket HASH"), + "{error}" + ); + assert_eq!(latest(&table).await.id(), snapshot.id()); + } +} diff --git a/crates/paimon/tests/pk_compact/merge_engines.rs b/crates/paimon/tests/pk_compact/merge_engines.rs new file mode 100644 index 000000000..37ec8fc12 --- /dev/null +++ b/crates/paimon/tests/pk_compact/merge_engines.rs @@ -0,0 +1,356 @@ +// 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. + +use super::*; +use arrow_array::Int8Array; +use std::sync::Arc; + +#[path = "merge_engines_edges.rs"] +mod edges; + +async fn compact(table: &Table) -> paimon::table::CompactMetrics { + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert!(!plan.tasks().is_empty()); + builder + .new_commit() + .unwrap() + .commit_compact(execute(table, &plan).await) + .await + .unwrap() +} + +async fn all_rows(table: &Table) -> Vec<(i32, i32)> { + let builder = table.new_read_builder(); + let plan = builder + .new_scan() + .with_scan_all_files() + .plan() + .await + .unwrap(); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut result = Vec::new(); + for batch in batches { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + result.extend((0..batch.num_rows()).map(|i| (ids.value(i), values.value(i)))); + } + result.sort_unstable(); + result +} + +async fn partial(options: &[(&str, &str)]) -> Table { + let mut schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("version", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .column("other", DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "1") + .option("merge-engine", "partial-update"); + for (key, value) in options { + schema = schema.option(*key, *value); + } + table_with_schema(TableSchema::new(0, &schema.build().unwrap())).await +} + +fn update(version: Option, value: Option, other: Option, kind: i8) -> RecordBatch { + RecordBatch::try_new( + Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", arrow_schema::DataType::Int32, false), + arrow_schema::Field::new("version", arrow_schema::DataType::Int32, true), + arrow_schema::Field::new("value", arrow_schema::DataType::Int32, true), + arrow_schema::Field::new("other", arrow_schema::DataType::Int32, true), + arrow_schema::Field::new("_VALUE_KIND", arrow_schema::DataType::Int8, false), + ])), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![version])), + Arc::new(Int32Array::from(vec![value])), + Arc::new(Int32Array::from(vec![other])), + Arc::new(Int8Array::from(vec![kind])), + ], + ) + .unwrap() +} + +async fn state(table: &Table) -> Vec<(Option, Option, Option)> { + let mut result = Vec::new(); + for batch in read(table, None).await { + for row in 0..batch.num_rows() { + let get = |name| { + let array = batch + .column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (!array.is_null(row)).then(|| array.value(row)) + }; + result.push((get("version"), get("value"), get("other"))); + } + } + result.sort_unstable(); + result +} + +#[tokio::test] +async fn first_row_compaction_preserves_earliest_rows_and_future_writes() { + for format in ["parquet", "orc", "avro", "row"] { + let table = table(&[ + ("merge-engine", "first-row"), + ("file.format", format), + ("file.compression", "none"), + ]) + .await; + write_batch(&table, &make_batch(vec![1, 2], vec![10, 20])).await; + write_batch(&table, &make_batch(vec![1, 3], vec![11, 30])).await; + assert_eq!(all_rows(&table).await, vec![(1, 10), (2, 20), (3, 30)]); + let metrics = compact(&table).await; + assert_eq!(metrics.rows_written, 3); + assert_eq!( + rows(&table).await, + vec![(1, 10), (2, 20), (3, 30)], + "{format}" + ); + write_batch(&table, &make_batch(vec![1, 4], vec![99, 40])).await; + compact(&table).await; + assert_eq!( + rows(&table).await, + vec![(1, 10), (2, 20), (3, 30), (4, 40)], + "{format}" + ); + } +} + +#[tokio::test] +async fn partial_update_compaction_combines_sparse_columns_and_keeps_sequence_metadata() { + for format in ["parquet", "orc", "avro", "row"] { + let table = partial(&[("file.format", format), ("file.compression", "none")]).await; + write_batch(&table, &update(Some(1), Some(10), None, 0)).await; + write_batch(&table, &update(None, None, Some(20), 0)).await; + let before = sequences(&table).await; + compact(&table).await; + assert_eq!( + state(&table).await, + vec![(Some(1), Some(10), Some(20))], + "{format}" + ); + assert_eq!(sequences(&table).await, before); + write_batch(&table, &update(Some(2), None, Some(30), 0)).await; + compact(&table).await; + assert_eq!( + state(&table).await, + vec![(Some(2), Some(10), Some(30))], + "{format}" + ); + } +} + +#[tokio::test] +async fn partial_update_compaction_keeps_sequence_groups_and_field_aggregate_state() { + let table = partial(&[ + ("fields.version.sequence-group", "value"), + ("fields.value.aggregate-function", "sum"), + ]) + .await; + write_batch(&table, &update(Some(10), Some(100), Some(1), 0)).await; + write_batch(&table, &update(Some(9), Some(5), Some(2), 0)).await; + let expected = state(&table).await; + assert_eq!(expected, vec![(Some(10), Some(105), Some(2))]); + compact(&table).await; + assert_eq!(state(&table).await, expected); + write_batch(&table, &update(Some(8), Some(7), Some(3), 0)).await; + assert_eq!(state(&table).await, vec![(Some(10), Some(112), Some(3))]); + compact(&table).await; + write_batch(&table, &update(Some(11), Some(2), None, 0)).await; + assert_eq!(state(&table).await, vec![(Some(11), Some(114), Some(3))]); +} + +#[tokio::test] +async fn partial_update_compaction_keeps_sequence_group_delete_and_payload() { + let table = partial(&[ + ("fields.version.sequence-group", "value,other"), + ("partial-update.remove-record-on-sequence-group", "version"), + ]) + .await; + write_batch(&table, &update(Some(10), Some(100), Some(7), 0)).await; + write_batch(&table, &update(Some(11), Some(20), None, 3)).await; + assert!(state(&table).await.is_empty()); + assert_eq!(compact(&table).await.rows_written, 1); + assert!(state(&table).await.is_empty()); + write_batch(&table, &update(Some(9), Some(999), Some(999), 0)).await; + assert!(state(&table).await.is_empty()); + compact(&table).await; + write_batch(&table, &update(Some(12), Some(3), Some(8), 0)).await; + assert_eq!(state(&table).await, vec![(Some(12), Some(3), Some(8))]); + + let table = partial(&[("partial-update.remove-record-on-delete", "true")]).await; + write_batch(&table, &update(Some(1), Some(10), Some(7), 0)).await; + write_batch(&table, &update(None, Some(20), Some(30), 3)).await; + compact(&table).await; + assert!(state(&table).await.is_empty()); + write_batch(&table, &update(None, None, Some(40), 0)).await; + assert_eq!(state(&table).await, vec![(None, Some(20), Some(40))]); +} + +#[tokio::test] +async fn aggregation_compaction_keeps_accumulators_for_later_updates() { + for (function, before, after) in [ + ("sum", 30, 35), + ("product", 200, 1000), + ("min", 10, 5), + ("max", 20, 20), + ("first_value", 10, 10), + ("last_value", 20, 5), + ] { + let table = table(&[ + ("merge-engine", "aggregation"), + ("fields.value.aggregate-function", function), + ]) + .await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + write_batch(&table, &make_batch(vec![1], vec![20])).await; + assert_eq!(rows(&table).await, vec![(1, before)]); + assert_eq!(compact(&table).await.rows_written, 1); + assert_eq!(rows(&table).await, vec![(1, before)], "{function}"); + write_batch(&table, &make_batch(vec![1], vec![5])).await; + assert_eq!(rows(&table).await, vec![(1, after)], "{function}"); + compact(&table).await; + assert_eq!(rows(&table).await, vec![(1, after)], "{function}"); + } +} + +#[tokio::test] +async fn aggregation_compaction_keeps_singleton_retract_and_merged_delete_state() { + let table = table(&[ + ("merge-engine", "aggregation"), + ("fields.value.aggregate-function", "sum"), + ]) + .await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + write_batch( + &table, + &make_batch_with_kinds(vec![1, 2], vec![3, 40], vec![3, 3]), + ) + .await; + assert_eq!(rows(&table).await, vec![(1, 7)]); + compact(&table).await; + assert_eq!(rows(&table).await, vec![(1, 7)]); + write_batch(&table, &make_batch(vec![1, 2], vec![5, 3])).await; + assert_eq!(rows(&table).await, vec![(1, 12), (2, -37)]); + compact(&table).await; + assert_eq!(rows(&table).await, vec![(1, 12), (2, -37)]); + + let table = super::table(&[ + ("merge-engine", "aggregation"), + ("fields.value.aggregate-function", "sum"), + ("aggregation.remove-record-on-delete", "true"), + ]) + .await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + write_batch( + &table, + &make_batch_with_kinds(vec![1, 2], vec![5, 40], vec![3, 3]), + ) + .await; + compact(&table).await; + assert!(rows(&table).await.is_empty()); + write_batch(&table, &make_batch(vec![1, 2], vec![7, 3])).await; + assert_eq!(rows(&table).await, vec![(1, 12), (2, 43)]); +} + +#[tokio::test] +async fn first_row_legacy_retract_is_rejected_or_ignored_without_leaking_files() { + for ignore in [false, true] { + let table = table(&[ + ("merge-engine", "first-row"), + ("ignore-delete", if ignore { "true" } else { "false" }), + ("write-buffer-size", "1 B"), + ]) + .await; + let legacy = table.copy_with_options(std::collections::HashMap::from([ + ("merge-engine".into(), "deduplicate".into()), + ("ignore-delete".into(), "false".into()), + ])); + write_batch(&legacy, &make_batch(vec![1], vec![10])).await; + write_batch(&legacy, &make_batch_with_kinds(vec![2], vec![20], vec![3])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let before = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(); + let result = builder + .new_executor() + .unwrap() + .execute(&plan.tasks()[0]) + .await; + if ignore { + let result = result.unwrap(); + assert_eq!(result.metrics().rows_written, 1); + builder + .new_commit() + .unwrap() + .commit_compact(vec![result]) + .await + .unwrap(); + assert_eq!(rows(&table).await, vec![(1, 10)]); + } else { + assert!(result + .unwrap_err() + .to_string() + .contains("does not support DELETE")); + assert_eq!( + table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(), + before + ); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(2) + ); + } + } +} diff --git a/crates/paimon/tests/pk_compact/merge_engines_edges.rs b/crates/paimon/tests/pk_compact/merge_engines_edges.rs new file mode 100644 index 000000000..b9f7f55a5 --- /dev/null +++ b/crates/paimon/tests/pk_compact/merge_engines_edges.rs @@ -0,0 +1,390 @@ +// 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. + +use super::*; + +async fn assert_unsafe_aggregate_compaction_rejected(table: &Table) { + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let before = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(); + let error = builder + .new_executor() + .unwrap() + .execute(&plan.tasks()[0]) + .await + .unwrap_err(); + assert!( + error.to_string().contains("Compaction cannot preserve"), + "{error}" + ); + assert_eq!( + table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(), + before + ); +} + +#[tokio::test] +async fn partial_update_retract_only_aggregate_state_refuses_lossy_compaction() { + for aggregate_option in [ + "fields.value.aggregate-function", + "fields.default-aggregate-function", + ] { + let options = [ + ("fields.version.sequence-group", "value,other"), + (aggregate_option, "sum"), + ]; + let control = partial(&options).await; + let rewritten = partial(&options).await; + for table in [&control, &rewritten] { + write_batch(table, &update(Some(10), Some(10), None, 3)).await; + write_batch(table, &update(Some(11), Some(5), None, 3)).await; + } + let builder = rewritten.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let paths_before = rewritten + .file_io() + .list_status_recursive(rewritten.location()) + .await + .unwrap(); + let error = builder + .new_executor() + .unwrap() + .execute(&plan.tasks()[0]) + .await + .unwrap_err(); + assert!(error.to_string().contains("DELETE aggregate state")); + assert_eq!( + rewritten + .file_io() + .list_status_recursive(rewritten.location()) + .await + .unwrap() + .len(), + paths_before.len() + ); + for table in [&control, &rewritten] { + write_batch(table, &update(Some(12), Some(3), None, 0)).await; + } + assert_eq!(state(&control).await, vec![(Some(12), Some(-2), None)]); + assert_eq!(state(&rewritten).await, state(&control).await); + } +} + +#[tokio::test] +async fn partial_update_aggregate_delete_state_refuses_lossy_compaction() { + let options = [ + ("fields.version.sequence-group", "value,other"), + ("partial-update.remove-record-on-sequence-group", "version"), + ("fields.value.aggregate-function", "first_value"), + ]; + let control = partial(&options).await; + let rewritten = partial(&options).await; + for table in [&control, &rewritten] { + write_batch(table, &update(Some(1), Some(10), None, 0)).await; + write_batch(table, &update(Some(2), Some(20), None, 3)).await; + } + assert_unsafe_aggregate_compaction_rejected(&rewritten).await; + for table in [&control, &rewritten] { + write_batch(table, &update(Some(3), Some(30), None, 0)).await; + } + assert_eq!(state(&control).await, vec![(Some(3), Some(20), None)]); + assert_eq!(state(&rewritten).await, state(&control).await); +} + +#[tokio::test] +async fn aggregation_delete_state_refuses_lossy_compaction() { + let options = [ + ("merge-engine", "aggregation"), + ("fields.value.aggregate-function", "first_value"), + ("aggregation.remove-record-on-delete", "true"), + ]; + let control = table(&options).await; + let rewritten = table(&options).await; + for table in [&control, &rewritten] { + write_batch(table, &make_batch(vec![1], vec![10])).await; + write_batch(table, &make_batch_with_kinds(vec![1], vec![20], vec![3])).await; + } + assert_unsafe_aggregate_compaction_rejected(&rewritten).await; + for table in [&control, &rewritten] { + write_batch(table, &make_batch(vec![1], vec![30])).await; + } + assert_eq!(rows(&control).await, vec![(1, 20)]); + assert_eq!(rows(&rewritten).await, rows(&control).await); +} + +#[tokio::test] +async fn first_non_null_aggregate_rejects_retract_state_even_when_output_is_insert() { + for function in ["first_non_null_value", "first_not_null_value"] { + let table = table(&[ + ("merge-engine", "aggregation"), + ("fields.value.aggregate-function", function), + ("aggregation.remove-record-on-delete", "true"), + ]) + .await; + write_batch(&table, &make_batch_with_kinds(vec![1], vec![20], vec![3])).await; + let null_insert = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int32Array::from(vec![1])) as arrow_array::ArrayRef, + ), + ( + "value", + Arc::new(Int32Array::from(vec![None])) as arrow_array::ArrayRef, + ), + ]) + .unwrap(); + write_batch(&table, &null_insert).await; + assert_eq!(rows(&table).await, vec![(1, 20)]); + assert_unsafe_aggregate_compaction_rejected(&table).await; + write_batch(&table, &make_batch(vec![1], vec![30])).await; + assert_eq!(rows(&table).await, vec![(1, 30)]); + } +} + +#[tokio::test] +async fn aggregation_replay_guard_uses_effective_functions() { + for ignored in [ + "fields.default-aggregate-function", + "fields.id.aggregate-function", + ] { + let table = table(&[ + ("merge-engine", "aggregation"), + (ignored, "first_value"), + ("fields.value.aggregate-function", "sum"), + ]) + .await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + write_batch(&table, &make_batch_with_kinds(vec![1], vec![3], vec![1])).await; + compact(&table).await; + assert_eq!(rows(&table).await, vec![(1, 7)]); + write_batch(&table, &make_batch(vec![1], vec![5])).await; + assert_eq!(rows(&table).await, vec![(1, 12)]); + } +} + +#[tokio::test] +async fn listagg_rejects_raw_delete_state_that_an_insert_would_normalize() { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::VarChar(VarCharType::string_type())) + .primary_key(["id"]) + .option("bucket", "1") + .option("merge-engine", "aggregation") + .option("fields.value.aggregate-function", "listagg") + .option("aggregation.remove-record-on-delete", "true") + .build() + .unwrap(); + let table = table_with_schema(TableSchema::new(0, &schema)).await; + let batch = |value: Option<&str>, kind: i8| { + RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int32Array::from(vec![1])) as arrow_array::ArrayRef, + ), + ( + "value", + Arc::new(StringArray::from(vec![value])) as arrow_array::ArrayRef, + ), + ( + "_VALUE_KIND", + Arc::new(Int8Array::from(vec![kind])) as arrow_array::ArrayRef, + ), + ]) + .unwrap() + }; + write_batch(&table, &batch(Some(" "), 3)).await; + write_batch(&table, &batch(None, 0)).await; + assert_unsafe_aggregate_compaction_rejected(&table).await; + write_batch(&table, &batch(None, 0)).await; + let output = read(&table, None).await; + assert_eq!( + output[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + " " + ); +} + +#[tokio::test] +async fn partial_update_replay_guard_allows_ignored_legacy_retracts() { + let table = partial(&[ + ("fields.version.sequence-group", "value,other"), + ("fields.value.aggregate-function", "first_value"), + ("ignore-delete", "true"), + ]) + .await; + write_batch(&table, &update(Some(1), Some(10), None, 0)).await; + let legacy = table.copy_with_options(std::collections::HashMap::from([( + "ignore-delete".into(), + "false".into(), + )])); + write_batch(&legacy, &update(Some(2), Some(20), None, 3)).await; + compact(&table).await; + assert_eq!(state(&table).await, vec![(Some(1), Some(10), None)]); + write_batch(&table, &update(Some(3), Some(30), None, 0)).await; + assert_eq!(state(&table).await, vec![(Some(3), Some(10), None)]); +} + +#[tokio::test] +async fn ignored_legacy_retracts_can_compact_to_empty_and_accept_new_records() { + for engine in ["first-row", "partial-update"] { + let table = table(&[("merge-engine", engine), ("ignore-delete", "true")]).await; + let legacy = table.copy_with_options(std::collections::HashMap::from([ + ("merge-engine".into(), "deduplicate".into()), + ("ignore-delete".into(), "false".into()), + ])); + for kind in [1, 3] { + write_batch( + &legacy, + &make_batch_with_kinds(vec![1], vec![10], vec![kind]), + ) + .await; + } + assert!(all_rows(&table).await.is_empty()); + let metrics = compact(&table).await; + assert_eq!( + ( + metrics.files_removed, + metrics.files_added, + metrics.rows_written + ), + (2, 0, 0), + "{engine}" + ); + assert!(rows(&table).await.is_empty()); + write_batch(&table, &make_batch(vec![1], vec![20])).await; + compact(&table).await; + assert_eq!(rows(&table).await, vec![(1, 20)], "{engine}"); + } +} + +#[tokio::test] +async fn aggregation_compaction_keeps_listagg_state_across_formats() { + for format in ["parquet", "orc", "avro", "row"] { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::VarChar(VarCharType::string_type())) + .primary_key(["id"]) + .option("bucket", "1") + .option("merge-engine", "aggregation") + .option("fields.value.aggregate-function", "listagg") + .option("fields.value.list-agg-delimiter", "|") + .option("file.format", format) + .option("file.compression", "none") + .build() + .unwrap(); + let table = table_with_schema(TableSchema::new(0, &schema)).await; + let batch = |text: &str| { + RecordBatch::try_new( + Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", arrow_schema::DataType::Int32, false), + arrow_schema::Field::new("value", arrow_schema::DataType::Utf8, true), + ])), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec![text])), + ], + ) + .unwrap() + }; + write_batch(&table, &batch("a")).await; + write_batch(&table, &batch("b")).await; + compact(&table).await; + write_batch(&table, &batch("c")).await; + compact(&table).await; + let result = read(&table, None).await; + assert_eq!( + result.iter().map(|batch| batch.num_rows()).sum::(), + 1 + ); + assert_eq!( + result[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "a|b|c", + "{format}" + ); + } +} + +#[tokio::test] +async fn aggregation_unsupported_retract_aborts_already_written_outputs() { + let table = table(&[ + ("merge-engine", "aggregation"), + ("fields.value.aggregate-function", "min"), + ("write-buffer-size", "1 B"), + ("read.batch-size", "1"), + ]) + .await; + write_batch(&table, &make_batch(vec![1, 2], vec![10, 20])).await; + write_batch( + &table, + &make_batch_with_kinds(vec![1, 2], vec![11, 5], vec![0, 3]), + ) + .await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let before = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(); + let error = builder + .new_executor() + .unwrap() + .execute(&plan.tasks()[0]) + .await + .unwrap_err(); + assert!( + error.to_string().contains("does not support retract"), + "{error}" + ); + assert_eq!( + table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(), + before + ); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(2) + ); +} diff --git a/crates/paimon/tests/pk_compact_test.rs b/crates/paimon/tests/pk_compact_test.rs new file mode 100644 index 000000000..64541bcca --- /dev/null +++ b/crates/paimon/tests/pk_compact_test.rs @@ -0,0 +1,796 @@ +// 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. + +mod common; + +#[path = "pk_compact/merge_engines.rs"] +mod engine_tests; + +#[path = "pk_compact/dynamic_bucket.rs"] +mod dynamic_tests; + +use arrow_array::{Array, Int32Array, Int64Array, RecordBatch, StringArray}; +use common::incremental_helpers::{ + make_batch, make_batch_with_kinds, make_partitioned_batch, memory_table, persist_table_schema, + pk_schema, setup_dirs, write_batch, +}; +use futures::TryStreamExt; +use paimon::spec::{ + CommitKind, DataType, Datum, IntType, PredicateBuilder, Schema, TableSchema, VarCharType, +}; +use paimon::table::{CompactPlan, CompactResult, CompactTask, IncrementalScanMode, Table}; + +async fn table_with_schema(schema: TableSchema) -> Table { + let path = format!("memory:/pk_compact/{}", uuid::Uuid::new_v4()); + let (io, table) = memory_table(&path, schema); + setup_dirs(&io, &path).await; + persist_table_schema(&io, &path, table.schema()).await; + table +} + +async fn table(options: &[(&str, &str)]) -> Table { + table_with_schema(pk_schema(options)).await +} + +async fn read(table: &Table, projection: Option<&[&str]>) -> Vec { + let mut builder = table.new_read_builder(); + if let Some(projection) = projection { + builder.with_projection(projection).unwrap(); + } + let plan = builder.new_scan().plan().await.unwrap(); + builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap() +} + +async fn rows(table: &Table) -> Vec<(i32, i32)> { + let mut rows = Vec::new(); + for batch in read(table, None).await { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + rows.extend((0..batch.num_rows()).map(|i| (ids.value(i), values.value(i)))); + } + rows.sort_unstable(); + rows +} + +async fn sequences(table: &Table) -> Vec<(i32, i64)> { + let mut rows = Vec::new(); + for batch in read(table, Some(&["id", "_SEQUENCE_NUMBER"])).await { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let seqs = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + rows.extend((0..batch.num_rows()).map(|i| (ids.value(i), seqs.value(i)))); + } + rows.sort_unstable(); + rows +} + +async fn execute(table: &Table, plan: &CompactPlan) -> Vec { + let executor = table.new_compact_builder().new_executor().unwrap(); + let mut results = Vec::new(); + for task in plan.tasks() { + let task: CompactTask = serde_json::from_slice(&serde_json::to_vec(task).unwrap()).unwrap(); + results.push(executor.execute(&task).await.unwrap()); + } + serde_json::from_slice(&serde_json::to_vec(&results).unwrap()).unwrap() +} + +#[tokio::test] +async fn pk_compaction_rejects_concurrent_overlapping_positive_level_files() { + for (key, conflicts) in [(1, true), (256, false)] { + let table = table(&[]).await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + // A Java writer created before our input was committed can compact its + // own L0 files locally and publish L1 without replacing our L0 input. + let write = table.new_write_builder(); + let mut writer = write.new_write().unwrap(); + writer + .write_arrow(&[make_batch(vec![key], vec![20])], None) + .await + .unwrap(); + let mut messages = writer.prepare_commit().await.unwrap(); + for file in &mut messages[0].new_files { + file.level = 1; + } + write.new_commit().commit(messages).await.unwrap(); + let result = builder.new_commit().unwrap().commit_compact(results).await; + if conflicts { + assert!(result.unwrap_err().to_string().contains("LSM")); + assert_eq!(rows(&table).await, vec![(1, 20)]); + } else { + result.unwrap(); + assert_eq!(rows(&table).await, vec![(1, 10), (256, 20)]); + } + } +} + +#[tokio::test] +async fn pk_compaction_flushes_small_buffers_without_fragmenting_files() { + for format in ["parquet", "orc", "avro", "row"] { + let table = table(&[ + ("file.format", format), + ("file.compression", "none"), + ("read.batch-size", "1"), + ("write-buffer-size", "1b"), + ("target-file-size", "256mb"), + ("file-index.bitmap.columns", "id"), + ]) + .await; + write_batch( + &table, + &make_batch(vec![1, 2, 3, 4, 5], vec![10, 20, 30, 40, 50]), + ) + .await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.tasks()[0].files().len(), 1); + let before = sequences(&table).await; + let results = execute(&table, &plan).await; + assert_eq!( + results[0].files().len(), + 1, + "{format}: a buffer flush must not close the file" + ); + let file = &results[0].files()[0]; + for (bytes, expected) in [ + (file.key_stats.min_values(), 1), + (file.key_stats.max_values(), 5), + (file.min_key.as_slice(), 1), + (file.max_key.as_slice(), 5), + ] { + let key = paimon::spec::BinaryRow::from_serialized_bytes(bytes).unwrap(); + assert_eq!(key.get_int(0).unwrap(), expected, "{format}"); + } + assert_eq!(file.key_stats.null_counts(), &[Some(0)]); + assert!(file.embedded_index.is_some()); + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap(); + assert_eq!( + rows(&table).await, + vec![(1, 10), (2, 20), (3, 30), (4, 40), (5, 50)] + ); + assert_eq!(sequences(&table).await, before); + let mut reader = table.new_read_builder(); + reader.with_filter( + PredicateBuilder::new(table.schema().fields()) + .equal("id", Datum::Int(3)) + .unwrap(), + ); + let plan = reader.new_scan().plan().await.unwrap(); + let batches: Vec = reader + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 1, + "{format}" + ); + } +} + +#[tokio::test] +async fn pk_compaction_rolls_by_bytes_without_a_row_limit() { + let table = table(&[("target-file-size", "1b"), ("read.batch-size", "1")]).await; + write_batch(&table, &make_batch(vec![1, 2, 3], vec![10, 20, 30])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.tasks()[0].files().len(), 1); + let results = execute(&table, &plan).await; + assert_eq!(results[0].files().len(), 3); + assert!(results[0].files().iter().all(|file| file.row_count == 1)); + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap(); + assert_eq!(rows(&table).await, vec![(1, 10), (2, 20), (3, 30)]); +} + +#[tokio::test] +async fn pk_compaction_merges_versions_and_preserves_sequences_deletes_and_history() { + for format in ["parquet", "orc", "avro", "row"] { + let table = table(&[ + ("file.format", format), + ("file.compression", "none"), + ("changelog-producer", "input"), + ]) + .await; + write_batch(&table, &make_batch(vec![1, 2], vec![10, 20])).await; + write_batch(&table, &make_batch(vec![1, 3], vec![11, 30])).await; + write_batch( + &table, + &make_batch_with_kinds(vec![2, 4], vec![20, 40], vec![3, 0]), + ) + .await; + let expected = vec![(1, 11), (3, 30), (4, 40)]; + assert_eq!(rows(&table).await, expected); + let before_sequences = sequences(&table).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.tasks().len(), 1); + assert_eq!(plan.tasks()[0].files().len(), 3); + let results = execute(&table, &plan).await; + assert_eq!(results[0].files().len(), 1, "{format}"); + let file = &results[0].files()[0]; + assert_eq!(file.level, 1); + assert_eq!(file.delete_row_count, Some(1)); + assert_eq!((file.min_sequence_number, file.max_sequence_number), (2, 5)); + assert_eq!(file.file_source, Some(1)); + let committer = builder.new_commit().unwrap(); + let metrics = committer.commit_compact(results.clone()).await.unwrap(); + assert_eq!( + ( + metrics.files_removed, + metrics.files_added, + metrics.rows_written + ), + (3, 1, 4) + ); + assert_eq!(rows(&table).await, expected, "{format}"); + assert_eq!(sequences(&table).await, before_sequences, "{format}"); + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + assert_eq!(snapshot.commit_kind(), &CommitKind::COMPACT); + assert_eq!(snapshot.total_record_count(), Some(4)); + assert_eq!(snapshot.delta_record_count(), Some(-2)); + let old = table.copy_with_pinned_snapshot(Some( + &table.snapshot_manager().get_snapshot(3).await.unwrap(), + )); + assert_eq!(rows(&old).await, expected); + let reader = table.new_read_builder(); + let events = reader + .new_incremental_scan(IncrementalScanMode::Delta, 0, 4) + .plan_combined_delta() + .await + .unwrap(); + let batches: Vec = reader + .new_read() + .unwrap() + .to_arrow(events.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 6 + ); + // Duplicate filtering must precede validation of an expired read snapshot. + table + .file_io() + .delete_file(&table.snapshot_manager().snapshot_path(3)) + .await + .unwrap(); + committer.commit_compact(results).await.unwrap(); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(4) + ); + assert!(builder + .new_planner() + .unwrap() + .plan() + .await + .unwrap() + .tasks() + .is_empty()); + write_batch(&table, &make_batch(vec![1, 5], vec![12, 50])).await; + assert_eq!(rows(&table).await, vec![(1, 12), (3, 30), (4, 40), (5, 50)]); + assert!(sequences(&table) + .await + .iter() + .any(|&(id, seq)| id == 1 && seq > 5)); + } +} + +#[tokio::test] +async fn pk_compaction_keeps_tombstones_against_later_out_of_order_user_sequences() { + for (order, initial, deleted, stale, fresh) in [ + ("ascending", 100, 150, 120, 160), + ("descending", 100, 50, 80, 40), + ] { + let table = table(&[ + ("sequence.field", "value"), + ("sequence.field.sort-order", order), + ]) + .await; + write_batch(&table, &make_batch(vec![1], vec![initial])).await; + write_batch( + &table, + &make_batch_with_kinds(vec![1], vec![deleted], vec![3]), + ) + .await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + assert_eq!(results[0].files()[0].delete_row_count, Some(1)); + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap(); + assert!(rows(&table).await.is_empty()); + write_batch(&table, &make_batch(vec![1], vec![stale])).await; + assert!(rows(&table).await.is_empty(), "{order}"); + // Compact again with the stale value present; it must still lose. + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + builder + .new_commit() + .unwrap() + .commit_compact(execute(&table, &plan).await) + .await + .unwrap(); + write_batch(&table, &make_batch(vec![1], vec![deleted])).await; + assert_eq!( + rows(&table).await, + vec![(1, deleted)], + "equal user sequence must use the newer system sequence" + ); + write_batch(&table, &make_batch(vec![1], vec![fresh])).await; + assert_eq!(rows(&table).await, vec![(1, fresh)], "{order}"); + } +} + +#[tokio::test] +async fn pk_compaction_includes_large_files_and_whole_buckets_and_upgrades_single_l0() { + let table = table(&[("target-file-size", "1b")]).await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.tasks().len(), 1); + assert_eq!(plan.tasks()[0].files().len(), 1); + builder + .new_commit() + .unwrap() + .commit_compact(execute(&table, &plan).await) + .await + .unwrap(); + for value in 11..16 { + write_batch(&table, &make_batch(vec![1], vec![value])).await; + } + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.tasks().len(), 1); + assert_eq!(plan.tasks()[0].files().len(), 6); + assert!(plan.tasks()[0] + .files() + .iter() + .all(|file| file.file_size > 1)); + builder + .new_commit() + .unwrap() + .commit_compact(execute(&table, &plan).await) + .await + .unwrap(); + assert_eq!(rows(&table).await, vec![(1, 15)]); +} + +#[tokio::test] +async fn pk_compaction_keeps_writes_committed_after_planning_and_rejects_competing_rewrites() { + let table = table(&[]).await; + write_batch(&table, &make_batch(vec![1, 2], vec![10, 20])).await; + write_batch(&table, &make_batch(vec![1, 3], vec![11, 30])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let competitor = table.new_compact_builder(); + let competing = execute( + &table, + &competitor.new_planner().unwrap().plan().await.unwrap(), + ) + .await; + write_batch( + &table, + &make_batch_with_kinds(vec![1, 2], vec![12, 20], vec![0, 3]), + ) + .await; + // Execute after the append, using only the captured bucket state. + let results = execute(&table, &plan).await; + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap(); + assert_eq!(rows(&table).await, vec![(1, 12), (3, 30)]); + assert!(competitor + .new_commit() + .unwrap() + .commit_compact(competing) + .await + .is_err()); + assert_eq!(rows(&table).await, vec![(1, 12), (3, 30)]); +} + +#[tokio::test] +async fn pk_compaction_rolls_ordered_files_and_rebuilds_indexes_with_level_stats() { + let table = table(&[ + ("target-file-row-num", "2"), + ("read.batch-size", "4"), + ("write-buffer-size", "64 B"), + ("file-index.bitmap.columns", "id"), + ("file-index.in-manifest-threshold", "0"), + ("metadata.stats-mode-per-level", "0:none,1:full"), + ]) + .await; + write_batch( + &table, + &make_batch(vec![0, 1, 2, 3, 4], vec![0, 10, 20, 30, 40]), + ) + .await; + write_batch( + &table, + &make_batch_with_kinds(vec![1, 3], vec![11, 30], vec![0, 3]), + ) + .await; + let expected = rows(&table).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(&table, &plan).await; + assert_eq!(results[0].files().len(), 3); + assert!(results[0].files().iter().all(|file| file.level == 1 + && file.row_count <= 2 + && (file.embedded_index.is_some() || !file.extra_files.is_empty()) + && !file.value_stats.null_counts().is_empty())); + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap(); + assert_eq!(rows(&table).await, expected); + // Rolled files at one positive level already form a sorted run. Replanning + // must not continually rewrite that same run on every maintenance pass. + assert!(builder + .new_planner() + .unwrap() + .plan() + .await + .unwrap() + .tasks() + .is_empty()); + for id in 0..6 { + let mut read = table.new_read_builder(); + read.with_filter( + PredicateBuilder::new(table.schema().fields()) + .equal("id", Datum::Int(id)) + .unwrap(), + ); + let plan = read.new_scan().plan().await.unwrap(); + let batches: Vec = read + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + usize::from(id < 5 && id != 3) + ); + } +} + +#[tokio::test] +async fn pk_compaction_refuses_partial_buckets_and_unsupported_layouts() { + for options in [ + vec![ + ("merge-engine", "first-row"), + ("changelog-producer", "lookup"), + ], + vec![("changelog-producer", "full-compaction")], + vec![ + ("merge-engine", "first-row"), + ("pk-clustering-override", "true"), + ("clustering.columns", "value"), + ], + vec![ + ("record-level.expire-time", "1 d"), + ("record-level.time-field", "value"), + ], + ] { + let table = table(&options).await; + assert!(table.new_compact_builder().new_planner().is_err()); + } + let table = table(&[]).await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + write_batch(&table, &make_batch(vec![2], vec![20])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let mut task = serde_json::to_value(&plan.tasks()[0]).unwrap(); + task["files"].as_array_mut().unwrap().pop(); + let task: CompactTask = serde_json::from_value(task).unwrap(); + assert!(builder + .new_executor() + .unwrap() + .execute(&task) + .await + .unwrap_err() + .to_string() + .contains("complete captured bucket")); + let results = execute(&table, &plan).await; + let mut encoded = serde_json::to_value(&results).unwrap(); + encoded[0]["task"]["files"].as_array_mut().unwrap().pop(); + let partial: Vec = serde_json::from_value(encoded).unwrap(); + assert!(builder + .new_commit() + .unwrap() + .commit_compact(partial) + .await + .is_err()); + let mut encoded = serde_json::to_value(&results).unwrap(); + encoded[0]["files"][0]["_LEVEL"] = serde_json::json!(0); + let bad: Vec = serde_json::from_value(encoded).unwrap(); + assert!(builder + .new_commit() + .unwrap() + .commit_compact(bad) + .await + .is_err()); + assert_eq!(rows(&table).await, vec![(1, 10), (2, 20)]); +} + +#[tokio::test] +async fn pk_compaction_groups_composite_keys_by_partition_and_bucket() { + let schema = Schema::builder() + .column("pt", DataType::VarChar(VarCharType::string_type())) + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .primary_key(["pt", "id"]) + .partition_keys(["pt"]) + .option("bucket", "2") + .build() + .unwrap(); + let table = table_with_schema(TableSchema::new(0, &schema)).await; + for version in 0..2 { + for pt in ["a", "b"] { + for bucket in 0..2 { + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer + .write_arrow( + &[make_partitioned_batch( + vec![pt], + vec![bucket], + vec![version * 10 + bucket], + )], + Some(bucket), + ) + .await + .unwrap(); + builder + .new_commit() + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + } + } + } + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.tasks().len(), 4); + assert!(plan.tasks().iter().all(|task| task.files().len() == 2)); + let mut results = execute(&table, &plan).await; + let committer = builder.new_commit().unwrap(); + committer + .commit_compact(results.drain(..2).collect()) + .await + .unwrap(); + committer.commit_compact(results).await.unwrap(); + let mut actual = Vec::new(); + for batch in read(&table, None).await { + let pt = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let id = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let value = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + actual.extend( + (0..batch.num_rows()).map(|i| (pt.value(i).to_string(), id.value(i), value.value(i))), + ); + } + actual.sort(); + assert_eq!( + actual, + vec![ + ("a".into(), 0, 10), + ("a".into(), 1, 11), + ("b".into(), 0, 10), + ("b".into(), 1, 11) + ] + ); +} + +#[tokio::test] +async fn pk_compaction_preserves_a_null_winner_instead_of_an_older_non_null_value() { + let table = table(&[]).await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + let nullable = RecordBatch::try_new( + std::sync::Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", arrow_schema::DataType::Int32, false), + arrow_schema::Field::new("value", arrow_schema::DataType::Int32, true), + ])), + vec![ + std::sync::Arc::new(Int32Array::from(vec![1])), + std::sync::Arc::new(Int32Array::from(vec![None])), + ], + ) + .unwrap(); + write_batch(&table, &nullable).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + builder + .new_commit() + .unwrap() + .commit_compact(execute(&table, &plan).await) + .await + .unwrap(); + let batches = read(&table, None).await; + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 1 + ); + assert!(batches[0].column_by_name("value").unwrap().is_null(0)); +} + +#[tokio::test] +async fn pk_compaction_memory_failure_leaves_source_files_and_releases_reservations() { + let table = table(&[]).await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + write_batch(&table, &make_batch(vec![1], vec![11])).await; + let plan = table + .new_compact_builder() + .new_planner() + .unwrap() + .plan() + .await + .unwrap(); + let resources = paimon::resource::ResourceContext::builder() + .memory_limit(1) + .build() + .unwrap(); + let executor = table + .new_compact_builder() + .with_resources(resources.clone()) + .new_executor() + .unwrap(); + let before = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(); + let error = executor.execute(&plan.tasks()[0]).await.unwrap_err(); + assert!( + matches!(error, paimon::Error::ResourceExhausted { .. }), + "{error}" + ); + assert_eq!(resources.metrics().reserved_memory_bytes, 0); + assert_eq!( + table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(), + before + ); + assert_eq!(rows(&table).await, vec![(1, 11)]); +} + +#[tokio::test] +async fn pk_compaction_read_failure_removes_already_flushed_unsubmitted_files() { + let table = table(&[("write-buffer-size", "1 B"), ("read.batch-size", "1")]).await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + write_batch(&table, &make_batch(vec![2], vec![20])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let read = table.new_read_builder().new_scan().plan().await.unwrap(); + let bucket = read.splits()[0].bucket_path(); + let second = plan.tasks()[0] + .files() + .iter() + .find(|file| file.max_sequence_number == 1) + .unwrap(); + table + .file_io() + .delete_file(&format!("{bucket}/{}", second.file_name)) + .await + .unwrap(); + let before = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(); + assert!(builder + .new_executor() + .unwrap() + .execute(&plan.tasks()[0]) + .await + .is_err()); + assert_eq!( + table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(), + before + ); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + Some(2) + ); +} diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 6c3ec41ca..6f2ec7ad3 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -49,6 +49,7 @@ nav: - Home: index.md - Getting Started: getting-started.md - Batch Incremental Reading: incremental-reading.md + - Table Compaction: compaction.md - SQL Integration: sql.md - Performance: benchmark.md - C Integration: c-binding.md diff --git a/docs/src/compaction.md b/docs/src/compaction.md new file mode 100644 index 000000000..712621700 --- /dev/null +++ b/docs/src/compaction.md @@ -0,0 +1,194 @@ + + +# Table Compaction + +Rust supports small-file compaction for ordinary append tables and full-bucket +compaction for primary-key tables with fixed or dynamic buckets. `CompactBuilder` creates +a `CompactPlanner`, a `CompactExecutor`, and a `TableCommit`. Planning and +execution do not publish snapshots; `commit_compact` atomically publishes the +selected results in a `COMPACT` snapshot. + +```rust +let builder = table.new_compact_builder(); +let planner = builder.new_planner()?; +let executor = builder.new_executor()?; +let committer = builder.new_commit()?; + +let plan = planner.plan().await?; +let mut results = Vec::new(); +for task in plan.tasks() { + results.push(executor.execute(task).await?); +} +let metrics = committer.commit_compact(results).await?; +println!("Compacted {} files into {}", metrics.files_removed, metrics.files_added); +``` + +Compaction uses the table's `target-file-size` option for its output size target. +For unaware append tables (`bucket=-1`), task selection follows the batch packing rules in Java's +[`AppendCompactCoordinator`](https://github.com/apache/paimon/blob/master/paimon-core/src/main/java/org/apache/paimon/append/AppendCompactCoordinator.java): + +- Candidate files are strictly smaller than `target-file-size` multiplied by + `compaction.small-file-ratio` (default 0.7). This avoids repeatedly rewriting + output files that are slightly below the target after compression. +- Within each partition and bucket, candidates are sorted by file size. + Each file contributes `file_size + source.split.open-file-cost` to the task + weight. The open cost defaults to 4 MiB. +- A task is emitted once it contains at least two files and its weight reaches + twice `target-file-size`. A remaining group is emitted only when it has at + least `compaction.min.file-num` files (default 5, minimum 2). + +There is no per-task file-count cap. An underfilled tail remains available for +planning against a later snapshot. Each plan independently captures a snapshot; +the API does not maintain a streaming coordinator's aging state. + +Unordered bucketed append tables (`bucket>0`, `bucket-append-ordered=false`) +use Java's `BucketedAppendCompactManager` best-effort selection instead: +files are visited in minimum-sequence order, and a window is selected once it +reaches `compaction.min.file-num`. Before that count is reached, a window whose +physical bytes reach twice `target-file-size` drops its first file and keeps +sliding. The unaware-table small-file filter and open cost do not apply. +Ordered bucketed append compaction is unsupported: it requires sequence counters +that the ordinary Rust append writer does not yet maintain. + +Ordinary read splits use a separate strategy, matching Java's +[`AppendOnlySplitGenerator`](https://github.com/apache/paimon/blob/master/paimon-core/src/main/java/org/apache/paimon/table/source/AppendOnlySplitGenerator.java): +sort by minimum sequence number, then pack up to `source.split.target-size` +using `max(file_size, source.split.open-file-cost)` per file. Read split boundaries +do not determine compaction tasks, and changing `source.split.target-size` does +not change compaction grouping. + +File sizes are approximate targets; normal writer rolling and +`target-file-row-num` still apply. An empty plan or result list publishes no +snapshot. Use `with_partition_filter(predicate)` to restrict planning to +partitions and `with_resources(context)` to share read/write memory budgets. + +`CompactPlan`, `CompactTask` and `CompactResult` support Serde serialization. +Workers create their own executor with a table handle and local resources; +the task fixes the input snapshot, schema and output settings. Applications +can schedule independent tasks concurrently and submit completed subsets. +Results from different plans, duplicate tasks and overlapping inputs are +rejected. Executors stream input batches instead of collecting the whole task. + +Concurrent appends can commit alongside compaction. A rewrite or overwrite of +an input file, a metadata change to that file, or a table schema change causes +the compaction commit to fail. Create a new plan after such a conflict. +PK commits also reject overlapping key ranges within a positive level, including +ranges introduced by a concurrent compaction. Overlapping L0 appends are allowed. +If the latest snapshot is below the plan's captured snapshot after a rollback, +commit rejects the uncommitted plan; create a new plan. +Retrying the same plan and task subset with the same `commit_user` is +idempotent, including when result order changes. Recovery searches only snapshots +created after the plan's captured snapshot. To retry after restarting +the coordinator, configure `with_commit_user(...)` with the original value +and retain the serialized results. A commit error can mean publication is +uncertain: retain submitted files rather than deleting them on error. + +Append compaction preserves duplicate and null rows. Original files remain readable +through old snapshots, and append incremental scans skip `COMPACT` snapshots. +The removed-file metric counts manifest removals, not physical deletion; +snapshot expiration and storage reclamation are separate maintenance work. + +## Primary-key compaction + +The same API compacts primary-key tables with `deduplicate`, +`first-row`, `partial-update` or `aggregation` merge engines. Changelog producers +`none` and `input` are supported; `first-row` requires `none`. +The planner includes every active file of each selected partition/bucket, +including large files. Append size and count selection options do not split PK +buckets. A single level-0 file is also eligible. A bucket whose files already form one sorted +run at a single positive level needs no rewrite, including when the previous +compaction rolled multiple output files. Output is key ordered at the highest input level, or level 1 +when all inputs are level 0. + +Compaction uses each engine's existing merge function and persists its merge +state, including original sequence metadata and synthesized DELETE records: + +- `deduplicate` retains the winning record, including retracts that must still + suppress later out-of-order values. +- `first-row` retains the earliest record. Retracts are rejected unless + `ignore-delete=true`. Its normal batch reader still skips uncompacted L0 + files; compaction makes the selected records visible above L0. +- `partial-update` combines sparse columns using the configured sequence + groups and field aggregators. Delete payloads and group sequence values are + preserved so subsequent sparse updates can continue from that state. +- `aggregation` writes the engine's field accumulators. Singleton records + remain unreduced, including retracts; compound DELETE states are preserved + rather than being discarded as they are by the query reader. + +The existing engine rules for ignored deletes, unsupported retracts and +whole-record deletion also apply during compaction. If all records are +legitimately ignored, compaction can remove the selected files without +creating an output file. Key bounds, sequence bounds, delete counts and file +indexes are regenerated, using metadata statistics options for the output level. +Compound `partial-update` DELETE states with field aggregators are rejected. +For both reducer engines, compound groups containing retracts are also rejected +when they use `first_value`, `first_non_null_value`, `collect`, `listagg`, +`nested_update` or `nested_partial_update`. These states may carry initialization +flags or raw DELETE payloads that cannot be reconstructed from one output row, +even when that row is an INSERT. Execution fails without publishing or deleting +the source files. +Partial-update retracts ignored by `ignore-delete=true` do not trigger this +restriction. + +The engine's logical merged state remains unchanged. Physical record counts and +`rows_written` can decrease as older versions are removed; the latter includes +retained retract records. Append and input-changelog incremental reads do not +emit new logical events for a compaction snapshot. + +PK rewrites stream ordered physical KV batches into the same low-level file +writer used by ordinary PK writes. Sequence allocation and merge buffering stay +in the ordinary write path. Key bounds, sequence bounds, statistics and indexes +are collected incrementally without retaining all rows of an output file. +`write-buffer-size` flushes format buffers. Parquet, ORC, Avro and ROW keep the +same output file open across these flushes. Formats that retain their pending +rows after flushing (Mosaic/Vortex) close the file at the buffer threshold to +preserve bounded memory. Files also roll when the format's approximate byte +count reaches `target-file-size` or the row count reaches `target-file-row-num`. +Output chunks are materialized at row limits to avoid retaining and repeatedly +charging the entire source batch for small slices. +The merge currently supports +up to 256 overlapping sorted-run inputs, matching the query reader's limit. +There is no multipass merge yet; too many overlapping runs fail without +publishing a snapshot. The captured snapshot must remain available until each +task subset is successfully committed; retries of an already committed subset remain +idempotent after its read snapshot expires. + +## Dynamic buckets + +Dynamic bucket compaction (`bucket=-1`) rewrites each existing partition/bucket +without changing key assignments. HASH index files and their complete assignment +history, including deleted keys, are retained. A concurrent append may replace +the HASH index; compaction carries forward the latest manifest at commit time. +Cross-partition dynamic buckets are also supported: subsequent writers rebuild +their key-location index from the compacted data as usual. Compaction does not +rebalance buckets or reclaim HASH index capacity. + +## Current scope + +Compaction-generated changelogs (`full-compaction` / `lookup`) are unsupported. +Row tracking, Data Evolution, deletion vectors, query authorization, rowkind +fields, postpone buckets and dedicated blob/vector storage are unsupported. +Ordered bucketed append, append incremental clustering, primary-key clustering +override (`pk-clustering-override`) and record-level expiration +(`record-level.expire-time`) are also unsupported. +Snapshot index manifests may contain only dynamic-bucket HASH indexes. +Embedded and sidecar file indexes are regenerated according to the captured +table options. Language bindings and +background compaction scheduling are not yet provided.