From 9230361b453eb9caf85d823719e7a786e3374f0f Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 4 Jul 2026 12:48:35 +0800 Subject: [PATCH 1/8] feat: add sparse structural encoding --- docs/src/format/file/versioning.md | 2 +- protos/encodings_v2_1.proto | 62 + rust/lance-encoding/src/constants.rs | 2 + .../src/encodings/logical/list.rs | 126 +- .../src/encodings/logical/primitive.rs | 2071 ++++++++++++++++- .../src/encodings/physical/bitpacking.rs | 8 + rust/lance-encoding/src/format.rs | 25 + rust/lance-encoding/src/repdef.rs | 604 +++++ rust/lance-encoding/src/testing.rs | 6 + 9 files changed, 2804 insertions(+), 102 deletions(-) diff --git a/docs/src/format/file/versioning.md b/docs/src/format/file/versioning.md index 2add82ef2fa..aa352c2abde 100644 --- a/docs/src/format/file/versioning.md +++ b/docs/src/format/file/versioning.md @@ -21,7 +21,7 @@ The following values are supported: | 2.0 | 0.16.0 | Any | Rework of the Lance file format that removed row groups and introduced null support for lists, fixed size lists, and primitives | | 2.1 | 0.38.1 | Any | Enhances integer and string compression, adds support for nulls in struct fields, and improves random access performance with nested fields. | | 2.2 | None | Any | Adds support for newer nested type/encoding capabilities (including map support) and 2.2-era storage features. | -| 2.3 (unstable) | None | Any | Adds experimental encodings for upcoming features. | +| 2.3 (unstable) | None | Any | Adds experimental encodings for upcoming features, including sparse structural encoding. | | legacy | N/A | N/A | Alias for 0.1 | | stable | N/A | N/A | Alias for the default version for new datasets in the Lance release you are running. | | next | N/A | N/A | Alias for the latest unstable version in the Lance release you are running.| diff --git a/protos/encodings_v2_1.proto b/protos/encodings_v2_1.proto index 46fd012fb58..9a8b90b3a14 100644 --- a/protos/encodings_v2_1.proto +++ b/protos/encodings_v2_1.proto @@ -147,6 +147,66 @@ message FullZipLayout { repeated RepDefLayer layers = 8; } +// A layout used for sparse nested pages where the Arrow structure is discrete. +// +// Sparse layout stores structural layers as native slot-domain mappings instead of dense +// repetition / definition events: +// +// - validity layer: slot -> valid/null +// - list layer: parent slot -> child slot range, plus null parent slots if present +// - fixed-size-list layer: parent slot -> deterministic child slot range, plus null parent slots +// +// Value buffers are still mini-block compressed, but the structural side is a sparse index over +// Arrow structure. +message SparseLayout { + // Description of the compression of values. + CompressiveEncoding value_compression = 3; + // The number of value buffers in each chunk. This does NOT include repetition or definition + // buffers. + uint64 num_buffers = 5; + // The number of structural items in the page, including invisible structural items. + uint64 num_items = 7; + // The number of visible items in the page. + uint64 num_visible_items = 8; + // If true, chunk-local value buffer sizes are stored as u32 values. If false, they are stored + // as u16 values, matching mini-block chunk payload metadata. + bool has_large_chunk = 9; + // Structural layers, stored from outer-most to inner-most. + repeated SparseStructuralLayer structural_layers = 10; +} + +message SparseStructuralLayer { + enum Kind { + SPARSE_LAYER_UNSPECIFIED = 0; + // A nullable item / struct layer. Null slots are recorded in the parent slot domain. + SPARSE_LAYER_VALIDITY = 1; + // A list-like layer. Non-empty parent slots map to child ranges; missing valid slots are empty. + SPARSE_LAYER_LIST = 2; + // A fixed-size-list layer. Child ranges are deterministic from constant_count. + SPARSE_LAYER_FIXED_SIZE_LIST = 3; + } + + Kind kind = 1; + // Number of slots in this layer's parent domain. + uint64 num_slots = 2; + // Number of slots in this layer's child domain after this layer is applied. For validity + // layers this matches num_slots because Arrow child slot numbering is unchanged. + uint64 num_child_slots = 3; + // For LIST, this is the child count for every non-empty slot when count_compression is absent. + // For FIXED_SIZE_LIST, this is the fixed list dimension. + uint64 constant_count = 4; + // Delta-compressed u64 positions for non-empty LIST slots. Absent for non-list layers. + CompressiveEncoding non_empty_compression = 5; + // Number of non-empty LIST positions. + uint64 num_non_empty = 6; + // Optional compressed u64 child counts for non-empty LIST slots. + CompressiveEncoding count_compression = 7; + // Optional delta-compressed u64 null positions for nullable layers. + CompressiveEncoding null_compression = 8; + // Number of null positions. + uint64 num_nulls = 9; +} + // A layout used for pages where all (visible) values are the same scalar value. // // This generalizes the prior AllNullLayout semantics for file_version >= 2.2. @@ -206,6 +266,8 @@ message PageLayout { // A layout where large binary data is encoded externally // and only the descriptions are put in the page BlobLayout blob_layout = 4; + // A layout used for sparse nested pages with native structural layers (2.3+). + SparseLayout sparse_layout = 5; } } diff --git a/rust/lance-encoding/src/constants.rs b/rust/lance-encoding/src/constants.rs index c95b587a532..0bd31d676cc 100644 --- a/rust/lance-encoding/src/constants.rs +++ b/rust/lance-encoding/src/constants.rs @@ -55,6 +55,8 @@ pub const STRUCTURAL_ENCODING_META_KEY: &str = "lance-encoding:structural-encodi pub const STRUCTURAL_ENCODING_MINIBLOCK: &str = "miniblock"; /// Value for fullzip structural encoding pub const STRUCTURAL_ENCODING_FULLZIP: &str = "fullzip"; +/// Value for sparse structural encoding +pub const STRUCTURAL_ENCODING_SPARSE: &str = "sparse"; // Byte stream split metadata keys /// Metadata key for byte stream split encoding configuration diff --git a/rust/lance-encoding/src/encodings/logical/list.rs b/rust/lance-encoding/src/encodings/logical/list.rs index 250eb476671..9b33060d27c 100644 --- a/rust/lance-encoding/src/encodings/logical/list.rs +++ b/rust/lance-encoding/src/encodings/logical/list.rs @@ -267,20 +267,34 @@ mod tests { async fn try_encode_v22_pages( array: ArrayRef, ) -> lance_core::Result> { - try_encode_v22_pages_with_metadata(array, HashMap::new()).await + try_encode_pages(array, HashMap::new(), LanceFileVersion::V2_2).await } async fn try_encode_v22_pages_with_metadata( array: ArrayRef, field_metadata: HashMap, + ) -> lance_core::Result> { + try_encode_pages(array, field_metadata, LanceFileVersion::V2_2).await + } + + async fn try_encode_v23_pages( + array: ArrayRef, + ) -> lance_core::Result> { + try_encode_pages(array, HashMap::new(), LanceFileVersion::V2_3).await + } + + async fn try_encode_pages( + array: ArrayRef, + field_metadata: HashMap, + version: LanceFileVersion, ) -> lance_core::Result> { let arrow_field = Field::new("", array.data_type().clone(), true).with_metadata(field_metadata); let lance_field = lance_core::datatypes::Field::try_from(&arrow_field).unwrap(); - let encoding_strategy = crate::encoder::default_encoding_strategy(LanceFileVersion::V2_2); + let encoding_strategy = crate::encoder::default_encoding_strategy(version); let mut column_index_seq = crate::encoder::ColumnIndexSequence::default(); let encoding_options = crate::encoder::EncodingOptions { - version: LanceFileVersion::V2_2, + version, ..Default::default() }; let mut encoder = encoding_strategy @@ -317,6 +331,10 @@ mod tests { try_encode_v22_pages(array).await.unwrap() } + async fn encode_v23_pages(array: ArrayRef) -> Vec { + try_encode_v23_pages(array).await.unwrap() + } + fn assert_split_miniblock_layout( pages: &[crate::encoder::EncodedPage], expect_structural_only_page: bool, @@ -324,6 +342,7 @@ mod tests { let mut miniblock_pages = 0; let mut fullzip_pages = 0; let mut structural_only_pages = 0; + let mut sparse_pages = 0; for page in pages { let crate::decoder::PageEncoding::Structural(layout) = &page.description else { @@ -336,6 +355,9 @@ mod tests { crate::format::pb21::page_layout::Layout::FullZipLayout(_) => { fullzip_pages += 1; } + crate::format::pb21::page_layout::Layout::SparseLayout(_) => { + sparse_pages += 1; + } crate::format::pb21::page_layout::Layout::ConstantLayout(layout) => { if layout.inline_value.is_none() && (layout.num_rep_values > 0 || layout.num_def_values > 0) @@ -355,6 +377,10 @@ mod tests { fullzip_pages, 0, "split list pages should not fall back to full-zip" ); + assert_eq!( + sparse_pages, 0, + "forced mini-block pages should not use sparse layout" + ); if expect_structural_only_page { assert!( structural_only_pages > 0, @@ -363,6 +389,59 @@ mod tests { } } + fn assert_sparse_layout(pages: &[crate::encoder::EncodedPage]) { + let mut sparse_pages = 0; + let mut fullzip_pages = 0; + + for page in pages { + let crate::decoder::PageEncoding::Structural(layout) = &page.description else { + continue; + }; + match layout.layout.as_ref().unwrap() { + crate::format::pb21::page_layout::Layout::SparseLayout(_) => { + sparse_pages += 1; + } + crate::format::pb21::page_layout::Layout::FullZipLayout(_) => { + fullzip_pages += 1; + } + crate::format::pb21::page_layout::Layout::MiniBlockLayout(_) + | crate::format::pb21::page_layout::Layout::ConstantLayout(_) + | crate::format::pb21::page_layout::Layout::BlobLayout(_) => {} + } + } + + assert!(sparse_pages > 0, "expected at least one sparse page"); + assert_eq!( + fullzip_pages, 0, + "sparse list pages should not fall back to full-zip" + ); + } + + fn sparse_structural_list_stats(pages: &[crate::encoder::EncodedPage]) -> (u64, Vec) { + let mut num_non_empty_slots = 0; + let mut constant_counts = Vec::new(); + for page in pages { + let crate::decoder::PageEncoding::Structural(layout) = &page.description else { + continue; + }; + let crate::format::pb21::page_layout::Layout::SparseLayout(layout) = + layout.layout.as_ref().unwrap() + else { + continue; + }; + for layer in &layout.structural_layers { + if matches!( + crate::format::pb21::sparse_structural_layer::Kind::try_from(layer.kind), + Ok(crate::format::pb21::sparse_structural_layer::Kind::SparseLayerList) + ) { + num_non_empty_slots += layer.num_non_empty; + constant_counts.push(layer.constant_count); + } + } + } + (num_non_empty_slots, constant_counts) + } + fn assert_has_fullzip_layout(pages: &[crate::encoder::EncodedPage]) { let has_fullzip = pages.iter().any(|page| { let crate::decoder::PageEncoding::Structural(layout) = &page.description else { @@ -1393,13 +1472,9 @@ mod tests { /// Reproduces the HNSW-flush shape at v2.2 on the auto path (no /// `STRUCTURAL_ENCODING` metadata): a dense level-0 prefix followed by a - /// long tail of empty lists. The global levels/values ratio looks dense, - /// so this used to encode as a single mini-block page whose final chunk - /// absorbed every trailing empty list and overflowed the per-chunk `u16` - /// `num_levels`, corrupting the read. The structural page planner now - /// splits on top-level row boundaries: the dense prefix stays on - /// mini-block pages and the empty tail becomes structural-only pages, so - /// the round-trip is lossless without falling back to full-zip. + /// long tail of empty lists. Since v2.2 is stable, this path must keep the + /// pre-sparse behavior and split into mini-block-compatible pages instead + /// of emitting `SparseLayout`. #[test_log::test(tokio::test)] async fn test_list_hnsw_shape_splits_to_miniblock_v2_2() { let list_array = make_hnsw_shaped_list_u32(); @@ -1419,6 +1494,37 @@ mod tests { check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; } + /// Reproduces the HNSW-flush shape at v2.3 on the auto path (no + /// `STRUCTURAL_ENCODING` metadata): a dense level-0 prefix followed by a + /// long tail of empty lists. The global levels/values ratio looks dense, + /// so this used to encode as a single mini-block page whose final chunk + /// absorbed every trailing empty list and overflowed the per-chunk `u16` + /// `num_levels`, corrupting the read. The structural page planner now + /// uses sparse layout on the auto path, preserving mini-block-style + /// value compression while replacing dense rep/def structural events with + /// native sparse structural layers. + #[test_log::test(tokio::test)] + async fn test_list_hnsw_shape_uses_sparse_v2_3() { + let list_array = make_hnsw_shaped_list_u32(); + let dense_rows: u64 = 40_000; + let total_rows = list_array.len() as u64; + + 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_3) + .with_max_file_version(LanceFileVersion::V2_3); + let list_array = Arc::new(list_array) as ArrayRef; + let pages = encode_v23_pages(list_array.clone()).await; + assert_sparse_layout(&pages); + let (num_non_empty_rows, constant_counts) = sparse_structural_list_stats(&pages); + assert_eq!(num_non_empty_rows, dense_rows); + assert_eq!(constant_counts, vec![32]); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; + } + /// Companion to the auto-path test: even when the user explicitly requests /// `STRUCTURAL_ENCODING_MINIBLOCK`, the structural page planner splits the /// HNSW shape so every emitted page fits the mini-block per-chunk budget. diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 8b6eed5d757..aebb334782d 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -15,6 +15,7 @@ use std::{ use crate::{ constants::{ STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, + STRUCTURAL_ENCODING_SPARSE, }, data::DictionaryDataBlock, encodings::logical::primitive::blob::{BlobDescriptionPageScheduler, BlobPageScheduler}, @@ -23,7 +24,9 @@ use crate::{ pb21::{self, CompressiveEncoding, PageLayout, compressive_encoding::Compression}, }, }; -use arrow_array::{Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, types::UInt64Type}; +use arrow_array::{ + Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, new_empty_array, types::UInt64Type, +}; use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field as ArrowField}; use bytes::Bytes; @@ -58,7 +61,8 @@ use crate::{ use crate::{ repdef::{ CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation, - RepDefSlicer, SerializedRepDefs, StructuralPagePlan, build_control_word_iterator, + RepDefSlicer, SerializedRepDefs, SparseStructuralLayerPlan, SparseStructuralPlan, + StructuralPagePlan, build_control_word_iterator, }, utils::accumulation::AccumulationQueue, }; @@ -672,6 +676,956 @@ impl Clone for LoadedChunk { } } +#[derive(Debug, Clone)] +enum SparseLayerDecompressors { + Validity { + num_slots: u64, + null_decompressor: Option>, + num_nulls: u64, + }, + List { + num_slots: u64, + num_child_slots: u64, + non_empty_decompressor: Option>, + num_non_empty: u64, + count_decompressor: Option>, + constant_count: u64, + null_decompressor: Option>, + num_nulls: u64, + }, + FixedSizeList { + num_slots: u64, + dimension: u64, + null_decompressor: Option>, + num_nulls: u64, + }, +} + +#[derive(Debug)] +struct SparseStructuralCacheableState { + chunk_meta: Vec, + chunk_value_offsets: Arc<[u64]>, + plan: SparseStructuralPlan, +} + +impl DeepSizeOf for SparseStructuralCacheableState { + fn deep_size_of_children(&self, _context: &mut Context) -> usize { + let structural_size = self + .plan + .layers + .iter() + .map(|layer| match layer { + SparseStructuralLayerPlan::Validity { null_positions, .. } => { + null_positions.len() * std::mem::size_of::() + } + SparseStructuralLayerPlan::List { + non_empty_positions, + counts, + null_positions, + .. + } => { + (non_empty_positions.len() + counts.len() + null_positions.len()) + * std::mem::size_of::() + } + SparseStructuralLayerPlan::FixedSizeList { null_positions, .. } => { + null_positions.len() * std::mem::size_of::() + } + }) + .sum::(); + self.chunk_meta.len() * std::mem::size_of::() + + self.chunk_value_offsets.len() * std::mem::size_of::() + + structural_size + } +} + +impl CachedPageData for SparseStructuralCacheableState { + fn as_arc_any(self: Arc) -> Arc { + self + } +} + +#[derive(Debug)] +struct SparseStructuralScheduler { + buffer_offsets_and_sizes: Vec<(u64, u64)>, + priority: u64, + num_items: u64, + num_visible_items: u64, + num_buffers: u64, + value_decompressor: Arc, + layer_decompressors: Vec, + data_type: DataType, + page_meta: Option>, + has_large_chunk: bool, +} + +impl SparseStructuralScheduler { + fn try_new( + buffer_offsets_and_sizes: &[(u64, u64)], + priority: u64, + data_type: DataType, + layout: &pb21::SparseLayout, + decompressors: &dyn DecompressionStrategy, + ) -> Result { + let value_decompressor = decompressors.create_miniblock_decompressor( + layout.value_compression.as_ref().unwrap(), + decompressors, + )?; + let layer_decompressors = layout + .structural_layers + .iter() + .map(|layer| Self::layer_decompressors(layer, decompressors)) + .collect::>>()?; + + Ok(Self { + buffer_offsets_and_sizes: buffer_offsets_and_sizes.to_vec(), + priority, + num_items: layout.num_items, + num_visible_items: layout.num_visible_items, + num_buffers: layout.num_buffers, + value_decompressor: value_decompressor.into(), + layer_decompressors, + data_type, + page_meta: None, + has_large_chunk: layout.has_large_chunk, + }) + } + + fn layer_decompressors( + layer: &pb21::SparseStructuralLayer, + decompressors: &dyn DecompressionStrategy, + ) -> Result { + let null_decompressor = layer + .null_compression + .as_ref() + .map(|compression| { + decompressors + .create_block_decompressor(compression) + .map(Arc::from) + }) + .transpose()?; + let kind = pb21::sparse_structural_layer::Kind::try_from(layer.kind)?; + Ok(match kind { + pb21::sparse_structural_layer::Kind::SparseLayerValidity => { + SparseLayerDecompressors::Validity { + num_slots: layer.num_slots, + null_decompressor, + num_nulls: layer.num_nulls, + } + } + pb21::sparse_structural_layer::Kind::SparseLayerList => { + let non_empty_decompressor = layer + .non_empty_compression + .as_ref() + .map(|compression| { + decompressors + .create_block_decompressor(compression) + .map(Arc::from) + }) + .transpose()?; + let count_decompressor = layer + .count_compression + .as_ref() + .map(|compression| { + decompressors + .create_block_decompressor(compression) + .map(Arc::from) + }) + .transpose()?; + SparseLayerDecompressors::List { + num_slots: layer.num_slots, + num_child_slots: layer.num_child_slots, + non_empty_decompressor, + num_non_empty: layer.num_non_empty, + count_decompressor, + constant_count: layer.constant_count, + null_decompressor, + num_nulls: layer.num_nulls, + } + } + pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList => { + SparseLayerDecompressors::FixedSizeList { + num_slots: layer.num_slots, + dimension: layer.constant_count, + null_decompressor, + num_nulls: layer.num_nulls, + } + } + pb21::sparse_structural_layer::Kind::SparseLayerUnspecified => { + return Err(Error::invalid_input_source( + "Sparse structural layer kind is unspecified".into(), + )); + } + }) + } + + fn parse_chunk_meta(&self, meta_bytes: Bytes) -> Result> { + if !meta_bytes.len().is_multiple_of(8) { + return Err(Error::invalid_input_source( + format!( + "Sparse layout metadata length {} is not a multiple of 8", + meta_bytes.len() + ) + .into(), + )); + } + + let value_buf_position = self.buffer_offsets_and_sizes[1].0; + let mut rows_counter = 0; + let mut offset_bytes = value_buf_position; + let mut chunk_meta = Vec::with_capacity(meta_bytes.len() / 8); + for chunk in meta_bytes.chunks_exact(8) { + let divided_bytes_minus_one = + u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]); + let num_values = u32::from_le_bytes([chunk[4], chunk[5], chunk[6], chunk[7]]) as u64; + let num_bytes = (divided_bytes_minus_one as usize + 1) * MINIBLOCK_ALIGNMENT; + rows_counter += num_values; + chunk_meta.push(ChunkMeta { + num_values, + chunk_size_bytes: num_bytes as u64, + offset_bytes, + }); + offset_bytes += num_bytes as u64; + } + if rows_counter != self.num_visible_items { + return Err(Error::invalid_input_source( + format!( + "Sparse layout visible item count mismatch: metadata has {}, layout has {}", + rows_counter, self.num_visible_items + ) + .into(), + )); + } + Ok(chunk_meta) + } + + fn decode_u64_values( + decompressor: &dyn BlockDecompressor, + data: Bytes, + num_values: u64, + label: &str, + ) -> Result> { + let decoded = decompressor.decompress(LanceBuffer::from_bytes(data, 1), num_values)?; + let fixed = decoded.as_fixed_width().ok_or_else(|| { + Error::internal(format!( + "Sparse structural {label} did not decode to fixed width data" + )) + })?; + if fixed.bits_per_value != 64 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} decoded to {} bits per value, expected 64", + fixed.bits_per_value + ) + .into(), + )); + } + Ok(fixed.data.borrow_to_typed_slice::().to_vec()) + } + + fn decode_positions( + decompressor: Option<&Arc>, + data: Option, + num_positions: u64, + num_slots: u64, + label: &str, + ) -> Result> { + if num_positions == 0 { + return Ok(Vec::new()); + } + let decompressor = decompressor.ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} has positions but no compression").into(), + ) + })?; + let data = data.ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} is missing its position buffer").into(), + ) + })?; + let deltas = Self::decode_u64_values(decompressor.as_ref(), data, num_positions, label)?; + let mut positions = Vec::with_capacity(deltas.len()); + let mut current = 0_u64; + for (idx, delta) in deltas.into_iter().enumerate() { + if idx == 0 { + current = delta; + } else { + current = current.checked_add(delta).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} position overflow").into(), + ) + })?; + } + if current >= num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} position {} is outside layer with {} slots", + current, num_slots + ) + .into(), + )); + } + positions.push(current); + } + Ok(positions) + } + + fn decode_layer( + layer: &SparseLayerDecompressors, + buffers: &mut impl Iterator, + ) -> Result { + Ok(match layer { + SparseLayerDecompressors::Validity { + num_slots, + null_decompressor, + num_nulls, + } => SparseStructuralLayerPlan::Validity { + num_slots: *num_slots, + null_positions: Self::decode_positions( + null_decompressor.as_ref(), + (*num_nulls > 0).then(|| buffers.next().unwrap()), + *num_nulls, + *num_slots, + "validity null positions", + )?, + }, + SparseLayerDecompressors::List { + num_slots, + num_child_slots, + non_empty_decompressor, + num_non_empty, + count_decompressor, + constant_count, + null_decompressor, + num_nulls, + } => { + let non_empty_positions = Self::decode_positions( + non_empty_decompressor.as_ref(), + (*num_non_empty > 0).then(|| buffers.next().unwrap()), + *num_non_empty, + *num_slots, + "list non-empty positions", + )?; + let counts = if let Some(count_decompressor) = count_decompressor { + Self::decode_u64_values( + count_decompressor.as_ref(), + buffers.next().unwrap(), + *num_non_empty, + "list counts", + )? + } else { + vec![*constant_count; *num_non_empty as usize] + }; + let null_positions = Self::decode_positions( + null_decompressor.as_ref(), + (*num_nulls > 0).then(|| buffers.next().unwrap()), + *num_nulls, + *num_slots, + "list null positions", + )?; + let actual_child_slots = counts.iter().sum::(); + if actual_child_slots != *num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list count sum {} does not match declared child slots {}", + actual_child_slots, num_child_slots + ) + .into(), + )); + } + SparseStructuralLayerPlan::List { + num_slots: *num_slots, + num_child_slots: *num_child_slots, + non_empty_positions, + counts, + constant_count: (*constant_count > 0).then_some(*constant_count), + null_positions, + } + } + SparseLayerDecompressors::FixedSizeList { + num_slots, + dimension, + null_decompressor, + num_nulls, + } => SparseStructuralLayerPlan::FixedSizeList { + num_slots: *num_slots, + dimension: *dimension, + null_positions: Self::decode_positions( + null_decompressor.as_ref(), + (*num_nulls > 0).then(|| buffers.next().unwrap()), + *num_nulls, + *num_slots, + "fixed-size-list null positions", + )?, + }, + }) + } + + fn lookup_value_chunks(&self, chunk_indices: &[usize]) -> Vec { + let page_meta = self.page_meta.as_ref().unwrap(); + chunk_indices + .iter() + .map(|&chunk_idx| { + let chunk_meta = &page_meta.chunk_meta[chunk_idx]; + let bytes_start = chunk_meta.offset_bytes; + let bytes_end = bytes_start + chunk_meta.chunk_size_bytes; + LoadedChunk { + byte_range: bytes_start..bytes_end, + items_in_chunk: chunk_meta.num_values, + chunk_idx, + data: LanceBuffer::empty(), + } + }) + .collect() + } + + fn value_chunk_range(chunk_value_offsets: &[u64], value_range: Range) -> Range { + let start = chunk_value_offsets.partition_point(|&offset| offset <= value_range.start) - 1; + let end = chunk_value_offsets + .partition_point(|&offset| offset < value_range.end) + .max(start + 1); + start..end + } +} + +impl StructuralPageScheduler for SparseStructuralScheduler { + fn initialize<'a>( + &'a mut self, + io: &Arc, + ) -> BoxFuture<'a, Result>> { + let (meta_buf_position, meta_buf_size) = self.buffer_offsets_and_sizes[0]; + let mut required_ranges = Vec::new(); + required_ranges.push(meta_buf_position..meta_buf_position + meta_buf_size); + for (position, size) in self.buffer_offsets_and_sizes.iter().skip(2) { + required_ranges.push(*position..*position + *size); + } + let io_req = io.submit_request(required_ranges, 0); + + async move { + let mut buffers = io_req.await?.into_iter(); + let meta_bytes = buffers.next().unwrap(); + + let chunk_meta = self.parse_chunk_meta(meta_bytes)?; + let mut chunk_value_offsets = Vec::with_capacity(chunk_meta.len() + 1); + let mut value_offset = 0; + chunk_value_offsets.push(value_offset); + for chunk in &chunk_meta { + value_offset += chunk.num_values; + chunk_value_offsets.push(value_offset); + } + + let layers = self + .layer_decompressors + .iter() + .map(|layer| Self::decode_layer(layer, &mut buffers)) + .collect::>>()?; + let plan = SparseStructuralPlan { + layers, + num_items: self.num_items, + num_visible_items: self.num_visible_items, + }; + + let page_meta = Arc::new(SparseStructuralCacheableState { + chunk_meta, + chunk_value_offsets: chunk_value_offsets.into(), + plan, + }); + self.page_meta = Some(page_meta.clone()); + Ok(page_meta as Arc) + } + .boxed() + } + + fn load(&mut self, data: &Arc) { + self.page_meta = Some( + data.clone() + .as_arc_any() + .downcast::() + .unwrap(), + ); + } + + fn schedule_ranges( + &self, + ranges: &[Range], + io: &Arc, + ) -> Result> { + let num_rows = ranges.iter().map(|range| range.end - range.start).sum(); + let page_meta = self.page_meta.as_ref().unwrap(); + + let mut chunks_needed = Vec::new(); + let selection = slice_sparse_plan(&page_meta.plan, ranges)?; + for value_range in &selection.leaf_ranges { + chunks_needed.extend(Self::value_chunk_range( + &page_meta.chunk_value_offsets, + value_range.clone(), + )); + } + chunks_needed.sort_unstable(); + chunks_needed.dedup(); + + let mut loaded_chunks = self.lookup_value_chunks(&chunks_needed); + let chunk_ranges = loaded_chunks + .iter() + .map(|chunk| chunk.byte_range.clone()) + .collect::>(); + let loaded_chunk_data = io.submit_request(chunk_ranges, self.priority); + let ranges = VecDeque::from(ranges.to_vec()); + let value_decompressor = self.value_decompressor.clone(); + let data_type = self.data_type.clone(); + let page_meta = page_meta.clone(); + let num_buffers = self.num_buffers; + let has_large_chunk = self.has_large_chunk; + + let res = async move { + let loaded_chunk_data = loaded_chunk_data.await?; + for (loaded_chunk, chunk_data) in loaded_chunks.iter_mut().zip(loaded_chunk_data) { + loaded_chunk.data = LanceBuffer::from_bytes(chunk_data, 1); + } + + Ok(Box::new(SparseStructuralDecoder { + value_decompressor, + data_type, + page_meta, + loaded_chunks: Arc::new(loaded_chunks), + ranges, + offset_in_current_range: 0, + num_rows, + num_buffers, + has_large_chunk, + }) as Box) + } + .boxed(); + Ok(vec![PageLoadTask { + decoder_fut: res, + num_rows, + }]) + } +} + +#[derive(Debug)] +struct SparseStructuralDecoder { + value_decompressor: Arc, + data_type: DataType, + page_meta: Arc, + loaded_chunks: Arc>, + ranges: VecDeque>, + offset_in_current_range: u64, + num_rows: u64, + num_buffers: u64, + has_large_chunk: bool, +} + +impl SparseStructuralDecoder { + fn drain_ranges(&mut self, mut rows_desired: u64) -> Vec> { + let mut ranges = Vec::new(); + while rows_desired > 0 { + let range = self.ranges.front().unwrap(); + let start = range.start + self.offset_in_current_range; + let rows_available = range.end - start; + let rows_to_take = rows_available.min(rows_desired); + ranges.push(start..start + rows_to_take); + rows_desired -= rows_to_take; + self.offset_in_current_range += rows_to_take; + if self.offset_in_current_range == range.end - range.start { + self.offset_in_current_range = 0; + self.ranges.pop_front(); + } + } + ranges + } +} + +impl StructuralPageDecoder for SparseStructuralDecoder { + fn drain(&mut self, num_rows: u64) -> Result> { + Ok(Box::new(DecodeSparseStructuralTask { + row_ranges: self.drain_ranges(num_rows), + value_decompressor: self.value_decompressor.clone(), + data_type: self.data_type.clone(), + page_meta: self.page_meta.clone(), + loaded_chunks: self.loaded_chunks.clone(), + num_buffers: self.num_buffers, + has_large_chunk: self.has_large_chunk, + })) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +#[derive(Debug)] +struct DecodeSparseStructuralTask { + row_ranges: Vec>, + value_decompressor: Arc, + data_type: DataType, + page_meta: Arc, + loaded_chunks: Arc>, + num_buffers: u64, + has_large_chunk: bool, +} + +impl DecodeSparseStructuralTask { + fn loaded_chunk(&self, chunk_idx: usize) -> Result<&LoadedChunk> { + self.loaded_chunks + .binary_search_by_key(&chunk_idx, |chunk| chunk.chunk_idx) + .map(|idx| &self.loaded_chunks[idx]) + .map_err(|_| { + Error::internal(format!( + "Sparse structural decode missing loaded value chunk {}", + chunk_idx + )) + }) + } + + fn decode_value_chunk(&self, chunk: &LoadedChunk) -> Result { + let buf = &chunk.data; + let mut offset = 0; + let num_levels = u16::from_le_bytes([buf[offset], buf[offset + 1]]); + offset += 2; + if num_levels != 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk unexpectedly contains {} rep/def levels", + num_levels + ) + .into(), + )); + } + + let buffer_sizes = if self.has_large_chunk { + DecodeMiniBlockTask::read_buffer_sizes::(buf, &mut offset, self.num_buffers) + } else { + DecodeMiniBlockTask::read_buffer_sizes::(buf, &mut offset, self.num_buffers) + }; + + offset += pad_bytes::(offset); + let buffers = buffer_sizes + .into_iter() + .map(|buf_size| { + let buffer = buf.slice_with_length(offset, buf_size as usize); + offset += buf_size as usize; + offset += pad_bytes::(offset); + buffer + }) + .collect::>(); + + self.value_decompressor + .decompress(buffers, chunk.items_in_chunk) + } + + fn append_value_range( + &self, + value_range: Range, + data_builder: &mut DataBlockBuilder, + chunk_cache: &mut Option<(usize, DataBlock)>, + ) -> Result<()> { + let mut value_start = value_range.start; + while value_start < value_range.end { + let chunk_idx = self + .page_meta + .chunk_value_offsets + .partition_point(|&offset| offset <= value_start) + - 1; + let chunk_value_start = self.page_meta.chunk_value_offsets[chunk_idx]; + let chunk_value_end = self.page_meta.chunk_value_offsets[chunk_idx + 1]; + let take_end = value_range.end.min(chunk_value_end); + + if !matches!(chunk_cache, Some((cached_idx, _)) if *cached_idx == chunk_idx) { + let chunk = self.loaded_chunk(chunk_idx)?; + *chunk_cache = Some((chunk_idx, self.decode_value_chunk(chunk)?)); + } + let values = &chunk_cache.as_ref().unwrap().1; + data_builder.append( + values, + value_start - chunk_value_start..take_end - chunk_value_start, + ); + value_start = take_end; + } + Ok(()) + } +} + +impl DecodePageTask for DecodeSparseStructuralTask { + fn decode(self: Box) -> Result { + let selection = slice_sparse_plan(&self.page_meta.plan, &self.row_ranges)?; + let estimated_size_bytes = self + .loaded_chunks + .iter() + .map(|chunk| chunk.data.len()) + .sum::() + * 2; + let mut data_builder = + DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64); + let mut chunk_cache: Option<(usize, DataBlock)> = None; + let mut appended_values = false; + for value_range in &selection.leaf_ranges { + self.append_value_range(value_range.clone(), &mut data_builder, &mut chunk_cache)?; + appended_values = true; + } + + let data = if appended_values { + data_builder.finish() + } else { + DataBlock::from_array(new_empty_array(&self.data_type)) + }; + let unraveler = RepDefUnraveler::new_sparse(selection.plan); + Ok(DecodedPage { + data, + repdef: unraveler, + }) + } +} + +struct SparseStructuralSelection { + plan: SparseStructuralPlan, + leaf_ranges: Vec>, +} + +fn translate_positions(positions: &[u64], ranges: &[Range], label: &str) -> Result> { + let mut translated = Vec::new(); + let mut output_base = 0_u64; + for range in ranges { + let mut idx = positions.partition_point(|position| *position < range.start); + while idx < positions.len() && positions[idx] < range.end { + translated.push(output_base + positions[idx] - range.start); + idx += 1; + } + output_base += range.end - range.start; + } + if translated.windows(2).any(|window| window[0] >= window[1]) { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} positions are not sorted after slicing").into(), + )); + } + Ok(translated) +} + +fn offsets_from_counts(counts: &[u64]) -> Vec { + let mut offsets = Vec::with_capacity(counts.len() + 1); + let mut offset = 0_u64; + offsets.push(offset); + for count in counts { + offset += count; + offsets.push(offset); + } + offsets +} + +fn coalesce_ranges(ranges: Vec>) -> Vec> { + let mut coalesced: Vec> = Vec::with_capacity(ranges.len()); + for range in ranges { + if range.is_empty() { + continue; + } + if let Some(last) = coalesced.last_mut() + && last.end == range.start + { + last.end = range.end; + continue; + } + coalesced.push(range); + } + coalesced +} + +fn slice_list_layer( + num_slots: u64, + num_child_slots: u64, + non_empty_positions: &[u64], + counts: &[u64], + constant_count: Option, + null_positions: &[u64], + ranges: &[Range], +) -> Result<(SparseStructuralLayerPlan, Vec>)> { + if non_empty_positions.len() != counts.len() { + return Err(Error::invalid_input_source( + "Sparse structural list has mismatched non-empty positions and counts".into(), + )); + } + if ranges + .iter() + .any(|range| range.start > range.end || range.end > num_slots) + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list slice is outside layer with {} slots", + num_slots + ) + .into(), + )); + } + + let mut out_non_empty_positions = Vec::new(); + let mut out_counts = Vec::new(); + let mut matched_indices = Vec::new(); + let mut output_base = 0_u64; + for range in ranges { + let mut idx = non_empty_positions.partition_point(|position| *position < range.start); + while idx < non_empty_positions.len() && non_empty_positions[idx] < range.end { + out_non_empty_positions.push(output_base + non_empty_positions[idx] - range.start); + out_counts.push(counts[idx]); + matched_indices.push(idx); + idx += 1; + } + output_base += range.end - range.start; + } + + let child_ranges = if matched_indices.is_empty() { + Vec::new() + } else if let Some(constant_count) = constant_count { + let expected_child_slots = + constant_count + .checked_mul(counts.len() as u64) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list constant counts overflow child slots".into(), + ) + })?; + if expected_child_slots != num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list constant count sum {} does not match child slots {}", + expected_child_slots, num_child_slots + ) + .into(), + )); + } + matched_indices + .iter() + .map(|idx| { + let start = *idx as u64 * constant_count; + start..start + constant_count + }) + .collect() + } else { + let value_offsets = offsets_from_counts(counts); + if *value_offsets.last().unwrap() != num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list count sum {} does not match child slots {}", + value_offsets.last().unwrap(), + num_child_slots + ) + .into(), + )); + } + matched_indices + .iter() + .map(|idx| value_offsets[*idx]..value_offsets[*idx + 1]) + .collect() + }; + + let out_null_positions = translate_positions(null_positions, ranges, "list null")?; + let out_num_slots = ranges + .iter() + .map(|range| range.end - range.start) + .sum::(); + let out_num_child_slots = out_counts.iter().sum::(); + let constant_count = out_counts + .first() + .copied() + .filter(|first| out_counts.iter().all(|count| count == first)); + Ok(( + SparseStructuralLayerPlan::List { + num_slots: out_num_slots, + num_child_slots: out_num_child_slots, + non_empty_positions: out_non_empty_positions, + counts: out_counts, + constant_count, + null_positions: out_null_positions, + }, + child_ranges, + )) +} + +fn slice_sparse_plan( + plan: &SparseStructuralPlan, + row_ranges: &[Range], +) -> Result { + let mut selected_ranges = row_ranges.to_vec(); + let mut sliced_layers = Vec::with_capacity(plan.layers.len()); + + for layer in &plan.layers { + match layer { + SparseStructuralLayerPlan::Validity { + num_slots: _, + null_positions, + } => { + let out_num_slots = selected_ranges + .iter() + .map(|range| range.end - range.start) + .sum::(); + sliced_layers.push(SparseStructuralLayerPlan::Validity { + num_slots: out_num_slots, + null_positions: translate_positions( + null_positions, + &selected_ranges, + "validity null", + )?, + }); + } + SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + constant_count, + null_positions, + .. + } => { + let (sliced_layer, child_ranges) = slice_list_layer( + *num_slots, + *num_child_slots, + non_empty_positions, + counts, + *constant_count, + null_positions, + &selected_ranges, + )?; + sliced_layers.push(sliced_layer); + selected_ranges = child_ranges; + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots: _, + dimension, + null_positions, + } => { + let out_num_slots = selected_ranges + .iter() + .map(|range| range.end - range.start) + .sum::(); + let child_ranges = selected_ranges + .iter() + .map(|range| range.start * dimension..range.end * dimension) + .collect::>(); + sliced_layers.push(SparseStructuralLayerPlan::FixedSizeList { + num_slots: out_num_slots, + dimension: *dimension, + null_positions: translate_positions( + null_positions, + &selected_ranges, + "fixed-size-list null", + )?, + }); + selected_ranges = child_ranges; + } + } + } + + let num_visible_items = selected_ranges + .iter() + .map(|range| range.end - range.start) + .sum::(); + Ok(SparseStructuralSelection { + plan: SparseStructuralPlan { + layers: sliced_layers, + num_items: num_visible_items, + num_visible_items, + }, + leaf_ranges: coalesce_ranges(selected_ranges), + }) +} + /// Decodes mini-block formatted data. See [`PrimitiveStructuralEncoder`] for more /// details on the different layouts. #[derive(Debug)] @@ -1206,6 +2160,11 @@ struct MiniBlockSchedulerDictionary { num_dictionary_items: u64, } +#[derive(Debug, Clone, Copy)] +enum ChunkMetadataKind { + MiniBlock { has_large_chunk: bool }, +} + /// Individual block metadata within a MiniBlock repetition index. #[derive(Debug)] struct MiniBlockRepIndexBlock { @@ -1376,7 +2335,7 @@ pub struct MiniBlockScheduler { dictionary: Option, // This is set after initialization page_meta: Option>, - has_large_chunk: bool, + chunk_metadata_kind: ChunkMetadataKind, } impl MiniBlockScheduler { @@ -1460,7 +2419,9 @@ impl MiniBlockScheduler { dictionary, def_meaning: def_meaning.into(), page_meta: None, - has_large_chunk: layout.has_large_chunk, + chunk_metadata_kind: ChunkMetadataKind::MiniBlock { + has_large_chunk: layout.has_large_chunk, + }, }) } @@ -1875,35 +2836,40 @@ impl StructuralPageScheduler for MiniBlockScheduler { let rep_index_bytes = buffers.next(); // Parse the metadata and build the chunk meta - let words = Words::from_bytes(meta_bytes, self.has_large_chunk)?; - let mut chunk_meta = Vec::with_capacity(words.len()); - let mut rows_counter = 0; let mut offset_bytes = value_buf_position; - for (word_idx, word) in words.iter().enumerate() { - let log_num_values = word & 0x0F; - let divided_bytes = word >> 4; - let num_bytes = (divided_bytes as usize + 1) * MINIBLOCK_ALIGNMENT; - debug_assert!(num_bytes > 0); - let num_values = if word_idx < words.len() - 1 { - debug_assert!(log_num_values > 0); - 1 << log_num_values - } else { - debug_assert!( - log_num_values == 0 - || (1 << log_num_values) == (self.items_in_page - rows_counter) - ); - self.items_in_page - rows_counter - }; - rows_counter += num_values; + let chunk_meta = match self.chunk_metadata_kind { + ChunkMetadataKind::MiniBlock { has_large_chunk } => { + let words = Words::from_bytes(meta_bytes, has_large_chunk)?; + let mut chunk_meta = Vec::with_capacity(words.len()); + + for (word_idx, word) in words.iter().enumerate() { + let log_num_values = word & 0x0F; + let divided_bytes = word >> 4; + let num_bytes = (divided_bytes as usize + 1) * MINIBLOCK_ALIGNMENT; + debug_assert!(num_bytes > 0); + let num_values = if word_idx < words.len() - 1 { + debug_assert!(log_num_values > 0); + 1 << log_num_values + } else { + debug_assert!( + log_num_values == 0 + || (1 << log_num_values) == (self.items_in_page - rows_counter) + ); + self.items_in_page - rows_counter + }; + rows_counter += num_values; - chunk_meta.push(ChunkMeta { - num_values, - chunk_size_bytes: num_bytes as u64, - offset_bytes, - }); - offset_bytes += num_bytes as u64; - } + chunk_meta.push(ChunkMeta { + num_values, + chunk_size_bytes: num_bytes as u64, + offset_bytes, + }); + offset_bytes += num_bytes as u64; + } + chunk_meta + } + }; // Build the repetition index let rep_index = if let Some(rep_index_data) = rep_index_bytes { @@ -1985,7 +2951,9 @@ impl StructuralPageScheduler for MiniBlockScheduler { let def_decompressor = self.def_decompressor.clone(); let value_decompressor = self.value_decompressor.clone(); let num_buffers = self.num_buffers; - let has_large_chunk = self.has_large_chunk; + let has_large_chunk = match self.chunk_metadata_kind { + ChunkMetadataKind::MiniBlock { has_large_chunk } => has_large_chunk, + }; let dictionary = page_meta .dictionary .as_ref() @@ -3349,6 +4317,13 @@ impl StructuralPrimitiveFieldScheduler { mini_block, decompressors, )?), + Layout::SparseLayout(sparse) => Box::new(SparseStructuralScheduler::try_new( + &page_info.buffer_offsets_and_sizes, + page_info.priority, + target_field.data_type(), + sparse, + decompressors, + )?), Layout::FullZipLayout(full_zip) => { let mut scheduler = FullZipScheduler::try_new( &page_info.buffer_offsets_and_sizes, @@ -3775,6 +4750,22 @@ struct CompressedLevels { rep_index: Option, } +struct SparseMiniBlockChunk { + buffer_sizes: Vec, + num_values: u32, +} + +struct SparseMiniBlockCompressed { + data: Vec, + chunks: Vec, +} + +#[derive(Debug)] +struct SparseStructuralEncoded { + layers: Vec, + buffers: Vec, +} + struct SerializedMiniBlockPage { num_buffers: u64, data: LanceBuffer, @@ -3799,6 +4790,10 @@ struct PrimitivePageData { num_rows: u64, // Present when one top-level row is too large for one miniblock rep/def chunk. unsplittable_miniblock_levels: Option, + // Structural split plan from the original accumulated page. + structural_plan: StructuralPagePlan, + // Native sparse structural mapping for the original accumulated page. + sparse_structural_plan: Option, } // Immutable encoder state shared by per-page encode tasks. @@ -3826,6 +4821,10 @@ struct PrimitiveEncodeContext { } impl PrimitiveStructuralEncoder { + fn supports_sparse_layout(version: LanceFileVersion) -> bool { + version.resolve() >= LanceFileVersion::V2_3 + } + pub fn try_new( options: &EncodingOptions, compression_strategy: Arc, @@ -4233,6 +5232,149 @@ impl PrimitiveStructuralEncoder { }) } + fn serialize_sparse_miniblocks( + miniblocks: SparseMiniBlockCompressed, + rep: Option>, + def: Option>, + support_large_chunk: bool, + ) -> Result { + let bytes_rep = rep + .as_ref() + .map(|rep| rep.iter().map(|r| r.data.len()).sum::()) + .unwrap_or(0); + let bytes_def = def + .as_ref() + .map(|def| def.iter().map(|d| d.data.len()).sum::()) + .unwrap_or(0); + let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::(); + let mut num_buffers = miniblocks.data.len(); + if rep.is_some() { + num_buffers += 1; + } + if def.is_some() { + num_buffers += 1; + } + let max_extra = 9 * num_buffers; + let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra); + let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * 8); + + let mut rep_iter = rep.map(|r| r.into_iter()); + let mut def_iter = def.map(|d| d.into_iter()); + + let mut buffer_offsets = vec![0; miniblocks.data.len()]; + for chunk in miniblocks.chunks { + if chunk.buffer_sizes.len() != miniblocks.data.len() { + return Err(Error::internal(format!( + "Sparse chunk has {} value buffer sizes, expected {}", + chunk.buffer_sizes.len(), + miniblocks.data.len() + ))); + } + + let start_pos = data_buffer.len(); + debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0); + + let rep = rep_iter.as_mut().map(|r| r.next().unwrap()); + let def = def_iter.as_mut().map(|d| d.next().unwrap()); + + let num_levels = rep + .as_ref() + .map(|r| r.num_levels) + .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0)); + data_buffer.extend_from_slice(&num_levels.to_le_bytes()); + + if let Some(rep) = rep.as_ref() { + let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| { + Error::internal(format!( + "Sparse repetition buffer size ({} bytes) too large", + rep.data.len() + )) + })?; + data_buffer.extend_from_slice(&bytes_rep.to_le_bytes()); + } + if let Some(def) = def.as_ref() { + let bytes_def = u16::try_from(def.data.len()).map_err(|_| { + Error::internal(format!( + "Sparse definition buffer size ({} bytes) too large", + def.data.len() + )) + })?; + data_buffer.extend_from_slice(&bytes_def.to_le_bytes()); + } + + if support_large_chunk { + for &buffer_size in &chunk.buffer_sizes { + data_buffer.extend_from_slice(&buffer_size.to_le_bytes()); + } + } else { + for &buffer_size in &chunk.buffer_sizes { + let buffer_size = u16::try_from(buffer_size).map_err(|_| { + Error::internal(format!( + "Sparse value buffer size ({} bytes) too large for 16-bit metadata", + buffer_size + )) + })?; + data_buffer.extend_from_slice(&buffer_size.to_le_bytes()); + } + } + + let add_padding = |data_buffer: &mut Vec| { + let pad = pad_bytes::(data_buffer.len()); + data_buffer.extend(iter::repeat_n(FILL_BYTE, pad)); + }; + add_padding(&mut data_buffer); + + if let Some(rep) = rep.as_ref() { + data_buffer.extend_from_slice(&rep.data); + add_padding(&mut data_buffer); + } + if let Some(def) = def.as_ref() { + data_buffer.extend_from_slice(&def.data); + add_padding(&mut data_buffer); + } + for (buffer_size, (buffer, buffer_offset)) in chunk + .buffer_sizes + .iter() + .zip(miniblocks.data.iter().zip(buffer_offsets.iter_mut())) + { + let start = *buffer_offset; + let end = start + *buffer_size as usize; + *buffer_offset += *buffer_size as usize; + data_buffer.extend_from_slice(&buffer[start..end]); + add_padding(&mut data_buffer); + } + + let chunk_bytes = data_buffer.len() - start_pos; + if chunk_bytes == 0 || chunk_bytes > (u32::MAX as usize + 1) * MINIBLOCK_ALIGNMENT { + return Err(Error::internal(format!( + "Sparse chunk size {} bytes exceeds the metadata limit", + chunk_bytes + ))); + } + if chunk_bytes % MINIBLOCK_ALIGNMENT != 0 { + return Err(Error::internal(format!( + "Sparse chunk size {} bytes is not aligned to {} bytes", + chunk_bytes, MINIBLOCK_ALIGNMENT + ))); + } + let divided_bytes_minus_one = u32::try_from(chunk_bytes / MINIBLOCK_ALIGNMENT - 1) + .map_err(|_| { + Error::internal(format!( + "Sparse chunk size {} bytes exceeds the metadata limit", + chunk_bytes + )) + })?; + meta_buffer.extend_from_slice(÷d_bytes_minus_one.to_le_bytes()); + meta_buffer.extend_from_slice(&chunk.num_values.to_le_bytes()); + } + + Ok(SerializedMiniBlockPage { + num_buffers: miniblocks.data.len() as u64, + data: LanceBuffer::from(data_buffer), + metadata: LanceBuffer::from(meta_buffer), + }) + } + fn encode_simple_all_null( column_idx: u32, num_rows: u64, @@ -4565,6 +5707,274 @@ impl PrimitiveStructuralEncoder { Ok(dict_values_field) } + fn sparse_miniblock_compressed( + compressed: MiniBlockCompressed, + ) -> Result { + let mut values_counter = 0_u64; + let mut chunks = Vec::with_capacity(compressed.chunks.len()); + for chunk in compressed.chunks { + let num_values = chunk.num_values(values_counter, compressed.num_values); + values_counter += num_values; + let num_values = u32::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse value chunk has {} visible values, which exceeds the u32 metadata limit", + num_values + ) + .into(), + ) + })?; + chunks.push(SparseMiniBlockChunk { + buffer_sizes: chunk.buffer_sizes, + num_values, + }); + } + if values_counter != compressed.num_values { + return Err(Error::internal(format!( + "Sparse value chunks describe {} values, expected {}", + values_counter, compressed.num_values + ))); + } + Ok(SparseMiniBlockCompressed { + data: compressed.data, + chunks, + }) + } + + fn encode_sparse_u64_values( + values: Vec, + compression_strategy: &dyn CompressionStrategy, + ) -> Result<(LanceBuffer, CompressiveEncoding)> { + let num_values = values.len() as u64; + let mut block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(values), + bits_per_value: 64, + num_values, + block_info: BlockInfo::new(), + }); + block.compute_stat(); + let field = Field::new_arrow("", DataType::UInt64, false)?; + let (compressor, encoding) = + compression_strategy.create_block_compressor(&field, &block)?; + Ok((compressor.compress(block)?, encoding)) + } + + fn positions_to_deltas(positions: &[u64], label: &str) -> Result> { + let mut previous_row = 0_u64; + positions + .iter() + .copied() + .enumerate() + .map(|(idx, row_id)| { + if idx > 0 && row_id <= previous_row { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} positions must be strictly increasing") + .into(), + )); + } + let delta = if idx == 0 { + row_id + } else { + row_id - previous_row + }; + previous_row = row_id; + Ok(delta) + }) + .collect() + } + + fn encode_sparse_positions( + positions: &[u64], + compression_strategy: &dyn CompressionStrategy, + label: &str, + ) -> Result> { + if positions.is_empty() { + return Ok(None); + } + let deltas = Self::positions_to_deltas(positions, label)?; + Self::encode_sparse_u64_values(deltas, compression_strategy).map(Some) + } + + fn encode_sparse_structural_plan( + plan: &SparseStructuralPlan, + compression_strategy: &dyn CompressionStrategy, + ) -> Result { + let mut layers = Vec::with_capacity(plan.layers.len()); + let mut buffers = Vec::new(); + + for layer in &plan.layers { + match layer { + SparseStructuralLayerPlan::Validity { + num_slots, + null_positions, + } => { + let null_compression = Self::encode_sparse_positions( + null_positions, + compression_strategy, + "validity null", + )? + .map(|(buffer, encoding)| { + buffers.push(buffer); + encoding + }); + layers.push(pb21::SparseStructuralLayer { + kind: pb21::sparse_structural_layer::Kind::SparseLayerValidity as i32, + num_slots: *num_slots, + num_child_slots: *num_slots, + constant_count: 0, + non_empty_compression: None, + num_non_empty: 0, + count_compression: None, + null_compression, + num_nulls: null_positions.len() as u64, + }); + } + SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + constant_count, + null_positions, + } => { + if non_empty_positions.len() != counts.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list has {} non-empty positions but {} counts", + non_empty_positions.len(), + counts.len() + ) + .into(), + )); + } + let non_empty_compression = Self::encode_sparse_positions( + non_empty_positions, + compression_strategy, + "list non-empty", + )? + .map(|(buffer, encoding)| { + buffers.push(buffer); + encoding + }); + let count_compression = if constant_count.is_some() || counts.is_empty() { + None + } else { + let (buffer, encoding) = + Self::encode_sparse_u64_values(counts.clone(), compression_strategy)?; + buffers.push(buffer); + Some(encoding) + }; + let null_compression = Self::encode_sparse_positions( + null_positions, + compression_strategy, + "list null", + )? + .map(|(buffer, encoding)| { + buffers.push(buffer); + encoding + }); + layers.push(pb21::SparseStructuralLayer { + kind: pb21::sparse_structural_layer::Kind::SparseLayerList as i32, + num_slots: *num_slots, + num_child_slots: *num_child_slots, + constant_count: constant_count.unwrap_or(0), + non_empty_compression, + num_non_empty: non_empty_positions.len() as u64, + count_compression, + null_compression, + num_nulls: null_positions.len() as u64, + }); + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots, + dimension, + null_positions, + } => { + let null_compression = Self::encode_sparse_positions( + null_positions, + compression_strategy, + "fixed-size-list null", + )? + .map(|(buffer, encoding)| { + buffers.push(buffer); + encoding + }); + let num_child_slots = num_slots.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", + num_slots, dimension + ) + .into(), + ) + })?; + layers.push(pb21::SparseStructuralLayer { + kind: pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList as i32, + num_slots: *num_slots, + num_child_slots, + constant_count: *dimension, + non_empty_compression: None, + num_non_empty: 0, + count_compression: None, + null_compression, + num_nulls: null_positions.len() as u64, + }); + } + } + } + + Ok(SparseStructuralEncoded { layers, buffers }) + } + + #[allow(clippy::too_many_arguments)] + fn encode_sparse_structural_page( + column_idx: u32, + field: &Field, + compression_strategy: &dyn CompressionStrategy, + data: DataBlock, + sparse_plan: SparseStructuralPlan, + row_number: u64, + num_rows: u64, + support_large_chunk: bool, + ) -> Result { + if sparse_plan.num_visible_items != data.num_values() { + return Err(Error::internal(format!( + "Sparse structural plan has {} visible items but data has {} values", + sparse_plan.num_visible_items, + data.num_values() + ))); + } + + let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; + let (compressed_data, value_encoding) = compressor.compress(data)?; + let compressed_data = Self::sparse_miniblock_compressed(compressed_data)?; + let serialized = + Self::serialize_sparse_miniblocks(compressed_data, None, None, support_large_chunk)?; + let structural = Self::encode_sparse_structural_plan(&sparse_plan, compression_strategy)?; + + let description = ProtobufUtils21::sparse_layout( + value_encoding, + serialized.num_buffers, + sparse_plan.num_items, + sparse_plan.num_visible_items, + support_large_chunk, + structural.layers, + ); + + let mut data = Vec::with_capacity(2 + structural.buffers.len()); + data.push(serialized.metadata); + data.push(serialized.data); + data.extend(structural.buffers); + + Ok(EncodedPage { + num_rows, + column_idx, + data, + description: PageEncoding::Structural(description), + row_number, + }) + } + #[allow(clippy::too_many_arguments)] fn encode_miniblock( column_idx: u32, @@ -5306,6 +6716,8 @@ impl PrimitiveStructuralEncoder { row_number, num_rows, unsplittable_miniblock_levels: None, + structural_plan: StructuralPagePlan::Fits, + sparse_structural_plan: None, }]); } if let StructuralPagePlan::UnsplittableOverBudget(num_levels) = plan { @@ -5315,6 +6727,8 @@ impl PrimitiveStructuralEncoder { row_number, num_rows, unsplittable_miniblock_levels: Some(num_levels), + structural_plan: StructuralPagePlan::UnsplittableOverBudget(num_levels), + sparse_structural_plan: None, }]); } @@ -5332,6 +6746,8 @@ impl PrimitiveStructuralEncoder { row_number: row_number + split.row_start, num_rows: split.num_rows, unsplittable_miniblock_levels: None, + structural_plan: StructuralPagePlan::Fits, + sparse_structural_plan: None, }); } Ok(pages) @@ -5354,6 +6770,8 @@ impl PrimitiveStructuralEncoder { row_number, num_rows, unsplittable_miniblock_levels, + structural_plan, + sparse_structural_plan, } = page; let num_values = arrays.iter().map(|arr| arr.len() as u64).sum(); @@ -5421,8 +6839,30 @@ impl PrimitiveStructuralEncoder { } let data_block = DataBlock::from_arrays(&arrays, num_values); + let requested_encoding = encoding_metadata + .get(STRUCTURAL_ENCODING_META_KEY) + .map(|requested| requested.to_lowercase()); + let supports_sparse_layout = Self::supports_sparse_layout(version); + let should_encode_sparse = requested_encoding.as_deref() + == Some(STRUCTURAL_ENCODING_SPARSE) + || (supports_sparse_layout + && requested_encoding.is_none() + && matches!(&structural_plan, StructuralPagePlan::Split(_))); + + if requested_encoding.as_deref() == Some(STRUCTURAL_ENCODING_SPARSE) + && !supports_sparse_layout + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural encoding requires Lance file format 2.3+, current version: {}", + version.resolve() + ) + .into(), + )); + } if version.resolve() >= LanceFileVersion::V2_2 + && !should_encode_sparse && let Some(scalar) = Self::find_constant_scalar(&arrays, leaf_validity.as_ref())? { log::debug!( @@ -5436,10 +6876,73 @@ impl PrimitiveStructuralEncoder { ); } + let requires_full_zip_packed_struct = + if let DataBlock::Struct(ref struct_data_block) = data_block { + struct_data_block.has_variable_width_child() + } else { + false + }; + + if let DataBlock::Dictionary(dict) = data_block { + if should_encode_sparse { + return Err(Error::not_supported_source( + "Sparse layout does not support dictionary data blocks".into(), + )); + } + log::debug!( + "Encoding column {} with {} items using dictionary encoding (already dictionary encoded)", + column_idx, + num_values + ); + let (mut indices_data_block, dictionary_data_block) = dict.into_parts(); + // TODO: https://github.com/lancedb/lance/issues/4809 + // If we compute stats on dictionary_data_block => panic. + // If we don't compute stats on indices_data_block => panic. + // This is messy. Don't make me call compute_stat ever. + indices_data_block.compute_stat(); + return Self::encode_miniblock( + column_idx, + &field, + compression_strategy.as_ref(), + indices_data_block, + repdef, + row_number, + Some(dictionary_data_block), + num_rows, + support_large_chunk, + ); + } + + if should_encode_sparse { + if requires_full_zip_packed_struct { + return Err(Error::not_supported_source( + "Sparse layout does not support packed struct value blocks".into(), + )); + } + let sparse_structural_plan = sparse_structural_plan.ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural encoding requires native structural layers".into(), + ) + })?; + log::debug!( + "Encoding column {} with {} visible items ({} rows) using sparse layout", + column_idx, + num_values, + num_rows + ); + return Self::encode_sparse_structural_page( + column_idx, + &field, + compression_strategy.as_ref(), + data_block, + sparse_structural_plan, + row_number, + num_rows, + support_large_chunk, + ); + } + if let Some(num_levels) = unsplittable_miniblock_levels { - let requested_encoding = encoding_metadata - .get(STRUCTURAL_ENCODING_META_KEY) - .map(|requested| requested.to_lowercase()); let fullzip_error = match &data_block { DataBlock::FixedWidth(fixed) if !fixed.bits_per_value.is_multiple_of(8) => { Some(format!( @@ -5539,13 +7042,6 @@ impl PrimitiveStructuralEncoder { } } - let requires_full_zip_packed_struct = - if let DataBlock::Struct(ref struct_data_block) = data_block { - struct_data_block.has_variable_width_child() - } else { - false - }; - if requires_full_zip_packed_struct { log::debug!( "Encoding column {} with {} items using full-zip packed struct layout", @@ -5563,31 +7059,6 @@ impl PrimitiveStructuralEncoder { ); } - if let DataBlock::Dictionary(dict) = data_block { - log::debug!( - "Encoding column {} with {} items using dictionary encoding (already dictionary encoded)", - column_idx, - num_values - ); - let (mut indices_data_block, dictionary_data_block) = dict.into_parts(); - // TODO: https://github.com/lancedb/lance/issues/4809 - // If we compute stats on dictionary_data_block => panic. - // If we don't compute stats on indices_data_block => panic. - // This is messy. Don't make me call compute_stat ever. - indices_data_block.compute_stat(); - return Self::encode_miniblock( - column_idx, - &field, - compression_strategy.as_ref(), - indices_data_block, - repdef, - row_number, - Some(dictionary_data_block), - num_rows, - support_large_chunk, - ); - } - // Try dictionary encoding first if applicable. If encoding aborts, fall back to the // preferred structural encoding. let dict_result = Self::should_dictionary_encode(&data_block, &field, version).and_then(|budget| { @@ -5663,19 +7134,62 @@ impl PrimitiveStructuralEncoder { let num_values = arrays.iter().map(|arr| arr.len() as u64).sum(); let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity()); let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty()); + let sparse_structural_plan = if Self::supports_sparse_layout(self.version) { + RepDefBuilder::to_sparse_structural_plan(&repdefs, num_values)? + } else { + None + }; let (repdef, structural_plan) = RepDefBuilder::serialize_with_structural_plan( repdefs, miniblock::max_repdef_levels_per_chunk, num_rows, num_values, )?; - let pages = Self::split_structural_pages_for_miniblock_budget( - arrays, - repdef, - structural_plan, - row_number, - num_rows, - )?; + let requested_encoding = self + .encoding_metadata + .get(STRUCTURAL_ENCODING_META_KEY) + .map(|requested| requested.to_lowercase()); + + if requested_encoding.as_deref() == Some(STRUCTURAL_ENCODING_SPARSE) + && !Self::supports_sparse_layout(self.version) + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural encoding requires Lance file format 2.3+, current version: {}", + self.version.resolve() + ) + .into(), + )); + } + + let should_use_sparse_plan = requested_encoding.as_deref() + == Some(STRUCTURAL_ENCODING_SPARSE) + || (Self::supports_sparse_layout(self.version) + && requested_encoding.is_none() + && matches!(&structural_plan, StructuralPagePlan::Split(_))); + let pages = if should_use_sparse_plan { + let unsplittable_miniblock_levels = match &structural_plan { + StructuralPagePlan::UnsplittableOverBudget(num_levels) => Some(*num_levels), + _ => None, + }; + vec![PrimitivePageData { + arrays, + repdef, + row_number, + num_rows, + unsplittable_miniblock_levels, + structural_plan, + sparse_structural_plan, + }] + } else { + Self::split_structural_pages_for_miniblock_budget( + arrays, + repdef, + structural_plan, + row_number, + num_rows, + )? + }; let mut tasks = Vec::with_capacity(pages.len()); let ctx = PrimitiveEncodeContext { @@ -5801,8 +7315,8 @@ mod tests { use crate::compression::DefaultDecompressionStrategy; use crate::constants::{ COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY, - DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY, - STRUCTURAL_ENCODING_MINIBLOCK, + DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_FULLZIP, + STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_SPARSE, }; use crate::data::BlockInfo; use crate::decoder::{PageEncoding, StructuralFieldDecoder}; @@ -7649,11 +9163,391 @@ mod tests { ); } + fn sparse_i32_list_array(num_rows: i32, visible_stride: i32) -> ArrayRef { + use arrow_array::builder::{Int32Builder, ListBuilder}; + + let mut builder = ListBuilder::new(Int32Builder::new()); + for row in 0..num_rows { + if row % visible_stride == 0 { + builder.values().append_value(row); + } + builder.append(true); + } + Arc::new(builder.finish()) + } + + fn sparse_layout_from_page( + page: &crate::encoder::EncodedPage, + ) -> &crate::format::pb21::SparseLayout { + let PageEncoding::Structural(layout) = &page.description else { + panic!("Expected structural page encoding"); + }; + let pb21::page_layout::Layout::SparseLayout(layout) = layout.layout.as_ref().unwrap() + else { + panic!("Expected sparse layout"); + }; + layout + } + + fn sparse_metadata_visible_values(page: &crate::encoder::EncodedPage) -> Vec { + page.data[0] + .chunks_exact(8) + .map(|chunk| u32::from_le_bytes([chunk[4], chunk[5], chunk[6], chunk[7]])) + .collect() + } + + fn sparse_list_layers( + sparse: &crate::format::pb21::SparseLayout, + ) -> Vec<&crate::format::pb21::SparseStructuralLayer> { + sparse + .structural_layers + .iter() + .filter(|layer| { + matches!( + crate::format::pb21::sparse_structural_layer::Kind::try_from(layer.kind), + Ok(crate::format::pb21::sparse_structural_layer::Kind::SparseLayerList) + ) + }) + .collect() + } + + #[tokio::test] + async fn test_explicit_sparse_layout_roundtrip() { + let arr = sparse_i32_list_array(4096, 64); + let mut metadata = HashMap::new(); + metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_SPARSE.to_string(), + ); + let field = arrow_schema::Field::new("c", arr.data_type().clone(), true) + .with_metadata(metadata.clone()); + + let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_3).await; + let sparse = sparse_layout_from_page(&page); + assert_eq!(sparse.num_visible_items, 64); + assert!(sparse.num_items > sparse.num_visible_items); + let list_layers = sparse_list_layers(sparse); + assert_eq!(list_layers.len(), 1); + assert_eq!(list_layers[0].num_slots, 4096); + assert_eq!(list_layers[0].num_non_empty, 64); + assert_eq!(list_layers[0].constant_count, 1); + assert!(list_layers[0].non_empty_compression.is_some()); + assert!(list_layers[0].count_compression.is_none()); + let chunk_values = sparse_metadata_visible_values(&page); + assert_eq!( + chunk_values.iter().map(|v| *v as u64).sum::(), + sparse.num_visible_items + ); + + let test_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3); + check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; + } + + #[tokio::test] + async fn test_explicit_sparse_layout_requires_v2_3() { + let arr = sparse_i32_list_array(4096, 64); + let mut metadata = HashMap::new(); + metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_SPARSE.to_string(), + ); + let field = + arrow_schema::Field::new("c", arr.data_type().clone(), true).with_metadata(metadata); + + let err = try_encode_first_page(field, arr, LanceFileVersion::V2_2) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("Sparse structural encoding requires Lance file format 2.3+"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn test_sparse_layout_auto_for_overbudget_structural_page() { + let arr = sparse_i32_list_array(70_000, 65_535); + let field = arrow_schema::Field::new("c", arr.data_type().clone(), true); + + let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_3).await; + let sparse = sparse_layout_from_page(&page); + assert_eq!(sparse.num_visible_items, 2); + assert!(sparse.num_items > 65_535); + let list_layers = sparse_list_layers(sparse); + assert_eq!(list_layers.len(), 1); + assert_eq!(list_layers[0].num_non_empty, 2); + assert_eq!(list_layers[0].constant_count, 1); + + let test_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3); + check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; + } + + fn deep_sparse_struct_array(num_rows: usize, stride: usize) -> ArrayRef { + use arrow_array::{Int32Array, ListArray, StructArray}; + use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::Fields; + + fn offsets(values: Vec) -> OffsetBuffer { + OffsetBuffer::new(ScalarBuffer::from(values)) + } + + fn validity(values: Vec) -> Option { + Some(NullBuffer::new(BooleanBuffer::from(values))) + } + + let mut event_rows = Vec::new(); + let mut event_offsets = Vec::with_capacity(num_rows + 1); + let mut event_validity = Vec::with_capacity(num_rows); + let mut event_child_count = 0_i32; + event_offsets.push(event_child_count); + for row in 0..num_rows { + let is_null = row % 4096 == 7; + let is_non_empty = row % stride == 0; + event_validity.push(!is_null); + if !is_null && is_non_empty { + event_rows.push(row as i32); + event_child_count += 1; + } + event_offsets.push(event_child_count); + } + + let score_values = Arc::new(Int32Array::from(event_rows.clone())) as ArrayRef; + let mut sample_offsets = Vec::with_capacity(event_rows.len() + 1); + let mut sample_values = Vec::with_capacity(event_rows.len() * 2); + sample_offsets.push(0_i32); + for row in &event_rows { + sample_values.push(*row); + sample_values.push(*row + 1); + sample_offsets.push(sample_values.len() as i32); + } + let sample_values = Arc::new(Int32Array::from(sample_values)) as ArrayRef; + let sample_field = Arc::new(ArrowField::new("item", DataType::Int32, true)); + let samples = Arc::new( + ListArray::try_new(sample_field, offsets(sample_offsets), sample_values, None).unwrap(), + ) as ArrayRef; + let event_fields = Fields::from(vec![ + ArrowField::new("score", DataType::Int32, true), + ArrowField::new("samples", samples.data_type().clone(), true), + ]); + let events_struct = Arc::new(StructArray::new( + event_fields, + vec![score_values, samples], + None, + )) as ArrayRef; + let event_field = Arc::new(ArrowField::new( + "item", + events_struct.data_type().clone(), + true, + )); + let events = Arc::new( + ListArray::try_new( + event_field, + offsets(event_offsets), + events_struct, + validity(event_validity), + ) + .unwrap(), + ) as ArrayRef; + let profile_fields = Fields::from(vec![ArrowField::new( + "events", + events.data_type().clone(), + true, + )]); + let profile = Arc::new(StructArray::new(profile_fields, vec![events], None)) as ArrayRef; + + let mut payload_rows = Vec::new(); + let mut payload_offsets = Vec::with_capacity(num_rows + 1); + let mut payload_validity = Vec::with_capacity(num_rows); + let mut payload_child_count = 0_i32; + payload_offsets.push(payload_child_count); + for row in 0..num_rows { + let is_null = row % 4096 == 11; + let is_non_empty = row % stride == 0; + payload_validity.push(!is_null); + if !is_null && is_non_empty { + payload_rows.push(row as i32); + payload_child_count += 1; + } + payload_offsets.push(payload_child_count); + } + + let mut payload_inner_offsets = Vec::with_capacity(payload_rows.len() + 1); + let mut payload_values = Vec::with_capacity(payload_rows.len() * 2); + payload_inner_offsets.push(0_i32); + for row in &payload_rows { + payload_values.push(*row); + payload_values.push(*row + 3); + payload_inner_offsets.push(payload_values.len() as i32); + } + let payload_values = Arc::new(Int32Array::from(payload_values)) as ArrayRef; + let payload_inner_field = Arc::new(ArrowField::new("item", DataType::Int32, true)); + let payload_inner = Arc::new( + ListArray::try_new( + payload_inner_field, + offsets(payload_inner_offsets), + payload_values, + None, + ) + .unwrap(), + ) as ArrayRef; + let payload_outer_field = Arc::new(ArrowField::new( + "item", + payload_inner.data_type().clone(), + true, + )); + let payload = Arc::new( + ListArray::try_new( + payload_outer_field, + offsets(payload_offsets), + payload_inner, + validity(payload_validity), + ) + .unwrap(), + ) as ArrayRef; + + Arc::new(StructArray::new( + Fields::from(vec![ + ArrowField::new("profile", profile.data_type().clone(), true), + ArrowField::new("payload", payload.data_type().clone(), true), + ]), + vec![profile, payload], + None, + )) + } + + fn encoded_page_bytes(pages: &[crate::encoder::EncodedPage]) -> usize { + pages + .iter() + .flat_map(|page| page.data.iter()) + .map(|buffer| buffer.len()) + .sum() + } + + fn sparse_page_count(pages: &[crate::encoder::EncodedPage]) -> usize { + pages + .iter() + .filter(|page| { + let PageEncoding::Structural(layout) = &page.description else { + return false; + }; + matches!( + layout.layout.as_ref().unwrap(), + pb21::page_layout::Layout::SparseLayout(_) + ) + }) + .count() + } + + async fn encode_pages_with_structural_mode( + array: ArrayRef, + mode: &str, + ) -> Vec { + let metadata = + HashMap::from([(STRUCTURAL_ENCODING_META_KEY.to_string(), mode.to_string())]); + let field = + arrow_schema::Field::new("c", array.data_type().clone(), true).with_metadata(metadata); + try_encode_pages(field, array, LanceFileVersion::V2_3) + .await + .unwrap() + } + + #[tokio::test] + async fn test_sparse_layout_deep_nested_struct_beats_miniblock_fullzip() { + let arr = deep_sparse_struct_array(80_000, 512); + let sparse_pages = + encode_pages_with_structural_mode(arr.clone(), STRUCTURAL_ENCODING_SPARSE).await; + let miniblock_pages = + encode_pages_with_structural_mode(arr.clone(), STRUCTURAL_ENCODING_MINIBLOCK).await; + let fullzip_pages = + encode_pages_with_structural_mode(arr.clone(), STRUCTURAL_ENCODING_FULLZIP).await; + + let sparse_bytes = encoded_page_bytes(&sparse_pages); + let miniblock_bytes = encoded_page_bytes(&miniblock_pages); + let fullzip_bytes = encoded_page_bytes(&fullzip_pages); + + assert!(sparse_page_count(&sparse_pages) > 0); + assert!( + sparse_pages.iter().any(|page| { + let PageEncoding::Structural(layout) = &page.description else { + return false; + }; + let pb21::page_layout::Layout::SparseLayout(sparse) = + layout.layout.as_ref().unwrap() + else { + return false; + }; + sparse.structural_layers.len() >= 4 + }), + "expected at least one sparse leaf page with deep structural layers" + ); + assert!( + sparse_pages.len() < miniblock_pages.len(), + "expected sparse to emit fewer pages than miniblock: sparse={}, miniblock={}", + sparse_pages.len(), + miniblock_pages.len() + ); + assert!( + sparse_pages.len() <= fullzip_pages.len(), + "expected sparse page count to be no worse than fullzip: sparse={}, fullzip={}", + sparse_pages.len(), + fullzip_pages.len() + ); + assert!( + sparse_bytes < miniblock_bytes, + "expected sparse bytes < miniblock bytes: sparse={}, miniblock={}", + sparse_bytes, + miniblock_bytes + ); + assert!( + sparse_bytes < fullzip_bytes, + "expected sparse bytes < fullzip bytes: sparse={}, fullzip={}", + sparse_bytes, + fullzip_bytes + ); + + let metadata = HashMap::from([( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_SPARSE.to_string(), + )]); + let test_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_range(0..128) + .with_range(511..514) + .with_range(4096..4112) + .with_indices(vec![0, 1, 511, 512, 4096, 79_999]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; + } + async fn encode_first_page( field: arrow_schema::Field, array: ArrayRef, version: LanceFileVersion, ) -> crate::encoder::EncodedPage { + try_encode_first_page(field, array, version).await.unwrap() + } + + async fn try_encode_first_page( + field: arrow_schema::Field, + array: ArrayRef, + version: LanceFileVersion, + ) -> lance_core::Result { + Ok(try_encode_pages(field, array, version) + .await? + .into_iter() + .next() + .unwrap()) + } + + async fn try_encode_pages( + field: arrow_schema::Field, + array: ArrayRef, + version: LanceFileVersion, + ) -> lance_core::Result> { use crate::encoder::{ ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers, default_encoding_strategy, @@ -7671,29 +9565,24 @@ mod tests { version, }; - let mut encoder = encoding_strategy - .create_field_encoder( - encoding_strategy.as_ref(), - &lance_field, - &mut column_index_seq, - &encoding_options, - ) - .unwrap(); + let mut encoder = encoding_strategy.create_field_encoder( + encoding_strategy.as_ref(), + &lance_field, + &mut column_index_seq, + &encoding_options, + )?; let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT); let repdef = RepDefBuilder::default(); let num_rows = array.len() as u64; let mut pages = Vec::new(); - for task in encoder - .maybe_encode(array, &mut external_buffers, repdef, 0, num_rows) - .unwrap() - { - pages.push(task.await.unwrap()); + for task in encoder.maybe_encode(array, &mut external_buffers, repdef, 0, num_rows)? { + pages.push(task.await?); } - for task in encoder.flush(&mut external_buffers).unwrap() { - pages.push(task.await.unwrap()); + for task in encoder.flush(&mut external_buffers)? { + pages.push(task.await?); } - pages.into_iter().next().unwrap() + Ok(pages) } #[tokio::test] diff --git a/rust/lance-encoding/src/encodings/physical/bitpacking.rs b/rust/lance-encoding/src/encodings/physical/bitpacking.rs index 8ebdcc13c56..218066da0b2 100644 --- a/rust/lance-encoding/src/encodings/physical/bitpacking.rs +++ b/rust/lance-encoding/src/encodings/physical/bitpacking.rs @@ -241,6 +241,14 @@ impl MiniBlockDecompressor for InlineBitpacking { fn decompress(&self, data: Vec, num_values: u64) -> Result { assert_eq!(data.len(), 1); let data = data.into_iter().next().unwrap(); + if num_values == 0 { + return Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::empty(), + bits_per_value: self.uncompressed_bit_width, + num_values: 0, + block_info: BlockInfo::new(), + })); + } match self.uncompressed_bit_width { 8 => Self::unchunk::(data, num_values), 16 => Self::unchunk::(data, num_values), diff --git a/rust/lance-encoding/src/format.rs b/rust/lance-encoding/src/format.rs index f37e69b0216..0f438f75d7b 100644 --- a/rust/lance-encoding/src/format.rs +++ b/rust/lance-encoding/src/format.rs @@ -643,6 +643,31 @@ macro_rules! impl_common_protobuf_utils { ) } + #[allow(clippy::too_many_arguments)] + pub fn sparse_layout( + value_encoding: crate::format::$module::CompressiveEncoding, + num_buffers: u64, + num_items: u64, + num_visible_items: u64, + has_large_chunk: bool, + structural_layers: Vec, + ) -> crate::format::$module::PageLayout { + crate::format::$module::PageLayout { + layout: Some( + crate::format::$module::page_layout::Layout::SparseLayout( + crate::format::$module::SparseLayout { + value_compression: Some(value_encoding), + num_buffers, + num_items, + num_visible_items, + has_large_chunk, + structural_layers, + }, + ), + ), + } + } + pub fn blob_layout( inner_layout: crate::format::$module::PageLayout, def_meaning: &[DefinitionInterpretation], diff --git a/rust/lance-encoding/src/repdef.rs b/rust/lance-encoding/src/repdef.rs index b418b906de8..afdc534eb90 100644 --- a/rust/lance-encoding/src/repdef.rs +++ b/rust/lance-encoding/src/repdef.rs @@ -122,6 +122,445 @@ use crate::buffer::LanceBuffer; pub type LevelBuffer = Vec; +/// Native sparse structural representation used by the 2.3 sparse layout. +/// +/// Layers are stored from outer-most to inner-most, matching the order Arrow structural +/// encoders record offsets and validity. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct SparseStructuralPlan { + pub(crate) layers: Vec, + pub(crate) num_items: u64, + pub(crate) num_visible_items: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum SparseStructuralLayerPlan { + Validity { + num_slots: u64, + null_positions: Vec, + }, + List { + num_slots: u64, + num_child_slots: u64, + non_empty_positions: Vec, + counts: Vec, + constant_count: Option, + null_positions: Vec, + }, + FixedSizeList { + num_slots: u64, + dimension: u64, + null_positions: Vec, + }, +} + +impl SparseStructuralPlan { + pub(crate) fn to_repdef(&self) -> Result { + let mut builder = RepDefBuilder::default(); + for layer in &self.layers { + match layer { + SparseStructuralLayerPlan::Validity { + num_slots, + null_positions, + } => { + if null_positions.is_empty() { + builder.add_no_null(*num_slots as usize); + } else { + let validity = validity_from_null_positions(*num_slots, null_positions)?; + builder.add_validity_bitmap(NullBuffer::new(validity)); + } + } + SparseStructuralLayerPlan::List { + num_slots, + non_empty_positions, + counts, + null_positions, + .. + } => { + let (offsets, validity) = list_offsets_and_validity( + *num_slots, + non_empty_positions, + counts, + null_positions, + )?; + builder.add_offsets(offsets, validity); + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots, + dimension, + null_positions, + } => { + let validity = if null_positions.is_empty() { + None + } else { + Some(NullBuffer::new(validity_from_null_positions( + *num_slots, + null_positions, + )?)) + }; + builder.add_fsl(validity, *dimension as usize, *num_slots as usize); + } + } + } + Ok(RepDefBuilder::serialize(vec![builder])) + } +} + +fn validity_from_null_positions(num_slots: u64, null_positions: &[u64]) -> Result { + let mut validity = BooleanBufferBuilder::new(num_slots as usize); + let mut null_iter = null_positions.iter().copied().peekable(); + for slot in 0..num_slots { + let is_null = null_iter.peek().is_some_and(|null_pos| *null_pos == slot); + if is_null { + null_iter.next(); + } + validity.append(!is_null); + } + if let Some(extra_null) = null_iter.next() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural null position {} is outside layer with {} slots", + extra_null, num_slots + ) + .into(), + )); + } + Ok(validity.finish()) +} + +fn list_offsets_and_validity( + num_slots: u64, + non_empty_positions: &[u64], + counts: &[u64], + null_positions: &[u64], +) -> Result<(OffsetBuffer, Option)> { + if non_empty_positions.len() != counts.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list has {} non-empty positions but {} counts", + non_empty_positions.len(), + counts.len() + ) + .into(), + )); + } + + let mut offsets = Vec::with_capacity(num_slots as usize + 1); + let mut validity = + (!null_positions.is_empty()).then(|| BooleanBufferBuilder::new(num_slots as usize)); + let mut non_empty_iter = non_empty_positions + .iter() + .copied() + .zip(counts.iter().copied()) + .peekable(); + let mut null_iter = null_positions.iter().copied().peekable(); + let mut current_offset = 0_i64; + offsets.push(current_offset); + for slot in 0..num_slots { + let is_null = null_iter.peek().is_some_and(|null_pos| *null_pos == slot); + if is_null { + null_iter.next(); + } + if let Some(validity) = validity.as_mut() { + validity.append(!is_null); + } + + if non_empty_iter + .peek() + .is_some_and(|(non_empty_pos, _)| *non_empty_pos == slot) + { + let (_, count) = non_empty_iter.next().unwrap(); + if is_null { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list slot {} is both null and non-empty", + slot + ) + .into(), + )); + } + current_offset = current_offset + .checked_add(i64::try_from(count).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural list count {} exceeds i64::MAX", count).into(), + ) + })?) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list offsets overflow i64".into(), + ) + })?; + } + offsets.push(current_offset); + } + if let Some((extra_pos, _)) = non_empty_iter.next() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural non-empty position {} is outside layer with {} slots", + extra_pos, num_slots + ) + .into(), + )); + } + if let Some(extra_null) = null_iter.next() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural null position {} is outside layer with {} slots", + extra_null, num_slots + ) + .into(), + )); + } + + Ok(( + OffsetBuffer::new(ScalarBuffer::::from(offsets)), + validity.map(|mut validity| NullBuffer::new(validity.finish())), + )) +} + +#[derive(Debug)] +struct SparseStructuralUnraveler { + layers: Vec, + next_layer: usize, + num_items: u64, +} + +impl SparseStructuralUnraveler { + fn new(plan: SparseStructuralPlan) -> Self { + let next_layer = plan.layers.len(); + Self { + layers: plan.layers, + next_layer, + num_items: plan.num_visible_items, + } + } + + fn current_layer(&self) -> Option<&SparseStructuralLayerPlan> { + self.next_layer.checked_sub(1).map(|idx| &self.layers[idx]) + } + + fn consume_current_layer(&mut self) { + debug_assert!(self.next_layer > 0); + self.next_layer -= 1; + } + + fn is_all_valid(&self) -> bool { + match self.current_layer() { + Some(SparseStructuralLayerPlan::Validity { null_positions, .. }) + | Some(SparseStructuralLayerPlan::FixedSizeList { null_positions, .. }) + | Some(SparseStructuralLayerPlan::List { null_positions, .. }) => { + null_positions.is_empty() + } + None => true, + } + } + + fn max_lists(&self) -> usize { + match self.current_layer() { + Some(SparseStructuralLayerPlan::List { num_slots, .. }) => *num_slots as usize, + _ => 0, + } + } + + fn skip_validity(&mut self) { + match self.current_layer() { + Some(SparseStructuralLayerPlan::Validity { .. }) + | Some(SparseStructuralLayerPlan::FixedSizeList { .. }) => { + self.consume_current_layer(); + } + None => {} + Some(SparseStructuralLayerPlan::List { .. }) => { + unreachable!("Sparse structural list layer cannot be skipped as validity") + } + } + } + + fn append_validity( + validity: &mut BooleanBufferBuilder, + num_slots: u64, + null_positions: &[u64], + ) { + if null_positions.is_empty() { + validity.append_n(num_slots as usize, true); + return; + } + + let mut null_iter = null_positions.iter().copied().peekable(); + for slot in 0..num_slots { + let is_null = null_iter.peek().is_some_and(|null_pos| *null_pos == slot); + if is_null { + null_iter.next(); + } + validity.append(!is_null); + } + } + + fn unravel_validity(&mut self, validity: &mut BooleanBufferBuilder) { + match self.current_layer() { + Some(SparseStructuralLayerPlan::Validity { + num_slots, + null_positions, + }) + | Some(SparseStructuralLayerPlan::FixedSizeList { + num_slots, + null_positions, + .. + }) => { + Self::append_validity(validity, *num_slots, null_positions); + self.consume_current_layer(); + } + None => { + validity.append_n(self.num_items as usize, true); + } + Some(SparseStructuralLayerPlan::List { .. }) => { + unreachable!("Sparse structural list layer cannot be unraveled as validity") + } + } + } + + fn decimate(&mut self, dimension: usize) { + if let Some(SparseStructuralLayerPlan::FixedSizeList { + dimension: actual_dimension, + .. + }) = self.current_layer() + { + debug_assert_eq!(*actual_dimension as usize, dimension); + } + } + + fn to_offset(value: u64) -> Result { + let value = usize::try_from(value).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural offset {} exceeds usize::MAX", value).into(), + ) + })?; + T::from_usize(value).ok_or_else(|| { + Error::invalid_input( + "A single batch had more than i32::MAX values and so a large container type is required", + ) + }) + } + + fn unravel_offsets( + &mut self, + offsets: &mut Vec, + validity: Option<&mut BooleanBufferBuilder>, + ) -> Result<()> { + let Some(SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + null_positions, + .. + }) = self.current_layer() + else { + return Err(Error::internal( + "Expected sparse structural list layer while unraveling offsets".to_string(), + )); + }; + + if non_empty_positions.len() != counts.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list has {} non-empty positions but {} counts", + non_empty_positions.len(), + counts.len() + ) + .into(), + )); + } + let actual_child_slots = counts.iter().sum::(); + if actual_child_slots != *num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list count sum {} does not match child slots {}", + actual_child_slots, num_child_slots + ) + .into(), + )); + } + + let mut current_offset = offsets + .last() + .map(|offset| offset.as_usize() as u64) + .unwrap_or(0); + if offsets.is_empty() { + offsets.push(Self::to_offset(current_offset)?); + } + + if non_empty_positions.is_empty() { + if let Some(validity) = validity { + Self::append_validity(validity, *num_slots, null_positions); + } + let offset = Self::to_offset(current_offset)?; + offsets.resize(offsets.len() + *num_slots as usize, offset); + self.consume_current_layer(); + return Ok(()); + } + + let mut non_empty_iter = non_empty_positions + .iter() + .copied() + .zip(counts.iter().copied()) + .peekable(); + let mut null_iter = null_positions.iter().copied().peekable(); + let mut validity = validity; + for slot in 0..*num_slots { + let is_null = null_iter.peek().is_some_and(|null_pos| *null_pos == slot); + if is_null { + null_iter.next(); + } + if let Some(validity) = validity.as_mut() { + validity.append(!is_null); + } + + if non_empty_iter + .peek() + .is_some_and(|(non_empty_pos, _)| *non_empty_pos == slot) + { + let (_, count) = non_empty_iter.next().unwrap(); + if is_null { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list slot {} is both null and non-empty", + slot + ) + .into(), + )); + } + current_offset = current_offset.checked_add(count).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list offsets overflow u64".into(), + ) + })?; + } + offsets.push(Self::to_offset(current_offset)?); + } + if let Some((extra_pos, _)) = non_empty_iter.next() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural non-empty position {} is outside layer with {} slots", + extra_pos, num_slots + ) + .into(), + )); + } + if let Some(extra_null) = null_iter.next() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural null position {} is outside layer with {} slots", + extra_null, num_slots + ) + .into(), + )); + } + + self.consume_current_layer(); + Ok(()) + } +} + /// A contiguous top-level-row range that can be encoded as one structural page. #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct StructuralPageSplit { @@ -1427,6 +1866,134 @@ impl RepDefBuilder { ) } + pub(crate) fn to_sparse_structural_plan( + builders: &[Self], + num_visible_items: u64, + ) -> Result> { + if builders.is_empty() || builders.iter().all(|builder| builder.is_empty()) { + return Ok(None); + } + + let num_layers = builders[0].num_layers(); + if num_layers == 0 { + return Ok(None); + } + if builders + .iter() + .any(|builder| builder.num_layers() != num_layers) + { + return Err(Error::internal( + "Cannot build sparse structural plan from builders with different layer counts", + )); + } + + let combined_layers = (0..num_layers) + .map(|layer_index| { + Self::concat_layers( + builders.iter().map(|builder| &builder.repdefs[layer_index]), + builders.len(), + ) + }) + .collect::>(); + + let mut sparse_layers = Vec::with_capacity(combined_layers.len()); + for layer in combined_layers { + sparse_layers.push(Self::raw_layer_to_sparse(layer)?); + } + + let mut plan = SparseStructuralPlan { + layers: sparse_layers, + num_items: num_visible_items, + num_visible_items, + }; + let repdef = plan.to_repdef()?; + plan.num_items = repdef + .repetition_levels + .as_ref() + .map(|levels| levels.len() as u64) + .or_else(|| { + repdef + .definition_levels + .as_ref() + .map(|levels| levels.len() as u64) + }) + .unwrap_or(num_visible_items); + Ok(Some(plan)) + } + + fn null_positions(validity: Option<&BooleanBuffer>, num_values: usize) -> Vec { + if let Some(validity) = validity { + validity + .iter() + .enumerate() + .filter_map(|(idx, is_valid)| (!is_valid).then_some(idx as u64)) + .collect() + } else { + Vec::with_capacity(num_values) + } + } + + fn raw_layer_to_sparse(layer: RawRepDef) -> Result { + Ok(match layer { + RawRepDef::Validity(ValidityDesc { + validity, + num_values, + }) => SparseStructuralLayerPlan::Validity { + num_slots: num_values as u64, + null_positions: Self::null_positions(validity.as_ref(), num_values), + }, + RawRepDef::Fsl(FslDesc { + validity, + num_values, + dimension, + }) => SparseStructuralLayerPlan::FixedSizeList { + num_slots: num_values as u64, + dimension: dimension as u64, + null_positions: Self::null_positions(validity.as_ref(), num_values), + }, + RawRepDef::Offsets(OffsetDesc { + offsets, + validity, + num_values, + .. + }) => { + let mut non_empty_positions = Vec::new(); + let mut counts = Vec::new(); + let mut null_positions = Vec::new(); + for slot in 0..num_values { + let is_valid = validity + .as_ref() + .is_none_or(|validity| validity.value(slot)); + if !is_valid { + null_positions.push(slot as u64); + } + let len = offsets[slot + 1] - offsets[slot]; + if is_valid && len > 0 { + non_empty_positions.push(slot as u64); + counts.push(u64::try_from(len).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural list has negative length {}", len) + .into(), + ) + })?); + } + } + let constant_count = counts + .first() + .copied() + .filter(|first| counts.iter().all(|count| count == first)); + SparseStructuralLayerPlan::List { + num_slots: num_values as u64, + num_child_slots: *offsets.last().unwrap() as u64, + non_empty_positions, + counts, + constant_count, + null_positions, + } + } + }) + } + fn serialize_builders(builders: Vec) -> (SerializerContext, Option) { assert!(!builders.is_empty()); if builders.iter().all(|b| b.is_empty()) { @@ -1514,6 +2081,7 @@ impl RepDefBuilder { /// This is used during decoding to create the necessary arrow structures #[derive(Debug)] pub struct RepDefUnraveler { + sparse: Option, rep_levels: Option, def_levels: Option, // Maps from definition level to the rep level at which that definition level is visible @@ -1567,6 +2135,7 @@ impl RepDefUnraveler { } } Self { + sparse: None, rep_levels, def_levels, current_def_cmp: 0, @@ -1578,7 +2147,24 @@ impl RepDefUnraveler { } } + pub(crate) fn new_sparse(plan: SparseStructuralPlan) -> Self { + Self { + sparse: Some(SparseStructuralUnraveler::new(plan)), + rep_levels: None, + def_levels: None, + current_def_cmp: 0, + current_rep_cmp: 0, + levels_to_rep: Vec::new(), + current_layer: 0, + def_meaning: Arc::new([]), + num_items: 0, + } + } + pub fn is_all_valid(&self) -> bool { + if let Some(sparse) = &self.sparse { + return sparse.is_all_valid(); + } self.def_levels.is_none() || self.def_meaning[self.current_layer].is_all_valid() } @@ -1588,6 +2174,9 @@ impl RepDefUnraveler { /// This is not valid to call when the current level is a struct/primitive layer because /// in some cases there may be no rep or def information to know this. pub fn max_lists(&self) -> usize { + if let Some(sparse) = &self.sparse { + return sparse.max_lists(); + } debug_assert!( self.def_meaning[self.current_layer] != DefinitionInterpretation::NullableItem ); @@ -1607,6 +2196,9 @@ impl RepDefUnraveler { offsets: &mut Vec, validity: Option<&mut BooleanBufferBuilder>, ) -> Result<()> { + if let Some(sparse) = self.sparse.as_mut() { + return sparse.unravel_offsets(offsets, validity); + } let rep_levels = self .rep_levels .as_mut() @@ -1758,12 +2350,20 @@ impl RepDefUnraveler { } pub fn skip_validity(&mut self) { + if let Some(sparse) = self.sparse.as_mut() { + sparse.skip_validity(); + return; + } debug_assert!(self.is_all_valid()); self.current_layer += 1; } /// Unravels a layer of validity from the definition levels pub fn unravel_validity(&mut self, validity: &mut BooleanBufferBuilder) { + if let Some(sparse) = self.sparse.as_mut() { + sparse.unravel_validity(validity); + return; + } let meaning = self.def_meaning[self.current_layer]; if meaning == DefinitionInterpretation::AllValidItem || self.def_levels.is_none() { self.current_layer += 1; @@ -1789,6 +2389,10 @@ impl RepDefUnraveler { } pub fn decimate(&mut self, dimension: usize) { + if let Some(sparse) = self.sparse.as_mut() { + sparse.decimate(dimension); + return; + } if self.rep_levels.is_some() { // If we need to support this then I think we need to walk through the rep def levels to find // the spots at which we keep. E.g. if we have: diff --git a/rust/lance-encoding/src/testing.rs b/rust/lance-encoding/src/testing.rs index 176083d6d64..2e11d729b17 100644 --- a/rust/lance-encoding/src/testing.rs +++ b/rust/lance-encoding/src/testing.rs @@ -658,6 +658,12 @@ fn collect_page_encoding(layout: &PageLayout, actual_chain: &mut Vec) -> actual_chain.extend(chain); } } + Layout::SparseLayout(sparse) => { + if let Some(ref value_comp) = sparse.value_compression { + let chain = extract_array_encoding_chain(value_comp); + actual_chain.extend(chain); + } + } Layout::ConstantLayout(_) => { // Constant layout does not describe a value encoding chain. } From f039eea668a747b2e9cc387d28e4a9d82535e595 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 4 Jul 2026 13:11:42 +0800 Subject: [PATCH 2/8] fix: validate malformed sparse layouts --- docs/src/format/file/encoding.md | 51 ++ .../src/encodings/logical/primitive.rs | 601 ++++++++++++++++-- 2 files changed, 595 insertions(+), 57 deletions(-) diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index 4ca053d4fa6..12e32222b53 100644 --- a/docs/src/format/file/encoding.md +++ b/docs/src/format/file/encoding.md @@ -327,6 +327,57 @@ The protobuf for the full zip layout describes the compression of the data buffe size of the control words and how many bits we have per value (for fixed-width data) or how many bits we have per offset (for variable-width data). +### Sparse Page Layout + +The sparse page layout is a 2.3+ structural layout for sparse nested data. It is selected explicitly with +`lance-encoding:structural-encoding=sparse`, or automatically in 2.3+ when the writer would otherwise need to split a +page because the structural information is much larger than the visible value stream. Writers must not emit this layout +for file versions before 2.3. + +Sparse layout stores Arrow structure as native slot-domain mappings instead of dense repetition / definition events. +The value buffers are still mini-block compressed. The structural layers are stored from outer-most to inner-most: + +- validity layer: null slots in the parent slot domain +- list layer: non-empty parent slots, child counts, and null parent slots +- fixed-size-list layer: the fixed dimension and null parent slots + +Valid empty list slots are not stored in the non-empty positions. They are represented by their absence from both the +non-empty positions and null positions for that list layer. + +#### Buffers + +Sparse pages always have a metadata buffer followed by one value buffer, then zero or more structural buffers. + +| Buffer | Meaning | +| ------ | ------- | +| 0 | Value chunk metadata | +| 1 | Mini-block encoded value chunks | +| 2+ | Structural position/count buffers, in `SparseStructuralLayer` order | + +Buffer 0 stores one 8-byte entry per value chunk. The first 4 bytes store the chunk size divided by 8 minus one. The +second 4 bytes store the number of visible values in the chunk. The sum of chunk value counts must match +`SparseLayout.num_visible_items`, and the sum of chunk byte sizes must exactly match the value buffer size. + +Structural position buffers store delta-encoded `u64` positions. Positions must be strictly increasing after delta +decoding and must be within the layer's parent slot domain. List count buffers store `u64` child counts for each +non-empty list slot. If all non-empty lists have the same length then `constant_count` is used instead of a count +buffer. + +#### Protobuf + +```protobuf +%%% proto.message.SparseLayout %%% +``` + +```protobuf +%%% proto.message.SparseStructuralLayer %%% +``` + +The `value_compression` field is required. The physical page buffer list must match the structural layer metadata: +each non-zero position/count cardinality requires exactly one corresponding compressed structural buffer, and zero +cardinality layers must not include an unused compression descriptor. Readers should reject malformed sparse pages with +a format error. + ### Constant Page Layout This layout is used when all (visible) values in the page are the same scalar value. diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index aebb334782d..db1547f2e26 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -695,6 +695,7 @@ enum SparseLayerDecompressors { }, FixedSizeList { num_slots: u64, + num_child_slots: u64, dimension: u64, null_decompressor: Option>, num_nulls: u64, @@ -766,10 +767,23 @@ impl SparseStructuralScheduler { layout: &pb21::SparseLayout, decompressors: &dyn DecompressionStrategy, ) -> Result { - let value_decompressor = decompressors.create_miniblock_decompressor( - layout.value_compression.as_ref().unwrap(), - decompressors, - )?; + let value_compression = layout.value_compression.as_ref().ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing value compression".into()) + })?; + Self::validate_compression(value_compression, "value compression")?; + let expected_buffers = 2 + Self::structural_buffer_count(&layout.structural_layers)?; + if buffer_offsets_and_sizes.len() != expected_buffers { + return Err(Error::invalid_input_source( + format!( + "Sparse layout has {} buffers, expected {}", + buffer_offsets_and_sizes.len(), + expected_buffers + ) + .into(), + )); + } + let value_decompressor = + decompressors.create_miniblock_decompressor(value_compression, decompressors)?; let layer_decompressors = layout .structural_layers .iter() @@ -790,6 +804,130 @@ impl SparseStructuralScheduler { }) } + fn validate_compression<'a>( + compression: &'a CompressiveEncoding, + label: &str, + ) -> Result<&'a CompressiveEncoding> { + compression.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} is missing compression details").into(), + ) + })?; + Ok(compression) + } + + fn metadata_buffer(&self) -> Result<(u64, u64)> { + self.buffer_offsets_and_sizes + .first() + .copied() + .ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing metadata buffer".into()) + }) + } + + fn value_buffer(&self) -> Result<(u64, u64)> { + self.buffer_offsets_and_sizes + .get(1) + .copied() + .ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing value buffer".into()) + }) + } + + fn count_optional_buffer( + compression: &Option, + count: u64, + label: &str, + ) -> Result { + if count == 0 { + if compression.is_some() { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} has compression but no values").into(), + )); + } + return Ok(0); + } + let compression = compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} has values but no compression").into(), + ) + })?; + Self::validate_compression(compression, label)?; + Ok(1) + } + + fn structural_buffer_count(layers: &[pb21::SparseStructuralLayer]) -> Result { + layers.iter().try_fold(0_usize, |mut count, layer| { + let kind = pb21::sparse_structural_layer::Kind::try_from(layer.kind)?; + match kind { + pb21::sparse_structural_layer::Kind::SparseLayerValidity => { + count += Self::count_optional_buffer( + &layer.null_compression, + layer.num_nulls, + "validity null positions", + )?; + } + pb21::sparse_structural_layer::Kind::SparseLayerList => { + count += Self::count_optional_buffer( + &layer.non_empty_compression, + layer.num_non_empty, + "list non-empty positions", + )?; + if layer.num_non_empty > 0 + && layer.count_compression.is_none() + && layer.constant_count == 0 + { + return Err(Error::invalid_input_source( + "Sparse structural list has non-empty positions but no count compression or constant count".into(), + )); + } + if layer.count_compression.is_some() { + let compression = layer.count_compression.as_ref().expect_ok()?; + Self::validate_compression(compression, "list counts")?; + count += 1; + } + count += Self::count_optional_buffer( + &layer.null_compression, + layer.num_nulls, + "list null positions", + )?; + } + pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList => { + let expected_child_slots = + layer.num_slots.checked_mul(layer.constant_count).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", + layer.num_slots, layer.constant_count + ) + .into(), + ) + })?; + if expected_child_slots != layer.num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count {} does not match slots {} * dimension {}", + layer.num_child_slots, layer.num_slots, layer.constant_count + ) + .into(), + )); + } + count += Self::count_optional_buffer( + &layer.null_compression, + layer.num_nulls, + "fixed-size-list null positions", + )?; + } + pb21::sparse_structural_layer::Kind::SparseLayerUnspecified => { + return Err(Error::invalid_input_source( + "Sparse structural layer kind is unspecified".into(), + )); + } + } + Ok(count) + }) + } + fn layer_decompressors( layer: &pb21::SparseStructuralLayer, decompressors: &dyn DecompressionStrategy, @@ -798,6 +936,7 @@ impl SparseStructuralScheduler { .null_compression .as_ref() .map(|compression| { + let compression = Self::validate_compression(compression, "null positions")?; decompressors .create_block_decompressor(compression) .map(Arc::from) @@ -817,6 +956,8 @@ impl SparseStructuralScheduler { .non_empty_compression .as_ref() .map(|compression| { + let compression = + Self::validate_compression(compression, "list non-empty positions")?; decompressors .create_block_decompressor(compression) .map(Arc::from) @@ -826,6 +967,7 @@ impl SparseStructuralScheduler { .count_compression .as_ref() .map(|compression| { + let compression = Self::validate_compression(compression, "list counts")?; decompressors .create_block_decompressor(compression) .map(Arc::from) @@ -845,6 +987,7 @@ impl SparseStructuralScheduler { pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList => { SparseLayerDecompressors::FixedSizeList { num_slots: layer.num_slots, + num_child_slots: layer.num_child_slots, dimension: layer.constant_count, null_decompressor, num_nulls: layer.num_nulls, @@ -869,8 +1012,13 @@ impl SparseStructuralScheduler { )); } - let value_buf_position = self.buffer_offsets_and_sizes[1].0; - let mut rows_counter = 0; + let (value_buf_position, value_buf_size) = self.value_buffer()?; + let value_buf_end = value_buf_position + .checked_add(value_buf_size) + .ok_or_else(|| { + Error::invalid_input_source("Sparse layout value buffer range overflows".into()) + })?; + let mut rows_counter = 0_u64; let mut offset_bytes = value_buf_position; let mut chunk_meta = Vec::with_capacity(meta_bytes.len() / 8); for chunk in meta_bytes.chunks_exact(8) { @@ -878,13 +1026,17 @@ impl SparseStructuralScheduler { u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]); let num_values = u32::from_le_bytes([chunk[4], chunk[5], chunk[6], chunk[7]]) as u64; let num_bytes = (divided_bytes_minus_one as usize + 1) * MINIBLOCK_ALIGNMENT; - rows_counter += num_values; + rows_counter = rows_counter.checked_add(num_values).ok_or_else(|| { + Error::invalid_input_source("Sparse layout visible item count overflows".into()) + })?; chunk_meta.push(ChunkMeta { num_values, chunk_size_bytes: num_bytes as u64, offset_bytes, }); - offset_bytes += num_bytes as u64; + offset_bytes = offset_bytes.checked_add(num_bytes as u64).ok_or_else(|| { + Error::invalid_input_source("Sparse layout value chunk byte range overflows".into()) + })?; } if rows_counter != self.num_visible_items { return Err(Error::invalid_input_source( @@ -895,6 +1047,16 @@ impl SparseStructuralScheduler { .into(), )); } + if offset_bytes != value_buf_end { + return Err(Error::invalid_input_source( + format!( + "Sparse layout chunk metadata describes {} value bytes, but value buffer has {} bytes", + offset_bytes - value_buf_position, + value_buf_size + ) + .into(), + )); + } Ok(chunk_meta) } @@ -949,6 +1111,12 @@ impl SparseStructuralScheduler { if idx == 0 { current = delta; } else { + if delta == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} positions must be strictly increasing") + .into(), + )); + } current = current.checked_add(delta).ok_or_else(|| { Error::invalid_input_source( format!("Sparse structural {label} position overflow").into(), @@ -969,6 +1137,29 @@ impl SparseStructuralScheduler { Ok(positions) } + fn next_structural_buffer( + buffers: &mut impl Iterator, + label: &str, + ) -> Result { + buffers.next().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} is missing its buffer").into(), + ) + }) + } + + fn optional_structural_buffer( + buffers: &mut impl Iterator, + needed: bool, + label: &str, + ) -> Result> { + if needed { + Self::next_structural_buffer(buffers, label).map(Some) + } else { + Ok(None) + } + } + fn decode_layer( layer: &SparseLayerDecompressors, buffers: &mut impl Iterator, @@ -982,7 +1173,11 @@ impl SparseStructuralScheduler { num_slots: *num_slots, null_positions: Self::decode_positions( null_decompressor.as_ref(), - (*num_nulls > 0).then(|| buffers.next().unwrap()), + Self::optional_structural_buffer( + buffers, + *num_nulls > 0, + "validity null positions", + )?, *num_nulls, *num_slots, "validity null positions", @@ -1000,7 +1195,11 @@ impl SparseStructuralScheduler { } => { let non_empty_positions = Self::decode_positions( non_empty_decompressor.as_ref(), - (*num_non_empty > 0).then(|| buffers.next().unwrap()), + Self::optional_structural_buffer( + buffers, + *num_non_empty > 0, + "list non-empty positions", + )?, *num_non_empty, *num_slots, "list non-empty positions", @@ -1008,7 +1207,7 @@ impl SparseStructuralScheduler { let counts = if let Some(count_decompressor) = count_decompressor { Self::decode_u64_values( count_decompressor.as_ref(), - buffers.next().unwrap(), + Self::next_structural_buffer(buffers, "list counts")?, *num_non_empty, "list counts", )? @@ -1017,12 +1216,22 @@ impl SparseStructuralScheduler { }; let null_positions = Self::decode_positions( null_decompressor.as_ref(), - (*num_nulls > 0).then(|| buffers.next().unwrap()), + Self::optional_structural_buffer( + buffers, + *num_nulls > 0, + "list null positions", + )?, *num_nulls, *num_slots, "list null positions", )?; - let actual_child_slots = counts.iter().sum::(); + let actual_child_slots = counts.iter().try_fold(0_u64, |sum, count| { + sum.checked_add(*count).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list count sum overflows".into(), + ) + }) + })?; if actual_child_slots != *num_child_slots { return Err(Error::invalid_input_source( format!( @@ -1043,47 +1252,118 @@ impl SparseStructuralScheduler { } SparseLayerDecompressors::FixedSizeList { num_slots, + num_child_slots, dimension, null_decompressor, num_nulls, - } => SparseStructuralLayerPlan::FixedSizeList { - num_slots: *num_slots, - dimension: *dimension, - null_positions: Self::decode_positions( - null_decompressor.as_ref(), - (*num_nulls > 0).then(|| buffers.next().unwrap()), - *num_nulls, - *num_slots, - "fixed-size-list null positions", - )?, - }, + } => { + let expected_child_slots = num_slots.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", + num_slots, dimension + ) + .into(), + ) + })?; + if expected_child_slots != *num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count {} does not match slots {} * dimension {}", + num_child_slots, num_slots, dimension + ) + .into(), + )); + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots: *num_slots, + dimension: *dimension, + null_positions: Self::decode_positions( + null_decompressor.as_ref(), + Self::optional_structural_buffer( + buffers, + *num_nulls > 0, + "fixed-size-list null positions", + )?, + *num_nulls, + *num_slots, + "fixed-size-list null positions", + )?, + } + } }) } - fn lookup_value_chunks(&self, chunk_indices: &[usize]) -> Vec { + fn lookup_value_chunks(&self, chunk_indices: &[usize]) -> Result> { let page_meta = self.page_meta.as_ref().unwrap(); chunk_indices .iter() .map(|&chunk_idx| { - let chunk_meta = &page_meta.chunk_meta[chunk_idx]; + let chunk_meta = page_meta.chunk_meta.get(chunk_idx).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse layout missing value chunk metadata for chunk {chunk_idx}") + .into(), + ) + })?; let bytes_start = chunk_meta.offset_bytes; let bytes_end = bytes_start + chunk_meta.chunk_size_bytes; - LoadedChunk { + Ok(LoadedChunk { byte_range: bytes_start..bytes_end, items_in_chunk: chunk_meta.num_values, chunk_idx, data: LanceBuffer::empty(), - } + }) }) .collect() } - fn value_chunk_range(chunk_value_offsets: &[u64], value_range: Range) -> Range { - let start = chunk_value_offsets.partition_point(|&offset| offset <= value_range.start) - 1; + fn value_chunk_index(chunk_value_offsets: &[u64], value: u64) -> Result { + let total_values = chunk_value_offsets.last().copied().ok_or_else(|| { + Error::invalid_input_source("Sparse layout has no value chunk offsets".into()) + })?; + if chunk_value_offsets.len() < 2 || value >= total_values { + return Err(Error::invalid_input_source( + format!( + "Sparse layout value index {} is outside {} visible items", + value, total_values + ) + .into(), + )); + } + chunk_value_offsets + .partition_point(|&offset| offset <= value) + .checked_sub(1) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse layout value index {value} is before the first chunk").into(), + ) + }) + } + + fn value_chunk_range( + chunk_value_offsets: &[u64], + value_range: Range, + ) -> Result> { + if value_range.is_empty() { + return Ok(0..0); + } + let total_values = chunk_value_offsets.last().copied().ok_or_else(|| { + Error::invalid_input_source("Sparse layout has no value chunk offsets".into()) + })?; + if value_range.start > value_range.end || value_range.end > total_values { + return Err(Error::invalid_input_source( + format!( + "Sparse layout value range {}..{} is outside {} visible items", + value_range.start, value_range.end, total_values + ) + .into(), + )); + } + let start = Self::value_chunk_index(chunk_value_offsets, value_range.start)?; let end = chunk_value_offsets .partition_point(|&offset| offset < value_range.end) .max(start + 1); - start..end + Ok(start..end) } } @@ -1092,7 +1372,10 @@ impl StructuralPageScheduler for SparseStructuralScheduler { &'a mut self, io: &Arc, ) -> BoxFuture<'a, Result>> { - let (meta_buf_position, meta_buf_size) = self.buffer_offsets_and_sizes[0]; + let (meta_buf_position, meta_buf_size) = match self.metadata_buffer() { + Ok(buffer) => buffer, + Err(err) => return std::future::ready(Err(err)).boxed(), + }; let mut required_ranges = Vec::new(); required_ranges.push(meta_buf_position..meta_buf_position + meta_buf_size); for (position, size) in self.buffer_offsets_and_sizes.iter().skip(2) { @@ -1102,14 +1385,20 @@ impl StructuralPageScheduler for SparseStructuralScheduler { async move { let mut buffers = io_req.await?.into_iter(); - let meta_bytes = buffers.next().unwrap(); + let meta_bytes = buffers.next().ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing chunk metadata buffer".into()) + })?; let chunk_meta = self.parse_chunk_meta(meta_bytes)?; let mut chunk_value_offsets = Vec::with_capacity(chunk_meta.len() + 1); - let mut value_offset = 0; + let mut value_offset = 0_u64; chunk_value_offsets.push(value_offset); for chunk in &chunk_meta { - value_offset += chunk.num_values; + value_offset = value_offset.checked_add(chunk.num_values).ok_or_else(|| { + Error::invalid_input_source( + "Sparse layout visible item offset overflows".into(), + ) + })?; chunk_value_offsets.push(value_offset); } @@ -1158,12 +1447,12 @@ impl StructuralPageScheduler for SparseStructuralScheduler { chunks_needed.extend(Self::value_chunk_range( &page_meta.chunk_value_offsets, value_range.clone(), - )); + )?); } chunks_needed.sort_unstable(); chunks_needed.dedup(); - let mut loaded_chunks = self.lookup_value_chunks(&chunks_needed); + let mut loaded_chunks = self.lookup_value_chunks(&chunks_needed)?; let chunk_ranges = loaded_chunks .iter() .map(|chunk| chunk.byte_range.clone()) @@ -1280,6 +1569,16 @@ impl DecodeSparseStructuralTask { fn decode_value_chunk(&self, chunk: &LoadedChunk) -> Result { let buf = &chunk.data; let mut offset = 0; + if buf.len() < 2 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} is too small for its header: {} bytes", + chunk.chunk_idx, + buf.len() + ) + .into(), + )); + } let num_levels = u16::from_le_bytes([buf[offset], buf[offset + 1]]); offset += 2; if num_levels != 0 { @@ -1292,6 +1591,37 @@ impl DecodeSparseStructuralTask { )); } + let size_width = if self.has_large_chunk { 4 } else { 2 }; + let num_buffers = usize::try_from(self.num_buffers).map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk has too many buffers: {}", + self.num_buffers + ) + .into(), + ) + })?; + let sizes_len = num_buffers.checked_mul(size_width).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural value chunk buffer-size header overflows".into(), + ) + })?; + let header_len = offset.checked_add(sizes_len).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural value chunk header length overflows".into(), + ) + })?; + if buf.len() < header_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} is too small for {} buffer sizes: {} bytes", + chunk.chunk_idx, + self.num_buffers, + buf.len() + ) + .into(), + )); + } let buffer_sizes = if self.has_large_chunk { DecodeMiniBlockTask::read_buffer_sizes::(buf, &mut offset, self.num_buffers) } else { @@ -1299,15 +1629,52 @@ impl DecodeSparseStructuralTask { }; offset += pad_bytes::(offset); + if offset > buf.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} is missing padding after its header", + chunk.chunk_idx + ) + .into(), + )); + } let buffers = buffer_sizes .into_iter() .map(|buf_size| { - let buffer = buf.slice_with_length(offset, buf_size as usize); - offset += buf_size as usize; + let buf_size = buf_size as usize; + let end = offset.checked_add(buf_size).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} buffer size overflows", + chunk.chunk_idx + ) + .into(), + ) + })?; + if end > buf.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} buffer extends past chunk end", + chunk.chunk_idx + ) + .into(), + )); + } + let buffer = buf.slice_with_length(offset, buf_size); + offset = end; offset += pad_bytes::(offset); - buffer + if offset > buf.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} padding extends past chunk end", + chunk.chunk_idx + ) + .into(), + )); + } + Ok(buffer) }) - .collect::>(); + .collect::>>()?; self.value_decompressor .decompress(buffers, chunk.items_in_chunk) @@ -1321,11 +1688,10 @@ impl DecodeSparseStructuralTask { ) -> Result<()> { let mut value_start = value_range.start; while value_start < value_range.end { - let chunk_idx = self - .page_meta - .chunk_value_offsets - .partition_point(|&offset| offset <= value_start) - - 1; + let chunk_idx = SparseStructuralScheduler::value_chunk_index( + &self.page_meta.chunk_value_offsets, + value_start, + )?; let chunk_value_start = self.page_meta.chunk_value_offsets[chunk_idx]; let chunk_value_end = self.page_meta.chunk_value_offsets[chunk_idx + 1]; let take_end = value_range.end.min(chunk_value_end); @@ -1334,7 +1700,10 @@ impl DecodeSparseStructuralTask { let chunk = self.loaded_chunk(chunk_idx)?; *chunk_cache = Some((chunk_idx, self.decode_value_chunk(chunk)?)); } - let values = &chunk_cache.as_ref().unwrap().1; + let values = &chunk_cache + .as_ref() + .ok_or_else(|| Error::internal("Sparse structural chunk cache is empty"))? + .1; data_builder.append( values, value_start - chunk_value_start..take_end - chunk_value_start, @@ -1400,15 +1769,17 @@ fn translate_positions(positions: &[u64], ranges: &[Range], label: &str) -> Ok(translated) } -fn offsets_from_counts(counts: &[u64]) -> Vec { +fn offsets_from_counts(counts: &[u64]) -> Result> { let mut offsets = Vec::with_capacity(counts.len() + 1); let mut offset = 0_u64; offsets.push(offset); for count in counts { - offset += count; + offset = offset.checked_add(*count).ok_or_else(|| { + Error::invalid_input_source("Sparse structural list count offsets overflow".into()) + })?; offsets.push(offset); } - offsets + Ok(offsets) } fn coalesce_ranges(ranges: Vec>) -> Vec> { @@ -1498,13 +1869,13 @@ fn slice_list_layer( }) .collect() } else { - let value_offsets = offsets_from_counts(counts); - if *value_offsets.last().unwrap() != num_child_slots { + let value_offsets = offsets_from_counts(counts)?; + let last_offset = *value_offsets.last().expect_ok()?; + if last_offset != num_child_slots { return Err(Error::invalid_input_source( format!( "Sparse structural list count sum {} does not match child slots {}", - value_offsets.last().unwrap(), - num_child_slots + last_offset, num_child_slots ) .into(), )); @@ -7308,8 +7679,8 @@ mod tests { ChunkInstructions, DataBlock, DecodeMiniBlockTask, FixedPerValueDecompressor, FixedWidthDataBlock, FullZipCacheableState, FullZipDecodeDetails, FullZipReadSource, FullZipRepIndexDetails, FullZipScheduler, MiniBlockChunk, MiniBlockCompressed, - MiniBlockRepIndex, PerValueDecompressor, PreambleAction, StructuralPageScheduler, - VariableFullZipDecoder, + MiniBlockRepIndex, PerValueDecompressor, PreambleAction, SparseStructuralScheduler, + StructuralPageScheduler, VariableFullZipDecoder, }; use crate::buffer::LanceBuffer; use crate::compression::DefaultDecompressionStrategy; @@ -7327,10 +7698,12 @@ mod tests { use crate::format::pb21; use crate::format::pb21::compressive_encoding::Compression; use crate::repdef::build_control_word_iterator; - use crate::testing::{TestCases, check_round_trip_encoding_of_data}; + use crate::testing::{SimulatedScheduler, TestCases, check_round_trip_encoding_of_data}; use crate::version::LanceFileVersion; use arrow_array::{ArrayRef, Int8Array, StringArray}; use arrow_schema::{DataType, Field as ArrowField}; + use bytes::Bytes; + use lance_core::error::Error; use std::collections::HashMap; use std::{collections::VecDeque, sync::Arc}; @@ -9211,6 +9584,120 @@ mod tests { .collect() } + fn sparse_test_layout() -> crate::format::pb21::SparseLayout { + crate::format::pb21::SparseLayout { + value_compression: Some(ProtobufUtils21::flat(32, None)), + num_buffers: 1, + num_items: 1, + num_visible_items: 1, + has_large_chunk: false, + structural_layers: Vec::new(), + } + } + + fn assert_invalid_input_contains(err: Error, expected: &str) { + let message = err.to_string(); + assert!( + message.contains(expected), + "expected error to contain {expected:?}, got {message}" + ); + } + + #[test] + fn test_sparse_layout_rejects_missing_value_compression() { + let mut layout = sparse_test_layout(); + layout.value_compression = None; + let decompressors = DefaultDecompressionStrategy::default(); + + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + + assert_invalid_input_contains(err, "Sparse layout is missing value compression"); + } + + #[test] + fn test_sparse_layout_rejects_missing_value_buffer() { + let layout = sparse_test_layout(); + let decompressors = DefaultDecompressionStrategy::default(); + + let err = SparseStructuralScheduler::try_new( + &[(0, 0)], + 0, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + + assert_invalid_input_contains(err, "Sparse layout has 1 buffers, expected 2"); + } + + #[test] + fn test_sparse_layout_rejects_missing_structural_buffer() { + let mut layout = sparse_test_layout(); + layout + .structural_layers + .push(crate::format::pb21::SparseStructuralLayer { + kind: crate::format::pb21::sparse_structural_layer::Kind::SparseLayerValidity + as i32, + num_slots: 1, + num_child_slots: 1, + constant_count: 0, + non_empty_compression: None, + num_non_empty: 0, + count_compression: None, + null_compression: Some(ProtobufUtils21::flat(64, None)), + num_nulls: 1, + }); + let decompressors = DefaultDecompressionStrategy::default(); + + let err = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + + assert_invalid_input_contains(err, "Sparse layout has 2 buffers, expected 3"); + } + + #[tokio::test] + async fn test_sparse_layout_rejects_value_metadata_past_value_buffer() { + let layout = sparse_test_layout(); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + + let mut data = Vec::new(); + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&[0; 8]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + + let Err(err) = scheduler.initialize(&io).await else { + panic!("expected sparse layout initialization to fail"); + }; + + assert_invalid_input_contains( + err, + "Sparse layout chunk metadata describes 16 value bytes, but value buffer has 8 bytes", + ); + } + #[tokio::test] async fn test_explicit_sparse_layout_roundtrip() { let arr = sparse_i32_list_array(4096, 64); From 42d00ab15fe700d4bf8dada8cc96cfef2c076d62 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 4 Jul 2026 13:17:01 +0800 Subject: [PATCH 3/8] docs: document sparse layout selection --- docs/src/format/file/encoding.md | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index 12e32222b53..1a74a89977f 100644 --- a/docs/src/format/file/encoding.md +++ b/docs/src/format/file/encoding.md @@ -344,6 +344,18 @@ The value buffers are still mini-block compressed. The structural layers are sto Valid empty list slots are not stored in the non-empty positions. They are represented by their absence from both the non-empty positions and null positions for that list layer. +#### Selection and Compatibility + +Writers may emit sparse layout only for file versions 2.3 and later. Explicit sparse selection is requested with field +metadata `lance-encoding:structural-encoding=sparse`; the request is invalid for older file versions. In 2.3 and later, +the automatic writer path may select sparse layout when the page has native sparse structural layers and the equivalent +mini-block repetition / definition stream would exceed the mini-block structural page budget. When sparse is not requested +and the structural budget is satisfied, writers use the normal mini-block / full-zip selection path. + +Sparse layout is a page layout, not a column-level promise. Readers must identify it from `PageLayout.sparse_layout`. +The `lance-encoding:structural-encoding` field metadata is a writer input only and must not be used by readers to infer +the layout of existing pages. + #### Buffers Sparse pages always have a metadata buffer followed by one value buffer, then zero or more structural buffers. @@ -600,7 +612,7 @@ options. However, they can also be set in the field metadata in the schema. | `lance-encoding:dict-values-compression-level` | Integers (scheme dependent) | Varies by scheme | Compression level for dictionary values general compression | | `lance-encoding:general` | `off`, `on` | `off` | Whether to apply general compression. | | `lance-encoding:packed` | Any string | Not set | Whether to apply packed struct encoding (see above). | -| `lance-encoding:structural-encoding` | `miniblock`, `fullzip` | Not set | Force a particular structural encoding to be applied (only useful for testing purposes) | +| `lance-encoding:structural-encoding` | `miniblock`, `fullzip`, `sparse` | Not set | Force a structural encoding. `sparse` requires file version 2.3+. | ### Configuration Details From 5e556ad10a1f13689c148b9e063cbeabf63cf3fd Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 4 Jul 2026 13:35:23 +0800 Subject: [PATCH 4/8] fix: handle sparse nested structural round trips --- .../src/encodings/logical/primitive.rs | 437 +++++++++++++++++- rust/lance-encoding/src/repdef.rs | 2 +- 2 files changed, 429 insertions(+), 10 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index db1547f2e26..c33fe0dd53b 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -707,6 +707,7 @@ struct SparseStructuralCacheableState { chunk_meta: Vec, chunk_value_offsets: Arc<[u64]>, plan: SparseStructuralPlan, + row_domain: u64, } impl DeepSizeOf for SparseStructuralCacheableState { @@ -749,6 +750,7 @@ impl CachedPageData for SparseStructuralCacheableState { struct SparseStructuralScheduler { buffer_offsets_and_sizes: Vec<(u64, u64)>, priority: u64, + row_domain: u64, num_items: u64, num_visible_items: u64, num_buffers: u64, @@ -763,6 +765,7 @@ impl SparseStructuralScheduler { fn try_new( buffer_offsets_and_sizes: &[(u64, u64)], priority: u64, + row_domain: u64, data_type: DataType, layout: &pb21::SparseLayout, decompressors: &dyn DecompressionStrategy, @@ -793,6 +796,7 @@ impl SparseStructuralScheduler { Ok(Self { buffer_offsets_and_sizes: buffer_offsets_and_sizes.to_vec(), priority, + row_domain, num_items: layout.num_items, num_visible_items: layout.num_visible_items, num_buffers: layout.num_buffers, @@ -1417,6 +1421,7 @@ impl StructuralPageScheduler for SparseStructuralScheduler { chunk_meta, chunk_value_offsets: chunk_value_offsets.into(), plan, + row_domain: self.row_domain, }); self.page_meta = Some(page_meta.clone()); Ok(page_meta as Arc) @@ -1442,7 +1447,7 @@ impl StructuralPageScheduler for SparseStructuralScheduler { let page_meta = self.page_meta.as_ref().unwrap(); let mut chunks_needed = Vec::new(); - let selection = slice_sparse_plan(&page_meta.plan, ranges)?; + let selection = slice_sparse_plan(&page_meta.plan, ranges, page_meta.row_domain)?; for value_range in &selection.leaf_ranges { chunks_needed.extend(Self::value_chunk_range( &page_meta.chunk_value_offsets, @@ -1716,7 +1721,11 @@ impl DecodeSparseStructuralTask { impl DecodePageTask for DecodeSparseStructuralTask { fn decode(self: Box) -> Result { - let selection = slice_sparse_plan(&self.page_meta.plan, &self.row_ranges)?; + let selection = slice_sparse_plan( + &self.page_meta.plan, + &self.row_ranges, + self.page_meta.row_domain, + )?; let estimated_size_bytes = self .loaded_chunks .iter() @@ -1909,11 +1918,38 @@ fn slice_list_layer( )) } +fn fsl_parent_ranges_from_child_ranges( + ranges: &[Range], + dimension: u64, + num_child_slots: u64, +) -> Result>> { + if ranges.iter().any(|range| { + range.start > range.end + || range.end > num_child_slots + || range.start % dimension != 0 + || range.end % dimension != 0 + }) { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child-domain slice is not aligned to dimension {}", + dimension + ) + .into(), + )); + } + Ok(ranges + .iter() + .map(|range| range.start / dimension..range.end / dimension) + .collect()) +} + fn slice_sparse_plan( plan: &SparseStructuralPlan, row_ranges: &[Range], + row_domain: u64, ) -> Result { let mut selected_ranges = row_ranges.to_vec(); + let mut selected_domain = row_domain; let mut sliced_layers = Vec::with_capacity(plan.layers.len()); for layer in &plan.layers { @@ -1955,30 +1991,65 @@ fn slice_sparse_plan( )?; sliced_layers.push(sliced_layer); selected_ranges = child_ranges; + selected_domain = *num_child_slots; } SparseStructuralLayerPlan::FixedSizeList { - num_slots: _, + num_slots, dimension, null_positions, } => { - let out_num_slots = selected_ranges + let num_child_slots = num_slots.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", + num_slots, dimension + ) + .into(), + ) + })?; + let already_in_child_domain = + selected_domain == num_child_slots && selected_domain != *num_slots; + let fsl_ranges = if already_in_child_domain { + fsl_parent_ranges_from_child_ranges( + &selected_ranges, + *dimension, + num_child_slots, + )? + } else { + if selected_domain != *num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list slice domain {} does not match slots {} or child slots {}", + selected_domain, num_slots, num_child_slots + ) + .into(), + )); + } + selected_ranges.clone() + }; + let out_num_slots = fsl_ranges .iter() .map(|range| range.end - range.start) .sum::(); - let child_ranges = selected_ranges - .iter() - .map(|range| range.start * dimension..range.end * dimension) - .collect::>(); + let child_ranges = if already_in_child_domain { + selected_ranges.clone() + } else { + selected_ranges + .iter() + .map(|range| range.start * dimension..range.end * dimension) + .collect::>() + }; sliced_layers.push(SparseStructuralLayerPlan::FixedSizeList { num_slots: out_num_slots, dimension: *dimension, null_positions: translate_positions( null_positions, - &selected_ranges, + &fsl_ranges, "fixed-size-list null", )?, }); selected_ranges = child_ranges; + selected_domain = num_child_slots; } } } @@ -4691,6 +4762,7 @@ impl StructuralPrimitiveFieldScheduler { Layout::SparseLayout(sparse) => Box::new(SparseStructuralScheduler::try_new( &page_info.buffer_offsets_and_sizes, page_info.priority, + page_info.num_rows, target_field.data_type(), sparse, decompressors, @@ -9584,6 +9656,252 @@ mod tests { .collect() } + fn sparse_layouts_from_pages( + pages: &[crate::encoder::EncodedPage], + ) -> Vec<&crate::format::pb21::SparseLayout> { + pages + .iter() + .filter_map(|page| { + let PageEncoding::Structural(layout) = &page.description else { + return None; + }; + let pb21::page_layout::Layout::SparseLayout(sparse) = + layout.layout.as_ref().unwrap() + else { + return None; + }; + Some(sparse) + }) + .collect() + } + + fn sparse_nested_coverage_array() -> ArrayRef { + use arrow_array::{FixedSizeListArray, Int32Array, ListArray, StructArray}; + use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::Fields; + + fn offsets(values: Vec) -> OffsetBuffer { + OffsetBuffer::new(ScalarBuffer::from(values)) + } + + fn validity(values: Vec) -> Option { + Some(NullBuffer::new(BooleanBuffer::from(values))) + } + + let event_lengths = [ + Some(2), + Some(0), + None, + Some(0), + Some(1), + Some(3), + Some(0), + None, + Some(1), + Some(0), + Some(2), + Some(0), + Some(1), + Some(0), + None, + Some(2), + ]; + let mut event_offsets = Vec::with_capacity(event_lengths.len() + 1); + let mut event_validity = Vec::with_capacity(event_lengths.len()); + let mut event_child_count = 0_i32; + event_offsets.push(event_child_count); + for length in event_lengths { + match length { + Some(length) => { + event_child_count += length; + event_validity.push(true); + } + None => event_validity.push(false), + } + event_offsets.push(event_child_count); + } + + let id_values = Arc::new(Int32Array::from(vec![ + Some(10), + Some(11), + None, + Some(13), + Some(14), + Some(15), + None, + Some(17), + Some(18), + Some(19), + Some(20), + Some(21), + ])) as ArrayRef; + + let tag_values = Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(5), + Some(6), + None, + Some(8), + Some(9), + Some(11), + Some(12), + ])) as ArrayRef; + let tags = Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + offsets(vec![0, 2, 2, 2, 2, 2, 3, 6, 6, 6, 7, 7, 9]), + tag_values, + validity(vec![ + true, true, false, true, true, true, true, true, true, true, false, true, + ]), + ) + .unwrap(), + ) as ArrayRef; + + let pair_values = Arc::new(Int32Array::from(vec![ + Some(0), + None, + Some(2), + Some(3), + None, + Some(5), + Some(6), + Some(7), + Some(8), + Some(9), + Some(10), + None, + Some(12), + Some(13), + Some(14), + Some(15), + Some(16), + Some(17), + None, + Some(19), + Some(20), + Some(21), + Some(22), + None, + ])) as ArrayRef; + let pair = Arc::new( + FixedSizeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + 2, + pair_values, + validity(vec![ + true, true, true, false, true, true, true, true, true, true, false, true, + ]), + ) + .unwrap(), + ) as ArrayRef; + + let event_struct = Arc::new(StructArray::new( + Fields::from(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("tags", tags.data_type().clone(), true), + ArrowField::new("pair", pair.data_type().clone(), true), + ]), + vec![id_values, tags, pair], + validity(vec![ + true, false, true, true, false, true, true, true, false, true, true, true, + ]), + )) as ArrayRef; + + let events = Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new( + "item", + event_struct.data_type().clone(), + true, + )), + offsets(event_offsets), + event_struct, + validity(event_validity), + ) + .unwrap(), + ) as ArrayRef; + + Arc::new(StructArray::new( + Fields::from(vec![ArrowField::new( + "events", + events.data_type().clone(), + true, + )]), + vec![events], + validity(vec![ + true, true, true, false, true, true, true, true, true, true, true, false, true, + true, true, true, + ]), + )) + } + + fn sparse_fsl_struct_coverage_array() -> ArrayRef { + use arrow_array::{FixedSizeListArray, Int32Array, ListArray, StructArray}; + use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::Fields; + + fn offsets(values: Vec) -> OffsetBuffer { + OffsetBuffer::new(ScalarBuffer::from(values)) + } + + fn validity(values: Vec) -> Option { + Some(NullBuffer::new(BooleanBuffer::from(values))) + } + + let values = Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + Some(4), + Some(5), + Some(6), + None, + Some(8), + Some(9), + Some(10), + ])) as ArrayRef; + let lists = Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + offsets(vec![0, 2, 2, 2, 3, 3, 4, 4, 6, 6, 7, 7, 7, 8, 8, 10, 10]), + values, + validity(vec![ + true, true, false, true, true, true, true, true, true, false, true, true, true, + true, true, true, + ]), + ) + .unwrap(), + ) as ArrayRef; + let item_struct = Arc::new(StructArray::new( + Fields::from(vec![ArrowField::new( + "values", + lists.data_type().clone(), + true, + )]), + vec![lists], + validity(vec![ + true, true, true, false, true, true, true, true, false, true, true, true, true, + true, true, true, + ]), + )) as ArrayRef; + + Arc::new( + FixedSizeListArray::try_new( + Arc::new(ArrowField::new( + "item", + item_struct.data_type().clone(), + true, + )), + 2, + item_struct, + validity(vec![true, true, false, true, true, true, false, true]), + ) + .unwrap(), + ) + } + fn sparse_test_layout() -> crate::format::pb21::SparseLayout { crate::format::pb21::SparseLayout { value_compression: Some(ProtobufUtils21::flat(32, None)), @@ -9612,6 +9930,7 @@ mod tests { let err = SparseStructuralScheduler::try_new( &[(0, 0), (0, 0)], 0, + layout.num_items, DataType::Int32, &layout, &decompressors, @@ -9629,6 +9948,7 @@ mod tests { let err = SparseStructuralScheduler::try_new( &[(0, 0)], 0, + layout.num_items, DataType::Int32, &layout, &decompressors, @@ -9660,6 +9980,7 @@ mod tests { let err = SparseStructuralScheduler::try_new( &[(0, 8), (8, 8)], 0, + layout.num_items, DataType::Int32, &layout, &decompressors, @@ -9676,6 +9997,7 @@ mod tests { let mut scheduler = SparseStructuralScheduler::try_new( &[(0, 8), (8, 8)], 0, + layout.num_items, DataType::Int32, &layout, &decompressors, @@ -10010,6 +10332,103 @@ mod tests { check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; } + #[tokio::test] + async fn test_sparse_layout_nested_structural_states_roundtrip() { + let arr = sparse_nested_coverage_array(); + let sparse_pages = + encode_pages_with_structural_mode(arr.clone(), STRUCTURAL_ENCODING_SPARSE).await; + let sparse_layouts = sparse_layouts_from_pages(&sparse_pages); + assert!( + !sparse_layouts.is_empty(), + "expected forced sparse encoding to emit sparse pages" + ); + + let has_variable_list_layer = sparse_layouts.iter().any(|sparse| { + sparse.structural_layers.iter().any(|layer| { + matches!( + pb21::sparse_structural_layer::Kind::try_from(layer.kind), + Ok(pb21::sparse_structural_layer::Kind::SparseLayerList) + ) && layer.num_non_empty > 0 + && layer.num_nulls > 0 + && layer.count_compression.is_some() + }) + }); + assert!( + has_variable_list_layer, + "expected a sparse list layer with non-empty positions, counts, and null positions" + ); + + let has_validity_layer = sparse_layouts.iter().any(|sparse| { + sparse.structural_layers.iter().any(|layer| { + matches!( + pb21::sparse_structural_layer::Kind::try_from(layer.kind), + Ok(pb21::sparse_structural_layer::Kind::SparseLayerValidity) + ) && layer.num_nulls > 0 + }) + }); + assert!( + has_validity_layer, + "expected a sparse validity layer for nullable struct or primitive values" + ); + + let metadata = HashMap::from([( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_SPARSE.to_string(), + )]); + let test_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_range(0..16) + .with_range(1..5) + .with_range(2..8) + .with_range(10..16) + .with_indices(vec![0, 1, 2, 3, 4, 5, 6, 7]) + .with_indices(vec![2]) + .with_indices(vec![3, 11, 15]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; + } + + #[tokio::test] + async fn test_sparse_layout_fixed_size_list_struct_roundtrip() { + let arr = sparse_fsl_struct_coverage_array(); + let sparse_pages = + encode_pages_with_structural_mode(arr.clone(), STRUCTURAL_ENCODING_SPARSE).await; + let sparse_layouts = sparse_layouts_from_pages(&sparse_pages); + assert!( + !sparse_layouts.is_empty(), + "expected forced sparse encoding to emit sparse pages" + ); + + let has_fixed_size_list_layer = sparse_layouts.iter().any(|sparse| { + sparse.structural_layers.iter().any(|layer| { + matches!( + pb21::sparse_structural_layer::Kind::try_from(layer.kind), + Ok(pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList) + ) && layer.constant_count == 2 + && layer.num_nulls > 0 + }) + }); + assert!( + has_fixed_size_list_layer, + "expected a sparse fixed-size-list layer with dimension and null positions" + ); + + let metadata = HashMap::from([( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_SPARSE.to_string(), + )]); + let test_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_range(0..8) + .with_range(1..4) + .with_range(5..8) + .with_indices(vec![0, 1, 2, 3, 6, 7]) + .with_indices(vec![2]) + .with_indices(vec![6, 7]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; + } + async fn encode_first_page( field: arrow_schema::Field, array: ArrayRef, diff --git a/rust/lance-encoding/src/repdef.rs b/rust/lance-encoding/src/repdef.rs index afdc534eb90..67f32c1185a 100644 --- a/rust/lance-encoding/src/repdef.rs +++ b/rust/lance-encoding/src/repdef.rs @@ -1122,7 +1122,7 @@ impl SerializerContext { debug_assert!( self.current_len == 0 || self.current_len == validity.len() + self.current_num_specials ); - self.current_len = validity.len(); + self.current_len = validity.len() + self.current_num_specials; let mut def_read_itr = self.def_levels.iter().copied(); let mut def_write_itr = self.spare_def.iter_mut(); From ec367bb70735e0e2262aa80a7ad70638bc6a0066 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 4 Jul 2026 13:44:10 +0800 Subject: [PATCH 5/8] fix: auto sparse for unsplittable structural pages --- docs/src/format/file/encoding.md | 5 +- .../src/encodings/logical/primitive.rs | 119 ++++++++++++++++-- 2 files changed, 113 insertions(+), 11 deletions(-) diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index 1a74a89977f..52cd9d86e61 100644 --- a/docs/src/format/file/encoding.md +++ b/docs/src/format/file/encoding.md @@ -330,9 +330,8 @@ have per offset (for variable-width data). ### Sparse Page Layout The sparse page layout is a 2.3+ structural layout for sparse nested data. It is selected explicitly with -`lance-encoding:structural-encoding=sparse`, or automatically in 2.3+ when the writer would otherwise need to split a -page because the structural information is much larger than the visible value stream. Writers must not emit this layout -for file versions before 2.3. +`lance-encoding:structural-encoding=sparse`, or automatically in 2.3+ when dense mini-block structural information +would exceed the mini-block structural page budget. Writers must not emit this layout for file versions before 2.3. Sparse layout stores Arrow structure as native slot-domain mappings instead of dense repetition / definition events. The value buffers are still mini-block compressed. The structural layers are stored from outer-most to inner-most: diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index c33fe0dd53b..9c1b8e2ccac 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -5268,6 +5268,25 @@ impl PrimitiveStructuralEncoder { version.resolve() >= LanceFileVersion::V2_3 } + fn structural_plan_needs_sparse(plan: &StructuralPagePlan) -> bool { + matches!( + plan, + StructuralPagePlan::Split(_) | StructuralPagePlan::UnsplittableOverBudget(_) + ) + } + + fn should_auto_sparse_layout( + version: LanceFileVersion, + requested_encoding: Option<&str>, + structural_plan: &StructuralPagePlan, + sparse_structural_plan: Option<&SparseStructuralPlan>, + ) -> bool { + Self::supports_sparse_layout(version) + && requested_encoding.is_none() + && Self::structural_plan_needs_sparse(structural_plan) + && sparse_structural_plan.is_some() + } + pub fn try_new( options: &EncodingOptions, compression_strategy: Arc, @@ -7285,15 +7304,17 @@ impl PrimitiveStructuralEncoder { let requested_encoding = encoding_metadata .get(STRUCTURAL_ENCODING_META_KEY) .map(|requested| requested.to_lowercase()); - let supports_sparse_layout = Self::supports_sparse_layout(version); let should_encode_sparse = requested_encoding.as_deref() == Some(STRUCTURAL_ENCODING_SPARSE) - || (supports_sparse_layout - && requested_encoding.is_none() - && matches!(&structural_plan, StructuralPagePlan::Split(_))); + || Self::should_auto_sparse_layout( + version, + requested_encoding.as_deref(), + &structural_plan, + sparse_structural_plan.as_ref(), + ); if requested_encoding.as_deref() == Some(STRUCTURAL_ENCODING_SPARSE) - && !supports_sparse_layout + && !Self::supports_sparse_layout(version) { return Err(Error::invalid_input_source( format!( @@ -7605,11 +7626,27 @@ impl PrimitiveStructuralEncoder { )); } + if requested_encoding.is_none() + && Self::supports_sparse_layout(self.version) + && matches!( + structural_plan, + StructuralPagePlan::UnsplittableOverBudget(_) + ) + && sparse_structural_plan.is_none() + { + return Err(Error::invalid_input_source( + "Mini-block cannot split the structural page and sparse structural encoding requires native structural layers".into(), + )); + } + let should_use_sparse_plan = requested_encoding.as_deref() == Some(STRUCTURAL_ENCODING_SPARSE) - || (Self::supports_sparse_layout(self.version) - && requested_encoding.is_none() - && matches!(&structural_plan, StructuralPagePlan::Split(_))); + || Self::should_auto_sparse_layout( + self.version, + requested_encoding.as_deref(), + &structural_plan, + sparse_structural_plan.as_ref(), + ); let pages = if should_use_sparse_plan { let unsplittable_miniblock_levels = match &structural_plan { StructuralPagePlan::UnsplittableOverBudget(num_levels) => Some(*num_levels), @@ -9621,6 +9658,42 @@ mod tests { Arc::new(builder.finish()) } + fn unsplittable_sparse_i32_nested_list_array(num_inner_lists: usize) -> ArrayRef { + use arrow_array::{Int32Array, ListArray}; + use arrow_buffer::{OffsetBuffer, ScalarBuffer}; + + let mut inner_offsets = Vec::with_capacity(num_inner_lists + 1); + let mut value_count = 0_i32; + inner_offsets.push(value_count); + for inner_idx in 0..num_inner_lists { + if inner_idx + 1 == num_inner_lists { + value_count += 1; + } + inner_offsets.push(value_count); + } + + let values = Arc::new(Int32Array::from(vec![42])) as ArrayRef; + let inner = Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + OffsetBuffer::new(ScalarBuffer::from(inner_offsets)), + values, + None, + ) + .unwrap(), + ) as ArrayRef; + + Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new("item", inner.data_type().clone(), true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, num_inner_lists as i32])), + inner, + None, + ) + .unwrap(), + ) + } + fn sparse_layout_from_page( page: &crate::encoder::EncodedPage, ) -> &crate::format::pb21::SparseLayout { @@ -10095,6 +10168,36 @@ mod tests { check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; } + #[tokio::test] + async fn test_sparse_layout_auto_for_unsplittable_overbudget_structural_page() { + let arr = unsplittable_sparse_i32_nested_list_array(70_000); + let field = arrow_schema::Field::new("c", arr.data_type().clone(), true); + + let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_3).await; + let sparse = sparse_layout_from_page(&page); + assert_eq!(sparse.num_visible_items, 1); + assert!( + sparse.num_items > u16::MAX as u64, + "expected unsplittable dense rep/def stream, got {} levels", + sparse.num_items + ); + let list_layers = sparse_list_layers(sparse); + assert_eq!(list_layers.len(), 2); + assert_eq!(list_layers[0].num_slots, 1); + assert_eq!(list_layers[0].num_non_empty, 1); + assert_eq!(list_layers[0].constant_count, 70_000); + assert_eq!(list_layers[1].num_slots, 70_000); + assert_eq!(list_layers[1].num_non_empty, 1); + assert_eq!(list_layers[1].constant_count, 1); + + let test_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_range(0..1) + .with_indices(vec![0]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; + } + fn deep_sparse_struct_array(num_rows: usize, stride: usize) -> ArrayRef { use arrow_array::{Int32Array, ListArray, StructArray}; use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; From d42c8fbb73f93ddb0eae0e3d1a9ad9ee521f2b75 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 4 Jul 2026 16:57:31 +0800 Subject: [PATCH 6/8] refactor: rename mini-block repdef budget --- .../src/encodings/logical/primitive.rs | 131 +++++++++--------- rust/lance-encoding/src/repdef.rs | 73 +++++----- 2 files changed, 100 insertions(+), 104 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 9c1b8e2ccac..8d14e853bd7 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -61,8 +61,8 @@ use crate::{ use crate::{ repdef::{ CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation, - RepDefSlicer, SerializedRepDefs, SparseStructuralLayerPlan, SparseStructuralPlan, - StructuralPagePlan, build_control_word_iterator, + MiniBlockRepDefBudget, RepDefSlicer, SerializedRepDefs, SparseStructuralLayerPlan, + SparseStructuralPlan, build_control_word_iterator, }, utils::accumulation::AccumulationQueue, }; @@ -5231,10 +5231,10 @@ struct PrimitivePageData { row_number: u64, // Number of top-level rows in this page. num_rows: u64, - // Present when one top-level row is too large for one miniblock rep/def chunk. - unsplittable_miniblock_levels: Option, - // Structural split plan from the original accumulated page. - structural_plan: StructuralPagePlan, + // Present when one top-level row is too large for one mini-block rep/def page. + single_row_miniblock_repdef_levels: Option, + // Dense mini-block rep/def budget result from the original accumulated page. + miniblock_repdef_budget: MiniBlockRepDefBudget, // Native sparse structural mapping for the original accumulated page. sparse_structural_plan: Option, } @@ -5268,22 +5268,23 @@ impl PrimitiveStructuralEncoder { version.resolve() >= LanceFileVersion::V2_3 } - fn structural_plan_needs_sparse(plan: &StructuralPagePlan) -> bool { + fn miniblock_repdef_budget_needs_sparse(budget: &MiniBlockRepDefBudget) -> bool { matches!( - plan, - StructuralPagePlan::Split(_) | StructuralPagePlan::UnsplittableOverBudget(_) + budget, + MiniBlockRepDefBudget::RequiresPageSplit(_) + | MiniBlockRepDefBudget::SingleRowOverBudget(_) ) } fn should_auto_sparse_layout( version: LanceFileVersion, requested_encoding: Option<&str>, - structural_plan: &StructuralPagePlan, + miniblock_repdef_budget: &MiniBlockRepDefBudget, sparse_structural_plan: Option<&SparseStructuralPlan>, ) -> bool { Self::supports_sparse_layout(version) && requested_encoding.is_none() - && Self::structural_plan_needs_sparse(structural_plan) + && Self::miniblock_repdef_budget_needs_sparse(miniblock_repdef_budget) && sparse_structural_plan.is_some() } @@ -7164,55 +7165,50 @@ impl PrimitiveStructuralEncoder { Ok(sliced) } - fn split_structural_pages_for_miniblock_budget( + fn split_pages_for_miniblock_repdef_budget( arrays: Vec, repdef: SerializedRepDefs, - plan: StructuralPagePlan, + budget: MiniBlockRepDefBudget, row_number: u64, num_rows: u64, ) -> Result> { - if plan == StructuralPagePlan::Fits { - return Ok(vec![PrimitivePageData { + match budget { + MiniBlockRepDefBudget::WithinBudget => Ok(vec![PrimitivePageData { arrays, repdef, row_number, num_rows, - unsplittable_miniblock_levels: None, - structural_plan: StructuralPagePlan::Fits, + single_row_miniblock_repdef_levels: None, + miniblock_repdef_budget: MiniBlockRepDefBudget::WithinBudget, sparse_structural_plan: None, - }]); - } - if let StructuralPagePlan::UnsplittableOverBudget(num_levels) = plan { - return Ok(vec![PrimitivePageData { + }]), + MiniBlockRepDefBudget::SingleRowOverBudget(num_levels) => Ok(vec![PrimitivePageData { arrays, repdef, row_number, num_rows, - unsplittable_miniblock_levels: Some(num_levels), - structural_plan: StructuralPagePlan::UnsplittableOverBudget(num_levels), - sparse_structural_plan: None, - }]); - } - - let StructuralPagePlan::Split(splits) = plan else { - unreachable!(); - }; - - let mut pages = Vec::with_capacity(splits.len()); - for split in splits { - let arrays = Self::slice_arrays(&arrays, split.value_start, split.num_values)?; - let repdef = Self::slice_repdef(&repdef, split.level_range); - pages.push(PrimitivePageData { - arrays, - repdef, - row_number: row_number + split.row_start, - num_rows: split.num_rows, - unsplittable_miniblock_levels: None, - structural_plan: StructuralPagePlan::Fits, + single_row_miniblock_repdef_levels: Some(num_levels), + miniblock_repdef_budget: MiniBlockRepDefBudget::SingleRowOverBudget(num_levels), sparse_structural_plan: None, - }); + }]), + MiniBlockRepDefBudget::RequiresPageSplit(splits) => { + let mut pages = Vec::with_capacity(splits.len()); + for split in splits { + let arrays = Self::slice_arrays(&arrays, split.value_start, split.num_values)?; + let repdef = Self::slice_repdef(&repdef, split.level_range); + pages.push(PrimitivePageData { + arrays, + repdef, + row_number: row_number + split.row_start, + num_rows: split.num_rows, + single_row_miniblock_repdef_levels: None, + miniblock_repdef_budget: MiniBlockRepDefBudget::WithinBudget, + sparse_structural_plan: None, + }); + } + Ok(pages) + } } - Ok(pages) } fn encode_page(ctx: PrimitiveEncodeContext, page: PrimitivePageData) -> Result { @@ -7231,8 +7227,8 @@ impl PrimitiveStructuralEncoder { repdef, row_number, num_rows, - unsplittable_miniblock_levels, - structural_plan, + single_row_miniblock_repdef_levels, + miniblock_repdef_budget, sparse_structural_plan, } = page; let num_values = arrays.iter().map(|arr| arr.len() as u64).sum(); @@ -7309,7 +7305,7 @@ impl PrimitiveStructuralEncoder { || Self::should_auto_sparse_layout( version, requested_encoding.as_deref(), - &structural_plan, + &miniblock_repdef_budget, sparse_structural_plan.as_ref(), ); @@ -7406,7 +7402,7 @@ impl PrimitiveStructuralEncoder { ); } - if let Some(num_levels) = unsplittable_miniblock_levels { + if let Some(num_levels) = single_row_miniblock_repdef_levels { let fullzip_error = match &data_block { DataBlock::FixedWidth(fixed) if !fixed.bits_per_value.is_multiple_of(8) => { Some(format!( @@ -7603,12 +7599,13 @@ impl PrimitiveStructuralEncoder { } else { None }; - let (repdef, structural_plan) = RepDefBuilder::serialize_with_structural_plan( - repdefs, - miniblock::max_repdef_levels_per_chunk, - num_rows, - num_values, - )?; + let (repdef, miniblock_repdef_budget) = + RepDefBuilder::serialize_with_miniblock_repdef_budget( + repdefs, + miniblock::max_repdef_levels_per_chunk, + num_rows, + num_values, + )?; let requested_encoding = self .encoding_metadata .get(STRUCTURAL_ENCODING_META_KEY) @@ -7629,8 +7626,8 @@ impl PrimitiveStructuralEncoder { if requested_encoding.is_none() && Self::supports_sparse_layout(self.version) && matches!( - structural_plan, - StructuralPagePlan::UnsplittableOverBudget(_) + miniblock_repdef_budget, + MiniBlockRepDefBudget::SingleRowOverBudget(_) ) && sparse_structural_plan.is_none() { @@ -7644,12 +7641,12 @@ impl PrimitiveStructuralEncoder { || Self::should_auto_sparse_layout( self.version, requested_encoding.as_deref(), - &structural_plan, + &miniblock_repdef_budget, sparse_structural_plan.as_ref(), ); let pages = if should_use_sparse_plan { - let unsplittable_miniblock_levels = match &structural_plan { - StructuralPagePlan::UnsplittableOverBudget(num_levels) => Some(*num_levels), + let single_row_miniblock_repdef_levels = match &miniblock_repdef_budget { + MiniBlockRepDefBudget::SingleRowOverBudget(num_levels) => Some(*num_levels), _ => None, }; vec![PrimitivePageData { @@ -7657,15 +7654,15 @@ impl PrimitiveStructuralEncoder { repdef, row_number, num_rows, - unsplittable_miniblock_levels, - structural_plan, + single_row_miniblock_repdef_levels, + miniblock_repdef_budget, sparse_structural_plan, }] } else { - Self::split_structural_pages_for_miniblock_budget( + Self::split_pages_for_miniblock_repdef_budget( arrays, repdef, - structural_plan, + miniblock_repdef_budget, row_number, num_rows, )? @@ -9658,7 +9655,7 @@ mod tests { Arc::new(builder.finish()) } - fn unsplittable_sparse_i32_nested_list_array(num_inner_lists: usize) -> ArrayRef { + fn single_row_overbudget_sparse_i32_nested_list_array(num_inner_lists: usize) -> ArrayRef { use arrow_array::{Int32Array, ListArray}; use arrow_buffer::{OffsetBuffer, ScalarBuffer}; @@ -10169,8 +10166,8 @@ mod tests { } #[tokio::test] - async fn test_sparse_layout_auto_for_unsplittable_overbudget_structural_page() { - let arr = unsplittable_sparse_i32_nested_list_array(70_000); + async fn test_sparse_layout_auto_for_single_row_overbudget_repdef_page() { + let arr = single_row_overbudget_sparse_i32_nested_list_array(70_000); let field = arrow_schema::Field::new("c", arr.data_type().clone(), true); let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_3).await; @@ -10178,7 +10175,7 @@ mod tests { assert_eq!(sparse.num_visible_items, 1); assert!( sparse.num_items > u16::MAX as u64, - "expected unsplittable dense rep/def stream, got {} levels", + "expected single-row over-budget dense rep/def stream, got {} levels", sparse.num_items ); let list_layers = sparse_list_layers(sparse); diff --git a/rust/lance-encoding/src/repdef.rs b/rust/lance-encoding/src/repdef.rs index 67f32c1185a..6a14fabb805 100644 --- a/rust/lance-encoding/src/repdef.rs +++ b/rust/lance-encoding/src/repdef.rs @@ -561,9 +561,9 @@ impl SparseStructuralUnraveler { } } -/// A contiguous top-level-row range that can be encoded as one structural page. +/// A top-level-row range whose dense rep/def stream fits one mini-block page. #[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct StructuralPageSplit { +pub(crate) struct MiniBlockRepDefSplit { /// Top-level row offset, relative to the original unsplit page. pub(crate) row_start: u64, /// Number of top-level rows in this split. @@ -576,15 +576,15 @@ pub(crate) struct StructuralPageSplit { pub(crate) num_values: u64, } -/// Planner result for structural page budget handling. +/// Dense mini-block rep/def budget result for one accumulated page. #[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) enum StructuralPagePlan { - /// The original page can be encoded as-is. - Fits, - /// The original page should be split on top-level row boundaries. - Split(Vec), - /// One top-level row is larger than the requested structural page budget. - UnsplittableOverBudget(u64), +pub(crate) enum MiniBlockRepDefBudget { + /// The dense rep/def stream fits one mini-block structural page. + WithinBudget, + /// The dense rep/def stream fits after splitting on top-level row boundaries. + RequiresPageSplit(Vec), + /// A single top-level row has this many rep/def levels and exceeds the budget. + SingleRowOverBudget(u64), } // As we build def levels we add this to special values to indicate that they @@ -1245,21 +1245,20 @@ impl SerializerContext { max_levels_per_page: Option, num_rows: u64, num_values: u64, - ) -> Result { + ) -> Result { // Extremely sparse lists can have many rep/def levels for very few // visible leaf values. If this ratio becomes too skewed then a - // miniblock structural chunk can exceed its packed rep/def metadata - // budget even though the value buffers are small. We detect that case - // while normalizing special def levels and split the structural page on - // top-level row boundaries so each emitted page stays within the - // miniblock structural budget. + // mini-block rep/def chunk can exceed its packed metadata budget even + // though the value buffers are small. We detect that case while + // normalizing special def levels and split on top-level row boundaries + // so each emitted dense mini-block page stays within the budget. if self.def_levels.is_empty() { - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); } if self.rep_levels.is_empty() { self.normalize_specials(); - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); } if self.rep_levels.len() != self.def_levels.len() { @@ -1272,12 +1271,12 @@ impl SerializerContext { let Some(max_levels_per_page) = max_levels_per_page else { self.normalize_specials(); - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); }; if num_values == 0 { self.normalize_specials(); - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); } let max_schema_rep = def_meaning.iter().filter(|level| level.is_list()).count() as u16; @@ -1286,7 +1285,7 @@ impl SerializerContext { if !should_plan { self.normalize_specials(); - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); } let max_visible_level = max_visible_level.unwrap(); @@ -1294,7 +1293,7 @@ impl SerializerContext { let mut counted_rows = 0u64; let mut counted_values = 0u64; let mut saw_structural_overhead = false; - let mut unsplittable_over_budget = None; + let mut single_row_over_budget_levels = None; let mut current_row_level_start = None; let mut current_row_num_values = 0u64; @@ -1315,14 +1314,14 @@ impl SerializerContext { saw_structural_overhead |= row_has_structural_overhead; if row_has_structural_overhead && row_num_levels > max_levels_per_page { - unsplittable_over_budget = Some(row_num_levels); + single_row_over_budget_levels = Some(row_num_levels); } if current_page_num_rows > 0 && (current_page_has_structural_overhead || row_has_structural_overhead) && current_page_num_levels + row_num_levels > max_levels_per_page { - splits.push(StructuralPageSplit { + splits.push(MiniBlockRepDefSplit { row_start: current_page_row_start, num_rows: current_page_num_rows, level_range: current_page_level_start..current_page_level_end, @@ -1405,14 +1404,14 @@ impl SerializerContext { ))); } if !saw_structural_overhead { - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); } - if let Some(row_num_levels) = unsplittable_over_budget { - return Ok(StructuralPagePlan::UnsplittableOverBudget(row_num_levels)); + if let Some(row_num_levels) = single_row_over_budget_levels { + return Ok(MiniBlockRepDefBudget::SingleRowOverBudget(row_num_levels)); } if current_page_num_rows > 0 { - splits.push(StructuralPageSplit { + splits.push(MiniBlockRepDefSplit { row_start: current_page_row_start, num_rows: current_page_num_rows, level_range: current_page_level_start..current_page_level_end, @@ -1422,9 +1421,9 @@ impl SerializerContext { } if splits.len() > 1 { - Ok(StructuralPagePlan::Split(splits)) + Ok(MiniBlockRepDefBudget::RequiresPageSplit(splits)) } else { - Ok(StructuralPagePlan::Fits) + Ok(MiniBlockRepDefBudget::WithinBudget) } } @@ -1462,12 +1461,12 @@ impl SerializerContext { ) } - fn build_with_structural_plan( + fn build_with_miniblock_repdef_budget( mut self, max_levels_per_page: Option, num_rows: u64, num_values: u64, - ) -> Result<(SerializedRepDefs, StructuralPagePlan)> { + ) -> Result<(SerializedRepDefs, MiniBlockRepDefBudget)> { if self.current_len == 0 { return Ok(( SerializedRepDefs::new_with_fixed_size_list_levels( @@ -1476,7 +1475,7 @@ impl SerializerContext { self.def_meaning, self.has_fsl, ), - StructuralPagePlan::Fits, + MiniBlockRepDefBudget::WithinBudget, )); } @@ -1851,15 +1850,15 @@ impl RepDefBuilder { Self::serialize_builders(builders).0.build() } - /// Converts gathered structural buffers into rep/def levels and an encode-time plan. - pub(crate) fn serialize_with_structural_plan( + /// Converts gathered structural buffers into rep/def levels and a mini-block budget result. + pub(crate) fn serialize_with_miniblock_repdef_budget( builders: Vec, max_levels_for_bits: impl FnOnce(u64) -> u64, num_rows: u64, num_values: u64, - ) -> Result<(SerializedRepDefs, StructuralPagePlan)> { + ) -> Result<(SerializedRepDefs, MiniBlockRepDefBudget)> { let (context, bits_per_level) = Self::serialize_builders(builders); - context.build_with_structural_plan( + context.build_with_miniblock_repdef_budget( bits_per_level.map(max_levels_for_bits), num_rows, num_values, From 0d18eb03d76163bb690cdc57b0a234c43750ebad Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 4 Jul 2026 17:08:38 +0800 Subject: [PATCH 7/8] fix: harden sparse layout validation --- .../src/encodings/logical/primitive.rs | 429 +++++++++++++++--- 1 file changed, 361 insertions(+), 68 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 8d14e853bd7..52a93acc749 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -860,11 +860,51 @@ impl SparseStructuralScheduler { Ok(1) } + fn reject_optional_buffer( + compression: &Option, + label: &str, + ) -> Result<()> { + if compression.is_some() { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} is not valid for this layer").into(), + )); + } + Ok(()) + } + + fn checked_buffer_range(position: u64, size: u64, label: &str) -> Result> { + let end = position.checked_add(size).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} buffer range overflows").into(), + ) + })?; + Ok(position..end) + } + fn structural_buffer_count(layers: &[pb21::SparseStructuralLayer]) -> Result { layers.iter().try_fold(0_usize, |mut count, layer| { let kind = pb21::sparse_structural_layer::Kind::try_from(layer.kind)?; match kind { pb21::sparse_structural_layer::Kind::SparseLayerValidity => { + if layer.num_child_slots != layer.num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural validity child slots {} do not match slots {}", + layer.num_child_slots, layer.num_slots + ) + .into(), + )); + } + if layer.constant_count != 0 || layer.num_non_empty != 0 { + return Err(Error::invalid_input_source( + "Sparse structural validity layer has list-only metadata".into(), + )); + } + Self::reject_optional_buffer( + &layer.non_empty_compression, + "validity non-empty positions", + )?; + Self::reject_optional_buffer(&layer.count_compression, "validity counts")?; count += Self::count_optional_buffer( &layer.null_compression, layer.num_nulls, @@ -885,10 +925,46 @@ impl SparseStructuralScheduler { "Sparse structural list has non-empty positions but no count compression or constant count".into(), )); } - if layer.count_compression.is_some() { - let compression = layer.count_compression.as_ref().expect_ok()?; - Self::validate_compression(compression, "list counts")?; - count += 1; + if layer.num_non_empty == 0 && layer.constant_count != 0 { + return Err(Error::invalid_input_source( + "Sparse structural list has constant count but no non-empty positions" + .into(), + )); + } + if layer.count_compression.is_some() && layer.constant_count != 0 { + return Err(Error::invalid_input_source( + "Sparse structural list has both count compression and constant count" + .into(), + )); + } + if layer.constant_count != 0 { + let expected_child_slots = layer + .num_non_empty + .checked_mul(layer.constant_count) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural list child slot count overflows: non_empty={}, constant_count={}", + layer.num_non_empty, layer.constant_count + ) + .into(), + ) + })?; + if expected_child_slots != layer.num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list child slot count {} does not match non-empty slots {} * constant count {}", + layer.num_child_slots, layer.num_non_empty, layer.constant_count + ) + .into(), + )); + } + } else { + count += Self::count_optional_buffer( + &layer.count_compression, + layer.num_non_empty, + "list counts", + )?; } count += Self::count_optional_buffer( &layer.null_compression, @@ -897,6 +973,20 @@ impl SparseStructuralScheduler { )?; } pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList => { + if layer.num_non_empty != 0 { + return Err(Error::invalid_input_source( + "Sparse structural fixed-size-list layer has list-only metadata" + .into(), + )); + } + Self::reject_optional_buffer( + &layer.non_empty_compression, + "fixed-size-list non-empty positions", + )?; + Self::reject_optional_buffer( + &layer.count_compression, + "fixed-size-list counts", + )?; let expected_child_slots = layer.num_slots.checked_mul(layer.constant_count).ok_or_else(|| { Error::invalid_input_source( @@ -1085,7 +1175,49 @@ impl SparseStructuralScheduler { .into(), )); } - Ok(fixed.data.borrow_to_typed_slice::().to_vec()) + if fixed.num_values != num_values { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} decoded {} values, expected {}", + fixed.num_values, num_values + ) + .into(), + )); + } + let num_values_usize = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} value count exceeds usize::MAX").into(), + ) + })?; + let expected_len = num_values_usize + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} decoded byte length overflows").into(), + ) + })?; + if fixed.data.len() != expected_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} decoded {} bytes, expected {}", + fixed.data.len(), + expected_len + ) + .into(), + )); + } + let values = fixed.data.borrow_to_typed_slice::(); + if values.len() != num_values_usize { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} decoded {} u64 values, expected {}", + values.len(), + num_values + ) + .into(), + )); + } + Ok(values.to_vec()) } fn decode_positions( @@ -1380,11 +1512,21 @@ impl StructuralPageScheduler for SparseStructuralScheduler { Ok(buffer) => buffer, Err(err) => return std::future::ready(Err(err)).boxed(), }; - let mut required_ranges = Vec::new(); - required_ranges.push(meta_buf_position..meta_buf_position + meta_buf_size); - for (position, size) in self.buffer_offsets_and_sizes.iter().skip(2) { - required_ranges.push(*position..*position + *size); - } + let required_ranges = match (|| -> Result>> { + let mut required_ranges = Vec::new(); + required_ranges.push(Self::checked_buffer_range( + meta_buf_position, + meta_buf_size, + "metadata", + )?); + for (position, size) in self.buffer_offsets_and_sizes.iter().skip(2) { + required_ranges.push(Self::checked_buffer_range(*position, *size, "structural")?); + } + Ok(required_ranges) + })() { + Ok(ranges) => ranges, + Err(err) => return std::future::ready(Err(err)).boxed(), + }; let io_req = io.submit_request(required_ranges, 0); async move { @@ -5697,33 +5839,14 @@ impl PrimitiveStructuralEncoder { fn serialize_sparse_miniblocks( miniblocks: SparseMiniBlockCompressed, - rep: Option>, - def: Option>, support_large_chunk: bool, ) -> Result { - let bytes_rep = rep - .as_ref() - .map(|rep| rep.iter().map(|r| r.data.len()).sum::()) - .unwrap_or(0); - let bytes_def = def - .as_ref() - .map(|def| def.iter().map(|d| d.data.len()).sum::()) - .unwrap_or(0); let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::(); - let mut num_buffers = miniblocks.data.len(); - if rep.is_some() { - num_buffers += 1; - } - if def.is_some() { - num_buffers += 1; - } + let num_buffers = miniblocks.data.len(); let max_extra = 9 * num_buffers; - let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra); + let mut data_buffer = Vec::with_capacity(bytes_data + max_extra); let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * 8); - let mut rep_iter = rep.map(|r| r.into_iter()); - let mut def_iter = def.map(|d| d.into_iter()); - let mut buffer_offsets = vec![0; miniblocks.data.len()]; for chunk in miniblocks.chunks { if chunk.buffer_sizes.len() != miniblocks.data.len() { @@ -5737,33 +5860,7 @@ impl PrimitiveStructuralEncoder { let start_pos = data_buffer.len(); debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0); - let rep = rep_iter.as_mut().map(|r| r.next().unwrap()); - let def = def_iter.as_mut().map(|d| d.next().unwrap()); - - let num_levels = rep - .as_ref() - .map(|r| r.num_levels) - .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0)); - data_buffer.extend_from_slice(&num_levels.to_le_bytes()); - - if let Some(rep) = rep.as_ref() { - let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| { - Error::internal(format!( - "Sparse repetition buffer size ({} bytes) too large", - rep.data.len() - )) - })?; - data_buffer.extend_from_slice(&bytes_rep.to_le_bytes()); - } - if let Some(def) = def.as_ref() { - let bytes_def = u16::try_from(def.data.len()).map_err(|_| { - Error::internal(format!( - "Sparse definition buffer size ({} bytes) too large", - def.data.len() - )) - })?; - data_buffer.extend_from_slice(&bytes_def.to_le_bytes()); - } + data_buffer.extend_from_slice(&0_u16.to_le_bytes()); if support_large_chunk { for &buffer_size in &chunk.buffer_sizes { @@ -5787,14 +5884,6 @@ impl PrimitiveStructuralEncoder { }; add_padding(&mut data_buffer); - if let Some(rep) = rep.as_ref() { - data_buffer.extend_from_slice(&rep.data); - add_padding(&mut data_buffer); - } - if let Some(def) = def.as_ref() { - data_buffer.extend_from_slice(&def.data); - add_padding(&mut data_buffer); - } for (buffer_size, (buffer, buffer_offset)) in chunk .buffer_sizes .iter() @@ -6411,8 +6500,7 @@ impl PrimitiveStructuralEncoder { let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; let (compressed_data, value_encoding) = compressor.compress(data)?; let compressed_data = Self::sparse_miniblock_compressed(compressed_data)?; - let serialized = - Self::serialize_sparse_miniblocks(compressed_data, None, None, support_large_chunk)?; + let serialized = Self::serialize_sparse_miniblocks(compressed_data, support_large_chunk)?; let structural = Self::encode_sparse_structural_plan(&sparse_plan, compression_strategy)?; let description = ProtobufUtils21::sparse_layout( @@ -10060,6 +10148,211 @@ mod tests { assert_invalid_input_contains(err, "Sparse layout has 2 buffers, expected 3"); } + #[test] + fn test_sparse_layout_rejects_validity_child_slot_mismatch() { + let mut layout = sparse_test_layout(); + layout.structural_layers.push(pb21::SparseStructuralLayer { + kind: pb21::sparse_structural_layer::Kind::SparseLayerValidity as i32, + num_slots: 1, + num_child_slots: 2, + constant_count: 0, + non_empty_compression: None, + num_non_empty: 0, + count_compression: None, + null_compression: None, + num_nulls: 0, + }); + let decompressors = DefaultDecompressionStrategy::default(); + + let err = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + layout.num_items, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + + assert_invalid_input_contains(err, "validity child slots"); + } + + #[test] + fn test_sparse_layout_rejects_wrong_kind_structural_compression() { + let mut layout = sparse_test_layout(); + layout.structural_layers.push(pb21::SparseStructuralLayer { + kind: pb21::sparse_structural_layer::Kind::SparseLayerValidity as i32, + num_slots: 1, + num_child_slots: 1, + constant_count: 0, + non_empty_compression: Some(ProtobufUtils21::flat(64, None)), + num_non_empty: 0, + count_compression: None, + null_compression: None, + num_nulls: 0, + }); + let decompressors = DefaultDecompressionStrategy::default(); + + let err = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + layout.num_items, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + + assert_invalid_input_contains(err, "validity non-empty positions is not valid"); + } + + #[test] + fn test_sparse_layout_rejects_zero_count_list_count_compression() { + let mut layout = sparse_test_layout(); + layout.structural_layers.push(pb21::SparseStructuralLayer { + kind: pb21::sparse_structural_layer::Kind::SparseLayerList as i32, + num_slots: 1, + num_child_slots: 0, + constant_count: 0, + non_empty_compression: None, + num_non_empty: 0, + count_compression: Some(ProtobufUtils21::flat(64, None)), + null_compression: None, + num_nulls: 0, + }); + let decompressors = DefaultDecompressionStrategy::default(); + + let err = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + layout.num_items, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + + assert_invalid_input_contains(err, "list counts has compression but no values"); + } + + #[test] + fn test_sparse_layout_rejects_list_constant_and_count_compression() { + let mut layout = sparse_test_layout(); + layout.structural_layers.push(pb21::SparseStructuralLayer { + kind: pb21::sparse_structural_layer::Kind::SparseLayerList as i32, + num_slots: 1, + num_child_slots: 1, + constant_count: 1, + non_empty_compression: Some(ProtobufUtils21::flat(64, None)), + num_non_empty: 1, + count_compression: Some(ProtobufUtils21::flat(64, None)), + null_compression: None, + num_nulls: 0, + }); + let decompressors = DefaultDecompressionStrategy::default(); + + let err = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + layout.num_items, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + + assert_invalid_input_contains(err, "both count compression and constant count"); + } + + #[test] + fn test_sparse_layout_rejects_list_constant_child_slot_mismatch() { + let mut layout = sparse_test_layout(); + layout.structural_layers.push(pb21::SparseStructuralLayer { + kind: pb21::sparse_structural_layer::Kind::SparseLayerList as i32, + num_slots: 2, + num_child_slots: 3, + constant_count: 2, + non_empty_compression: Some(ProtobufUtils21::flat(64, None)), + num_non_empty: 2, + count_compression: None, + null_compression: None, + num_nulls: 0, + }); + let decompressors = DefaultDecompressionStrategy::default(); + + let err = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8), (16, 8)], + 0, + layout.num_items, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + + assert_invalid_input_contains(err, "list child slot count"); + } + + #[tokio::test] + async fn test_sparse_layout_rejects_short_structural_buffer() { + let mut layout = sparse_test_layout(); + layout.structural_layers.push(pb21::SparseStructuralLayer { + kind: pb21::sparse_structural_layer::Kind::SparseLayerValidity as i32, + num_slots: 1, + num_child_slots: 1, + constant_count: 0, + non_empty_compression: None, + num_non_empty: 0, + count_compression: None, + null_compression: Some(ProtobufUtils21::flat(64, None)), + num_nulls: 1, + }); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8), (16, 0)], + 0, + layout.num_items, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + + let mut data = Vec::new(); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&[0; 8]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + + let Err(err) = scheduler.initialize(&io).await else { + panic!("expected sparse layout to reject short structural buffer"); + }; + + assert_invalid_input_contains(err, "decoded 0 bytes, expected 8"); + } + + #[tokio::test] + async fn test_sparse_layout_rejects_metadata_buffer_range_overflow() { + let layout = sparse_test_layout(); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(u64::MAX, 1), (0, 0)], + 0, + layout.num_items, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::new())); + + let Err(err) = scheduler.initialize(&io).await else { + panic!("expected sparse layout to reject metadata buffer range overflow"); + }; + + assert_invalid_input_contains(err, "metadata buffer range overflows"); + } + #[tokio::test] async fn test_sparse_layout_rejects_value_metadata_past_value_buffer() { let layout = sparse_test_layout(); From 6489916aaf54c2ca9662fb7a806699089785e1aa Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sun, 5 Jul 2026 02:09:10 +0800 Subject: [PATCH 8/8] feat: use semantic sparse structural sets --- docs/src/format/file/encoding.md | 31 +- protos/encodings_v2_1.proto | 68 +- .../src/encodings/logical/list.rs | 12 +- .../src/encodings/logical/primitive.rs | 1634 ++++++++++++----- rust/lance-encoding/src/repdef.rs | 289 ++- 5 files changed, 1509 insertions(+), 525 deletions(-) diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index 52cd9d86e61..9060a310439 100644 --- a/docs/src/format/file/encoding.md +++ b/docs/src/format/file/encoding.md @@ -343,6 +343,19 @@ The value buffers are still mini-block compressed. The structural layers are sto Valid empty list slots are not stored in the non-empty positions. They are represented by their absence from both the non-empty positions and null positions for that list layer. +Sparse structural sets are semantic, not necessarily materialized position arrays. Position sets are encoded as one of: + +- empty: no slot is selected +- all: every slot in the layer domain is selected +- range: one contiguous slot range is selected +- explicit: a delta-compressed `u64` position buffer + +List counts are encoded as one of: + +- empty: there are no non-empty list slots +- constant: every non-empty list slot has the same child count +- explicit: a compressed `u64` count buffer + #### Selection and Compatibility Writers may emit sparse layout only for file versions 2.3 and later. Explicit sparse selection is requested with field @@ -369,10 +382,14 @@ Buffer 0 stores one 8-byte entry per value chunk. The first 4 bytes store the ch second 4 bytes store the number of visible values in the chunk. The sum of chunk value counts must match `SparseLayout.num_visible_items`, and the sum of chunk byte sizes must exactly match the value buffer size. -Structural position buffers store delta-encoded `u64` positions. Positions must be strictly increasing after delta -decoding and must be within the layer's parent slot domain. List count buffers store `u64` child counts for each -non-empty list slot. If all non-empty lists have the same length then `constant_count` is used instead of a count -buffer. +Structural position buffers are present only for explicit position sets. They store delta-encoded `u64` positions. +Positions must be strictly increasing after delta decoding and must be within the layer's parent slot domain. List count +buffers are present only for explicit count sets. They store `u64` child counts for each non-empty list slot. + +`SparseStructuralLayer.non_empty_positions`, `SparseStructuralLayer.counts`, and +`SparseStructuralLayer.null_positions` are the canonical structural representation. `LIST` layers require +`non_empty_positions`, `counts`, and `null_positions`; `VALIDITY` layers require only `null_positions`; +`FIXED_SIZE_LIST` layers require `null_positions` and `fixed_size_list_dimension`. #### Protobuf @@ -385,9 +402,9 @@ buffer. ``` The `value_compression` field is required. The physical page buffer list must match the structural layer metadata: -each non-zero position/count cardinality requires exactly one corresponding compressed structural buffer, and zero -cardinality layers must not include an unused compression descriptor. Readers should reject malformed sparse pages with -a format error. +each explicit position/count set requires exactly one corresponding compressed structural buffer, and empty/all/range or +constant sets must not include an unused compression descriptor. Readers should reject malformed sparse pages with a +format error. ### Constant Page Layout diff --git a/protos/encodings_v2_1.proto b/protos/encodings_v2_1.proto index 9a8b90b3a14..8b28da2c341 100644 --- a/protos/encodings_v2_1.proto +++ b/protos/encodings_v2_1.proto @@ -182,7 +182,7 @@ message SparseStructuralLayer { SPARSE_LAYER_VALIDITY = 1; // A list-like layer. Non-empty parent slots map to child ranges; missing valid slots are empty. SPARSE_LAYER_LIST = 2; - // A fixed-size-list layer. Child ranges are deterministic from constant_count. + // A fixed-size-list layer. Child ranges are deterministic from fixed_size_list_dimension. SPARSE_LAYER_FIXED_SIZE_LIST = 3; } @@ -192,19 +192,59 @@ message SparseStructuralLayer { // Number of slots in this layer's child domain after this layer is applied. For validity // layers this matches num_slots because Arrow child slot numbering is unchanged. uint64 num_child_slots = 3; - // For LIST, this is the child count for every non-empty slot when count_compression is absent. - // For FIXED_SIZE_LIST, this is the fixed list dimension. - uint64 constant_count = 4; - // Delta-compressed u64 positions for non-empty LIST slots. Absent for non-list layers. - CompressiveEncoding non_empty_compression = 5; - // Number of non-empty LIST positions. - uint64 num_non_empty = 6; - // Optional compressed u64 child counts for non-empty LIST slots. - CompressiveEncoding count_compression = 7; - // Optional delta-compressed u64 null positions for nullable layers. - CompressiveEncoding null_compression = 8; - // Number of null positions. - uint64 num_nulls = 9; + // Semantic representation for non-empty LIST slots. Required for LIST and absent otherwise. + SparsePositionSet non_empty_positions = 4; + // Semantic representation for LIST child counts. Required for LIST and absent otherwise. + SparseCountSet counts = 5; + // Semantic representation for null slots. Required for all sparse structural layers. + SparsePositionSet null_positions = 6; + // Fixed list dimension. Required for FIXED_SIZE_LIST and zero otherwise. + uint64 fixed_size_list_dimension = 7; +} + +message SparsePositionEmpty {} + +message SparsePositionAll {} + +message SparsePositionRange { + // First position in the layer slot domain. + uint64 start = 1; + // Number of positions in the contiguous range. + uint64 length = 2; +} + +message SparsePositionSet { + oneof positions { + // Delta-compressed u64 positions. Cardinality comes from num_positions. + CompressiveEncoding explicit = 1; + // A contiguous non-empty position range. + SparsePositionRange range = 2; + // Every slot in the layer domain is present. + SparsePositionAll all = 3; + // No slots are present. + SparsePositionEmpty empty = 4; + } + // Number of selected positions. For explicit sets this is the number of compressed u64 + // positions. For empty/all/range sets this must match the semantic cardinality. + uint64 num_positions = 5; +} + +message SparseCountEmpty {} + +message SparseCountConstant { + // Child count for every selected non-empty LIST slot. + uint64 value = 1; +} + +message SparseCountSet { + oneof counts { + // Compressed u64 child counts. Cardinality is the containing LIST layer's non-empty position count. + CompressiveEncoding explicit = 1; + // A single child count for every non-empty LIST slot. + SparseCountConstant constant = 2; + // No counts; only valid when the containing LIST layer has no non-empty positions. + SparseCountEmpty empty = 3; + } } // A layout used for pages where all (visible) values are the same scalar value. diff --git a/rust/lance-encoding/src/encodings/logical/list.rs b/rust/lance-encoding/src/encodings/logical/list.rs index 9b33060d27c..e7b1f5d8910 100644 --- a/rust/lance-encoding/src/encodings/logical/list.rs +++ b/rust/lance-encoding/src/encodings/logical/list.rs @@ -434,8 +434,16 @@ mod tests { crate::format::pb21::sparse_structural_layer::Kind::try_from(layer.kind), Ok(crate::format::pb21::sparse_structural_layer::Kind::SparseLayerList) ) { - num_non_empty_slots += layer.num_non_empty; - constant_counts.push(layer.constant_count); + num_non_empty_slots += layer + .non_empty_positions + .as_ref() + .map(|set| set.num_positions) + .unwrap_or(0); + if let Some(crate::format::pb21::sparse_count_set::Counts::Constant(constant)) = + layer.counts.as_ref().and_then(|set| set.counts.as_ref()) + { + constant_counts.push(constant.value); + } } } } diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 52a93acc749..eec156e47c6 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -61,8 +61,8 @@ use crate::{ use crate::{ repdef::{ CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation, - MiniBlockRepDefBudget, RepDefSlicer, SerializedRepDefs, SparseStructuralLayerPlan, - SparseStructuralPlan, build_control_word_iterator, + MiniBlockRepDefBudget, RepDefSlicer, SerializedRepDefs, SparseCountSet, SparsePositionSet, + SparseStructuralLayerPlan, SparseStructuralPlan, build_control_word_iterator, }, utils::accumulation::AccumulationQueue, }; @@ -676,29 +676,54 @@ impl Clone for LoadedChunk { } } +#[derive(Debug, Clone)] +enum SparsePositionSetDecoder { + Empty, + All { + len: u64, + }, + Range { + start: u64, + len: u64, + }, + Explicit { + decompressor: Arc, + count: u64, + domain_len: u64, + }, +} + +#[derive(Debug, Clone)] +enum SparseCountSetDecoder { + Empty, + Constant { + value: u64, + len: u64, + }, + Explicit { + decompressor: Arc, + count: u64, + }, +} + #[derive(Debug, Clone)] enum SparseLayerDecompressors { Validity { num_slots: u64, - null_decompressor: Option>, - num_nulls: u64, + null_positions: SparsePositionSetDecoder, }, List { num_slots: u64, num_child_slots: u64, - non_empty_decompressor: Option>, - num_non_empty: u64, - count_decompressor: Option>, - constant_count: u64, - null_decompressor: Option>, - num_nulls: u64, + non_empty_positions: SparsePositionSetDecoder, + counts: SparseCountSetDecoder, + null_positions: SparsePositionSetDecoder, }, FixedSizeList { num_slots: u64, num_child_slots: u64, dimension: u64, - null_decompressor: Option>, - num_nulls: u64, + null_positions: SparsePositionSetDecoder, }, } @@ -718,7 +743,7 @@ impl DeepSizeOf for SparseStructuralCacheableState { .iter() .map(|layer| match layer { SparseStructuralLayerPlan::Validity { null_positions, .. } => { - null_positions.len() * std::mem::size_of::() + null_positions.deep_size() } SparseStructuralLayerPlan::List { non_empty_positions, @@ -726,11 +751,12 @@ impl DeepSizeOf for SparseStructuralCacheableState { null_positions, .. } => { - (non_empty_positions.len() + counts.len() + null_positions.len()) - * std::mem::size_of::() + non_empty_positions.deep_size() + + counts.deep_size() + + null_positions.deep_size() } SparseStructuralLayerPlan::FixedSizeList { null_positions, .. } => { - null_positions.len() * std::mem::size_of::() + null_positions.deep_size() } }) .sum::(); @@ -838,47 +864,301 @@ impl SparseStructuralScheduler { }) } - fn count_optional_buffer( - compression: &Option, - count: u64, - label: &str, - ) -> Result { - if count == 0 { - if compression.is_some() { - return Err(Error::invalid_input_source( - format!("Sparse structural {label} has compression but no values").into(), - )); - } - return Ok(0); - } - let compression = compression.as_ref().ok_or_else(|| { + fn checked_buffer_range(position: u64, size: u64, label: &str) -> Result> { + let end = position.checked_add(size).ok_or_else(|| { Error::invalid_input_source( - format!("Sparse structural {label} has values but no compression").into(), + format!("Sparse structural {label} buffer range overflows").into(), ) })?; - Self::validate_compression(compression, label)?; - Ok(1) + Ok(position..end) } - fn reject_optional_buffer( - compression: &Option, + fn require_position_set<'a>( + set: &'a Option, label: &str, - ) -> Result<()> { - if compression.is_some() { + ) -> Result<&'a pb21::SparsePositionSet> { + set.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} position set is required").into(), + ) + }) + } + + fn require_count_set<'a>( + set: &'a Option, + label: &str, + ) -> Result<&'a pb21::SparseCountSet> { + set.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} count set is required").into(), + ) + }) + } + + fn reject_position_set(set: &Option, label: &str) -> Result<()> { + if set.is_some() { return Err(Error::invalid_input_source( - format!("Sparse structural {label} is not valid for this layer").into(), + format!("Sparse structural {label} position set is not valid for this layer") + .into(), )); } Ok(()) } - fn checked_buffer_range(position: u64, size: u64, label: &str) -> Result> { - let end = position.checked_add(size).ok_or_else(|| { + fn reject_count_set(set: &Option, label: &str) -> Result<()> { + if set.is_some() { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} count set is not valid for this layer").into(), + )); + } + Ok(()) + } + + fn position_cardinality( + position_set: &pb21::SparsePositionSet, + domain_len: u64, + label: &str, + ) -> Result { + let positions = position_set.positions.as_ref().ok_or_else(|| { Error::invalid_input_source( - format!("Sparse structural {label} buffer range overflows").into(), + format!("Sparse structural {label} position set is missing its variant").into(), ) })?; - Ok(position..end) + let cardinality = position_set.num_positions; + if cardinality > domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} cardinality {} exceeds domain {}", + cardinality, domain_len + ) + .into(), + )); + } + match positions { + pb21::sparse_position_set::Positions::Empty(_) => { + if cardinality != 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} empty set has cardinality {}", + cardinality + ) + .into(), + )); + } + } + pb21::sparse_position_set::Positions::All(_) => { + if cardinality != domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} all set has cardinality {}, expected {}", + cardinality, domain_len + ) + .into(), + )); + } + } + pb21::sparse_position_set::Positions::Range(range) => { + let end = range.start.checked_add(range.length).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} range overflows").into(), + ) + })?; + if range.length != cardinality || end > domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} range {}..{} does not match cardinality {} in domain {}", + range.start, end, cardinality, domain_len + ) + .into(), + )); + } + } + pb21::sparse_position_set::Positions::Explicit(compression) => { + if cardinality == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} has compression but no values").into(), + )); + } + Self::validate_compression(compression, label)?; + } + } + Ok(cardinality) + } + + fn position_buffer_count( + position_set: &pb21::SparsePositionSet, + domain_len: u64, + label: &str, + ) -> Result { + Self::position_cardinality(position_set, domain_len, label)?; + let positions = position_set.positions.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} position set is missing its variant").into(), + ) + })?; + Ok(usize::from(matches!( + positions, + pb21::sparse_position_set::Positions::Explicit(_) + ))) + } + + fn count_buffer_count( + count_set: &pb21::SparseCountSet, + cardinality: u64, + label: &str, + ) -> Result { + let counts = count_set.counts.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} count set is missing its variant").into(), + ) + })?; + match counts { + pb21::sparse_count_set::Counts::Empty(_) => { + if cardinality != 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} empty count set has cardinality {}", + cardinality + ) + .into(), + )); + } + Ok(0) + } + pb21::sparse_count_set::Counts::Constant(constant) => { + if cardinality == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} constant count has no values").into(), + )); + } + if constant.value == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} constant count is zero").into(), + )); + } + Ok(0) + } + pb21::sparse_count_set::Counts::Explicit(compression) => { + if cardinality == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} has compression but no values").into(), + )); + } + Self::validate_compression(compression, label)?; + Ok(1) + } + } + } + + fn count_set_child_slots( + count_set: &pb21::SparseCountSet, + cardinality: u64, + label: &str, + ) -> Result> { + let counts = count_set.counts.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} count set is missing its variant").into(), + ) + })?; + match counts { + pb21::sparse_count_set::Counts::Empty(_) => Ok(Some(0)), + pb21::sparse_count_set::Counts::Constant(constant) => constant + .value + .checked_mul(cardinality) + .map(Some) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} constant count sum overflows: value={}, len={}", + constant.value, cardinality + ) + .into(), + ) + }), + pb21::sparse_count_set::Counts::Explicit(_) => Ok(None), + } + } + + fn create_position_decompressor( + compression: &CompressiveEncoding, + label: &str, + decompressors: &dyn DecompressionStrategy, + ) -> Result> { + let compression = Self::validate_compression(compression, label)?; + decompressors + .create_block_decompressor(compression) + .map(Arc::from) + } + + fn position_set_decoder( + position_set: &pb21::SparsePositionSet, + domain_len: u64, + label: &str, + decompressors: &dyn DecompressionStrategy, + ) -> Result<(SparsePositionSetDecoder, u64)> { + let cardinality = Self::position_cardinality(position_set, domain_len, label)?; + let positions = position_set.positions.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} position set is missing its variant").into(), + ) + })?; + Ok(( + match positions { + pb21::sparse_position_set::Positions::Empty(_) => SparsePositionSetDecoder::Empty, + pb21::sparse_position_set::Positions::All(_) => { + SparsePositionSetDecoder::All { len: domain_len } + } + pb21::sparse_position_set::Positions::Range(range) => { + SparsePositionSetDecoder::Range { + start: range.start, + len: range.length, + } + } + pb21::sparse_position_set::Positions::Explicit(compression) => { + SparsePositionSetDecoder::Explicit { + decompressor: Self::create_position_decompressor( + compression, + label, + decompressors, + )?, + count: cardinality, + domain_len, + } + } + }, + cardinality, + )) + } + + fn count_set_decoder( + count_set: &pb21::SparseCountSet, + cardinality: u64, + label: &str, + decompressors: &dyn DecompressionStrategy, + ) -> Result { + Self::count_buffer_count(count_set, cardinality, label)?; + let counts = count_set.counts.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} count set is missing its variant").into(), + ) + })?; + Ok(match counts { + pb21::sparse_count_set::Counts::Empty(_) => SparseCountSetDecoder::Empty, + pb21::sparse_count_set::Counts::Constant(constant) => SparseCountSetDecoder::Constant { + value: constant.value, + len: cardinality, + }, + pb21::sparse_count_set::Counts::Explicit(compression) => { + SparseCountSetDecoder::Explicit { + decompressor: Self::create_position_decompressor( + compression, + label, + decompressors, + )?, + count: cardinality, + } + } + }) } fn structural_buffer_count(layers: &[pb21::SparseStructuralLayer]) -> Result { @@ -895,104 +1175,78 @@ impl SparseStructuralScheduler { .into(), )); } - if layer.constant_count != 0 || layer.num_non_empty != 0 { + if layer.fixed_size_list_dimension != 0 { return Err(Error::invalid_input_source( - "Sparse structural validity layer has list-only metadata".into(), + "Sparse structural validity layer has fixed-size-list dimension".into(), )); } - Self::reject_optional_buffer( - &layer.non_empty_compression, + Self::reject_position_set( + &layer.non_empty_positions, "validity non-empty positions", )?; - Self::reject_optional_buffer(&layer.count_compression, "validity counts")?; - count += Self::count_optional_buffer( - &layer.null_compression, - layer.num_nulls, + Self::reject_count_set(&layer.counts, "validity counts")?; + let null_positions = + Self::require_position_set(&layer.null_positions, "validity null")?; + count += Self::position_buffer_count( + null_positions, + layer.num_slots, "validity null positions", )?; } pb21::sparse_structural_layer::Kind::SparseLayerList => { - count += Self::count_optional_buffer( - &layer.non_empty_compression, - layer.num_non_empty, - "list non-empty positions", - )?; - if layer.num_non_empty > 0 - && layer.count_compression.is_none() - && layer.constant_count == 0 - { + if layer.fixed_size_list_dimension != 0 { return Err(Error::invalid_input_source( - "Sparse structural list has non-empty positions but no count compression or constant count".into(), + "Sparse structural list layer has fixed-size-list dimension".into(), )); } - if layer.num_non_empty == 0 && layer.constant_count != 0 { - return Err(Error::invalid_input_source( - "Sparse structural list has constant count but no non-empty positions" - .into(), - )); - } - if layer.count_compression.is_some() && layer.constant_count != 0 { + let non_empty_positions = + Self::require_position_set(&layer.non_empty_positions, "list non-empty")?; + let num_non_empty = Self::position_cardinality( + non_empty_positions, + layer.num_slots, + "list non-empty positions", + )?; + count += Self::position_buffer_count( + non_empty_positions, + layer.num_slots, + "list non-empty positions", + )?; + let counts = Self::require_count_set(&layer.counts, "list counts")?; + count += Self::count_buffer_count(counts, num_non_empty, "list counts")?; + if let Some(child_slots) = + Self::count_set_child_slots(counts, num_non_empty, "list counts")? + && child_slots != layer.num_child_slots + { return Err(Error::invalid_input_source( - "Sparse structural list has both count compression and constant count" - .into(), + format!( + "Sparse structural list count sum {} does not match child slots {}", + child_slots, layer.num_child_slots + ) + .into(), )); } - if layer.constant_count != 0 { - let expected_child_slots = layer - .num_non_empty - .checked_mul(layer.constant_count) - .ok_or_else(|| { - Error::invalid_input_source( - format!( - "Sparse structural list child slot count overflows: non_empty={}, constant_count={}", - layer.num_non_empty, layer.constant_count - ) - .into(), - ) - })?; - if expected_child_slots != layer.num_child_slots { - return Err(Error::invalid_input_source( - format!( - "Sparse structural list child slot count {} does not match non-empty slots {} * constant count {}", - layer.num_child_slots, layer.num_non_empty, layer.constant_count - ) - .into(), - )); - } - } else { - count += Self::count_optional_buffer( - &layer.count_compression, - layer.num_non_empty, - "list counts", - )?; - } - count += Self::count_optional_buffer( - &layer.null_compression, - layer.num_nulls, + let null_positions = + Self::require_position_set(&layer.null_positions, "list null")?; + count += Self::position_buffer_count( + null_positions, + layer.num_slots, "list null positions", )?; } pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList => { - if layer.num_non_empty != 0 { - return Err(Error::invalid_input_source( - "Sparse structural fixed-size-list layer has list-only metadata" - .into(), - )); - } - Self::reject_optional_buffer( - &layer.non_empty_compression, + Self::reject_position_set( + &layer.non_empty_positions, "fixed-size-list non-empty positions", )?; - Self::reject_optional_buffer( - &layer.count_compression, - "fixed-size-list counts", - )?; - let expected_child_slots = - layer.num_slots.checked_mul(layer.constant_count).ok_or_else(|| { + Self::reject_count_set(&layer.counts, "fixed-size-list counts")?; + let expected_child_slots = layer + .num_slots + .checked_mul(layer.fixed_size_list_dimension) + .ok_or_else(|| { Error::invalid_input_source( format!( "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", - layer.num_slots, layer.constant_count + layer.num_slots, layer.fixed_size_list_dimension ) .into(), ) @@ -1001,14 +1255,18 @@ impl SparseStructuralScheduler { return Err(Error::invalid_input_source( format!( "Sparse structural fixed-size-list child slot count {} does not match slots {} * dimension {}", - layer.num_child_slots, layer.num_slots, layer.constant_count + layer.num_child_slots, layer.num_slots, layer.fixed_size_list_dimension ) .into(), )); } - count += Self::count_optional_buffer( - &layer.null_compression, - layer.num_nulls, + let null_positions = Self::require_position_set( + &layer.null_positions, + "fixed-size-list null", + )?; + count += Self::position_buffer_count( + null_positions, + layer.num_slots, "fixed-size-list null positions", )?; } @@ -1026,65 +1284,67 @@ impl SparseStructuralScheduler { layer: &pb21::SparseStructuralLayer, decompressors: &dyn DecompressionStrategy, ) -> Result { - let null_decompressor = layer - .null_compression - .as_ref() - .map(|compression| { - let compression = Self::validate_compression(compression, "null positions")?; - decompressors - .create_block_decompressor(compression) - .map(Arc::from) - }) - .transpose()?; let kind = pb21::sparse_structural_layer::Kind::try_from(layer.kind)?; Ok(match kind { pb21::sparse_structural_layer::Kind::SparseLayerValidity => { + let null_positions = + Self::require_position_set(&layer.null_positions, "validity null")?; + let (null_positions, _) = Self::position_set_decoder( + null_positions, + layer.num_slots, + "validity null positions", + decompressors, + )?; SparseLayerDecompressors::Validity { num_slots: layer.num_slots, - null_decompressor, - num_nulls: layer.num_nulls, + null_positions, } } pb21::sparse_structural_layer::Kind::SparseLayerList => { - let non_empty_decompressor = layer - .non_empty_compression - .as_ref() - .map(|compression| { - let compression = - Self::validate_compression(compression, "list non-empty positions")?; - decompressors - .create_block_decompressor(compression) - .map(Arc::from) - }) - .transpose()?; - let count_decompressor = layer - .count_compression - .as_ref() - .map(|compression| { - let compression = Self::validate_compression(compression, "list counts")?; - decompressors - .create_block_decompressor(compression) - .map(Arc::from) - }) - .transpose()?; + let non_empty_positions = + Self::require_position_set(&layer.non_empty_positions, "list non-empty")?; + let (non_empty_positions, num_non_empty) = Self::position_set_decoder( + non_empty_positions, + layer.num_slots, + "list non-empty positions", + decompressors, + )?; + let counts = Self::count_set_decoder( + Self::require_count_set(&layer.counts, "list counts")?, + num_non_empty, + "list counts", + decompressors, + )?; + let null_positions = + Self::require_position_set(&layer.null_positions, "list null")?; + let (null_positions, _) = Self::position_set_decoder( + null_positions, + layer.num_slots, + "list null positions", + decompressors, + )?; SparseLayerDecompressors::List { num_slots: layer.num_slots, num_child_slots: layer.num_child_slots, - non_empty_decompressor, - num_non_empty: layer.num_non_empty, - count_decompressor, - constant_count: layer.constant_count, - null_decompressor, - num_nulls: layer.num_nulls, + non_empty_positions, + counts, + null_positions, } } pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList => { + let null_positions = + Self::require_position_set(&layer.null_positions, "fixed-size-list null")?; + let (null_positions, _) = Self::position_set_decoder( + null_positions, + layer.num_slots, + "fixed-size-list null positions", + decompressors, + )?; SparseLayerDecompressors::FixedSizeList { num_slots: layer.num_slots, num_child_slots: layer.num_child_slots, - dimension: layer.constant_count, - null_decompressor, - num_nulls: layer.num_nulls, + dimension: layer.fixed_size_list_dimension, + null_positions, } } pb21::sparse_structural_layer::Kind::SparseLayerUnspecified => { @@ -1220,26 +1480,13 @@ impl SparseStructuralScheduler { Ok(values.to_vec()) } - fn decode_positions( - decompressor: Option<&Arc>, - data: Option, + fn decode_explicit_positions( + decompressor: &Arc, + data: Bytes, num_positions: u64, num_slots: u64, label: &str, - ) -> Result> { - if num_positions == 0 { - return Ok(Vec::new()); - } - let decompressor = decompressor.ok_or_else(|| { - Error::invalid_input_source( - format!("Sparse structural {label} has positions but no compression").into(), - ) - })?; - let data = data.ok_or_else(|| { - Error::invalid_input_source( - format!("Sparse structural {label} is missing its position buffer").into(), - ) - })?; + ) -> Result { let deltas = Self::decode_u64_values(decompressor.as_ref(), data, num_positions, label)?; let mut positions = Vec::with_capacity(deltas.len()); let mut current = 0_u64; @@ -1270,7 +1517,57 @@ impl SparseStructuralScheduler { } positions.push(current); } - Ok(positions) + SparsePositionSet::from_positions(positions, num_slots, label) + } + + fn decode_position_set( + decoder: &SparsePositionSetDecoder, + buffers: &mut impl Iterator, + label: &str, + ) -> Result { + match decoder { + SparsePositionSetDecoder::Empty => Ok(SparsePositionSet::empty()), + SparsePositionSetDecoder::All { len } => Ok(SparsePositionSet::all(*len)), + SparsePositionSetDecoder::Range { start, len } => { + Ok(SparsePositionSet::range(*start, *len)) + } + SparsePositionSetDecoder::Explicit { + decompressor, + count, + domain_len, + } => Self::decode_explicit_positions( + decompressor, + Self::next_structural_buffer(buffers, label)?, + *count, + *domain_len, + label, + ), + } + } + + fn decode_count_set( + decoder: &SparseCountSetDecoder, + buffers: &mut impl Iterator, + label: &str, + ) -> Result { + match decoder { + SparseCountSetDecoder::Empty => Ok(SparseCountSet::Empty), + SparseCountSetDecoder::Constant { value, len } => { + Ok(SparseCountSet::constant(*value, *len)) + } + SparseCountSetDecoder::Explicit { + decompressor, + count, + } => { + let counts = Self::decode_u64_values( + decompressor.as_ref(), + Self::next_structural_buffer(buffers, label)?, + *count, + label, + )?; + Ok(SparseCountSet::from_counts(counts)) + } + } } fn next_structural_buffer( @@ -1284,18 +1581,6 @@ impl SparseStructuralScheduler { }) } - fn optional_structural_buffer( - buffers: &mut impl Iterator, - needed: bool, - label: &str, - ) -> Result> { - if needed { - Self::next_structural_buffer(buffers, label).map(Some) - } else { - Ok(None) - } - } - fn decode_layer( layer: &SparseLayerDecompressors, buffers: &mut impl Iterator, @@ -1303,71 +1588,31 @@ impl SparseStructuralScheduler { Ok(match layer { SparseLayerDecompressors::Validity { num_slots, - null_decompressor, - num_nulls, + null_positions, } => SparseStructuralLayerPlan::Validity { num_slots: *num_slots, - null_positions: Self::decode_positions( - null_decompressor.as_ref(), - Self::optional_structural_buffer( - buffers, - *num_nulls > 0, - "validity null positions", - )?, - *num_nulls, - *num_slots, + null_positions: Self::decode_position_set( + null_positions, + buffers, "validity null positions", )?, }, SparseLayerDecompressors::List { num_slots, num_child_slots, - non_empty_decompressor, - num_non_empty, - count_decompressor, - constant_count, - null_decompressor, - num_nulls, + non_empty_positions, + counts, + null_positions, } => { - let non_empty_positions = Self::decode_positions( - non_empty_decompressor.as_ref(), - Self::optional_structural_buffer( - buffers, - *num_non_empty > 0, - "list non-empty positions", - )?, - *num_non_empty, - *num_slots, + let non_empty_positions = Self::decode_position_set( + non_empty_positions, + buffers, "list non-empty positions", )?; - let counts = if let Some(count_decompressor) = count_decompressor { - Self::decode_u64_values( - count_decompressor.as_ref(), - Self::next_structural_buffer(buffers, "list counts")?, - *num_non_empty, - "list counts", - )? - } else { - vec![*constant_count; *num_non_empty as usize] - }; - let null_positions = Self::decode_positions( - null_decompressor.as_ref(), - Self::optional_structural_buffer( - buffers, - *num_nulls > 0, - "list null positions", - )?, - *num_nulls, - *num_slots, - "list null positions", - )?; - let actual_child_slots = counts.iter().try_fold(0_u64, |sum, count| { - sum.checked_add(*count).ok_or_else(|| { - Error::invalid_input_source( - "Sparse structural list count sum overflows".into(), - ) - }) - })?; + let counts = Self::decode_count_set(counts, buffers, "list counts")?; + let null_positions = + Self::decode_position_set(null_positions, buffers, "list null positions")?; + let actual_child_slots = counts.sum()?; if actual_child_slots != *num_child_slots { return Err(Error::invalid_input_source( format!( @@ -1382,7 +1627,6 @@ impl SparseStructuralScheduler { num_child_slots: *num_child_slots, non_empty_positions, counts, - constant_count: (*constant_count > 0).then_some(*constant_count), null_positions, } } @@ -1390,8 +1634,7 @@ impl SparseStructuralScheduler { num_slots, num_child_slots, dimension, - null_decompressor, - num_nulls, + null_positions, } => { let expected_child_slots = num_slots.checked_mul(*dimension).ok_or_else(|| { Error::invalid_input_source( @@ -1414,15 +1657,9 @@ impl SparseStructuralScheduler { SparseStructuralLayerPlan::FixedSizeList { num_slots: *num_slots, dimension: *dimension, - null_positions: Self::decode_positions( - null_decompressor.as_ref(), - Self::optional_structural_buffer( - buffers, - *num_nulls > 0, - "fixed-size-list null positions", - )?, - *num_nulls, - *num_slots, + null_positions: Self::decode_position_set( + null_positions, + buffers, "fixed-size-list null positions", )?, } @@ -1901,23 +2138,175 @@ struct SparseStructuralSelection { leaf_ranges: Vec>, } -fn translate_positions(positions: &[u64], ranges: &[Range], label: &str) -> Result> { - let mut translated = Vec::new(); - let mut output_base = 0_u64; +struct SparsePositionSelection { + positions: SparsePositionSet, + ordinal_ranges: Vec>, +} + +fn validate_slice_ranges(ranges: &[Range], domain_len: u64, label: &str) -> Result { + let mut total = 0_u64; for range in ranges { - let mut idx = positions.partition_point(|position| *position < range.start); - while idx < positions.len() && positions[idx] < range.end { - translated.push(output_base + positions[idx] - range.start); - idx += 1; + if range.start > range.end || range.end > domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} slice {}..{} is outside domain {}", + range.start, range.end, domain_len + ) + .into(), + )); + } + total = total.checked_add(range.end - range.start).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} slice length overflows").into(), + ) + })?; + } + Ok(total) +} + +fn push_coalesced_range(ranges: &mut Vec>, range: Range) { + if range.is_empty() { + return; + } + if let Some(last) = ranges.last_mut() + && last.end == range.start + { + last.end = range.end; + return; + } + ranges.push(range); +} + +fn position_segments_to_set( + segments: Vec>, + output_domain: u64, + label: &str, +) -> Result { + let segments = coalesce_ranges(segments); + if segments.is_empty() { + return Ok(SparsePositionSet::empty()); + } + if segments.len() == 1 { + let segment = &segments[0]; + if segment.start == 0 && segment.end == output_domain { + return Ok(SparsePositionSet::all(output_domain)); + } + return Ok(SparsePositionSet::range( + segment.start, + segment.end - segment.start, + )); + } + + let total_len = segments + .iter() + .map(|range| range.end - range.start) + .try_fold(0_u64, |sum, len| { + sum.checked_add(len).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} segment length overflows").into(), + ) + }) + })?; + let total_len = usize::try_from(total_len).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} segment length exceeds usize::MAX").into(), + ) + })?; + let mut positions = Vec::with_capacity(total_len); + for segment in segments { + positions.extend(segment); + } + SparsePositionSet::from_positions(positions, output_domain, label) +} + +fn select_position_set( + positions: &SparsePositionSet, + ranges: &[Range], + domain_len: u64, + label: &str, +) -> Result { + let output_domain = validate_slice_ranges(ranges, domain_len, label)?; + let mut segments = Vec::new(); + let mut ordinal_ranges = Vec::new(); + let mut output_base = 0_u64; + + match positions { + SparsePositionSet::Empty => {} + SparsePositionSet::All { len } => { + if *len != domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} all set length {} does not match domain {}", + len, domain_len + ) + .into(), + )); + } + for range in ranges { + let range_len = range.end - range.start; + push_coalesced_range(&mut segments, output_base..output_base + range_len); + push_coalesced_range(&mut ordinal_ranges, range.clone()); + output_base += range_len; + } + } + SparsePositionSet::Range { start, len } => { + let source_end = start.checked_add(*len).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} range overflows").into(), + ) + })?; + if source_end > domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} range {}..{} is outside domain {}", + start, source_end, domain_len + ) + .into(), + )); + } + for range in ranges { + let intersect_start = range.start.max(*start); + let intersect_end = range.end.min(source_end); + if intersect_start < intersect_end { + let out_start = output_base + intersect_start - range.start; + let out_end = output_base + intersect_end - range.start; + push_coalesced_range(&mut segments, out_start..out_end); + push_coalesced_range( + &mut ordinal_ranges, + intersect_start - *start..intersect_end - *start, + ); + } + output_base += range.end - range.start; + } + } + SparsePositionSet::Explicit(source_positions) => { + let mut out_positions = Vec::new(); + for range in ranges { + let idx_start = + source_positions.partition_point(|position| *position < range.start); + let idx_end = source_positions.partition_point(|position| *position < range.end); + if idx_start < idx_end { + out_positions.extend( + source_positions[idx_start..idx_end] + .iter() + .map(|position| output_base + *position - range.start), + ); + push_coalesced_range(&mut ordinal_ranges, idx_start as u64..idx_end as u64); + } + output_base += range.end - range.start; + } + let positions = SparsePositionSet::from_positions(out_positions, output_domain, label)?; + return Ok(SparsePositionSelection { + positions, + ordinal_ranges, + }); } - output_base += range.end - range.start; - } - if translated.windows(2).any(|window| window[0] >= window[1]) { - return Err(Error::invalid_input_source( - format!("Sparse structural {label} positions are not sorted after slicing").into(), - )); } - Ok(translated) + + Ok(SparsePositionSelection { + positions: position_segments_to_set(segments, output_domain, label)?, + ordinal_ranges, + }) } fn offsets_from_counts(counts: &[u64]) -> Result> { @@ -1953,10 +2342,9 @@ fn coalesce_ranges(ranges: Vec>) -> Vec> { fn slice_list_layer( num_slots: u64, num_child_slots: u64, - non_empty_positions: &[u64], - counts: &[u64], - constant_count: Option, - null_positions: &[u64], + non_empty_positions: &SparsePositionSet, + counts: &SparseCountSet, + null_positions: &SparsePositionSet, ranges: &[Range], ) -> Result<(SparseStructuralLayerPlan, Vec>)> { if non_empty_positions.len() != counts.len() { @@ -1964,102 +2352,177 @@ fn slice_list_layer( "Sparse structural list has mismatched non-empty positions and counts".into(), )); } - if ranges - .iter() - .any(|range| range.start > range.end || range.end > num_slots) - { + if counts.sum()? != num_child_slots { return Err(Error::invalid_input_source( format!( - "Sparse structural list slice is outside layer with {} slots", - num_slots + "Sparse structural list count sum {} does not match child slots {}", + counts.sum()?, + num_child_slots ) .into(), )); } - let mut out_non_empty_positions = Vec::new(); - let mut out_counts = Vec::new(); - let mut matched_indices = Vec::new(); - let mut output_base = 0_u64; - for range in ranges { - let mut idx = non_empty_positions.partition_point(|position| *position < range.start); - while idx < non_empty_positions.len() && non_empty_positions[idx] < range.end { - out_non_empty_positions.push(output_base + non_empty_positions[idx] - range.start); - out_counts.push(counts[idx]); - matched_indices.push(idx); - idx += 1; - } - output_base += range.end - range.start; - } - - let child_ranges = if matched_indices.is_empty() { - Vec::new() - } else if let Some(constant_count) = constant_count { - let expected_child_slots = - constant_count - .checked_mul(counts.len() as u64) - .ok_or_else(|| { - Error::invalid_input_source( - "Sparse structural list constant counts overflow child slots".into(), - ) - })?; - if expected_child_slots != num_child_slots { - return Err(Error::invalid_input_source( - format!( - "Sparse structural list constant count sum {} does not match child slots {}", - expected_child_slots, num_child_slots - ) - .into(), - )); - } - matched_indices - .iter() - .map(|idx| { - let start = *idx as u64 * constant_count; - start..start + constant_count - }) - .collect() - } else { - let value_offsets = offsets_from_counts(counts)?; - let last_offset = *value_offsets.last().expect_ok()?; - if last_offset != num_child_slots { - return Err(Error::invalid_input_source( - format!( - "Sparse structural list count sum {} does not match child slots {}", - last_offset, num_child_slots - ) - .into(), - )); - } - matched_indices - .iter() - .map(|idx| value_offsets[*idx]..value_offsets[*idx + 1]) - .collect() - }; - - let out_null_positions = translate_positions(null_positions, ranges, "list null")?; + let non_empty_selection = + select_position_set(non_empty_positions, ranges, num_slots, "list non-empty")?; + let out_counts = select_count_set( + counts, + &non_empty_selection.ordinal_ranges, + non_empty_selection.positions.len(), + )?; + let child_ranges = + child_ranges_from_counts(counts, num_child_slots, &non_empty_selection.ordinal_ranges)?; + let out_null_positions = + select_position_set(null_positions, ranges, num_slots, "list null")?.positions; let out_num_slots = ranges .iter() .map(|range| range.end - range.start) .sum::(); - let out_num_child_slots = out_counts.iter().sum::(); - let constant_count = out_counts - .first() - .copied() - .filter(|first| out_counts.iter().all(|count| count == first)); + let out_num_child_slots = out_counts.sum()?; Ok(( SparseStructuralLayerPlan::List { num_slots: out_num_slots, num_child_slots: out_num_child_slots, - non_empty_positions: out_non_empty_positions, + non_empty_positions: non_empty_selection.positions, counts: out_counts, - constant_count, null_positions: out_null_positions, }, child_ranges, )) } +fn select_count_set( + counts: &SparseCountSet, + ordinal_ranges: &[Range], + selected_len: u64, +) -> Result { + match counts { + SparseCountSet::Empty => { + if selected_len != 0 { + return Err(Error::invalid_input_source( + "Sparse structural selected non-empty positions but counts are empty".into(), + )); + } + Ok(SparseCountSet::Empty) + } + SparseCountSet::Constant { value, len } => { + validate_ordinal_ranges(ordinal_ranges, *len, "constant list counts")?; + Ok(SparseCountSet::constant(*value, selected_len)) + } + SparseCountSet::Explicit(source_counts) => { + validate_ordinal_ranges( + ordinal_ranges, + source_counts.len() as u64, + "explicit list counts", + )?; + let selected_len = usize::try_from(selected_len).map_err(|_| { + Error::invalid_input_source( + "Sparse structural selected count length exceeds usize::MAX".into(), + ) + })?; + let mut out_counts = Vec::with_capacity(selected_len); + for range in ordinal_ranges { + let start = usize::try_from(range.start).map_err(|_| { + Error::invalid_input_source( + "Sparse structural count ordinal exceeds usize::MAX".into(), + ) + })?; + let end = usize::try_from(range.end).map_err(|_| { + Error::invalid_input_source( + "Sparse structural count ordinal exceeds usize::MAX".into(), + ) + })?; + out_counts.extend_from_slice(&source_counts[start..end]); + } + Ok(SparseCountSet::from_counts(out_counts)) + } + } +} + +fn validate_ordinal_ranges(ranges: &[Range], len: u64, label: &str) -> Result<()> { + for range in ranges { + if range.start > range.end || range.end > len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} ordinal range {}..{} is outside {} values", + range.start, range.end, len + ) + .into(), + )); + } + } + Ok(()) +} + +fn child_ranges_from_counts( + counts: &SparseCountSet, + num_child_slots: u64, + ordinal_ranges: &[Range], +) -> Result>> { + if ordinal_ranges.is_empty() { + return Ok(Vec::new()); + } + let child_ranges = match counts { + SparseCountSet::Empty => { + return Err(Error::invalid_input_source( + "Sparse structural selected non-empty positions but counts are empty".into(), + )); + } + SparseCountSet::Constant { value, len } => { + let expected_child_slots = value.checked_mul(*len).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list constant count sum overflows child slots".into(), + ) + })?; + if expected_child_slots != num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list constant count sum {} does not match child slots {}", + expected_child_slots, num_child_slots + ) + .into(), + )); + } + validate_ordinal_ranges(ordinal_ranges, *len, "constant list counts")?; + ordinal_ranges + .iter() + .map(|range| range.start * value..range.end * value) + .collect() + } + SparseCountSet::Explicit(counts) => { + let value_offsets = offsets_from_counts(counts)?; + let last_offset = *value_offsets.last().expect_ok()?; + if last_offset != num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list count sum {} does not match child slots {}", + last_offset, num_child_slots + ) + .into(), + )); + } + validate_ordinal_ranges(ordinal_ranges, counts.len() as u64, "explicit list counts")?; + ordinal_ranges + .iter() + .map(|range| { + let start = usize::try_from(range.start).map_err(|_| { + Error::invalid_input_source( + "Sparse structural count ordinal exceeds usize::MAX".into(), + ) + })?; + let end = usize::try_from(range.end).map_err(|_| { + Error::invalid_input_source( + "Sparse structural count ordinal exceeds usize::MAX".into(), + ) + })?; + Ok(value_offsets[start]..value_offsets[end]) + }) + .collect::>>()? + } + }; + Ok(coalesce_ranges(child_ranges)) +} + fn fsl_parent_ranges_from_child_ranges( ranges: &[Range], dimension: u64, @@ -2097,7 +2560,7 @@ fn slice_sparse_plan( for layer in &plan.layers { match layer { SparseStructuralLayerPlan::Validity { - num_slots: _, + num_slots, null_positions, } => { let out_num_slots = selected_ranges @@ -2106,11 +2569,13 @@ fn slice_sparse_plan( .sum::(); sliced_layers.push(SparseStructuralLayerPlan::Validity { num_slots: out_num_slots, - null_positions: translate_positions( + null_positions: select_position_set( null_positions, &selected_ranges, + *num_slots, "validity null", - )?, + )? + .positions, }); } SparseStructuralLayerPlan::List { @@ -2118,7 +2583,6 @@ fn slice_sparse_plan( num_child_slots, non_empty_positions, counts, - constant_count, null_positions, .. } => { @@ -2127,7 +2591,6 @@ fn slice_sparse_plan( *num_child_slots, non_empty_positions, counts, - *constant_count, null_positions, &selected_ranges, )?; @@ -2184,11 +2647,13 @@ fn slice_sparse_plan( sliced_layers.push(SparseStructuralLayerPlan::FixedSizeList { num_slots: out_num_slots, dimension: *dimension, - null_positions: translate_positions( + null_positions: select_position_set( null_positions, &fsl_ranges, + *num_slots, "fixed-size-list null", - )?, + )? + .positions, }); selected_ranges = child_ranges; selected_domain = num_child_slots; @@ -6347,6 +6812,88 @@ impl PrimitiveStructuralEncoder { Self::encode_sparse_u64_values(deltas, compression_strategy).map(Some) } + fn encode_sparse_position_set( + positions: &SparsePositionSet, + compression_strategy: &dyn CompressionStrategy, + label: &str, + ) -> Result<(Option, pb21::SparsePositionSet)> { + let num_positions = positions.len(); + let (buffer, positions) = match positions { + SparsePositionSet::Empty => ( + None, + pb21::sparse_position_set::Positions::Empty(pb21::SparsePositionEmpty {}), + ), + SparsePositionSet::All { .. } => ( + None, + pb21::sparse_position_set::Positions::All(pb21::SparsePositionAll {}), + ), + SparsePositionSet::Range { start, len } => ( + None, + pb21::sparse_position_set::Positions::Range(pb21::SparsePositionRange { + start: *start, + length: *len, + }), + ), + SparsePositionSet::Explicit(positions) => { + let (buffer, encoding) = + Self::encode_sparse_positions(positions, compression_strategy, label)? + .ok_or_else(|| { + Error::internal(format!( + "Sparse structural {label} explicit set is empty" + )) + })?; + ( + Some(buffer), + pb21::sparse_position_set::Positions::Explicit(encoding), + ) + } + }; + Ok(( + buffer, + pb21::SparsePositionSet { + positions: Some(positions), + num_positions, + }, + )) + } + + fn encode_sparse_count_set( + counts: &SparseCountSet, + compression_strategy: &dyn CompressionStrategy, + ) -> Result<(Option, pb21::SparseCountSet)> { + let (buffer, counts) = match counts { + SparseCountSet::Empty => ( + None, + pb21::sparse_count_set::Counts::Empty(pb21::SparseCountEmpty {}), + ), + SparseCountSet::Constant { value, .. } => ( + None, + pb21::sparse_count_set::Counts::Constant(pb21::SparseCountConstant { + value: *value, + }), + ), + SparseCountSet::Explicit(counts) => { + if counts.is_empty() { + return Err(Error::internal( + "Sparse structural explicit count set is empty".to_string(), + )); + } + let (buffer, encoding) = + Self::encode_sparse_u64_values(counts.clone(), compression_strategy)?; + ( + Some(buffer), + pb21::sparse_count_set::Counts::Explicit(encoding), + ) + } + }; + Ok(( + buffer, + pb21::SparseCountSet { + counts: Some(counts), + }, + )) + } + fn encode_sparse_structural_plan( plan: &SparseStructuralPlan, compression_strategy: &dyn CompressionStrategy, @@ -6360,25 +6907,22 @@ impl PrimitiveStructuralEncoder { num_slots, null_positions, } => { - let null_compression = Self::encode_sparse_positions( + let (null_buffer, null_positions_pb) = Self::encode_sparse_position_set( null_positions, compression_strategy, "validity null", - )? - .map(|(buffer, encoding)| { + )?; + if let Some(buffer) = null_buffer { buffers.push(buffer); - encoding - }); + } layers.push(pb21::SparseStructuralLayer { kind: pb21::sparse_structural_layer::Kind::SparseLayerValidity as i32, num_slots: *num_slots, num_child_slots: *num_slots, - constant_count: 0, - non_empty_compression: None, - num_non_empty: 0, - count_compression: None, - null_compression, - num_nulls: null_positions.len() as u64, + non_empty_positions: None, + counts: None, + null_positions: Some(null_positions_pb), + fixed_size_list_dimension: 0, }); } SparseStructuralLayerPlan::List { @@ -6386,7 +6930,6 @@ impl PrimitiveStructuralEncoder { num_child_slots, non_empty_positions, counts, - constant_count, null_positions, } => { if non_empty_positions.len() != counts.len() { @@ -6399,42 +6942,36 @@ impl PrimitiveStructuralEncoder { .into(), )); } - let non_empty_compression = Self::encode_sparse_positions( - non_empty_positions, - compression_strategy, - "list non-empty", - )? - .map(|(buffer, encoding)| { + let (non_empty_buffer, non_empty_positions_pb) = + Self::encode_sparse_position_set( + non_empty_positions, + compression_strategy, + "list non-empty", + )?; + if let Some(buffer) = non_empty_buffer { buffers.push(buffer); - encoding - }); - let count_compression = if constant_count.is_some() || counts.is_empty() { - None - } else { - let (buffer, encoding) = - Self::encode_sparse_u64_values(counts.clone(), compression_strategy)?; + } + let (count_buffer, counts_pb) = + Self::encode_sparse_count_set(counts, compression_strategy)?; + if let Some(buffer) = count_buffer { buffers.push(buffer); - Some(encoding) - }; - let null_compression = Self::encode_sparse_positions( + } + let (null_buffer, null_positions_pb) = Self::encode_sparse_position_set( null_positions, compression_strategy, "list null", - )? - .map(|(buffer, encoding)| { + )?; + if let Some(buffer) = null_buffer { buffers.push(buffer); - encoding - }); + } layers.push(pb21::SparseStructuralLayer { kind: pb21::sparse_structural_layer::Kind::SparseLayerList as i32, num_slots: *num_slots, num_child_slots: *num_child_slots, - constant_count: constant_count.unwrap_or(0), - non_empty_compression, - num_non_empty: non_empty_positions.len() as u64, - count_compression, - null_compression, - num_nulls: null_positions.len() as u64, + non_empty_positions: Some(non_empty_positions_pb), + counts: Some(counts_pb), + null_positions: Some(null_positions_pb), + fixed_size_list_dimension: 0, }); } SparseStructuralLayerPlan::FixedSizeList { @@ -6442,15 +6979,14 @@ impl PrimitiveStructuralEncoder { dimension, null_positions, } => { - let null_compression = Self::encode_sparse_positions( + let (null_buffer, null_positions_pb) = Self::encode_sparse_position_set( null_positions, compression_strategy, "fixed-size-list null", - )? - .map(|(buffer, encoding)| { + )?; + if let Some(buffer) = null_buffer { buffers.push(buffer); - encoding - }); + } let num_child_slots = num_slots.checked_mul(*dimension).ok_or_else(|| { Error::invalid_input_source( format!( @@ -6464,12 +7000,10 @@ impl PrimitiveStructuralEncoder { kind: pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList as i32, num_slots: *num_slots, num_child_slots, - constant_count: *dimension, - non_empty_compression: None, - num_non_empty: 0, - count_compression: None, - null_compression, - num_nulls: null_positions.len() as u64, + non_empty_positions: None, + counts: None, + null_positions: Some(null_positions_pb), + fixed_size_list_dimension: *dimension, }); } } @@ -9743,6 +10277,27 @@ mod tests { Arc::new(builder.finish()) } + fn dense_prefix_sparse_i32_list_array( + num_rows: i32, + dense_rows: i32, + values_per_row: i32, + ) -> ArrayRef { + use arrow_array::builder::{Int32Builder, ListBuilder}; + + let mut builder = ListBuilder::new(Int32Builder::new()); + for row in 0..num_rows { + if row < dense_rows { + for value_idx in 0..values_per_row { + builder + .values() + .append_value(row * values_per_row + value_idx); + } + } + builder.append(true); + } + Arc::new(builder.finish()) + } + fn single_row_overbudget_sparse_i32_nested_list_array(num_inner_lists: usize) -> ArrayRef { use arrow_array::{Int32Array, ListArray}; use arrow_buffer::{OffsetBuffer, ScalarBuffer}; @@ -9814,6 +10369,111 @@ mod tests { .collect() } + fn sparse_position_set( + positions: pb21::sparse_position_set::Positions, + num_positions: u64, + ) -> Option { + Some(pb21::SparsePositionSet { + positions: Some(positions), + num_positions, + }) + } + + fn sparse_position_empty() -> Option { + sparse_position_set( + pb21::sparse_position_set::Positions::Empty(pb21::SparsePositionEmpty {}), + 0, + ) + } + + fn sparse_position_all(num_positions: u64) -> Option { + sparse_position_set( + pb21::sparse_position_set::Positions::All(pb21::SparsePositionAll {}), + num_positions, + ) + } + + fn sparse_position_explicit(num_positions: u64) -> Option { + sparse_position_set( + pb21::sparse_position_set::Positions::Explicit(ProtobufUtils21::flat(64, None)), + num_positions, + ) + } + + fn sparse_count_set(counts: pb21::sparse_count_set::Counts) -> Option { + Some(pb21::SparseCountSet { + counts: Some(counts), + }) + } + + fn sparse_count_constant(value: u64) -> Option { + sparse_count_set(pb21::sparse_count_set::Counts::Constant( + pb21::SparseCountConstant { value }, + )) + } + + fn sparse_count_explicit() -> Option { + sparse_count_set(pb21::sparse_count_set::Counts::Explicit( + ProtobufUtils21::flat(64, None), + )) + } + + fn sparse_position_num(position_set: &Option) -> u64 { + position_set + .as_ref() + .map(|set| set.num_positions) + .unwrap_or(0) + } + + fn assert_sparse_position_explicit(position_set: &Option) { + let positions = position_set + .as_ref() + .and_then(|set| set.positions.as_ref()) + .unwrap_or_else(|| panic!("expected sparse position set")); + assert!( + matches!(positions, pb21::sparse_position_set::Positions::Explicit(_)), + "expected explicit sparse position set, got {positions:?}" + ); + } + + fn assert_sparse_position_range( + position_set: &Option, + start: u64, + length: u64, + ) { + let positions = position_set + .as_ref() + .and_then(|set| set.positions.as_ref()) + .unwrap_or_else(|| panic!("expected sparse position set")); + let pb21::sparse_position_set::Positions::Range(range) = positions else { + panic!("expected range sparse position set, got {positions:?}"); + }; + assert_eq!(range.start, start); + assert_eq!(range.length, length); + } + + fn assert_sparse_position_all(position_set: &Option) { + let positions = position_set + .as_ref() + .and_then(|set| set.positions.as_ref()) + .unwrap_or_else(|| panic!("expected sparse position set")); + assert!( + matches!(positions, pb21::sparse_position_set::Positions::All(_)), + "expected all sparse position set, got {positions:?}" + ); + } + + fn assert_sparse_count_constant(count_set: &Option, value: u64) { + let counts = count_set + .as_ref() + .and_then(|set| set.counts.as_ref()) + .unwrap_or_else(|| panic!("expected sparse count set")); + let pb21::sparse_count_set::Counts::Constant(constant) = counts else { + panic!("expected constant sparse count set, got {counts:?}"); + }; + assert_eq!(constant.value, value); + } + fn sparse_layouts_from_pages( pages: &[crate::encoder::EncodedPage], ) -> Vec<&crate::format::pb21::SparseLayout> { @@ -10126,12 +10786,10 @@ mod tests { as i32, num_slots: 1, num_child_slots: 1, - constant_count: 0, - non_empty_compression: None, - num_non_empty: 0, - count_compression: None, - null_compression: Some(ProtobufUtils21::flat(64, None)), - num_nulls: 1, + non_empty_positions: None, + counts: None, + null_positions: sparse_position_explicit(1), + fixed_size_list_dimension: 0, }); let decompressors = DefaultDecompressionStrategy::default(); @@ -10155,12 +10813,10 @@ mod tests { kind: pb21::sparse_structural_layer::Kind::SparseLayerValidity as i32, num_slots: 1, num_child_slots: 2, - constant_count: 0, - non_empty_compression: None, - num_non_empty: 0, - count_compression: None, - null_compression: None, - num_nulls: 0, + non_empty_positions: None, + counts: None, + null_positions: sparse_position_empty(), + fixed_size_list_dimension: 0, }); let decompressors = DefaultDecompressionStrategy::default(); @@ -10184,12 +10840,10 @@ mod tests { kind: pb21::sparse_structural_layer::Kind::SparseLayerValidity as i32, num_slots: 1, num_child_slots: 1, - constant_count: 0, - non_empty_compression: Some(ProtobufUtils21::flat(64, None)), - num_non_empty: 0, - count_compression: None, - null_compression: None, - num_nulls: 0, + non_empty_positions: sparse_position_empty(), + counts: None, + null_positions: sparse_position_empty(), + fixed_size_list_dimension: 0, }); let decompressors = DefaultDecompressionStrategy::default(); @@ -10203,22 +10857,23 @@ mod tests { ) .unwrap_err(); - assert_invalid_input_contains(err, "validity non-empty positions is not valid"); + assert_invalid_input_contains( + err, + "validity non-empty positions position set is not valid", + ); } #[test] - fn test_sparse_layout_rejects_zero_count_list_count_compression() { + fn test_sparse_layout_rejects_zero_count_explicit_list_counts() { let mut layout = sparse_test_layout(); layout.structural_layers.push(pb21::SparseStructuralLayer { kind: pb21::sparse_structural_layer::Kind::SparseLayerList as i32, num_slots: 1, num_child_slots: 0, - constant_count: 0, - non_empty_compression: None, - num_non_empty: 0, - count_compression: Some(ProtobufUtils21::flat(64, None)), - null_compression: None, - num_nulls: 0, + non_empty_positions: sparse_position_empty(), + counts: sparse_count_explicit(), + null_positions: sparse_position_empty(), + fixed_size_list_dimension: 0, }); let decompressors = DefaultDecompressionStrategy::default(); @@ -10236,18 +10891,16 @@ mod tests { } #[test] - fn test_sparse_layout_rejects_list_constant_and_count_compression() { + fn test_sparse_layout_rejects_missing_list_counts() { let mut layout = sparse_test_layout(); layout.structural_layers.push(pb21::SparseStructuralLayer { kind: pb21::sparse_structural_layer::Kind::SparseLayerList as i32, num_slots: 1, num_child_slots: 1, - constant_count: 1, - non_empty_compression: Some(ProtobufUtils21::flat(64, None)), - num_non_empty: 1, - count_compression: Some(ProtobufUtils21::flat(64, None)), - null_compression: None, - num_nulls: 0, + non_empty_positions: sparse_position_all(1), + counts: None, + null_positions: sparse_position_empty(), + fixed_size_list_dimension: 0, }); let decompressors = DefaultDecompressionStrategy::default(); @@ -10261,7 +10914,7 @@ mod tests { ) .unwrap_err(); - assert_invalid_input_contains(err, "both count compression and constant count"); + assert_invalid_input_contains(err, "list counts count set is required"); } #[test] @@ -10271,12 +10924,10 @@ mod tests { kind: pb21::sparse_structural_layer::Kind::SparseLayerList as i32, num_slots: 2, num_child_slots: 3, - constant_count: 2, - non_empty_compression: Some(ProtobufUtils21::flat(64, None)), - num_non_empty: 2, - count_compression: None, - null_compression: None, - num_nulls: 0, + non_empty_positions: sparse_position_all(2), + counts: sparse_count_constant(2), + null_positions: sparse_position_empty(), + fixed_size_list_dimension: 0, }); let decompressors = DefaultDecompressionStrategy::default(); @@ -10290,7 +10941,7 @@ mod tests { ) .unwrap_err(); - assert_invalid_input_contains(err, "list child slot count"); + assert_invalid_input_contains(err, "list count sum 4 does not match child slots 3"); } #[tokio::test] @@ -10300,12 +10951,10 @@ mod tests { kind: pb21::sparse_structural_layer::Kind::SparseLayerValidity as i32, num_slots: 1, num_child_slots: 1, - constant_count: 0, - non_empty_compression: None, - num_non_empty: 0, - count_compression: None, - null_compression: Some(ProtobufUtils21::flat(64, None)), - num_nulls: 1, + non_empty_positions: None, + counts: None, + null_positions: sparse_position_explicit(1), + fixed_size_list_dimension: 0, }); let decompressors = DefaultDecompressionStrategy::default(); let mut scheduler = SparseStructuralScheduler::try_new( @@ -10401,10 +11050,9 @@ mod tests { let list_layers = sparse_list_layers(sparse); assert_eq!(list_layers.len(), 1); assert_eq!(list_layers[0].num_slots, 4096); - assert_eq!(list_layers[0].num_non_empty, 64); - assert_eq!(list_layers[0].constant_count, 1); - assert!(list_layers[0].non_empty_compression.is_some()); - assert!(list_layers[0].count_compression.is_none()); + assert_eq!(sparse_position_num(&list_layers[0].non_empty_positions), 64); + assert_sparse_position_explicit(&list_layers[0].non_empty_positions); + assert_sparse_count_constant(&list_layers[0].counts, 1); let chunk_values = sparse_metadata_visible_values(&page); assert_eq!( chunk_values.iter().map(|v| *v as u64).sum::(), @@ -10417,6 +11065,45 @@ mod tests { check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; } + #[tokio::test] + async fn test_sparse_layout_uses_range_position_set_for_dense_prefix() { + let arr = dense_prefix_sparse_i32_list_array(4096, 1024, 32); + let metadata = HashMap::from([( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_SPARSE.to_string(), + )]); + let field = arrow_schema::Field::new("c", arr.data_type().clone(), true) + .with_metadata(metadata.clone()); + + let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_3).await; + let sparse = sparse_layout_from_page(&page); + assert_eq!( + page.data.len(), + 2, + "range/constant sparse structural metadata should not write structural buffers" + ); + assert_eq!(sparse.num_visible_items, 1024 * 32); + let list_layers = sparse_list_layers(sparse); + assert_eq!(list_layers.len(), 1); + assert_eq!(list_layers[0].num_slots, 4096); + assert_eq!( + sparse_position_num(&list_layers[0].non_empty_positions), + 1024 + ); + assert_eq!(sparse_position_num(&list_layers[0].null_positions), 0); + assert_sparse_position_range(&list_layers[0].non_empty_positions, 0, 1024); + assert_sparse_count_constant(&list_layers[0].counts, 32); + + let test_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_range(0..16) + .with_range(1020..1030) + .with_range(3000..3010) + .with_indices(vec![0, 1, 1023, 1024, 4095]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; + } + #[tokio::test] async fn test_explicit_sparse_layout_requires_v2_3() { let arr = sparse_i32_list_array(4096, 64); @@ -10449,8 +11136,9 @@ mod tests { assert!(sparse.num_items > 65_535); let list_layers = sparse_list_layers(sparse); assert_eq!(list_layers.len(), 1); - assert_eq!(list_layers[0].num_non_empty, 2); - assert_eq!(list_layers[0].constant_count, 1); + assert_eq!(sparse_position_num(&list_layers[0].non_empty_positions), 2); + assert_sparse_position_explicit(&list_layers[0].non_empty_positions); + assert_sparse_count_constant(&list_layers[0].counts, 1); let test_cases = TestCases::default() .with_min_file_version(LanceFileVersion::V2_3) @@ -10474,11 +11162,13 @@ mod tests { let list_layers = sparse_list_layers(sparse); assert_eq!(list_layers.len(), 2); assert_eq!(list_layers[0].num_slots, 1); - assert_eq!(list_layers[0].num_non_empty, 1); - assert_eq!(list_layers[0].constant_count, 70_000); + assert_eq!(sparse_position_num(&list_layers[0].non_empty_positions), 1); + assert_sparse_position_all(&list_layers[0].non_empty_positions); + assert_sparse_count_constant(&list_layers[0].counts, 70_000); assert_eq!(list_layers[1].num_slots, 70_000); - assert_eq!(list_layers[1].num_non_empty, 1); - assert_eq!(list_layers[1].constant_count, 1); + assert_eq!(sparse_position_num(&list_layers[1].non_empty_positions), 1); + assert_sparse_position_range(&list_layers[1].non_empty_positions, 69_999, 1); + assert_sparse_count_constant(&list_layers[1].counts, 1); let test_cases = TestCases::default() .with_min_file_version(LanceFileVersion::V2_3) @@ -10741,9 +11431,14 @@ mod tests { matches!( pb21::sparse_structural_layer::Kind::try_from(layer.kind), Ok(pb21::sparse_structural_layer::Kind::SparseLayerList) - ) && layer.num_non_empty > 0 - && layer.num_nulls > 0 - && layer.count_compression.is_some() + ) && sparse_position_num(&layer.non_empty_positions) > 0 + && sparse_position_num(&layer.null_positions) > 0 + && layer.non_empty_positions.is_some() + && layer.null_positions.is_some() + && matches!( + layer.counts.as_ref().and_then(|set| set.counts.as_ref()), + Some(pb21::sparse_count_set::Counts::Explicit(_)) + ) }) }); assert!( @@ -10756,7 +11451,7 @@ mod tests { matches!( pb21::sparse_structural_layer::Kind::try_from(layer.kind), Ok(pb21::sparse_structural_layer::Kind::SparseLayerValidity) - ) && layer.num_nulls > 0 + ) && sparse_position_num(&layer.null_positions) > 0 }) }); assert!( @@ -10797,8 +11492,9 @@ mod tests { matches!( pb21::sparse_structural_layer::Kind::try_from(layer.kind), Ok(pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList) - ) && layer.constant_count == 2 - && layer.num_nulls > 0 + ) && layer.fixed_size_list_dimension == 2 + && sparse_position_num(&layer.null_positions) > 0 + && layer.null_positions.is_some() }) }); assert!( diff --git a/rust/lance-encoding/src/repdef.rs b/rust/lance-encoding/src/repdef.rs index 6a14fabb805..a2c2aa07d4b 100644 --- a/rust/lance-encoding/src/repdef.rs +++ b/rust/lance-encoding/src/repdef.rs @@ -133,24 +133,219 @@ pub(crate) struct SparseStructuralPlan { pub(crate) num_visible_items: u64, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum SparsePositionSet { + Empty, + All { len: u64 }, + Range { start: u64, len: u64 }, + Explicit(Vec), +} + +impl SparsePositionSet { + pub(crate) fn from_positions( + positions: Vec, + domain_len: u64, + label: &str, + ) -> Result { + if positions.is_empty() { + return Ok(Self::Empty); + } + for window in positions.windows(2) { + if window[0] >= window[1] { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} positions must be strictly increasing") + .into(), + )); + } + } + if let Some(position) = positions.iter().find(|position| **position >= domain_len) { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} position {} is outside layer with {} slots", + position, domain_len + ) + .into(), + )); + } + + let len = positions.len() as u64; + let first = positions[0]; + let last = positions[positions.len() - 1]; + if first == 0 && len == domain_len && last == domain_len - 1 { + return Ok(Self::All { len: domain_len }); + } + if last - first + 1 == len { + return Ok(Self::Range { start: first, len }); + } + Ok(Self::Explicit(positions)) + } + + pub(crate) fn empty() -> Self { + Self::Empty + } + + pub(crate) fn all(len: u64) -> Self { + if len == 0 { + Self::Empty + } else { + Self::All { len } + } + } + + pub(crate) fn range(start: u64, len: u64) -> Self { + if len == 0 { + Self::Empty + } else { + Self::Range { start, len } + } + } + + pub(crate) fn len(&self) -> u64 { + match self { + Self::Empty => 0, + Self::All { len } | Self::Range { len, .. } => *len, + Self::Explicit(positions) => positions.len() as u64, + } + } + + pub(crate) fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub(crate) fn deep_size(&self) -> usize { + match self { + Self::Explicit(positions) => positions.len() * std::mem::size_of::(), + Self::Empty | Self::All { .. } | Self::Range { .. } => 0, + } + } + + pub(crate) fn materialize(&self) -> Result> { + match self { + Self::Empty => Ok(Vec::new()), + Self::All { len } => Self::materialize_range(0, *len), + Self::Range { start, len } => Self::materialize_range(*start, *len), + Self::Explicit(positions) => Ok(positions.clone()), + } + } + + fn materialize_range(start: u64, len: u64) -> Result> { + let len_usize = usize::try_from(len).map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse structural position range length {} exceeds usize::MAX", + len + ) + .into(), + ) + })?; + let end = start.checked_add(len).ok_or_else(|| { + Error::invalid_input_source("Sparse structural position range overflows".into()) + })?; + let mut positions = Vec::with_capacity(len_usize); + positions.extend(start..end); + Ok(positions) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum SparseCountSet { + Empty, + Constant { value: u64, len: u64 }, + Explicit(Vec), +} + +impl SparseCountSet { + pub(crate) fn from_counts(counts: Vec) -> Self { + if counts.is_empty() { + return Self::Empty; + } + if let Some(first) = counts.first().copied() + && counts.iter().all(|count| *count == first) + { + return Self::Constant { + value: first, + len: counts.len() as u64, + }; + } + Self::Explicit(counts) + } + + pub(crate) fn constant(value: u64, len: u64) -> Self { + if len == 0 { + Self::Empty + } else { + Self::Constant { value, len } + } + } + + pub(crate) fn len(&self) -> u64 { + match self { + Self::Empty => 0, + Self::Constant { len, .. } => *len, + Self::Explicit(counts) => counts.len() as u64, + } + } + + pub(crate) fn deep_size(&self) -> usize { + match self { + Self::Explicit(counts) => counts.len() * std::mem::size_of::(), + Self::Empty | Self::Constant { .. } => 0, + } + } + + pub(crate) fn materialize(&self) -> Result> { + match self { + Self::Empty => Ok(Vec::new()), + Self::Constant { value, len } => { + let len = usize::try_from(*len).map_err(|_| { + Error::invalid_input_source( + "Sparse structural count set length exceeds usize::MAX".into(), + ) + })?; + Ok(vec![*value; len]) + } + Self::Explicit(counts) => Ok(counts.clone()), + } + } + + pub(crate) fn sum(&self) -> Result { + match self { + Self::Empty => Ok(0), + Self::Constant { value, len } => value.checked_mul(*len).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural constant count sum overflows: value={}, len={}", + value, len + ) + .into(), + ) + }), + Self::Explicit(counts) => counts.iter().try_fold(0_u64, |sum, count| { + sum.checked_add(*count).ok_or_else(|| { + Error::invalid_input_source("Sparse structural list count sum overflows".into()) + }) + }), + } + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum SparseStructuralLayerPlan { Validity { num_slots: u64, - null_positions: Vec, + null_positions: SparsePositionSet, }, List { num_slots: u64, num_child_slots: u64, - non_empty_positions: Vec, - counts: Vec, - constant_count: Option, - null_positions: Vec, + non_empty_positions: SparsePositionSet, + counts: SparseCountSet, + null_positions: SparsePositionSet, }, FixedSizeList { num_slots: u64, dimension: u64, - null_positions: Vec, + null_positions: SparsePositionSet, }, } @@ -206,7 +401,11 @@ impl SparseStructuralPlan { } } -fn validity_from_null_positions(num_slots: u64, null_positions: &[u64]) -> Result { +fn validity_from_null_positions( + num_slots: u64, + null_positions: &SparsePositionSet, +) -> Result { + let null_positions = null_positions.materialize()?; let mut validity = BooleanBufferBuilder::new(num_slots as usize); let mut null_iter = null_positions.iter().copied().peekable(); for slot in 0..num_slots { @@ -230,10 +429,13 @@ fn validity_from_null_positions(num_slots: u64, null_positions: &[u64]) -> Resul fn list_offsets_and_validity( num_slots: u64, - non_empty_positions: &[u64], - counts: &[u64], - null_positions: &[u64], + non_empty_positions: &SparsePositionSet, + counts: &SparseCountSet, + null_positions: &SparsePositionSet, ) -> Result<(OffsetBuffer, Option)> { + let non_empty_positions = non_empty_positions.materialize()?; + let counts = counts.materialize()?; + let null_positions = null_positions.materialize()?; if non_empty_positions.len() != counts.len() { return Err(Error::invalid_input_source( format!( @@ -378,20 +580,27 @@ impl SparseStructuralUnraveler { fn append_validity( validity: &mut BooleanBufferBuilder, num_slots: u64, - null_positions: &[u64], + null_positions: &SparsePositionSet, ) { - if null_positions.is_empty() { - validity.append_n(num_slots as usize, true); - return; - } - - let mut null_iter = null_positions.iter().copied().peekable(); - for slot in 0..num_slots { - let is_null = null_iter.peek().is_some_and(|null_pos| *null_pos == slot); - if is_null { - null_iter.next(); + match null_positions { + SparsePositionSet::Empty => validity.append_n(num_slots as usize, true), + SparsePositionSet::All { .. } => validity.append_n(num_slots as usize, false), + SparsePositionSet::Range { start, len } => { + validity.append_n(*start as usize, true); + validity.append_n(*len as usize, false); + let end = start + len; + validity.append_n((num_slots - end) as usize, true); + } + SparsePositionSet::Explicit(null_positions) => { + let mut null_iter = null_positions.iter().copied().peekable(); + for slot in 0..num_slots { + let is_null = null_iter.peek().is_some_and(|null_pos| *null_pos == slot); + if is_null { + null_iter.next(); + } + validity.append(!is_null); + } } - validity.append(!is_null); } } @@ -470,7 +679,7 @@ impl SparseStructuralUnraveler { .into(), )); } - let actual_child_slots = counts.iter().sum::(); + let actual_child_slots = counts.sum()?; if actual_child_slots != *num_child_slots { return Err(Error::invalid_input_source( format!( @@ -499,6 +708,9 @@ impl SparseStructuralUnraveler { return Ok(()); } + let non_empty_positions = non_empty_positions.materialize()?; + let counts = counts.materialize()?; + let null_positions = null_positions.materialize()?; let mut non_empty_iter = non_empty_positions .iter() .copied() @@ -1939,7 +2151,11 @@ impl RepDefBuilder { num_values, }) => SparseStructuralLayerPlan::Validity { num_slots: num_values as u64, - null_positions: Self::null_positions(validity.as_ref(), num_values), + null_positions: SparsePositionSet::from_positions( + Self::null_positions(validity.as_ref(), num_values), + num_values as u64, + "validity null", + )?, }, RawRepDef::Fsl(FslDesc { validity, @@ -1948,7 +2164,11 @@ impl RepDefBuilder { }) => SparseStructuralLayerPlan::FixedSizeList { num_slots: num_values as u64, dimension: dimension as u64, - null_positions: Self::null_positions(validity.as_ref(), num_values), + null_positions: SparsePositionSet::from_positions( + Self::null_positions(validity.as_ref(), num_values), + num_values as u64, + "fixed-size-list null", + )?, }, RawRepDef::Offsets(OffsetDesc { offsets, @@ -1977,17 +2197,20 @@ impl RepDefBuilder { })?); } } - let constant_count = counts - .first() - .copied() - .filter(|first| counts.iter().all(|count| count == first)); SparseStructuralLayerPlan::List { num_slots: num_values as u64, num_child_slots: *offsets.last().unwrap() as u64, - non_empty_positions, - counts, - constant_count, - null_positions, + non_empty_positions: SparsePositionSet::from_positions( + non_empty_positions, + num_values as u64, + "list non-empty", + )?, + counts: SparseCountSet::from_counts(counts), + null_positions: SparsePositionSet::from_positions( + null_positions, + num_values as u64, + "list null", + )?, } } })