diff --git a/crates/paimon/src/spec/binary_row.rs b/crates/paimon/src/spec/binary_row.rs index 15a3048c3..7f430a392 100644 --- a/crates/paimon/src/spec/binary_row.rs +++ b/crates/paimon/src/spec/binary_row.rs @@ -804,6 +804,97 @@ pub fn deserialize_binary_array_str(data: &[u8]) -> crate::Result> { Ok(out) } +/// Read a `BinaryArray` of non-null `int` (Java `array`): 4-byte element +/// slots after the header, so the layout differs from the 8-byte slots the +/// `bigint` and pointer forms use. +pub(crate) fn deserialize_binary_array_int(data: &[u8]) -> crate::Result> { + let n = read_binary_array_len(data)?; + let header = binary_array_header(n); + // Bound the count against the buffer before multiplying it, so a forged one + // cannot wrap the offset arithmetic below -- the reason + // `check_binary_array_fits` is written the way it is, which this cannot use + // because it assumes the 8-byte slot width. + if n > data.len().saturating_sub(header) / 4 { + return Err(bin_arr_err( + "int array element region exceeds buffer length", + )); + } + // An int array has no variable-length part, so the writer's own size is the + // fixed region rounded up to a word -- exactly, not at least. Requiring the + // equality rejects both an unpadded region and trailing bytes, neither of + // which the writer can emit. + if round_to_word(header + n * 4) != data.len() { + return Err(bin_arr_err( + "int array size is not its word-padded fixed region", + )); + } + let mut out = Vec::with_capacity(n); + for k in 0..n { + if data.get(4 + k / 8).is_some_and(|b| b & (1 << (k % 8)) != 0) { + return Err(bin_arr_err("int element must not be null")); + } + let eo = header + k * 4; + let slot = data + .get(eo..eo + 4) + .ok_or_else(|| bin_arr_err("int element slot out of range"))?; + out.push(i32::from_le_bytes(slot.try_into().unwrap())); + } + Ok(out) +} + +/// Read a `BinaryArray` whose elements are rows, returning each element's raw +/// bytes. Rows are addressed the way variable-length fields are, by a packed +/// offset and length, so the caller decodes each slice with the arity its own +/// schema fixes. +pub(crate) fn deserialize_binary_array_rows(data: &[u8]) -> crate::Result> { + let n = read_binary_array_len(data)?; + let header = binary_array_header(n); + // Bounds the count before `n * 8` is computed, so a forged one cannot wrap it. + check_binary_array_fits(n, header, data.len())?; + let fixed_part = round_to_word(header + n * 8); + if fixed_part > data.len() { + return Err(bin_arr_err( + "row array element region exceeds buffer length", + )); + } + // The writer appends each row after the last, word-padded, so an element body + // starts exactly where the previous one ended and the last ends at the array's + // own end. Requiring that leaves no layout the writer cannot emit: no gap, no + // trailing bytes, and no two elements sharing a body -- which would let a + // small array drive a decode many times its own size. + let mut next = fixed_part; + let mut out = Vec::with_capacity(n); + for k in 0..n { + if data.get(4 + k / 8).is_some_and(|b| b & (1 << (k % 8)) != 0) { + return Err(bin_arr_err("row element must not be null")); + } + let eo = header + k * 8; + let slot = data + .get(eo..eo + 8) + .ok_or_else(|| bin_arr_err("row element slot out of range"))?; + let encoded = u64::from_le_bytes(slot.try_into().unwrap()); + let offset = (encoded >> 32) as usize; + let length = (encoded & 0xFFFF_FFFF) as usize; + if offset != next { + return Err(bin_arr_err( + "row element body must start where the previous element ended", + )); + } + let end = offset + .checked_add(length) + .ok_or_else(|| bin_arr_err("row element bytes out of range"))?; + out.push( + data.get(offset..end) + .ok_or_else(|| bin_arr_err("row element bytes out of range"))?, + ); + next = round_to_word(end); + } + if next != data.len() { + return Err(bin_arr_err("row array has bytes after its last element")); + } + Ok(out) +} + /// Reverse of [`serialize_binary_array_long`]. pub fn deserialize_binary_array_long(data: &[u8]) -> crate::Result>> { let n = read_binary_array_len(data)?; @@ -1576,6 +1667,163 @@ mod tests { use super::*; use crate::variant::GenericVariant; + /// Java writes `array` with 4-byte element slots, so the layout differs + /// from the 8-byte forms; build one by hand and read it back. + #[test] + fn deserialize_binary_array_int_reads_four_byte_slots() { + let values: [i32; 3] = [3, 5, -9]; + // count + one null-bitset word. The writer rounds the fixed region up to + // a word, so three 4-byte slots occupy 16 bytes, not 12. + let header = 4 + 4; + let mut data = vec![0u8; round_to_word(header + values.len() * 4)]; + data[0..4].copy_from_slice(&(values.len() as i32).to_le_bytes()); + for (k, v) in values.iter().enumerate() { + let offset = header + k * 4; + data[offset..offset + 4].copy_from_slice(&v.to_le_bytes()); + } + assert_eq!( + deserialize_binary_array_int(&data).unwrap(), + values.to_vec() + ); + } + + /// A fixed region the writer would have padded is not a layout it can emit. + #[test] + fn deserialize_binary_array_int_rejects_an_unpadded_fixed_region() { + let header = 4 + 4; + let mut data = vec![0u8; header + 3 * 4]; + data[0..4].copy_from_slice(&3i32.to_le_bytes()); + assert!(deserialize_binary_array_int(&data).is_err()); + } + + #[test] + fn deserialize_binary_array_int_rejects_a_null_element() { + let header = 4 + 4; + let mut data = vec![0u8; round_to_word(header + 4)]; + data[0..4].copy_from_slice(&1i32.to_le_bytes()); + data[4] = 1; // element 0 is null + let error = deserialize_binary_array_int(&data).unwrap_err(); + assert!( + error.to_string().contains("int element must not be null"), + "unexpected error: {error}" + ); + } + + #[test] + fn deserialize_binary_array_int_rejects_a_count_past_the_buffer() { + let mut data = vec![0u8; 8]; + data[0..4].copy_from_slice(&i32::MAX.to_le_bytes()); + assert!(deserialize_binary_array_int(&data).is_err()); + } + + /// Row elements are addressed like variable-length fields, by a packed offset + /// and length, and the writer appends each body after the last, word-padded. + fn binary_array_of_rows(bodies: &[&[u8]], alias_all_at: Option) -> Vec { + let header = 4 + 4; + let mut data = vec![0u8; round_to_word(header + bodies.len() * 8)]; + data[0..4].copy_from_slice(&(bodies.len() as i32).to_le_bytes()); + let mut offsets = Vec::new(); + for body in bodies { + offsets.push(data.len()); + data.extend_from_slice(body); + let pad = (8 - (body.len() % 8)) % 8; + data.extend(std::iter::repeat_n(0u8, pad)); + } + for (k, body) in bodies.iter().enumerate() { + let (offset, length) = match alias_all_at { + Some(shared) => (shared, bodies[0].len()), + None => (offsets[k], body.len()), + }; + let encoded = ((offset as u64) << 32) | (length as u64); + let slot = header + k * 8; + data[slot..slot + 8].copy_from_slice(&encoded.to_le_bytes()); + } + data + } + + #[test] + fn deserialize_binary_array_rows_returns_element_slices() { + let data = binary_array_of_rows(&[&[0xAA; 3], &[0xBB; 5]], None); + let elements = deserialize_binary_array_rows(&data).unwrap(); + assert_eq!( + elements, + vec![[0xAAu8; 3].as_slice(), [0xBBu8; 5].as_slice()] + ); + } + + /// The writer never points two elements at one body. Accepting it would let a + /// small array drive a decode many times its own size. + #[test] + fn deserialize_binary_array_rows_rejects_aliased_elements() { + let bodies: [&[u8]; 3] = [&[0xAA; 8], &[0xBB; 8], &[0xCC; 8]]; + let first_body = round_to_word(4 + 4 + bodies.len() * 8); + let data = binary_array_of_rows(&bodies, Some(first_body)); + let error = deserialize_binary_array_rows(&data).unwrap_err(); + assert!( + error.to_string().contains("must start where the previous"), + "unexpected error: {error}" + ); + } + + /// The writer leaves no gap between bodies either, so a body one word late is + /// as much a forgery as one that overlaps. + #[test] + fn deserialize_binary_array_rows_rejects_a_gap_before_an_element() { + let header = 4 + 4; + let fixed_part = round_to_word(header + 8); + // One word of gap, then a body that is itself inside the buffer, so this + // can only fail on the gap. + let mut data = vec![0u8; fixed_part + 8 + 8]; + data[0..4].copy_from_slice(&1i32.to_le_bytes()); + let encoded = (((fixed_part + 8) as u64) << 32) | 8; + data[header..header + 8].copy_from_slice(&encoded.to_le_bytes()); + let error = deserialize_binary_array_rows(&data).unwrap_err(); + assert!( + error.to_string().contains("must start where the previous"), + "unexpected error: {error}" + ); + } + + /// The array ends where its last element ends, so trailing bytes are a layout + /// the writer cannot emit. + #[test] + fn deserialize_binary_array_rows_rejects_trailing_bytes() { + let mut data = binary_array_of_rows(&[&[0xAA; 8]], None); + data.push(0); + let error = deserialize_binary_array_rows(&data).unwrap_err(); + assert!( + error.to_string().contains("after its last element"), + "unexpected error: {error}" + ); + } + + /// Counterpart to the int reader's test: a forged count must be rejected + /// before it reserves anything. + #[test] + fn deserialize_binary_array_rows_rejects_a_count_past_the_buffer() { + let mut data = vec![0u8; 8]; + data[0..4].copy_from_slice(&i32::MAX.to_le_bytes()); + assert!(deserialize_binary_array_rows(&data).is_err()); + } + + /// An element body inside the fixed part is likewise not a layout the writer + /// can emit. + #[test] + fn deserialize_binary_array_rows_rejects_a_body_in_the_fixed_part() { + let data = binary_array_of_rows(&[&[0xAA; 8]], Some(0)); + assert!(deserialize_binary_array_rows(&data).is_err()); + } + + #[test] + fn deserialize_binary_array_rows_rejects_an_element_past_the_buffer() { + let header = 4 + 4; + let mut data = vec![0u8; header + 8]; + data[0..4].copy_from_slice(&1i32.to_le_bytes()); + let encoded = ((data.len() as u64) << 32) | 16; + data[header..header + 8].copy_from_slice(&encoded.to_le_bytes()); + assert!(deserialize_binary_array_rows(&data).is_err()); + } + #[test] fn test_empty_binary_row() { let row = BinaryRow::new(0); diff --git a/crates/paimon/src/table/goldens/bucket_vector_search_split_v1.bin b/crates/paimon/src/table/goldens/bucket_vector_search_split_v1.bin new file mode 100644 index 000000000..12b9d100b Binary files /dev/null and b/crates/paimon/src/table/goldens/bucket_vector_search_split_v1.bin differ diff --git a/crates/paimon/src/table/goldens/bucket_vector_search_split_v1_deletion_vectors.bin b/crates/paimon/src/table/goldens/bucket_vector_search_split_v1_deletion_vectors.bin new file mode 100644 index 000000000..cad97fb9e Binary files /dev/null and b/crates/paimon/src/table/goldens/bucket_vector_search_split_v1_deletion_vectors.bin differ diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index b415f0dd4..f8b99f9fa 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -68,6 +68,7 @@ mod pk_full_text_read; mod pk_full_text_scan; mod pk_search_position; mod pk_search_ranker; +mod pk_vector_bucket_split; mod pk_vector_data_file_reader; mod pk_vector_indexed_split_read; mod pk_vector_orchestrator; @@ -126,6 +127,7 @@ pub use incremental_scan::{ }; pub use lumina_index_build_builder::LuminaIndexBuildBuilder; pub use partition_stat::PartitionStat; +pub use pk_vector_bucket_split::{BucketVectorPayload, BucketVectorSearchSplit}; pub use postpone_bucket_plan::{PostponeBucketPlan, POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD}; pub use postpone_fixed_bucket_write::{ PostponeFixedBucketTableCommit, PostponeFixedBucketTableWrite, diff --git a/crates/paimon/src/table/pk_vector_bucket_split.rs b/crates/paimon/src/table/pk_vector_bucket_split.rs new file mode 100644 index 000000000..0c50b675b --- /dev/null +++ b/crates/paimon/src/table/pk_vector_bucket_split.rs @@ -0,0 +1,927 @@ +// 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. + +//! Reader for the byte form of Java's `BucketVectorSearchSplit`. +//! +//! A primary-key vector search is planned per bucket, because deciding which ANN +//! segments are current needs the bucket's complete active-file set rather than +//! the arbitrary file subset an ordinary table-scan split carries. Java plans +//! those splits and serializes them in this form; this module decodes them so +//! the search can then run outside the JVM. + +use std::collections::HashMap; + +use indexmap::IndexMap; + +use crate::spec::{ + deserialize_binary_array_int, deserialize_binary_array_rows, BinaryRow, DeletionVectorMeta, + GlobalIndexMeta, +}; +use crate::table::source::{read_i32, read_i64, read_java_utf}; +use crate::table::{DataSplit, RowRange}; + +/// `"PKVSPLIT"` in ASCII. +const MAGIC: i64 = 0x504B_5653_504C_4954; +const VERSION: i32 = 1; + +/// Field counts of the rows nested in a payload. They come from the writer's +/// schema rather than the bytes, so both sides have to agree on them; a change +/// to either schema is what `VERSION` exists to signal. +const PAYLOAD_ARITY: i32 = 7; +const GLOBAL_INDEX_ARITY: i32 = 6; +const DELETION_VECTOR_ARITY: i32 = 4; + +fn data_invalid(message: impl Into) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} + +/// All active data files and primary-key vector index payloads for one snapshot +/// bucket. +/// +/// Reference: [org.apache.paimon.table.source.BucketVectorSearchSplit](https://github.com/apache/paimon/blob/master/paimon-core/src/main/java/org/apache/paimon/table/source/BucketVectorSearchSplit.java) +#[derive(Debug, Clone, PartialEq)] +pub struct BucketVectorSearchSplit { + data_split: DataSplit, + payload_files: Vec, + /// Rows to keep per data file, as inclusive positions local to that file. + /// Kept in the order the message carries them -- the writer sorts by file + /// name -- so iterating a decoded split is reproducible; lookups are still + /// by name. + row_ranges_by_file: IndexMap>, +} + +/// One primary-key vector index file carried by a [`BucketVectorSearchSplit`]. +/// +/// A distinct type from [`crate::spec::IndexFileMeta`], which models the +/// manifest form: that one has no external path and narrows the row count to +/// `i32`, while the schema behind these bytes has both. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BucketVectorPayload { + index_type: String, + file_name: String, + file_size: i64, + row_count: i64, + deletion_vectors_ranges: Option>, + external_path: Option, + global_index_meta: GlobalIndexMeta, +} + +impl BucketVectorPayload { + pub fn index_type(&self) -> &str { + &self.index_type + } + + pub fn file_name(&self) -> &str { + &self.file_name + } + + pub fn file_size(&self) -> i64 { + self.file_size + } + + pub fn row_count(&self) -> i64 { + self.row_count + } + + /// Deletion vector ranges. Always absent on a vector payload, but the schema + /// can express them, so they are carried rather than dropped. + pub fn deletion_vectors_ranges(&self) -> Option<&IndexMap> { + self.deletion_vectors_ranges.as_ref() + } + + pub fn external_path(&self) -> Option<&str> { + self.external_path.as_deref() + } + + /// Global index metadata. The schema makes the row nullable, but a + /// `BucketVectorSearchSplit` cannot hold a payload without one, so a message + /// that omits it is rejected rather than decoded into an absence callers + /// would have to handle. + pub fn global_index_meta(&self) -> &GlobalIndexMeta { + &self.global_index_meta + } + + /// The `_SOURCE_META` blob, which maps ANN ordinals back to rows. Always + /// present, for the same reason. + pub fn source_meta(&self) -> &[u8] { + self.global_index_meta + .source_meta + .as_deref() + .expect("a decoded payload always carries source metadata") + } +} + +impl BucketVectorSearchSplit { + /// The bucket's data files. Its own `row_ranges` are always absent here -- + /// this form carries them per file in [`Self::row_ranges_by_file`] -- so + /// re-serializing it on its own would drop them. + pub fn data_split(&self) -> &DataSplit { + &self.data_split + } + + pub fn payload_files(&self) -> &[BucketVectorPayload] { + &self.payload_files + } + + pub fn row_ranges_by_file(&self) -> &IndexMap> { + &self.row_ranges_by_file + } + + /// Parse a Java `BucketVectorSearchSplit#serialize` message. + /// + /// Integers are big-endian and file names are Java modified UTF-8, following + /// the split formats this one nests. The layout is: + /// + /// ```text + /// i64 magic = "PKVSPLIT"; i32 version = 1 + /// DataSplit // DataSplit.serialize, inline + /// i32 payloadCount; (i32 rowLength, IndexFileMeta row)* + /// i32 rangeFileCount; (utf fileName, i32 rangeCount, (i64 from, i64 to)*)* + /// ``` + /// + /// Each payload is one `BinaryRow` of Java's `IndexFileMeta.SCHEMA`, which + /// carries no version of its own. `VERSION` therefore pins the layout of + /// what is nested as well as the envelope, and has to move when either + /// schema does -- unlike the `DataSplit`, which carries its own version and + /// so may change underneath a message still labelled version 1. + /// + /// Consumes the entire buffer; trailing bytes are an error. + pub fn deserialize(bytes: &[u8]) -> crate::Result { + let mut input = bytes; + let cur = &mut input; + + let magic = read_i64(cur)?; + if magic != MAGIC { + return Err(data_invalid(format!( + "invalid BucketVectorSearchSplit magic: {magic:#018x}" + ))); + } + let version = read_i32(cur)?; + if version != VERSION { + return Err(crate::Error::Unsupported { + message: format!("BucketVectorSearchSplit version {version} is not supported"), + }); + } + + let data_split = DataSplit::read_body(cur).map_err(nested_error)?; + + let payload_count = read_count(cur, "payload file")?; + let mut payload_files = Vec::new(); + for _ in 0..payload_count { + payload_files.push(read_payload(cur)?); + } + + // Row ranges are only meaningful against the data file they name, so index + // the bucket's files before reading them. Two files sharing a name would + // leave no way to say which of them a range belongs to; a bucket cannot + // hold such a pair and the planner rejects it, but these bytes are + // untrusted. + let mut row_counts: HashMap<&str, i64> = HashMap::new(); + for file in data_split.data_files() { + if row_counts + .insert(file.file_name.as_str(), file.row_count) + .is_some() + { + return Err(data_invalid(format!( + "duplicate data file in the bucket split: {}", + file.file_name + ))); + } + } + + let range_file_count = read_count(cur, "row-range file")?; + let mut row_ranges_by_file = IndexMap::new(); + for _ in 0..range_file_count { + let file_name = read_java_utf(cur)?; + let ranges = read_row_ranges(cur, &file_name, &row_counts)?; + if row_ranges_by_file + .insert(file_name.clone(), ranges) + .is_some() + { + return Err(data_invalid(format!( + "duplicate row-range file entry: {file_name}" + ))); + } + } + + if !cur.is_empty() { + return Err(data_invalid(format!( + "{} trailing bytes after BucketVectorSearchSplit", + cur.len() + ))); + } + + Ok(Self { + data_split, + payload_files, + row_ranges_by_file, + }) + } +} + +/// Report a malformed nested structure as invalid data. The decoders reached +/// from here signal a short buffer as an unexpected error, which reads as an +/// internal fault; these bytes are untrusted, so the caller has to be able to +/// tell a bad message from a bug. A version that cannot be read stays +/// unsupported. +fn nested_error(error: crate::Error) -> crate::Error { + match error { + crate::Error::Unsupported { .. } | crate::Error::DataInvalid { .. } => error, + other => crate::Error::DataInvalid { + message: "invalid nested structure in BucketVectorSearchSplit".to_string(), + source: Some(Box::new(other)), + }, + } +} + +/// Read the rows to keep in one data file, as inclusive positions local to that +/// file. Java writes ranges its planner produced and re-checks nothing, so the +/// checks here are what a reader of untrusted bytes needs rather than a mirror +/// of the writer: `RowRange` cannot represent a descending pair at all, and a +/// range outside its file would read rows that are not there. The file's own row +/// count is checked first, so a forged one cannot lift that bound. +fn read_row_ranges( + cur: &mut &[u8], + file_name: &str, + row_counts: &HashMap<&str, i64>, +) -> crate::Result> { + let row_count = *row_counts.get(file_name).ok_or_else(|| { + data_invalid(format!( + "row ranges reference data file not present in the bucket split: {file_name}" + )) + })?; + // A negative row count is forged by construction. Rejecting it, rather than + // skipping the bound check for it, is what keeps the bound below meaningful: + // otherwise a forged count would lift it entirely. + if row_count < 0 { + return Err(data_invalid(format!( + "data file {file_name} has a negative row count: {row_count}" + ))); + } + + let count = read_count(cur, "row range")?; + let mut ranges: Vec = Vec::new(); + for _ in 0..count { + let from = read_i64(cur)?; + let to = read_i64(cur)?; + if from > to { + return Err(data_invalid(format!( + "invalid row range [{from}, {to}] for file {file_name}" + ))); + } + if from < 0 { + return Err(data_invalid(format!( + "negative row range [{from}, {to}] for file {file_name}" + ))); + } + if to >= row_count { + return Err(data_invalid(format!( + "row range [{from}, {to}] for file {file_name} is outside [0, {row_count})" + ))); + } + ranges.push(RowRange::new(from, to)); + } + Ok(ranges) +} + +/// Read one payload: an `i32` byte length followed by a `BinaryRow` of Java's +/// `IndexFileMeta.SCHEMA`, the framing `ObjectSerializer` writes a record with. +fn read_payload(cur: &mut &[u8]) -> crate::Result { + let row = read_nested_row(cur, PAYLOAD_ARITY, "payload")?; + // Fields the schema declares NOT NULL. A null slot is zeroed rather than + // absent, so without this an empty name or a zero size would come back as a + // value the writer never wrote. + require_present(&row, &[0, 1, 2, 3], "payload")?; + + let deletion_vectors_ranges = if row.is_null_at(4) { + None + } else { + Some(read_deletion_vector_ranges( + row.get_binary(4).map_err(nested_error)?, + )?) + }; + let file_name = row.get_string(1).map_err(nested_error)?.to_string(); + if row.is_null_at(6) { + return Err(data_invalid(format!( + "PK-vector payload {file_name} has no global index metadata" + ))); + } + let global_index_meta = read_global_index_meta(row.get_binary(6).map_err(nested_error)?)?; + if global_index_meta.source_meta.is_none() { + return Err(data_invalid(format!( + "PK-vector payload {file_name} has no source metadata" + ))); + } + + Ok(BucketVectorPayload { + index_type: row.get_string(0).map_err(nested_error)?.to_string(), + file_name, + file_size: row.get_long(2).map_err(nested_error)?, + row_count: row.get_long(3).map_err(nested_error)?, + deletion_vectors_ranges, + external_path: if row.is_null_at(5) { + None + } else { + Some(row.get_string(5).map_err(nested_error)?.to_string()) + }, + global_index_meta, + }) +} + +fn read_global_index_meta(data: &[u8]) -> crate::Result { + let row = nested_row(data, GLOBAL_INDEX_ARITY, "global index metadata")?; + require_present(&row, &[0, 1, 2], "global index metadata")?; + Ok(GlobalIndexMeta { + row_range_start: row.get_long(0).map_err(nested_error)?, + row_range_end: row.get_long(1).map_err(nested_error)?, + index_field_id: row.get_int(2).map_err(nested_error)?, + extra_field_ids: if row.is_null_at(3) { + None + } else { + Some(deserialize_binary_array_int( + row.get_binary(3).map_err(nested_error)?, + )?) + }, + index_meta: if row.is_null_at(4) { + None + } else { + Some(row.get_binary(4).map_err(nested_error)?.to_vec()) + }, + source_meta: if row.is_null_at(5) { + None + } else { + Some(row.get_binary(5).map_err(nested_error)?.to_vec()) + }, + }) +} + +fn read_deletion_vector_ranges(data: &[u8]) -> crate::Result> { + let mut ranges = IndexMap::new(); + for element in deserialize_binary_array_rows(data)? { + let row = nested_row(element, DELETION_VECTOR_ARITY, "deletion vector metadata")?; + require_present(&row, &[0, 1, 2], "deletion vector metadata")?; + // The data file name is field 0 of the row and also the map key, the way + // Java rebuilds this map. + let file_name = row.get_string(0).map_err(nested_error)?.to_string(); + let meta = DeletionVectorMeta { + offset: row.get_int(1).map_err(nested_error)?, + length: row.get_int(2).map_err(nested_error)?, + cardinality: if row.is_null_at(3) { + None + } else { + Some(row.get_long(3).map_err(nested_error)?) + }, + }; + ranges.insert(file_name, meta); + } + Ok(ranges) +} + +/// Read an `i32`-framed row body off the cursor. +fn read_nested_row(cur: &mut &[u8], arity: i32, what: &str) -> crate::Result { + let length = read_i32(cur)?; + if length < 0 { + return Err(data_invalid(format!( + "negative {what} row length: {length}" + ))); + } + let length = length as usize; + if length > cur.len() { + return Err(data_invalid(format!( + "{what} row length {length} exceeds {} remaining bytes", + cur.len() + ))); + } + nested_row(crate::table::source::take(cur, length)?, arity, what) +} + +/// Reject a null in a field the schema declares NOT NULL. +fn require_present(row: &BinaryRow, fields: &[usize], what: &str) -> crate::Result<()> { + for &field in fields { + if row.is_null_at(field) { + return Err(data_invalid(format!( + "{what} field {field} must not be null" + ))); + } + } + Ok(()) +} + +/// View bytes that are already delimited as a row of the given arity. +fn nested_row(data: &[u8], arity: i32, what: &str) -> crate::Result { + let fixed_part = BinaryRow::cal_fix_part_size_in_bytes(arity) as usize; + if data.len() < fixed_part { + return Err(data_invalid(format!( + "{what} row of {} bytes is shorter than its {fixed_part}-byte fixed part", + data.len() + ))); + } + Ok(BinaryRow::from_bytes(arity, data.to_vec())) +} + +/// Read an element count, bounded by the bytes that can still follow it. Every +/// element of every repetition here costs at least four bytes, so this keeps a +/// forged count from driving work the message could not contain. +fn read_count(cur: &mut &[u8], element: &str) -> crate::Result { + let count = read_i32(cur)?; + if count < 0 { + return Err(data_invalid(format!("negative {element} count: {count}"))); + } + let count = count as usize; + if count > cur.len() / 4 { + return Err(data_invalid(format!( + "{element} count {count} exceeds the maximum allowed by {} remaining bytes", + cur.len() + ))); + } + Ok(count) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::PrimaryKeyIndexSourceMeta; + + /// Byte-for-byte copy of what Java's `BucketVectorSearchSplit#serialize` + /// produces, so a change on either side that is not mirrored shows up as a + /// decode failure here. Dumped from `BucketVectorSearchSplitTest.split()` in + /// apache/paimon at `088d4880ff` (#9386); the Java side keeps no golden bytes + /// of its own, so a change to that builder has to be dumped here again by + /// hand -- this test is the only thing that notices. + const GOLDEN: &[u8] = include_bytes!("goldens/bucket_vector_search_split_v1.bin"); + + fn golden() -> Vec { + GOLDEN.to_vec() + } + + /// Offsets into the fixture. The nested `DataSplit` is written inline and is + /// variable-length, so what follows it cannot be located by a formula over + /// the header; these come from the fixture's own shape, and + /// `fixture_offsets_are_current` fails if it changes. + const DATA_SPLIT_VERSION_OFFSET: usize = 8 + 4 + 8; + const PAYLOAD_COUNT_OFFSET: usize = 643; + + /// The `DataSplit`'s data-file count, which follows its deprecated + /// before-files count and before-deletion-files flag. + const DATA_FILE_COUNT_OFFSET: usize = 80; + + /// `_ROW_COUNT` inside the only data file's row: field 2 of a 21-field + /// `BinaryRow`, so past its 8-byte null region and two 8-byte slots. Written + /// little-endian, unlike the big-endian envelope around it. + const DATA_FILE_ROW_COUNT_OFFSET: usize = DATA_FILE_COUNT_OFFSET + 4 + 4 + 8 + 2 * 8; + + /// The row-range section closes the message: one modified-UTF-8 file name + /// (`u16` length + 10 bytes), its range count, and two `[from, to]` pairs. + const RANGE_SECTION_BYTES: usize = 2 + 10 + 4 + 4 * 8; + + fn payload_row_length_offset() -> usize { + PAYLOAD_COUNT_OFFSET + 4 + } + + fn range_file_name_offset() -> usize { + GOLDEN.len() - RANGE_SECTION_BYTES + 2 + } + + fn range_bound_offset(index: usize) -> usize { + GOLDEN.len() - 4 * 8 + index * 8 + } + + fn read_i32_at(bytes: &[u8], offset: usize) -> i32 { + i32::from_be_bytes(bytes[offset..offset + 4].try_into().unwrap()) + } + + /// The offsets above are hand-derived, so check them against the fixture + /// rather than letting a stale one make another test pass for the wrong + /// reason. + #[test] + fn fixture_offsets_are_current() { + assert_eq!(read_i32_at(GOLDEN, DATA_SPLIT_VERSION_OFFSET), 9); + assert_eq!( + i64::from_le_bytes( + GOLDEN[DATA_FILE_ROW_COUNT_OFFSET..DATA_FILE_ROW_COUNT_OFFSET + 8] + .try_into() + .unwrap() + ), + 6, + "expected the data file's row count" + ); + assert_eq!(read_i32_at(GOLDEN, PAYLOAD_COUNT_OFFSET), 1); + assert_eq!(read_i32_at(GOLDEN, DATA_FILE_COUNT_OFFSET), 1); + assert_eq!( + read_i32_at(GOLDEN, GOLDEN.len() - RANGE_SECTION_BYTES - 4), + 1, + "expected one row-range file entry" + ); + assert_eq!( + &GOLDEN[range_file_name_offset()..range_file_name_offset() + 10], + b"data-1.orc" + ); + } + + fn put_i32(bytes: &mut [u8], offset: usize, value: i32) { + bytes[offset..offset + 4].copy_from_slice(&value.to_be_bytes()); + } + + fn put_i64(bytes: &mut [u8], offset: usize, value: i64) { + bytes[offset..offset + 8].copy_from_slice(&value.to_be_bytes()); + } + + fn single_string_row(value: &str) -> Vec { + let mut builder = crate::spec::BinaryRowBuilder::new(1); + builder.write_bytes(0, value.as_bytes()); + builder.build_serialized() + } + + fn single_int_row(value: i32) -> Vec { + let mut builder = crate::spec::BinaryRowBuilder::new(1); + builder.write_int(0, value); + builder.build_serialized() + } + + #[test] + fn deserialize_matches_java_golden() { + let split = BucketVectorSearchSplit::deserialize(GOLDEN).unwrap(); + + let data_split = split.data_split(); + assert_eq!(data_split.snapshot_id(), 11); + assert_eq!(data_split.bucket(), 2); + assert_eq!(data_split.bucket_path(), "bucket-2"); + assert_eq!(data_split.total_buckets(), 8); + assert!(data_split.data_deletion_files().is_none()); + assert!(!data_split.raw_convertible()); + // The nested rows are the part of the message most easily misread, so the + // fixture carries real ones rather than empty rows -- and they are read + // back as values, not just as bytes. + assert_eq!( + data_split.partition().to_serialized_bytes(), + single_int_row(20_250_826) + ); + assert_eq!(data_split.partition().get_int(0).unwrap(), 20_250_826); + + let files = data_split.data_files(); + assert_eq!(files.len(), 1); + let file = &files[0]; + assert_eq!(file.file_name, "data-1.orc"); + assert_eq!(file.file_size, 1_234); + assert_eq!(file.row_count, 6); + assert_eq!(file.min_key, single_string_row("min_key")); + assert_eq!(file.max_key, single_string_row("max_key")); + for stats in [&file.key_stats, &file.value_stats] { + assert_eq!(stats.min_values(), &single_string_row("min_value")); + assert_eq!(stats.max_values(), &single_string_row("max_value")); + assert_eq!(stats.null_counts(), &[Some(0)]); + } + assert_eq!(file.min_sequence_number, 3); + assert_eq!(file.max_sequence_number, 9); + assert_eq!(file.schema_id, 7); + assert_eq!(file.level, 1); + assert!(file.extra_files.is_empty()); + assert_eq!( + file.creation_time.unwrap().timestamp_millis(), + 1_700_000_000_000 + ); + assert_eq!(file.delete_row_count, Some(0)); + assert_eq!(file.embedded_index, None); + // FileSource.COMPACT + assert_eq!(file.file_source, Some(1)); + assert_eq!(file.value_stats_cols, None); + assert_eq!(file.external_path, None); + assert_eq!(file.first_row_id, Some(40)); + assert_eq!( + file.write_cols.as_deref(), + Some(["k".to_string(), "v".to_string()].as_slice()) + ); + assert_eq!( + file.column_max_sequence_numbers.as_deref(), + Some([3i64, 9].as_slice()) + ); + + assert_eq!(split.payload_files().len(), 1); + let payload = &split.payload_files()[0]; + assert_eq!(payload.index_type(), "ivf-pq"); + assert_eq!(payload.file_name(), "ann-0.idx"); + assert_eq!(payload.file_size(), 5_000_000_000); + assert_eq!(payload.row_count(), 6); + assert_eq!(payload.deletion_vectors_ranges(), None); + assert_eq!( + payload.external_path(), + Some("s3://vector-bucket/ann-0.idx") + ); + let global = payload.global_index_meta(); + assert_eq!(global.row_range_start, 40); + assert_eq!(global.row_range_end, 45); + assert_eq!(global.index_field_id, 7); + assert_eq!( + global.extra_field_ids.as_deref(), + Some([3i32, 5].as_slice()) + ); + assert_eq!(global.index_meta.as_deref(), Some([1u8, 2, 3].as_slice())); + + let source_meta = PrimaryKeyIndexSourceMeta::from_global_index_meta(global).unwrap(); + assert_eq!(source_meta.data_level(), 1); + assert_eq!(source_meta.source_files().len(), 1); + assert_eq!(source_meta.source_files()[0].file_name(), "data-1.orc"); + assert_eq!(source_meta.source_files()[0].row_count(), 6); + + let mut expected_ranges = IndexMap::new(); + expected_ranges.insert( + "data-1.orc".to_string(), + vec![RowRange::new(0, 1), RowRange::new(4, 5)], + ); + assert_eq!(split.row_ranges_by_file(), &expected_ranges); + } + + /// A second fixture, for what the first cannot reach: deletion vector ranges + /// (the only array-of-rows in this format), more than one payload, and a + /// payload with every optional the schema allows to be absent left absent. + /// + /// No committed Java builder produces this one, so its shape is recorded here + /// to keep it reproducible. It is `BucketVectorSearchSplitTest.split()`'s + /// `dataSplit` and `rowRangesByFile` with two payloads in place of one: + /// + /// ```java + /// LinkedHashMap dv = new LinkedHashMap<>(); + /// dv.put("data-1.orc", new DeletionVectorMeta("data-1.orc", 0, 8, 2L)); + /// dv.put("data-2.orc", new DeletionVectorMeta("data-2.orc", 8, 16, null)); + /// new IndexFileMeta("ivf-pq", "ann-0.idx", 5_000_000_000L, 6, dv, + /// "s3://vector-bucket/ann-0.idx", + /// new GlobalIndexMeta(40, 45, 7, new int[] {3, 5, 9}, + /// new byte[] {1, 2, 3}, sourceMeta)); + /// new IndexFileMeta("flat", "ann-1.idx", 0, 0, null, null, + /// new GlobalIndexMeta(0, 0, 1, null, null, sourceMeta)); + /// ``` + /// + /// where `sourceMeta` is the same blob both payloads carry, the one + /// `split()` builds: `PrimaryKeyIndexSourceMeta(1, [("data-1.orc", 6)])`. + const GOLDEN_DELETION_VECTORS: &[u8] = + include_bytes!("goldens/bucket_vector_search_split_v1_deletion_vectors.bin"); + + #[test] + fn deserialize_matches_java_golden_with_deletion_vectors() { + let split = BucketVectorSearchSplit::deserialize(GOLDEN_DELETION_VECTORS).unwrap(); + assert_eq!(split.payload_files().len(), 2); + + // The nested DataSplit and the row ranges are the same as the other + // fixture's, so a change that only breaks one of the two shows up here. + assert_eq!(split.data_split().snapshot_id(), 11); + assert_eq!(split.data_split().data_files()[0].file_name, "data-1.orc"); + assert_eq!( + split.row_ranges_by_file()["data-1.orc"], + vec![RowRange::new(0, 1), RowRange::new(4, 5)] + ); + + let full = &split.payload_files()[0]; + assert_eq!(full.index_type(), "ivf-pq"); + assert_eq!(full.file_name(), "ann-0.idx"); + assert_eq!(full.file_size(), 5_000_000_000); + assert_eq!(full.row_count(), 6); + assert_eq!(full.external_path(), Some("s3://vector-bucket/ann-0.idx")); + let dv_ranges = full.deletion_vectors_ranges().unwrap(); + assert_eq!( + dv_ranges.keys().collect::>(), + ["data-1.orc", "data-2.orc"] + ); + assert_eq!( + dv_ranges["data-1.orc"], + DeletionVectorMeta { + offset: 0, + length: 8, + cardinality: Some(2), + } + ); + // Java writes a null cardinality; it must not come back as a value. + assert_eq!( + dv_ranges["data-2.orc"], + DeletionVectorMeta { + offset: 8, + length: 16, + cardinality: None, + } + ); + let global = full.global_index_meta(); + assert_eq!(global.row_range_start, 40); + assert_eq!(global.row_range_end, 45); + assert_eq!(global.index_field_id, 7); + assert_eq!( + global.extra_field_ids.as_deref(), + Some([3i32, 5, 9].as_slice()) + ); + assert_eq!(global.index_meta.as_deref(), Some([1u8, 2, 3].as_slice())); + // Parse each payload's own source metadata rather than comparing the two + // payloads': the fixture gives them the same blob, so a comparison passes + // even when one side is read off the wrong payload. + assert_source_meta(full); + + let minimal = &split.payload_files()[1]; + assert_eq!(minimal.index_type(), "flat"); + assert_eq!(minimal.file_name(), "ann-1.idx"); + assert_eq!(minimal.file_size(), 0); + assert_eq!(minimal.row_count(), 0); + assert_eq!(minimal.deletion_vectors_ranges(), None); + assert_eq!(minimal.external_path(), None); + let minimal_global = minimal.global_index_meta(); + assert_eq!(minimal_global.row_range_start, 0); + assert_eq!(minimal_global.row_range_end, 0); + assert_eq!(minimal_global.index_field_id, 1); + assert_eq!(minimal_global.extra_field_ids, None); + assert_eq!(minimal_global.index_meta, None); + assert_source_meta(minimal); + } + + /// Both payloads carry the same `_SOURCE_META`, so each is checked by parsing + /// its own rather than by comparing them. + fn assert_source_meta(payload: &BucketVectorPayload) { + let source = PrimaryKeyIndexSourceMeta::deserialize(payload.source_meta()).unwrap(); + assert_eq!(source.data_level(), 1); + assert_eq!(source.source_files().len(), 1); + assert_eq!(source.source_files()[0].file_name(), "data-1.orc"); + assert_eq!(source.source_files()[0].row_count(), 6); + } + + #[test] + fn rejects_invalid_magic() { + let mut bytes = golden(); + put_i64(&mut bytes, 0, MAGIC + 1); + assert_error_contains(&bytes, "invalid BucketVectorSearchSplit magic"); + } + + #[test] + fn rejects_unsupported_version() { + let mut bytes = golden(); + put_i32(&mut bytes, 8, VERSION + 1); + let error = BucketVectorSearchSplit::deserialize(&bytes).unwrap_err(); + assert!( + matches!(error, crate::Error::Unsupported { .. }), + "unexpected error: {error}" + ); + } + + #[test] + fn rejects_truncated_header() { + assert_error_contains(&GOLDEN[..8], "underrun"); + } + + #[test] + fn rejects_trailing_bytes() { + let mut bytes = golden(); + bytes.push(0); + assert_error_contains(&bytes, "1 trailing bytes after BucketVectorSearchSplit"); + } + + /// The nested `DataSplit` carries its own version, and one this reader cannot + /// read has to stay distinguishable from a corrupt message. + #[test] + fn rejects_unsupported_nested_data_split_version() { + let mut bytes = golden(); + put_i32(&mut bytes, DATA_SPLIT_VERSION_OFFSET, 99); + let error = BucketVectorSearchSplit::deserialize(&bytes).unwrap_err(); + assert!( + matches!(error, crate::Error::Unsupported { .. }), + "unexpected error: {error}" + ); + } + + /// A malformed nested structure is invalid data, not an internal fault: + /// callers that treat the two differently have to be able to tell them apart. + #[test] + fn reports_a_malformed_nested_row_as_invalid_data() { + let mut bytes = golden(); + // Shrink the payload row so its fixed part no longer fits. + put_i32(&mut bytes, payload_row_length_offset(), 4); + let error = BucketVectorSearchSplit::deserialize(&bytes).unwrap_err(); + assert!( + matches!(error, crate::Error::DataInvalid { .. }), + "unexpected error: {error}" + ); + } + + #[test] + fn rejects_negative_payload_row_length() { + let mut bytes = golden(); + put_i32(&mut bytes, payload_row_length_offset(), -1); + assert_error_contains(&bytes, "negative payload row length: -1"); + } + + #[test] + fn rejects_payload_row_longer_than_the_message() { + let mut bytes = golden(); + put_i32(&mut bytes, payload_row_length_offset(), i32::MAX); + assert_error_contains(&bytes, "payload row length"); + } + + #[test] + fn rejects_negative_payload_count() { + let mut bytes = golden(); + put_i32(&mut bytes, PAYLOAD_COUNT_OFFSET, -1); + assert_error_contains(&bytes, "negative payload file count: -1"); + } + + /// A count is only trusted up to what the remaining bytes could hold, so an + /// inflated one is rejected instead of driving work the message cannot + /// contain. + #[test] + fn rejects_payload_count_larger_than_the_message() { + let mut bytes = golden(); + put_i32(&mut bytes, PAYLOAD_COUNT_OFFSET, i32::MAX); + assert_error_contains(&bytes, "payload file count"); + } + + #[test] + fn rejects_descending_row_range() { + let mut bytes = golden(); + put_i64(&mut bytes, range_bound_offset(0), 42); + assert_error_contains(&bytes, "invalid row range [42, 1] for file data-1.orc"); + } + + #[test] + fn rejects_row_range_past_the_end_of_its_file() { + let mut bytes = golden(); + put_i64(&mut bytes, range_bound_offset(3), 6); + assert_error_contains(&bytes, "is outside [0, 6)"); + } + + #[test] + fn rejects_row_ranges_for_an_unknown_data_file() { + let mut bytes = golden(); + let offset = range_file_name_offset(); + assert_eq!(&bytes[offset..offset + 10], b"data-1.orc"); + bytes[offset..offset + 10].copy_from_slice(b"data-2.orc"); + assert_error_contains( + &bytes, + "row ranges reference data file not present in the bucket split: data-2.orc", + ); + } + + /// Row ranges key on the file name, so the names have to identify one file + /// each for the mapping to mean anything. + #[test] + fn rejects_duplicate_data_files() { + let mut bytes = golden(); + // Point the second data file entry at the first file's row by cloning it: + // simplest here is to rewrite the count and append a copy of the row. + let row_start = DATA_FILE_COUNT_OFFSET + 4; + let row_length = read_i32_at(&bytes, row_start) as usize; + let row = bytes[row_start..row_start + 4 + row_length].to_vec(); + put_i32(&mut bytes, DATA_FILE_COUNT_OFFSET, 2); + let insert_at = row_start + 4 + row_length; + bytes.splice(insert_at..insert_at, row); + assert_error_contains( + &bytes, + "duplicate data file in the bucket split: data-1.orc", + ); + } + + /// A null in a field the schema declares NOT NULL is zeroed rather than + /// absent, so it has to be caught by the bit and not by the value. + #[test] + fn rejects_a_null_in_a_required_payload_field() { + let mut bytes = golden(); + // Field 0 of the payload row: its null bit is bit 8, the first bit of the + // second byte of the row's null region. + let row = payload_row_length_offset() + 4; + bytes[row + 1] |= 1; + assert_error_contains(&bytes, "payload field 0 must not be null"); + } + + /// A forged row count must not be usable to lift the range bound. + #[test] + fn rejects_a_negative_data_file_row_count() { + let mut bytes = golden(); + bytes[DATA_FILE_ROW_COUNT_OFFSET..DATA_FILE_ROW_COUNT_OFFSET + 8] + .copy_from_slice(&(-1i64).to_le_bytes()); + assert_error_contains(&bytes, "data file data-1.orc has a negative row count: -1"); + } + + #[test] + fn rejects_negative_row_range() { + let mut bytes = golden(); + put_i64(&mut bytes, range_bound_offset(0), -1); + assert_error_contains(&bytes, "negative row range [-1, 1] for file data-1.orc"); + } + + fn assert_error_contains(bytes: &[u8], expected: &str) { + let error = BucketVectorSearchSplit::deserialize(bytes).unwrap_err(); + let message = error.to_string(); + assert!( + message.contains(expected), + "expected {expected:?} in {message:?}" + ); + } +} diff --git a/crates/paimon/src/table/source.rs b/crates/paimon/src/table/source.rs index 38bfdb896..aaaf66cfc 100644 --- a/crates/paimon/src/table/source.rs +++ b/crates/paimon/src/table/source.rs @@ -729,7 +729,7 @@ impl DataSplit { /// Read a `DataSplit` body from the cursor, leaving it positioned after the body /// (used both by `deserialize` and the SPLIT_V1 frame reader). - fn read_body(cur: &mut &[u8]) -> crate::Result { + pub(super) fn read_body(cur: &mut &[u8]) -> crate::Result { let magic = read_i64(cur)?; if magic != SPLIT_MAGIC { return Err(crate::Error::DataInvalid { @@ -1041,10 +1041,10 @@ fn write_deletion_list( } /// Advances `cur` by `n` bytes, returning the consumed slice. Errors on underrun. -fn take<'a>(cur: &mut &'a [u8], n: usize) -> crate::Result<&'a [u8]> { +pub(super) fn take<'a>(cur: &mut &'a [u8], n: usize) -> crate::Result<&'a [u8]> { if cur.len() < n { return Err(crate::Error::DataInvalid { - message: format!("split buffer underrun: need {n}, have {}", cur.len()), + message: format!("buffer underrun: need {n}, have {}", cur.len()), source: None, }); } @@ -1061,11 +1061,11 @@ fn read_i16(cur: &mut &[u8]) -> crate::Result { Ok(i16::from_be_bytes(take(cur, 2)?.try_into().unwrap())) } -fn read_i32(cur: &mut &[u8]) -> crate::Result { +pub(super) fn read_i32(cur: &mut &[u8]) -> crate::Result { Ok(i32::from_be_bytes(take(cur, 4)?.try_into().unwrap())) } -fn read_i64(cur: &mut &[u8]) -> crate::Result { +pub(super) fn read_i64(cur: &mut &[u8]) -> crate::Result { Ok(i64::from_be_bytes(take(cur, 8)?.try_into().unwrap())) } @@ -1079,7 +1079,7 @@ fn utf_err() -> crate::Error { /// Reverse of [`write_java_utf`]: `u16` byte-length prefix + modified UTF-8. Each UTF-16 code /// unit is encoded independently, so a supplementary char arrives as two 3-byte surrogate units; /// collect the raw `u16` units and let [`String::from_utf16`] pair the surrogates. -fn read_java_utf(cur: &mut &[u8]) -> crate::Result { +pub(super) fn read_java_utf(cur: &mut &[u8]) -> crate::Result { let len = read_i16(cur)? as u16 as usize; let bytes = take(cur, len)?; let mut units: Vec = Vec::new();