From bdf0a924ff74ffc6f4945a6e5ade8043bf80f96f Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Wed, 26 Aug 2026 11:34:54 +0800 Subject: [PATCH] feat(table): decode the BucketVectorSearchSplit byte form A primary-key vector search has to be planned per bucket: deciding which ANN segments are current needs the bucket's complete active-file set, which an ordinary table-scan split does not carry. Java plans those splits and serializes them in a form a reader outside the JVM can consume; add that reader, so the search can then run natively. The message nests two things this crate could not read yet. The `DataSplit` is written inline and already had a decoder. Each payload is a `BinaryRow` of Java's `IndexFileMeta.SCHEMA`, which needed two array readers: `array` uses 4-byte element slots rather than the 8-byte ones the existing `bigint` and string readers assume, and the deletion vector ranges are an array of rows, addressed the way variable-length fields are. Nested rows themselves reuse the existing pattern of reading a field's bytes and viewing them at the arity its schema fixes. Row ranges are checked against the data file they name, and a payload without global index metadata is rejected: the Java constructor enforces that, so it is not a state a real split reaches, but these bytes are untrusted. A malformed nested structure is reported as invalid data rather than as an internal fault, for the same reason. The checks are exact wherever the writer's layout is exact: an element body in an array of rows starts where the previous one ended and the array ends where its last element does, an `array` is its word-padded fixed region and nothing more, and a negative row count on a data file is rejected rather than taken as a reason to skip the bound it sets. Element counts are bounded against the buffer before they are multiplied, so a forged one cannot wrap the offset arithmetic that follows. Pinned against two messages generated by the Java writer, one of them carrying the deletion vector ranges and the absent optionals the other cannot reach. --- crates/paimon/src/spec/binary_row.rs | 248 +++++ .../goldens/bucket_vector_search_split_v1.bin | Bin 0 -> 919 bytes ...ector_search_split_v1_deletion_vectors.bin | Bin 0 -> 1235 bytes crates/paimon/src/table/mod.rs | 2 + .../src/table/pk_vector_bucket_split.rs | 927 ++++++++++++++++++ crates/paimon/src/table/source.rs | 12 +- 6 files changed, 1183 insertions(+), 6 deletions(-) create mode 100644 crates/paimon/src/table/goldens/bucket_vector_search_split_v1.bin create mode 100644 crates/paimon/src/table/goldens/bucket_vector_search_split_v1_deletion_vectors.bin create mode 100644 crates/paimon/src/table/pk_vector_bucket_split.rs 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 0000000000000000000000000000000000000000..12b9d100b1371cd14a874b238d7e91ffba288215 GIT binary patch literal 919 zcmc&yJx{_=6uqsWhzo;*CJqb^##lg&iE(1o=s==zH5E!|s(=U$Y#I|C{4@Rp6E_zV zCMGvG2jb-DIq&rW7I#l_-u=Gs>y{2rk4r}fCn93Kzw`?m%b!)~^B}PRE`k;f;#GEx zZcFOCQx3egvqe`*#j(J!*F=glbfm#2u%DnMGOLAZ`V5`T2jC4jjARoXR55IGkU8jt z*_^u#Y1BU69*vU;5&m$g2o9#Og!LOqlo?y3;ZLC`FM)TGb7&Czi1b4lwF;D&g literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..cad97fb9e59c4c504137c7eaf07f1e1885a21694 GIT binary patch literal 1235 zcmc&zyGlbr5Z!xCG{z={g@Q#23lVQ#B8XK=MT>wUc5=;wgv2DqB*acpRPaYE{RF|z zMo494XJ-+tQhLtroh0Ugr3cQ=>^$}{JDb`)+)wT894Mu%kC%3AZSgAyeHJ9IOZVy`HboT4EJK7nl?XG%>Qp%;CI&c)xsOR(d~Dpt_R zn9D(CpcAq=ck6JH|9pE8BLh5)F=-DDrm=vriy+F7tyIRXyG}e(siln{V{eJmLixCu zy}0u{eY@Klv}^Y=#xO{^4Tiuf2=wxJC)&Ag(=9c!^!70thkb_QdO`g>t^N1@D5XZx zh^wY2#XkJ5#tpBPH>y&FX3nXesauvQN#nPFmDiQlxrd$tmq4}-%i9Ira@mPR3Yivk zLblw5z5r4(Pr%lccAzcbIV5kSCA#w`^U!>oEMhhm?}i56V}8%5oV7M{dC~KB^c-AXjqh5I4PAYwP6U2rll; l3&aO%tqVLRZBmORll{{gjAWVZkS literal 0 HcmV?d00001 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();