From 17111c8ca8979a314af7fbffae686b4acbd6bef2 Mon Sep 17 00:00:00 2001 From: Heng Ge Date: Thu, 28 May 2026 18:41:09 -0700 Subject: [PATCH 1/3] fix(encoding): detect per-chunk num_levels u16 overflow in miniblock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The miniblock structural codec stores each chunk's num_levels as a u16 in the v2.2 on-disk header. For shapes where a single chunk packs more than 65 535 levels (e.g. the HNSW persistence pattern: a dense level-0 prefix followed by ~6x as many mostly-empty higher-level rows, where the final chunk via slice_rest absorbs every trailing empty), num_chunk_levels as u16 silently truncated and the decoder short-read the chunk, corrupting the round trip (struct.rs:382 length-mismatch panic on read). repdef_too_sparse_for_miniblock previously only looked at the global levels/values ratio, which is healthy for this shape (~1.19) so the heuristic fell through to miniblock and the corruption fired. Extend the heuristic with any_chunk_levels_overflow_u16, a single linear pass that simulates the encoder's chunking (MAX_MINIBLOCK_VALUES visible values per non-last chunk; final chunk via slice_rest absorbs trailing empties) and returns true if any chunk would carry more than u16::MAX levels. Early-exit when the rep/def levels are absent (no list nesting) or when the total def-level count itself fits in u16, so dense / non-list paths pay zero added cost (encoder bench delta vs baseline: <0.5%, within noise). As a codec safety net, replace the silent num_chunk_levels as u16 cast at the chunk-build site with a u16::try_from that errors with a clear 'fullzip required' message — dormant in normal operation now that the heuristic catches the shape upstream, but prevents future shapes from silently corrupting data if they sneak past the heuristic. v2.2 is a stable on-disk format (is_unstable() = self >= Next, and V2_2 < Next), so widening the header isn't an option; the proper fix at the codec layer is to route affected shapes away from miniblock. --- .../src/encodings/logical/list.rs | 85 +++++++++++- .../src/encodings/logical/primitive.rs | 124 +++++++++++++++--- 2 files changed, 193 insertions(+), 16 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/list.rs b/rust/lance-encoding/src/encodings/logical/list.rs index 01422ae5dc0..6242cc94845 100644 --- a/rust/lance-encoding/src/encodings/logical/list.rs +++ b/rust/lance-encoding/src/encodings/logical/list.rs @@ -242,7 +242,9 @@ mod tests { use arrow_array::{ Array, ArrayRef, BooleanArray, DictionaryArray, LargeStringArray, ListArray, StructArray, UInt8Array, UInt64Array, - builder::{Int32Builder, Int64Builder, LargeListBuilder, ListBuilder, StringBuilder}, + builder::{ + Int32Builder, Int64Builder, LargeListBuilder, ListBuilder, StringBuilder, UInt32Builder, + }, }; use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; @@ -946,4 +948,85 @@ mod tests { check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) .await; } + + /// Builds the HNSW-flush repro shape: a dense prefix where every row has + /// `NEIGHBORS_PER_ROW` distinct values, followed by a long tail of empty + /// lists. Mirrors `HNSW::schema()` `__neighbors` / `__dists` columns: + /// dense level-0 lists, then ~6× as many mostly-empty higher-level rows. + fn make_hnsw_shaped_list_u32() -> ListArray { + const DENSE_ROWS: u32 = 40_000; + const NEIGHBORS_PER_ROW: u32 = 32; + const EMPTY_TAIL_ROWS: u32 = 240_000; + + let mut list_builder = ListBuilder::new(UInt32Builder::new()); + let mut next_val: u32 = 0; + for _ in 0..DENSE_ROWS { + for _ in 0..NEIGHBORS_PER_ROW { + list_builder.values().append_value(next_val); + next_val = next_val.wrapping_add(1); + } + list_builder.append(true); + } + for _ in 0..EMPTY_TAIL_ROWS { + list_builder.append(true); + } + list_builder.finish() + } + + /// Reproduces the HNSW-shaped variable-length `List` miniblock bug at v2.2 + /// **on the auto-routing path** (no `STRUCTURAL_ENCODING` metadata): a + /// dense prefix followed by a long tail of empty lists. Globally the data + /// looks dense, so the unfixed `repdef_too_sparse_for_miniblock` heuristic + /// picks miniblock; the final chunk's level count then overflows the u16 + /// stored in the chunk header and the read drops rows. After the heuristic + /// fix, this shape correctly routes to fullzip and the round-trip is + /// lossless. + #[test_log::test(tokio::test)] + async fn test_list_hnsw_shape_auto_routes_around_miniblock_overflow_v2_2() { + let list_array = make_hnsw_shaped_list_u32(); + let dense_rows: u64 = 40_000; + let total_rows = list_array.len() as u64; + + let field_metadata = HashMap::new(); + + let test_cases = TestCases::default() + .with_range(0..1000) + .with_range(dense_rows.saturating_sub(8)..(dense_rows + 8)) + .with_range(0..total_rows) + .with_indices(vec![0, dense_rows - 1, dense_rows, total_rows - 1]) + .with_min_file_version(LanceFileVersion::V2_2) + .with_max_file_version(LanceFileVersion::V2_2); + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) + .await; + } + + /// Companion to the auto-routing test: even when the user *explicitly* + /// sets `STRUCTURAL_ENCODING_MINIBLOCK` on the HNSW shape, the heuristic + /// must still detect the per-chunk `num_levels: u16` overflow and override + /// the request with fullzip — silently overriding a corrupt-encoding + /// preference is the safe behaviour given v2.2 is a stable on-disk format + /// whose chunk header width cannot widen. + /// + /// (The codec-side `u16::try_from` safety net at the chunk-build site is + /// dormant on this shape precisely because the heuristic intercepts it; + /// the safety net is there in case a future shape sneaks past the + /// heuristic.) + #[test_log::test(tokio::test)] + async fn test_forced_miniblock_hnsw_shape_routed_to_fullzip_v2_2() { + let list_array = make_hnsw_shaped_list_u32(); + let total_rows = list_array.len() as u64; + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_MINIBLOCK.into(), + ); + + let test_cases = TestCases::default() + .with_range(0..total_rows) + .with_min_file_version(LanceFileVersion::V2_2) + .with_max_file_version(LanceFileVersion::V2_2); + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) + .await; + } } diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 8ae0377d07a..0f9a95d69a2 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -3835,16 +3835,29 @@ impl PrimitiveStructuralEncoder { /// Checks if the rep/def levels are too sparse for miniblock encoding. /// - /// Miniblock chunks are limited to ~32KiB total. Data can use up to ~16KiB, - /// leaving ~16KiB for both rep and def buffers combined. Each chunk has at most - /// MAX_MINIBLOCK_VALUES (4096) data values, but when data has many empty/null - /// lists, the number of rep/def levels can far exceed the number of data values - /// (each empty list adds a level entry with no corresponding data value). + /// Two distinct conditions both force a fallback to fullzip: /// - /// We estimate the compressed bits per level by computing the max value in each - /// buffer and taking ceil(log2(max_val + 1)) — the minimum bits needed to - /// bitpack each level. We then calculate the maximum number of levels that fit - /// in 16KiB and compare against the actual levels-to-values ratio. + /// 1. **Compressed rep+def bytes overflow the per-chunk budget**: each + /// miniblock chunk only has ~16KiB for rep+def buffers combined, and + /// if the levels-to-values ratio is high enough that a chunk's + /// bit-packed levels would exceed that budget, we can't use miniblock. + /// This is a *global* ratio check that captures uniformly-sparse data + /// (e.g., 2.5M rows with 100 non-empty lists scattered across). + /// + /// 2. **A single chunk's `num_levels` overflows the `u16` header field** + /// on disk. v2.2 is a stable format whose per-chunk metadata stores + /// `num_levels` as a `u16`, capping each chunk at 65 535 levels. The + /// *global* ratio above can look healthy (HNSW: ~1.19 levels per value) + /// while one chunk — typically the final chunk consumed via + /// `slice_rest`, which absorbs every trailing empty list — packs far + /// more than 65 535 levels. We catch this with a single linear pass + /// over the def levels that simulates the encoder's chunking + /// (`MAX_MINIBLOCK_VALUES` visible values per non-last chunk, last + /// chunk takes the rest) and checks the worst-case per-chunk level + /// count against `u16::MAX`. + /// + /// Cheap on the dense path: when there are no rep/def levels (no nesting, + /// no nullability above the leaf) we early-exit without scanning anything. fn repdef_too_sparse_for_miniblock( repdef: &crate::repdef::SerializedRepDefs, num_values: u64, @@ -3881,16 +3894,81 @@ impl PrimitiveStructuralEncoder { return false; } - // 16KiB budget for rep+def combined (half the ~32KiB chunk limit) + // (1) Global rep+def budget check. const REPDEF_BUDGET_BITS: u64 = 16 * 1024 * 8; let max_levels_per_chunk = REPDEF_BUDGET_BITS / bits_per_level; - - // A chunk has at most MAX_MINIBLOCK_VALUES data values. The levels-to-values - // ratio tells us how many levels a chunk of that size would need. let levels_per_chunk = (num_levels as f64 / num_values as f64) * *miniblock::MAX_MINIBLOCK_VALUES as f64; + if levels_per_chunk > max_levels_per_chunk as f64 { + return true; + } + + // (2) Per-chunk num_levels u16-overflow check. The HNSW persistence + // shape (dense level-0 lists then 5x as many empty higher-level rows) + // has a low global ratio (e.g. 1.19) so (1) passes, but the trailing + // empties cluster into the final chunk via `slice_rest` and overflow. + Self::any_chunk_levels_overflow_u16(repdef, num_values) + } - levels_per_chunk > max_levels_per_chunk as f64 + /// Simulates the miniblock chunking the encoder is about to perform and + /// returns `true` if any chunk would carry more levels than the on-disk + /// `u16 num_levels` header can express (`> u16::MAX`). + /// + /// Mirrors [`RepDefSlicer::slice_next`] semantics: each non-last chunk + /// consumes levels until it has accumulated `MAX_MINIBLOCK_VALUES` + /// *visible* values (def level ≤ max_visible), leaving any boundary + /// invisibles after the chunk's last visible value for the next chunk; the + /// final chunk is via `slice_rest`, which scoops up every remaining level + /// (including any trailing all-empty rows that no value chunk would + /// otherwise have claimed). + fn any_chunk_levels_overflow_u16( + repdef: &crate::repdef::SerializedRepDefs, + num_values: u64, + ) -> bool { + // Without rep levels the encoder uses a 1:1 levels-to-values mapping, + // so each chunk has at most MAX_MINIBLOCK_VALUES levels — well under + // u16::MAX. Nothing to check on the dense / nullable-leaf-only path. + let Some(max_visible) = repdef.max_visible_level else { + return false; + }; + let Some(def_levels) = repdef.definition_levels.as_ref() else { + return false; + }; + // If even the *total* level count fits in u16, no chunk can overflow, + // so the linear scan would always return false. Cheap O(1) shortcut on + // small/medium pages. + let u16_max = u16::MAX as u64; + if (def_levels.len() as u64) <= u16_max { + return false; + } + let chunk_val_size = *miniblock::MAX_MINIBLOCK_VALUES; + + let mut visible_in_chunk: u64 = 0; + let mut levels_in_chunk: u64 = 0; + let mut visible_seen: u64 = 0; + for &def_lvl in def_levels.iter() { + levels_in_chunk += 1; + // Short-circuit: if the current chunk already exceeds u16, no + // further scanning is needed. + if levels_in_chunk > u16_max { + return true; + } + if def_lvl <= max_visible { + visible_in_chunk += 1; + visible_seen += 1; + let is_last_value = visible_seen == num_values; + if visible_in_chunk == chunk_val_size && !is_last_value { + // Encoder would call slice_next here. Boundary invisibles + // (none yet — we just finished a visible value) carry to + // the next chunk. + visible_in_chunk = 0; + levels_in_chunk = 0; + } + } + } + // Final chunk (slice_rest) takes everything left in the current + // accumulator, including any trailing empties after the last value. + levels_in_chunk > u16_max } fn prefers_fullzip(encoding_metadata: &HashMap) -> bool { @@ -4194,9 +4272,25 @@ impl PrimitiveStructuralEncoder { chunk_fixed_width.compute_stat(); let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width); let compressed_levels = compressor.compress(chunk_levels_block)?; + // `num_levels` is a `u16` in the v2.2 on-disk per-chunk header + // (`primitive.rs:467` reads it as `u16::from_le_bytes`). Silently + // truncating here used to corrupt the round-trip when chunking + // packed >65 535 levels into one chunk (HNSW persistence shape; + // see `repdef_too_sparse_for_miniblock`). The heuristic should + // route such shapes to fullzip before reaching this point; if it + // doesn't (e.g. the caller explicitly forces `MINIBLOCK`), surface + // an error rather than a silent corruption. + let num_levels_u16 = u16::try_from(num_chunk_levels).map_err(|_| { + Error::internal(format!( + "miniblock chunk has {} levels but per-chunk header is u16 \ + (max {}); this shape needs fullzip encoding", + num_chunk_levels, + u16::MAX, + )) + })?; level_chunks.push(CompressedLevelsChunk { data: compressed_levels, - num_levels: num_chunk_levels as u16, + num_levels: num_levels_u16, }); } debug_assert_eq!(levels.num_levels_remaining(), 0); From cd89af7adc226aecda6c95d211365ff4476ae718 Mon Sep 17 00:00:00 2001 From: Heng Ge Date: Thu, 28 May 2026 18:56:51 -0700 Subject: [PATCH 2/3] =?UTF-8?q?fix(encoding):=20review=20polish=20?= =?UTF-8?q?=E2=80=94=20clearer=20names,=20error=20type,=20warn=20on=20over?= =?UTF-8?q?ride?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address cosmetic review items from the round-0 cross reviews: - Correct '5x' doc comment to '~6x' (matches the 240k/40k test ratio). - Rename local 'chunk_val_size' to 'max_visibles_per_chunk' for clarity. - Use Error::invalid_input_source instead of Error::internal at the chunk-levels try_from site — this is a user-shape error, not an internal invariant violation. - Emit a log::warn when the user explicitly requested STRUCTURAL_ENCODING_MINIBLOCK but the heuristic vetoed it and routed to fullzip, so the override is observable rather than silent. No behavioural change for already-supported shapes; only metadata/log output and the error variant differ. --- .../src/encodings/logical/primitive.rs | 69 ++++++++++++------- 1 file changed, 45 insertions(+), 24 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 0f9a95d69a2..f068d18c094 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -3904,9 +3904,10 @@ impl PrimitiveStructuralEncoder { } // (2) Per-chunk num_levels u16-overflow check. The HNSW persistence - // shape (dense level-0 lists then 5x as many empty higher-level rows) - // has a low global ratio (e.g. 1.19) so (1) passes, but the trailing - // empties cluster into the final chunk via `slice_rest` and overflow. + // shape (~40k dense level-0 lists followed by ~6x as many empty + // higher-level rows) has a low global ratio (~1.19) so (1) passes, + // but the trailing empties cluster into the final chunk via + // `slice_rest` and overflow. Self::any_chunk_levels_overflow_u16(repdef, num_values) } @@ -3941,34 +3942,34 @@ impl PrimitiveStructuralEncoder { if (def_levels.len() as u64) <= u16_max { return false; } - let chunk_val_size = *miniblock::MAX_MINIBLOCK_VALUES; + let max_visibles_per_chunk = *miniblock::MAX_MINIBLOCK_VALUES; - let mut visible_in_chunk: u64 = 0; - let mut levels_in_chunk: u64 = 0; - let mut visible_seen: u64 = 0; + let mut visibles_in_current_chunk: u64 = 0; + let mut levels_in_current_chunk: u64 = 0; + let mut visibles_seen_total: u64 = 0; for &def_lvl in def_levels.iter() { - levels_in_chunk += 1; + levels_in_current_chunk += 1; // Short-circuit: if the current chunk already exceeds u16, no // further scanning is needed. - if levels_in_chunk > u16_max { + if levels_in_current_chunk > u16_max { return true; } if def_lvl <= max_visible { - visible_in_chunk += 1; - visible_seen += 1; - let is_last_value = visible_seen == num_values; - if visible_in_chunk == chunk_val_size && !is_last_value { + visibles_in_current_chunk += 1; + visibles_seen_total += 1; + let is_last_value = visibles_seen_total == num_values; + if visibles_in_current_chunk == max_visibles_per_chunk && !is_last_value { // Encoder would call slice_next here. Boundary invisibles // (none yet — we just finished a visible value) carry to // the next chunk. - visible_in_chunk = 0; - levels_in_chunk = 0; + visibles_in_current_chunk = 0; + levels_in_current_chunk = 0; } } } // Final chunk (slice_rest) takes everything left in the current // accumulator, including any trailing empties after the last value. - levels_in_chunk > u16_max + levels_in_current_chunk > u16_max } fn prefers_fullzip(encoding_metadata: &HashMap) -> bool { @@ -4278,15 +4279,18 @@ impl PrimitiveStructuralEncoder { // packed >65 535 levels into one chunk (HNSW persistence shape; // see `repdef_too_sparse_for_miniblock`). The heuristic should // route such shapes to fullzip before reaching this point; if it - // doesn't (e.g. the caller explicitly forces `MINIBLOCK`), surface - // an error rather than a silent corruption. + // doesn't (e.g. a future shape sneaks past the heuristic), surface + // a clear input-shape error rather than a silent corruption. let num_levels_u16 = u16::try_from(num_chunk_levels).map_err(|_| { - Error::internal(format!( - "miniblock chunk has {} levels but per-chunk header is u16 \ - (max {}); this shape needs fullzip encoding", - num_chunk_levels, - u16::MAX, - )) + Error::invalid_input_source( + format!( + "miniblock chunk has {} levels but per-chunk header is u16 \ + (max {}); this shape needs fullzip encoding", + num_chunk_levels, + u16::MAX, + ) + .into(), + ) })?; level_chunks.push(CompressedLevelsChunk { data: compressed_levels, @@ -5479,6 +5483,23 @@ impl PrimitiveStructuralEncoder { support_large_chunk, ) } else if too_sparse || Self::prefers_fullzip(encoding_metadata.as_ref()) { + // If the user explicitly asked for miniblock but the + // heuristic vetoed it, warn so the override is observable + // rather than silent. + let user_forced_miniblock = encoding_metadata + .get(STRUCTURAL_ENCODING_META_KEY) + .map(|v| v.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK) + .unwrap_or(false); + if too_sparse && user_forced_miniblock { + log::warn!( + "Field '{}' requested {}={} but the rep/def shape would \ + overflow miniblock per-chunk limits; encoding as fullzip \ + instead to avoid corruption.", + field.name, + STRUCTURAL_ENCODING_META_KEY, + STRUCTURAL_ENCODING_MINIBLOCK, + ); + } log::debug!( "Encoding column {} with {} items using full-zip layout", column_idx, From faa568c0b88581ec2698cfb10e967b872f8e0fc3 Mon Sep 17 00:00:00 2001 From: Heng Ge Date: Thu, 28 May 2026 19:13:28 -0700 Subject: [PATCH 3/3] fix(encoding): rename visibles -> visible to satisfy typos linter The plural 'visibles' is flagged by crate-ci/typos as a misspelling of 'visible'. Use the singular adjective form for the local variable names in any_chunk_levels_overflow_u16. --- .../src/encodings/logical/primitive.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index f068d18c094..aa6e0744191 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -3942,11 +3942,11 @@ impl PrimitiveStructuralEncoder { if (def_levels.len() as u64) <= u16_max { return false; } - let max_visibles_per_chunk = *miniblock::MAX_MINIBLOCK_VALUES; + let max_visible_per_chunk = *miniblock::MAX_MINIBLOCK_VALUES; - let mut visibles_in_current_chunk: u64 = 0; + let mut visible_in_chunk: u64 = 0; let mut levels_in_current_chunk: u64 = 0; - let mut visibles_seen_total: u64 = 0; + let mut visible_seen_total: u64 = 0; for &def_lvl in def_levels.iter() { levels_in_current_chunk += 1; // Short-circuit: if the current chunk already exceeds u16, no @@ -3955,14 +3955,14 @@ impl PrimitiveStructuralEncoder { return true; } if def_lvl <= max_visible { - visibles_in_current_chunk += 1; - visibles_seen_total += 1; - let is_last_value = visibles_seen_total == num_values; - if visibles_in_current_chunk == max_visibles_per_chunk && !is_last_value { + visible_in_chunk += 1; + visible_seen_total += 1; + let is_last_value = visible_seen_total == num_values; + if visible_in_chunk == max_visible_per_chunk && !is_last_value { // Encoder would call slice_next here. Boundary invisibles // (none yet — we just finished a visible value) carry to // the next chunk. - visibles_in_current_chunk = 0; + visible_in_chunk = 0; levels_in_current_chunk = 0; } }