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,