diff --git a/crates/paimon/src/btree/composite_predicate.rs b/crates/paimon/src/btree/composite_predicate.rs index cb96f5208..0727c9254 100644 --- a/crates/paimon/src/btree/composite_predicate.rs +++ b/crates/paimon/src/btree/composite_predicate.rs @@ -22,7 +22,10 @@ use super::CompositeKeyCodec; use crate::spec::{DataField, Datum, Predicate, PredicateOperator}; use crate::Result; use std::cmp::Ordering; +use std::collections::HashSet; +// Bound Cartesian expansion, while allowing a single IN domain to grow +// linearly with its input, including constant prefixes and suffixes. pub(crate) const MAX_INTERVALS: usize = 256; type Component = Option>; @@ -56,7 +59,7 @@ impl CompositeBound { }) } - fn compare(&self, other: &Self, codec: &CompositeKeyCodec) -> Result { + pub(super) fn compare(&self, other: &Self, codec: &CompositeKeyCodec) -> Result { for i in 0..self.values.len().min(other.values.len()) { let order = codec.compare_component( i, @@ -175,49 +178,69 @@ impl CompositePlan { } else { literals.iter().cloned().map(Some).collect() }; + // Other IN predicates on this column must also avoid a linear + // literal scan for every candidate. Normalized component bytes + // preserve comparison equality, including canonical NaNs. + let in_domains = column + .iter() + .map(|leaf| match leaf { + Predicate::Leaf { + op: PredicateOperator::In, + literals, + .. + } if !std::ptr::eq(leaf, selected) => Some( + literals + .iter() + .map(|value| codec.serialize_component(i, value)) + .collect::>(), + ), + _ => None, + }) + .collect::>(); let mut domain: Vec> = Vec::new(); + let mut seen = HashSet::new(); for candidate in candidates { let encoded = candidate.as_ref().map(|v| codec.serialize_component(i, v)); let mut matches = true; - for leaf in &column { - if !matches_leaf(&codec, i, encoded.as_deref(), leaf)? { + for (leaf, in_domain) in column.iter().zip(&in_domains) { + // Candidates already come from this leaf. Rechecking an + // IN against itself would make planning quadratic. + if std::ptr::eq(leaf, selected) { + continue; + } + let satisfies = if let Some(values) = in_domain { + encoded.as_ref().is_some_and(|value| values.contains(value)) + } else { + matches_leaf(&codec, i, encoded.as_deref(), leaf)? + }; + if !satisfies { matches = false; break; } } - if matches { - let mut duplicate = false; - for value in &domain { - let value = value.as_ref().map(|v| codec.serialize_component(i, v)); - if codec.compare_component(i, encoded.as_deref(), value.as_deref())? - == Ordering::Equal - { - duplicate = true; - break; - } - } - if !duplicate { - domain.push(candidate); - } - if domain.len() > MAX_INTERVALS { - break; - } + if matches && seen.insert(encoded) { + domain.push(candidate); } } - equal_columns += 1; if domain.is_empty() { return Ok(Some(Self { fields: fields.to_vec(), codec, intervals: Vec::new(), - bound_columns: equal_columns, - equal_columns, + bound_columns: equal_columns + 1, + equal_columns: equal_columns + 1, })); } - if prefixes.len() * domain.len() > MAX_INTERVALS { - return Ok(None); + if prefixes.len() > 1 + && domain.len() > 1 + && prefixes.len() > MAX_INTERVALS / domain.len() + { + // Keep the bound prefix and leave this column and its suffix + // for residual filtering, without expanding a Cartesian product. + break; } - let mut expanded = Vec::new(); + equal_columns += 1; + let mut expanded = Vec::with_capacity(prefixes.len() * domain.len()); for prefix in prefixes { for value in &domain { let mut next = prefix.clone(); diff --git a/crates/paimon/src/btree/composite_tests.rs b/crates/paimon/src/btree/composite_tests.rs index ffc090af2..cfe85e844 100644 --- a/crates/paimon/src/btree/composite_tests.rs +++ b/crates/paimon/src/btree/composite_tests.rs @@ -318,12 +318,14 @@ fn composite_interval_expansion_and_residual_columns() { .len(), 1 ); - let too_many = b.is_in("item", (0..257).map(Datum::Int).collect()).unwrap(); + let many = b.is_in("item", (0..257).map(Datum::Int).collect()).unwrap(); let predicate = Predicate::and(vec![ b.equal("category", Datum::String("a".into())).unwrap(), - too_many, + many, ]); - assert!(CompositePlan::plan(&fields, &predicate).unwrap().is_none()); + let plan = CompositePlan::plan(&fields, &predicate).unwrap().unwrap(); + assert_eq!(plan.intervals.len(), 257); + assert_eq!(plan.equal_columns, 2); let predicate = Predicate::and(vec![ b.greater_than("category", Datum::String("a".into())) .unwrap(), @@ -348,6 +350,429 @@ fn composite_interval_expansion_and_residual_columns() { .is_empty()); } +fn integer_fields() -> Vec { + (0..3) + .map(|i| DataField::new(i, format!("k{i}"), DataType::Int(IntType::new()))) + .collect() +} + +async fn reader_for_tuples( + fields: &[DataField], + tuples: &[Vec>], +) -> BTreeIndexReader { + let codec = CompositeKeyCodec::new(fields); + let mut entries = tuples + .iter() + .enumerate() + .map(|(id, tuple)| (codec.serialize(tuple).unwrap(), id as i64)) + .collect::>(); + entries.sort_by(|(left, _), (right, _)| codec.compare_keys(left, right).unwrap()); + let key_type = DataType::Row(RowType::new(fields.to_vec())); + let buf = VecFileWrite::new(); + let mut writer = BTreeIndexWriter::with_comparator_and_options( + Box::new(buf.clone()), + 256, + BlockCompressionType::None, + 1, + false, + make_key_comparator(&key_type), + ); + for (key, id) in entries { + writer.write(Some(&key), id).await.unwrap(); + } + let result = writer.finish().await.unwrap(); + let bytes = Bytes::from(buf.to_vec()); + let size = bytes.len() as u64; + BTreeIndexReader::open( + Box::new(BytesFileRead(bytes)), + size, + &result.meta, + make_key_comparator(&key_type), + ) + .await + .unwrap() +} + +fn row_matches(fields: &[DataField], tuple: &[Option], predicate: &Predicate) -> bool { + let row = BinaryRow::from_datums( + &tuple + .iter() + .zip(fields) + .map(|(value, field)| (value.as_ref(), field.data_type())) + .collect::>(), + ); + crate::spec::eval_row(predicate, &row).unwrap() +} + +async fn assert_prefix_candidates_and_residual( + fields: &[DataField], + tuples: &[Vec>], + predicate: &Predicate, + prefix: &Predicate, + expected_prefixes: usize, +) { + let plan = CompositePlan::plan(fields, predicate).unwrap().unwrap(); + assert_eq!(plan.bound_columns, 1); + assert_eq!(plan.equal_columns, 1); + assert_eq!(plan.intervals.len(), expected_prefixes); + assert!(!plan.is_point_lookup()); + assert!(plan + .intervals + .iter() + .all(|interval| interval.point_key.is_none())); + let reader = reader_for_tuples(fields, tuples).await; + let candidates = reader.query_composite(&plan).await.unwrap(); + let expected = tuples + .iter() + .enumerate() + .filter(|(_, tuple)| row_matches(fields, tuple, predicate)) + .map(|(id, _)| id as u64) + .collect::(); + let expected_candidates = tuples + .iter() + .enumerate() + .filter(|(_, tuple)| row_matches(fields, tuple, prefix)) + .map(|(id, _)| id as u64) + .collect::(); + assert!(!expected.is_empty()); + assert!(expected.is_subset(&candidates)); + assert!(candidates.len() > expected.len()); + assert_eq!(candidates, expected_candidates); + let actual = candidates + .iter() + .filter(|id| row_matches(fields, &tuples[*id as usize], predicate)) + .collect::(); + assert_eq!(actual, expected); +} + +#[tokio::test] +async fn composite_large_in_keeps_every_point_and_null_suffix() { + let fields = integer_fields(); + let b = PredicateBuilder::new(&fields); + // An odd multiplier permutes the complete domain; duplicates must not + // inflate expansion and the final values must not be silently truncated. + let mut literals = (0..4096) + .map(|i| Datum::Int((i * 1021) % 4096)) + .collect::>(); + literals.extend((0..1000).map(Datum::Int)); + let predicate = Predicate::and(vec![ + b.is_in("k0", literals).unwrap(), + b.is_null("k1").unwrap(), + b.equal("k2", Datum::Int(7)).unwrap(), + ]); + let plan = CompositePlan::plan(&fields, &predicate).unwrap().unwrap(); + assert_eq!(plan.intervals.len(), 4096); + assert_eq!(plan.bound_columns, 3); + assert_eq!(plan.equal_columns, 3); + assert!(plan.is_point_lookup()); + assert!(plan + .intervals + .iter() + .all(|interval| interval.point_key.is_some())); + let mut tuples = (0..4096) + .map(|i| vec![Some(Datum::Int(i)), None, Some(Datum::Int(7))]) + .collect::>(); + tuples.extend([ + vec![ + Some(Datum::Int(4095)), + Some(Datum::Int(0)), + Some(Datum::Int(7)), + ], + vec![Some(Datum::Int(4095)), None, Some(Datum::Int(8))], + vec![Some(Datum::Int(5000)), None, Some(Datum::Int(7))], + ]); + let reader = reader_for_tuples(&fields, &tuples).await; + let actual = reader.query_composite(&plan).await.unwrap(); + assert_eq!(actual, (0..4096).collect::()); + assert!(actual.contains(4095)); +} + +#[tokio::test] +async fn composite_cartesian_limit_keeps_left_prefix_and_residual() { + let fields = integer_fields(); + let b = PredicateBuilder::new(&fields); + let prefix = b + .is_in("k0", (0..17).rev().map(Datum::Int).collect()) + .unwrap(); + let predicate = Predicate::and(vec![ + prefix.clone(), + b.is_in("k1", (0..17).map(Datum::Int).collect()).unwrap(), + b.equal("k2", Datum::Int(5)).unwrap(), + ]); + let mut tuples = Vec::new(); + for first in -1..19 { + for second in -1..19 { + for suffix in [4, 5] { + tuples.push(vec![ + Some(Datum::Int(first)), + Some(Datum::Int(second)), + Some(Datum::Int(suffix)), + ]); + } + } + } + assert_prefix_candidates_and_residual(&fields, &tuples, &predicate, &prefix, 17).await; +} + +#[tokio::test] +async fn composite_oversized_prefix_is_preserved_before_next_in() { + let fields = integer_fields(); + let b = PredicateBuilder::new(&fields); + let prefix = b + .is_in("k0", (0..513).rev().map(Datum::Int).collect()) + .unwrap(); + let predicate = Predicate::and(vec![ + prefix.clone(), + b.is_in("k1", vec![Datum::Int(0), Datum::Int(1)]).unwrap(), + b.equal("k2", Datum::Int(5)).unwrap(), + ]); + let mut tuples = Vec::new(); + for first in 0..514 { + for second in [0, 1, 2] { + for suffix in [4, 5] { + tuples.push(vec![ + Some(Datum::Int(first)), + Some(Datum::Int(second)), + Some(Datum::Int(suffix)), + ]); + } + } + } + assert_prefix_candidates_and_residual(&fields, &tuples, &predicate, &prefix, 513).await; +} + +#[tokio::test] +async fn composite_large_in_intersects_other_constraints() { + let fields = integer_fields(); + let b = PredicateBuilder::new(&fields); + let first = b + .is_in("k0", (0..4096).rev().map(Datum::Int).collect()) + .unwrap(); + let mut tuples = (0..8192) + .map(|i| { + vec![ + Some(Datum::Int(i)), + Some(Datum::Int(1)), + Some(Datum::Int(1)), + ] + }) + .collect::>(); + tuples.push(vec![None, Some(Datum::Int(1)), Some(Datum::Int(1))]); + let reader = reader_for_tuples(&fields, &tuples).await; + for (constraint, expected) in [ + ( + Predicate::and(vec![ + b.greater_or_equal("k0", Datum::Int(4089)).unwrap(), + b.less_than("k0", Datum::Int(4096)).unwrap(), + ]), + 4089..4096, + ), + ( + b.is_in("k0", (2048..6144).map(Datum::Int).collect()) + .unwrap(), + 2048..4096, + ), + ] { + let predicate = Predicate::and(vec![ + first.clone(), + constraint, + b.equal("k1", Datum::Int(1)).unwrap(), + b.equal("k2", Datum::Int(1)).unwrap(), + ]); + let plan = CompositePlan::plan(&fields, &predicate).unwrap().unwrap(); + assert!(plan.is_point_lookup()); + assert_eq!( + plan.intervals.len(), + (expected.end - expected.start) as usize + ); + let actual = reader.query_composite(&plan).await.unwrap(); + assert_eq!(actual, expected.collect::()); + } +} + +#[tokio::test] +async fn composite_point_domain_identity_matches_encoded_keys() { + use crate::spec::{DecimalType, DoubleType, FloatType}; + + let fields = vec![ + DataField::new(0, "f".into(), DataType::Float(FloatType::new())), + DataField::new(1, "d".into(), DataType::Double(DoubleType::new())), + DataField::new( + 2, + "n".into(), + DataType::Decimal(DecimalType::new(18, 2).unwrap()), + ), + ]; + let b = PredicateBuilder::new(&fields); + let decimal = |unscaled, scale| Datum::Decimal { + unscaled, + precision: 18, + scale, + }; + let predicate = Predicate::and(vec![ + b.is_in( + "f", + vec![ + Datum::Float(f32::NAN), + Datum::Float(f32::from_bits(0xffc00001)), + Datum::Float(-0.0), + Datum::Float(0.0), + ], + ) + .unwrap(), + b.is_in( + "d", + vec![ + Datum::Double(f64::NAN), + Datum::Double(f64::from_bits(0xfff8000000000001)), + Datum::Double(-0.0), + Datum::Double(0.0), + ], + ) + .unwrap(), + b.is_in("n", vec![decimal(1, 0), decimal(100, 2), decimal(1000, 3)]) + .unwrap(), + ]); + let plan = CompositePlan::plan(&fields, &predicate).unwrap().unwrap(); + assert!(plan.is_point_lookup()); + // Each floating domain has one canonical NaN and two distinct zeros; + // all three decimal literals normalize to the same indexed value. + assert_eq!(plan.intervals.len(), 9); + let mut tuples = Vec::new(); + for float in [f32::NAN, -0.0, 0.0] { + for double in [f64::NAN, -0.0, 0.0] { + tuples.push(vec![ + Some(Datum::Float(float)), + Some(Datum::Double(double)), + Some(decimal(100, 2)), + ]); + } + } + let reader = reader_for_tuples(&fields, &tuples).await; + assert_eq!( + reader.query_composite(&plan).await.unwrap(), + (0..9).collect::() + ); + let narrowed = Predicate::and(vec![ + predicate, + b.is_in( + "f", + vec![Datum::Float(0.0), Datum::Float(f32::from_bits(0xffc00001))], + ) + .unwrap(), + b.is_in("n", vec![decimal(1000, 3)]).unwrap(), + ]); + let plan = CompositePlan::plan(&fields, &narrowed).unwrap().unwrap(); + assert_eq!( + reader.query_composite(&plan).await.unwrap(), + [0, 1, 2, 6, 7, 8] + .into_iter() + .collect::() + ); +} + +#[tokio::test] +async fn composite_lookup_batches_reads_without_block_cache() { + use crate::io::FileRead; + use std::collections::HashMap; + use std::sync::{Arc, Mutex}; + + struct CountedRead { + bytes: Bytes, + reads: Arc>>, + } + + #[async_trait::async_trait] + impl FileRead for CountedRead { + async fn read(&self, range: std::ops::Range) -> crate::Result { + *self + .reads + .lock() + .unwrap() + .entry((range.start, range.end)) + .or_default() += 1; + Ok(self.bytes.slice(range.start as usize..range.end as usize)) + } + } + + let fields = integer_fields(); + let b = PredicateBuilder::new(&fields); + // Cover complete points, prefixes in one block, and shuffled prefixes + // spanning many blocks, with some intervals spanning blocks themselves. + for (complete, block_size) in [(true, 1_000_000), (false, 1_000_000), (false, 512)] { + let predicate = Predicate::and(vec![ + b.is_in( + "k0", + (0..4096).map(|i| Datum::Int((i * 2053) % 4096)).collect(), + ) + .unwrap(), + b.is_null("k1").unwrap(), + if complete { + b.equal("k2", Datum::Int(7)).unwrap() + } else { + b.between("k2", Datum::Int(20), Datum::Int(70)).unwrap() + }, + ]); + let plan = CompositePlan::plan(&fields, &predicate).unwrap().unwrap(); + let codec = CompositeKeyCodec::new(&fields); + let key_type = DataType::Row(RowType::new(fields.clone())); + let buf = VecFileWrite::new(); + let mut writer = BTreeIndexWriter::with_comparator_and_options( + Box::new(buf.clone()), + block_size, + BlockCompressionType::None, + 1, + false, + make_key_comparator(&key_type), + ); + let mut expected = roaring::RoaringTreemap::new(); + let mut row_id = 0; + for id in 0..4096 { + let suffixes = if id % 256 == 0 { 0..101 } else { 7..8 }; + for suffix in suffixes { + let key = codec + .serialize(&[Some(Datum::Int(id)), None, Some(Datum::Int(suffix))]) + .unwrap(); + writer.write(Some(&key), row_id).await.unwrap(); + if (complete && suffix == 7) || (!complete && (20..=70).contains(&suffix)) { + expected.insert(row_id as u64); + } + row_id += 1; + } + } + let result = writer.finish().await.unwrap(); + let bytes = Bytes::from(buf.to_vec()); + let size = bytes.len() as u64; + let reads = Arc::new(Mutex::new(HashMap::new())); + let reader = BTreeIndexReader::open_with_data_block_cache( + Box::new(CountedRead { + bytes, + reads: Arc::clone(&reads), + }), + size, + &result.meta, + make_key_comparator(&key_type), + Arc::new(BTreeDataBlockCache::new(0)), + Arc::from("memory:/composite-batch.btree"), + ) + .await + .unwrap(); + reads.lock().unwrap().clear(); + assert_eq!(reader.query_composite(&plan).await.unwrap(), expected); + let reads = reads.lock().unwrap(); + assert!(!reads.is_empty()); + assert!( + reads.values().all(|count| *count == 1), + "a target block was read repeatedly: complete={complete}, block_size={block_size}" + ); + if block_size == 1_000_000 { + assert_eq!(reads.len(), 1); + } else { + assert!(reads.len() > 1); + } + } +} + fn predicates(fields: &[DataField]) -> Vec { let b = PredicateBuilder::new(fields); let mut predicates = vec![ diff --git a/crates/paimon/src/btree/reader.rs b/crates/paimon/src/btree/reader.rs index 44fe15b60..0cf4f009e 100644 --- a/crates/paimon/src/btree/reader.rs +++ b/crates/paimon/src/btree/reader.rs @@ -226,12 +226,29 @@ impl crate::Result> BTreeIndexReader { &self, plan: &super::CompositePlan, ) -> io::Result { + if plan.is_point_lookup() { + let keys = plan + .intervals + .iter() + .map(|interval| interval.point_key.as_deref().expect("complete tuple point")) + .collect::>(); + return match keys.as_slice() { + [key] => self.query_equal(key).await, + _ => self.query_in(&keys).await, + }; + } + // The planner emits disjoint prefix intervals in input literal order. + // Sort them so scans only move forwards, retaining their last block + // across intervals even when the shared block cache is disabled. + let mut intervals = plan.intervals.iter().collect::>(); + try_sort_by(&mut intervals, |left, right| { + left.lower + .compare(&right.lower, &plan.codec) + .map_err(key_comparison_io_error) + })?; + let mut last_block: Option<(BlockHandle, Arc)> = None; let mut result = RoaringTreemap::new(); - for interval in &plan.intervals { - if let Some(key) = &interval.point_key { - result |= self.query_equal(key).await?; - continue; - } + for interval in intervals { let cmp = |key: &[u8], _: &[u8]| { interval .lower @@ -242,7 +259,18 @@ impl crate::Result> BTreeIndexReader { let mut first = true; while let Some((_, handle_bytes)) = index_iter.next() { let handle = BlockHandle::decode(handle_bytes)?; - let block = self.read_data_block(&handle).await?; + let block = match &last_block { + Some((previous, block)) + if previous.offset == handle.offset && previous.size == handle.size => + { + Arc::clone(block) + } + _ => { + let block = self.read_data_block(&handle).await?; + last_block = Some((handle, Arc::clone(&block))); + block + } + }; let mut offset = if first { first = false; block.seek_and_iter(&[], &cmp)?.1.offset diff --git a/crates/paimon/src/table/global_index_freshness.rs b/crates/paimon/src/table/global_index_freshness.rs new file mode 100644 index 000000000..3430374d2 --- /dev/null +++ b/crates/paimon/src/table/global_index_freshness.rs @@ -0,0 +1,735 @@ +// 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. + +//! Verify that active global indexes still describe the indexed column values. + +use std::collections::{HashMap, HashSet}; + +use super::global_index_types::normalize_queryable_global_index_type; +use super::key_normalization::canonical_partition_key; +use super::{DataSplit, Table}; +use crate::spec::{ + BinaryRow, DataEvolutionIndexSourceMeta, DataField, DataFileMeta, + GlobalIndexColumnUpdateAction, IndexManifest, Snapshot, +}; + +/// Under IGNORE, prove that indexed columns have not changed since the index's +/// source snapshot. Other current update policies bypass this check. +/// `partitions` contains canonical logical partition keys. +pub(super) async fn indexed_columns_unchanged( + table: &Table, + snapshot: &Snapshot, + columns: &[String], + partitions: &HashSet>, +) -> crate::Result { + if table + .schema() + .core_options() + .global_index_column_update_action()? + != GlobalIndexColumnUpdateAction::Ignore + { + return Ok(true); + } + let Some(manifest) = snapshot.index_manifest() else { + return Ok(true); + }; + let path = table.snapshot_manager().manifest_path(manifest); + let entries = IndexManifest::read(table.file_io(), &path).await?; + let fields = table.schema().fields(); + let partition_fields = table.schema().partition_fields(); + let mut indexes = Vec::new(); + for entry in &entries { + if entry.kind != crate::spec::FileKind::Add + || normalize_queryable_global_index_type(&entry.index_file.index_type).is_none() + { + continue; + } + let Some(meta) = &entry.index_file.global_index_meta else { + continue; + }; + let index_partition = if partition_fields.is_empty() { + Vec::new() + } else { + canonical_partition_key( + &BinaryRow::from_serialized_bytes(&entry.partition)?, + &partition_fields, + )? + }; + if !partition_fields.is_empty() && !partitions.contains(&index_partition) { + continue; + } + if meta.row_range_start < 0 || meta.row_range_end < meta.row_range_start { + return Ok(false); + } + let keys = fields + .iter() + .filter(|field| { + columns.iter().any(|column| column == field.name()) + && (field.id() == meta.index_field_id + || meta + .extra_field_ids + .as_ref() + .is_some_and(|ids| ids.contains(&field.id()))) + }) + .collect::>(); + if keys.is_empty() { + continue; + } + let Some(source_id) = meta + .source_meta + .as_deref() + .and_then(|bytes| DataEvolutionIndexSourceMeta::deserialize(bytes).ok()) + .map(|meta| meta.scan_snapshot_id()) + .filter(|id| *id <= snapshot.id()) + else { + return Ok(false); + }; + indexes.push((entry, index_partition, keys, source_id)); + } + if indexes.is_empty() { + return Ok(true); + } + // Inspect every active column file before key predicates or indexes can + // prune away evidence of historical updates. This reads metadata only. + let unindexed = table + .copy_with_pinned_snapshot(Some(snapshot)) + .copy_with_options(HashMap::from([( + "global-index.enabled".into(), + "false".into(), + )])); + let plan = unindexed.new_read_builder().new_scan().plan_all().await?; + let index_groups = indexes + .iter() + .map(|(entry, partition, _, _)| (partition.clone(), entry.bucket)) + .collect(); + let files_by_group = file_range_lookups(plan.splits(), &partition_fields, &index_groups)?; + for (entry, index_partition, keys, source_id) in indexes { + let meta = entry.index_file.global_index_meta.as_ref().unwrap(); + let Some(files) = files_by_group.get(&(index_partition, entry.bucket)) else { + continue; + }; + if !files.all_overlapping_unchanged( + meta.row_range_start, + meta.row_range_end, + table.schema(), + &keys, + source_id, + ) { + return Ok(false); + } + } + Ok(true) +} + +fn file_range_lookups<'a>( + splits: &'a [DataSplit], + partition_fields: &[DataField], + index_groups: &HashSet<(Vec, i32)>, +) -> crate::Result, i32), FileRangeLookup<'a>>> { + let buckets = index_groups + .iter() + .map(|(_, bucket)| *bucket) + .collect::>(); + let mut groups: HashMap<_, Vec<_>> = HashMap::new(); + for split in splits { + if !buckets.contains(&split.bucket()) { + continue; + } + let partition = if partition_fields.is_empty() { + Vec::new() + } else { + canonical_partition_key(split.partition(), partition_fields)? + }; + if !index_groups.contains(&(partition.clone(), split.bucket())) { + continue; + } + groups + .entry((partition, split.bucket())) + .or_default() + .extend(split.data_files()); + } + Ok(groups + .into_iter() + .map(|(group, files)| (group, FileRangeLookup::new(files))) + .collect()) +} + +struct FileRangeLookup<'a> { + ranges: Option>>, + unknown_ranges: Vec<&'a DataFileMeta>, +} + +impl<'a> FileRangeLookup<'a> { + fn new(files: Vec<&'a DataFileMeta>) -> Self { + let mut ranges = Vec::new(); + let mut unknown_ranges = Vec::new(); + for file in files { + if let Some((start, end)) = file.row_id_range() { + ranges.push(FileRange { start, end, file }); + } else { + unknown_ranges.push(file); + } + } + ranges.sort_unstable_by_key(|range| (range.start, range.end)); + Self { + ranges: FileRangeTree::from_sorted(ranges), + unknown_ranges, + } + } + + fn all_overlapping_unchanged( + &self, + start: i64, + end: i64, + schema: &crate::spec::TableSchema, + keys: &[&DataField], + source_id: i64, + ) -> bool { + // A missing range cannot prove disjointness. Retain the existing + // column-level check, which can still accept a payload-only file. + self.unknown_ranges.iter().all(|file| { + file.schema_id == schema.id() + && file_columns_unchanged(file, schema, keys, source_id, false) + }) && self.ranges.as_ref().is_none_or(|ranges| { + ranges.all_overlapping(start, end, &mut |file| { + file.schema_id == schema.id() + && file_columns_unchanged(file, schema, keys, source_id, true) + }) + }) + } +} + +struct FileRange<'a> { + start: i64, + end: i64, + file: &'a DataFileMeta, +} + +/// A centered interval tree reports overlaps in O(log(files) + overlaps). +/// A sorted prefix alone would rescan every preceding file when even one +/// long-lived column file spans many disjoint index shards. +struct FileRangeTree<'a> { + center: i64, + by_start: Vec>, + by_end: Vec, + left: Option>, + right: Option>, + #[cfg(test)] + inspected: std::cell::Cell, +} + +impl<'a> FileRangeTree<'a> { + fn from_sorted(ranges: Vec>) -> Option> { + if ranges.is_empty() { + return None; + } + let center = ranges[ranges.len() / 2].start; + let (mut left, mut right, mut by_start) = (Vec::new(), Vec::new(), Vec::new()); + for range in ranges { + if range.end < center { + left.push(range); + } else if range.start > center { + right.push(range); + } else { + by_start.push(range); + } + } + let mut by_end = (0..by_start.len()).collect::>(); + by_end.sort_unstable_by_key(|index| std::cmp::Reverse(by_start[*index].end)); + Some(Box::new(Self { + center, + by_start, + by_end, + left: Self::from_sorted(left), + right: Self::from_sorted(right), + #[cfg(test)] + inspected: std::cell::Cell::new(0), + })) + } + + fn all_overlapping( + &self, + start: i64, + end: i64, + check: &mut impl FnMut(&DataFileMeta) -> bool, + ) -> bool { + if end < self.center { + for range in &self.by_start { + #[cfg(test)] + self.inspected.set(self.inspected.get() + 1); + if range.start > end { + break; + } + if !check(range.file) { + return false; + } + } + } else if start > self.center { + for index in &self.by_end { + #[cfg(test)] + self.inspected.set(self.inspected.get() + 1); + let range = &self.by_start[*index]; + if range.end < start { + break; + } + if !check(range.file) { + return false; + } + } + } else { + for range in &self.by_start { + #[cfg(test)] + self.inspected.set(self.inspected.get() + 1); + if !check(range.file) { + return false; + } + } + } + (start > self.center + || self + .left + .as_ref() + .is_none_or(|left| left.all_overlapping(start, end, check))) + && (end < self.center + || self + .right + .as_ref() + .is_none_or(|right| right.all_overlapping(start, end, check))) + } + + #[cfg(test)] + fn inspections(&self) -> usize { + self.inspected.get() + + self.left.as_ref().map_or(0, |tree| tree.inspections()) + + self.right.as_ref().map_or(0, |tree| tree.inspections()) + } +} + +fn file_columns_unchanged( + file: &DataFileMeta, + schema: &crate::spec::TableSchema, + keys: &[&DataField], + source_id: i64, + has_row_range: bool, +) -> bool { + // Java aligns column sequences with the physical schema: omitted write + // columns and row-tracking metadata do not occupy a sequence slot. + let fields = schema.data_file_fields(file.write_cols.as_deref()); + if fields + .iter() + .map(DataField::id) + .collect::>() + .len() + != fields.len() + || file.write_cols.as_ref().is_some_and(|columns| { + columns.iter().any(|name| { + !crate::spec::is_row_tracking_column(name) + && name != crate::spec::VALUE_KIND_FIELD_NAME + && !schema.fields().iter().any(|field| field.name() == name) + }) + }) + { + return false; + } + keys.iter().all(|key| { + let Some(position) = fields.iter().position(|field| field.id() == key.id()) else { + return true; + }; + let sequence = file + .column_max_sequence_numbers + .as_ref() + .filter(|sequences| sequences.len() == fields.len()) + .and_then(|sequences| sequences.get(position)) + .copied() + .filter(|sequence| *sequence > 0 && *sequence <= file.max_sequence_number) + .unwrap_or(file.max_sequence_number); + has_row_range && sequence > 0 && sequence <= source_id + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::stats::BinaryTableStats; + use crate::spec::{ + BinaryRowBuilder, DataFileMeta, DataType as PaimonType, IntType, Schema as PaimonSchema, + TableSchema, + }; + + fn test_schema() -> TableSchema { + let schema = PaimonSchema::builder() + .column("key", PaimonType::Int(IntType::new())) + .column("value", PaimonType::Int(IntType::new())) + .build() + .unwrap(); + TableSchema::new(0, &schema) + } + + fn test_file(name: &str, start: i64, end: i64) -> DataFileMeta { + DataFileMeta { + file_name: name.into(), + file_size: 128, + row_count: end - start + 1, + min_key: vec![], + max_key: vec![], + key_stats: BinaryTableStats::empty(), + value_stats: BinaryTableStats::empty(), + min_sequence_number: 1, + max_sequence_number: 1, + schema_id: 0, + level: 0, + extra_files: vec![], + creation_time: None, + delete_row_count: None, + embedded_index: None, + first_row_id: Some(start), + write_cols: None, + external_path: None, + file_source: None, + value_stats_cols: None, + column_max_sequence_numbers: None, + } + } + + #[test] + fn file_range_lookup_preserves_inclusive_nested_and_duplicate_overlaps() { + let intervals = [ + (20, 29), + (0, 100), + (10, 19), + (10, 19), + (5, 9), + (30, 30), + (0, 0), + (20, 20), + (29, 40), + ]; + let files = intervals + .iter() + .enumerate() + .map(|(id, (start, end))| test_file(&id.to_string(), *start, *end)) + .collect::>(); + let lookup = FileRangeLookup::new(files.iter().collect()); + let tree = lookup.ranges.as_ref().unwrap(); + for (start, end) in [ + (0, 0), + (4, 5), + (9, 10), + (19, 20), + (29, 30), + (31, 39), + (101, 102), + (i64::MIN, -1), + ] { + let mut expected = files + .iter() + .filter(|file| { + let (file_start, file_end) = file.row_id_range().unwrap(); + file_start <= end && file_end >= start + }) + .map(|file| file.file_name.clone()) + .collect::>(); + let mut actual = Vec::new(); + assert!(tree.all_overlapping(start, end, &mut |file| { + actual.push(file.file_name.clone()); + true + })); + actual.sort(); + expected.sort(); + assert_eq!(actual, expected, "query=[{start}, {end}]"); + if !expected.is_empty() { + let mut checks = 0; + assert!(!tree.all_overlapping(start, end, &mut |_| { + checks += 1; + false + })); + assert_eq!(checks, 1, "stop at the first stale overlap"); + } + } + assert!(FileRangeLookup::new(Vec::new()).ranges.is_none()); + } + + #[test] + fn file_range_lookup_checks_only_relevant_columns_ranges_and_sources() { + let schema = test_schema(); + let key = &schema.fields()[0]; + let value = &schema.fields()[1]; + let mut updated = test_file("updated", 10, 19); + updated.max_sequence_number = 3; + updated.column_max_sequence_numbers = Some(vec![1, 3]); + let mut malformed = test_file("malformed", 30, 39); + malformed.schema_id = 99; + let lookup = FileRangeLookup::new(vec![&updated, &malformed]); + assert!(lookup.all_overlapping_unchanged(0, 9, &schema, &[key], 1)); + assert!(lookup.all_overlapping_unchanged(10, 19, &schema, &[key], 1)); + assert!(!lookup.all_overlapping_unchanged(10, 19, &schema, &[value], 2)); + assert!(lookup.all_overlapping_unchanged(10, 19, &schema, &[value], 3)); + assert!(!lookup.all_overlapping_unchanged(29, 30, &schema, &[key], 3)); + + let mut unknown = test_file("unknown", 0, 9); + unknown.first_row_id = None; + unknown.write_cols = Some(vec!["value".into()]); + unknown.max_sequence_number = 3; + assert!( + FileRangeLookup::new(vec![&unknown]).all_overlapping_unchanged( + 100, + 109, + &schema, + &[key], + 1 + ) + ); + assert!( + !FileRangeLookup::new(vec![&unknown]).all_overlapping_unchanged( + 100, + 109, + &schema, + &[value], + 3 + ) + ); + unknown.write_cols = None; + assert!( + !FileRangeLookup::new(vec![&unknown]).all_overlapping_unchanged( + 100, + 109, + &schema, + &[key], + 3 + ) + ); + for (first_row_id, row_count) in [(Some(i64::MAX), 2), (Some(0), 0)] { + unknown.first_row_id = first_row_id; + unknown.row_count = row_count; + assert!( + !FileRangeLookup::new(vec![&unknown]).all_overlapping_unchanged( + 100, + 109, + &schema, + &[key], + 3 + ) + ); + } + } + + #[test] + fn file_range_lookup_keeps_partition_bucket_groups_separate() { + let field = DataField::new(0, "part".into(), PaimonType::Int(IntType::new())); + let partition = |value| { + let mut builder = BinaryRowBuilder::new(1); + builder.write_int(0, value); + builder.build() + }; + let split = |part, bucket, file| { + DataSplit::builder() + .with_snapshot(1) + .with_partition(partition(part)) + .with_bucket(bucket) + .with_bucket_path(format!("memory:/part={part}/bucket-{bucket}")) + .with_data_files(vec![file]) + .build() + .unwrap() + }; + let good = test_file("good", 0, 9); + let mut unknown = test_file("unknown", 0, 9); + unknown.first_row_id = None; + let splits = vec![ + split(1, 0, good.clone()), + split(1, 0, good), + split(1, 1, unknown.clone()), + split(2, 0, unknown), + ]; + let index_groups = HashSet::from([ + (partition(1).to_serialized_bytes(), 0), + (partition(1).to_serialized_bytes(), 1), + (partition(2).to_serialized_bytes(), 0), + ]); + let groups = file_range_lookups(&splits, &[field], &index_groups).unwrap(); + let schema = test_schema(); + let keys = [&schema.fields()[0]]; + assert_eq!(groups.len(), 3); + assert!(groups[&(partition(1).to_serialized_bytes(), 0)] + .all_overlapping_unchanged(0, 9, &schema, &keys, 1)); + for group in [ + (partition(1).to_serialized_bytes(), 1), + (partition(2).to_serialized_bytes(), 0), + ] { + assert!(!groups[&group].all_overlapping_unchanged(0, 9, &schema, &keys, 1)); + } + let wanted = HashSet::from([(partition(1).to_serialized_bytes(), 0)]); + let groups = file_range_lookups( + &splits, + &[DataField::new( + 0, + "part".into(), + PaimonType::Int(IntType::new()), + )], + &wanted, + ) + .unwrap(); + assert_eq!(groups.len(), 1); + assert!(groups[&(partition(1).to_serialized_bytes(), 0)] + .all_overlapping_unchanged(0, 9, &schema, &keys, 1)); + } + + #[test] + fn file_range_lookup_avoids_quadratic_work_with_one_covering_file() { + const SHARDS: usize = 5_000; + let mut files = (0..SHARDS) + .rev() + .map(|id| test_file(&id.to_string(), (id * 4) as i64, (id * 4 + 1) as i64)) + .collect::>(); + files.push(test_file("covering", 0, (SHARDS * 4) as i64)); + let lookup = FileRangeLookup::new(files.iter().collect()); + let tree = lookup.ranges.as_ref().unwrap(); + let schema = test_schema(); + let mut overlaps = 0; + for id in 0..SHARDS { + let start = (id * 4) as i64; + assert!(lookup.all_overlapping_unchanged( + start, + start + 1, + &schema, + &[&schema.fields()[0]], + 1, + )); + assert!(tree.all_overlapping(start, start + 1, &mut |_| { + overlaps += 1; + true + })); + } + assert_eq!(overlaps, SHARDS * 2); + assert!( + tree.inspections() < SHARDS * 2 * 32, + "{} range inspections for {SHARDS} shards must stay below quadratic work", + tree.inspections() + ); + // Many ranges crossing one center also need bounded scans at either + // edge; checking all crossing ranges would reintroduce quadratic work. + for left_edge in [true, false] { + let files = (0..SHARDS) + .map(|id| { + let (start, end) = if left_edge { + ((id * 4) as i64, (SHARDS * 4) as i64) + } else { + (0, (id * 4 + 1) as i64) + }; + test_file(&id.to_string(), start, end) + }) + .collect::>(); + let lookup = FileRangeLookup::new(files.iter().collect()); + let tree = lookup.ranges.as_ref().unwrap(); + let query = if left_edge { + 0 + } else { + (SHARDS * 4 - 3) as i64 + }; + let mut overlaps = 0; + for _ in 0..SHARDS { + assert!(tree.all_overlapping(query, query, &mut |_| { + overlaps += 1; + true + })); + } + assert_eq!(overlaps, SHARDS); + assert!(tree.inspections() <= SHARDS * 2, "left_edge={left_edge}"); + } + } + + #[test] + fn indexed_columns_use_compacted_column_sequences() { + let schema = PaimonSchema::builder() + .column("key", PaimonType::Int(IntType::new())) + .column("value", PaimonType::Int(IntType::new())) + .column("keep", PaimonType::Int(IntType::new())) + .build() + .unwrap(); + let schema = TableSchema::new(0, &schema); + let mut file = DataFileMeta { + file_name: "data.parquet".into(), + file_size: 128, + row_count: 3, + min_key: vec![], + max_key: vec![], + key_stats: BinaryTableStats::empty(), + value_stats: BinaryTableStats::empty(), + min_sequence_number: 1, + max_sequence_number: 1, + schema_id: 0, + level: 0, + extra_files: vec![], + creation_time: None, + delete_row_count: None, + embedded_index: None, + first_row_id: Some(0), + write_cols: None, + external_path: None, + file_source: None, + value_stats_cols: None, + column_max_sequence_numbers: None, + }; + file.write_cols = Some(vec!["key".into(), "value".into()]); + file.column_max_sequence_numbers = Some(vec![1, 3]); + file.max_sequence_number = 3; + file.file_source = Some(1); + let key = &schema.fields()[0]; + assert!(file_columns_unchanged(&file, &schema, &[key], 2, true)); + file.column_max_sequence_numbers = Some(vec![3, 3]); + assert!(!file_columns_unchanged(&file, &schema, &[key], 2, true)); + for sequences in [None, Some(vec![1]), Some(vec![0, 3]), Some(vec![4, 3])] { + file.column_max_sequence_numbers = sequences; + assert!(!file_columns_unchanged(&file, &schema, &[key], 2, true)); + } + file.max_sequence_number = 1; + file.column_max_sequence_numbers = None; + assert!(file_columns_unchanged(&file, &schema, &[key], 2, true)); + file.max_sequence_number = 3; + file.write_cols = None; + file.column_max_sequence_numbers = Some(vec![1, 3, 1]); + assert!(file_columns_unchanged(&file, &schema, &[key], 2, true)); + file.write_cols = Some(vec!["_ROW_ID".into(), "value".into(), "key".into()]); + file.column_max_sequence_numbers = Some(vec![3, 1]); + assert!(file_columns_unchanged(&file, &schema, &[key], 2, true)); + for columns in [ + vec!["key".into(), "key".into()], + vec!["key".into(), "unknown".into()], + ] { + file.write_cols = Some(columns); + file.column_max_sequence_numbers = Some(vec![1, 1]); + assert!(!file_columns_unchanged(&file, &schema, &[key], 2, true)); + } + + let schema = PaimonSchema::builder() + .column("blob", PaimonType::Blob(crate::spec::BlobType::new())) + .column("key", PaimonType::Int(IntType::new())) + .column("value", PaimonType::Int(IntType::new())) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .option("data-evolution.write-cols-optimization.enabled", "true") + .build() + .unwrap(); + let schema = TableSchema::new(0, &schema); + let key = &schema.fields()[1]; + file.write_cols = None; + file.column_max_sequence_numbers = Some(vec![1, 3]); + assert!(file_columns_unchanged(&file, &schema, &[key], 2, true)); + file.write_cols = Some(vec!["blob".into()]); + file.column_max_sequence_numbers = Some(vec![3]); + assert!(file_columns_unchanged(&file, &schema, &[key], 2, true)); + } +} diff --git a/crates/paimon/src/table/global_index_scanner/composite.rs b/crates/paimon/src/table/global_index_scanner/composite.rs index b84f6e31b..c99926673 100644 --- a/crates/paimon/src/table/global_index_scanner/composite.rs +++ b/crates/paimon/src/table/global_index_scanner/composite.rs @@ -184,7 +184,14 @@ impl GlobalIndexScanner { continue; } return Ok(Some(CompositeIndexScanResult { - key_field_ids: plan.fields.iter().map(|field| field.id()).collect(), + // Unbound suffix predicates can still use independent scalar + // indexes after the composite prefix has been consumed. + key_field_ids: plan + .fields + .iter() + .take(plan.bound_columns) + .map(|field| field.id()) + .collect(), result: GlobalIndexScanResult { row_ranges: merge_row_ranges(ranges), // Remaining predicates are always retained by the table reader. diff --git a/crates/paimon/src/table/global_index_scanner/tests.rs b/crates/paimon/src/table/global_index_scanner/tests.rs index c00c99673..89773c87b 100644 --- a/crates/paimon/src/table/global_index_scanner/tests.rs +++ b/crates/paimon/src/table/global_index_scanner/tests.rs @@ -2121,6 +2121,103 @@ async fn test_composite_prefers_fully_covering_leading_bitmap() { assert_eq!(result.unwrap(), vec![RowRange::new(105, 106)]); } +#[tokio::test] +async fn test_composite_interval_limit_keeps_suffix_scalar_index() { + let fields = two_field_schema_fields(); + let key_type = DataType::Row(crate::spec::RowType::new(fields.clone())); + let composite_output = VecFileWrite::new(); + let mut composite_writer = BTreeIndexWriter::with_comparator( + Box::new(composite_output.clone()), + 256, + BlockCompressionType::None, + make_key_comparator(&key_type), + ); + let tuples = [[0, 0], [0, 1], [1, 0], [1, 1]]; + for (id, [first, second]) in tuples.iter().enumerate() { + let key = crate::btree::serialize_composite_key( + &[Some(Datum::Int(*first)), Some(Datum::Int(*second))], + &fields, + ) + .unwrap(); + composite_writer.write(Some(&key), id as i64).await.unwrap(); + } + let composite_meta = composite_writer.finish().await.unwrap().meta; + let scalar_output = VecFileWrite::new(); + let mut scalar_writer = BTreeIndexWriter::with_comparator( + Box::new(scalar_output.clone()), + 256, + BlockCompressionType::None, + make_key_comparator(fields[1].data_type()), + ); + for (value, ids) in [(0, [0, 2]), (1, [1, 3])] { + for id in ids { + scalar_writer + .write(Some(&le_int_key(value)), id) + .await + .unwrap(); + } + } + let scalar_meta = scalar_writer.finish().await.unwrap().meta; + let tmp = tempfile::tempdir().unwrap(); + std::fs::create_dir(tmp.path().join("index")).unwrap(); + let composite_bytes = composite_output.to_vec(); + let scalar_bytes = scalar_output.to_vec(); + std::fs::write(tmp.path().join("index/composite.btree"), &composite_bytes).unwrap(); + std::fs::write(tmp.path().join("index/suffix.btree"), &scalar_bytes).unwrap(); + let mut composite = make_global_index_entry("composite.btree", 1, 0, 3, &composite_meta); + composite.index_file.file_size = composite_bytes.len() as i64; + composite + .index_file + .global_index_meta + .as_mut() + .unwrap() + .extra_field_ids = Some(vec![2]); + let mut scalar = make_global_index_entry("suffix.btree", 2, 0, 3, &scalar_meta); + scalar.index_file.file_size = scalar_bytes.len() as i64; + let entries = [composite, scalar]; + let io = crate::io::FileIOBuilder::new("file").build().unwrap(); + let path = format!("file://{}", tmp.path().display()); + let mut scanner = + GlobalIndexScanner::create(&io, &path, 2, i64::MAX, i64::MAX, &entries, &fields) + .unwrap() + .unwrap(); + let probe = Arc::new(QueryIoProbe::default()); + scanner.query_io_probe = Some(Arc::clone(&probe)); + let b = crate::spec::PredicateBuilder::new(&fields); + // 17 x 17 stops the tuple plan after the first column. Its unbound suffix + // must remain available for the independent scalar index to narrow rows. + let predicate = Predicate::and(vec![ + b.is_in( + "id", + [0, 1].into_iter().chain(100..115).map(Datum::Int).collect(), + ) + .unwrap(), + b.is_in( + "value", + [1].into_iter().chain(100..116).map(Datum::Int).collect(), + ) + .unwrap(), + ]); + let expected = tuples + .iter() + .enumerate() + .filter(|(_, tuple)| { + let datums = tuple.map(Datum::Int); + let row = crate::spec::BinaryRow::from_datums(&[ + (Some(&datums[0]), fields[0].data_type()), + (Some(&datums[1]), fields[1].data_type()), + ]); + crate::spec::eval_row(&predicate, &row).unwrap() + }) + .map(|(id, _)| id as u64) + .collect::(); + assert_eq!(expected.iter().collect::>(), vec![1, 3]); + let result = scanner.evaluate(&predicate).await.unwrap().unwrap(); + assert_eq!(result.row_ranges, bitmap_to_ranges(&expected)); + assert_eq!(result.evaluated_field_ids, HashSet::from([1, 2])); + assert_eq!(probe.predicate_queries.load(TestOrdering::SeqCst), 1); +} + async fn assert_bitmap_int_fixture(file_name: &str) { let data_type = DataType::Int(crate::spec::IntType::new()); let (file_io, table_path, file_name, meta, _tmp) = setup_bitmap_testdata_table(file_name); diff --git a/crates/paimon/src/table/key_normalization.rs b/crates/paimon/src/table/key_normalization.rs index 01264bb42..8622f8688 100644 --- a/crates/paimon/src/table/key_normalization.rs +++ b/crates/paimon/src/table/key_normalization.rs @@ -22,12 +22,29 @@ use arrow_schema::DataType; use super::update_input::{cast_update_value, encoded, CastMode}; use super::write_batch_normalize::normalize_write_array; +use crate::spec::{extract_datum, BinaryRow, BinaryRowBuilder, DataField}; mod numeric; #[cfg(test)] mod encoded_tests; +// Java may reserve variable-length storage for NULL timestamps/decimals. +// Compare logical partitions using one encoding, independent of that padding. +pub(super) fn canonical_partition_key( + row: &BinaryRow, + fields: &[DataField], +) -> crate::Result> { + let mut builder = BinaryRowBuilder::new(fields.len() as i32); + for (index, field) in fields.iter().enumerate() { + match extract_datum(row, index, field.data_type())? { + Some(value) => builder.write_datum(index, &value, field.data_type()), + None => builder.set_null_at(index), + } + } + Ok(builder.build().to_serialized_bytes()) +} + /// Unify Arrow representations without changing key identity. Dictionary /// indices are batch-local; numeric widths, decimal scales and temporal units /// may differ only when conversion is exact. Binary keys compare their bytes. @@ -147,6 +164,7 @@ pub(super) fn normalize_match_key( #[cfg(test)] mod tests { use super::*; + use crate::spec::{DataType as PaimonType, DecimalType, TimestampType, VarCharType}; use arrow_array::types::Int8Type; use arrow_array::{ Array, BinaryArray, Decimal128Array, DictionaryArray, FixedSizeBinaryArray, Float32Array, @@ -365,4 +383,45 @@ mod tests { normalize_match_key(&local, &DataType::Timestamp(TimeUnit::Microsecond, None)).is_err() ); } + + #[test] + fn partition_identity_ignores_java_null_variable_storage() { + let fields = vec![ + DataField::new( + 0, + "ts".into(), + PaimonType::Timestamp(TimestampType::new(6).unwrap()), + ), + DataField::new( + 1, + "amount".into(), + PaimonType::Decimal(DecimalType::new(20, 2).unwrap()), + ), + DataField::new( + 2, + "label".into(), + PaimonType::VarChar(VarCharType::new(100).unwrap()), + ), + ]; + let mut builder = BinaryRowBuilder::new(3); + builder.set_null_at(0); + builder.set_null_at(1); + builder.write_string(2, "after-null-fields"); + let canonical = builder.build(); + let mut java = canonical.data()[..32].to_vec(); + // BinaryWriter.writeTimestamp(NULL, 6) reserves 8 bytes; + // writeDecimal(NULL, 20) reserves another 16, even though both are NULL. + java.extend_from_slice(&[0; 24]); + java.extend_from_slice(&canonical.data()[32..]); + java[8..16].copy_from_slice(&(32_u64 << 32).to_le_bytes()); + java[16..24].copy_from_slice(&(40_u64 << 32).to_le_bytes()); + let offset_and_size = u64::from_le_bytes(canonical.data()[24..32].try_into().unwrap()); + java[24..32].copy_from_slice(&(offset_and_size + (24_u64 << 32)).to_le_bytes()); + let java = BinaryRow::from_bytes(3, java); + assert_ne!(java.to_serialized_bytes(), canonical.to_serialized_bytes()); + assert_eq!( + canonical_partition_key(&java, &fields).unwrap(), + canonical.to_serialized_bytes() + ); + } } diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 0fee5af87..87e89f56e 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -80,6 +80,7 @@ mod full_text_search_builder; mod global_index_build_builder; pub(crate) mod global_index_build_common; mod global_index_drop_builder; +mod global_index_freshness; mod global_index_result; pub(crate) mod global_index_scanner; mod global_index_search_filter; diff --git a/crates/paimon/src/table/table_upsert.rs b/crates/paimon/src/table/table_upsert.rs index 73d626470..e18521ea2 100644 --- a/crates/paimon/src/table/table_upsert.rs +++ b/crates/paimon/src/table/table_upsert.rs @@ -18,7 +18,7 @@ //! Table-level append-only upsert by user-specified keys. use std::borrow::Cow; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::ops::Range; use std::sync::Arc; @@ -28,11 +28,14 @@ use arrow_select::{concat::concat_batches, take::take}; use futures::TryStreamExt; use indexmap::IndexMap; +use super::global_index_freshness::indexed_columns_unchanged; +use super::key_normalization::canonical_partition_key; use super::update_input::unique_column_index; use super::upsert_key_matcher::UpsertKeyMatcher; use super::write_batch_normalize::{normalize_blob_write_field, normalize_write_array}; use crate::spec::{ - batch_to_serialized_bytes, extract_datum, BinaryRow, BinaryRowBuilder, CoreOptions, DataField, + batch_to_serialized_bytes, extract_datum_from_arrow, BinaryRow, CoreOptions, Predicate, + PredicateBuilder, }; use crate::table::source::{DataSplit, Plan}; use crate::table::{CommitMessage, Table}; @@ -64,19 +67,6 @@ fn selected_rows(batch: &RecordBatch, rows: &[usize]) -> crate::Result crate::Result> { - let mut builder = BinaryRowBuilder::new(fields.len() as i32); - for (index, field) in fields.iter().enumerate() { - match extract_datum(row, index, field.data_type())? { - Some(value) => builder.write_datum(index, &value, field.data_type()), - None => builder.set_null_at(index), - } - } - Ok(builder.build().to_serialized_bytes()) -} - /// Upsert Arrow rows into a data-evolution table without primary keys. /// Internal executor for `TableUpdate::upsert_by_arrow_with_key`. Existing /// keys are updated by row ID; new keys are appended. @@ -379,34 +369,130 @@ impl TableUpsert { let fields = schema.partition_fields(); let mut splits = Vec::new(); for split in plan.splits() { - if partitions.contains(&partition_key(split.partition(), &fields)?) { + if partitions.contains(&canonical_partition_key(split.partition(), &fields)?) { splits.push(split.clone()); } } Ok(Cow::Owned(splits)) } - /// Match source keys against the target snapshot, stage row-ID updates and - /// new rows, then return both sets of commit messages as one operation. - #[must_use = "commit messages must be passed to TableCommit"] - pub(super) async fn prepare_commit(self) -> crate::Result> { - let source = self.key_rows()?; - if source.num_rows() == 0 { - return Err(invalid("Input data is empty")); + /// Each column domain is a superset of the input tuples. The matcher still + /// checks complete keys, including NULLs, after index and reader pruning. + fn lookup_filter(&self, source: &RecordBatch) -> crate::Result> { + let builder = PredicateBuilder::new(self.table.schema().fields()); + let mut predicates = Vec::new(); + for key in &self.keys { + // Partition sets compare raw BinaryRows, which may use different + // variable-length storage for Java NULL timestamps and decimals. + // matching_splits checks their canonical logical identity instead. + if self.table.schema().partition_keys().contains(key) { + continue; + } + let column = unique_column_index(&source.schema(), key)?; + // Arrow row encoding preserves NaN payloads and signed zero. SQL + // equality need not share that identity, so retain the full domain. + if matches!( + source.column(column).data_type(), + DataType::Float32 | DataType::Float64 + ) { + continue; + } + let field = self + .table + .schema() + .fields() + .iter() + .find(|field| field.name() == key) + .unwrap(); + // BLOB keys can be matched as Arrow bytes but have no scalar Datum. + if field.data_type().is_blob_type() { + continue; + } + let mut seen = HashSet::new(); + let mut values = Vec::new(); + let mut has_null = false; + for row in 0..source.num_rows() { + match extract_datum_from_arrow(source, row, column, field.data_type())? { + Some(value) => { + if seen.insert(crate::btree::serialize_datum(&value, field.data_type())) { + values.push(value); + } + } + None => has_null = true, + } + } + let mut alternatives = Vec::new(); + if values.len() == 1 { + alternatives.push(builder.equal(key.as_str(), values.remove(0))?); + } else if !values.is_empty() { + alternatives.push(builder.is_in(key.as_str(), values)?); + } + if has_null { + alternatives.push(builder.is_null(key.as_str())?); + } + if !alternatives.is_empty() { + predicates.push(Predicate::or(alternatives)); + } } - let mut matcher = UpsertKeyMatcher::new(&source, self.keys.clone())?; + Ok((!predicates.is_empty()).then(|| Predicate::and(predicates))) + } - let mut read_builder = self.table.new_read_builder(); + async fn match_existing( + &self, + source: &RecordBatch, + matcher: &mut UpsertKeyMatcher, + ) -> crate::Result { + // Mutation lookups must include unindexed rows even when queries use + // fast mode, or a matching target would incorrectly become an insert. + let lookup_options = HashMap::from([("scalar-index.search-mode".into(), "full".into())]); + let mut lookup_table = self.table.copy_with_options(lookup_options); + if lookup_table.schema().core_options().global_index_enabled() { + let snapshot = super::time_travel::resolve_snapshot(&lookup_table).await?; + lookup_table = lookup_table.copy_with_pinned_snapshot(snapshot.as_ref()); + if let Some(snapshot) = snapshot { + let partitions = self.serialized_partitions(source)?.into_iter().collect(); + if !indexed_columns_unchanged(&lookup_table, &snapshot, &self.keys, &partitions) + .await + .unwrap_or(false) + { + lookup_table = lookup_table.copy_with_options(HashMap::from([( + "global-index.enabled".into(), + "false".into(), + )])); + } + } + } + let mut scan_builder = lookup_table.new_read_builder(); + if let Some(filter) = self.lookup_filter(source)? { + scan_builder.with_filter(filter); + } + let plan = scan_builder.new_scan().plan_all().await?; + // Split row ranges retain index pruning. Exact tuple matching below + // avoids evaluating a potentially large IN list again for every row. + let mut read_builder = lookup_table.new_read_builder(); let mut projection = self.keys.iter().map(String::as_str).collect::>(); projection.push(ROW_ID); read_builder.with_projection(&projection)?; - let plan = read_builder.new_scan().plan_all().await?; let read = read_builder.new_read()?.for_update(); - let splits = self.matching_splits(&source, &plan)?; + let splits = self.matching_splits(source, &plan)?; let mut stream = read.to_arrow(splits.as_ref())?; while let Some(batch) = stream.try_next().await? { matcher.add_existing_batch(&batch)?; } + Ok(plan) + } + + /// Match source keys against the target snapshot, stage row-ID updates and + /// new rows, then return both sets of commit messages as one operation. + #[must_use = "commit messages must be passed to TableCommit"] + pub(super) async fn prepare_commit(self) -> crate::Result> { + let source = self.key_rows()?; + if source.num_rows() == 0 { + return Err(invalid("Input data is empty")); + } + let mut matcher = UpsertKeyMatcher::new(&source, self.keys.clone())?; + + let plan = self.match_existing(&source, &mut matcher).await?; let (matched_indices, row_ids, new_indices) = matcher.finish(); // Validate append field sets before staging any updates or new files. @@ -475,7 +561,10 @@ mod tests { use super::*; use crate::catalog::Identifier; use crate::io::FileIOBuilder; - use crate::spec::{DataType as PaimonType, DecimalType, TimestampType, VarCharType}; + use crate::spec::{ + DataEvolutionIndexSourceMeta, DataType as PaimonType, IndexManifest, TimestampType, + VarCharType, + }; use crate::spec::{IntType, Schema as PaimonSchema, TableSchema}; use crate::table::TableCommit; use arrow_array::{Array, Int32Array, StringArray}; @@ -572,6 +661,758 @@ mod tests { rows } + #[tokio::test] + async fn upsert_key_lookup_uses_btree_row_ranges() { + let table = numeric_table("memory:/upsert_btree_lookup").await; + table + .new_global_index_build_builder() + .with_index_column("key") + .execute() + .await + .unwrap(); + let mut upsert = TableUpsert::new( + &table, + "upsert".into(), + vec!["key".into()], + vec!["value".into()], + ) + .unwrap(); + upsert + .add_batch(numeric_source( + vec![1], + Arc::new(Int32Array::from(vec![42])), + )) + .unwrap(); + let source = upsert.key_rows().unwrap(); + let mut matcher = UpsertKeyMatcher::new(&source, upsert.keys.clone()).unwrap(); + let plan = upsert.match_existing(&source, &mut matcher).await.unwrap(); + assert_eq!(plan.snapshot_id(), Some(2)); + assert!(plan + .splits() + .iter() + .all(|split| { split.row_ranges() == Some(&[crate::table::RowRange::new(0, 1)][..]) })); + let (indices, mut row_ids, new) = matcher.finish(); + row_ids.sort(); + assert_eq!(indices, vec![0, 0]); + assert_eq!(row_ids, vec![0, 1]); + assert!(new.is_empty()); + TableCommit::new(table.clone(), "upsert".into()) + .commit(upsert.prepare_commit().await.unwrap()) + .await + .unwrap(); + assert_eq!( + numeric_rows(&table).await, + vec![ + (0, 1, Some(42), Some(10)), + (1, 1, Some(42), Some(11)), + (2, 2, Some(0), Some(12)), + ] + ); + commit_upsert( + &table, + vec![numeric_source( + vec![9], + Arc::new(Int32Array::from(vec![49])), + )], + vec!["key".into()], + ) + .await; + assert_eq!(numeric_rows(&table).await[3], (3, 9, Some(49), None)); + } + + async fn commit_upsert(table: &Table, batches: Vec, keys: Vec) { + let mut update = table.new_write_builder().new_update().unwrap(); + update.with_update_type(vec!["value".into()]).unwrap(); + let messages = update + .upsert_by_arrow_with_key(batches, keys) + .await + .unwrap(); + TableCommit::new(table.clone(), "upsert".into()) + .commit(messages) + .await + .unwrap(); + } + + async fn values_by_row_id(table: &Table) -> Vec<(i64, i32)> { + let mut builder = table.new_read_builder(); + builder.with_projection(&[ROW_ID, "value"]).unwrap(); + let plan = builder.new_scan().plan_all().await.unwrap(); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut rows = Vec::new(); + for batch in batches { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + rows.extend((0..batch.num_rows()).map(|row| (ids.value(row), values.value(row)))); + } + rows.sort(); + rows + } + + #[tokio::test] + async fn upsert_key_lookup_includes_unindexed_rows_in_fast_mode() { + for indexed in [false, true] { + let table = numeric_table(&format!("memory:/upsert_partial_index_{indexed}")).await; + let table = table.copy_with_options(HashMap::from([( + "scalar-index.search-mode".into(), + "fast".into(), + )])); + if indexed { + table + .new_global_index_build_builder() + .with_index_column("key") + .execute() + .await + .unwrap(); + } + let appended = RecordBatch::try_from_iter([ + ("key", Arc::new(Int32Array::from(vec![1, 3])) as ArrayRef), + ("value", Arc::new(Int32Array::from(vec![0, 0])) as ArrayRef), + ("keep", Arc::new(Int32Array::from(vec![13, 14])) as ArrayRef), + ]) + .unwrap(); + let mut write = table.new_write_builder().new_write().unwrap(); + write.write_arrow_batch(&appended).await.unwrap(); + TableCommit::new(table.clone(), "append".into()) + .commit(write.prepare_commit().await.unwrap()) + .await + .unwrap(); + write.close().await; + commit_upsert( + &table, + vec![numeric_source( + vec![1, 3, 4], + Arc::new(Int32Array::from(vec![41, 43, 44])), + )], + vec!["key".into()], + ) + .await; + assert_eq!( + numeric_rows(&table).await, + vec![ + (0, 1, Some(41), Some(10)), + (1, 1, Some(41), Some(11)), + (2, 2, Some(0), Some(12)), + (3, 1, Some(41), Some(13)), + (4, 3, Some(43), Some(14)), + (5, 4, Some(44), None), + ], + "indexed={indexed}" + ); + } + } + + #[tokio::test] + async fn upsert_ignore_policy_keeps_unchanged_key_indexes() { + for update_payload in [false, true] { + let table = + numeric_table(&format!("memory:/upsert_ignore_fresh_{update_payload}")).await; + let table = table.copy_with_options(HashMap::from([( + "global-index.column-update-action".into(), + "IGNORE".into(), + )])); + table + .new_global_index_build_builder() + .with_index_column("key") + .execute() + .await + .unwrap(); + if update_payload { + let batch = RecordBatch::try_from_iter([ + (ROW_ID, Arc::new(Int64Array::from(vec![0, 1])) as ArrayRef), + ( + "value", + Arc::new(Int32Array::from(vec![11, 12])) as ArrayRef, + ), + ]) + .unwrap(); + let update = table.new_write_builder().new_update().unwrap(); + TableCommit::new(table.clone(), "payload".into()) + .commit( + update + .update_by_arrow_with_row_id(vec![batch]) + .await + .unwrap(), + ) + .await + .unwrap(); + } + let (plan, ids, new) = numeric_lookup(&table, 1).await; + assert!(!plan.splits().is_empty()); + assert!( + plan.splits().iter().all(|split| { + split.row_ranges() == Some(&[crate::table::RowRange::new(0, 1)][..]) + }), + "payload update={update_payload}" + ); + assert_eq!(ids, vec![0, 1]); + assert!(new.is_empty()); + } + } + + async fn numeric_lookup(table: &Table, key: i32) -> (Plan, Vec, Vec) { + let mut upsert = TableUpsert::new( + table, + "upsert".into(), + vec!["key".into()], + vec!["value".into()], + ) + .unwrap(); + upsert + .add_batch(numeric_source( + vec![key], + Arc::new(Int32Array::from(vec![43])), + )) + .unwrap(); + let source = upsert.key_rows().unwrap(); + let mut matcher = UpsertKeyMatcher::new(&source, upsert.keys.clone()).unwrap(); + let plan = upsert.match_existing(&source, &mut matcher).await.unwrap(); + let (_, mut ids, new) = matcher.finish(); + ids.sort(); + (plan, ids, new) + } + + #[tokio::test] + async fn upsert_checks_index_sources_only_under_ignore() { + for action in ["IGNORE", "THROW_ERROR", "DROP_PARTITION_INDEX"] { + for (label, source) in [ + ("missing", None), + ("malformed", Some(vec![1, 2, 3])), + ( + "future", + Some(DataEvolutionIndexSourceMeta::new(100).unwrap().serialize()), + ), + ] { + let table = numeric_table(&format!("memory:/upsert_source_{action}_{label}")).await; + let table = table.copy_with_options(HashMap::from([( + "global-index.column-update-action".into(), + action.into(), + )])); + table + .new_global_index_build_builder() + .with_index_column("key") + .execute() + .await + .unwrap(); + let manager = table.snapshot_manager(); + let snapshot = manager.get_latest_snapshot().await.unwrap().unwrap(); + let path = manager.manifest_path(snapshot.index_manifest().unwrap()); + let mut entries = IndexManifest::read(table.file_io(), &path).await.unwrap(); + for entry in &mut entries { + entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .source_meta = source.clone(); + } + IndexManifest::write(table.file_io(), &path, &entries) + .await + .unwrap(); + let (plan, ids, new) = numeric_lookup(&table, 1).await; + assert!(!plan.splits().is_empty()); + if action == "IGNORE" { + assert!( + plan.splits() + .iter() + .all(|split| split.row_ranges().is_none()), + "action={action}, source={label}" + ); + } else { + assert!( + plan.splits().iter().all(|split| { + split.row_ranges() == Some(&[crate::table::RowRange::new(0, 1)][..]) + }), + "action={action}, source={label}" + ); + } + assert_eq!(ids, vec![0, 1]); + assert!(new.is_empty()); + } + } + } + + #[tokio::test] + async fn upsert_key_lookup_ignores_stale_indexes_under_ignore_policy() { + let table = numeric_table("memory:/upsert_ignored_index_updates").await; + let table = table.copy_with_options(HashMap::from([( + "global-index.column-update-action".into(), + "IGNORE".into(), + )])); + table + .new_global_index_build_builder() + .with_index_column("key") + .execute() + .await + .unwrap(); + let changed_keys = RecordBatch::try_from_iter([ + (ROW_ID, Arc::new(Int64Array::from(vec![0, 1])) as ArrayRef), + ("key", Arc::new(Int32Array::from(vec![3, 3])) as ArrayRef), + ]) + .unwrap(); + let update = table.new_write_builder().new_update().unwrap(); + TableCommit::new(table.clone(), "change-keys".into()) + .commit( + update + .update_by_arrow_with_row_id(vec![changed_keys]) + .await + .unwrap(), + ) + .await + .unwrap(); + commit_upsert( + &table, + vec![numeric_source( + vec![3, 4], + Arc::new(Int32Array::from(vec![43, 44])), + )], + vec!["key".into()], + ) + .await; + assert_eq!( + numeric_rows(&table).await, + vec![ + (0, 3, Some(43), Some(10)), + (1, 3, Some(43), Some(11)), + (2, 2, Some(0), Some(12)), + (3, 4, Some(44), None), + ] + ); + // Dropping and rebuilding makes the updated keys safe to probe again. + table + .new_global_index_drop_builder() + .with_index_column("key") + .execute() + .await + .unwrap(); + table + .new_global_index_build_builder() + .with_index_column("key") + .execute() + .await + .unwrap(); + let (plan, ids, new) = numeric_lookup(&table, 3).await; + assert!(!plan.splits().is_empty()); + assert!(plan + .splits() + .iter() + .all(|split| { split.row_ranges() == Some(&[crate::table::RowRange::new(0, 1)][..]) })); + assert_eq!(ids, vec![0, 1]); + assert!(new.is_empty()); + } + + #[tokio::test] + async fn upsert_key_lookup_preserves_updated_keys_and_deleted_rows() { + let table = numeric_table("memory:/upsert_changed_keys").await; + let table = table.copy_with_options(HashMap::from([ + ("deletion-vectors.enabled".into(), "true".into()), + ( + "global-index.column-update-action".into(), + "DROP_PARTITION_INDEX".into(), + ), + ])); + table + .new_global_index_build_builder() + .with_index_column("key") + .execute() + .await + .unwrap(); + let changed_key = RecordBatch::try_from_iter([ + (ROW_ID, Arc::new(Int64Array::from(vec![2])) as ArrayRef), + ("key", Arc::new(Int32Array::from(vec![3])) as ArrayRef), + ]) + .unwrap(); + let update = table.new_write_builder().new_update().unwrap(); + TableCommit::new(table.clone(), "change-key".into()) + .commit( + update + .update_by_arrow_with_row_id(vec![changed_key]) + .await + .unwrap(), + ) + .await + .unwrap(); + commit_upsert( + &table, + vec![numeric_source( + vec![3], + Arc::new(Int32Array::from(vec![43])), + )], + vec!["key".into()], + ) + .await; + assert_eq!(numeric_rows(&table).await[2], (2, 3, Some(43), Some(12))); + // Rebuild so the next lookup encounters an index entry for a deleted row. + table + .new_global_index_build_builder() + .with_index_column("key") + .execute() + .await + .unwrap(); + TableCommit::new(table.clone(), "delete".into()) + .commit(update.delete_by_row_id(vec![0, 1]).await.unwrap()) + .await + .unwrap(); + commit_upsert( + &table, + vec![numeric_source( + vec![1], + Arc::new(Int32Array::from(vec![41])), + )], + vec!["key".into()], + ) + .await; + assert_eq!( + numeric_rows(&table).await, + vec![(2, 3, Some(43), Some(12)), (3, 1, Some(41), None)] + ); + } + + #[tokio::test] + async fn upsert_composite_btree_preserves_nulls_partitions_and_tuple_identity() { + let schema = PaimonSchema::builder() + .column("part", PaimonType::VarChar(VarCharType::string_type())) + .column("key", PaimonType::Int(IntType::new())) + .column("tag", PaimonType::VarChar(VarCharType::string_type())) + .column("value", PaimonType::Int(IntType::new())) + .partition_keys(["part"]) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .build() + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("db", "upsert_composite"), + "memory:/upsert_composite_btree".into(), + TableSchema::new(0, &schema), + None, + ); + let batch = |parts: Vec<&str>, + keys: Vec>, + tags: Vec>, + values: Vec| { + RecordBatch::try_from_iter([ + ("part", Arc::new(StringArray::from(parts)) as ArrayRef), + ("key", Arc::new(Int32Array::from(keys)) as ArrayRef), + ("tag", Arc::new(StringArray::from(tags)) as ArrayRef), + ("value", Arc::new(Int32Array::from(values)) as ArrayRef), + ]) + .unwrap() + }; + let initial = batch( + vec!["a", "a", "a", "a", "a", "b"], + vec![Some(1), Some(2), None, Some(1), Some(1), Some(1)], + vec![Some("x"), Some("y"), Some("x"), None, Some("y"), Some("x")], + vec![0; 6], + ); + let mut write = table.new_write_builder().new_write().unwrap(); + write.write_arrow_batch(&initial).await.unwrap(); + TableCommit::new(table.clone(), "seed".into()) + .commit(write.prepare_commit().await.unwrap()) + .await + .unwrap(); + write.close().await; + table + .new_global_index_build_builder() + .with_index_columns(&["key", "tag"]) + .execute() + .await + .unwrap(); + commit_upsert( + &table, + vec![batch( + vec!["a"; 6], + vec![Some(1), Some(2), None, Some(1), Some(1), Some(9)], + vec![Some("x"), Some("y"), Some("x"), None, Some("x"), Some("x")], + vec![10, 20, 30, 40, 50, 60], + )], + vec!["key".into(), "tag".into()], + ) + .await; + let mut builder = table.new_read_builder(); + builder + .with_projection(&["part", "key", "tag", "value"]) + .unwrap(); + let plan = builder.new_scan().plan_all().await.unwrap(); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut rows = Vec::new(); + for batch in batches { + let parts = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let keys = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let tags = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + rows.extend((0..batch.num_rows()).map(|row| { + ( + parts.value(row).to_string(), + (!keys.is_null(row)).then(|| keys.value(row)), + (!tags.is_null(row)).then(|| tags.value(row).to_string()), + values.value(row), + ) + })); + } + rows.sort(); + // (1, y) is only a cross-product candidate. Partition b must not match. + let expected = vec![ + ("a", None, Some("x"), 30), + ("a", Some(1), None, 40), + ("a", Some(1), Some("x"), 50), + ("a", Some(1), Some("y"), 0), + ("a", Some(2), Some("y"), 20), + ("a", Some(9), Some("x"), 60), + ("b", Some(1), Some("x"), 0), + ] + .into_iter() + .map(|(part, key, tag, value)| (part.to_string(), key, tag.map(str::to_string), value)) + .collect::>(); + assert_eq!(rows, expected); + } + + #[tokio::test] + async fn upsert_large_composite_domains_keep_index_pruning_and_exact_matches() { + for large_domain in [false, true] { + let schema = PaimonSchema::builder() + .column("key", PaimonType::Int(IntType::new())) + .column("tag", PaimonType::Int(IntType::new())) + .column("value", PaimonType::Int(IntType::new())) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .build() + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("db", "large_composite_upsert"), + format!("memory:/upsert_large_composite_domains_{large_domain}"), + TableSchema::new(0, &schema), + None, + ); + let batch = |keys: Vec, tags: Vec, value: i32| { + let count = keys.len(); + RecordBatch::try_from_iter([ + ("key", Arc::new(Int32Array::from(keys)) as ArrayRef), + ("tag", Arc::new(Int32Array::from(tags)) as ArrayRef), + ( + "value", + Arc::new(Int32Array::from(vec![value; count])) as ArrayRef, + ), + ]) + .unwrap() + }; + let (keys, tags, count, stride, matches, candidates) = if large_domain { + (512, 8, 257, 2, 256, 256) + } else { + (128, 32, 17, 4, 16, 17 * 32) + }; + let mut write = table.new_write_builder().new_write().unwrap(); + write + .write_arrow_batch(&batch( + (0..keys) + .flat_map(|key| std::iter::repeat_n(key, tags as usize)) + .collect(), + (0..keys).flat_map(|_| 0..tags).collect(), + 0, + )) + .await + .unwrap(); + TableCommit::new(table.clone(), "seed".into()) + .commit(write.prepare_commit().await.unwrap()) + .await + .unwrap(); + write.close().await; + table + .new_global_index_build_builder() + .with_index_columns(&["key", "tag"]) + .execute() + .await + .unwrap(); + + let mut upsert = TableUpsert::new( + &table, + "upsert".into(), + vec!["key".into(), "tag".into()], + vec!["value".into()], + ) + .unwrap(); + // Exercise a large linear IN and a 17 x 17 Cartesian domain. + // The last tuple is absent; prefix candidates with other suffixes + // must not be updated or cause an existing match to become an insert. + upsert + .add_batch(batch( + (0..count).map(|key| key * stride).collect(), + (0..count) + .map(|i| { + if large_domain { + 7 + } else if i < matches { + i + } else { + 99 + } + }) + .collect(), + 42, + )) + .unwrap(); + let source = upsert.key_rows().unwrap(); + let mut matcher = UpsertKeyMatcher::new(&source, upsert.keys.clone()).unwrap(); + let plan = upsert.match_existing(&source, &mut matcher).await.unwrap(); + let mut actual_candidates = 0; + for split in plan.splits() { + for range in split + .row_ranges() + .expect("large domains must retain index pruning") + { + actual_candidates += range.to() - range.from() + 1; + } + } + assert_eq!(actual_candidates, candidates); + let (_, mut matched, new) = matcher.finish(); + matched.sort(); + let expected = (0..matches) + .map(|key| i64::from(key * stride * tags + if large_domain { 7 } else { key })) + .collect::>(); + assert_eq!(matched, expected); + assert_eq!(new, vec![(count - 1) as usize]); + TableCommit::new(table.clone(), "upsert".into()) + .commit(upsert.prepare_commit().await.unwrap()) + .await + .unwrap(); + let rows = values_by_row_id(&table).await; + assert_eq!(rows.len(), 4097); + for (row_id, value) in rows { + let updated = expected.contains(&row_id) || row_id == 4096; + assert_eq!( + value, + if updated { 42 } else { 0 }, + "row {row_id}, large={large_domain}" + ); + } + } + } + + #[tokio::test] + async fn upsert_blob_keys_fall_back_to_byte_matching() { + use crate::spec::BlobType; + use arrow_array::LargeBinaryArray; + + let schema = PaimonSchema::builder() + .column("key", PaimonType::Blob(BlobType::new())) + .column("value", PaimonType::Int(IntType::new())) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .build() + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("db", "upsert_blob"), + "memory:/upsert_blob_key".into(), + TableSchema::new(0, &schema), + None, + ); + for value in [41, 42] { + let source = RecordBatch::try_from_iter([ + ( + "key", + Arc::new(LargeBinaryArray::from_vec(vec![b"blob-key"])) as ArrayRef, + ), + ("value", Arc::new(Int32Array::from(vec![value])) as ArrayRef), + ]) + .unwrap(); + commit_upsert(&table, vec![source], vec!["key".into()]).await; + } + assert_eq!(values_by_row_id(&table).await, vec![(0, 42)]); + } + + #[tokio::test] + async fn upsert_float_keys_keep_arrow_nan_and_signed_zero_identity() { + use crate::spec::DoubleType; + use arrow_array::Float64Array; + + let schema = PaimonSchema::builder() + .column("key", PaimonType::Double(DoubleType::new())) + .column("value", PaimonType::Int(IntType::new())) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .build() + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("db", "upsert_float"), + "memory:/upsert_float_btree".into(), + TableSchema::new(0, &schema), + None, + ); + let batch = |keys: Vec, values: Vec| { + RecordBatch::try_from_iter([ + ("key", Arc::new(Float64Array::from(keys)) as ArrayRef), + ("value", Arc::new(Int32Array::from(values)) as ArrayRef), + ]) + .unwrap() + }; + let mut write = table.new_write_builder().new_write().unwrap(); + write + .write_arrow_batch(&batch(vec![f64::NAN, -0.0, 0.0, 1.0], vec![0; 4])) + .await + .unwrap(); + TableCommit::new(table.clone(), "seed".into()) + .commit(write.prepare_commit().await.unwrap()) + .await + .unwrap(); + write.close().await; + table + .new_global_index_build_builder() + .with_index_column("key") + .execute() + .await + .unwrap(); + commit_upsert( + &table, + vec![batch(vec![f64::NAN, -0.0], vec![41, 42])], + vec!["key".into()], + ) + .await; + assert_eq!( + values_by_row_id(&table).await, + vec![(0, 41), (1, 42), (2, 0), (3, 0)] + ); + } + #[tokio::test] async fn upsert_casts_surviving_matched_values_and_ignores_other_payloads() { let table = numeric_table("memory:/upsert_selected").await; @@ -898,44 +1739,58 @@ mod tests { ); } - #[test] - fn partition_identity_ignores_java_null_variable_storage() { - let fields = vec![ - DataField::new( - 0, - "ts".into(), + #[tokio::test] + async fn upsert_finds_java_padded_null_timestamp_partition() { + let schema = PaimonSchema::builder() + .column( + "part", PaimonType::Timestamp(TimestampType::new(6).unwrap()), - ), - DataField::new( - 1, - "amount".into(), - PaimonType::Decimal(DecimalType::new(20, 2).unwrap()), - ), - DataField::new( - 2, - "label".into(), - PaimonType::VarChar(VarCharType::new(100).unwrap()), - ), - ]; - let mut builder = BinaryRowBuilder::new(3); - builder.set_null_at(0); - builder.set_null_at(1); - builder.write_string(2, "after-null-fields"); - let canonical = builder.build(); - let mut java = canonical.data()[..32].to_vec(); - // BinaryWriter.writeTimestamp(NULL, 6) reserves 8 bytes; - // writeDecimal(NULL, 20) reserves another 16, even though both are NULL. - java.extend_from_slice(&[0; 24]); - java.extend_from_slice(&canonical.data()[32..]); - java[8..16].copy_from_slice(&(32_u64 << 32).to_le_bytes()); - java[16..24].copy_from_slice(&(40_u64 << 32).to_le_bytes()); - let offset_and_size = u64::from_le_bytes(canonical.data()[24..32].try_into().unwrap()); - java[24..32].copy_from_slice(&(offset_and_size + (24_u64 << 32)).to_le_bytes()); - let java = BinaryRow::from_bytes(3, java); - assert_ne!(java.to_serialized_bytes(), canonical.to_serialized_bytes()); - assert_eq!( - partition_key(&java, &fields).unwrap(), - canonical.to_serialized_bytes() + ) + .column("key", PaimonType::Int(IntType::new())) + .column("value", PaimonType::Int(IntType::new())) + .partition_keys(["part"]) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .build() + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("db", "padded_null_partition"), + "memory:/upsert_padded_null_partition".into(), + TableSchema::new(0, &schema), + None, ); + let batch = |value| { + RecordBatch::try_from_iter([ + ( + "part", + Arc::new(arrow_array::TimestampMicrosecondArray::from(vec![ + None::, + ])) as ArrayRef, + ), + ("key", Arc::new(Int32Array::from(vec![7])) as ArrayRef), + ("value", Arc::new(Int32Array::from(vec![value])) as ArrayRef), + ]) + .unwrap() + }; + let mut write = table.new_write_builder().new_write().unwrap(); + write.write_arrow_batch(&batch(10)).await.unwrap(); + let mut messages = write.prepare_commit().await.unwrap(); + for message in &mut messages { + let row = BinaryRow::from_serialized_bytes(&message.partition).unwrap(); + let mut padded = row.data().to_vec(); + // Java reserves a variable slot for NULL TIMESTAMP(6), although + // the null bit makes its offset irrelevant to the logical value. + padded.extend_from_slice(&[0; 8]); + padded[8..16].copy_from_slice(&(16_u64 << 32).to_le_bytes()); + message.partition = BinaryRow::from_bytes(1, padded).to_serialized_bytes(); + } + TableCommit::new(table.clone(), "seed".into()) + .commit(messages) + .await + .unwrap(); + write.close().await; + commit_upsert(&table, vec![batch(42)], vec!["key".into()]).await; + assert_eq!(values_by_row_id(&table).await, vec![(0, 42)]); } }