From e31c24dca5919acf4fc46a20ae6d9d5ea0dfd7be Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 29 Oct 2026 23:22:17 +0800 Subject: [PATCH 1/4] feat(table): cache lookup history for incremental compaction --- Cargo.lock | 1 + crates/paimon/Cargo.toml | 1 + crates/paimon/src/btree/mod.rs | 2 + crates/paimon/src/btree/sst_file.rs | 12 + crates/paimon/src/io/cache/generated.rs | 223 +++++ crates/paimon/src/io/cache/mod.rs | 1 + crates/paimon/src/table/compact.rs | 56 +- .../paimon/src/table/compact/primary_key.rs | 31 +- .../table/compact/primary_key/lookup_file.rs | 843 ++++++++++++++++++ .../table/compact/primary_key/lookup_store.rs | 223 +++++ .../src/table/compact/primary_key/planner.rs | 2 +- crates/paimon/src/table/kv_file_reader.rs | 13 + .../src/table/merge_tree_split_generator.rs | 4 + crates/paimon/src/table/sort_merge.rs | 42 + crates/paimon/src/table/table_commit.rs | 99 +- .../tests/pk_compact/incremental_lookup.rs | 410 +++++++++ crates/paimon/tests/pk_compact_test.rs | 3 + docs/src/compaction.md | 32 +- 18 files changed, 1928 insertions(+), 70 deletions(-) create mode 100644 crates/paimon/src/io/cache/generated.rs create mode 100644 crates/paimon/src/table/compact/primary_key/lookup_file.rs create mode 100644 crates/paimon/src/table/compact/primary_key/lookup_store.rs create mode 100644 crates/paimon/tests/pk_compact/incremental_lookup.rs diff --git a/Cargo.lock b/Cargo.lock index a2fb5f151..01748f3ce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4654,6 +4654,7 @@ dependencies = [ "arrow-array", "arrow-buffer", "arrow-cast", + "arrow-ipc", "arrow-json", "arrow-ord", "arrow-row", diff --git a/crates/paimon/Cargo.toml b/crates/paimon/Cargo.toml index b655d2607..3d11744c3 100644 --- a/crates/paimon/Cargo.toml +++ b/crates/paimon/Cargo.toml @@ -131,6 +131,7 @@ arrow-arith = { workspace = true } arrow-buffer = { workspace = true } arrow-cast = { workspace = true } arrow-json = "58.0" +arrow-ipc = "58.0" arrow-ord = { workspace = true } arrow-row = { workspace = true } arrow-schema = { workspace = true } diff --git a/crates/paimon/src/btree/mod.rs b/crates/paimon/src/btree/mod.rs index 644883332..23b11c5ba 100644 --- a/crates/paimon/src/btree/mod.rs +++ b/crates/paimon/src/btree/mod.rs @@ -54,6 +54,7 @@ mod writer; pub use block::BlockCompressionType; pub(crate) use block::{ compress_block, compress_codec_block, compute_crc32, decompress_block, decompress_codec_block, + BlockHandle, BlockReader, }; pub use composite_key::serialize_composite_key; pub(crate) use composite_key::CompositeKeyCodec; @@ -64,6 +65,7 @@ pub use key_serde::{make_key_comparator, serialize_datum}; pub use meta::BTreeIndexMeta; pub use query::IndexQuery; pub use reader::BTreeIndexReader; +pub(crate) use sst_file::{read_block_from_bytes, SstFileWriter}; pub use writer::BTreeIndexWriter; #[cfg(test)] diff --git a/crates/paimon/src/btree/sst_file.rs b/crates/paimon/src/btree/sst_file.rs index 5695a21a6..2b6a0308a 100644 --- a/crates/paimon/src/btree/sst_file.rs +++ b/crates/paimon/src/btree/sst_file.rs @@ -92,6 +92,18 @@ impl SstFileWriter { writer } + pub(crate) fn retained_bytes(&self) -> usize { + self.data_block_writer + .memory() + .saturating_add(self.index_block_writer.memory()) + .saturating_add( + self.bloom_hashes + .as_ref() + .map_or(0, |hashes| hashes.capacity() * 4), + ) + .saturating_add(self.last_key.as_ref().map_or(0, Vec::capacity)) + } + /// Current write position in the output. pub fn position(&self) -> u64 { self.bytes_written diff --git a/crates/paimon/src/io/cache/generated.rs b/crates/paimon/src/io/cache/generated.rs new file mode 100644 index 000000000..c1e8b8874 --- /dev/null +++ b/crates/paimon/src/io/cache/generated.rs @@ -0,0 +1,223 @@ +// 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. + +//! Disposable immutable files stored in the ordinary cache backend and budget. + +use super::LocalCache; +use crate::io::{FileRead, FileWrite}; +use crate::{Error, Result}; +use bytes::{Bytes, BytesMut}; +use std::ops::Range; +use std::sync::Arc; + +impl LocalCache { + pub(crate) fn generated_block_size(&self) -> u64 { + self.block_size + } + + pub(crate) async fn generated_get(&self, path: &str, index: u64, len: usize) -> Option { + let token = self.read_token(path); + self.get_block(&self.block_key(path, index), len, &token) + .await + .filter(|block| block.len() == len) + } + + pub(crate) async fn generated_put(&self, path: &str, index: u64, data: Bytes) -> Result<()> { + let token = self.read_token(path); + let key = self.block_key(path, index); + self.put_block(&key, data.clone(), &token).await; + if self.get_block(&key, data.len(), &token).await.as_ref() != Some(&data) { + return Err(Error::ResourceExhausted { + message: "Local cache could not retain a Lookup File block; check local-cache.max-size and cache storage".into(), + }); + } + Ok(()) + } + + /// Coalesce process-local builds by their exact head key. Do not hold a + /// shared download stripe while the build reads its source: a source block + /// may use the same stripe and would wait for its own build to finish. + pub(crate) async fn generated_load_lock(&self, path: &str) -> tokio::sync::OwnedMutexGuard<()> { + let key = self.block_key(path, 0); + self.block_load_lock(&key).await.lock_owned().await + } + + pub(crate) async fn generated_publish(&self, path: &str, descriptor: Bytes) -> Result<()> { + // Only publication holds the shared stripe. No source reads occur + // here, and another publisher cannot overwrite before verification. + let _guard = self.shared_load_lock(&self.block_key(path, 0)).await; + self.generated_put(path, 0, descriptor).await + } +} + +#[cfg(test)] +pub(crate) fn colliding_cached_source( + cache: Arc, + head: &str, + data: Bytes, +) -> Arc { + struct Source(Bytes); + #[async_trait::async_trait] + impl FileRead for Source { + async fn read(&self, range: Range) -> Result { + Ok(self.0.slice(range.start as usize..range.end as usize)) + } + } + // Shared downloads use 256 lock stripes, not exact block locks. Find a + // genuine source-block collision rather than relying on scheduling luck. + let head_key = cache.block_key(head, 0).cache_relative_path(); + let path = (0..) + .map(|index| format!("lookup-source-{index}.parquet")) + .find(|path| { + cache + .block_key(path, 0) + .cache_relative_path() + .components() + .next() + == head_key.components().next() + }) + .unwrap(); + assert!(cache.is_cacheable(&path)); + assert!(data.len() as u64 <= cache.generated_block_size()); + Arc::new(super::reader::CachedFileReader::new( + Arc::new(Source(data.clone())), + path, + data.len() as u64, + cache, + )) +} + +pub(crate) struct GeneratedFileRead { + cache: Arc, + path: String, + size: u64, +} + +impl GeneratedFileRead { + pub(crate) fn new(cache: Arc, path: String, size: u64) -> Self { + Self { cache, path, size } + } +} + +#[async_trait::async_trait] +impl FileRead for GeneratedFileRead { + async fn read(&self, range: Range) -> Result { + if range.start > range.end || range.end > self.size { + return Err(invalid("Generated cache range exceeds its file")); + } + let len = usize::try_from(range.end - range.start) + .map_err(|_| invalid("Generated cache range exceeds the address space"))?; + let mut output = BytesMut::with_capacity(len); + let mut offset = range.start; + let size = self.cache.block_size; + while offset < range.end { + let index = offset / size; + let start = index * size; + let end = start.saturating_add(size).min(self.size); + let len = usize::try_from(end - start) + .map_err(|_| invalid("Generated cache block exceeds the address space"))?; + let block = self + .cache + .generated_get(&self.path, index, len) + .await + .ok_or_else(|| invalid("Lookup File block was evicted or corrupt"))?; + let take_end = range.end.min(end); + output + .extend_from_slice(&block[(offset - start) as usize..(take_end - start) as usize]); + offset = take_end; + } + Ok(output.freeze()) + } +} + +pub(crate) struct GeneratedFileWrite { + cache: Arc, + path: String, + buffer: BytesMut, + index: u64, + length: u64, +} + +impl GeneratedFileWrite { + pub(crate) fn new(cache: Arc, path: String) -> Self { + Self { + cache, + path, + buffer: BytesMut::new(), + index: 0, + length: 0, + } + } + + async fn flush_block(&mut self) -> Result<()> { + if !self.buffer.is_empty() { + self.cache + .generated_put(&self.path, self.index, self.buffer.split().freeze()) + .await?; + self.index += 1; + } + Ok(()) + } +} + +#[async_trait::async_trait] +impl FileWrite for GeneratedFileWrite { + async fn write(&mut self, mut bytes: Bytes) -> Result<()> { + self.length = self + .length + .checked_add(bytes.len() as u64) + .ok_or_else(|| invalid("Generated cache file length overflow"))?; + let size = usize::try_from(self.cache.block_size) + .map_err(|_| invalid("Generated cache block exceeds the address space"))?; + while !bytes.is_empty() { + let take = (size - self.buffer.len()).min(bytes.len()); + self.buffer.extend_from_slice(&bytes.split_to(take)); + if self.buffer.len() == size { + self.flush_block().await?; + } + } + Ok(()) + } + + async fn close(&mut self) -> Result<()> { + self.flush_block().await?; + for index in 0..self.index { + let length = + (self.length - index * self.cache.block_size).min(self.cache.block_size) as usize; + if self + .cache + .generated_get(&self.path, index, length) + .await + .is_none() + { + return Err(Error::ResourceExhausted { + message: + "Lookup File was evicted during construction; increase local-cache.max-size" + .into(), + }); + } + } + Ok(()) + } +} + +fn invalid(message: &str) -> Error { + Error::DataInvalid { + message: message.into(), + source: None, + } +} diff --git a/crates/paimon/src/io/cache/mod.rs b/crates/paimon/src/io/cache/mod.rs index 4fa399534..5a58ad94a 100644 --- a/crates/paimon/src/io/cache/mod.rs +++ b/crates/paimon/src/io/cache/mod.rs @@ -17,6 +17,7 @@ mod disk; mod file_type; +pub(crate) mod generated; mod memory; mod priority_lru; mod reader; diff --git a/crates/paimon/src/table/compact.rs b/crates/paimon/src/table/compact.rs index d06ebd302..cad166b29 100644 --- a/crates/paimon/src/table/compact.rs +++ b/crates/paimon/src/table/compact.rs @@ -111,14 +111,10 @@ impl<'a> CompactBuilder<'a> { /// Executor resources are local; input files and output settings come from each task. pub fn new_executor(&self) -> Result { ensure_supported(self.table)?; - if self - .table - .schema() - .core_options() - .try_changelog_producer()? - == ChangelogProducer::Lookup - && self.table.file_io().local_cache().is_none() - { + let options = self.table.schema().core_options(); + let lookup = options.try_changelog_producer()? == ChangelogProducer::Lookup; + let local_cache = self.table.file_io().local_cache(); + if lookup && local_cache.is_none() { return Err(Error::ConfigInvalid { message: "Lookup compaction requires local-cache.enabled=true on the catalog/FileIO" @@ -128,6 +124,14 @@ impl<'a> CompactBuilder<'a> { Ok(CompactExecutor { table: self.table.clone(), resources: self.resources.clone(), + lookup_cache: if lookup && options.merge_engine()? == MergeEngine::Deduplicate { + Some(Arc::new(primary_key::lookup_file::LookupFileCache::new( + local_cache.expect("lookup cache validated above"), + self.table.schema().options(), + )?)) + } else { + None + }, }) } @@ -282,6 +286,7 @@ impl CompactPlanner { pub struct CompactExecutor { table: Table, resources: Option, + lookup_cache: Option>, } /// Prepared files owned by the caller until submitted to `TableCommit::commit_compact`. @@ -394,8 +399,33 @@ impl CompactExecutor { }); } if task.is_primary_key() { + let lookup = if let Some(cache) = &self.lookup_cache { + let history = active + .iter() + .filter(|entry| { + entry.partition() == task.partition + && entry.bucket() == task.bucket + && entry.total_buckets() == task.total_buckets + && entry.file().level > task.output_level() + }) + .map(|entry| entry.file().clone()) + .collect::>(); + if history.is_empty() { + None + } else { + Some(primary_key::lookup_store::provider( + &table, + &split, + history, + cache.clone(), + self.resources.clone(), + )?) + } + } else { + None + }; let (files, changelog_files) = - primary_key::execute(&table, task, &split, self.resources.clone()).await?; + primary_key::execute(&table, task, &split, self.resources.clone(), lookup).await?; return Ok(CompactResult { task: task.clone(), files, @@ -765,11 +795,17 @@ pub(super) fn prepare_results(table: &Table, results: &[CompactResult]) -> Resul )); } let mut changelog_rows = 0u64; + let changelog_level = output_level.unwrap_or_else(|| { + result + .files + .first() + .map_or_else(|| task.output_level(), |file| file.level) + }); for file in &result.changelog_files { validate_file(file, true)?; if file.schema_id != task.context.schema.id() || file.file_source != Some(1) - || file.level != task.output_level() + || file.level != changelog_level || file .delete_row_count .is_none_or(|count| count > file.row_count) diff --git a/crates/paimon/src/table/compact/primary_key.rs b/crates/paimon/src/table/compact/primary_key.rs index 5f7ae1119..9f65eb6df 100644 --- a/crates/paimon/src/table/compact/primary_key.rs +++ b/crates/paimon/src/table/compact/primary_key.rs @@ -16,6 +16,8 @@ // under the License. pub(super) mod lookup; +pub(super) mod lookup_file; +pub(super) mod lookup_store; pub(in crate::table) mod planner; use std::sync::Arc; @@ -35,7 +37,7 @@ 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, + compare_sequence_order, BufferedBatch, LookupProvider, MergeFunction, MergeResult, MergeRow, CHANGELOG_BATCH_METADATA_KEY, }; use crate::table::table_read::{ @@ -44,7 +46,7 @@ use crate::table::table_read::{ use crate::table::{DataSplit, Plan, Table}; use crate::Result; -/// Use incremental sorted runs when merging subsets does not need history lookup. +/// Select incremental sorted runs for supported merge modes. pub(super) struct PrimaryKeyCompactPlanner; impl PrimaryKeyCompactPlanner { @@ -263,23 +265,17 @@ fn read_type(table: &Table) -> Vec { fields } -pub(super) async fn execute( - table: &Table, - task: &CompactTask, - split: &DataSplit, - resources: Option, -) -> Result<(Vec, Vec)> { +fn new_reader(table: &Table, resources: Option) -> Result { let schema = table.schema(); let options = schema.core_options(); let merge_engine = options.merge_engine()?; - let lookup_changelog = options.try_changelog_producer()? == ChangelogProducer::Lookup; let read_type = read_type(table); 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( + Ok(KeyValueFileReader::new( table.file_io().clone(), KeyValueReadConfig { table_name: table.identifier().full_name(), @@ -303,7 +299,20 @@ pub(super) async fn execute( parquet_read_budget: Some(budget), mosaic_prefetch: configured_mosaic_prefetch(table)?, }, - ); + )) +} + +pub(super) async fn execute( + table: &Table, + task: &CompactTask, + split: &DataSplit, + resources: Option, + lookup: Option>, +) -> Result<(Vec, Vec)> { + let schema = table.schema(); + let options = schema.core_options(); + let lookup_changelog = options.try_changelog_producer()? == ChangelogProducer::Lookup; + let reader = new_reader(table, resources.clone())?.with_lookup(lookup); let mut batches = reader.read_with_merge_function(std::slice::from_ref(split), move |config, fields| { if lookup_changelog { diff --git a/crates/paimon/src/table/compact/primary_key/lookup_file.rs b/crates/paimon/src/table/compact/primary_key/lookup_file.rs new file mode 100644 index 000000000..574828fdb --- /dev/null +++ b/crates/paimon/src/table/compact/primary_key/lookup_file.rs @@ -0,0 +1,843 @@ +// 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. + +//! Disposable, versioned local SST files containing complete physical KV rows. + +use crate::btree::{ + compress_block, decompress_block, read_block_from_bytes, BTreeFileFooter, BlockCompressionType, + BlockHandle, BlockReader, SstFileWriter, +}; +use crate::common::options::parse_memory_size; +use crate::io::cache::generated::{GeneratedFileRead, GeneratedFileWrite}; +use crate::io::cache::LocalCache; +use crate::io::FileRead; +use crate::table::ArrowRecordBatchStream; +use crate::{Error, Result}; +use arrow_array::{RecordBatch, UInt64Array}; +use arrow_ipc::{reader::StreamReader, writer::StreamWriter}; +use arrow_row::{RowConverter, SortField}; +use bytes::Bytes; +use futures::TryStreamExt; +use std::collections::HashMap; +use std::io::{Cursor, Write}; +use std::sync::{Arc, Mutex}; + +const PREFIX: &str = "__paimon_generated__/lookup-sst-v1"; +const HEAD_LEN: usize = 24; +const POINTER_LEN: usize = 37; + +struct Payload { + path: String, + offset: u64, + batch: Arc, +} + +pub(in crate::table::compact) struct LookupFileCache { + local: Arc, + limit: usize, + compression: BlockCompressionType, + // One executor cache owns this working budget. Serialize both builds and + // warm index/IPC reads so different files cannot multiply that budget. + working: tokio::sync::Mutex<()>, + // One decoded IPC batch across all files, never one cache per file/key. + payload: Mutex>, +} + +impl LookupFileCache { + pub(in crate::table::compact) fn new( + local: Arc, + options: &HashMap, + ) -> Result { + if options.contains_key("lookup.cache-max-disk-size") { + return Err(invalid( + "Lookup Files use local-cache.max-size, not lookup.cache-max-disk-size", + )); + } + let memory = options + .get("lookup.cache-max-memory-size") + .map(|value| { + parse_memory_size(value) + .map_err(|_| invalid("Invalid lookup.cache-max-memory-size")) + }) + .transpose()? + .unwrap_or(256 * 1024 * 1024); + if memory == 0 { + return Err(Error::ConfigInvalid { + message: "lookup.cache-max-memory-size must be positive".into(), + }); + } + // Bound input, encoded/decoded payloads, SST construction and reads + // separately, leaving room for copies and Arrow sort-key construction. + let limit = usize::try_from(memory / 8).map_err(invalid)?; + if limit < 64 * 1024 || local.generated_block_size() > limit as u64 { + return Err(exhausted( + "Lookup memory budget is too small for local-cache.block-size and SST metadata", + )); + } + let compression = match options + .get("lookup.cache-spill-compression") + .map(String::as_str) + .unwrap_or("zstd") + { + "none" => BlockCompressionType::None, + "zstd" => BlockCompressionType::Zstd, + "lz4" => BlockCompressionType::Lz4, + "lzo" => BlockCompressionType::Lzo, + _ => return Err(invalid("Invalid lookup.cache-spill-compression")), + }; + Ok(Self { + local, + limit, + compression, + working: tokio::sync::Mutex::new(()), + payload: Mutex::new(None), + }) + } + + pub(super) async fn lookup( + &self, + identity: &str, + key: &[u8], + keys: &[usize], + input: F, + ) -> Result> + where + F: Fn() -> Result + Send, + { + // Lock order is cache working set, then the shared per-file head. + let _working = self.working.lock().await; + let head = format!("{PREFIX}/{identity}/head"); + let _guard = self.local.generated_load_lock(&head).await; + let mut published = self.local.generated_get(&head, 0, HEAD_LEN).await; + for attempt in 0..2 { + let descriptor = match published.take() { + Some(bytes) => bytes, + None => self.build(&head, input()?, keys).await?, + }; + let generation = uuid::Uuid::from_slice(&descriptor[..16]).map_err(invalid)?; + let path = format!("{PREFIX}/{identity}/{generation}"); + let size = u64::from_le_bytes(descriptor[16..].try_into().unwrap()); + match self.read(&path, size, key).await { + Ok(row) => return Ok(row), + Err(error) if attempt == 1 => return Err(error), + Err(_) => { + *self.payload.lock().unwrap() = None; + // A new immutable generation replaces the head only after + // close. Do not invalidate a shared backend's global epoch. + } + } + } + unreachable!() + } + + async fn build( + &self, + head: &str, + mut input: ArrowRecordBatchStream, + keys: &[usize], + ) -> Result { + let generation = uuid::Uuid::new_v4(); + let path = format!("{}/{generation}", head.strip_suffix("/head").unwrap()); + let mut writer = SstFileWriter::with_compression_level( + Box::new(GeneratedFileWrite::new(self.local.clone(), path)), + 32 * 1024, + self.compression, + 1, + ); + let mut previous = None::>; + let mut converter = None; + while let Some(batch) = input.try_next().await? { + if batch.num_rows() == 0 { + continue; + } + if batch.get_array_memory_size() > self.limit { + return Err(exhausted("Lookup input batch exceeds lookup.cache-max-memory-size; reduce read.batch-size")); + } + if keys.is_empty() || keys.iter().any(|&index| index >= batch.num_columns()) { + return Err(invalid("Invalid Lookup File primary-key projection")); + } + if converter.is_none() { + converter = Some( + RowConverter::new( + keys.iter() + .map(|&key| { + SortField::new(batch.schema().field(key).data_type().clone()) + }) + .collect(), + ) + .map_err(invalid)?, + ); + } + let rows = converter + .as_ref() + .unwrap() + .convert_columns( + &keys + .iter() + .map(|&key| batch.column(key).clone()) + .collect::>(), + ) + .map_err(invalid)?; + if rows.size() > self.limit { + return Err(exhausted("Lookup sort keys exceed the memory budget")); + } + let mut ipc = LimitedBuffer { + bytes: Vec::new(), + limit: self.limit, + exceeded: false, + }; + let encoded = (|| { + let mut stream = StreamWriter::try_new(&mut ipc, batch.schema().as_ref())?; + stream.write(&batch)?; + stream.finish() + })(); + if let Err(error) = encoded { + if ipc.exceeded { + return Err(exhausted(error)); + } + return Err(invalid(error)); + } + let (bytes, compression) = + compress_block(&ipc.bytes, self.compression, 1).map_err(invalid)?; + if bytes.len() > self.limit { + return Err(exhausted( + "Lookup encoded payload exceeds the memory budget", + )); + } + let offset = writer.position(); + writer.write_raw(&bytes).await.map_err(exhausted)?; + let pointer = Pointer { + offset, + length: bytes.len() as u64, + raw_length: ipc.bytes.len() as u64, + row: 0, + crc: crc32fast::hash(&bytes), + compression, + }; + for row in 0..batch.num_rows() { + let key = rows.row(row); + if previous.as_deref().is_some_and(|old| old >= key.as_ref()) { + return Err(invalid( + "Lookup history must have strictly increasing primary keys", + )); + } + previous = Some(key.as_ref().to_vec()); + if writer + .retained_bytes() + .saturating_add(key.as_ref().len().saturating_mul(3)) + .saturating_add(POINTER_LEN) + > self.limit + { + return Err(exhausted( + "Lookup SST construction exceeds lookup.cache-max-memory-size", + )); + } + writer + .put( + key.as_ref(), + &Pointer { + row: row as u64, + ..pointer + } + .encode(), + ) + .await + .map_err(exhausted)?; + if writer.retained_bytes() > self.limit { + return Err(exhausted( + "Lookup SST construction exceeds lookup.cache-max-memory-size", + )); + } + } + } + writer.flush().await.map_err(exhausted)?; + if writer.retained_bytes() > self.limit { + return Err(exhausted( + "Lookup SST construction exceeds lookup.cache-max-memory-size", + )); + } + let index = writer.write_index_block().await.map_err(exhausted)?; + writer + .write_raw(&BTreeFileFooter::new(None, index, None).write_footer()) + .await + .map_err(exhausted)?; + let size = writer.position(); + writer.close().await.map_err(exhausted)?; + let mut descriptor = generation.as_bytes().to_vec(); + descriptor.extend_from_slice(&size.to_le_bytes()); + let descriptor = Bytes::from(descriptor); + self.local + .generated_publish(head, descriptor.clone()) + .await?; + Ok(descriptor) + } + + async fn block( + &self, + reader: &GeneratedFileRead, + size: u64, + handle: BlockHandle, + ) -> Result { + let length = u64::from(handle.size) + 5; + if length > self.limit as u64 { + return Err(exhausted("Lookup index block exceeds the memory budget")); + } + let end = handle + .offset + .checked_add(length) + .filter(|&end| end <= size) + .ok_or_else(|| invalid("Lookup block is outside its file"))?; + let bytes = reader.read(handle.offset..end).await?; + read_block_from_bytes(&bytes, handle.size).map_err(invalid) + } + + async fn read(&self, path: &str, size: u64, key: &[u8]) -> Result> { + let footer_start = size + .checked_sub(BTreeFileFooter::ENCODED_LENGTH as u64) + .ok_or_else(|| invalid("Truncated Lookup File"))?; + let reader = GeneratedFileRead::new(self.local.clone(), path.into(), size); + let footer = BTreeFileFooter::read_footer(&reader.read(footer_start..size).await?) + .map_err(invalid)?; + let index = self + .block(&reader, footer_start, footer.index_block_handle) + .await?; + let cmp = |a: &[u8], b: &[u8]| Ok(a.cmp(b)); + let (_, mut candidates) = index.seek_and_iter(key, &cmp).map_err(invalid)?; + let Some((_, bytes)) = candidates.next() else { + return Ok(None); + }; + let block = self + .block( + &reader, + footer_start, + BlockHandle::decode(bytes).map_err(invalid)?, + ) + .await?; + let (found, mut entries) = block.seek_and_iter(key, &cmp).map_err(invalid)?; + let Some((_, bytes)) = entries.next().filter(|_| found) else { + return Ok(None); + }; + let pointer = Pointer::decode(bytes)?; + if pointer.length > self.limit as u64 || pointer.raw_length > self.limit as u64 { + return Err(exhausted("Lookup IPC block exceeds the memory budget")); + } + let cached = self + .payload + .lock() + .unwrap() + .as_ref() + .filter(|p| p.path == path && p.offset == pointer.offset) + .map(|p| p.batch.clone()); + let batch = match cached { + Some(batch) => batch, + None => { + let end = pointer + .offset + .checked_add(pointer.length) + .filter(|&end| end <= footer_start) + .ok_or_else(|| invalid("Lookup payload is outside its file"))?; + let bytes = reader.read(pointer.offset..end).await?; + if crc32fast::hash(&bytes) != pointer.crc { + return Err(invalid("Lookup IPC checksum mismatch")); + } + let decoded = decompress_block(&bytes, pointer.compression).map_err(invalid)?; + if decoded.len() as u64 != pointer.raw_length { + return Err(invalid("Lookup IPC length mismatch")); + } + let mut stream = + StreamReader::try_new(Cursor::new(decoded), None).map_err(invalid)?; + let batch = Arc::new( + stream + .next() + .ok_or_else(|| invalid("Empty Lookup IPC payload"))? + .map_err(invalid)?, + ); + if stream.next().is_some() { + return Err(invalid("Lookup IPC payload contains multiple batches")); + } + if batch.get_array_memory_size() > self.limit { + return Err(exhausted("Decoded Lookup IPC exceeds the memory budget")); + } + *self.payload.lock().unwrap() = Some(Payload { + path: path.into(), + offset: pointer.offset, + batch: batch.clone(), + }); + batch + } + }; + if pointer.row >= batch.num_rows() as u64 { + return Err(invalid("Lookup row is outside its IPC batch")); + } + arrow_select::take::take_record_batch(&batch, &UInt64Array::from(vec![pointer.row])) + .map(Some) + .map_err(invalid) + } +} + +#[derive(Clone, Copy)] +struct Pointer { + offset: u64, + length: u64, + raw_length: u64, + row: u64, + crc: u32, + compression: BlockCompressionType, +} +impl Pointer { + fn encode(self) -> Vec { + let mut bytes = Vec::with_capacity(POINTER_LEN); + for value in [self.offset, self.length, self.raw_length, self.row] { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes.extend_from_slice(&self.crc.to_le_bytes()); + bytes.push(self.compression as u8); + bytes + } + fn decode(bytes: &[u8]) -> Result { + if bytes.len() != POINTER_LEN { + return Err(invalid("Invalid Lookup IPC pointer")); + } + let number = |offset| u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()); + Ok(Self { + offset: number(0), + length: number(8), + raw_length: number(16), + row: number(24), + crc: u32::from_le_bytes(bytes[32..36].try_into().unwrap()), + compression: BlockCompressionType::from_persistent_id(bytes[36]).map_err(invalid)?, + }) + } +} + +struct LimitedBuffer { + bytes: Vec, + limit: usize, + exceeded: bool, +} +impl Write for LimitedBuffer { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + if self.bytes.len().saturating_add(bytes.len()) > self.limit { + self.exceeded = true; + return Err(std::io::Error::other( + "Lookup IPC exceeds the memory budget", + )); + } + self.bytes.extend_from_slice(bytes); + Ok(bytes.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} +fn invalid(message: impl ToString) -> Error { + Error::DataInvalid { + message: message.to_string(), + source: None, + } +} +fn exhausted(message: impl ToString) -> Error { + Error::ResourceExhausted { + message: message.to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{ArrayRef, Int32Array, Int64Array, Int8Array, StringArray}; + use futures::{stream, StreamExt}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + fn local(directory: Option<&std::path::Path>, shared: bool, capacity: &str) -> Arc { + local_with_whitelist(directory, shared, capacity, "meta") + } + + fn local_with_whitelist( + directory: Option<&std::path::Path>, + shared: bool, + capacity: &str, + whitelist: &str, + ) -> Arc { + let mut options = crate::Options::new(); + for (key, value) in [ + ("local-cache.enabled", "true"), + ("local-cache.max-size", capacity), + ("local-cache.block-size", "8kb"), + ("local-cache.whitelist", whitelist), + ] { + options.set(key, value); + } + if let Some(directory) = directory { + options.set("local-cache.dir", directory.to_str().unwrap()); + } + if shared { + options.set("local-cache.shared", "true"); + } + crate::io::cache::create_local_cache(&options) + .unwrap() + .unwrap() + } + fn batch() -> RecordBatch { + RecordBatch::try_from_iter([ + ( + "_SEQUENCE_NUMBER", + Arc::new(Int64Array::from(vec![9, 10, 11])) as ArrayRef, + ), + ( + "_VALUE_KIND", + Arc::new(Int8Array::from(vec![3, 2, 0])) as ArrayRef, + ), + ("id", Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef), + ( + "value", + Arc::new(StringArray::from(vec![ + Some("deleted payload"), + None, + Some("variable payload"), + ])) as ArrayRef, + ), + ]) + .unwrap() + } + fn key(value: i32) -> Vec { + RowConverter::new(vec![SortField::new(arrow_schema::DataType::Int32)]) + .unwrap() + .convert_columns(&[Arc::new(Int32Array::from(vec![value]))]) + .unwrap() + .row(0) + .as_ref() + .to_vec() + } + fn source( + counter: Arc, + ) -> impl Fn() -> Result + Send + Sync { + move || { + counter.fetch_add(1, Ordering::Relaxed); + Ok(stream::once(async { + tokio::task::yield_now().await; + Ok(batch()) + }) + .boxed()) + } + } + async fn descriptor(cache: &LookupFileCache, identity: &str) -> Option { + cache + .local + .generated_get(&format!("{PREFIX}/{identity}/head"), 0, HEAD_LEN) + .await + } + + #[tokio::test] + async fn cold_hot_and_concurrent_lookups_retain_complete_kvs_on_all_backends() { + for backend in ["memory", "disk", "shared"] { + let directory = tempfile::tempdir().unwrap(); + let local = local( + (backend != "memory").then_some(directory.path()), + backend == "shared", + "4mb", + ); + let cache = LookupFileCache::new(local.clone(), &HashMap::new()).unwrap(); + let calls = Arc::new(AtomicUsize::new(0)); + let input = source(calls.clone()); + let results = futures::future::join_all((0..8).map(|_| async { + cache + .lookup("full-kv", &key(2), &[2], &input) + .await + .unwrap() + })) + .await; + assert!(results + .iter() + .all(|row| row.as_ref() == Some(&batch().slice(1, 1)))); + for value in [1, 3, 2] { + assert_eq!( + cache + .lookup("full-kv", &key(value), &[2], &input) + .await + .unwrap(), + Some(batch().slice((value - 1) as usize, 1)) + ); + } + assert!(cache + .lookup("full-kv", &key(99), &[2], &input) + .await + .unwrap() + .is_none()); + assert_eq!(calls.load(Ordering::Relaxed), 1, "{backend}"); + let active = Arc::new(AtomicUsize::new(0)); + let peak = Arc::new(AtomicUsize::new(0)); + let distinct_input = || { + let active = active.clone(); + let peak = peak.clone(); + Ok(async_stream::try_stream! { + let concurrent = active.fetch_add(1, Ordering::Relaxed) + 1; + peak.fetch_max(concurrent, Ordering::Relaxed); + yield batch(); + // Keep the source open after one decoded batch, while + // other identities have a chance to start building. + tokio::task::yield_now().await; + active.fetch_sub(1, Ordering::Relaxed); + } + .boxed()) + }; + let distinct = futures::future::join_all((0..8).map(|id| { + let cache = &cache; + let input = &distinct_input; + async move { + cache + .lookup(&format!("distinct-{id}"), &key(2), &[2], input) + .await + .unwrap() + } + })) + .await; + assert!(distinct + .iter() + .all(|row| row.as_ref() == Some(&batch().slice(1, 1)))); + assert_eq!(peak.load(Ordering::Relaxed), 1, "{backend}"); + assert_eq!(active.load(Ordering::Relaxed), 0); + // A fresh cache object has no decoded payload, but can reuse the + // published local file without calling the source closure. + drop(cache); + let cache = LookupFileCache::new(local, &HashMap::new()).unwrap(); + assert_eq!( + cache + .lookup("full-kv", &key(1), &[2], &input) + .await + .unwrap(), + Some(batch().slice(0, 1)) + ); + assert_eq!(calls.load(Ordering::Relaxed), 1); + } + } + + #[tokio::test] + async fn shared_head_build_does_not_hold_a_colliding_source_download_stripe() { + let directory = tempfile::tempdir().unwrap(); + let local = local_with_whitelist(Some(directory.path()), true, "4mb", "data"); + let mut bytes = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut bytes, batch().schema().as_ref()).unwrap(); + writer.write(&batch()).unwrap(); + writer.finish().unwrap(); + } + let length = bytes.len() as u64; + let reader = crate::io::cache::generated::colliding_cached_source( + local.clone(), + &format!("{PREFIX}/source-collision/head"), + Bytes::from(bytes), + ); + let input = || { + let reader = reader.clone(); + Ok(async_stream::try_stream! { + let bytes = reader.read(0..length).await?; + let mut stream = StreamReader::try_new(Cursor::new(bytes), None).map_err(invalid)?; + yield stream.next().unwrap().map_err(invalid)?; + } + .boxed()) + }; + let cache = LookupFileCache::new(local, &HashMap::new()).unwrap(); + let lookup_key = key(2); + let row = tokio::time::timeout( + std::time::Duration::from_secs(5), + cache.lookup("source-collision", &lookup_key, &[2], input), + ) + .await + .expect("a cold Lookup File must not wait for its own source download stripe") + .unwrap(); + assert_eq!(row, Some(batch().slice(1, 1))); + } + + #[tokio::test] + async fn interleaved_payloads_and_sst_blocks_support_point_lookups() { + let options = HashMap::from([("lookup.cache-spill-compression".into(), "none".into())]); + let cache = LookupFileCache::new(local(None, false, "4mb"), &options).unwrap(); + let calls = AtomicUsize::new(0); + let input = || { + calls.fetch_add(1, Ordering::Relaxed); + let batches = (0..20) + .map(|batch| { + let ids = (batch * 100..(batch + 1) * 100).collect::>(); + let values = ids + .iter() + .map(|id| format!("payload-{id:08}")) + .collect::>(); + RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(ids)) as ArrayRef), + ("value", Arc::new(StringArray::from(values)) as ArrayRef), + ]) + .map_err(invalid) + }) + .collect::>(); + Ok(stream::iter(batches).boxed()) + }; + for id in [0, 99, 100, 975, 1501, 1999] { + let row = cache + .lookup("blocks", &key(id), &[0], &input) + .await + .unwrap() + .unwrap(); + assert_eq!( + row.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + id + ); + assert_eq!( + row.column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + format!("payload-{id:08}") + ); + } + assert_eq!(calls.load(Ordering::Relaxed), 1); + } + + #[tokio::test] + async fn missing_and_corrupt_artifacts_are_rebuilt_instead_of_becoming_key_misses() { + for corrupt in [false, true] { + let local = local(None, false, "64kb"); + let cache = LookupFileCache::new(local.clone(), &HashMap::new()).unwrap(); + let calls = Arc::new(AtomicUsize::new(0)); + let input = source(calls.clone()); + cache.lookup("repair", &key(1), &[2], &input).await.unwrap(); + if corrupt { + let head = descriptor(&cache, "repair").await.unwrap(); + let generation = uuid::Uuid::from_slice(&head[..16]).unwrap(); + let path = format!("{PREFIX}/repair/{generation}"); + let size = u64::from_le_bytes(head[16..].try_into().unwrap()); + let mut bytes = local + .generated_get(&path, 0, size.min(8192) as usize) + .await + .unwrap() + .to_vec(); + bytes[0] ^= 0xff; + local + .generated_put(&path, 0, Bytes::from(bytes)) + .await + .unwrap(); + *cache.payload.lock().unwrap() = None; + } else { + for block in 0..16 { + local + .generated_put("cache-pressure", block, Bytes::from(vec![0; 8192])) + .await + .unwrap(); + } + assert!(descriptor(&cache, "repair").await.is_none()); + } + assert_eq!( + cache.lookup("repair", &key(1), &[2], &input).await.unwrap(), + Some(batch().slice(0, 1)) + ); + assert_eq!(calls.load(Ordering::Relaxed), 2); + } + } + + #[tokio::test] + async fn source_failure_and_cancellation_never_publish_partial_files() { + let local = local(None, false, "4mb"); + let cache = LookupFileCache::new(local, &HashMap::new()).unwrap(); + let failed = || Ok(stream::iter(vec![Ok(batch()), Err(invalid("source failed"))]).boxed()); + assert!(cache.lookup("failed", &key(1), &[2], failed).await.is_err()); + assert!(descriptor(&cache, "failed").await.is_none()); + let ready = Arc::new(tokio::sync::Notify::new()); + let stalled = || { + let ready = ready.clone(); + Ok(stream::iter(vec![Ok(batch())]) + .chain(stream::once(async move { + ready.notify_one(); + futures::future::pending::>().await + })) + .boxed()) + }; + let lookup_key = key(1); + tokio::select! { + result = cache.lookup("cancelled", &lookup_key, &[2], stalled) => panic!("unexpected completion: {result:?}"), + _ = ready.notified() => {} + } + assert!(descriptor(&cache, "cancelled").await.is_none()); + let calls = Arc::new(AtomicUsize::new(0)); + assert!(cache + .lookup("cancelled", &key(1), &[2], source(calls.clone())) + .await + .unwrap() + .is_some()); + assert_eq!(calls.load(Ordering::Relaxed), 1); + } + + #[tokio::test] + async fn capacity_input_and_index_limits_are_enforced_without_publishing() { + let cache = LookupFileCache::new(local(None, false, "1b"), &HashMap::new()).unwrap(); + assert!(matches!( + cache + .lookup( + "capacity", + &key(1), + &[2], + source(Arc::new(AtomicUsize::new(0))) + ) + .await, + Err(Error::ResourceExhausted { .. }) + )); + assert!(descriptor(&cache, "capacity").await.is_none()); + let options = HashMap::from([ + ("lookup.cache-max-memory-size".into(), "1mb".into()), + ("lookup.cache-spill-compression".into(), "none".into()), + ]); + let cache = LookupFileCache::new(local(None, false, "4mb"), &options).unwrap(); + let oversized = || { + let large = "x".repeat(256 * 1024); + Ok(stream::iter(vec![Ok(RecordBatch::try_from_iter([( + "id", + Arc::new(StringArray::from(vec![large])) as ArrayRef, + )]) + .unwrap())]) + .boxed()) + }; + assert!(matches!( + cache.lookup("input-limit", &[], &[0], oversized).await, + Err(Error::ResourceExhausted { .. }) + )); + assert!(descriptor(&cache, "input-limit").await.is_none()); + let wide_keys = || { + let batches = (b'a'..=b'f') + .map(|prefix| { + let key = format!("{}{}", prefix as char, "x".repeat(40_000)); + RecordBatch::try_from_iter([( + "id", + Arc::new(StringArray::from(vec![key])) as ArrayRef, + )]) + .map_err(invalid) + }) + .collect::>(); + Ok(stream::iter(batches).boxed()) + }; + let error = cache + .lookup("index-limit", &[], &[0], wide_keys) + .await + .unwrap_err(); + assert!(matches!(&error, Error::ResourceExhausted { .. }), "{error}"); + assert!(error.to_string().contains("SST construction"), "{error}"); + assert!(descriptor(&cache, "index-limit").await.is_none()); + } +} diff --git a/crates/paimon/src/table/compact/primary_key/lookup_store.rs b/crates/paimon/src/table/compact/primary_key/lookup_store.rs new file mode 100644 index 000000000..128518461 --- /dev/null +++ b/crates/paimon/src/table/compact/primary_key/lookup_store.rs @@ -0,0 +1,223 @@ +// 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::lookup_file::LookupFileCache; +use crate::io::FileRead; +use crate::resource::ResourceContext; +use crate::spec::{BinaryRowBuilder, DataField, DataFileMeta}; +use crate::table::merge_tree_split_generator::KeyComparator; +use crate::table::sort_merge::{LookupProvider, FILE_LEVEL_METADATA_KEY}; +use crate::table::{DataSplit, Table}; +use crate::{Error, Result}; +use arrow_array::RecordBatch; +use arrow_row::{RowConverter, SortField}; +use sha2::{Digest, Sha256}; +use std::collections::BTreeMap; +use std::sync::{Arc, Mutex}; + +struct HistoryFile { + meta: DataFileMeta, + identity: tokio::sync::OnceCell, +} + +struct LevelLookup { + table: Table, + split: DataSplit, + levels: BTreeMap>, + cache: Arc, + keys: Vec, + key_fields: Vec, + converter: Mutex, + comparator: Option, + resources: Option, + schema_identity: Vec, +} + +pub(in crate::table::compact) fn provider( + table: &Table, + split: &DataSplit, + files: Vec, + cache: Arc, + resources: Option, +) -> Result> { + let fields = super::read_type(table); + let keys = table + .schema() + .trimmed_primary_keys() + .iter() + .map(|name| { + fields + .iter() + .position(|field| field.name() == name) + .unwrap() + }) + .collect::>(); + let key_fields = keys + .iter() + .map(|&index| fields[index].clone()) + .collect::>(); + let schema = crate::arrow::build_read_arrow_schema(&key_fields)?; + let converter = RowConverter::new( + schema + .fields() + .iter() + .map(|field| SortField::new(field.data_type().clone())) + .collect(), + ) + .map_err(invalid)?; + let mut levels = BTreeMap::<_, Vec<_>>::new(); + for meta in files { + levels.entry(meta.level).or_default().push(HistoryFile { + meta, + identity: tokio::sync::OnceCell::new(), + }); + } + let options = table.schema().options().iter().collect::>(); + let schema_identity = serde_json::to_vec(&( + table.schema().fields(), + table.schema().primary_keys(), + table.schema().partition_keys(), + options, + )) + .map_err(invalid)?; + Ok(Arc::new(LevelLookup { + table: table.clone(), + split: split.clone(), + levels, + cache, + keys, + key_fields, + converter: Mutex::new(converter), + comparator: KeyComparator::from_table_schema(table.schema()), + resources, + schema_identity, + })) +} + +#[async_trait::async_trait] +impl LookupProvider for LevelLookup { + async fn lookup(&self, key: &[u8]) -> Result> { + let binary_key = { + let converter = self.converter.lock().unwrap(); + let parser = converter.parser(); + let columns = converter + .convert_rows(std::iter::once(parser.parse(key))) + .map_err(invalid)?; + let schema = crate::arrow::build_read_arrow_schema(&self.key_fields)?; + let batch = RecordBatch::try_new(schema, columns).map_err(invalid)?; + let mut row = BinaryRowBuilder::new(self.key_fields.len() as i32); + for (index, field) in self.key_fields.iter().enumerate() { + if let Some(value) = + crate::spec::extract_datum_from_arrow(&batch, 0, index, field.data_type())? + { + row.write_datum(index, &value, field.data_type()); + } else { + row.set_null_at(index); + } + } + row.build_serialized() + }; + for files in self.levels.values() { + let mut found = None; + for file in files { + // Old key types may have been promoted. Unverifiable ranges + // must not prune a possible history match. + if file.meta.schema_id == self.table.schema().id() { + if let Some(compare) = &self.comparator { + if compare.compare(&file.meta.max_key, &binary_key) + == Some(std::cmp::Ordering::Less) + || compare.compare(&file.meta.min_key, &binary_key) + == Some(std::cmp::Ordering::Greater) + { + continue; + } + } + } + let identity = file + .identity + .get_or_try_init(|| async { + let reader = self + .table + .file_io() + .new_input(&self.split.data_file_path(&file.meta))? + .reader_with_size(file.meta.file_size as u64) + .await?; + let namespace = reader + .cache_key() + .map(str::to_string) + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let mut physical = file.meta.clone(); + physical.level = 0; + let mut digest = Sha256::new(); + digest.update(b"paimon-lookup-sst-v1"); + digest.update((namespace.len() as u64).to_le_bytes()); + digest.update(namespace.as_bytes()); + digest.update(serde_json::to_vec(&physical).map_err(invalid)?); + digest.update(&self.schema_identity); + Ok::<_, Error>(hex::encode(digest.finalize())) + }) + .await?; + let row = self + .cache + .lookup(identity, key, &self.keys, || { + let split = DataSplit::builder() + .with_snapshot(self.split.snapshot_id()) + .with_partition(self.split.partition().clone()) + .with_bucket(self.split.bucket()) + .with_total_buckets(self.split.total_buckets()) + .with_bucket_path(self.split.bucket_path().to_string()) + .with_data_files(vec![file.meta.clone()]) + .build()?; + // Reuse the normal KV decoder for schema evolution, file + // formats, and preservation of system sequence/row kind. + super::new_reader(&self.table, self.resources.clone())? + .read_with_merge_function(&[split], |_, _| { + Ok(Box::new(super::CompactDeduplicate)) + }) + }) + .await?; + if let Some(batch) = row { + if found.is_some() { + return Err(invalid( + "Lookup history has overlapping keys within a positive level", + )); + } + let mut metadata = batch.schema().metadata().clone(); + metadata.insert(FILE_LEVEL_METADATA_KEY.into(), file.meta.level.to_string()); + found = Some( + RecordBatch::try_new( + Arc::new(batch.schema().as_ref().clone().with_metadata(metadata)), + batch.columns().to_vec(), + ) + .map_err(invalid)?, + ); + } + } + if found.is_some() { + return Ok(found); + } + } + Ok(None) + } +} + +fn invalid(message: impl ToString) -> Error { + Error::DataInvalid { + message: message.to_string(), + source: None, + } +} diff --git a/crates/paimon/src/table/compact/primary_key/planner.rs b/crates/paimon/src/table/compact/primary_key/planner.rs index f519949d2..5a09538ab 100644 --- a/crates/paimon/src/table/compact/primary_key/planner.rs +++ b/crates/paimon/src/table/compact/primary_key/planner.rs @@ -29,7 +29,7 @@ pub(in crate::table::compact) fn enabled(schema: &TableSchema) -> Result { && !options.deletion_vectors_enabled() && matches!( options.try_changelog_producer()?, - ChangelogProducer::None | ChangelogProducer::Input + ChangelogProducer::None | ChangelogProducer::Input | ChangelogProducer::Lookup )) } diff --git a/crates/paimon/src/table/kv_file_reader.rs b/crates/paimon/src/table/kv_file_reader.rs index 09c373bf0..6b2bda833 100644 --- a/crates/paimon/src/table/kv_file_reader.rs +++ b/crates/paimon/src/table/kv_file_reader.rs @@ -64,6 +64,7 @@ pub(crate) struct KeyValueFileReader { /// filter using the full `config.predicates` instead. pushdown_predicates: Vec, include_file_levels: bool, + lookup: Option>, #[cfg(test)] input_batch_sizes: Option>>>, } @@ -360,12 +361,22 @@ impl KeyValueFileReader { file_io, config, include_file_levels, + lookup: None, pushdown_predicates, #[cfg(test)] input_batch_sizes: None, } } + pub(super) fn with_lookup( + mut self, + lookup: Option>, + ) -> Self { + self.include_file_levels |= lookup.is_some(); + self.lookup = lookup; + self + } + #[cfg(test)] fn with_input_batch_sizes( mut self, @@ -680,6 +691,7 @@ impl KeyValueFileReader { }; let file_io = self.file_io; let include_file_levels = self.include_file_levels; + let lookup = self.lookup; let config = self.config; let table_schema_id = config.table_schema_id; let pushdown_predicates = self.pushdown_predicates; @@ -865,6 +877,7 @@ impl KeyValueFileReader { merge_output_schema.clone(), merge_function(&config, &merge_output_fields)?, ) + .with_lookup(lookup.clone()) .with_user_sequence_descending( !CoreOptions::new(&config.table_options) .sequence_field_sort_order_is_ascending(), diff --git a/crates/paimon/src/table/merge_tree_split_generator.rs b/crates/paimon/src/table/merge_tree_split_generator.rs index 6eee84028..acd472623 100644 --- a/crates/paimon/src/table/merge_tree_split_generator.rs +++ b/crates/paimon/src/table/merge_tree_split_generator.rs @@ -72,6 +72,10 @@ impl KeyComparator { Some(Self::new(key_types)) } + pub(crate) fn compare(&self, left: &[u8], right: &[u8]) -> Option { + Some(compare_decoded(&self.decode(left)?, &self.decode(right)?)) + } + /// Decode a serialized min/max key. Returns `None` when the key is empty, /// malformed, or contains NaN, letting callers degrade to the safe "treat /// everything as overlapping" path instead of failing the scan. NaN must diff --git a/crates/paimon/src/table/sort_merge.rs b/crates/paimon/src/table/sort_merge.rs index 492dcd81a..f9068ef51 100644 --- a/crates/paimon/src/table/sort_merge.rs +++ b/crates/paimon/src/table/sort_merge.rs @@ -166,6 +166,12 @@ pub(crate) trait MergeFunction: Send + Sync { } } +/// Point lookup into retained sorted runs during incremental compaction. +#[async_trait::async_trait] +pub(crate) trait LookupProvider: Send + Sync { + async fn lookup(&self, key: &[u8]) -> crate::Result>; +} + /// Internal batch metadata used only by lookup compaction. Source levels /// survive physical reads; changelog batches are routed to a separate writer. pub(super) const FILE_LEVEL_METADATA_KEY: &str = "paimon.compact.file-level"; @@ -1344,6 +1350,7 @@ pub(crate) struct SortMergeReaderBuilder { output_schema: SchemaRef, merge_function: Box, batch_size: usize, + lookup: Option>, } impl SortMergeReaderBuilder { @@ -1371,9 +1378,15 @@ impl SortMergeReaderBuilder { output_schema, merge_function, batch_size: 1024, + lookup: None, } } + pub(crate) fn with_lookup(mut self, lookup: Option>) -> Self { + self.lookup = lookup; + self + } + pub(crate) fn with_batch_size(mut self, batch_size: usize) -> Self { self.batch_size = batch_size; self @@ -1431,6 +1444,7 @@ impl SortMergeReaderBuilder { self.output_schema, self.merge_function, self.batch_size, + self.lookup.map(|lookup| (lookup, self.input_schema)), ) } } @@ -1504,6 +1518,7 @@ fn sort_merge_stream( output_schema: SchemaRef, merge_function: Box, batch_size: usize, + lookup: Option<(Arc, SchemaRef)>, ) -> crate::Result { let num_streams = streams.len(); if num_streams == 0 { @@ -1640,6 +1655,33 @@ fn sort_merge_stream( tree.update(|a, b| compare_cursors(&cursors, a, b).then_with(|| a.cmp(&b)).is_gt()); } + if let Some((lookup, schema)) = &lookup { + let has_base = same_key_rows.iter().any(|row| match &batch_buffer[row.batch_idx] { + BufferedBatch::Source(batch) => batch.schema().metadata().get(FILE_LEVEL_METADATA_KEY) + .and_then(|level| level.parse::().ok()).is_some_and(|level| level > 0), + _ => false, + }); + if !has_base { + if let Some(batch) = lookup.lookup(&winner_key).await? { + if batch.num_rows() != 1 || batch.schema().fields() != schema.fields() { + Err(Error::DataInvalid { message: "Lookup must return one KV row matching the merge schema".into(), source: None })?; + } + let keys = convert_batch_keys(&batch, &key_indices, &mut row_converter)?; + if keys.row(0).as_ref() != winner_key.as_slice() { + Err(Error::DataInvalid { message: "Lookup returned a different primary key".into(), source: None })?; + } + let user = convert_batch_user_sequences(&batch, &user_sequence_indices, user_sequence_converter.as_mut())?; + let sequence_number = batch.column(seq_index).as_any().downcast_ref::() + .ok_or_else(|| Error::DataInvalid { message: "Invalid lookup sequence".into(), source: None })?.value(0); + let value_kind = batch.column(value_kind_index).as_any().downcast_ref::() + .ok_or_else(|| Error::DataInvalid { message: "Invalid lookup row kind".into(), source: None })?.value(0); + let batch_idx = batch_buffer.len(); + batch_buffer.push(BufferedBatch::Source(batch)); + same_key_rows.insert(0, MergeRow { batch_idx, row_idx: 0, sequence_number, value_kind, + user_sequence: user.map(|rows| rows.row(0).owned()) }); + } + } + } let (merged, changelog) = merge_function.merge_with_changelog( &same_key_rows, &batch_buffer, &source_output_col_indices, &output_schema, )?; diff --git a/crates/paimon/src/table/table_commit.rs b/crates/paimon/src/table/table_commit.rs index 007fb1b74..3558763df 100644 --- a/crates/paimon/src/table/table_commit.rs +++ b/crates/paimon/src/table/table_commit.rs @@ -1937,6 +1937,13 @@ impl TableCommit { }; validate_file_entries(entries.iter())?; validate_file_entries(changelog_entries.iter())?; + let lookup = self + .table + .schema() + .core_options() + .try_changelog_producer()? + == crate::spec::ChangelogProducer::Lookup; + let mut lookup_history = Vec::new(); // 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 { @@ -1947,7 +1954,7 @@ impl TableCommit { source: None, }); } - let captured_files = super::snapshot_manifest::read_snapshot_entries( + let mut captured_files = super::snapshot_manifest::read_snapshot_entries( &self.snapshot_manager, &snapshot, None, @@ -1986,7 +1993,7 @@ impl TableCommit { ) .await?; } - let captured_files = captured_files + let captured_inputs = captured_files .iter() .filter(|e| { buckets.contains(&(e.partition(), e.bucket())) @@ -1994,9 +2001,9 @@ impl TableCommit { }) .map(|entry| (commit_file_identifier(entry), entry)) .collect::>(); - if captured_files.len() != deletes.len() + if captured_inputs.len() != deletes.len() || deletes.iter().any(|deleted| { - !captured_files + !captured_inputs .get(&commit_file_identifier(deleted)) .is_some_and(|file| { file.file() == deleted.file() @@ -2013,6 +2020,14 @@ impl TableCommit { source: None, }); } + drop(captured_inputs); + if lookup { + captured_files.retain(|entry| { + entry.file().level > 0 + && buckets.contains(&(entry.partition(), entry.bucket())) + }); + lookup_history = captured_files; + } } let base_data_files = self .detect_commit_conflicts( @@ -2057,49 +2072,49 @@ impl TableCommit { base_data_files.as_deref().unwrap_or_default(), entries, )?; - if self - .table - .schema() - .core_options() - .try_changelog_producer()? - == crate::spec::ChangelogProducer::Lookup - { + if lookup { let buckets = entries .iter() .filter(|entry| *entry.kind() == FileKind::Delete) .map(|entry| (entry.partition(), entry.bucket())) .collect::>(); - for (partition, bucket) in buckets { - let expected = entries - .iter() - .filter(|entry| { - *entry.kind() == FileKind::Delete - && entry.partition() == partition - && entry.bucket() == bucket - && entry.file().level > 0 - }) - .collect::>(); - let current = base_data_files - .as_deref() - .unwrap_or_default() - .iter() - .filter(|entry| { - entry.partition() == partition - && entry.bucket() == bucket - && entry.file().level > 0 - }) - .collect::>(); - if expected.len() != current.len() - || expected.iter().any(|file| { - !current.iter().any(|entry| entry.file() == file.file()) - }) - { - return Err(crate::Error::DataInvalid { - message: "Lookup compaction history changed; create a new plan" - .into(), - source: None, - }); - } + let current = base_data_files + .as_deref() + .unwrap_or_default() + .iter() + .filter(|entry| { + buckets.contains(&(entry.partition(), entry.bucket())) + && entry.file().level > 0 + }) + .map(|entry| { + ( + ( + entry.partition(), + entry.bucket(), + entry.file().file_name.as_str(), + ), + entry, + ) + }) + .collect::>(); + if lookup_history.len() != current.len() + || lookup_history.iter().any(|file| { + current + .get(&( + file.partition(), + file.bucket(), + file.file().file_name.as_str(), + )) + .is_none_or(|entry| { + entry.file() != file.file() + || entry.total_buckets() != file.total_buckets() + }) + }) + { + return Err(crate::Error::DataInvalid { + message: "Lookup compaction history changed; create a new plan".into(), + source: None, + }); } } Ok(ResolvedCommit { diff --git a/crates/paimon/tests/pk_compact/incremental_lookup.rs b/crates/paimon/tests/pk_compact/incremental_lookup.rs new file mode 100644 index 000000000..ce3869117 --- /dev/null +++ b/crates/paimon/tests/pk_compact/incremental_lookup.rs @@ -0,0 +1,410 @@ +// 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::{ + cast::AsArray, + types::{Int32Type, Int64Type}, +}; +use opendal::Operator; +use paimon::io::FileIOProvider; +use paimon::table::CommitMessage; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +#[derive(Debug)] +struct HistoryProbe { + operator: Operator, + history: Mutex, + opens: AtomicUsize, +} + +#[async_trait::async_trait] +impl FileIOProvider for HistoryProbe { + async fn create(&self, path: &str) -> paimon::Result<(Operator, String)> { + if path.rsplit('/').next() == Some(self.history.lock().unwrap().as_str()) { + self.opens.fetch_add(1, Ordering::SeqCst); + } + Ok(( + self.operator.clone(), + path.strip_prefix("memory:/").unwrap().into(), + )) + } + + async fn create_with_cache_namespace( + &self, + path: &str, + ) -> paimon::Result<(Operator, String, Option)> { + let (operator, path) = self.create(path).await?; + let namespace = format!("{:p}", Arc::as_ptr(operator.service())); + Ok((operator, path, Some(namespace))) + } +} + +async fn probed_table() -> (Table, Arc) { + let path = format!("memory:/incremental_lookup/{}", uuid::Uuid::new_v4()); + let mut options = paimon::Options::new(); + options.set("warehouse", &path); + options.set("local-cache.enabled", "true"); + // Ordinary data blocks cannot explain the warm lookup result. + options.set("local-cache.whitelist", "meta"); + let catalog = paimon::catalog::FileSystemCatalog::new(options).unwrap(); + let probe = Arc::new(HistoryProbe { + operator: Operator::from_config(opendal::services::MemoryConfig::default()).unwrap(), + history: Mutex::new(String::new()), + opens: AtomicUsize::new(0), + }); + let io = catalog.file_io().clone().with_provider(probe.clone()); + let table = Table::new( + io.clone(), + paimon::catalog::Identifier::new("default", "incremental_lookup"), + path.clone(), + pk_schema(&[("changelog-producer", "lookup")]), + None, + ); + setup_dirs(&io, &path).await; + persist_table_schema(&io, &path, table.schema()).await; + (table, probe) +} + +async fn compact(table: &Table) -> Vec { + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let results = execute(table, &plan).await; + builder + .new_commit() + .unwrap() + .commit_compact(results.clone()) + .await + .unwrap(); + results +} + +async fn seed(table: &Table) -> paimon::spec::DataFileMeta { + let keys = (0..2000).step_by(2).collect::>(); + write_batch(table, &make_batch(keys.clone(), keys)).await; + let anchor = compact(table).await[0].files()[0].clone(); + assert_eq!(anchor.level, 5); + anchor +} + +async fn latest(table: &Table) -> i64 { + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap() + .unwrap() +} + +async fn events(table: &Table, start: i64, end: i64) -> Vec<(String, i32, i64, i64)> { + let value_name = table.schema().fields()[1].name(); + let mut reader = table.new_read_builder(); + reader + .with_projection(&["id", value_name, "_SEQUENCE_NUMBER"]) + .unwrap(); + let plan = reader + .new_incremental_scan(IncrementalScanMode::Auto, start, end) + .plan_combined() + .await + .unwrap(); + let batches: Vec = reader + .new_read() + .unwrap() + .to_arrow_with_row_kind(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut result = Vec::new(); + for batch in batches { + let kinds = batch.column(0).as_string::(); + let column = |name| batch.column_by_name(name).unwrap(); + let ids = column("id").as_primitive::(); + let values = arrow_cast::cast(column(value_name), &arrow_schema::DataType::Int64).unwrap(); + let values = values.as_primitive::(); + let seqs = column("_SEQUENCE_NUMBER").as_primitive::(); + result.extend((0..batch.num_rows()).map(|i| { + ( + kinds.value(i).to_string(), + ids.value(i), + values.value(i), + seqs.value(i), + ) + })); + } + result.sort(); + result +} + +async fn cold_and_hot( + table: &Table, + plan: &CompactPlan, + probe: &HistoryProbe, + anchor: &paimon::spec::DataFileMeta, +) -> (Vec, Vec) { + probe.opens.store(0, Ordering::SeqCst); + let cold = execute(table, plan).await; + assert!( + probe.opens.load(Ordering::SeqCst) > 1, + "cold lookup loads history data" + ); + let history_path = format!("{}/bucket-0/{}", table.location(), anchor.file_name); + let relative = history_path.strip_prefix("memory:/").unwrap(); + let backup = probe.operator.read(relative).await.unwrap(); + probe.operator.delete(relative).await.unwrap(); + probe.opens.store(0, Ordering::SeqCst); + // A new executor must consume the generated lookup file: the original data + // object is absent, and ordinary data caching is disabled by the whitelist. + let hot = execute(table, plan).await; + assert_eq!( + probe.opens.load(Ordering::SeqCst), + 1, + "only the identity opener may resolve history" + ); + probe.operator.write(relative, backup).await.unwrap(); + (cold, hot) +} + +#[tokio::test] +async fn incremental_lookup_consumes_cached_history_across_executors_and_retries() { + let (table, probe) = probed_table().await; + let anchor = seed(&table).await; + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let historical = table.copy_with_pinned_snapshot(Some(&snapshot)); + let initial = (0..1000).map(|row| (2 * row, 2 * row)).collect::>(); + *probe.history.lock().unwrap() = anchor.file_name.clone(); + write_batch( + &table, + &make_batch_with_kinds(vec![42, 44, 43], vec![4200, 999, 4300], vec![0, 3, 0]), + ) + .await; + let builder = table + .new_compact_builder() + .with_commit_user("lookup-cache-retry") + .unwrap(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert_eq!(plan.tasks().len(), 1); + assert!(plan.tasks()[0].files().iter().all(|file| file.level == 0)); + let (cold, hot) = cold_and_hot(&table, &plan, &probe, &anchor).await; + assert_eq!(cold[0].metrics().changelog_rows_written, 4); + assert_eq!(hot[0].metrics().changelog_rows_written, 4); + assert!(hot[0].files().iter().all(|file| file.level == 4)); + write_batch(&table, &make_batch(vec![42], vec![4201])).await; + let commit = builder.new_commit().unwrap(); + let mut legacy = serde_json::to_value(&hot).unwrap(); + for result in legacy.as_array_mut().unwrap() { + result["task"] + .as_object_mut() + .unwrap() + .remove("output_level"); + } + commit + .commit_compact(serde_json::from_value(legacy).unwrap()) + .await + .unwrap(); + let published = latest(&table).await; + commit.commit_compact(cold).await.unwrap(); + commit.commit_compact(hot).await.unwrap(); + assert_eq!(latest(&table).await, published); + let mut expected = vec![ + ("-U".into(), 42, 42, 21), + ("+U".into(), 42, 4200, 1000), + ("-D".into(), 44, 44, 22), + ("+I".into(), 43, 4300, 1002), + ]; + expected.sort(); + assert_eq!(events(&table, snapshot.id(), published).await, expected); + let mut final_rows = initial.clone(); + final_rows.retain(|(key, _)| *key != 44); + final_rows.iter_mut().find(|(key, _)| *key == 42).unwrap().1 = 4201; + final_rows.push((43, 4300)); + final_rows.sort(); + assert_eq!(rows(&table).await, final_rows); + assert_eq!(rows(&historical).await, initial); + let scan = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan() + .await + .unwrap(); + assert!(scan + .splits() + .iter() + .flat_map(|split| split.data_files()) + .any(|file| file == &anchor)); + compact(&table).await; + let last = latest(&table).await; + assert_eq!( + events(&table, published, last).await, + vec![("+U".into(), 42, 4201, 1003), ("-U".into(), 42, 4200, 1000)] + ); + assert_eq!(rows(&table).await, final_rows); +} + +#[tokio::test] +async fn incremental_lookup_casts_old_history_after_rename_and_type_promotion() { + use paimon::spec::{BigIntType, SchemaChange}; + let (original, probe) = probed_table().await; + let anchor = seed(&original).await; + *probe.history.lock().unwrap() = anchor.file_name.clone(); + write_batch(&original, &make_batch(vec![42], vec![4200])).await; + compact(&original).await; // Prime the same physical history under schema 0. + let start = latest(&original).await; + let schema = original + .schema() + .apply_changes(vec![ + SchemaChange::rename_column("value".into(), "renamed".into()), + SchemaChange::update_column_type("renamed".into(), DataType::BigInt(BigIntType::new())), + ]) + .unwrap(); + persist_table_schema(original.file_io(), original.location(), &schema).await; + let table = original.copy_with_resolved_schema(schema, "main").unwrap(); + write_batch( + &table, + &RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int32Array::from(vec![42, 43])) as arrow_array::ArrayRef, + ), + ( + "renamed", + Arc::new(Int64Array::from(vec![4_300_000_000, 4301])) as arrow_array::ArrayRef, + ), + ]) + .unwrap(), + ) + .await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + let (cold, hot) = cold_and_hot(&table, &plan, &probe, &anchor).await; + assert_eq!(cold[0].metrics().changelog_rows_written, 3); + assert_eq!(hot[0].metrics().changelog_rows_written, 3); + assert!(hot[0] + .files() + .iter() + .all(|file| file.schema_id == table.schema().id())); + builder + .new_commit() + .unwrap() + .commit_compact(hot) + .await + .unwrap(); + let mut expected = vec![ + ("-U".into(), 42, 4200, 1000), + ("+U".into(), 42, 4_300_000_000, 1001), + ("+I".into(), 43, 4301, 1002), + ]; + expected.sort(); + assert_eq!(events(&table, start, latest(&table).await).await, expected); + let mut actual = Vec::new(); + for batch in read(&table, None).await { + assert!(batch.column_by_name("value").is_none()); + let ids = batch + .column_by_name("id") + .unwrap() + .as_primitive::(); + let values = batch + .column_by_name("renamed") + .unwrap() + .as_primitive::(); + actual.extend((0..batch.num_rows()).map(|row| (ids.value(row), values.value(row)))); + } + actual.sort(); + let mut expected = (0..1000) + .map(|row| (2 * row, i64::from(2 * row))) + .collect::>(); + expected[21].1 = 4_300_000_000; + expected.push((43, 4301)); + expected.sort(); + assert_eq!(actual, expected); +} + +#[tokio::test] +async fn incremental_lookup_fences_positive_history_for_hits_and_absent_keys() { + for key in [42, 43] { + for mutation in ["add", "remove", "metadata"] { + let table = table(&[("changelog-producer", "lookup")]).await; + let anchor = seed(&table).await; + write_batch(&table, &make_batch(vec![key], vec![9000])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert!(plan.tasks()[0].files().iter().all(|file| file.level == 0)); + let results = execute(&table, &plan).await; + let writer_builder = table.new_write_builder(); + let mut changes = Vec::new(); + if mutation == "add" { + let mut writer = writer_builder.new_write().unwrap(); + writer + .write_arrow_batch(&make_batch(vec![key], vec![7777])) + .await + .unwrap(); + changes = writer.prepare_commit().await.unwrap(); + for file in changes + .iter_mut() + .flat_map(|message| &mut message.new_files) + { + file.level = 3; + } + writer.close().await; + } else { + let mut removal = + CommitMessage::new(plan.tasks()[0].partition().to_vec(), 0, vec![]); + removal.deleted_files.push(anchor.clone()); + changes.push(removal); + if mutation == "metadata" { + // The public writer API rejects same-identity DELETE+ADD + // in one commit. Remove and re-add across two snapshots so + // the latest state contains the same file with new metadata. + table + .new_write_builder() + .new_commit() + .commit(std::mem::take(&mut changes)) + .await + .unwrap(); + let mut modified = anchor.clone(); + assert!(modified.creation_time.take().is_some()); + changes.push(CommitMessage::new( + plan.tasks()[0].partition().to_vec(), + 0, + vec![modified], + )); + } + } + writer_builder.new_commit().commit(changes).await.unwrap(); + let before = latest(&table).await; + assert!( + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap_err() + .to_string() + .contains("history changed"), + "{mutation}, key {key}" + ); + assert_eq!(latest(&table).await, before, "{mutation}, key {key}"); + } + } +} diff --git a/crates/paimon/tests/pk_compact_test.rs b/crates/paimon/tests/pk_compact_test.rs index 4021bf428..776aaddcf 100644 --- a/crates/paimon/tests/pk_compact_test.rs +++ b/crates/paimon/tests/pk_compact_test.rs @@ -26,6 +26,9 @@ mod lookup_changelog_tests; #[path = "pk_compact/incremental_planning.rs"] mod incremental_planning_tests; +#[path = "pk_compact/incremental_lookup.rs"] +mod incremental_lookup_tests; + #[path = "pk_compact/dynamic_bucket.rs"] mod dynamic_tests; diff --git a/docs/src/compaction.md b/docs/src/compaction.md index d2ac5ab9b..f90dc68ef 100644 --- a/docs/src/compaction.md +++ b/docs/src/compaction.md @@ -20,7 +20,7 @@ under the License. # Table Compaction Rust supports small-file compaction for ordinary append tables, normal-file -compaction for Data Evolution tables, and full-bucket compaction for primary-key +compaction for Data Evolution tables, and sorted-run 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 @@ -115,8 +115,8 @@ snapshot expiration and storage reclamation are separate maintenance work. The same API compacts primary-key tables with `deduplicate`, `first-row`, `partial-update` or `aggregation` merge engines. Changelog producers `none`, `input` and `lookup` are supported; `first-row` supports `none` or `lookup`. -For `deduplicate` with `none` or `input`, the planner chooses incremental sorted -runs using Java's universal-compaction size and count policy. Each L0 file is a +For `deduplicate` with `none`, `input` or `lookup`, the planner chooses incremental +sorted runs using Java's universal-compaction size and count policy. Each L0 file is a run; all files at one positive level form a single run. Normal size and run-count triggers take precedence. If they select no work, explicit compaction forces progress on every L0 run and absorbs adjacent runs according to their sizes. A @@ -133,7 +133,7 @@ otherwise it goes immediately below the next unselected run. Explicit full compaction is not exposed. Outputs retain winning retracts, so older values in untouched files do not reappear after a delete or out-of-order update. -`lookup`, `first-row`, `partial-update` and `aggregation` retain their current +`first-row`, `partial-update` and `aggregation` retain their current selection behavior: every active file in a selected bucket participates, and output uses the highest input level or level 1 for all-L0 inputs. Their state and changelog generation must see history before incremental selection can be @@ -246,14 +246,34 @@ Changelog files roll at the ordinary size and row limits. `CompactResult` expose `changelog_files_added`, `changelog_rows_written` and `changelog_bytes_written`. The existing data metrics continue to count data files and rows only. +For `deduplicate`, unselected higher-level history is read through point lookups +in disposable local Lookup Files. Each file stores ordered keys and complete KV +rows, including row kinds and sequence metadata. The first access builds it by +streaming the source through the normal schema-evolution decoder. Later keys and +executors reuse it; a missing or corrupt cache block causes a bounded rebuild +instead of being treated as a missing key. Only selected runs are rewritten. + The executor streams both outputs and aborts both writers on execution failure. Prepared results include both file lists for serialization, subset commits and -idempotent retries. A concurrent positive-level addition in a selected bucket -rejects the commit because it may change the changelog's history. Concurrent L0 +idempotent retries. Any concurrent positive-level file addition, removal or +metadata change in a selected bucket rejects the commit because it may change +the changelog's history, including previous lookup misses. Concurrent L0 appends remain eligible for a later compaction and produce no premature events. The executor requires an actual catalog/FileIO cache (`local-cache.enabled=true`); setting the flag in table properties alone does not attach a cache. +Deduplicate Lookup Files use that cache's memory, disk or shared-disk backend, +namespace, eviction policy and `local-cache.max-size` budget. Set `local-cache.dir` to keep +them on disk. They do not require `data` in `local-cache.whitelist`. +`lookup.cache-max-disk-size` is rejected; use `local-cache.max-size` instead. + +For deduplicate lookup caching, `lookup.cache-max-memory-size` defaults to +`256 MiB` per executor and bounds working components such as input batches, +IPC payloads and SST indexes. Oversized inputs or indexes, +or a cache unable to retain the generated file, fail execution without publishing +a snapshot. Reduce `read.batch-size` for oversized input batches or increase the +appropriate budget. Cache storage is charged separately to `local-cache.max-size`. +`lookup.cache-spill-compression` accepts `none`, `zstd` (default), `lz4` or `lzo`. ## Dynamic buckets From 5ebfa962789b93845cb206e63fb6eb8fc428bf9f Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 29 Oct 2026 23:35:34 +0800 Subject: [PATCH 2/4] fix(table): align lookup commit conflicts with Java --- crates/paimon/src/table/table_commit.rs | 68 +--------- .../tests/pk_compact/incremental_lookup.rs | 120 ++++++++---------- .../tests/pk_compact/lookup_changelog.rs | 75 +++++++---- docs/src/compaction.md | 9 +- 4 files changed, 113 insertions(+), 159 deletions(-) diff --git a/crates/paimon/src/table/table_commit.rs b/crates/paimon/src/table/table_commit.rs index 3558763df..ae8bb4ddd 100644 --- a/crates/paimon/src/table/table_commit.rs +++ b/crates/paimon/src/table/table_commit.rs @@ -1937,13 +1937,6 @@ impl TableCommit { }; validate_file_entries(entries.iter())?; validate_file_entries(changelog_entries.iter())?; - let lookup = self - .table - .schema() - .core_options() - .try_changelog_producer()? - == crate::spec::ChangelogProducer::Lookup; - let mut lookup_history = Vec::new(); // 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 { @@ -1954,7 +1947,7 @@ impl TableCommit { source: None, }); } - let mut captured_files = super::snapshot_manifest::read_snapshot_entries( + let captured_files = super::snapshot_manifest::read_snapshot_entries( &self.snapshot_manager, &snapshot, None, @@ -1993,7 +1986,7 @@ impl TableCommit { ) .await?; } - let captured_inputs = captured_files + let captured_files = captured_files .iter() .filter(|e| { buckets.contains(&(e.partition(), e.bucket())) @@ -2001,9 +1994,9 @@ impl TableCommit { }) .map(|entry| (commit_file_identifier(entry), entry)) .collect::>(); - if captured_inputs.len() != deletes.len() + if captured_files.len() != deletes.len() || deletes.iter().any(|deleted| { - !captured_inputs + !captured_files .get(&commit_file_identifier(deleted)) .is_some_and(|file| { file.file() == deleted.file() @@ -2020,14 +2013,6 @@ impl TableCommit { source: None, }); } - drop(captured_inputs); - if lookup { - captured_files.retain(|entry| { - entry.file().level > 0 - && buckets.contains(&(entry.partition(), entry.bucket())) - }); - lookup_history = captured_files; - } } let base_data_files = self .detect_commit_conflicts( @@ -2072,51 +2057,6 @@ impl TableCommit { base_data_files.as_deref().unwrap_or_default(), entries, )?; - if lookup { - let buckets = entries - .iter() - .filter(|entry| *entry.kind() == FileKind::Delete) - .map(|entry| (entry.partition(), entry.bucket())) - .collect::>(); - let current = base_data_files - .as_deref() - .unwrap_or_default() - .iter() - .filter(|entry| { - buckets.contains(&(entry.partition(), entry.bucket())) - && entry.file().level > 0 - }) - .map(|entry| { - ( - ( - entry.partition(), - entry.bucket(), - entry.file().file_name.as_str(), - ), - entry, - ) - }) - .collect::>(); - if lookup_history.len() != current.len() - || lookup_history.iter().any(|file| { - current - .get(&( - file.partition(), - file.bucket(), - file.file().file_name.as_str(), - )) - .is_none_or(|entry| { - entry.file() != file.file() - || entry.total_buckets() != file.total_buckets() - }) - }) - { - return Err(crate::Error::DataInvalid { - message: "Lookup compaction history changed; create a new plan".into(), - source: None, - }); - } - } Ok(ResolvedCommit { entries: entries.clone(), changelog_entries: changelog_entries.clone(), diff --git a/crates/paimon/tests/pk_compact/incremental_lookup.rs b/crates/paimon/tests/pk_compact/incremental_lookup.rs index ce3869117..30e407371 100644 --- a/crates/paimon/tests/pk_compact/incremental_lookup.rs +++ b/crates/paimon/tests/pk_compact/incremental_lookup.rs @@ -341,70 +341,62 @@ async fn incremental_lookup_casts_old_history_after_rename_and_type_promotion() } #[tokio::test] -async fn incremental_lookup_fences_positive_history_for_hits_and_absent_keys() { - for key in [42, 43] { - for mutation in ["add", "remove", "metadata"] { - let table = table(&[("changelog-producer", "lookup")]).await; - let anchor = seed(&table).await; - write_batch(&table, &make_batch(vec![key], vec![9000])).await; - let builder = table.new_compact_builder(); - let plan = builder.new_planner().unwrap().plan().await.unwrap(); - assert!(plan.tasks()[0].files().iter().all(|file| file.level == 0)); - let results = execute(&table, &plan).await; - let writer_builder = table.new_write_builder(); - let mut changes = Vec::new(); - if mutation == "add" { - let mut writer = writer_builder.new_write().unwrap(); - writer - .write_arrow_batch(&make_batch(vec![key], vec![7777])) - .await - .unwrap(); - changes = writer.prepare_commit().await.unwrap(); - for file in changes - .iter_mut() - .flat_map(|message| &mut message.new_files) - { - file.level = 3; - } - writer.close().await; - } else { - let mut removal = - CommitMessage::new(plan.tasks()[0].partition().to_vec(), 0, vec![]); - removal.deleted_files.push(anchor.clone()); - changes.push(removal); - if mutation == "metadata" { - // The public writer API rejects same-identity DELETE+ADD - // in one commit. Remove and re-add across two snapshots so - // the latest state contains the same file with new metadata. - table - .new_write_builder() - .new_commit() - .commit(std::mem::take(&mut changes)) - .await - .unwrap(); - let mut modified = anchor.clone(); - assert!(modified.creation_time.take().is_some()); - changes.push(CommitMessage::new( - plan.tasks()[0].partition().to_vec(), - 0, - vec![modified], - )); - } - } - writer_builder.new_commit().commit(changes).await.unwrap(); - let before = latest(&table).await; - assert!( - builder - .new_commit() - .unwrap() - .commit_compact(results) - .await - .unwrap_err() - .to_string() - .contains("history changed"), - "{mutation}, key {key}" - ); - assert_eq!(latest(&table).await, before, "{mutation}, key {key}"); +async fn incremental_lookup_allows_safe_rewrites_of_unselected_history() { + for metadata_only in [false, true] { + let table = table(&[("changelog-producer", "lookup")]).await; + let anchor = seed(&table).await; + write_batch(&table, &make_batch(vec![42], vec![9000])).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert!(plan.tasks()[0].files().iter().all(|file| file.level == 0)); + let results = execute(&table, &plan).await; + assert!(results[0].files().iter().all(|file| file.level == 4)); + let mut replacement = anchor.clone(); + let mut change = CommitMessage::new(plan.tasks()[0].partition().to_vec(), 0, vec![]); + if metadata_only { + // Same-identity replacement uses two snapshots with the public API. + change.deleted_files.push(anchor.clone()); + table + .new_write_builder() + .new_commit() + .commit(vec![change]) + .await + .unwrap(); + change = CommitMessage::new(plan.tasks()[0].partition().to_vec(), 0, vec![]); + assert!(replacement.creation_time.take().is_some()); + } else { + replacement.file_name = format!("data-{}-0.parquet", uuid::Uuid::new_v4()); + let bucket = format!("{}/bucket-0", table.location()); + table + .file_io() + .copy_file( + &format!("{bucket}/{}", anchor.file_name), + &format!("{bucket}/{}", replacement.file_name), + ) + .await + .unwrap(); + change.deleted_files.push(anchor); } + change.new_files.push(replacement); + table + .new_write_builder() + .new_commit() + .commit(vec![change]) + .await + .unwrap(); + let expected = rows(&table).await; + let before = latest(&table).await; + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap(); + assert_eq!(latest(&table).await, before + 1); + assert_eq!(rows(&table).await, expected); + assert_eq!( + expected.iter().find(|(key, _)| *key == 42), + Some(&(42, 9000)) + ); } } diff --git a/crates/paimon/tests/pk_compact/lookup_changelog.rs b/crates/paimon/tests/pk_compact/lookup_changelog.rs index 197311666..78666aeea 100644 --- a/crates/paimon/tests/pk_compact/lookup_changelog.rs +++ b/crates/paimon/tests/pk_compact/lookup_changelog.rs @@ -736,36 +736,57 @@ async fn lookup_compaction_subsets_and_retries_publish_each_changelog_once_with_ } #[tokio::test] -async fn lookup_compaction_rejects_concurrent_positive_level_history_at_another_level() { - let table = lookup_table(&[]).await; - write_batch(&table, &make_batch(vec![1], vec![10])).await; - compact(&table).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; - append_at_level(&table, 12, 3).await; - let snapshot = table - .snapshot_manager() - .get_latest_snapshot_id() - .await - .unwrap(); - assert!(builder - .new_commit() - .unwrap() - .commit_compact(results) - .await - .unwrap_err() - .to_string() - .contains("history changed")); - assert_eq!( - table +async fn lookup_compaction_allows_disjoint_history_and_rejects_output_level_overlap() { + for overlaps in [false, true] { + let table = lookup_table(&[]).await; + write_batch(&table, &make_batch(vec![1], vec![10])).await; + compact(&table).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 output_level = results[0].files()[0].level; + assert!(output_level > 1); + let write = table.new_write_builder(); + let mut writer = write.new_write().unwrap(); + writer + .write_arrow_batch(&make_batch(vec![if overlaps { 1 } else { 256 }], vec![12])) + .await + .unwrap(); + let mut messages = writer.prepare_commit().await.unwrap(); + for file in messages + .iter_mut() + .flat_map(|message| &mut message.new_files) + { + file.level = if overlaps { + output_level + } else { + output_level - 1 + }; + } + write.new_commit().commit(messages).await.unwrap(); + writer.close().await; + let snapshot = table .snapshot_manager() .get_latest_snapshot_id() .await - .unwrap(), - snapshot - ); + .unwrap(); + let result = builder.new_commit().unwrap().commit_compact(results).await; + if overlaps { + assert!(result.unwrap_err().to_string().contains("LSM")); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap(), + snapshot + ); + } else { + result.unwrap(); + assert_eq!(rows(&table).await, vec![(1, 11), (256, 12)]); + } + } } #[tokio::test] diff --git a/docs/src/compaction.md b/docs/src/compaction.md index f90dc68ef..03b59ddc3 100644 --- a/docs/src/compaction.md +++ b/docs/src/compaction.md @@ -255,10 +255,11 @@ instead of being treated as a missing key. Only selected runs are rewritten. The executor streams both outputs and aborts both writers on execution failure. Prepared results include both file lists for serialization, subset commits and -idempotent retries. Any concurrent positive-level file addition, removal or -metadata change in a selected bucket rejects the commit because it may change -the changelog's history, including previous lookup misses. Concurrent L0 -appends remain eligible for a later compaction and produce no premature events. +idempotent retries. Commit validation checks that selected inputs remain valid +and that files in each positive level have disjoint key ranges, matching Java's +LSM conflict checks. Concurrent L0 appends and changes to unselected files are +allowed when these checks pass. New L0 files remain eligible for a later +compaction and produce no premature events. The executor requires an actual catalog/FileIO cache (`local-cache.enabled=true`); setting the flag in table properties alone does not attach a cache. From 1494a7a6d5b33384ef620b9e3b104d3e1ba836d4 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Fri, 30 Oct 2026 00:09:13 +0800 Subject: [PATCH 3/4] fix(table): preserve lookup pruning across schema evolution --- crates/paimon/src/table/compact.rs | 17 +-- .../table/compact/primary_key/lookup_store.rs | 104 +++++++++++++++--- .../tests/pk_compact/incremental_lookup.rs | 93 +++++++++++++++- 3 files changed, 188 insertions(+), 26 deletions(-) diff --git a/crates/paimon/src/table/compact.rs b/crates/paimon/src/table/compact.rs index cad166b29..3ec88e043 100644 --- a/crates/paimon/src/table/compact.rs +++ b/crates/paimon/src/table/compact.rs @@ -413,13 +413,16 @@ impl CompactExecutor { if history.is_empty() { None } else { - Some(primary_key::lookup_store::provider( - &table, - &split, - history, - cache.clone(), - self.resources.clone(), - )?) + Some( + primary_key::lookup_store::provider( + &table, + &split, + history, + cache.clone(), + self.resources.clone(), + ) + .await?, + ) } } else { None diff --git a/crates/paimon/src/table/compact/primary_key/lookup_store.rs b/crates/paimon/src/table/compact/primary_key/lookup_store.rs index 128518461..f215df20f 100644 --- a/crates/paimon/src/table/compact/primary_key/lookup_store.rs +++ b/crates/paimon/src/table/compact/primary_key/lookup_store.rs @@ -26,18 +26,25 @@ use crate::{Error, Result}; use arrow_array::RecordBatch; use arrow_row::{RowConverter, SortField}; use sha2::{Digest, Sha256}; +use std::cmp::Ordering; use std::collections::BTreeMap; use std::sync::{Arc, Mutex}; struct HistoryFile { meta: DataFileMeta, + range_compatible: bool, identity: tokio::sync::OnceCell, } +struct HistoryLevel { + files: Vec, + indexed: bool, +} + struct LevelLookup { table: Table, split: DataSplit, - levels: BTreeMap>, + levels: BTreeMap, cache: Arc, keys: Vec, key_fields: Vec, @@ -47,7 +54,7 @@ struct LevelLookup { schema_identity: Vec, } -pub(in crate::table::compact) fn provider( +pub(in crate::table::compact) async fn provider( table: &Table, split: &DataSplit, files: Vec, @@ -79,12 +86,59 @@ pub(in crate::table::compact) fn provider( .collect(), ) .map_err(invalid)?; - let mut levels = BTreeMap::<_, Vec<_>>::new(); + let comparator = KeyComparator::from_table_schema(table.schema()); + let mut schemas = BTreeMap::from([(table.schema().id(), true)]); + let mut levels = BTreeMap::<_, HistoryLevel>::new(); for meta in files { - levels.entry(meta.level).or_default().push(HistoryFile { - meta, - identity: tokio::sync::OnceCell::new(), - }); + let compatible = if let Some(&compatible) = schemas.get(&meta.schema_id) { + compatible + } else { + let schema = table.schema_manager().schema(meta.schema_id).await?; + let old_keys = schema.trimmed_primary_keys(); + let compatible = old_keys.len() == key_fields.len() + && old_keys.iter().zip(&key_fields).all(|(name, current)| { + schema + .fields() + .iter() + .find(|field| field.name() == name) + .is_some_and(|old| { + old.id() == current.id() + && old.data_type().equals_ignore_nullable(current.data_type()) + }) + }); + schemas.insert(meta.schema_id, compatible); + compatible + }; + let range_compatible = compatible + && comparator.as_ref().is_some_and(|compare| { + compare + .compare(&meta.min_key, &meta.max_key) + .is_some_and(|order| !order.is_gt()) + }); + levels + .entry(meta.level) + .or_insert_with(|| HistoryLevel { + files: Vec::new(), + indexed: false, + }) + .files + .push(HistoryFile { + meta, + range_compatible, + identity: tokio::sync::OnceCell::new(), + }); + } + for level in levels.values_mut() { + if level.files.iter().all(|file| file.range_compatible) { + let compare = comparator.as_ref().unwrap(); + level + .files + .sort_by(|a, b| compare.compare(&a.meta.min_key, &b.meta.min_key).unwrap()); + level.indexed = level.files.windows(2).all(|pair| { + compare.compare(&pair[0].meta.max_key, &pair[1].meta.min_key) + == Some(Ordering::Less) + }); + } } let options = table.schema().options().iter().collect::>(); let schema_identity = serde_json::to_vec(&( @@ -102,7 +156,7 @@ pub(in crate::table::compact) fn provider( keys, key_fields, converter: Mutex::new(converter), - comparator: KeyComparator::from_table_schema(table.schema()), + comparator, resources, schema_identity, })) @@ -131,19 +185,37 @@ impl LookupProvider for LevelLookup { } row.build_serialized() }; - for files in self.levels.values() { + let key_comparable = self + .comparator + .as_ref() + .is_some_and(|compare| compare.compare(&binary_key, &binary_key).is_some()); + for level in self.levels.values() { + let files = if level.indexed && key_comparable { + let compare = self.comparator.as_ref().unwrap(); + let index = level.files.partition_point(|file| { + compare.compare(&file.meta.max_key, &binary_key) == Some(Ordering::Less) + }); + level + .files + .get(index) + .map(std::slice::from_ref) + .unwrap_or_default() + } else { + level.files.as_slice() + }; let mut found = None; for file in files { // Old key types may have been promoted. Unverifiable ranges // must not prune a possible history match. - if file.meta.schema_id == self.table.schema().id() { + if file.range_compatible { if let Some(compare) = &self.comparator { - if compare.compare(&file.meta.max_key, &binary_key) - == Some(std::cmp::Ordering::Less) - || compare.compare(&file.meta.min_key, &binary_key) - == Some(std::cmp::Ordering::Greater) - { - continue; + if let (Some(min), Some(max)) = ( + compare.compare(&file.meta.min_key, &binary_key), + compare.compare(&file.meta.max_key, &binary_key), + ) { + if min.is_gt() || max.is_lt() { + continue; + } } } } diff --git a/crates/paimon/tests/pk_compact/incremental_lookup.rs b/crates/paimon/tests/pk_compact/incremental_lookup.rs index 30e407371..abc7e44a5 100644 --- a/crates/paimon/tests/pk_compact/incremental_lookup.rs +++ b/crates/paimon/tests/pk_compact/incremental_lookup.rs @@ -23,6 +23,7 @@ use arrow_array::{ use opendal::Operator; use paimon::io::FileIOProvider; use paimon::table::CommitMessage; +use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; @@ -31,11 +32,20 @@ struct HistoryProbe { operator: Operator, history: Mutex, opens: AtomicUsize, + data_opens: Mutex>, } #[async_trait::async_trait] impl FileIOProvider for HistoryProbe { async fn create(&self, path: &str) -> paimon::Result<(Operator, String)> { + if path.ends_with(".parquet") { + *self + .data_opens + .lock() + .unwrap() + .entry(path.into()) + .or_default() += 1; + } if path.rsplit('/').next() == Some(self.history.lock().unwrap().as_str()) { self.opens.fetch_add(1, Ordering::SeqCst); } @@ -56,6 +66,10 @@ impl FileIOProvider for HistoryProbe { } async fn probed_table() -> (Table, Arc) { + probed_table_with_options(&[]).await +} + +async fn probed_table_with_options(properties: &[(&str, &str)]) -> (Table, Arc) { let path = format!("memory:/incremental_lookup/{}", uuid::Uuid::new_v4()); let mut options = paimon::Options::new(); options.set("warehouse", &path); @@ -67,13 +81,16 @@ async fn probed_table() -> (Table, Arc) { operator: Operator::from_config(opendal::services::MemoryConfig::default()).unwrap(), history: Mutex::new(String::new()), opens: AtomicUsize::new(0), + data_opens: Mutex::new(HashMap::new()), }); let io = catalog.file_io().clone().with_provider(probe.clone()); + let mut configured = vec![("changelog-producer", "lookup")]; + configured.extend_from_slice(properties); let table = Table::new( io.clone(), paimon::catalog::Identifier::new("default", "incremental_lookup"), path.clone(), - pk_schema(&[("changelog-producer", "lookup")]), + pk_schema(&configured), None, ); setup_dirs(&io, &path).await; @@ -112,10 +129,11 @@ async fn latest(table: &Table) -> i64 { } async fn events(table: &Table, start: i64, end: i64) -> Vec<(String, i32, i64, i64)> { + let id_name = table.schema().fields()[0].name(); let value_name = table.schema().fields()[1].name(); let mut reader = table.new_read_builder(); reader - .with_projection(&["id", value_name, "_SEQUENCE_NUMBER"]) + .with_projection(&[id_name, value_name, "_SEQUENCE_NUMBER"]) .unwrap(); let plan = reader .new_incremental_scan(IncrementalScanMode::Auto, start, end) @@ -134,7 +152,7 @@ async fn events(table: &Table, start: i64, end: i64) -> Vec<(String, i32, i64, i for batch in batches { let kinds = batch.column(0).as_string::(); let column = |name| batch.column_by_name(name).unwrap(); - let ids = column("id").as_primitive::(); + let ids = column(id_name).as_primitive::(); let values = arrow_cast::cast(column(value_name), &arrow_schema::DataType::Int64).unwrap(); let values = values.as_primitive::(); let seqs = column("_SEQUENCE_NUMBER").as_primitive::(); @@ -400,3 +418,72 @@ async fn incremental_lookup_allows_safe_rewrites_of_unselected_history() { ); } } + +#[tokio::test] +async fn incremental_lookup_prunes_old_schema_history_by_unchanged_keys() { + use paimon::spec::SchemaChange; + for change in [ + SchemaChange::rename_column("value".into(), "renamed".into()), + SchemaChange::add_column("extra".into(), DataType::Int(IntType::new())), + SchemaChange::rename_column("id".into(), "renamed_id".into()), + ] { + let (original, probe) = probed_table_with_options(&[("target-file-row-num", "100")]).await; + let keys = (0..1000).collect::>(); + write_batch(&original, &make_batch(keys.clone(), keys)).await; + let history = compact(&original).await[0].files().to_vec(); + assert!( + history.len() >= 5, + "regression requires rolled positive-level files" + ); + assert!(history.iter().all(|file| file.level == 5)); + let start = latest(&original).await; + let schema = original.schema().apply_changes(vec![change]).unwrap(); + persist_table_schema(original.file_io(), original.location(), &schema).await; + let table = original.copy_with_resolved_schema(schema, "main").unwrap(); + let mut columns = vec![ + ( + table.schema().fields()[0].name(), + Arc::new(Int32Array::from(vec![42])) as arrow_array::ArrayRef, + ), + ( + table.schema().fields()[1].name(), + Arc::new(Int32Array::from(vec![9000])) as arrow_array::ArrayRef, + ), + ]; + for field in &table.schema().fields()[2..] { + columns.push(( + field.name(), + Arc::new(Int32Array::from(vec![None])) as arrow_array::ArrayRef, + )); + } + write_batch(&table, &RecordBatch::try_from_iter(columns).unwrap()).await; + let builder = table.new_compact_builder(); + let plan = builder.new_planner().unwrap().plan().await.unwrap(); + assert!(plan.tasks()[0].files().iter().all(|file| file.level == 0)); + probe.data_opens.lock().unwrap().clear(); + let results = execute(&table, &plan).await; + let touched = { + let opens = probe.data_opens.lock().unwrap(); + history + .iter() + .filter(|file| { + opens.contains_key(&format!("{}/bucket-0/{}", table.location(), file.file_name)) + }) + .count() + }; + assert_eq!( + touched, 1, + "unchanged key layouts must only read the candidate history file" + ); + builder + .new_commit() + .unwrap() + .commit_compact(results) + .await + .unwrap(); + assert_eq!( + events(&table, start, latest(&table).await).await, + vec![("+U".into(), 42, 9000, 1000), ("-U".into(), 42, 42, 42)] + ); + } +} From 4e97a7735cbc208af94ecc34f6c9dd5af65ef690 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Fri, 30 Oct 2026 00:09:23 +0800 Subject: [PATCH 4/4] fix(table): reuse lookup indexes and reclaim failed cache builds --- crates/paimon/src/io/cache/generated.rs | 242 +++++++++++++-- crates/paimon/src/io/cache/shared.rs | 39 +++ .../table/compact/primary_key/lookup_file.rs | 289 ++++++++++++++++-- docs/src/compaction.md | 4 + 4 files changed, 533 insertions(+), 41 deletions(-) diff --git a/crates/paimon/src/io/cache/generated.rs b/crates/paimon/src/io/cache/generated.rs index c1e8b8874..3bb5845c8 100644 --- a/crates/paimon/src/io/cache/generated.rs +++ b/crates/paimon/src/io/cache/generated.rs @@ -17,11 +17,12 @@ //! Disposable immutable files stored in the ordinary cache backend and budget. -use super::LocalCache; +use super::{CacheBackend, LocalCache}; use crate::io::{FileRead, FileWrite}; use crate::{Error, Result}; use bytes::{Bytes, BytesMut}; use std::ops::Range; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; impl LocalCache { @@ -62,6 +63,106 @@ impl LocalCache { let _guard = self.shared_load_lock(&self.block_key(path, 0)).await; self.generated_put(path, 0, descriptor).await } + + async fn generated_remove(&self, path: &str, blocks: u64) { + let _guard = self + .coordinator + .begin_path_invalidation(&self.namespace, path) + .await; + match &self.backend { + CacheBackend::Memory(cache) => cache.invalidate_path(&self.namespace, path), + CacheBackend::Disk(cache) => { + // Retry by exact key: path invalidation drops its in-memory + // index before unlinking, so cancelling it could hide orphans. + for index in 0..blocks { + cache.remove_block(&self.block_key(path, index)).await; + } + } + CacheBackend::SharedDisk(cache) => { + if let Err(error) = cache + .remove_generated(self.block_key(path, 0), blocks) + .await + { + log::debug!("Failed to remove unpublished Lookup File: {error}"); + } + } + CacheBackend::External(cache) => cache.invalidate_path(path).await, + } + } +} + +struct GeneratedFileState { + cache: Arc, + path: String, + write_gate: Arc>, + attempted_blocks: AtomicU64, + #[cfg(test)] + cleaned: Arc, +} + +impl GeneratedFileState { + async fn cleanup(&self) { + // A cancelled caller must wait for any already-started backend put. + let _guard = self.write_gate.lock().await; + self.cache + .generated_remove(&self.path, self.attempted_blocks.load(Ordering::SeqCst)) + .await; + #[cfg(test)] + self.cleaned.notify_one(); + } +} + +/// Owns an unpublished generation until its head has been published. +pub(crate) struct GeneratedFile { + state: Arc, + runtime: Option, + settled: bool, +} + +impl GeneratedFile { + pub(crate) fn new(cache: Arc, path: String) -> Self { + Self { + state: Arc::new(GeneratedFileState { + cache, + path, + write_gate: Arc::new(tokio::sync::Mutex::new(())), + attempted_blocks: AtomicU64::new(0), + #[cfg(test)] + cleaned: Arc::new(tokio::sync::Notify::new()), + }), + runtime: tokio::runtime::Handle::try_current().ok(), + settled: false, + } + } + + pub(crate) fn writer(&self) -> GeneratedFileWrite { + GeneratedFileWrite { + state: self.state.clone(), + buffer: BytesMut::new(), + index: 0, + length: 0, + } + } + + pub(crate) fn published(&mut self) { + self.settled = true; + } + + pub(crate) async fn cleanup(&mut self) { + self.state.cleanup().await; + self.settled = true; + } +} + +impl Drop for GeneratedFile { + fn drop(&mut self) { + if !self.settled { + let state = self.state.clone(); + if let Some(runtime) = &self.runtime { + runtime.spawn(async move { state.cleanup().await }); + } + } + } } #[cfg(test)] @@ -145,30 +246,47 @@ impl FileRead for GeneratedFileRead { } pub(crate) struct GeneratedFileWrite { - cache: Arc, - path: String, + state: Arc, buffer: BytesMut, index: u64, length: u64, } impl GeneratedFileWrite { - pub(crate) fn new(cache: Arc, path: String) -> Self { - Self { - cache, - path, - buffer: BytesMut::new(), - index: 0, - length: 0, - } - } - async fn flush_block(&mut self) -> Result<()> { if !self.buffer.is_empty() { - self.cache - .generated_put(&self.path, self.index, self.buffer.split().freeze()) - .await?; + let index = self.index; self.index += 1; + // Include the in-flight block even if cancellation happens before + // the backend's publication/verification returns. + self.state + .attempted_blocks + .store(self.index, Ordering::SeqCst); + let data = self.buffer.split().freeze(); + let state = self.state.clone(); + // Cancelled callers may leave their current backend put running. + // Reuse the cache's bounded fill slots instead of accumulating tasks. + let slot = state + .cache + .fill_slots + .clone() + .acquire_owned() + .await + .map_err(|error| Error::UnexpectedError { + message: "Lookup File write capacity was closed".into(), + source: Some(Box::new(error)), + })?; + let gate = state.write_gate.clone().lock_owned().await; + tokio::spawn(async move { + let _slot = slot; + let _gate = gate; + state.cache.generated_put(&state.path, index, data).await + }) + .await + .map_err(|error| Error::UnexpectedError { + message: "Lookup File cache write task failed".into(), + source: Some(Box::new(error)), + })??; } Ok(()) } @@ -181,7 +299,7 @@ impl FileWrite for GeneratedFileWrite { .length .checked_add(bytes.len() as u64) .ok_or_else(|| invalid("Generated cache file length overflow"))?; - let size = usize::try_from(self.cache.block_size) + let size = usize::try_from(self.state.cache.block_size) .map_err(|_| invalid("Generated cache block exceeds the address space"))?; while !bytes.is_empty() { let take = (size - self.buffer.len()).min(bytes.len()); @@ -196,11 +314,12 @@ impl FileWrite for GeneratedFileWrite { async fn close(&mut self) -> Result<()> { self.flush_block().await?; for index in 0..self.index { - let length = - (self.length - index * self.cache.block_size).min(self.cache.block_size) as usize; + let length = (self.length - index * self.state.cache.block_size) + .min(self.state.cache.block_size) as usize; if self + .state .cache - .generated_get(&self.path, index, length) + .generated_get(&self.state.path, index, length) .await .is_none() { @@ -221,3 +340,86 @@ fn invalid(message: &str) -> Error { source: None, } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn unpublished_blocks_are_reclaimed_without_invalidating_other_files_or_the_epoch() { + for backend in ["memory", "disk", "shared"] { + let directory = tempfile::tempdir().unwrap(); + let mut options = crate::Options::new(); + options.set("local-cache.enabled", "true"); + options.set("local-cache.block-size", "8kb"); + if backend != "memory" { + options.set("local-cache.dir", directory.path().to_str().unwrap()); + } + if backend == "shared" { + options.set("local-cache.shared", "true"); + options.set("local-cache.max-size", "4mb"); + } + // Memory and ordinary disk deliberately have no capacity limit. + let local = super::super::create_local_cache(&options).unwrap().unwrap(); + let bytes = Bytes::from(vec![7; 8192]); + let mut published = GeneratedFile::new(local.clone(), "published".into()); + let mut writer = published.writer(); + writer.write(bytes.clone()).await.unwrap(); + writer.close().await.unwrap(); + local + .generated_publish("published/head", Bytes::from(vec![0; 24])) + .await + .unwrap(); + published.published(); + drop(published); + let epoch = local.read_token("published").shared_generation; + let usage = directory.path().join("paimon-shared-local-cache-v1/usage"); + let used = std::fs::read_to_string(&usage).ok(); + for attempt in 0..3 { + let path = format!("unpublished-{attempt}"); + let mut file = GeneratedFile::new(local.clone(), path.clone()); + let cleaned = file.state.cleaned.clone(); + let mut writer = file.writer(); + writer + .write(Bytes::from(vec![9; 3 * 8192 + 1])) + .await + .unwrap(); + for index in 0..3 { + assert!(local.generated_get(&path, index, 8192).await.is_some()); + } + drop(writer); + if attempt == 1 { + // Cancellation drops the generation owner after full blocks + // have already reached its backend. + std::thread::spawn(move || drop(file)).join().unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(5), cleaned.notified()) + .await + .unwrap(); + } else { + file.cleanup().await; + } + for index in 0..3 { + assert!( + local.generated_get(&path, index, 8192).await.is_none(), + "{backend}" + ); + if backend == "disk" { + assert!(!directory + .path() + .join("paimon-local-cache-v2") + .join(local.block_key(&path, index).cache_relative_path()) + .exists()); + } + } + assert_eq!( + local.generated_get("published", 0, 8192).await, + Some(bytes.clone()) + ); + assert!(local.generated_get("published/head", 0, 24).await.is_some()); + assert_eq!(local.read_token("published").shared_generation, epoch); + assert_eq!(std::fs::read_to_string(&usage).ok(), used); + assert_eq!(local.fill_slots.available_permits(), 8); + } + } + } +} diff --git a/crates/paimon/src/io/cache/shared.rs b/crates/paimon/src/io/cache/shared.rs index 56db9c2ad..4ced74492 100644 --- a/crates/paimon/src/io/cache/shared.rs +++ b/crates/paimon/src/io/cache/shared.rs @@ -480,6 +480,45 @@ impl SharedDiskCache { } } + /// Remove only one unpublished generated file, without rotating the shared epoch. + pub(crate) async fn remove_generated(&self, mut key: BlockKey, blocks: u64) -> io::Result<()> { + let gate = lock(self.root.join("directory.lock")) + .await + .ok_or_else(|| { + io::Error::other("Shared generated-file cleanup could not acquire directory lock") + })?; + let root = self.root.clone(); + let max_size = self.max_size; + let ratio = self.high_priority_ratio; + tokio::task::spawn_blocking(move || -> io::Result<()> { + let _gate = gate; + let Some(generation) = read_generation(&root) else { + return Ok(()); + }; + let mut used = clean_usage(&root, max_size) + .map(Ok) + .unwrap_or_else(|| recover_usage(&root, max_size, ratio))?; + std::fs::write(root.join("dirty"), b"")?; + for index in 0..blocks { + key.block_index = index; + let path = root.join(generation_key(&key, &generation).cache_relative_path()); + let size = match std::fs::metadata(&path) { + Ok(meta) => meta.len(), + Err(error) if error.kind() == io::ErrorKind::NotFound => continue, + Err(error) => return Err(error), + }; + match std::fs::remove_file(path) { + Ok(()) => used = used.saturating_sub(size), + Err(error) if error.kind() == io::ErrorKind::NotFound => (), + Err(error) => return Err(error), + } + } + finish_usage(&root, used) + }) + .await + .map_err(io::Error::other)? + } + pub(crate) async fn invalidate(&self) -> io::Result<()> { let gate = lock(self.root.join("directory.lock")) .await diff --git a/crates/paimon/src/table/compact/primary_key/lookup_file.rs b/crates/paimon/src/table/compact/primary_key/lookup_file.rs index 574828fdb..2c8d94326 100644 --- a/crates/paimon/src/table/compact/primary_key/lookup_file.rs +++ b/crates/paimon/src/table/compact/primary_key/lookup_file.rs @@ -22,7 +22,7 @@ use crate::btree::{ BlockHandle, BlockReader, SstFileWriter, }; use crate::common::options::parse_memory_size; -use crate::io::cache::generated::{GeneratedFileRead, GeneratedFileWrite}; +use crate::io::cache::generated::{GeneratedFile, GeneratedFileRead}; use crate::io::cache::LocalCache; use crate::io::FileRead; use crate::table::ArrowRecordBatchStream; @@ -46,6 +46,13 @@ struct Payload { batch: Arc, } +struct ParsedIndex { + path: String, + size: u64, + footer_start: u64, + block: BlockReader, +} + pub(in crate::table::compact) struct LookupFileCache { local: Arc, limit: usize, @@ -53,6 +60,8 @@ pub(in crate::table::compact) struct LookupFileCache { // One executor cache owns this working budget. Serialize both builds and // warm index/IPC reads so different files cannot multiply that budget. working: tokio::sync::Mutex<()>, + // One parsed index, charged to the same working budget as the IPC payload. + index: Mutex>>, // One decoded IPC batch across all files, never one cache per file/key. payload: Mutex>, } @@ -104,6 +113,7 @@ impl LookupFileCache { limit, compression, working: tokio::sync::Mutex::new(()), + index: Mutex::new(None), payload: Mutex::new(None), }) } @@ -136,6 +146,7 @@ impl LookupFileCache { Err(error) if attempt == 1 => return Err(error), Err(_) => { *self.payload.lock().unwrap() = None; + *self.index.lock().unwrap() = None; // A new immutable generation replaces the head only after // close. Do not invalidate a shared backend's global epoch. } @@ -147,13 +158,33 @@ impl LookupFileCache { async fn build( &self, head: &str, - mut input: ArrowRecordBatchStream, + input: ArrowRecordBatchStream, keys: &[usize], ) -> Result { + *self.index.lock().unwrap() = None; + *self.payload.lock().unwrap() = None; let generation = uuid::Uuid::new_v4(); let path = format!("{}/{generation}", head.strip_suffix("/head").unwrap()); + let mut file = GeneratedFile::new(self.local.clone(), path); + let result = self.build_file(head, input, keys, &file, generation).await; + if result.is_ok() { + file.published(); + } else { + file.cleanup().await; + } + result + } + + async fn build_file( + &self, + head: &str, + mut input: ArrowRecordBatchStream, + keys: &[usize], + file: &GeneratedFile, + generation: uuid::Uuid, + ) -> Result { let mut writer = SstFileWriter::with_compression_level( - Box::new(GeneratedFileWrite::new(self.local.clone(), path)), + Box::new(file.writer()), 32 * 1024, self.compression, 1, @@ -305,18 +336,46 @@ impl LookupFileCache { read_block_from_bytes(&bytes, handle.size).map_err(invalid) } - async fn read(&self, path: &str, size: u64, key: &[u8]) -> Result> { + async fn parsed_index(&self, path: &str, size: u64) -> Result> { + { + let mut cached = self.index.lock().unwrap(); + if let Some(index) = cached + .as_ref() + .filter(|index| index.path == path && index.size == size) + { + return Ok(index.clone()); + } + // Drop the previous file before allocating its replacement. + *cached = None; + } let footer_start = size .checked_sub(BTreeFileFooter::ENCODED_LENGTH as u64) .ok_or_else(|| invalid("Truncated Lookup File"))?; let reader = GeneratedFileRead::new(self.local.clone(), path.into(), size); let footer = BTreeFileFooter::read_footer(&reader.read(footer_start..size).await?) .map_err(invalid)?; - let index = self + let block = self .block(&reader, footer_start, footer.index_block_handle) .await?; + if block.retained_bytes() > self.limit { + return Err(exhausted("Lookup parsed index exceeds the memory budget")); + } + let index = Arc::new(ParsedIndex { + path: path.into(), + size, + footer_start, + block, + }); + *self.index.lock().unwrap() = Some(index.clone()); + Ok(index) + } + + async fn read(&self, path: &str, size: u64, key: &[u8]) -> Result> { + let index = self.parsed_index(path, size).await?; + let footer_start = index.footer_start; + let reader = GeneratedFileRead::new(self.local.clone(), path.into(), size); let cmp = |a: &[u8], b: &[u8]| Ok(a.cmp(b)); - let (_, mut candidates) = index.seek_and_iter(key, &cmp).map_err(invalid)?; + let (_, mut candidates) = index.block.seek_and_iter(key, &cmp).map_err(invalid)?; let Some((_, bytes)) = candidates.next() else { return Ok(None); }; @@ -459,9 +518,93 @@ fn exhausted(message: impl ToString) -> Error { #[cfg(test)] mod tests { use super::*; + use crate::io::{FileBlockCache, FileIOBuilder}; use arrow_array::{ArrayRef, Int32Array, Int64Array, Int8Array, StringArray}; use futures::{stream, StreamExt}; - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + #[derive(Debug, Default)] + struct CountingBlocks { + blocks: Arc>>, + gets: AtomicUsize, + invalidations: AtomicUsize, + invalidated: tokio::sync::Notify, + reject_head: AtomicBool, + delay_put: AtomicBool, + put_started: Arc, + resume_put: Arc, + } + + #[async_trait::async_trait] + impl FileBlockCache for CountingBlocks { + async fn get(&self, path: &str, range: std::ops::Range) -> Option { + self.gets.fetch_add(1, Ordering::SeqCst); + self.blocks + .lock() + .unwrap() + .get(&(path.into(), range.start)) + .cloned() + .filter(|bytes| bytes.len() as u64 == range.end - range.start) + } + async fn put(&self, path: &str, offset: u64, data: Bytes) { + if path.ends_with("/head") && self.reject_head.load(Ordering::SeqCst) { + return; + } + if self.delay_put.swap(false, Ordering::SeqCst) { + let blocks = self.blocks.clone(); + let key = (path.to_owned(), offset); + let started = self.put_started.clone(); + let resume = self.resume_put.clone(); + // Model a backend write that survives cancellation of its caller. + tokio::spawn(async move { + started.notify_one(); + resume.notified().await; + blocks.lock().unwrap().insert(key, data); + }) + .await + .unwrap(); + } else { + self.blocks + .lock() + .unwrap() + .insert((path.into(), offset), data); + } + } + async fn invalidate_path(&self, path: &str) { + self.blocks + .lock() + .unwrap() + .retain(|(stored, _), _| stored != path); + self.invalidations.fetch_add(1, Ordering::SeqCst); + self.invalidated.notify_one(); + } + async fn invalidate_prefix(&self, prefix: &str) { + self.blocks + .lock() + .unwrap() + .retain(|(stored, _), _| !stored.starts_with(prefix)); + } + } + + fn external() -> (Arc, Arc) { + let blocks = Arc::new(CountingBlocks::default()); + let io = FileIOBuilder::new("memory") + .build() + .unwrap() + .with_file_block_cache(blocks.clone(), 8192, "meta") + .unwrap(); + (io.local_cache().unwrap(), blocks) + } + + async fn wait_for_cleanup(blocks: &CountingBlocks, count: usize) { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while blocks.invalidations.load(Ordering::SeqCst) < count { + blocks.invalidated.notified().await; + } + }) + .await + .unwrap(); + } fn local(directory: Option<&std::path::Path>, shared: bool, capacity: &str) -> Arc { local_with_whitelist(directory, shared, capacity, "meta") @@ -542,6 +685,44 @@ mod tests { .await } + #[tokio::test] + async fn parsed_index_is_reused_and_missing_data_still_rebuilds_its_generation() { + let (local, blocks) = external(); + let cache = LookupFileCache::new(local, &HashMap::new()).unwrap(); + let calls = Arc::new(AtomicUsize::new(0)); + let input = source(calls.clone()); + cache.lookup("index", &key(1), &[2], &input).await.unwrap(); + let before = descriptor(&cache, "index").await.unwrap(); + for rebuilt in [false, true] { + blocks.gets.store(0, Ordering::SeqCst); + for value in [2, 3, 1] { + assert_eq!( + cache + .lookup("index", &key(value), &[2], &input) + .await + .unwrap(), + Some(batch().slice((value - 1) as usize, 1)) + ); + } + // Only the head and target data block are read for each key; the + // footer/index and decoded IPC payload remain resident. + assert_eq!(blocks.gets.load(Ordering::SeqCst), 6); + if !rebuilt { + blocks + .blocks + .lock() + .unwrap() + .retain(|(path, _), _| path.ends_with("/head")); + assert_eq!( + cache.lookup("index", &key(1), &[2], &input).await.unwrap(), + Some(batch().slice(0, 1)) + ); + assert_ne!(descriptor(&cache, "index").await.unwrap(), before); + } + } + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + #[tokio::test] async fn cold_hot_and_concurrent_lookups_retain_complete_kvs_on_all_backends() { for backend in ["memory", "disk", "shared"] { @@ -755,16 +936,71 @@ mod tests { #[tokio::test] async fn source_failure_and_cancellation_never_publish_partial_files() { - let local = local(None, false, "4mb"); - let cache = LookupFileCache::new(local, &HashMap::new()).unwrap(); - let failed = || Ok(stream::iter(vec![Ok(batch()), Err(invalid("source failed"))]).boxed()); - assert!(cache.lookup("failed", &key(1), &[2], failed).await.is_err()); - assert!(descriptor(&cache, "failed").await.is_none()); + let (local, blocks) = external(); + let cache = LookupFileCache::new( + local, + &HashMap::from([("lookup.cache-spill-compression".into(), "none".into())]), + ) + .unwrap(); + let input = source(Arc::new(AtomicUsize::new(0))); + cache + .lookup("survivor", &key(1), &[2], &input) + .await + .unwrap(); + let survivor = descriptor(&cache, "survivor").await.unwrap(); + let large = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int32Array::from((0..4096).collect::>())) as ArrayRef, + ), + ( + "value", + Arc::new(StringArray::from(vec!["x".repeat(64); 4096])) as ArrayRef, + ), + ]) + .unwrap(); + let failed = + || Ok(stream::iter(vec![Ok(large.clone()), Err(invalid("source failed"))]).boxed()); + let only_survivor = || { + assert!(blocks + .blocks + .lock() + .unwrap() + .keys() + .all(|(path, _)| path.starts_with(&format!("{PREFIX}/survivor/")))) + }; + for _ in 0..2 { + assert!(cache + .lookup("failed", &key(1), &[0], &failed) + .await + .is_err()); + assert!(descriptor(&cache, "failed").await.is_none()); + only_survivor(); + } + // Closing the SST is insufficient: a rejected head must reclaim it too. + blocks.reject_head.store(true, Ordering::SeqCst); + assert!(cache + .lookup("head-failed", &key(1), &[2], &input) + .await + .is_err()); + blocks.reject_head.store(false, Ordering::SeqCst); + only_survivor(); let ready = Arc::new(tokio::sync::Notify::new()); let stalled = || { let ready = ready.clone(); - Ok(stream::iter(vec![Ok(batch())]) + let blocks = blocks.clone(); + Ok(stream::iter(vec![Ok(large.clone())]) .chain(stream::once(async move { + assert!( + blocks + .blocks + .lock() + .unwrap() + .keys() + .filter(|(path, _)| path.starts_with(&format!("{PREFIX}/cancelled/"))) + .count() + > 2 + ); ready.notify_one(); futures::future::pending::>().await })) @@ -772,17 +1008,28 @@ mod tests { }; let lookup_key = key(1); tokio::select! { - result = cache.lookup("cancelled", &lookup_key, &[2], stalled) => panic!("unexpected completion: {result:?}"), + result = cache.lookup("cancelled", &lookup_key, &[0], stalled) => panic!("unexpected completion: {result:?}"), _ = ready.notified() => {} } + wait_for_cleanup(&blocks, 4).await; assert!(descriptor(&cache, "cancelled").await.is_none()); - let calls = Arc::new(AtomicUsize::new(0)); - assert!(cache - .lookup("cancelled", &key(1), &[2], source(calls.clone())) - .await - .unwrap() - .is_some()); - assert_eq!(calls.load(Ordering::Relaxed), 1); + only_survivor(); + blocks.delay_put.store(true, Ordering::SeqCst); + tokio::select! { + result = cache.lookup("late-put", &lookup_key, &[0], || Ok(stream::iter(vec![Ok(large.clone())]).boxed())) => panic!("unexpected completion: {result:?}"), + _ = blocks.put_started.notified() => {} + } + blocks.resume_put.notify_one(); + wait_for_cleanup(&blocks, 5).await; + only_survivor(); + assert_eq!(descriptor(&cache, "survivor").await.unwrap(), survivor); + assert_eq!( + cache + .lookup("survivor", &key(2), &[2], &input) + .await + .unwrap(), + Some(batch().slice(1, 1)) + ); } #[tokio::test] diff --git a/docs/src/compaction.md b/docs/src/compaction.md index 03b59ddc3..b936547bc 100644 --- a/docs/src/compaction.md +++ b/docs/src/compaction.md @@ -252,6 +252,10 @@ rows, including row kinds and sequence metadata. The first access builds it by streaming the source through the normal schema-evolution decoder. Later keys and executors reuse it; a missing or corrupt cache block causes a bounded rebuild instead of being treated as a missing key. Only selected runs are rewritten. +Lookup selects candidate files by key range, retaining range pruning across +schema changes that preserve the primary-key field IDs, order and types. The +reader retains one bounded parsed SST index and one decoded IPC batch. +Failed or cancelled builds clean up their unpublished cache blocks. The executor streams both outputs and aborts both writers on execution failure. Prepared results include both file lists for serialization, subset commits and