diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index eda7a27c6..7802dfa3a 100644 --- a/crates/paimon/src/table/kv_file_reader.rs +++ b/crates/paimon/src/table/kv_file_reader.rs @@ -39,14 +39,14 @@ use crate::file_index::file_index_result::FileIndexResult; use crate::io::FileIO; use crate::spec::{ is_variant_extraction_row, BigIntType, CoreOptions, DataField, DataFileMeta, - DataType as PaimonDataType, MergeEngine, PartialUpdateConfig, Predicate, TinyIntType, + DataType as PaimonDataType, MergeEngine, PartialUpdateConfig, Predicate, RowKind, TinyIntType, KEY_FIELD_ID_START, SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, VALUE_KIND_FIELD_ID, VALUE_KIND_FIELD_NAME, }; use crate::table::schema_manager::SchemaManager; use crate::table::ArrowRecordBatchStream; use crate::{DataSplit, DeletionFile, Error}; -use arrow_array::{RecordBatch, RecordBatchOptions}; +use arrow_array::{Array, BooleanArray, Int8Array, RecordBatch, RecordBatchOptions}; use async_stream::try_stream; use futures::StreamExt; @@ -184,6 +184,43 @@ fn ensure_merge_input_limit(input_stream_count: usize, limit: Option) -> Ok(()) } +/// Handles retract rows of an unmerged single run the way the merge path +/// handles a lone version of a key: first-row without `ignore-delete` rejects +/// them, otherwise they are omitted. +fn drop_lone_retracts( + batch: RecordBatch, + value_kind_index: usize, + reject: bool, +) -> crate::Result { + // The merge cursor reads a missing or NULL `_VALUE_KIND` as INSERT. + let Some(kinds) = batch + .column(value_kind_index) + .as_any() + .downcast_ref::() + else { + return Ok(batch); + }; + if kinds.null_count() == 0 && kinds.values().iter().all(|&kind| kind == 0 || kind == 2) { + return Ok(batch); + } + let mut keep = Vec::with_capacity(kinds.len()); + for kind in kinds.iter() { + let is_add = RowKind::from_value(kind.unwrap_or(0))?.is_add(); + if !is_add && reject { + return Err(Error::Unsupported { + message: "merge-engine=first-row does not support DELETE or UPDATE_BEFORE rows; set ignore-delete=true to ignore them".to_string(), + }); + } + keep.push(is_add); + } + arrow_select::filter::filter_record_batch(&batch, &BooleanArray::from(keep)).map_err(|e| { + Error::DataInvalid { + message: format!("Failed to drop retract rows: {e}"), + source: Some(Box::new(e)), + } + }) +} + /// Java's `KeyValueFieldsExtractor` builds the physical file layout from the /// schema that wrote the file, including its historical key names and types. pub(super) fn key_value_fields_extractor( @@ -397,15 +434,28 @@ impl KeyValueFileReader { } pub fn read(self, data_splits: &[DataSplit]) -> crate::Result { - self.read_with_merge_function(data_splits, Self::new_merge_function) + self.read_merged(data_splits, Self::new_merge_function, true) } + /// Always merges, so a custom merge function (e.g. the audit log, which + /// keeps retract rows) sees every row. pub(super) fn read_with_merge_function( self, data_splits: &[DataSplit], merge_function: impl Fn(&KeyValueReadConfig, &[DataField]) -> crate::Result> + Send + 'static, + ) -> crate::Result { + self.read_merged(data_splits, merge_function, false) + } + + fn read_merged( + self, + data_splits: &[DataSplit], + merge_function: impl Fn(&KeyValueReadConfig, &[DataField]) -> crate::Result> + + Send + + 'static, + skip_single_runs: bool, ) -> crate::Result { // A projected `_ROW_ID` is synthesized as all-nulls here, so the residual // would silently drop every row rather than hit its missing-column guard. @@ -711,6 +761,42 @@ impl KeyValueFileReader { .map(|budget| Arc::new(budget.without_prefetch())) }; let mut file_streams: Vec = Vec::new(); + // One sorted run never repeats a key, so under these engines + // every key is a lone version and needs no merge. Retract + // rows are still dropped per batch below: older Rust writers + // recorded `delete_row_count = 0` for files holding them. + let mut skip_merge = skip_single_runs + && input_stream_count == 1 + && matches!( + config.merge_engine, + MergeEngine::Deduplicate | MergeEngine::FirstRow + ); + // A file written under another engine may hold several + // versions of a key (the aggregation writer kept them before + // #839), and ALTER can switch the engine afterwards. + if skip_merge { + 'files: for run in &merge_group { + for file in &run.files { + let schema_id = file.file.schema_id; + if schema_id == table_schema_id { + continue; + } + let schema = + config.schema_manager.schema(schema_id).await?; + let engine = + CoreOptions::new(schema.options()).merge_engine()?; + if !matches!( + engine, + MergeEngine::Deduplicate | MergeEngine::FirstRow + ) { + skip_merge = false; + break 'files; + } + } + } + } + let reject_retracts = config.merge_engine == MergeEngine::FirstRow + && !CoreOptions::new(&config.table_options).ignore_delete(); for MergeRun { files } in merge_group { let reader = DataFileReader::new( @@ -814,25 +900,43 @@ impl KeyValueFileReader { file_streams.push(run_stream); } - // Always go through sort-merge even for a single file: files - // written before the writer merged key groups at flush may - // still contain duplicate keys. - let mut merge_stream = SortMergeReaderBuilder::new( - file_streams, - internal_schema.clone(), - key_indices.clone(), - seq_index, - value_kind_index, - user_sequence_indices.clone(), - value_indices.clone(), - merge_output_schema.clone(), - merge_function(&config, &merge_output_fields)?, - ) - .with_user_sequence_descending( - !CoreOptions::new(&config.table_options) - .sequence_field_sort_order_is_ascending(), - ) - .build()?; + let mut merge_stream = if skip_merge { + let input = file_streams.pop().expect("one input stream"); + let columns: Vec = + key_indices.iter().chain(&value_indices).copied().collect(); + let schema = merge_output_schema.clone(); + input + .map(move |batch| { + let batch = + drop_lone_retracts(batch?, value_kind_index, reject_retracts)?; + let columns = + columns.iter().map(|&i| batch.column(i).clone()).collect(); + RecordBatch::try_new(schema.clone(), columns).map_err(|e| { + Error::UnexpectedError { + message: format!("Failed to project unmerged batch: {e}"), + source: Some(Box::new(e)), + } + }) + }) + .boxed() + } else { + SortMergeReaderBuilder::new( + file_streams, + internal_schema.clone(), + key_indices.clone(), + seq_index, + value_kind_index, + user_sequence_indices.clone(), + value_indices.clone(), + merge_output_schema.clone(), + merge_function(&config, &merge_output_fields)?, + ) + .with_user_sequence_descending( + !CoreOptions::new(&config.table_options) + .sequence_field_sort_order_is_ascending(), + ) + .build()? + }; while let Some(batch) = merge_stream.next().await { let batch = batch?; @@ -2011,7 +2115,9 @@ mod tests { .await .expect("the first output batch") .expect("a late deletion vector must be loaded lazily"); - assert_eq!(first.num_rows(), 128); + // These delete-free files form one run, which streams without the + // merge, so the first batch is one input batch. + assert!(first.num_rows() > 0); let err = stream.try_collect::>().await.unwrap_err(); assert!( @@ -2693,6 +2799,242 @@ mod tests { ); } + /// Writes one KV file with the given keys and row kinds, recorded with + /// `delete_row_count = Some(0)` as Rust writers did before #318. + async fn write_legacy_kv_file( + file_io: &FileIO, + table_path: &str, + file_name: &str, + rows: &[(i32, i8)], + ) -> DataFileMeta { + let schema = crate::arrow::build_target_arrow_schema(&[ + DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::new()), + ), + DataField::new( + VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME.to_string(), + DataType::TinyInt(TinyIntType::new()), + ), + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new(1, "value".to_string(), DataType::Int(IntType::new())), + ]) + .unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from_value(0, rows.len())), + Arc::new(Int8Array::from_iter_values(rows.iter().map(|row| row.1))), + Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.0))), + Arc::new(Int32Array::from_value(0, rows.len())), + ], + ) + .unwrap(); + let mut bytes = Vec::new(); + { + let mut writer = AsyncArrowWriter::try_new(&mut bytes, schema, None).unwrap(); + writer.write(&batch).await.unwrap(); + writer.close().await.unwrap(); + } + let bucket_path = format!("{table_path}/bucket-0"); + file_io.mkdirs(&format!("{bucket_path}/")).await.unwrap(); + file_io + .new_output(&format!("{bucket_path}/{file_name}")) + .unwrap() + .write(bytes::Bytes::from(bytes.clone())) + .await + .unwrap(); + + let mut file = dummy_data_file(file_name.to_string()); + file.file_size = bytes.len() as i64; + file.row_count = rows.len() as i64; + file.min_key = int_key(rows[0].0); + file.max_key = int_key(rows[rows.len() - 1].0); + file + } + + #[tokio::test] + async fn single_run_read_drops_retracts_hidden_by_legacy_metadata() { + let file_io = test_file_io(); + let adds = [(1, 0), (2, 0)]; + let retract = [(3, 3)]; + for (options, expected) in [ + (vec![], Some(vec![1, 2])), + (vec![("merge-engine", "first-row")], None), + ( + vec![("merge-engine", "first-row"), ("ignore-delete", "true")], + Some(vec![1, 2]), + ), + ] { + let table_path = "memory:/kv_legacy_retract"; + let table = pk_table(&file_io, table_path, &options); + let a = write_legacy_kv_file(&file_io, table_path, "a.parquet", &adds).await; + let b = write_legacy_kv_file(&file_io, table_path, "b.parquet", &retract).await; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![a, b]) + .build() + .unwrap(); + let core_options = table.schema().core_options(); + let result = KeyValueFileReader::new( + table.file_io().clone(), + KeyValueReadConfig { + table_name: table.identifier().full_name(), + table_options: table.schema().options().clone(), + schema_manager: table.schema_manager().clone(), + table_schema_id: table.schema().id(), + table_fields: table.schema().fields().to_vec(), + read_type: table.schema().fields().to_vec(), + predicates: Vec::new(), + primary_keys: table.schema().trimmed_primary_keys(), + table_primary_keys: table.schema().primary_keys().to_vec(), + merge_engine: core_options.merge_engine().unwrap(), + sequence_fields: Vec::new(), + read_batch_size: core_options.read_batch_size().unwrap(), + merge_splits: false, + max_merge_input_streams: None, + parquet_read_budget: None, + mosaic_prefetch: MosaicPrefetchOptions::default(), + }, + ) + .read(&[split]) + .unwrap() + .try_collect::>() + .await; + match expected { + Some(ids) => assert_eq!(int_column(&result.unwrap(), "id"), ids, "{options:?}"), + None => assert!(result.is_err(), "{options:?}"), + } + } + } + + /// Writes one KV file with `(id, sequence, value)` rows, under schema 0. + async fn write_kv_versions( + file_io: &FileIO, + table_path: &str, + file_name: &str, + rows: &[(i32, i64, i32)], + ) -> DataFileMeta { + let schema = crate::arrow::build_target_arrow_schema(&[ + DataField::new( + SEQUENCE_NUMBER_FIELD_ID, + SEQUENCE_NUMBER_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::new()), + ), + DataField::new( + VALUE_KIND_FIELD_ID, + VALUE_KIND_FIELD_NAME.to_string(), + DataType::TinyInt(TinyIntType::new()), + ), + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new(1, "value".to_string(), DataType::Int(IntType::new())), + ]) + .unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from_iter_values(rows.iter().map(|row| row.1))), + Arc::new(Int8Array::from_value(0, rows.len())), + Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.0))), + Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.2))), + ], + ) + .unwrap(); + let mut bytes = Vec::new(); + { + let mut writer = AsyncArrowWriter::try_new(&mut bytes, schema, None).unwrap(); + writer.write(&batch).await.unwrap(); + writer.close().await.unwrap(); + } + let bucket_path = format!("{table_path}/bucket-0"); + file_io.mkdirs(&format!("{bucket_path}/")).await.unwrap(); + file_io + .new_output(&format!("{bucket_path}/{file_name}")) + .unwrap() + .write(bytes::Bytes::from(bytes.clone())) + .await + .unwrap(); + + let mut file = dummy_data_file(file_name.to_string()); + file.file_size = bytes.len() as i64; + file.row_count = rows.len() as i64; + file.min_key = int_key(rows[0].0); + file.max_key = int_key(rows[rows.len() - 1].0); + file.min_sequence_number = rows.iter().map(|row| row.1).min().unwrap(); + file.max_sequence_number = rows.iter().map(|row| row.1).max().unwrap(); + file + } + + #[tokio::test] + async fn single_run_read_merges_files_written_under_aggregation() { + // Before #839 the aggregation writer kept several versions of a key in + // one file. After ALTER to deduplicate, such a file still needs merging. + let file_io = test_file_io(); + let table_path = "memory:/kv_engine_change"; + let old = pk_table( + &file_io, + table_path, + &[ + ("merge-engine", "aggregation"), + ("fields.value.aggregate-function", "sum"), + ], + ); + write_schema_file(&old, old.schema()).await; + let mut json = serde_json::to_value(old.schema()).unwrap(); + json["id"] = 1.into(); + json["options"]["merge-engine"] = "deduplicate".into(); + let schema: TableSchema = serde_json::from_value(json).unwrap(); + write_schema_file(&old, &schema).await; + let table = old.copy_with_resolved_schema(schema, "main").unwrap(); + let a = + write_kv_versions(&file_io, table_path, "a.parquet", &[(1, 0, 10), (1, 1, 11)]).await; + let b = write_kv_versions(&file_io, table_path, "b.parquet", &[(2, 2, 20)]).await; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![a, b]) + .build() + .unwrap(); + let core_options = table.schema().core_options(); + let batches = KeyValueFileReader::new( + table.file_io().clone(), + KeyValueReadConfig { + table_name: table.identifier().full_name(), + table_options: table.schema().options().clone(), + schema_manager: table.schema_manager().clone(), + table_schema_id: table.schema().id(), + table_fields: table.schema().fields().to_vec(), + read_type: table.schema().fields().to_vec(), + predicates: Vec::new(), + primary_keys: table.schema().trimmed_primary_keys(), + table_primary_keys: table.schema().primary_keys().to_vec(), + merge_engine: core_options.merge_engine().unwrap(), + sequence_fields: Vec::new(), + read_batch_size: core_options.read_batch_size().unwrap(), + merge_splits: false, + max_merge_input_streams: None, + parquet_read_budget: None, + mosaic_prefetch: MosaicPrefetchOptions::default(), + }, + ) + .read(&[split]) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(int_column(&batches, "id"), vec![1, 2]); + assert_eq!(int_column(&batches, "value"), vec![11, 20]); + } + #[tokio::test] async fn sorted_run_read_matches_per_file_fan_out() { let file_io = test_file_io(); diff --git a/crates/paimon/tests/pk_single_run_read_test.rs b/crates/paimon/tests/pk_single_run_read_test.rs new file mode 100644 index 000000000..2dca3d568 --- /dev/null +++ b/crates/paimon/tests/pk_single_run_read_test.rs @@ -0,0 +1,129 @@ +// 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. + +//! Primary-key reads of a split whose sections each hold one sorted run. + +mod common; + +use std::sync::Arc; + +use arrow_array::{ArrayRef, Int32Array, RecordBatch}; +use common::incremental_helpers::{memory_table, persist_table_schema, setup_dirs, write_batch}; +use futures::TryStreamExt; +use paimon::spec::{DataType, Datum, IntType, PredicateBuilder, Schema, TableSchema}; +use paimon::table::Table; + +async fn pk_table() -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("v", DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "1") + .build() + .unwrap(); + let path = "memory:/pk_single_run_read"; + let (io, table) = memory_table(path, TableSchema::new(0, &schema)); + setup_dirs(&io, path).await; + persist_table_schema(&io, path, table.schema()).await; + table +} + +async fn write(table: &Table, ids: std::ops::Range, offset: i32) { + let ids: Vec = ids.collect(); + let values: Vec = ids.iter().map(|id| id + offset).collect(); + write_batch( + table, + &RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(ids)) as ArrayRef), + ("v", Arc::new(Int32Array::from(values)) as ArrayRef), + ]) + .unwrap(), + ) + .await; +} + +async fn read( + table: &Table, + projection: &[&str], + filter: Option, +) -> Vec> { + let mut builder = table.new_read_builder(); + builder.with_projection(projection).unwrap(); + if let Some(filter) = filter { + builder.with_filter(filter); + } + let plan = builder.new_scan().plan().await.unwrap(); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut rows: Vec> = batches + .iter() + .flat_map(|batch| { + (0..batch.num_rows()).map(move |row| { + (0..batch.num_columns()) + .map(|col| { + batch + .column(col) + .as_any() + .downcast_ref::() + .unwrap() + .value(row) + }) + .collect() + }) + }) + .collect(); + rows.sort_unstable(); + rows +} + +#[tokio::test] +async fn disjoint_files_read_without_merge_and_overlaps_still_merge() { + let table = pk_table().await; + // Three files with disjoint keys land in one split, one run per section. + write(&table, 0..100, 1000).await; + write(&table, 100..200, 1000).await; + write(&table, 200..300, 1000).await; + let plan = table.new_read_builder().new_scan().plan().await.unwrap(); + assert_eq!(plan.splits().len(), 1); + + let expected: Vec> = (0..300).map(|id| vec![id, id + 1000]).collect(); + assert_eq!(read(&table, &["id", "v"], None).await, expected); + // Value-only projection and a residual filter on the unmerged path. + let pb = PredicateBuilder::new(table.schema().fields()); + assert_eq!( + read( + &table, + &["v"], + Some(pb.greater_or_equal("v", Datum::Int(1250)).unwrap()) + ) + .await, + (1250..1300).map(|v| vec![v]).collect::>() + ); + + // Overlapping keys make a two-run section, which still merges. + write(&table, 150..160, 2000).await; + let expected: Vec> = (0..300) + .map(|id| vec![id, id + if (150..160).contains(&id) { 2000 } else { 1000 }]) + .collect(); + assert_eq!(read(&table, &["id", "v"], None).await, expected); +}