From 852616499db8330001cdc133da62bdc20a0a2d9b Mon Sep 17 00:00:00 2001 From: Zihan Dai Date: Tue, 3 Nov 2026 22:28:40 +1100 Subject: [PATCH] fix(table): move rows to their new bucket when UPDATE changes the bucket key A fixed-bucket append table places each row by hashing its bucket-key columns, and scans prune buckets by that hash. The copy-on-write writer used by UPDATE and MERGE INTO rewrote every touched file back into the file's original bucket, so a row whose bucket key was updated stayed in the bucket of its old value. Any query filtering on the new value read only the new value's bucket and missed the row. When the update touches a bucket-key column, compute each rewritten row's bucket with the table's bucket function and write the rows of each bucket separately, like the regular write path and Java's Spark writer. --- .../datafusion/tests/append_merge_into.rs | 52 +++++++++ .../datafusion/tests/update_tests.rs | 53 +++++++++ crates/paimon/src/table/cow_writer.rs | 101 +++++++++++++++--- 3 files changed, 191 insertions(+), 15 deletions(-) diff --git a/crates/integrations/datafusion/tests/append_merge_into.rs b/crates/integrations/datafusion/tests/append_merge_into.rs index 3cd3c97dd..9bfaa7aed 100644 --- a/crates/integrations/datafusion/tests/append_merge_into.rs +++ b/crates/integrations/datafusion/tests/append_merge_into.rs @@ -1191,3 +1191,55 @@ async fn test_partitioned_successive_merges() { vec![(1, 100, 1), (2, 20, 1), (3, 300, 2), (4, 40, 2),] ); } + +/// Rows of a fixed-bucket table are placed by hashing `bucket-key`, and a filter on +/// that column reads only the matching bucket. A matched update of the bucket key +/// must move the row, or such queries no longer find it. +#[tokio::test] +async fn test_update_bucket_key_moves_row_to_new_bucket() { + let (_tmp, sql_context) = setup( + "CREATE TABLE paimon.test_db.target (a INT, b INT, c VARCHAR) \ + WITH ('bucket' = '4', 'bucket-key' = 'c')", + ) + .await; + exec( + &sql_context, + "INSERT INTO paimon.test_db.target VALUES \ + (1, 10, 'alice'), (2, 20, 'bob'), (3, 30, 'carol'), (4, 40, 'dave')", + ) + .await; + exec( + &sql_context, + "CREATE TEMPORARY TABLE paimon.test_db.source (a INT, c VARCHAR) AS SELECT * FROM (VALUES (1, 'erin'), (3, 'frank')) AS t(a, c)", + ) + .await; + + exec( + &sql_context, + "MERGE INTO paimon.test_db.target t \ + USING paimon.test_db.source s ON t.a = s.a \ + WHEN MATCHED THEN UPDATE SET c = s.c", + ) + .await; + + let batches = sql_context + .sql("SELECT a, b, c FROM paimon.test_db.target WHERE c IN ('erin', 'frank') ORDER BY a") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + collect_int_int_str(&batches), + vec![(1, 10, "erin".into()), (3, 30, "frank".into())] + ); + assert_eq!( + query_abc(&sql_context).await, + vec![ + (1, 10, "erin".into()), + (2, 20, "bob".into()), + (3, 30, "frank".into()), + (4, 40, "dave".into()), + ] + ); +} diff --git a/crates/integrations/datafusion/tests/update_tests.rs b/crates/integrations/datafusion/tests/update_tests.rs index 2c451d8c8..99c75c46a 100644 --- a/crates/integrations/datafusion/tests/update_tests.rs +++ b/crates/integrations/datafusion/tests/update_tests.rs @@ -458,6 +458,59 @@ async fn test_update_with_null_values() { drop(tmp); } +// ======================= Fixed-bucket table ======================= + +/// Rows of a fixed-bucket table are placed by hashing `bucket-key`, and a filter on +/// that column reads only the matching bucket. An updated bucket key must move the +/// row, or such queries no longer find it. +#[tokio::test] +async fn test_update_bucket_key_moves_row_to_new_bucket() { + let (tmp, catalog) = create_test_env(); + let sql_context = create_sql_context(catalog).await; + sql_context + .sql("CREATE SCHEMA paimon.test_db") + .await + .unwrap(); + sql_context + .sql( + "CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, age INT) \ + WITH ('bucket' = '4', 'bucket-key' = 'name')", + ) + .await + .unwrap(); + exec( + &sql_context, + "INSERT INTO paimon.test_db.t VALUES \ + (1, 'alice', 10), (2, 'bob', 20), (3, 'carol', 30), (4, 'dave', 40)", + ) + .await; + + exec( + &sql_context, + "UPDATE paimon.test_db.t SET name = 'erin' WHERE id = 1", + ) + .await; + + assert_eq!( + query_int_str_int( + &sql_context, + "SELECT id, name, age FROM paimon.test_db.t WHERE name = 'erin'", + ) + .await, + vec![(1, "erin".into(), 10)] + ); + assert_eq!( + query(&sql_context).await, + vec![ + (1, "erin".into(), 10), + (2, "bob".into(), 20), + (3, "carol".into(), 30), + (4, "dave".into(), 40), + ] + ); + drop(tmp); +} + // ======================= Empty table ======================= #[tokio::test] diff --git a/crates/paimon/src/table/cow_writer.rs b/crates/paimon/src/table/cow_writer.rs index a0ff4304c..aac31e5ae 100644 --- a/crates/paimon/src/table/cow_writer.rs +++ b/crates/paimon/src/table/cow_writer.rs @@ -21,7 +21,10 @@ //! Unaffected files remain untouched. The writer produces [`CommitMessage`]s with //! `deleted_files` (old files) and `new_files` (rewritten files) for atomic commit. -use crate::spec::{BinaryRow, CoreOptions, DataFileMeta, PartitionComputer}; +use crate::spec::{ + BinaryRow, BucketFunctionType, CoreOptions, DataField, DataFileMeta, PartitionComputer, +}; +use crate::table::bucket_function::batch_bucket_ids; use crate::table::commit_message::CommitMessage; use crate::table::data_file_writer::DataFileWriter; use crate::table::partition_filter::PartitionFilter; @@ -31,7 +34,7 @@ use crate::Result; use arrow_array::RecordBatch; use arrow_select::concat::concat_batches; use futures::TryStreamExt; -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; /// Metadata for a single data file in the target table. pub struct FileInfo { @@ -244,6 +247,25 @@ impl CopyOnWriteMergeWriter { let schema_id = schema.id(); let write_fields = schema.fields().to_vec(); + // A fixed-bucket append table places each row by hashing its `bucket-key`, + // and scans prune buckets by that hash. A row whose bucket-key column is + // updated must move to its new bucket, as Java's writer does; left in the + // old bucket, it is skipped by every query that filters on the new value. + let total_buckets = core_options.bucket(); + let rebucket = match core_options.bucket_key() { + Some(keys) + if total_buckets > 1 + && keys.iter().any(|key| self.update_columns.contains(key)) => + { + let key_indices: Vec = keys + .iter() + .filter_map(|key| write_fields.iter().position(|field| field.name() == key)) + .collect(); + Some((key_indices, core_options.bucket_function_type()?)) + } + _ => None, + }; + let update_columns = &self.update_columns; let update_batches = &self.update_batches; let file_index = &self.file_index; @@ -253,6 +275,7 @@ impl CopyOnWriteMergeWriter { let write_fields = &write_fields; let file_compression = file_compression.as_str(); let file_format = file_format.as_str(); + let rebucket = rebucket.as_ref(); // Process each affected file in parallel let rewrite_futures: Vec<_> = self @@ -303,12 +326,27 @@ impl CopyOnWriteMergeWriter { let deleted_file = file_info.file_meta.clone(); - let new_files = if rewritten.num_rows() > 0 { + let groups = match rebucket { + Some((key_indices, function)) => split_by_bucket( + &rewritten, + key_indices, + write_fields, + *function, + total_buckets, + )?, + None => vec![(file_info.bucket, rewritten)], + }; + + let mut new_files = Vec::with_capacity(groups.len()); + for (bucket, batch) in groups { + if batch.num_rows() == 0 { + continue; + } let mut writer = DataFileWriter::new( table.file_io().clone(), table.location().to_string(), - partition_path, - file_info.bucket, + partition_path.clone(), + bucket, schema_id, target_file_size, file_compression.to_string(), @@ -322,11 +360,9 @@ impl CopyOnWriteMergeWriter { None, )? .with_target_file_row_num(target_file_row_num); - writer.write(&rewritten).await?; - writer.prepare_commit().await? - } else { - vec![] - }; + writer.write(&batch).await?; + new_files.push((bucket, writer.prepare_commit().await?)); + } Ok(Some(( file_info.partition.clone(), @@ -346,11 +382,18 @@ impl CopyOnWriteMergeWriter { for result in results.into_iter().flatten() { let (partition, bucket, deleted_file, new_files) = result; - let entry = grouped - .entry((partition, bucket)) - .or_insert_with(|| (Vec::new(), Vec::new())); - entry.0.push(deleted_file); - entry.1.extend(new_files); + grouped + .entry((partition.clone(), bucket)) + .or_default() + .0 + .push(deleted_file); + for (new_bucket, files) in new_files { + grouped + .entry((partition.clone(), new_bucket)) + .or_default() + .1 + .extend(files); + } } let mut messages = Vec::new(); @@ -363,6 +406,34 @@ impl CopyOnWriteMergeWriter { } } +/// Split `batch` by the fixed bucket each row hashes to, keeping row order within a bucket. +fn split_by_bucket( + batch: &RecordBatch, + key_indices: &[usize], + fields: &[DataField], + function: BucketFunctionType, + total_buckets: i32, +) -> Result> { + let buckets = batch_bucket_ids(batch, key_indices, fields, function, total_buckets)?; + let mut rows_by_bucket: BTreeMap> = BTreeMap::new(); + for (row, bucket) in buckets.into_iter().enumerate() { + rows_by_bucket.entry(bucket).or_default().push(row as u32); + } + rows_by_bucket + .into_iter() + .map(|(bucket, rows)| { + let rows = arrow_array::UInt32Array::from(rows); + let taken = arrow_select::take::take_record_batch(batch, &rows).map_err(|e| { + crate::Error::DataInvalid { + message: format!("Failed to split rows by bucket: {e}"), + source: None, + } + })?; + Ok((bucket, taken)) + }) + .collect() +} + /// Apply row operations (update/delete) to a batch, producing a new batch. fn apply_operations( original: &RecordBatch,