fix(table): move rows to their new bucket when UPDATE changes the bucket key

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()),
+        ]
+    );
+}

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]

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<usize> = 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<Vec<(i32, RecordBatch)>> {
+    let buckets = batch_bucket_ids(batch, key_indices, fields, function, total_buckets)?;
+    let mut rows_by_bucket: BTreeMap<i32, Vec<u32>> = 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,