diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index 4ca053d4fa6..3581ecb1edf 100644 --- a/docs/src/format/file/encoding.md +++ b/docs/src/format/file/encoding.md @@ -327,6 +327,144 @@ 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 + +Sparse pages require Lance 2.3. They represent flat or nested Arrow structure directly as slot-domain mappings instead +of dense repetition and definition events. Writers emit this layout only in files declared as 2.3. The layout is +identified only by `PageLayout`; field metadata does not identify the layout of an existing page. + +A domain is a layer-local integer coordinate space `[0, num_slots)`, and a slot is one element in that space. The +outer-most domain contains the page's top-level rows. Each layer maps its parent domain to the next layer's parent +domain, and the terminal child domain contains the leaf value slots stored in value chunks. + +Structural layers are ordered from outer-most to inner-most: + +- validity maps a nullable item or struct slot to valid or null +- list maps non-empty parent slots to variable-size child ranges +- fixed-size-list maps each parent slot to a child range of a fixed dimension + +The layer list may be empty for a flat, non-nullable leaf page. In that case the scheduling domain and +`num_visible_items` must be equal. The explicit writer currently emits its normalized all-valid layer even when a flat +page could use this shorter wire representation. + +A list slot that is valid and absent from `non_empty_positions` is an empty list. Maps use the same structural contract +as lists. The terminal child-domain size equals `SparseLayout.num_visible_items`. + +`SparseLayout.num_visible_items` is the number of leaf value slots encoded in value chunks. Null leaf slots count +because they still occupy positions in Arrow's leaf value buffer; a nullable primitive with 100 slots, including 30 +nulls, has 100 visible items. `SparseLayout.num_items` is the number of entries in the equivalent dense repetition and +definition stream. It equals `num_visible_items` plus one structural placeholder for every list slot without children. +The first layer's `num_slots` is the logical top-level row count used for projection. + +Position sets have four semantic representations: `empty`, `all`, one non-empty `range`, or an `explicit` +delta-compressed `u64` buffer. Count sets are `empty`, one positive `constant` value, or an `explicit` compressed `u64` +buffer. Every layer has a `SparseValiditySet` whose meaning is explicit: + +- `SPARSE_VALIDITY_NULL_POSITIONS`: stored positions are null and all other positions are valid +- `SPARSE_VALIDITY_VALID_POSITIONS`: stored positions are valid and all other positions are null + +The unspecified validity meaning is invalid. Both polarities are part of the wire contract and have identical Arrow +semantics after normalization. + +#### Writer Selection + +The Lance 2.3 writer selects this layout when a field sets `lance-encoding:structural-encoding=sparse`. The same request +is an input error for earlier file versions. This metadata controls writer selection only: readers always use +`PageLayout` to determine the layout of an encoded page and must not use field metadata for that decision. + +The default Lance 2.3 writer policy remains unchanged, as do explicit `miniblock` and `fullzip` requests. Writers +normalize Arrow validity and list structure once, then use that semantic structure for either dense +repetition/definition serialization or sparse position/count serialization. All-valid layers use null positions plus +`empty`; all-null layers use valid positions plus `empty`. Other layers choose the validity polarity with the lower +semantic encoded cost, with ties using null positions. + +Pages without a value payload keep the existing canonical `ConstantLayout`: structural-only types such as an empty +struct, and leaf pages whose visible values are all null, do not emit `SparseLayout`. An explicitly sparse page with +at least one non-null visible value does emit `SparseLayout`, even when all non-null values are equal. This boundary +avoids introducing a second structural-only representation without evidence that it improves the existing constant +encoding. + +#### Buffers and Selective Reads + +A sparse page contains the following physical buffers: + +| Buffer | Contents | +| ------ | -------- | +| 0 | Value chunk metadata, one 8-byte entry per chunk | +| 1 | Mini-block compressed value chunks without repetition or definition levels | +| 2+ | One buffer for each explicit position or count set, in structural-layer field order | + +Each value chunk metadata entry stores `(chunk_size / 8) - 1` as little-endian `u32`, followed by its visible value +count as little-endian `u32`. Chunk sizes must be positive multiples of 8 and fit this representation. The sum of +chunk sizes must equal buffer 1 exactly and the sum of chunk value counts must equal `num_visible_items`. A value chunk +contains at most 32,768 visible values. `num_buffers` describes the number of value buffers inside every chunk and +excludes the structural buffers. + +General-compressed sparse buffers use the existing length-prefixed LZ4 or Zstd representation and must not contain +another general-compression wrapper. SparseLayout does not impose additional size or descriptor-complexity limits on +otherwise representable buffers. + +Readers normalize structural metadata once, project requested top-level ranges through each layer, and read only value +chunks that intersect the resulting leaf ranges. When no leaf range remains, readers rebuild offsets and validity from +the structural plan without reading buffer 1. + +#### Caching and Point Reads + +Reader initialization loads buffer 0 and every explicit structural buffer, then validates and normalizes them into a +cached page plan. The cached state contains parsed value-chunk descriptors and prefix offsets, decoded semantic +position/count sets, validity, and the ordered structural layers. It does not contain value payload bytes from buffer +1. The plan is cached per field and page and reused by later scans, range reads, and takes. + +After that plan is cached, reading one primitive leaf value reads only the value chunk that contains it. A cold read +first loads the structural metadata and then the intersecting value chunk. Reading one top-level list or +fixed-size-list value may intersect multiple leaf chunks and reads each intersecting chunk. A selection whose projected +structure contains no leaf slots reads no value chunk. + +#### Validation + +Readers must reject malformed sparse metadata instead of inferring or repairing it. Required checks include: + +- physical buffer count, chunk-count bounds, and every checked offset/size range +- first-layer row domain, adjacent parent/child domain chaining, and terminal visible-value domain +- semantic set cardinality, explicit position ordering and bounds, and validity meaning +- exact `num_items` +- list non-empty positions being valid, count cardinality, positive counts, and child-count sum +- fixed-size-list dimension and checked child-domain multiplication +- value chunk byte/value sums, size representation and alignment, general-compression headers, descriptor buffer + count, and complete chunk consumption + +```protobuf +%%% proto.message.SparseLayout %%% +``` + +```protobuf +%%% proto.message.SparseStructuralLayer %%% +``` + +```protobuf +%%% proto.message.SparseValidityLayer %%% +``` + +```protobuf +%%% proto.message.SparseListLayer %%% +``` + +```protobuf +%%% proto.message.SparseFixedSizeListLayer %%% +``` + +```protobuf +%%% proto.message.SparseValiditySet %%% +``` + +```protobuf +%%% proto.message.SparsePositionSet %%% +``` + +```protobuf +%%% proto.message.SparseCountSet %%% +``` + ### Constant Page Layout This layout is used when all (visible) values in the page are the same scalar value. @@ -549,7 +687,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 | Select a structural encoding; `sparse` requires Lance 2.3. | ### Configuration Details diff --git a/docs/src/format/file/versioning.md b/docs/src/format/file/versioning.md index 2add82ef2fa..8024a46fbc8 100644 --- a/docs/src/format/file/versioning.md +++ b/docs/src/format/file/versioning.md @@ -5,10 +5,10 @@ major number is changed when the file format itself is modified while the minor strategy is modified. Newer versions will typically have better performance and compression but may not be readable by older versions of Lance. -In addition, the `next` alias points to an unstable format version and should not be used for production use cases. -Breaking changes could be made to unstable encodings and that would mean that files written with these encodings are -no longer readable by any newer versions of Lance. The `next` version should only be used for experimentation and -benchmarking upcoming features. +Any version explicitly labeled unstable, including the current 2.3 format and the `next` alias, should not be used for +production use cases. Unstable formats have no compatibility guarantee: breaking encoding changes may make files +written by one Lance build unreadable by later builds. They should only be used for experimentation and benchmarking +upcoming features. The `stable` and `next` aliases are resolved by the specific Lance release you are using. During a format rollout (for example, 2.3), prefer explicit version pinning for deterministic behavior across environments. @@ -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 | Unspecified | Adds sparse structural pages and other experimental encodings. | | 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..51427332063 100644 --- a/protos/encodings_v2_1.proto +++ b/protos/encodings_v2_1.proto @@ -147,6 +147,128 @@ message FullZipLayout { repeated RepDefLayer layers = 8; } +// A layout used for sparse flat or nested pages where Arrow structure is represented directly +// in layer-local slot domains instead of as dense repetition / definition events. +// +// Structural layers are ordered from outer-most to inner-most. Values remain mini-block +// compressed and are split into independently readable chunks. +message SparseLayout { + // Description of the compression of values. + CompressiveEncoding value_compression = 1; + // Number of value buffers in each mini-block chunk. This does not include structural buffers. + uint64 num_buffers = 2; + // Number of entries in the equivalent dense repetition / definition stream. This equals + // num_visible_items plus one structural placeholder for every list slot without children. + // Null leaf slots count as visible items because they still occupy positions in Arrow's + // leaf value buffer. For example, a nullable primitive with 100 slots, 30 of them null, + // has num_items = num_visible_items = 100. + uint64 num_items = 3; + // Number of leaf value slots encoded in the value chunks, including null leaf slots. + uint64 num_visible_items = 4; + // If true, chunk-local value buffer sizes use u32. Otherwise they use u16. + bool has_large_chunk = 5; + // Structural layers ordered from outer-most to inner-most. This may be empty for a flat, + // non-nullable leaf page whose scheduling domain equals num_visible_items. + repeated SparseStructuralLayer structural_layers = 6; +} + +// A domain is a layer-local integer coordinate space [0, num_slots). A slot is one +// element in that space. The outer-most domain is the page's top-level rows; each +// layer's child domain is the next layer's parent domain, and the terminal child +// domain contains num_visible_items leaf value slots. +message SparseStructuralLayer { + // Exactly one layer kind is required. + oneof layer { + SparseValidityLayer validity = 1; + SparseListLayer list = 2; + SparseFixedSizeListLayer fixed_size_list = 3; + } +} + +message SparseValidityLayer { + // Number of nullable item or struct slots in this layer's parent and child domain. + uint64 num_slots = 1; + // Validity for the slots in this layer. + SparseValiditySet validity = 2; +} + +message SparseListLayer { + // Number of list, large-list, or map slots in this layer's parent domain. + uint64 num_slots = 1; + // Number of slots in this layer's child domain. + uint64 num_child_slots = 2; + // Non-empty parent slots. Valid parent slots absent from this set are empty lists. + SparsePositionSet non_empty_positions = 3; + // Positive child counts corresponding one-for-one with non_empty_positions. + SparseCountSet counts = 4; + // Validity for the parent slots in this layer. + SparseValiditySet validity = 5; +} + +message SparseFixedSizeListLayer { + // Number of fixed-size-list slots in this layer's parent domain. + uint64 num_slots = 1; + // Number of children per parent slot. The child domain has num_slots * dimension slots. + uint64 dimension = 2; + // Validity for the parent slots in this layer. + SparseValiditySet validity = 3; +} + +message SparseValiditySet { + enum Meaning { + SPARSE_VALIDITY_UNSPECIFIED = 0; + // Stored positions are null; all other positions are valid. + SPARSE_VALIDITY_NULL_POSITIONS = 1; + // Stored positions are valid; all other positions are null. + SPARSE_VALIDITY_VALID_POSITIONS = 2; + } + + Meaning meaning = 1; + SparsePositionSet positions = 2; +} + +message SparsePositionEmpty {} + +message SparsePositionAll {} + +message SparsePositionRange { + uint64 start = 1; + uint64 length = 2; +} + +message SparsePositionSet { + oneof positions { + // Delta-compressed u64 positions. Cardinality is num_positions. + CompressiveEncoding explicit = 1; + // One contiguous, non-empty range. + SparsePositionRange range = 2; + // Every position in the domain. + SparsePositionAll all = 3; + // No positions in the domain. + SparsePositionEmpty empty = 4; + } + // Semantic cardinality of this set. + uint64 num_positions = 5; +} + +message SparseCountEmpty {} + +message SparseCountConstant { + // Child count shared by every non-empty list slot. + uint64 value = 1; +} + +message SparseCountSet { + oneof counts { + // Compressed u64 child counts. Cardinality comes from the containing position set. + CompressiveEncoding explicit = 1; + // One positive child count shared by every non-empty list slot. + SparseCountConstant constant = 2; + // No counts; valid only when there are no non-empty list slots. + SparseCountEmpty empty = 3; + } +} + // 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 +328,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 sparse structural layout. This variant requires file version 2.3 or later. + SparseLayout sparse_layout = 5; } } diff --git a/protos/file2.proto b/protos/file2.proto index da0b1d5e96c..650a1568da6 100644 --- a/protos/file2.proto +++ b/protos/file2.proto @@ -207,4 +207,4 @@ message ColumnMetadata { // // This file format is extremely minimal. It is a building block for // creating more useful readers and writers and not terribly useful by itself. -// Other protobuf files will describe how this can be extended. \ No newline at end of file +// Other protobuf files will describe how this can be extended. 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/decoder.rs b/rust/lance-encoding/src/decoder.rs index aea3575dcb1..108a848db64 100644 --- a/rust/lance-encoding/src/decoder.rs +++ b/rust/lance-encoding/src/decoder.rs @@ -2717,8 +2717,9 @@ pub trait DecodeArrayTask: Send { impl DecodeArrayTask for Box { fn decode(self: Box) -> Result<(ArrayRef, u64)> { - StructuralDecodeArrayTask::decode(*self) - .map(|decoded_array| (decoded_array.array, decoded_array.data_size)) + let decoded_array = StructuralDecodeArrayTask::decode(*self)?; + decoded_array.repdef.ensure_exhausted()?; + Ok((decoded_array.array, decoded_array.data_size)) } } @@ -2740,10 +2741,7 @@ impl NextDecodeTask { // suggesting the user try a smaller batch size. #[instrument(name = "task_to_batch", level = "debug", skip_all)] fn into_batch(self, emitted_batch_size_warning: Arc) -> Result<(RecordBatch, u64)> { - let (struct_arr, data_size) = self - .task - .decode() - .map_err(|e| Error::internal(format!("Error decoding batch: {}", e)))?; + let (struct_arr, data_size) = self.task.decode()?; let batch = RecordBatch::from(struct_arr.as_struct()); if data_size > BATCH_SIZE_BYTES_WARNING { emitted_batch_size_warning.call_once(|| { @@ -2983,6 +2981,25 @@ mod tests { } } + struct InvalidInputDecodeTask; + + impl DecodeArrayTask for InvalidInputDecodeTask { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + Err(Error::invalid_input_source("malformed sparse page".into())) + } + } + + #[test] + fn next_decode_task_preserves_invalid_input_errors() { + let err = NextDecodeTask { + task: Box::new(InvalidInputDecodeTask), + num_rows: 0, + } + .into_batch(Arc::new(Once::new())) + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + } + #[test] fn test_read_zero_dimension_fsl_errors_instead_of_panicking() { // Simulates reading a column whose stored schema declares a diff --git a/rust/lance-encoding/src/encodings/logical/fixed_size_list.rs b/rust/lance-encoding/src/encodings/logical/fixed_size_list.rs index 9e8e3e109ea..1308288f226 100644 --- a/rust/lance-encoding/src/encodings/logical/fixed_size_list.rs +++ b/rust/lance-encoding/src/encodings/logical/fixed_size_list.rs @@ -65,7 +65,7 @@ impl FieldEncoder for FixedSizeListStructuralEncoder { } else { deep_copy_nulls(array.nulls()) }; - repdef.add_fsl(validity.clone(), dimension, num_rows as usize); + repdef.add_fsl(validity.clone(), dimension, fsl_arr.len()); // FSL forces child elements to exist even under null rows. Normalize any // nested lists under null FSL rows to null empty lists. @@ -255,7 +255,7 @@ impl StructuralDecodeArrayTask for StructuralFixedSizeListDecodeTask { match &self.data_type { DataType::FixedSizeList(child_field, dimension) => { let num_rows = self.num_rows as usize; - let validity = repdef.unravel_fsl_validity(num_rows, *dimension as usize); + let validity = repdef.unravel_fsl_validity(num_rows, *dimension as usize)?; let fsl_array = arrow_array::FixedSizeListArray::try_new( child_field.clone(), *dimension, diff --git a/rust/lance-encoding/src/encodings/logical/list.rs b/rust/lance-encoding/src/encodings/logical/list.rs index 250eb476671..c9bdcc0bc87 100644 --- a/rust/lance-encoding/src/encodings/logical/list.rs +++ b/rust/lance-encoding/src/encodings/logical/list.rs @@ -344,6 +344,7 @@ mod tests { } } crate::format::pb21::page_layout::Layout::BlobLayout(_) => {} + crate::format::pb21::page_layout::Layout::SparseLayout(_) => {} } } diff --git a/rust/lance-encoding/src/encodings/logical/map.rs b/rust/lance-encoding/src/encodings/logical/map.rs index b5172d8c189..be153030105 100644 --- a/rust/lance-encoding/src/encodings/logical/map.rs +++ b/rust/lance-encoding/src/encodings/logical/map.rs @@ -223,7 +223,8 @@ impl StructuralDecodeArrayTask for StructuralMapDecodeTask { .clone(); // Build the MapArray from offsets, entries, validity, and keys_sorted - let map_array = MapArray::new(entries_field, offsets, entries, validity, keys_sorted); + let map_array = MapArray::try_new(entries_field, offsets, entries, validity, keys_sorted) + .map_err(|error| Error::invalid_input_source(error.to_string().into()))?; Ok(DecodedArray { array: Arc::new(map_array), @@ -244,7 +245,13 @@ mod tests { use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field, Fields}; + use crate::decoder::{DecodedArray, StructuralDecodeArrayTask}; use crate::encoder::{ColumnIndexSequence, EncodingOptions, default_encoding_strategy}; + use crate::encodings::logical::primitive::sparse::{ + SparseCountSet, SparsePositionSet, SparseStructuralLayerPlan, SparseStructuralPlan, + SparseValidityMeaning, SparseValiditySet, + }; + use crate::repdef::{CompositeRepDefUnraveler, RepDefUnraveler}; use crate::{ testing::{TestCases, check_round_trip_encoding_of_data}, version::LanceFileVersion, @@ -252,6 +259,8 @@ mod tests { use arrow_schema::Field as ArrowField; use lance_core::datatypes::Field as LanceField; + use super::StructuralMapDecodeTask; + fn make_map_type(key_type: DataType, value_type: DataType) -> DataType { // Note: Arrow MapBuilder uses "keys" and "values" as field names (plural) let entries = Field::new( @@ -265,6 +274,70 @@ mod tests { DataType::Map(Arc::new(entries), false) } + #[derive(Debug)] + struct StaticMapEntriesTask { + entries: StructArray, + repdef: CompositeRepDefUnraveler, + } + + impl StructuralDecodeArrayTask for StaticMapEntriesTask { + fn decode(self: Box) -> lance_core::Result { + let Self { entries, repdef } = *self; + Ok(DecodedArray { + array: Arc::new(entries), + repdef, + data_size: 0, + }) + } + } + + #[test] + fn malformed_sparse_map_entries_return_invalid_input() { + let entry_fields = Fields::from(vec![ + Field::new("keys", DataType::Int32, false), + Field::new("values", DataType::Int32, true), + ]); + let entries = StructArray::try_new( + entry_fields.clone(), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![2])), + ], + Some(NullBuffer::from(vec![false])), + ) + .unwrap(); + let validity = SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions: SparsePositionSet::Empty, + }; + let plan = SparseStructuralPlan { + layers: vec![SparseStructuralLayerPlan::List { + num_slots: 1, + num_child_slots: 1, + non_empty_positions: SparsePositionSet::All { len: 1 }, + counts: SparseCountSet::Constant { value: 1, len: 1 }, + validity, + }], + num_items: 1, + num_visible_items: 1, + }; + let child_task = StaticMapEntriesTask { + entries, + repdef: CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new_sparse(plan)]), + }; + let map_type = DataType::Map( + Arc::new(Field::new("entries", DataType::Struct(entry_fields), false)), + false, + ); + + let Err(err) = + Box::new(StructuralMapDecodeTask::new(Box::new(child_task), map_type)).decode() + else { + panic!("expected malformed map entries to be rejected"); + }; + assert!(matches!(err, lance_core::Error::InvalidInput { .. })); + } + #[test_log::test(tokio::test)] async fn test_simple_map() { // Create a simple Map diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 8b6eed5d757..5429067efe5 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}, @@ -92,6 +93,7 @@ pub mod constant; pub mod dict; pub mod fullzip; pub mod miniblock; +pub(crate) mod sparse; const FILL_BYTE: u8 = 0xFE; const DEFAULT_DICT_DIVISOR: u64 = 2; @@ -3349,6 +3351,16 @@ impl StructuralPrimitiveFieldScheduler { mini_block, decompressors, )?), + Layout::SparseLayout(sparse_layout) => { + Box::new(sparse::SparseStructuralScheduler::try_new( + &page_info.buffer_offsets_and_sizes, + page_info.priority, + page_info.num_rows, + target_field.data_type(), + sparse_layout, + decompressors, + )?) + } Layout::FullZipLayout(full_zip) => { let mut scheduler = FullZipScheduler::try_new( &page_info.buffer_offsets_and_sizes, @@ -3576,25 +3588,34 @@ impl StructuralCompositeDecodeArrayTask { fn restore_validity( array: Arc, unraveler: &mut CompositeRepDefUnraveler, - ) -> Arc { - let validity = unraveler.unravel_validity(array.len()); + ) -> Result> { + let validity = unraveler.unravel_validity(array.len())?; let Some(validity) = validity else { - return array; + return Ok(array); }; if array.data_type() == &DataType::Null { // We unravel from a null array but we don't add the null buffer because arrow-rs doesn't like it - return array; + return Ok(array); + } + if validity.len() != array.len() { + return Err(Error::invalid_input_source( + format!( + "Structural validity has {} entries for an array with {} values", + validity.len(), + array.len() + ) + .into(), + )); } - assert_eq!(validity.len(), array.len()); - // SAFETY: We've should have already asserted the buffers are all valid, we are just - // adding null buffers to the array here - make_array(unsafe { + // SAFETY: The array buffers have already been validated and the null buffer length + // matches the array. We are only attaching the null buffer here. + Ok(make_array(unsafe { array .to_data() .into_builder() .nulls(Some(validity)) .build_unchecked() - }) + })) } } @@ -3620,7 +3641,7 @@ impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask { let array = arrow_select::concat::concat(&array_refs)?; let mut repdef = CompositeRepDefUnraveler::new(unravelers); - let array = Self::restore_validity(array, &mut repdef); + let array = Self::restore_validity(array, &mut repdef)?; Ok(DecodedArray { array, @@ -3787,18 +3808,24 @@ struct DictEncodingBudget { max_encoded_size: usize, } +enum PrimitivePageStructure { + Dense { + repdef: SerializedRepDefs, + unsplittable_miniblock_levels: Option, + }, + Sparse(sparse::SparseStructuralPlan), +} + // A primitive page after optional structural splitting. struct PrimitivePageData { // Arrow leaf arrays that contain this page's visible values. arrays: Vec, - // Repetition / definition levels aligned to this page. - repdef: SerializedRepDefs, + // Structural representation aligned to this page. + structure: PrimitivePageStructure, // Top-level row number of the first row in this page. 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, } // Immutable encoder state shared by per-page encode tasks. @@ -3833,6 +3860,19 @@ impl PrimitiveStructuralEncoder { field: Field, encoding_metadata: Arc>, ) -> Result { + let requests_sparse = encoding_metadata + .get(STRUCTURAL_ENCODING_META_KEY) + .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE)); + if requests_sparse && options.version.resolve() < LanceFileVersion::V2_3 { + return Err(Error::invalid_input_source( + format!( + "Field '{}' requests sparse structural encoding, which requires Lance file format 2.3+; current version is {}", + field.name, + options.version.resolve() + ) + .into(), + )); + } Ok(Self { accumulation_queue: AccumulationQueue::new( options.cache_bytes_per_column, @@ -4362,7 +4402,7 @@ impl PrimitiveStructuralEncoder { return Ok(None); } let mut validity = BooleanBufferBuilder::new(num_values); - unraveler.unravel_validity(&mut validity); + unraveler.unravel_validity(&mut validity)?; Ok(Some(validity.finish())) } @@ -5302,19 +5342,23 @@ impl PrimitiveStructuralEncoder { if plan == StructuralPagePlan::Fits { return Ok(vec![PrimitivePageData { arrays, - repdef, + structure: PrimitivePageStructure::Dense { + repdef, + unsplittable_miniblock_levels: None, + }, row_number, num_rows, - unsplittable_miniblock_levels: None, }]); } if let StructuralPagePlan::UnsplittableOverBudget(num_levels) = plan { return Ok(vec![PrimitivePageData { arrays, - repdef, + structure: PrimitivePageStructure::Dense { + repdef, + unsplittable_miniblock_levels: Some(num_levels), + }, row_number, num_rows, - unsplittable_miniblock_levels: Some(num_levels), }]); } @@ -5328,10 +5372,12 @@ impl PrimitiveStructuralEncoder { let repdef = Self::slice_repdef(&repdef, split.level_range); pages.push(PrimitivePageData { arrays, - repdef, + structure: PrimitivePageStructure::Dense { + repdef, + unsplittable_miniblock_levels: None, + }, row_number: row_number + split.row_start, num_rows: split.num_rows, - unsplittable_miniblock_levels: None, }); } Ok(pages) @@ -5350,13 +5396,37 @@ impl PrimitiveStructuralEncoder { } = ctx; let PrimitivePageData { arrays, - repdef, + structure, row_number, num_rows, - unsplittable_miniblock_levels, } = page; let num_values = arrays.iter().map(|arr| arr.len() as u64).sum(); + let (repdef, unsplittable_miniblock_levels) = match structure { + PrimitivePageStructure::Dense { + repdef, + unsplittable_miniblock_levels, + } => (repdef, unsplittable_miniblock_levels), + PrimitivePageStructure::Sparse(plan) => { + log::debug!( + "Encoding column {} with {} visible items ({} rows) using sparse layout", + column_idx, + num_values, + num_rows + ); + return sparse::writer::encode_page( + column_idx, + &field, + compression_strategy.as_ref(), + DataBlock::from_arrays(&arrays, num_values), + plan, + row_number, + num_rows, + support_large_chunk, + ); + } + }; + if num_values == 0 { // This page contains only structural events, such as empty/null list rows. // The existing complex-null layout stores the rep/def stream without value buffers. @@ -5663,19 +5733,38 @@ 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 (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 normalized = RepDefBuilder::normalize(repdefs); + let requests_sparse = self + .encoding_metadata + .get(STRUCTURAL_ENCODING_META_KEY) + .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE)); + let sparse_plan = requests_sparse + .then(|| sparse::writer::plan(&normalized, num_values)) + .transpose()?; + let pages = match sparse_plan { + Some(plan) if !sparse::writer::uses_constant_layout(&plan, &self.field) => { + vec![PrimitivePageData { + arrays, + structure: PrimitivePageStructure::Sparse(plan), + row_number, + num_rows, + }] + } + _ => { + let (repdef, structural_plan) = normalized.serialize_with_structural_plan( + miniblock::max_repdef_levels_per_chunk, + num_rows, + num_values, + )?; + 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 { diff --git a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs index 1cf3b9bf581..edfba526670 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs @@ -19,7 +19,8 @@ use lance_core::Result; pub const MAX_MINIBLOCK_BYTES: u64 = 8 * 1024 - 6; const DEFAULT_MAX_MINIBLOCK_VALUES: u64 = 4096; -const MAX_CONFIGURABLE_MINIBLOCK_VALUES: u64 = 32768; +/// Maximum number of values that any mini-block decoder accepts from page metadata. +pub(crate) const MAX_CONFIGURABLE_MINIBLOCK_VALUES: u64 = 32768; fn parse_max_miniblock_values() -> u64 { let val = std::env::var("LANCE_MINIBLOCK_MAX_VALUES") diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs new file mode 100644 index 00000000000..68becc0014f --- /dev/null +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs @@ -0,0 +1,5614 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use super::*; +use arrow_array::new_empty_array; +use arrow_buffer::ArrowNativeType; + +fn invalid_enum( + value: std::result::Result, + label: &str, +) -> Result { + value.map_err(|error| { + Error::invalid_input_source( + format!("Sparse structural {label} has an invalid enum value: {error}").into(), + ) + }) +} + +pub(super) mod writer; + +fn usize_from_u64(value: u64, label: &str) -> Result { + usize::try_from(value).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} {value} exceeds usize::MAX").into(), + ) + }) +} + +/// 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 struct SparseStructuralPlan { + pub(crate) layers: Vec, + pub(crate) num_items: u64, + pub(crate) num_visible_items: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub 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) { + let [previous, current] = window else { + continue; + }; + if previous >= current { + 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 = u64::try_from(positions.len()).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} position count exceeds u64::MAX").into(), + ) + })?; + let first = positions.first().copied().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} positions are unexpectedly empty").into(), + ) + })?; + let last = positions.last().copied().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} positions are unexpectedly empty").into(), + ) + })?; + if first == 0 && len == domain_len && domain_len > 0 && 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) + } + + fn contains(&self, position: u64) -> bool { + match self { + Self::Empty => false, + Self::All { len } => position < *len, + Self::Range { start, len } => { + position >= *start && position < start.saturating_add(*len) + } + Self::Explicit(positions) => positions.binary_search(&position).is_ok(), + } + } + + fn is_subset_of(&self, other: &Self, domain_len: u64) -> Result { + self.validate_domain(domain_len, "subset")?; + other.validate_domain(domain_len, "superset")?; + Ok(match self { + Self::Empty => true, + Self::All { .. } => other.len() == domain_len, + Self::Range { start, len } => { + let end = start.checked_add(*len).ok_or_else(|| { + Error::invalid_input_source("Sparse structural subset range overflows".into()) + })?; + match other { + Self::All { .. } => true, + Self::Range { + start: other_start, + len: other_len, + } => { + let other_end = other_start.saturating_add(*other_len); + *start >= *other_start && end <= other_end + } + Self::Explicit(positions) => { + let first = positions.partition_point(|position| *position < *start); + let last = positions.partition_point(|position| *position < end); + u64::try_from(last.saturating_sub(first)).ok() == Some(*len) + } + Self::Empty => false, + } + } + Self::Explicit(positions) => positions.iter().all(|position| other.contains(*position)), + }) + } + + fn is_disjoint(&self, other: &Self, domain_len: u64) -> Result { + self.validate_domain(domain_len, "first disjoint set")?; + other.validate_domain(domain_len, "second disjoint set")?; + let (smaller, larger) = if self.len() <= other.len() { + (self, other) + } else { + (other, self) + }; + Ok(match smaller { + Self::Empty => true, + Self::All { .. } => larger.is_empty(), + Self::Range { start, len } => { + let end = start.saturating_add(*len); + match larger { + Self::Empty => true, + Self::All { .. } => false, + Self::Range { + start: other_start, + len: other_len, + } => end <= *other_start || other_start.saturating_add(*other_len) <= *start, + Self::Explicit(positions) => { + let index = positions.partition_point(|position| *position < *start); + positions.get(index).is_none_or(|position| *position >= end) + } + } + } + Self::Explicit(positions) => { + positions.iter().all(|position| !larger.contains(*position)) + } + }) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SparseValidityMeaning { + NullPositions, + ValidPositions, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SparseValiditySet { + pub(crate) meaning: SparseValidityMeaning, + pub(crate) positions: SparsePositionSet, +} + +impl SparseValiditySet { + pub(crate) fn deep_size(&self) -> usize { + self.positions.deep_size() + } + + fn contains_only_valid_positions( + &self, + positions: &SparsePositionSet, + num_slots: u64, + ) -> Result { + match self.meaning { + SparseValidityMeaning::NullPositions => { + positions.is_disjoint(&self.positions, num_slots) + } + SparseValidityMeaning::ValidPositions => { + positions.is_subset_of(&self.positions, num_slots) + } + } + } + + fn append_to(&self, validity: &mut BooleanBufferBuilder, num_slots: u64) -> Result<()> { + self.positions.validate_domain(num_slots, "validity")?; + let num_slots_usize = usize_from_u64(num_slots, "validity slot count")?; + match (self.meaning, &self.positions) { + (SparseValidityMeaning::NullPositions, SparsePositionSet::Empty) => { + validity.append_n(num_slots_usize, true); + } + (SparseValidityMeaning::ValidPositions, SparsePositionSet::Empty) => { + validity.append_n(num_slots_usize, false); + } + (SparseValidityMeaning::NullPositions, SparsePositionSet::All { .. }) => { + validity.append_n(num_slots_usize, false); + } + (SparseValidityMeaning::ValidPositions, SparsePositionSet::All { .. }) => { + validity.append_n(num_slots_usize, true); + } + (meaning, SparsePositionSet::Range { start, len }) => { + let range_end = start.checked_add(*len).ok_or_else(|| { + Error::invalid_input_source("Sparse structural validity range overflows".into()) + })?; + let default_valid = matches!(meaning, SparseValidityMeaning::NullPositions); + let range_valid = !default_valid; + validity.append_n( + usize_from_u64(*start, "validity range start")?, + default_valid, + ); + validity.append_n(usize_from_u64(*len, "validity range length")?, range_valid); + validity.append_n( + usize_from_u64(num_slots - range_end, "validity range tail")?, + default_valid, + ); + } + (_, SparsePositionSet::Explicit(_)) => { + let mut cursor = SparseValidityCursor::new(self, num_slots, "validity")?; + for slot in 0..num_slots { + validity.append(cursor.is_valid(slot)?); + } + cursor.finish()?; + } + } + Ok(()) + } +} + +impl SparsePositionSet { + fn validate_domain(&self, domain_len: u64, label: &str) -> Result<()> { + match self { + Self::Empty => {} + Self::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(), + )); + } + } + Self::Range { start, len } => { + let end = start.checked_add(*len).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} range overflows").into(), + ) + })?; + if end > domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} range {}..{} is outside domain {}", + start, end, domain_len + ) + .into(), + )); + } + } + Self::Explicit(positions) => { + for window in positions.windows(2) { + let [previous, current] = window else { + continue; + }; + if previous >= current { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} positions must be strictly increasing" + ) + .into(), + )); + } + } + if let Some(position) = positions.last() + && *position >= domain_len + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} position {} is outside layer with {} slots", + position, domain_len + ) + .into(), + )); + } + } + } + Ok(()) + } +} + +struct SparsePositionSetCursor<'a> { + set: &'a SparsePositionSet, + explicit: Option>>, +} + +impl<'a> SparsePositionSetCursor<'a> { + fn new(set: &'a SparsePositionSet, domain_len: u64, label: &str) -> Result { + set.validate_domain(domain_len, label)?; + let explicit = match set { + SparsePositionSet::Explicit(positions) => Some(positions.iter().peekable()), + _ => None, + }; + Ok(Self { set, explicit }) + } + + fn contains(&mut self, slot: u64) -> Result { + Ok(match self.set { + SparsePositionSet::Empty => false, + SparsePositionSet::All { .. } => true, + SparsePositionSet::Range { start, len } => { + slot >= *start && slot < start.saturating_add(*len) + } + SparsePositionSet::Explicit(_) => { + let iter = self.explicit.as_mut().ok_or_else(|| { + Error::internal("Sparse structural explicit cursor is missing".to_string()) + })?; + if let Some(position) = iter.peek() + && **position < slot + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural explicit position {} was skipped before slot {}", + **position, slot + ) + .into(), + )); + } + if iter.peek().is_some_and(|position| **position == slot) { + iter.next(); + true + } else { + false + } + } + }) + } + + fn finish(&mut self) -> Result<()> { + if let Some(iter) = self.explicit.as_mut() + && let Some(position) = iter.next() + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural explicit position {} was not consumed", + position + ) + .into(), + )); + } + Ok(()) + } +} + +struct SparseValidityCursor<'a> { + meaning: SparseValidityMeaning, + positions: SparsePositionSetCursor<'a>, +} + +impl<'a> SparseValidityCursor<'a> { + fn new(validity: &'a SparseValiditySet, domain_len: u64, label: &str) -> Result { + Ok(Self { + meaning: validity.meaning, + positions: SparsePositionSetCursor::new(&validity.positions, domain_len, label)?, + }) + } + + fn is_valid(&mut self, slot: u64) -> Result { + let is_stored = self.positions.contains(slot)?; + Ok(match self.meaning { + SparseValidityMeaning::NullPositions => !is_stored, + SparseValidityMeaning::ValidPositions => is_stored, + }) + } + + fn finish(&mut self) -> Result<()> { + self.positions.finish() + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SparseCountSet { + Empty, + Constant { + value: u64, + len: u64, + }, + Explicit { + counts: Arc<[u64]>, + offsets: Arc<[u64]>, + }, +} + +impl SparseCountSet { + pub(crate) fn from_counts(counts: Vec) -> Result { + if counts.is_empty() { + return Ok(Self::Empty); + } + if let Some(first) = counts.first().copied() + && counts.iter().all(|count| *count == first) + { + return Ok(Self::Constant { + value: first, + len: counts.len() as u64, + }); + } + let offsets = offsets_from_counts(&counts)?; + Ok(Self::Explicit { + counts: counts.into(), + offsets: offsets.into(), + }) + } + + 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, offsets } => { + counts.len() * std::mem::size_of::() + + offsets.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.to_vec()), + } + } + + 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 { offsets, .. } => offsets.last().copied().ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural explicit count offsets are empty".into(), + ) + }), + } + } + + fn validate_positive(&self) -> Result<()> { + let has_zero = match self { + Self::Empty => false, + Self::Constant { value, .. } => *value == 0, + Self::Explicit { counts, .. } => counts.contains(&0), + }; + if has_zero { + return Err(Error::invalid_input_source( + "Sparse structural non-empty list count is zero".into(), + )); + } + Ok(()) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SparseStructuralLayerPlan { + Validity { + num_slots: u64, + validity: SparseValiditySet, + }, + List { + num_slots: u64, + num_child_slots: u64, + non_empty_positions: SparsePositionSet, + counts: SparseCountSet, + validity: SparseValiditySet, + }, + FixedSizeList { + num_slots: u64, + dimension: u64, + validity: SparseValiditySet, + }, +} + +impl SparseStructuralPlan { + fn expected_num_items( + layers: &[SparseStructuralLayerPlan], + num_visible_items: u64, + ) -> Result { + layers.iter().try_fold(num_visible_items, |items, layer| { + let additional = match layer { + SparseStructuralLayerPlan::List { + num_slots, + non_empty_positions, + .. + } => num_slots + .checked_sub(non_empty_positions.len()) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list has more non-empty positions than slots".into(), + ) + })?, + SparseStructuralLayerPlan::Validity { .. } + | SparseStructuralLayerPlan::FixedSizeList { .. } => 0, + }; + items.checked_add(additional).ok_or_else(|| { + Error::invalid_input_source("Sparse structural item count overflows".into()) + }) + }) + } + + fn validate(&self, row_domain: u64) -> Result<()> { + usize_from_u64(self.num_visible_items, "visible item count")?; + let expected_num_items = Self::expected_num_items(&self.layers, self.num_visible_items)?; + if self.num_items != expected_num_items { + return Err(Error::invalid_input_source( + format!( + "Sparse structural item count {} does not match the {} items implied by its layers", + self.num_items, expected_num_items + ) + .into(), + )); + } + let mut expected_slots = row_domain; + for (layer_index, layer) in self.layers.iter().enumerate() { + let (num_slots, num_child_slots, validity) = match layer { + SparseStructuralLayerPlan::Validity { + num_slots, + validity, + } => (*num_slots, *num_slots, validity), + SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + } => { + counts.validate_positive()?; + if non_empty_positions.len() != counts.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list layer {} has {} non-empty positions but {} counts", + layer_index, + non_empty_positions.len(), + counts.len() + ) + .into(), + )); + } + if counts.sum()? != *num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list layer {} count sum does not match {} child slots", + layer_index, num_child_slots + ) + .into(), + )); + } + if !validity.contains_only_valid_positions(non_empty_positions, *num_slots)? { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list layer {} contains a non-empty null slot", + layer_index + ) + .into(), + )); + } + (*num_slots, *num_child_slots, validity) + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots, + dimension, + validity, + } => { + if *dimension == 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list layer {} has dimension zero", + layer_index + ) + .into(), + )); + } + let num_child_slots = num_slots.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural fixed-size-list child domain overflows".into(), + ) + })?; + (*num_slots, num_child_slots, validity) + } + }; + usize_from_u64(num_slots, "layer slot count")?; + usize_from_u64(num_child_slots, "layer child slot count")?; + if num_slots != expected_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural layer {} has {} slots, expected {}", + layer_index, num_slots, expected_slots + ) + .into(), + )); + } + validity.positions.validate_domain(num_slots, "validity")?; + expected_slots = num_child_slots; + } + if expected_slots != self.num_visible_items { + return Err(Error::invalid_input_source( + format!( + "Sparse structural terminal domain has {} slots, expected {} visible items", + expected_slots, self.num_visible_items + ) + .into(), + )); + } + Ok(()) + } +} + +#[derive(Debug)] +pub struct SparseStructuralUnraveler { + layers: Vec, + next_layer: usize, + pending_fixed_size_list: bool, +} + +impl SparseStructuralUnraveler { + pub(crate) fn new(plan: SparseStructuralPlan) -> Self { + let next_layer = plan.layers.len(); + Self { + layers: plan.layers, + next_layer, + pending_fixed_size_list: false, + } + } + + fn current_layer(&self) -> Option<&SparseStructuralLayerPlan> { + self.next_layer + .checked_sub(1) + .and_then(|idx| self.layers.get(idx)) + } + + fn consume_current_layer(&mut self) -> Result<()> { + self.next_layer = self.next_layer.checked_sub(1).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural metadata has fewer layers than the Arrow schema".into(), + ) + })?; + Ok(()) + } + + pub(crate) fn ensure_exhausted(&self) -> Result<()> { + if self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse structural metadata has an unconsumed fixed-size-list layer".into(), + )); + } + if self.next_layer != 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural metadata has {} unconsumed layer(s)", + self.next_layer + ) + .into(), + )); + } + Ok(()) + } + + pub(crate) fn is_all_valid(&self) -> bool { + match self.current_layer() { + Some(SparseStructuralLayerPlan::Validity { + num_slots, + validity, + }) + | Some(SparseStructuralLayerPlan::FixedSizeList { + num_slots, + validity, + .. + }) + | Some(SparseStructuralLayerPlan::List { + num_slots, + validity, + .. + }) => match validity.meaning { + SparseValidityMeaning::NullPositions => validity.positions.is_empty(), + SparseValidityMeaning::ValidPositions => validity.positions.len() == *num_slots, + }, + None => true, + } + } + + pub(crate) fn max_lists(&self) -> Result { + match self.current_layer() { + Some(SparseStructuralLayerPlan::List { num_slots, .. }) => { + usize_from_u64(*num_slots, "list slot count") + } + _ => Ok(0), + } + } + + pub(crate) fn skip_validity(&mut self) -> Result<()> { + match self.current_layer() { + Some(SparseStructuralLayerPlan::Validity { .. }) => { + if self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list schema does not match a validity layer".into(), + )); + } + self.consume_current_layer()?; + } + Some(SparseStructuralLayerPlan::FixedSizeList { .. }) => { + if !self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list layer does not match the Arrow schema".into(), + )); + } + self.pending_fixed_size_list = false; + self.consume_current_layer()?; + } + None => { + return Err(Error::invalid_input_source( + "Sparse structural metadata has fewer layers than the Arrow schema".into(), + )); + } + Some(SparseStructuralLayerPlan::List { .. }) => { + return Err(Error::invalid_input_source( + "Sparse structural list layer does not match an Arrow validity layer".into(), + )); + } + } + Ok(()) + } + + pub(crate) fn unravel_validity(&mut self, validity: &mut BooleanBufferBuilder) -> Result<()> { + match self.current_layer() { + Some(SparseStructuralLayerPlan::Validity { + num_slots, + validity: layer_validity, + }) => { + if self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list schema does not match a validity layer".into(), + )); + } + layer_validity.append_to(validity, *num_slots)?; + self.consume_current_layer()?; + } + Some(SparseStructuralLayerPlan::FixedSizeList { + num_slots, + validity: layer_validity, + .. + }) => { + if !self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list layer does not match the Arrow schema".into(), + )); + } + layer_validity.append_to(validity, *num_slots)?; + self.pending_fixed_size_list = false; + self.consume_current_layer()?; + } + None => { + return Err(Error::invalid_input_source( + "Sparse structural metadata has fewer layers than the Arrow schema".into(), + )); + } + Some(SparseStructuralLayerPlan::List { .. }) => { + return Err(Error::invalid_input_source( + "Sparse structural list layer does not match an Arrow validity layer".into(), + )); + } + } + Ok(()) + } + + pub(crate) fn decimate(&mut self, dimension: usize) -> Result<()> { + if self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list layer was decimated more than once".into(), + )); + } + let Some(SparseStructuralLayerPlan::FixedSizeList { + dimension: actual_dimension, + .. + }) = self.current_layer() + else { + return Err(Error::invalid_input_source( + "Sparse structural layer does not match an Arrow fixed-size-list layer".into(), + )); + }; + if usize_from_u64(*actual_dimension, "fixed-size-list dimension")? != dimension { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list dimension {} does not match Arrow dimension {}", + actual_dimension, dimension + ) + .into(), + )); + } + self.pending_fixed_size_list = true; + Ok(()) + } + + 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_source( + "Sparse structural offset does not fit the Arrow offset type".into(), + ) + }) + } + + pub(crate) 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, + validity: layer_validity, + .. + }) = self.current_layer() + else { + return Err(Error::invalid_input_source( + "Sparse structural layer does not match an Arrow list layer".into(), + )); + }; + if self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list schema does not match an Arrow list layer".into(), + )); + } + + 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.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 { + layer_validity.append_to(validity, *num_slots)?; + } + let offset = Self::to_offset(current_offset)?; + let new_len = offsets + .len() + .checked_add(usize_from_u64(*num_slots, "list slot count")?) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list offset length overflows usize".into(), + ) + })?; + offsets.resize(new_len, offset); + self.consume_current_layer()?; + return Ok(()); + } + + let non_empty_positions = non_empty_positions.materialize()?; + let counts = counts.materialize()?; + let mut non_empty_iter = non_empty_positions + .iter() + .copied() + .zip(counts.iter().copied()) + .peekable(); + let mut validity_cursor = + SparseValidityCursor::new(layer_validity, *num_slots, "list validity")?; + let mut validity = validity; + for slot in 0..*num_slots { + let is_valid = validity_cursor.is_valid(slot)?; + if let Some(validity) = validity.as_mut() { + validity.append(is_valid); + } + + if non_empty_iter + .peek() + .is_some_and(|(non_empty_pos, _)| *non_empty_pos == slot) + { + let (_, count) = non_empty_iter.next().ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list position is missing its child count".into(), + ) + })?; + if !is_valid { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list slot {} is both invalid 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)?); + } + validity_cursor.finish()?; + 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(), + )); + } + + self.consume_current_layer()?; + Ok(()) + } +} + +#[derive(Debug, Clone)] +enum SparsePositionSetDecoder { + Empty, + All { + len: u64, + }, + Range { + start: u64, + len: u64, + }, + Explicit { + decompressor: Arc, + encoding: CompressiveEncoding, + count: u64, + domain_len: u64, + }, +} + +#[derive(Debug, Clone)] +enum SparseCountSetDecoder { + Empty, + Constant { + value: u64, + len: u64, + }, + Explicit { + decompressor: Arc, + encoding: CompressiveEncoding, + count: u64, + }, +} + +#[derive(Debug, Clone)] +enum SparseLayerDecompressors { + Validity { + num_slots: u64, + validity: SparseValiditySetDecoder, + }, + List { + num_slots: u64, + num_child_slots: u64, + non_empty_positions: SparsePositionSetDecoder, + counts: SparseCountSetDecoder, + validity: SparseValiditySetDecoder, + }, + FixedSizeList { + num_slots: u64, + num_child_slots: u64, + dimension: u64, + validity: SparseValiditySetDecoder, + }, +} + +#[derive(Debug, Clone)] +struct SparseValiditySetDecoder { + meaning: SparseValidityMeaning, + positions: SparsePositionSetDecoder, +} + +#[derive(Debug)] +struct SparseStructuralCacheableState { + chunk_meta: Vec, + chunk_value_offsets: Arc<[u64]>, + plan: SparseStructuralPlan, + row_domain: u64, +} + +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 { validity, .. } => validity.deep_size(), + SparseStructuralLayerPlan::List { + non_empty_positions, + counts, + validity, + .. + } => non_empty_positions.deep_size() + counts.deep_size() + validity.deep_size(), + SparseStructuralLayerPlan::FixedSizeList { validity, .. } => validity.deep_size(), + }) + .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)] +pub(super) struct SparseStructuralScheduler { + buffer_offsets_and_sizes: Vec<(u64, u64)>, + priority: u64, + row_domain: u64, + row_scale: u64, + num_items: u64, + num_visible_items: u64, + num_buffers: u64, + value_encoding: CompressiveEncoding, + value_decompressor: Arc, + layer_decompressors: Vec, + data_type: DataType, + page_meta: Option>, + has_large_chunk: bool, +} + +impl SparseStructuralScheduler { + fn require_layer( + layer: &pb21::SparseStructuralLayer, + ) -> Result<&pb21::sparse_structural_layer::Layer> { + layer.layer.as_ref().ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural layer is missing its layer variant".into(), + ) + }) + } + + fn layer_num_slots(layer: &pb21::SparseStructuralLayer) -> Result { + Ok(match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::Validity(layer) => layer.num_slots, + pb21::sparse_structural_layer::Layer::List(layer) => layer.num_slots, + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => layer.num_slots, + }) + } + + fn layer_num_child_slots(layer: &pb21::SparseStructuralLayer) -> Result { + Ok(match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::Validity(layer) => layer.num_slots, + pb21::sparse_structural_layer::Layer::List(layer) => layer.num_child_slots, + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => layer + .num_slots + .checked_mul(layer.dimension) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", + layer.num_slots, layer.dimension + ) + .into(), + ) + })?, + }) + } + + pub(super) fn try_new( + buffer_offsets_and_sizes: &[(u64, u64)], + priority: u64, + encoded_row_domain: u64, + data_type: DataType, + layout: &pb21::SparseLayout, + decompressors: &dyn DecompressionStrategy, + ) -> Result { + let value_compression = layout.value_compression.as_ref().ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing value compression".into()) + })?; + let value_buffer_count = Self::validate_value_encoding(value_compression)?; + if layout.num_buffers != value_buffer_count { + return Err(Error::invalid_input_source( + format!( + "Sparse layout declares {} value buffers, but its compression descriptor requires {}", + layout.num_buffers, value_buffer_count + ) + .into(), + )); + } + let row_domain = match layout.structural_layers.first() { + Some(layer) => Self::layer_num_slots(layer)?, + None => encoded_row_domain, + }; + Self::validate_domain_chain( + &layout.structural_layers, + row_domain, + layout.num_items, + layout.num_visible_items, + )?; + 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(), + )); + } + Self::validate_page_buffers(buffer_offsets_and_sizes, layout.num_visible_items)?; + let row_scale = + layout + .structural_layers + .iter() + .try_fold(1_u64, |scale, layer| -> Result { + match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => { + scale.checked_mul(layer.dimension).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural fixed-size-list row scale overflows".into(), + ) + }) + } + pb21::sparse_structural_layer::Layer::Validity(_) + | pb21::sparse_structural_layer::Layer::List(_) => Ok(scale), + } + })?; + let expected_encoded_row_domain = row_domain.checked_mul(row_scale).ok_or_else(|| { + Error::invalid_input_source("Sparse structural encoded row domain overflows".into()) + })?; + if encoded_row_domain != expected_encoded_row_domain { + return Err(Error::invalid_input_source( + format!( + "Sparse structural encoded row domain {} does not match outer domain {} * fixed-size-list scale {}", + encoded_row_domain, row_domain, row_scale + ) + .into(), + )); + } + let value_decompressor = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + decompressors.create_miniblock_decompressor(value_compression, decompressors) + })) + .map_err(|_| { + Error::invalid_input_source( + "Sparse value compression descriptor caused decompressor construction to panic" + .into(), + ) + })? + .map_err(|error| { + Error::invalid_input_source( + format!("Sparse value decompressor construction failed: {error}").into(), + ) + })?; + 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, + row_domain, + row_scale, + num_items: layout.num_items, + num_visible_items: layout.num_visible_items, + num_buffers: layout.num_buffers, + value_encoding: value_compression.clone(), + value_decompressor: value_decompressor.into(), + layer_decompressors, + data_type, + page_meta: None, + has_large_chunk: layout.has_large_chunk, + }) + } + + 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 validate_buffer_compression( + compression: Option<&pb21::BufferCompression>, + label: &str, + ) -> Result<()> { + let Some(compression) = compression else { + return Ok(()); + }; + match invalid_enum( + pb21::CompressionScheme::try_from(compression.scheme), + "compression scheme", + )? { + pb21::CompressionScheme::CompressionAlgorithmUnspecified => { + Err(Error::invalid_input_source( + format!("Sparse structural {label} buffer compression is unspecified").into(), + )) + } + pb21::CompressionScheme::CompressionAlgorithmLz4 + | pb21::CompressionScheme::CompressionAlgorithmZstd => Ok(()), + } + } + + fn encoding_contains_general(root: &CompressiveEncoding) -> bool { + use pb21::compressive_encoding::Compression; + + let mut stack = vec![root]; + while let Some(encoding) = stack.pop() { + let Some(compression) = encoding.compression.as_ref() else { + continue; + }; + match compression { + Compression::General(_) => return true, + Compression::Variable(variable) => { + stack.extend(variable.offsets.as_deref()); + } + Compression::OutOfLineBitpacking(bitpacking) => { + stack.extend(bitpacking.values.as_deref()); + } + Compression::Fsst(fsst) => stack.extend(fsst.values.as_deref()), + Compression::Dictionary(dictionary) => { + stack.extend(dictionary.indices.as_deref()); + stack.extend(dictionary.items.as_deref()); + } + Compression::Rle(rle) => { + stack.extend(rle.values.as_deref()); + stack.extend(rle.run_lengths.as_deref()); + } + Compression::ByteStreamSplit(split) => { + stack.extend(split.values.as_deref()); + } + Compression::FixedSizeList(fsl) => stack.extend(fsl.values.as_deref()), + Compression::PackedStruct(packed) => stack.extend(packed.values.as_deref()), + Compression::VariablePackedStruct(packed) => { + stack.extend( + packed + .fields + .iter() + .filter_map(|field| field.value.as_ref()), + ); + } + Compression::Flat(_) + | Compression::Constant(_) + | Compression::InlineBitpacking(_) => {} + } + } + false + } + + fn require_encoding<'a>( + encoding: &'a Option>, + label: &str, + ) -> Result<&'a CompressiveEncoding> { + encoding.as_deref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} encoding is required").into(), + ) + }) + } + + fn validate_flat(flat: &pb21::Flat, label: &str) -> Result<()> { + if flat.bits_per_value == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} flat bit width is zero").into(), + )); + } + if flat.data.is_some() { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} uses unsupported leaf buffer compression") + .into(), + )); + } + Ok(()) + } + + fn validate_value_encoding(encoding: &CompressiveEncoding) -> Result { + use pb21::compressive_encoding::Compression; + + let compression = encoding.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + "Sparse value compression is missing compression details".into(), + ) + })?; + match compression { + Compression::Flat(flat) => { + Self::validate_flat(flat, "value")?; + Ok(1) + } + Compression::InlineBitpacking(bitpacking) => { + if !matches!(bitpacking.uncompressed_bits_per_value, 8 | 16 | 32 | 64) { + return Err(Error::invalid_input_source( + format!( + "Sparse inline bitpacking width {} is not supported", + bitpacking.uncompressed_bits_per_value + ) + .into(), + )); + } + if bitpacking.values.is_some() { + return Err(Error::invalid_input_source( + "Sparse inline bitpacking uses unsupported leaf buffer compression".into(), + )); + } + Ok(1) + } + Compression::Variable(variable) => { + let offsets = variable.offsets.as_deref().ok_or_else(|| { + Error::invalid_input_source( + "Sparse variable compression is missing offsets".into(), + ) + })?; + let Some(Compression::Flat(offsets)) = offsets.compression.as_ref() else { + return Err(Error::invalid_input_source( + "Sparse variable offsets must use flat compression".into(), + )); + }; + Self::validate_flat(offsets, "variable offsets")?; + if !matches!(offsets.bits_per_value, 32 | 64) { + return Err(Error::invalid_input_source( + format!( + "Sparse variable offset width {} is not supported", + offsets.bits_per_value + ) + .into(), + )); + } + if variable.values.is_some() { + return Err(Error::invalid_input_source( + "Sparse variable values use unsupported leaf buffer compression".into(), + )); + } + Ok(1) + } + Compression::Fsst(fsst) => { + if fsst.symbol_table.is_empty() { + return Err(Error::invalid_input_source( + "Sparse FSST compression has an empty symbol table".into(), + )); + } + let values = Self::require_encoding(&fsst.values, "FSST values")?; + if !matches!(values.compression.as_ref(), Some(Compression::Variable(_))) { + return Err(Error::invalid_input_source( + "Sparse FSST values must use variable compression".into(), + )); + } + Self::validate_value_encoding(values) + } + Compression::ByteStreamSplit(split) => { + let values = Self::require_encoding(&split.values, "byte-stream-split values")?; + let Some(Compression::Flat(flat)) = values.compression.as_ref() else { + return Err(Error::invalid_input_source( + "Sparse byte-stream-split values must use flat compression".into(), + )); + }; + Self::validate_flat(flat, "byte-stream-split values")?; + if !matches!(flat.bits_per_value, 32 | 64) { + return Err(Error::invalid_input_source( + format!( + "Sparse byte-stream-split width {} is not supported", + flat.bits_per_value + ) + .into(), + )); + } + Ok(1) + } + Compression::FixedSizeList(fsl) => Self::validate_fsl_value_encoding(fsl), + Compression::PackedStruct(packed) => Self::validate_packed_value_encoding(packed), + Compression::Rle(rle) => { + let values = Self::require_encoding(&rle.values, "RLE values")?; + let lengths = Self::require_encoding(&rle.run_lengths, "RLE run lengths")?; + Self::validate_block_encoding(values, "RLE values")?; + Self::validate_block_encoding(lengths, "RLE run lengths")?; + Ok(2) + } + Compression::General(general) => { + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + "Sparse general compression is missing its buffer compression".into(), + ) + })?; + Self::validate_buffer_compression(Some(compression), "general")?; + Self::validate_value_encoding(Self::require_encoding( + &general.values, + "general values", + )?) + } + Compression::Constant(_) + | Compression::OutOfLineBitpacking(_) + | Compression::Dictionary(_) + | Compression::VariablePackedStruct(_) => Err(Error::invalid_input_source( + "Sparse value compression uses an unsupported mini-block encoding".into(), + )), + } + } + + fn validate_fsl_value_encoding(fsl: &pb21::FixedSizeList) -> Result { + use pb21::compressive_encoding::Compression; + + if fsl.items_per_value == 0 { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list value compression has dimension zero".into(), + )); + } + let values = Self::require_encoding(&fsl.values, "fixed-size-list values")?; + let child_buffers = match values.compression.as_ref() { + Some(Compression::Flat(flat)) => { + Self::validate_flat(flat, "fixed-size-list values")?; + 1_u64 + } + Some(Compression::FixedSizeList(inner)) => Self::validate_fsl_value_encoding(inner)?, + _ => { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list values must use fixed-size-list or flat compression" + .into(), + )); + } + }; + child_buffers + .checked_add(u64::from(fsl.has_validity)) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse fixed-size-list value buffer count overflows".into(), + ) + }) + } + + fn validate_packed_value_encoding(packed: &pb21::PackedStruct) -> Result { + use pb21::compressive_encoding::Compression; + + if packed.bits_per_value.is_empty() + || packed + .bits_per_value + .iter() + .any(|bits| *bits == 0 || !bits.is_multiple_of(8)) + { + return Err(Error::invalid_input_source( + "Sparse packed-struct widths must be non-empty positive byte widths".into(), + )); + } + let values = Self::require_encoding(&packed.values, "packed-struct values")?; + let Some(Compression::Flat(flat)) = values.compression.as_ref() else { + return Err(Error::invalid_input_source( + "Sparse packed-struct values must use flat compression".into(), + )); + }; + Self::validate_flat(flat, "packed-struct values")?; + let total_bits = packed.bits_per_value.iter().try_fold(0_u64, |sum, bits| { + sum.checked_add(*bits).ok_or_else(|| { + Error::invalid_input_source("Sparse packed-struct bit width sum overflows".into()) + }) + })?; + if total_bits != flat.bits_per_value { + return Err(Error::invalid_input_source( + format!( + "Sparse packed-struct child widths sum to {}, but values use {} bits", + total_bits, flat.bits_per_value + ) + .into(), + )); + } + Ok(1) + } + + fn validate_block_encoding(encoding: &CompressiveEncoding, label: &str) -> Result<()> { + use pb21::compressive_encoding::Compression; + + match encoding.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} is missing compression details").into(), + ) + })? { + Compression::Flat(flat) => Self::validate_flat(flat, label), + Compression::InlineBitpacking(bitpacking) => { + if !matches!(bitpacking.uncompressed_bits_per_value, 8 | 16 | 32 | 64) { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} inline bitpacking width {} is unsupported", + bitpacking.uncompressed_bits_per_value + ) + .into(), + )); + } + if bitpacking.values.is_some() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} uses unsupported leaf buffer compression" + ) + .into(), + )); + } + Ok(()) + } + Compression::OutOfLineBitpacking(bitpacking) => { + if !matches!(bitpacking.uncompressed_bits_per_value, 8 | 16 | 32 | 64) { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} out-of-line bitpacking width {} is unsupported", + bitpacking.uncompressed_bits_per_value + ) + .into(), + )); + } + let values = Self::require_encoding(&bitpacking.values, label)?; + let Some(Compression::Flat(flat)) = values.compression.as_ref() else { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} bitpacked values must be flat").into(), + )); + }; + Self::validate_flat(flat, label) + } + Compression::Constant(constant) => { + if constant + .value + .as_ref() + .is_some_and(|value| value.len() != 8) + { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} constant must be 64 bits").into(), + )); + } + Ok(()) + } + Compression::General(general) => { + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} general compression is missing config") + .into(), + ) + })?; + Self::validate_buffer_compression(Some(compression), label)?; + Self::validate_block_encoding( + Self::require_encoding(&general.values, label)?, + label, + ) + } + Compression::Rle(rle) => { + Self::validate_block_encoding( + Self::require_encoding(&rle.values, "RLE values")?, + "RLE values", + )?; + Self::validate_block_encoding( + Self::require_encoding(&rle.run_lengths, "RLE run lengths")?, + "RLE run lengths", + ) + } + _ => Err(Error::invalid_input_source( + format!("Sparse structural {label} uses an unsupported block encoding").into(), + )), + } + } + + fn validate_domain_chain( + layers: &[pb21::SparseStructuralLayer], + row_domain: u64, + num_items: u64, + num_visible_items: u64, + ) -> Result<()> { + let expected_num_items = + layers + .iter() + .try_fold(num_visible_items, |items, layer| -> Result { + let additional = match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::List(layer) => { + let positions = Self::require_position_set( + &layer.non_empty_positions, + "list non-empty", + )?; + let num_non_empty = Self::position_cardinality( + positions, + layer.num_slots, + "list non-empty positions", + )?; + layer.num_slots.checked_sub(num_non_empty).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list has more non-empty positions than slots" + .into(), + ) + })? + } + pb21::sparse_structural_layer::Layer::Validity(_) + | pb21::sparse_structural_layer::Layer::FixedSizeList(_) => 0, + }; + items.checked_add(additional).ok_or_else(|| { + Error::invalid_input_source("Sparse structural item count overflows".into()) + }) + })?; + if num_items != expected_num_items { + return Err(Error::invalid_input_source( + format!( + "Sparse layout has {} structural items, but its layers imply {}", + num_items, expected_num_items + ) + .into(), + )); + } + let mut expected_slots = row_domain; + for (layer_index, layer) in layers.iter().enumerate() { + let num_slots = Self::layer_num_slots(layer)?; + let num_child_slots = Self::layer_num_child_slots(layer)?; + usize_from_u64(num_slots, "layer slot count")?; + usize_from_u64(num_child_slots, "layer child slot count")?; + if num_slots != expected_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural layer {} has {} slots, expected {} from the outer domain", + layer_index, num_slots, expected_slots + ) + .into(), + )); + } + expected_slots = num_child_slots; + } + if expected_slots != num_visible_items { + return Err(Error::invalid_input_source( + format!( + "Sparse structural terminal domain has {} slots, expected {} visible items", + expected_slots, num_visible_items + ) + .into(), + )); + } + Ok(()) + } + + 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 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 validate_page_buffers( + buffer_offsets_and_sizes: &[(u64, u64)], + num_visible_items: u64, + ) -> Result<()> { + let (_, metadata_size) = buffer_offsets_and_sizes.first().copied().ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing metadata buffer".into()) + })?; + if !metadata_size.is_multiple_of(8) { + return Err(Error::invalid_input_source( + format!( + "Sparse layout metadata buffer has {metadata_size} bytes, which is not a multiple of 8" + ) + .into(), + )); + } + let chunk_count = metadata_size / 8; + if num_visible_items == 0 { + if chunk_count != 0 { + return Err(Error::invalid_input_source( + "Sparse layout with no visible items must have empty chunk metadata".into(), + )); + } + } else { + let min_chunks = + num_visible_items.div_ceil(miniblock::MAX_CONFIGURABLE_MINIBLOCK_VALUES); + if chunk_count < min_chunks || chunk_count > num_visible_items { + return Err(Error::invalid_input_source( + format!( + "Sparse layout metadata declares {chunk_count} chunks for {num_visible_items} visible items, expected {min_chunks}..={num_visible_items}" + ) + .into(), + )); + } + } + + let (_, value_size) = buffer_offsets_and_sizes.get(1).copied().ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing value buffer".into()) + })?; + if (num_visible_items == 0) != (value_size == 0) { + return Err(Error::invalid_input_source( + format!( + "Sparse layout value buffer has {value_size} bytes for {num_visible_items} visible items" + ) + .into(), + )); + } + + for (position, size) in buffer_offsets_and_sizes.iter().skip(2) { + Self::checked_buffer_range(*position, *size, "structural")?; + } + + for (position, size) in buffer_offsets_and_sizes.iter().take(2) { + Self::checked_buffer_range(*position, *size, "page")?; + } + Ok(()) + } + + fn require_position_set<'a>( + set: &'a Option, + label: &str, + ) -> 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 require_validity_set<'a>( + set: &'a Option, + label: &str, + ) -> Result<&'a pb21::SparseValiditySet> { + set.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} validity set is required").into(), + ) + }) + } + + fn validity_meaning( + validity_set: &pb21::SparseValiditySet, + label: &str, + ) -> Result { + match invalid_enum( + pb21::sparse_validity_set::Meaning::try_from(validity_set.meaning), + "validity meaning", + )? { + pb21::sparse_validity_set::Meaning::SparseValidityUnspecified => { + Err(Error::invalid_input_source( + format!("Sparse structural {label} meaning is unspecified").into(), + )) + } + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions => { + Ok(SparseValidityMeaning::NullPositions) + } + pb21::sparse_validity_set::Meaning::SparseValidityValidPositions => { + Ok(SparseValidityMeaning::ValidPositions) + } + } + } + + fn validity_buffer_count( + validity_set: &pb21::SparseValiditySet, + domain_len: u64, + label: &str, + ) -> Result<(usize, SparseValidityMeaning, u64)> { + let meaning = Self::validity_meaning(validity_set, label)?; + let position_set = validity_set.positions.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} positions are required").into(), + ) + })?; + let cardinality = Self::position_cardinality(position_set, domain_len, label)?; + let buffer_count = Self::position_buffer_count(position_set, domain_len, label)?; + Ok((buffer_count, meaning, cardinality)) + } + + 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} position set is missing its variant").into(), + ) + })?; + 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 domain_len == 0 || 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 == 0 || 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)?; + Self::validate_block_encoding(compression, label)?; + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + decompressors.create_block_decompressor(compression) + })) + .map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse structural {label} descriptor caused decompressor construction to panic" + ) + .into(), + ) + })? + .map(Arc::from) + .map_err(|error| { + Error::invalid_input_source( + format!("Sparse structural {label} decompressor construction failed: {error}") + .into(), + ) + }) + } + + 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, + )?, + encoding: compression.clone(), + count: cardinality, + domain_len, + } + } + }, + cardinality, + )) + } + + fn validity_set_decoder( + validity_set: &pb21::SparseValiditySet, + domain_len: u64, + label: &str, + decompressors: &dyn DecompressionStrategy, + ) -> Result<(SparseValiditySetDecoder, u64)> { + let meaning = Self::validity_meaning(validity_set, label)?; + let position_set = validity_set.positions.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} positions are required").into(), + ) + })?; + let (positions, cardinality) = + Self::position_set_decoder(position_set, domain_len, label, decompressors)?; + Ok((SparseValiditySetDecoder { meaning, positions }, cardinality)) + } + + fn num_valid_slots( + meaning: SparseValidityMeaning, + cardinality: u64, + num_slots: u64, + label: &str, + ) -> Result { + match meaning { + SparseValidityMeaning::NullPositions => { + num_slots.checked_sub(cardinality).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} null cardinality {} exceeds slots {}", + cardinality, num_slots + ) + .into(), + ) + }) + } + SparseValidityMeaning::ValidPositions => Ok(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, + )?, + encoding: compression.clone(), + count: cardinality, + } + } + }) + } + + fn add_buffer_count(count: &mut usize, additional: usize) -> Result<()> { + *count = count.checked_add(additional).ok_or_else(|| { + Error::invalid_input_source("Sparse structural buffer count overflows".into()) + })?; + Ok(()) + } + + fn structural_buffer_count(layers: &[pb21::SparseStructuralLayer]) -> Result { + layers + .iter() + .try_fold(0_usize, |mut count, layer| -> Result { + match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::Validity(layer) => { + let (validity_buffers, _, _) = Self::validity_buffer_count( + Self::require_validity_set(&layer.validity, "validity")?, + layer.num_slots, + "validity positions", + )?; + Self::add_buffer_count(&mut count, validity_buffers)?; + } + pb21::sparse_structural_layer::Layer::List(layer) => { + 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", + )?; + let non_empty_buffers = Self::position_buffer_count( + non_empty_positions, + layer.num_slots, + "list non-empty positions", + )?; + Self::add_buffer_count(&mut count, non_empty_buffers)?; + let (validity_buffers, validity_meaning, validity_cardinality) = + Self::validity_buffer_count( + Self::require_validity_set(&layer.validity, "list")?, + layer.num_slots, + "list validity positions", + )?; + Self::add_buffer_count(&mut count, validity_buffers)?; + let num_valid_slots = Self::num_valid_slots( + validity_meaning, + validity_cardinality, + layer.num_slots, + "list validity", + )?; + if num_non_empty > num_valid_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list has {} non-empty slots but only {} valid slots", + num_non_empty, num_valid_slots + ) + .into(), + )); + } + let counts = Self::require_count_set(&layer.counts, "list counts")?; + let count_buffers = + Self::count_buffer_count(counts, num_non_empty, "list counts")?; + Self::add_buffer_count(&mut count, count_buffers)?; + 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( + format!( + "Sparse structural list count sum {} does not match child slots {}", + child_slots, layer.num_child_slots + ) + .into(), + )); + } + } + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => { + if layer.dimension == 0 { + return Err(Error::invalid_input_source( + "Sparse structural fixed-size-list dimension is zero".into(), + )); + } + layer.num_slots.checked_mul(layer.dimension).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", + layer.num_slots, layer.dimension + ) + .into(), + ) + })?; + let (validity_buffers, _, _) = Self::validity_buffer_count( + Self::require_validity_set(&layer.validity, "fixed-size-list")?, + layer.num_slots, + "fixed-size-list validity positions", + )?; + Self::add_buffer_count(&mut count, validity_buffers)?; + } + } + Ok(count) + }) + } + + fn layer_decompressors( + layer: &pb21::SparseStructuralLayer, + decompressors: &dyn DecompressionStrategy, + ) -> Result { + Ok(match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::Validity(layer) => { + let (validity, _) = Self::validity_set_decoder( + Self::require_validity_set(&layer.validity, "validity")?, + layer.num_slots, + "validity positions", + decompressors, + )?; + SparseLayerDecompressors::Validity { + num_slots: layer.num_slots, + validity, + } + } + pb21::sparse_structural_layer::Layer::List(layer) => { + 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 (validity, _) = Self::validity_set_decoder( + Self::require_validity_set(&layer.validity, "list")?, + layer.num_slots, + "list validity positions", + decompressors, + )?; + SparseLayerDecompressors::List { + num_slots: layer.num_slots, + num_child_slots: layer.num_child_slots, + non_empty_positions, + counts, + validity, + } + } + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => { + let (validity, _) = Self::validity_set_decoder( + Self::require_validity_set(&layer.validity, "fixed-size-list")?, + layer.num_slots, + "fixed-size-list validity positions", + decompressors, + )?; + SparseLayerDecompressors::FixedSizeList { + num_slots: layer.num_slots, + num_child_slots: layer.num_slots.checked_mul(layer.dimension).ok_or_else( + || { + Error::invalid_input_source( + "Sparse structural fixed-size-list child slot count overflows" + .into(), + ) + }, + )?, + dimension: layer.dimension, + validity, + } + } + }) + } + + 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, 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) { + let entry: [u8; 8] = chunk.try_into().map_err(|_| { + Error::invalid_input_source( + "Sparse layout chunk metadata entry is not 8 bytes".into(), + ) + })?; + let divided_bytes_minus_one = u32::from_le_bytes( + entry + .get(..4) + .and_then(|bytes| bytes.try_into().ok()) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse layout chunk byte-size field is malformed".into(), + ) + })?, + ); + let num_values = u64::from(u32::from_le_bytes( + entry + .get(4..) + .and_then(|bytes| bytes.try_into().ok()) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse layout chunk value-count field is malformed".into(), + ) + })?, + )); + if num_values == 0 { + return Err(Error::invalid_input_source( + "Sparse layout contains an empty value chunk".into(), + )); + } + if num_values > miniblock::MAX_CONFIGURABLE_MINIBLOCK_VALUES { + return Err(Error::invalid_input_source( + format!( + "Sparse layout value chunk has {} values, exceeding the mini-block limit {}", + num_values, + miniblock::MAX_CONFIGURABLE_MINIBLOCK_VALUES + ) + .into(), + )); + } + let num_bytes = u64::from(divided_bytes_minus_one) + .checked_add(1) + .and_then(|units| units.checked_mul(MINIBLOCK_ALIGNMENT as u64)) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse layout value chunk byte size overflows".into(), + ) + })?; + 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, + offset_bytes, + }); + offset_bytes = offset_bytes.checked_add(num_bytes).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( + format!( + "Sparse layout visible item count mismatch: metadata has {}, layout has {}", + rows_counter, self.num_visible_items + ) + .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) + } + + fn validate_general_buffer_header( + general: &pb21::General, + data: &[u8], + label: &str, + ) -> Result<()> { + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} general compression is missing config").into(), + ) + })?; + Self::validate_buffer_compression(Some(compression), label)?; + let values = Self::require_encoding(&general.values, label)?; + if Self::encoding_contains_general(values) { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} contains nested general compression, which is unsupported" + ) + .into(), + )); + } + + let scheme = invalid_enum( + pb21::CompressionScheme::try_from(compression.scheme), + "compression scheme", + )?; + match scheme { + pb21::CompressionScheme::CompressionAlgorithmLz4 => { + data.get(..4).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} LZ4 buffer is missing its length prefix" + ) + .into(), + ) + })?; + } + pb21::CompressionScheme::CompressionAlgorithmZstd => { + data.get(..8).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} Zstd buffer is missing its length prefix" + ) + .into(), + ) + })?; + } + pb21::CompressionScheme::CompressionAlgorithmUnspecified => { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} general compression scheme is unspecified") + .into(), + )); + } + } + Ok(()) + } + + fn validate_general_child_buffer( + encoding: &CompressiveEncoding, + data: &[u8], + label: &str, + ) -> Result<()> { + if let Some(pb21::compressive_encoding::Compression::General(general)) = + encoding.compression.as_ref() + { + Self::validate_general_buffer_header(general, data, label)?; + } + Ok(()) + } + + fn validate_structural_buffer_headers( + encoding: &CompressiveEncoding, + data: &[u8], + label: &str, + ) -> Result<()> { + use pb21::compressive_encoding::Compression; + + match encoding.compression.as_ref() { + Some(Compression::General(general)) => { + Self::validate_general_buffer_header(general, data, label) + } + Some(Compression::Rle(rle)) => { + let values_size = u64::from_le_bytes( + data.get(..8) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} RLE buffer is missing its header" + ) + .into(), + ) + })? + .try_into() + .map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} RLE header is malformed").into(), + ) + })?, + ); + let values_size = usize_from_u64(values_size, "RLE values buffer size")?; + let values_end = 8_usize.checked_add(values_size).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} RLE values range overflows").into(), + ) + })?; + let values_data = data.get(8..values_end).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} RLE values buffer is truncated").into(), + ) + })?; + let lengths_data = data.get(values_end..).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} RLE run-length buffer is missing") + .into(), + ) + })?; + Self::validate_general_child_buffer( + Self::require_encoding(&rle.values, "RLE values")?, + values_data, + "RLE values", + )?; + Self::validate_general_child_buffer( + Self::require_encoding(&rle.run_lengths, "RLE run lengths")?, + lengths_data, + "RLE run lengths", + ) + } + _ => Ok(()), + } + } + + fn decode_u64_values( + decompressor: &dyn BlockDecompressor, + encoding: &CompressiveEncoding, + data: Bytes, + num_values: u64, + label: &str, + ) -> Result> { + Self::validate_structural_buffer_headers(encoding, &data, label)?; + let decoded = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + decompressor.decompress(LanceBuffer::from_bytes(data, 1), num_values) + })) + .map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} decompression panicked").into(), + ) + })? + .map_err(|error| { + Error::invalid_input_source( + format!("Sparse structural {label} decompression failed: {error}").into(), + ) + })?; + let fixed = decoded.as_fixed_width().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} did not decode to fixed width data").into(), + ) + })?; + 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(), + )); + } + 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_explicit_positions( + decompressor: &Arc, + encoding: &CompressiveEncoding, + data: Bytes, + num_positions: u64, + num_slots: u64, + label: &str, + ) -> Result { + let deltas = + Self::decode_u64_values(decompressor.as_ref(), encoding, 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 { + 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(), + ) + })?; + } + 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); + } + 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, + encoding, + count, + domain_len, + } => Self::decode_explicit_positions( + decompressor, + encoding, + Self::next_structural_buffer(buffers, label)?, + *count, + *domain_len, + label, + ), + } + } + + fn decode_validity_set( + decoder: &SparseValiditySetDecoder, + buffers: &mut impl Iterator, + label: &str, + ) -> Result { + Ok(SparseValiditySet { + meaning: decoder.meaning, + positions: Self::decode_position_set(&decoder.positions, buffers, 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, + encoding, + count, + } => { + let counts = Self::decode_u64_values( + decompressor.as_ref(), + encoding, + Self::next_structural_buffer(buffers, label)?, + *count, + label, + )?; + SparseCountSet::from_counts(counts) + } + } + } + + 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 decode_layer( + layer: &SparseLayerDecompressors, + buffers: &mut impl Iterator, + ) -> Result { + Ok(match layer { + SparseLayerDecompressors::Validity { + num_slots, + validity, + } => SparseStructuralLayerPlan::Validity { + num_slots: *num_slots, + validity: Self::decode_validity_set(validity, buffers, "validity positions")?, + }, + SparseLayerDecompressors::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + } => { + let non_empty_positions = Self::decode_position_set( + non_empty_positions, + buffers, + "list non-empty positions", + )?; + let counts = Self::decode_count_set(counts, buffers, "list counts")?; + let validity = + Self::decode_validity_set(validity, buffers, "list validity positions")?; + let actual_child_slots = counts.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, + validity, + } + } + SparseLayerDecompressors::FixedSizeList { + num_slots, + num_child_slots, + dimension, + validity, + } => { + 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, + validity: Self::decode_validity_set( + validity, + buffers, + "fixed-size-list validity positions", + )?, + } + } + }) + } + + fn lookup_value_chunks(&self, chunk_indices: &[usize]) -> Result> { + let page_meta = self.page_meta.as_ref().ok_or_else(|| { + Error::internal("Sparse page scheduler has not been initialized".to_string()) + })?; + chunk_indices + .iter() + .map(|&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 + .checked_add(chunk_meta.chunk_size_bytes) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse layout value chunk {} byte range overflows", + chunk_idx + ) + .into(), + ) + })?; + Ok(LoadedChunk { + byte_range: bytes_start..bytes_end, + items_in_chunk: chunk_meta.num_values, + chunk_idx, + data: LanceBuffer::empty(), + }) + }) + .collect() + } + + 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); + Ok(start..end) + } +} + +impl StructuralPageScheduler for SparseStructuralScheduler { + fn initialize<'a>( + &'a mut self, + io: &Arc, + ) -> BoxFuture<'a, Result>> { + let (meta_buf_position, meta_buf_size) = match self.metadata_buffer() { + Ok(buffer) => buffer, + Err(err) => return std::future::ready(Err(err)).boxed(), + }; + 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 { + let mut buffers = io_req.await?.into_iter(); + 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_u64; + chunk_value_offsets.push(value_offset); + for chunk in &chunk_meta { + 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); + } + + 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, + }; + plan.validate(self.row_domain)?; + if buffers.next().is_some() { + return Err(Error::invalid_input_source( + "Sparse layout has unused structural buffers".into(), + )); + } + + let page_meta = Arc::new(SparseStructuralCacheableState { + 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) + } + .boxed() + } + + fn load(&mut self, data: &Arc) { + self.page_meta = data + .clone() + .as_arc_any() + .downcast::() + .ok(); + } + + fn schedule_ranges( + &self, + ranges: &[Range], + io: &Arc, + ) -> Result> { + let page_meta = self.page_meta.as_ref().ok_or_else(|| { + Error::internal("Sparse page scheduler has not been initialized".to_string()) + })?; + let encoded_row_domain = page_meta + .row_domain + .checked_mul(self.row_scale) + .ok_or_else(|| { + Error::invalid_input_source("Sparse structural encoded row domain overflows".into()) + })?; + let num_rows = validate_slice_ranges(ranges, encoded_row_domain, "row")?; + let ranges = ranges + .iter() + .map(|range| { + if !range.start.is_multiple_of(self.row_scale) + || !range.end.is_multiple_of(self.row_scale) + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural encoded row range {}..{} is not aligned to fixed-size-list scale {}", + range.start, range.end, self.row_scale + ) + .into(), + )); + } + Ok((range.start / self.row_scale)..(range.end / self.row_scale)) + }) + .collect::>>()?; + + let mut chunks_needed = Vec::new(); + 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, + 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); + let value_decompressor = self.value_decompressor.clone(); + let value_encoding = self.value_encoding.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 row_scale = self.row_scale; + + 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, + value_encoding, + data_type, + page_meta, + loaded_chunks: Arc::new(loaded_chunks), + ranges, + offset_in_current_range: 0, + num_rows, + row_scale, + num_buffers, + has_large_chunk, + }) as Box) + } + .boxed(); + Ok(vec![PageLoadTask { + decoder_fut: res, + num_rows, + }]) + } +} + +#[derive(Debug)] +struct SparseStructuralDecoder { + value_decompressor: Arc, + value_encoding: CompressiveEncoding, + data_type: DataType, + page_meta: Arc, + loaded_chunks: Arc>, + ranges: VecDeque>, + offset_in_current_range: u64, + num_rows: u64, + row_scale: u64, + num_buffers: u64, + has_large_chunk: bool, +} + +impl SparseStructuralDecoder { + fn drain_ranges(&mut self, mut rows_desired: u64) -> Result>> { + if !rows_desired.is_multiple_of(self.row_scale) { + return Err(Error::invalid_input_source( + format!( + "Sparse page decoder drain of {} encoded rows is not aligned to fixed-size-list scale {}", + rows_desired, self.row_scale + ) + .into(), + )); + } + rows_desired /= self.row_scale; + let mut ranges = Vec::new(); + while rows_desired > 0 { + let range = self.ranges.front().ok_or_else(|| { + Error::invalid_input_source( + "Sparse page decoder was asked to drain more rows than were scheduled".into(), + ) + })?; + let start = range + .start + .checked_add(self.offset_in_current_range) + .ok_or_else(|| { + Error::invalid_input_source("Sparse page decoder row offset overflows".into()) + })?; + let rows_available = range.end.checked_sub(start).ok_or_else(|| { + Error::invalid_input_source( + "Sparse page decoder range offset exceeds the scheduled range".into(), + ) + })?; + let rows_to_take = rows_available.min(rows_desired); + let end = start.checked_add(rows_to_take).ok_or_else(|| { + Error::invalid_input_source("Sparse page decoder row range overflows".into()) + })?; + ranges.push(start..end); + 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(); + } + } + Ok(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(), + value_encoding: self.value_encoding.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, + value_encoding: CompressiveEncoding, + data_type: DataType, + page_meta: Arc, + loaded_chunks: Arc>, + num_buffers: u64, + has_large_chunk: bool, +} + +impl DecodeSparseStructuralTask { + fn read_chunk_size( + buf: &[u8], + offset: &mut usize, + width: usize, + chunk_idx: usize, + ) -> Result { + let end = offset.checked_add(width).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural value chunk {chunk_idx} size header overflows").into(), + ) + })?; + let bytes = buf.get(*offset..end).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural value chunk {chunk_idx} has a truncated size header") + .into(), + ) + })?; + let size = match width { + 2 => u32::from(u16::from_le_bytes(bytes.try_into().map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural value chunk {chunk_idx} has a malformed u16 size") + .into(), + ) + })?)), + 4 => u32::from_le_bytes(bytes.try_into().map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural value chunk {chunk_idx} has a malformed u32 size") + .into(), + ) + })?), + _ => { + return Err(Error::internal(format!( + "Unsupported sparse value chunk size width {width}" + ))); + } + }; + *offset = end; + Ok(size) + } + + fn expected_fixed_bytes(num_values: u64, bits_per_value: u64, label: &str) -> Result { + let bits = num_values.checked_mul(bits_per_value).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} decoded bit length overflows").into(), + ) + })?; + usize_from_u64(bits.div_ceil(8), label) + } + + fn validate_fixed_buffer( + buffer: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + label: &str, + ) -> Result<()> { + let expected = Self::expected_fixed_bytes(num_values, bits_per_value, label)?; + if buffer.len() != expected { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} buffer has {} bytes, expected {}", + buffer.len(), + expected + ) + .into(), + )); + } + Ok(()) + } + + fn validate_variable_buffer( + buffer: &LanceBuffer, + num_values: u64, + bits_per_offset: u64, + ) -> Result<()> { + let width = usize_from_u64(bits_per_offset / 8, "variable offset width")?; + let offset_count = num_values.checked_add(1).ok_or_else(|| { + Error::invalid_input_source("Sparse variable offset count overflows".into()) + })?; + let table_len = usize_from_u64(offset_count, "variable offset count")? + .checked_mul(width) + .ok_or_else(|| { + Error::invalid_input_source("Sparse variable offset table size overflows".into()) + })?; + if buffer.len() < table_len { + return Err(Error::invalid_input_source( + format!( + "Sparse variable buffer has {} bytes, smaller than its {}-byte offset table", + buffer.len(), + table_len + ) + .into(), + )); + } + let mut previous = None; + for index in 0..usize_from_u64(offset_count, "variable offset count")? { + let start = index.checked_mul(width).ok_or_else(|| { + Error::invalid_input_source("Sparse variable offset index overflows".into()) + })?; + let end = start.checked_add(width).ok_or_else(|| { + Error::invalid_input_source("Sparse variable offset range overflows".into()) + })?; + let bytes = buffer.as_ref().get(start..end).ok_or_else(|| { + Error::invalid_input_source("Sparse variable offset table is truncated".into()) + })?; + let offset = match width { + 4 => u64::from(u32::from_le_bytes(bytes.try_into().map_err(|_| { + Error::invalid_input_source("Sparse variable u32 offset is malformed".into()) + })?)), + 8 => u64::from_le_bytes(bytes.try_into().map_err(|_| { + Error::invalid_input_source("Sparse variable u64 offset is malformed".into()) + })?), + _ => { + return Err(Error::invalid_input_source( + format!( + "Sparse variable offset width {} is unsupported", + bits_per_offset + ) + .into(), + )); + } + }; + if offset < table_len as u64 || offset > buffer.len() as u64 { + return Err(Error::invalid_input_source( + format!( + "Sparse variable offset {} is outside payload range {}..{}", + offset, + table_len, + buffer.len() + ) + .into(), + )); + } + if previous.is_some_and(|previous| offset < previous) { + return Err(Error::invalid_input_source( + "Sparse variable offsets are not monotonically increasing".into(), + )); + } + previous = Some(offset); + } + Ok(()) + } + + fn validate_fsl_buffers( + fsl: &pb21::FixedSizeList, + buffers: &[LanceBuffer], + num_values: u64, + buffer_index: &mut usize, + ) -> Result<()> { + use pb21::compressive_encoding::Compression; + + let child_values = num_values.checked_mul(fsl.items_per_value).ok_or_else(|| { + Error::invalid_input_source("Sparse fixed-size-list value count overflows".into()) + })?; + if fsl.has_validity { + let validity = buffers.get(*buffer_index).ok_or_else(|| { + Error::invalid_input_source( + "Sparse fixed-size-list value validity buffer is missing".into(), + ) + })?; + Self::validate_fixed_buffer(validity, child_values, 1, "fixed-size-list validity")?; + *buffer_index = buffer_index.checked_add(1).ok_or_else(|| { + Error::invalid_input_source("Sparse fixed-size-list buffer index overflows".into()) + })?; + } + let values = fsl.values.as_deref().ok_or_else(|| { + Error::invalid_input_source("Sparse fixed-size-list value encoding is missing".into()) + })?; + match values.compression.as_ref() { + Some(Compression::FixedSizeList(inner)) => { + Self::validate_fsl_buffers(inner, buffers, child_values, buffer_index) + } + Some(Compression::Flat(flat)) => { + let values = buffers.get(*buffer_index).ok_or_else(|| { + Error::invalid_input_source( + "Sparse fixed-size-list leaf value buffer is missing".into(), + ) + })?; + Self::validate_fixed_buffer( + values, + child_values, + flat.bits_per_value, + "fixed-size-list leaf values", + )?; + *buffer_index = buffer_index.checked_add(1).ok_or_else(|| { + Error::invalid_input_source( + "Sparse fixed-size-list buffer index overflows".into(), + ) + })?; + Ok(()) + } + _ => Err(Error::invalid_input_source( + "Sparse fixed-size-list value encoding is malformed".into(), + )), + } + } + + fn validate_value_buffers(&self, buffers: &[LanceBuffer], num_values: u64) -> Result<()> { + use pb21::compressive_encoding::Compression; + + let compression = self.value_encoding.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source("Sparse value compression is missing".into()) + })?; + match compression { + Compression::Flat(flat) => Self::validate_fixed_buffer( + buffers.first().ok_or_else(|| { + Error::invalid_input_source("Sparse flat value buffer is missing".into()) + })?, + num_values, + flat.bits_per_value, + "flat values", + ), + Compression::InlineBitpacking(bitpacking) => { + let buffer = buffers.first().ok_or_else(|| { + Error::invalid_input_source( + "Sparse inline-bitpacked value buffer is missing".into(), + ) + })?; + if num_values > 1024 { + return Err(Error::invalid_input_source( + format!( + "Sparse inline-bitpacked chunk has {} values, exceeding 1024", + num_values + ) + .into(), + )); + } + let word_bytes = usize_from_u64( + bitpacking.uncompressed_bits_per_value / 8, + "inline bitpacking word width", + )?; + let header = buffer.as_ref().get(..word_bytes).ok_or_else(|| { + Error::invalid_input_source( + "Sparse inline-bitpacked buffer is missing its bit-width header".into(), + ) + })?; + let bit_width = + header + .iter() + .enumerate() + .try_fold(0_u64, |value, (idx, byte)| { + let shift = u32::try_from(idx.checked_mul(8).ok_or_else(|| { + Error::invalid_input_source( + "Sparse inline bit-width shift overflows".into(), + ) + })?) + .map_err(|_| { + Error::invalid_input_source( + "Sparse inline bit-width shift exceeds u32".into(), + ) + })?; + Ok::<_, Error>(value | (u64::from(*byte) << shift)) + })?; + if bit_width > bitpacking.uncompressed_bits_per_value { + return Err(Error::invalid_input_source( + format!( + "Sparse inline bit width {} exceeds uncompressed width {}", + bit_width, bitpacking.uncompressed_bits_per_value + ) + .into(), + )); + } + let payload_bytes = usize_from_u64( + bit_width.checked_mul(1024).ok_or_else(|| { + Error::invalid_input_source( + "Sparse inline-bitpacked payload size overflows".into(), + ) + })? / 8, + "inline-bitpacked payload size", + )?; + let expected = word_bytes.checked_add(payload_bytes).ok_or_else(|| { + Error::invalid_input_source( + "Sparse inline-bitpacked buffer size overflows".into(), + ) + })?; + if buffer.len() != expected { + return Err(Error::invalid_input_source( + format!( + "Sparse inline-bitpacked buffer has {} bytes, expected {}", + buffer.len(), + expected + ) + .into(), + )); + } + Ok(()) + } + Compression::Variable(variable) => { + let offsets = variable + .offsets + .as_deref() + .and_then(|encoding| encoding.compression.as_ref()) + .and_then(|compression| match compression { + Compression::Flat(flat) => Some(flat), + _ => None, + }) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse variable offset encoding is malformed".into(), + ) + })?; + Self::validate_variable_buffer( + buffers.first().ok_or_else(|| { + Error::invalid_input_source( + "Sparse variable value buffer is missing".into(), + ) + })?, + num_values, + offsets.bits_per_value, + ) + } + Compression::Fsst(fsst) => { + let variable = fsst + .values + .as_deref() + .and_then(|encoding| encoding.compression.as_ref()) + .and_then(|compression| match compression { + Compression::Variable(variable) => Some(variable), + _ => None, + }) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse FSST value encoding is malformed".into(), + ) + })?; + let offsets = variable + .offsets + .as_deref() + .and_then(|encoding| encoding.compression.as_ref()) + .and_then(|compression| match compression { + Compression::Flat(flat) => Some(flat), + _ => None, + }) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse FSST offset encoding is malformed".into(), + ) + })?; + Self::validate_variable_buffer( + buffers.first().ok_or_else(|| { + Error::invalid_input_source("Sparse FSST value buffer is missing".into()) + })?, + num_values, + offsets.bits_per_value, + ) + } + Compression::ByteStreamSplit(split) => { + let bits = split + .values + .as_deref() + .and_then(|encoding| encoding.compression.as_ref()) + .and_then(|compression| match compression { + Compression::Flat(flat) => Some(flat.bits_per_value), + _ => None, + }) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse byte-stream-split encoding is malformed".into(), + ) + })?; + Self::validate_fixed_buffer( + buffers.first().ok_or_else(|| { + Error::invalid_input_source( + "Sparse byte-stream-split buffer is missing".into(), + ) + })?, + num_values, + bits, + "byte-stream-split values", + ) + } + Compression::FixedSizeList(fsl) => { + let mut buffer_index = 0; + Self::validate_fsl_buffers(fsl, buffers, num_values, &mut buffer_index)?; + if buffer_index != buffers.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse fixed-size-list descriptor consumed {} of {} buffers", + buffer_index, + buffers.len() + ) + .into(), + )); + } + Ok(()) + } + Compression::PackedStruct(packed) => { + let bits = packed.bits_per_value.iter().try_fold(0_u64, |sum, bits| { + sum.checked_add(*bits).ok_or_else(|| { + Error::invalid_input_source( + "Sparse packed-struct bit width sum overflows".into(), + ) + }) + })?; + Self::validate_fixed_buffer( + buffers.first().ok_or_else(|| { + Error::invalid_input_source( + "Sparse packed-struct value buffer is missing".into(), + ) + })?, + num_values, + bits, + "packed-struct values", + ) + } + Compression::Rle(rle) => { + if buffers.len() != 2 { + return Err(Error::invalid_input_source( + format!( + "Sparse RLE value chunk has {} buffers, expected 2", + buffers.len() + ) + .into(), + )); + } + SparseStructuralScheduler::validate_general_child_buffer( + SparseStructuralScheduler::require_encoding(&rle.values, "RLE values")?, + buffers + .first() + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse RLE value buffer is missing after count validation".into(), + ) + })? + .as_ref(), + "value chunk RLE values", + )?; + SparseStructuralScheduler::validate_general_child_buffer( + SparseStructuralScheduler::require_encoding( + &rle.run_lengths, + "RLE run lengths", + )?, + buffers + .get(1) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse RLE run-length buffer is missing after count validation" + .into(), + ) + })? + .as_ref(), + "value chunk RLE run lengths", + ) + } + Compression::General(general) => { + let buffer = buffers.first().ok_or_else(|| { + Error::invalid_input_source( + "Sparse general-compressed value chunk is missing its first buffer".into(), + ) + })?; + SparseStructuralScheduler::validate_general_buffer_header( + general, + buffer.as_ref(), + "value chunk", + ) + } + _ => Err(Error::invalid_input_source( + "Sparse value chunk uses an unsupported compression descriptor".into(), + )), + } + } + + fn loaded_chunk(&self, chunk_idx: usize) -> Result<&LoadedChunk> { + let index = self + .loaded_chunks + .binary_search_by_key(&chunk_idx, |chunk| chunk.chunk_idx) + .map_err(|_| { + Error::internal(format!( + "Sparse structural decode missing loaded value chunk {}", + chunk_idx + )) + })?; + self.loaded_chunks.get(index).ok_or_else(|| { + Error::internal(format!( + "Sparse structural loaded chunk index {} is missing", + index + )) + }) + } + + fn decode_value_chunk(&self, chunk: &LoadedChunk) -> Result { + let buf = &chunk.data; + 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.as_ref() + .get(..2) + .and_then(|bytes| bytes.try_into().ok()) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} has a malformed level header", + chunk.chunk_idx + ) + .into(), + ) + })?, + ); + let mut offset: usize = 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 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 = (0..num_buffers) + .map(|_| Self::read_chunk_size(buf, &mut offset, size_width, chunk.chunk_idx)) + .collect::>>()?; + + offset = offset + .checked_add(pad_bytes::(offset)) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} padded header overflows", + chunk.chunk_idx + ) + .into(), + ) + })?; + 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 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 = offset + .checked_add(pad_bytes::(offset)) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} padded buffer range overflows", + chunk.chunk_idx + ) + .into(), + ) + })?; + 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::>>()?; + + if offset != buf.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} consumed {} of {} bytes", + chunk.chunk_idx, + offset, + buf.len() + ) + .into(), + )); + } + + self.validate_value_buffers(&buffers, chunk.items_in_chunk)?; + + let decoded = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.value_decompressor + .decompress(buffers, chunk.items_in_chunk) + })) + .map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} decompression panicked", + chunk.chunk_idx + ) + .into(), + ) + })? + .map_err(|error| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} decompression failed: {error}", + chunk.chunk_idx + ) + .into(), + ) + })?; + if decoded.num_values() != chunk.items_in_chunk { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} decoded {} values, expected {}", + chunk.chunk_idx, + decoded.num_values(), + chunk.items_in_chunk + ) + .into(), + )); + } + Ok(decoded) + } + + 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 = SparseStructuralScheduler::value_chunk_index( + &self.page_meta.chunk_value_offsets, + value_start, + )?; + let chunk_value_start = *self + .page_meta + .chunk_value_offsets + .get(chunk_idx) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse value chunk {} has no start offset", chunk_idx).into(), + ) + })?; + let chunk_value_end = *self + .page_meta + .chunk_value_offsets + .get(chunk_idx.checked_add(1).ok_or_else(|| { + Error::invalid_input_source("Sparse value chunk index overflows".into()) + })?) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse value chunk {} has no end offset", chunk_idx).into(), + ) + })?; + let take_end = value_range.end.min(chunk_value_end); + if value_start < chunk_value_start || take_end <= value_start { + return Err(Error::invalid_input_source( + format!( + "Sparse value range {}..{} does not make progress in chunk {} covering {}..{}", + value_range.start, + value_range.end, + chunk_idx, + chunk_value_start, + chunk_value_end + ) + .into(), + )); + } + + 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() + .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, + ); + value_start = take_end; + } + Ok(()) + } + + fn decode_checked(self) -> Result { + 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() + .map(|chunk| chunk.data.len()) + .try_fold(0_usize, |total, len| total.checked_add(len)) + .and_then(|total| total.checked_mul(2)) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural decode size estimate overflows".into(), + ) + })?; + let mut data_builder = DataBlockBuilder::with_capacity_estimate( + u64::try_from(estimated_size_bytes).map_err(|_| { + Error::invalid_input_source("Sparse structural decode size exceeds u64::MAX".into()) + })?, + ); + 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, + }) + } +} + +impl DecodePageTask for DecodeSparseStructuralTask { + fn decode(self: Box) -> Result { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| (*self).decode_checked())) + .map_err(|_| { + Error::invalid_input_source( + "Sparse structural page decoding panicked on malformed input".into(), + ) + })? + } +} + +struct SparseStructuralSelection { + plan: SparseStructuralPlan, + leaf_ranges: Vec>, +} + +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 { + 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.first().ok_or_else(|| { + Error::internal("Sparse structural segment unexpectedly missing".to_string()) + })?; + 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; + let output_end = output_base.checked_add(range_len).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} output domain overflows").into(), + ) + })?; + push_coalesced_range(&mut segments, output_base..output_end); + push_coalesced_range(&mut ordinal_ranges, range.clone()); + output_base = output_end; + } + } + 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 + .checked_add(intersect_start - range.start) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} output position overflows") + .into(), + ) + })?; + let out_end = output_base + .checked_add(intersect_end - range.start) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} output position overflows") + .into(), + ) + })?; + push_coalesced_range(&mut segments, out_start..out_end); + push_coalesced_range( + &mut ordinal_ranges, + intersect_start - *start..intersect_end - *start, + ); + } + output_base = output_base + .checked_add(range.end - range.start) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} output domain overflows").into(), + ) + })?; + } + } + 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 { + let selected_positions = + source_positions.get(idx_start..idx_end).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} explicit position slice is invalid" + ) + .into(), + ) + })?; + for position in selected_positions { + out_positions.push( + output_base + .checked_add(*position - range.start) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} output position overflows" + ) + .into(), + ) + })?, + ); + } + push_coalesced_range(&mut ordinal_ranges, idx_start as u64..idx_end as u64); + } + output_base = output_base + .checked_add(range.end - range.start) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} output domain overflows").into(), + ) + })?; + } + let positions = SparsePositionSet::from_positions(out_positions, output_domain, label)?; + return Ok(SparsePositionSelection { + positions, + ordinal_ranges, + }); + } + } + + Ok(SparsePositionSelection { + positions: position_segments_to_set(segments, output_domain, label)?, + ordinal_ranges, + }) +} + +fn select_validity_set( + validity: &SparseValiditySet, + ranges: &[Range], + domain_len: u64, + label: &str, +) -> Result { + Ok(SparseValiditySet { + meaning: validity.meaning, + positions: select_position_set(&validity.positions, ranges, domain_len, label)?.positions, + }) +} + +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 = offset.checked_add(*count).ok_or_else(|| { + Error::invalid_input_source("Sparse structural list count offsets overflow".into()) + })?; + offsets.push(offset); + } + Ok(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: &SparsePositionSet, + counts: &SparseCountSet, + validity: &SparseValiditySet, + 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(), + )); + } + let count_sum = counts.sum()?; + if count_sum != num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list count sum {} does not match child slots {}", + count_sum, num_child_slots + ) + .into(), + )); + } + + 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_validity = select_validity_set(validity, ranges, num_slots, "list validity")?; + let out_num_slots = validate_slice_ranges(ranges, num_slots, "list")?; + 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: non_empty_selection.positions, + counts: out_counts, + validity: out_validity, + }, + 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 { + counts: 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.get(start..end).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural explicit count slice is invalid".into(), + ) + })?); + } + 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| { + let start = range.start.checked_mul(*value).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list child range start overflows".into(), + ) + })?; + let end = range.end.checked_mul(*value).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list child range end overflows".into(), + ) + })?; + Ok(start..end) + }) + .collect::>>()? + } + SparseCountSet::Explicit { + counts, + offsets: value_offsets, + } => { + let last_offset = *value_offsets.last().ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list count offsets are unexpectedly empty".into(), + ) + })?; + 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(), + ) + })?; + let start_offset = *value_offsets.get(start).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list start offset is missing".into(), + ) + })?; + let end_offset = *value_offsets.get(end).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list end offset is missing".into(), + ) + })?; + Ok(start_offset..end_offset) + }) + .collect::>>()? + } + }; + Ok(coalesce_ranges(child_ranges)) +} + +fn slice_sparse_plan( + plan: &SparseStructuralPlan, + row_ranges: &[Range], + row_domain: u64, +) -> Result { + plan.validate(row_domain)?; + 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 { + match layer { + SparseStructuralLayerPlan::Validity { + num_slots, + validity, + } => { + if selected_domain != *num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural validity slice domain {} does not match {} slots", + selected_domain, num_slots + ) + .into(), + )); + } + let out_num_slots = + validate_slice_ranges(&selected_ranges, *num_slots, "validity")?; + sliced_layers.push(SparseStructuralLayerPlan::Validity { + num_slots: out_num_slots, + validity: select_validity_set( + validity, + &selected_ranges, + *num_slots, + "validity", + )?, + }); + } + SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + .. + } => { + if selected_domain != *num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list slice domain {} does not match {} slots", + selected_domain, num_slots + ) + .into(), + )); + } + let (sliced_layer, child_ranges) = slice_list_layer( + *num_slots, + *num_child_slots, + non_empty_positions, + counts, + validity, + &selected_ranges, + )?; + sliced_layers.push(sliced_layer); + selected_ranges = child_ranges; + selected_domain = *num_child_slots; + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots, + dimension, + validity, + } => { + if selected_domain != *num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list slice domain {} does not match {} slots", + selected_domain, num_slots + ) + .into(), + )); + } + 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 out_num_slots = + validate_slice_ranges(&selected_ranges, *num_slots, "fixed-size-list")?; + let child_ranges = selected_ranges + .iter() + .map(|range| { + let start = range.start.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural fixed-size-list child range start overflows" + .into(), + ) + })?; + let end = range.end.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural fixed-size-list child range end overflows" + .into(), + ) + })?; + Ok(start..end) + }) + .collect::>>()?; + sliced_layers.push(SparseStructuralLayerPlan::FixedSizeList { + num_slots: out_num_slots, + dimension: *dimension, + validity: select_validity_set( + validity, + &selected_ranges, + *num_slots, + "fixed-size-list validity", + )?, + }); + selected_ranges = child_ranges; + selected_domain = num_child_slots; + } + } + } + + if selected_domain != plan.num_visible_items { + return Err(Error::invalid_input_source( + format!( + "Sparse structural selected terminal domain {} does not match {} visible items", + selected_domain, plan.num_visible_items + ) + .into(), + )); + } + let num_visible_items = + validate_slice_ranges(&selected_ranges, plan.num_visible_items, "visible value")?; + let num_items = SparseStructuralPlan::expected_num_items(&sliced_layers, num_visible_items)?; + Ok(SparseStructuralSelection { + plan: SparseStructuralPlan { + layers: sliced_layers, + num_items, + num_visible_items, + }, + leaf_ranges: coalesce_ranges(selected_ranges), + }) +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use crate::{ + compression::DefaultDecompressionStrategy, + encodings::physical::block::{CompressionConfig, CompressionScheme}, + testing::SimulatedScheduler, + }; + + use super::*; + + fn position_set( + positions: pb21::sparse_position_set::Positions, + num_positions: u64, + ) -> Option { + Some(pb21::SparsePositionSet { + positions: Some(positions), + num_positions, + }) + } + + fn position_empty() -> Option { + position_set( + pb21::sparse_position_set::Positions::Empty(pb21::SparsePositionEmpty {}), + 0, + ) + } + + fn position_all(num_positions: u64) -> Option { + position_set( + pb21::sparse_position_set::Positions::All(pb21::SparsePositionAll {}), + num_positions, + ) + } + + fn position_explicit(num_positions: u64) -> Option { + position_set( + pb21::sparse_position_set::Positions::Explicit(ProtobufUtils21::flat(64, None)), + num_positions, + ) + } + + fn general_lz4(values: CompressiveEncoding) -> CompressiveEncoding { + ProtobufUtils21::wrapped(CompressionConfig::new(CompressionScheme::Lz4, None), values) + .unwrap() + } + + fn validity( + meaning: pb21::sparse_validity_set::Meaning, + positions: Option, + ) -> Option { + Some(pb21::SparseValiditySet { + meaning: meaning as i32, + positions, + }) + } + + fn null_positions( + positions: Option, + ) -> Option { + validity( + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions, + positions, + ) + } + + fn count_empty() -> Option { + Some(pb21::SparseCountSet { + counts: Some(pb21::sparse_count_set::Counts::Empty( + pb21::SparseCountEmpty {}, + )), + }) + } + + fn count_constant(value: u64) -> Option { + Some(pb21::SparseCountSet { + counts: Some(pb21::sparse_count_set::Counts::Constant( + pb21::SparseCountConstant { value }, + )), + }) + } + + fn sparse_layout() -> pb21::SparseLayout { + 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 validity_layer( + num_slots: u64, + validity: Option, + ) -> pb21::SparseStructuralLayer { + pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::Validity( + pb21::SparseValidityLayer { + num_slots, + validity, + }, + )), + } + } + + fn list_layer( + num_slots: u64, + num_child_slots: u64, + non_empty_positions: Option, + counts: Option, + validity: Option, + ) -> pb21::SparseStructuralLayer { + pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::List( + pb21::SparseListLayer { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + }, + )), + } + } + + fn fixed_size_list_layer( + num_slots: u64, + dimension: u64, + validity: Option, + ) -> pb21::SparseStructuralLayer { + pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::FixedSizeList( + pb21::SparseFixedSizeListLayer { + num_slots, + dimension, + validity, + }, + )), + } + } + + fn assert_invalid_input_contains(err: Error, expected: &str) { + let message = err.to_string(); + assert!( + matches!(&err, Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + assert!( + message.contains(expected), + "expected error to contain {expected:?}, got {message}" + ); + } + + #[test] + fn rejects_missing_layer_variants_and_item_count_mismatches() { + let decompressors = DefaultDecompressionStrategy::default(); + + let mut layout = sparse_layout(); + layout + .structural_layers + .push(pb21::SparseStructuralLayer { layer: None }); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "missing its layer variant"); + + let mut layout = sparse_layout(); + layout.num_items = 2; + layout + .structural_layers + .push(validity_layer(1, null_positions(position_empty()))); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "layers imply 1"); + } + + #[test] + fn accepts_large_structural_descriptors() { + let explicit_values = 8_u64 * 1024 * 1024 + 1; + let metadata_size = explicit_values * 8; + let value_position = metadata_size; + let value_size = explicit_values * 16; + let structural_position = value_position + value_size; + let structural_size = explicit_values * std::mem::size_of::() as u64; + let mut layout = sparse_layout(); + layout.num_items = explicit_values; + layout.num_visible_items = explicit_values; + layout.structural_layers.push(validity_layer( + explicit_values, + null_positions(position_explicit(explicit_values)), + )); + + SparseStructuralScheduler::try_new( + &[ + (0, metadata_size), + (value_position, value_size), + (structural_position, structural_size), + ], + 0, + explicit_values, + DataType::Int32, + &layout, + &DefaultDecompressionStrategy::default(), + ) + .unwrap(); + } + + #[test] + fn accepts_deep_supported_value_encodings() { + let encoding = (0..300).fold(ProtobufUtils21::flat(32, None), |values, _| { + ProtobufUtils21::fsl(1, false, values) + }); + + assert_eq!( + SparseStructuralScheduler::validate_value_encoding(&encoding).unwrap(), + 1 + ); + } + + fn null_set(positions: SparsePositionSet) -> SparseValiditySet { + SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions, + } + } + + fn valid_set(positions: SparsePositionSet) -> SparseValiditySet { + SparseValiditySet { + meaning: SparseValidityMeaning::ValidPositions, + positions, + } + } + + #[test] + fn semantic_position_and_count_sets_project_without_materializing_ranges() { + let ranges = [1..4]; + assert_eq!( + select_position_set(&SparsePositionSet::Empty, &ranges, 5, "empty") + .unwrap() + .positions, + SparsePositionSet::Empty + ); + assert_eq!( + select_position_set(&SparsePositionSet::all(5), &ranges, 5, "all") + .unwrap() + .positions, + SparsePositionSet::all(3) + ); + assert_eq!( + select_position_set(&SparsePositionSet::range(1, 3), &ranges, 5, "range") + .unwrap() + .positions, + SparsePositionSet::all(3) + ); + assert_eq!( + select_position_set( + &SparsePositionSet::Explicit(vec![0, 2, 4]), + &[0..1, 4..5], + 5, + "explicit", + ) + .unwrap() + .positions, + SparsePositionSet::all(2) + ); + + assert_eq!( + select_count_set(&SparseCountSet::Empty, &[], 0).unwrap(), + SparseCountSet::Empty + ); + assert_eq!( + select_count_set(&SparseCountSet::constant(2, 3), &[0..1, 2..3], 2,).unwrap(), + SparseCountSet::constant(2, 2) + ); + assert_eq!( + select_count_set( + &SparseCountSet::from_counts(vec![1, 2, 3]).unwrap(), + &[0..1, 2..3], + 2, + ) + .unwrap(), + SparseCountSet::from_counts(vec![1, 3]).unwrap() + ); + } + + #[test] + fn validity_polarities_rebuild_the_same_arrow_domain() { + let mut null_builder = BooleanBufferBuilder::new(4); + null_set(SparsePositionSet::Explicit(vec![1, 3])) + .append_to(&mut null_builder, 4) + .unwrap(); + assert_eq!( + null_builder.finish().iter().collect::>(), + vec![true, false, true, false] + ); + + let mut valid_builder = BooleanBufferBuilder::new(4); + valid_set(SparsePositionSet::range(1, 2)) + .append_to(&mut valid_builder, 4) + .unwrap(); + assert_eq!( + valid_builder.finish().iter().collect::>(), + vec![false, true, true, false] + ); + } + + #[test] + fn schema_layer_mismatches_are_invalid_input() { + let mut missing = SparseStructuralUnraveler::new(SparseStructuralPlan { + layers: Vec::new(), + num_items: 1, + num_visible_items: 1, + }); + let mut validity = BooleanBufferBuilder::new(1); + let err = missing.unravel_validity(&mut validity).unwrap_err(); + assert_invalid_input_contains(err, "fewer layers than the Arrow schema"); + + let mut fixed_size_list = SparseStructuralUnraveler::new(SparseStructuralPlan { + layers: vec![SparseStructuralLayerPlan::FixedSizeList { + num_slots: 2, + dimension: 2, + validity: null_set(SparsePositionSet::empty()), + }], + num_items: 4, + num_visible_items: 4, + }); + let mut validity = BooleanBufferBuilder::new(2); + let err = fixed_size_list.unravel_validity(&mut validity).unwrap_err(); + assert_invalid_input_contains(err, "does not match the Arrow schema"); + + let extra = SparseStructuralUnraveler::new(SparseStructuralPlan { + layers: vec![SparseStructuralLayerPlan::Validity { + num_slots: 1, + validity: null_set(SparsePositionSet::empty()), + }], + num_items: 1, + num_visible_items: 1, + }); + let err = extra.ensure_exhausted().unwrap_err(); + assert_invalid_input_contains(err, "1 unconsumed layer"); + } + + fn nested_plan() -> SparseStructuralPlan { + SparseStructuralPlan { + layers: vec![ + SparseStructuralLayerPlan::Validity { + num_slots: 6, + validity: null_set(SparsePositionSet::Explicit(vec![1, 4])), + }, + SparseStructuralLayerPlan::List { + num_slots: 6, + num_child_slots: 5, + non_empty_positions: SparsePositionSet::Explicit(vec![0, 2, 5]), + counts: SparseCountSet::from_counts(vec![2, 1, 2]).unwrap(), + validity: null_set(SparsePositionSet::Explicit(vec![1, 4])), + }, + SparseStructuralLayerPlan::FixedSizeList { + num_slots: 5, + dimension: 2, + validity: valid_set(SparsePositionSet::range(1, 3)), + }, + ], + num_items: 13, + num_visible_items: 10, + } + } + + #[test] + fn discontiguous_projection_preserves_outer_to_inner_order() { + let selection = slice_sparse_plan(&nested_plan(), &[5..6, 0..1], 6).unwrap(); + assert_eq!(selection.leaf_ranges, vec![6..10, 0..4]); + assert_eq!(selection.plan.num_visible_items, 8); + selection.plan.validate(2).unwrap(); + + let SparseStructuralLayerPlan::List { + non_empty_positions, + counts, + .. + } = &selection.plan.layers[1] + else { + panic!("expected projected list layer"); + }; + assert_eq!(*non_empty_positions, SparsePositionSet::all(2)); + assert_eq!(*counts, SparseCountSet::constant(2, 2)); + } + + #[test] + fn no_value_projection_keeps_empty_and_null_list_structure() { + let selection = slice_sparse_plan(&nested_plan(), &[1..2, 3..4], 6).unwrap(); + assert!(selection.leaf_ranges.is_empty()); + assert_eq!(selection.plan.num_visible_items, 0); + selection.plan.validate(2).unwrap(); + + let SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + } = &selection.plan.layers[1] + else { + panic!("expected projected list layer"); + }; + assert_eq!((*num_slots, *num_child_slots), (2, 0)); + assert_eq!(*non_empty_positions, SparsePositionSet::Empty); + assert_eq!(*counts, SparseCountSet::Empty); + assert_eq!(*validity, null_set(SparsePositionSet::range(0, 1))); + } + + #[test] + fn rejects_missing_value_compression_and_buffer_count_mismatch() { + let decompressors = DefaultDecompressionStrategy::default(); + let mut layout = sparse_layout(); + layout.value_compression = None; + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "missing value compression"); + + let mut layout = sparse_layout(); + layout.num_buffers = 2; + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "declares 2 value buffers"); + } + + #[test] + fn rejects_inconsistent_chunk_count_before_io() { + let decompressors = DefaultDecompressionStrategy::default(); + + let layout = sparse_layout(); + let err = SparseStructuralScheduler::try_new( + &[(0, 16), (16, 8)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "declares 2 chunks for 1 visible items"); + } + + #[cfg(feature = "lz4")] + #[test] + fn accepts_large_general_decompression_headers() { + let encoding = general_lz4(ProtobufUtils21::flat(64, None)); + let Some(pb21::compressive_encoding::Compression::General(general)) = + encoding.compression.as_ref() + else { + panic!("expected General compression"); + }; + let declared_size = 65_u32 * 1024 * 1024; + + SparseStructuralScheduler::validate_general_buffer_header( + general, + &declared_size.to_le_bytes(), + "test", + ) + .unwrap(); + } + + #[cfg(feature = "lz4")] + #[tokio::test] + async fn rejects_malformed_general_structural_buffers_as_invalid_input() { + let mut layout = sparse_layout(); + layout.structural_layers.push(validity_layer( + 1, + null_positions(position_set( + pb21::sparse_position_set::Positions::Explicit(general_lz4(ProtobufUtils21::flat( + 64, None, + ))), + 1, + )), + )); + let decompressors = DefaultDecompressionStrategy::default(); + + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8), (16, 8)], + 0, + 1, + 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]); + data.extend_from_slice(&8_u32.to_le_bytes()); + data.extend_from_slice(&[0xff; 4]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + let Err(err) = scheduler.initialize(&io).await else { + panic!("expected malformed General buffer to be rejected"); + }; + assert_invalid_input_contains(err, "decompression failed"); + } + + #[cfg(feature = "lz4")] + #[tokio::test] + async fn rejects_malformed_general_value_buffer() { + let mut layout = sparse_layout(); + layout.value_compression = Some(general_lz4(ProtobufUtils21::flat(32, None))); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + 1, + 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_u16.to_le_bytes()); + data.extend_from_slice(&0_u16.to_le_bytes()); + data.extend_from_slice(&[0; 4]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + scheduler.initialize(&io).await.unwrap(); + let mut page_tasks = scheduler.schedule_ranges(&[0..1], &io).unwrap(); + let mut decoder = page_tasks.pop().unwrap().decoder_fut.await.unwrap(); + let Err(err) = decoder.drain(1).unwrap().decode() else { + panic!("expected malformed General value buffer to be rejected"); + }; + assert_invalid_input_contains(err, "missing its length prefix"); + } + + #[test] + fn rejects_layer_domain_and_fixed_size_list_mismatches() { + let decompressors = DefaultDecompressionStrategy::default(); + let mut layout = sparse_layout(); + layout.num_items = 2; + layout.num_visible_items = 2; + layout + .structural_layers + .push(validity_layer(1, null_positions(position_empty()))); + layout + .structural_layers + .push(validity_layer(2, null_positions(position_empty()))); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "layer 1 has 2 slots, expected 1"); + + let mut layout = sparse_layout(); + layout.num_items = 4; + layout.num_visible_items = 4; + layout.structural_layers.push(fixed_size_list_layer( + 2, + 3, + null_positions(position_empty()), + )); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 2, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "terminal domain has 6 slots"); + } + + #[test] + fn rejects_invalid_validity_and_list_count_semantics() { + let decompressors = DefaultDecompressionStrategy::default(); + let mut layout = sparse_layout(); + layout.structural_layers.push(validity_layer( + 1, + validity( + pb21::sparse_validity_set::Meaning::SparseValidityUnspecified, + position_empty(), + ), + )); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "meaning is unspecified"); + + let mut layout = sparse_layout(); + layout.num_items = 3; + layout.num_visible_items = 3; + layout.structural_layers.push(list_layer( + 2, + 3, + position_all(2), + count_constant(2), + null_positions(position_empty()), + )); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 2, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "count sum 4 does not match child slots 3"); + } + + #[tokio::test] + async fn rejects_unordered_explicit_positions_after_decompression() { + let mut layout = sparse_layout(); + layout.num_items = 4; + layout.num_visible_items = 4; + layout + .structural_layers + .push(validity_layer(4, null_positions(position_explicit(2)))); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8), (16, 16)], + 0, + 4, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + + let mut data = Vec::new(); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&4_u32.to_le_bytes()); + data.extend_from_slice(&[0; 8]); + data.extend_from_slice(&3_u64.to_le_bytes()); + data.extend_from_slice(&0_u64.to_le_bytes()); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + + let Err(err) = scheduler.initialize(&io).await else { + panic!("expected unordered sparse positions to be rejected"); + }; + assert_invalid_input_contains(err, "positions must be strictly increasing"); + } + + #[tokio::test] + async fn rejects_chunk_metadata_value_and_byte_sum_mismatches() { + let layout = sparse_layout(); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + 1, + 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 chunk byte sum mismatch"); + }; + assert_invalid_input_contains(err, "describes 16 value bytes"); + + let mut layout = sparse_layout(); + layout.num_items = 2; + layout.num_visible_items = 2; + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + 2, + 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 chunk value sum mismatch"); + }; + assert_invalid_input_contains(err, "metadata has 1, layout has 2"); + } + + #[tokio::test] + async fn rejects_malformed_value_chunk_before_decompression() { + let layout = sparse_layout(); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + 1, + 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))); + scheduler.initialize(&io).await.unwrap(); + let mut page_tasks = scheduler.schedule_ranges(&[0..1], &io).unwrap(); + let page_task = page_tasks.pop().unwrap(); + let mut decoder = page_task.decoder_fut.await.unwrap(); + let decode_task = decoder.drain(1).unwrap(); + let Err(err) = decode_task.decode() else { + panic!("expected malformed value chunk to be rejected"); + }; + assert_invalid_input_contains(err, "flat values buffer has 0 bytes, expected 4"); + } + + #[tokio::test] + async fn accepts_value_chunks_larger_than_64_mib() { + let chunk_size = 65_u64 * 1024 * 1024; + let layout = sparse_layout(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, chunk_size)], + 0, + 1, + DataType::Int32, + &layout, + &DefaultDecompressionStrategy::default(), + ) + .unwrap(); + let words_minus_one = u32::try_from(chunk_size / MINIBLOCK_ALIGNMENT as u64 - 1).unwrap(); + let mut metadata = Vec::new(); + metadata.extend_from_slice(&words_minus_one.to_le_bytes()); + metadata.extend_from_slice(&1_u32.to_le_bytes()); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(metadata))); + + scheduler.initialize(&io).await.unwrap(); + } + + #[tokio::test] + async fn rejects_value_chunks_above_the_miniblock_limit() { + let num_values = miniblock::MAX_CONFIGURABLE_MINIBLOCK_VALUES + 1; + let mut layout = sparse_layout(); + layout.num_items = num_values; + layout.num_visible_items = num_values; + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 16), (16, 16)], + 0, + num_values, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + let mut data = Vec::new(); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&(num_values as u32).to_le_bytes()); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&[0; 16]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + + let Err(err) = scheduler.initialize(&io).await else { + panic!("expected oversized value chunk to be rejected"); + }; + assert_invalid_input_contains(err, "exceeding the mini-block limit"); + } + + #[derive(Debug, Clone)] + struct RecordingIo { + data: Bytes, + calls: Arc>>>>, + } + + impl RecordingIo { + fn new(data: Bytes) -> Self { + Self { + data, + calls: Arc::new(Mutex::new(Vec::new())), + } + } + } + + impl EncodingsIo for RecordingIo { + fn submit_request( + &self, + ranges: Vec>, + _priority: u64, + ) -> BoxFuture<'static, Result>> { + self.calls.lock().unwrap().push(ranges.clone()); + let data = self.data.clone(); + async move { + ranges + .into_iter() + .map(|range| { + let start = usize_from_u64(range.start, "test range start")?; + let end = usize_from_u64(range.end, "test range end")?; + if start > end || end > data.len() { + return Err(Error::invalid_input_source( + "Test I/O range is outside fixture data".into(), + )); + } + Ok(data.slice(start..end)) + }) + .collect() + } + .boxed() + } + } + + #[tokio::test] + async fn selective_read_requests_only_intersecting_value_chunk() { + let mut layout = sparse_layout(); + layout.num_items = 4; + layout.num_visible_items = 4; + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 16), (16, 32)], + 0, + 4, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + + let mut data = Vec::new(); + for _ in 0..2 { + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&2_u32.to_le_bytes()); + } + for values in [[10_i32, 20], [30, 40]] { + data.extend_from_slice(&0_u16.to_le_bytes()); + data.extend_from_slice(&8_u16.to_le_bytes()); + data.extend_from_slice(&[0; 4]); + for value in values { + data.extend_from_slice(&value.to_le_bytes()); + } + } + + let io = Arc::new(RecordingIo::new(Bytes::from(data))); + let trait_io: Arc = io.clone(); + scheduler.initialize(&trait_io).await.unwrap(); + let mut page_tasks = scheduler.schedule_ranges(&[2..3], &trait_io).unwrap(); + let page_task = page_tasks.pop().unwrap(); + let mut decoder = page_task.decoder_fut.await.unwrap(); + let decoded = decoder.drain(1).unwrap().decode().unwrap(); + assert_eq!(decoded.data.num_values(), 1); + + let calls = io.calls.lock().unwrap(); + assert_eq!(calls.as_slice(), &[vec![0..16], vec![32..48]]); + } + + #[tokio::test] + async fn empty_leaf_selection_rebuilds_offsets_without_value_io() { + let mut layout = sparse_layout(); + layout.num_visible_items = 0; + layout.structural_layers.push(list_layer( + 1, + 0, + position_empty(), + count_empty(), + null_positions(position_empty()), + )); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + let io = Arc::new(RecordingIo::new(Bytes::new())); + let trait_io: Arc = io.clone(); + scheduler.initialize(&trait_io).await.unwrap(); + let mut page_tasks = scheduler.schedule_ranges(&[0..1], &trait_io).unwrap(); + let page_task = page_tasks.pop().unwrap(); + let mut decoder = page_task.decoder_fut.await.unwrap(); + let decoded = decoder.drain(1).unwrap().decode().unwrap(); + assert_eq!(decoded.data.num_values(), 0); + + let mut repdef = CompositeRepDefUnraveler::new(vec![decoded.repdef]); + let (offsets, validity) = repdef.unravel_offsets::().unwrap(); + assert_eq!(offsets.as_ref(), &[0, 0]); + assert!(validity.is_none()); + + let calls = io.calls.lock().unwrap(); + assert_eq!(calls.len(), 2); + assert!(calls[1].is_empty(), "value payload must not be requested"); + } +} diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs new file mode 100644 index 00000000000..9925e9cef79 --- /dev/null +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -0,0 +1,1811 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Sparse structural planning and serialization. + +use std::iter; + +use arrow_buffer::BooleanBuffer; +use lance_core::{Error, Result, datatypes::Field, utils::bit::pad_bytes}; + +use crate::{ + buffer::LanceBuffer, + compression::CompressionStrategy, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + decoder::PageEncoding, + encoder::EncodedPage, + format::pb21::{self, CompressiveEncoding}, + repdef::{NormalizedStructuralLayer, NormalizedStructuralPlan}, + statistics::ComputeStat, +}; + +use super::{ + SparseCountSet, SparsePositionSet, SparseStructuralLayerPlan, SparseStructuralPlan, + SparseValidityMeaning, SparseValiditySet, +}; +use crate::encodings::logical::primitive::{ + FILL_BYTE, MINIBLOCK_ALIGNMENT, miniblock::MiniBlockCompressed, +}; + +#[derive(Clone, Copy, Default)] +struct PositionSetStats { + count: u64, + first: u64, + last: u64, + is_contiguous: bool, +} + +impl PositionSetStats { + fn observe(&mut self, position: u64) { + if self.count == 0 { + self.first = position; + self.is_contiguous = true; + } else if self.last.checked_add(1) != Some(position) { + self.is_contiguous = false; + } + self.last = position; + self.count += 1; + } + + fn encoded_cost(&self, domain_len: u64) -> u64 { + if self.count == 0 || self.count == domain_len || self.is_contiguous { + 0 + } else { + self.count + } + } + + fn to_set( + self, + validity: &BooleanBuffer, + want_valid: bool, + domain_len: u64, + label: &str, + ) -> Result { + if self.count == 0 { + return Ok(SparsePositionSet::empty()); + } + if self.count == domain_len { + return Ok(SparsePositionSet::all(domain_len)); + } + if self.is_contiguous { + return Ok(SparsePositionSet::range(self.first, self.count)); + } + + let capacity = usize::try_from(self.count).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} positions exceed usize::MAX").into(), + ) + })?; + let mut positions = Vec::with_capacity(capacity); + for (index, is_valid) in validity.iter().enumerate() { + if is_valid == want_valid { + positions.push(u64::try_from(index).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} position exceeds u64::MAX").into(), + ) + })?); + } + } + SparsePositionSet::from_positions(positions, domain_len, label) + } +} + +fn usize_to_u64(value: usize, label: &str) -> Result { + u64::try_from(value).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} {value} exceeds u64::MAX").into(), + ) + }) +} + +fn validity_set( + validity: Option<&BooleanBuffer>, + num_slots: usize, + label: &str, +) -> Result { + let Some(validity) = validity else { + return Ok(SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions: SparsePositionSet::empty(), + }); + }; + if validity.len() != num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} validity length {} does not match {} slots", + validity.len(), + num_slots + ) + .into(), + )); + } + + let domain_len = usize_to_u64(num_slots, "validity domain")?; + let mut valid_stats = PositionSetStats::default(); + let mut null_stats = PositionSetStats::default(); + for (index, is_valid) in validity.iter().enumerate() { + let index = usize_to_u64(index, "validity position")?; + if is_valid { + valid_stats.observe(index); + } else { + null_stats.observe(index); + } + } + + if null_stats.count == 0 { + return Ok(SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions: SparsePositionSet::empty(), + }); + } + if valid_stats.count == 0 { + return Ok(SparseValiditySet { + meaning: SparseValidityMeaning::ValidPositions, + positions: SparsePositionSet::empty(), + }); + } + + let valid_cost = valid_stats.encoded_cost(domain_len); + let null_cost = null_stats.encoded_cost(domain_len); + if valid_cost < null_cost { + Ok(SparseValiditySet { + meaning: SparseValidityMeaning::ValidPositions, + positions: valid_stats.to_set(validity, true, domain_len, label)?, + }) + } else { + Ok(SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions: null_stats.to_set(validity, false, domain_len, label)?, + }) + } +} + +/// Builds the semantic sparse plan directly from the once-normalized Arrow layers. +pub(in crate::encodings::logical::primitive) fn plan( + normalized: &NormalizedStructuralPlan, + num_visible_items: u64, +) -> Result { + let mut layers = Vec::with_capacity(normalized.layers().len()); + let mut num_items = num_visible_items; + + for layer in normalized.layers() { + match layer { + NormalizedStructuralLayer::Validity { + validity, + num_slots, + } => { + layers.push(SparseStructuralLayerPlan::Validity { + num_slots: usize_to_u64(num_slots, "validity slot count")?, + validity: validity_set(validity, num_slots, "validity")?, + }); + } + NormalizedStructuralLayer::FixedSizeList { + validity, + dimension, + num_slots, + } => { + if dimension == 0 { + return Err(Error::invalid_input_source( + "Sparse structural fixed-size-list dimension is zero".into(), + )); + } + layers.push(SparseStructuralLayerPlan::FixedSizeList { + num_slots: usize_to_u64(num_slots, "fixed-size-list slot count")?, + dimension: usize_to_u64(dimension, "fixed-size-list dimension")?, + validity: validity_set(validity, num_slots, "fixed-size-list validity")?, + }); + } + NormalizedStructuralLayer::List { + offsets, + validity, + num_slots, + } => { + let expected_offsets = num_slots.checked_add(1).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list offset count overflows".into(), + ) + })?; + if offsets.len() != expected_offsets { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list has {} offsets for {} slots", + offsets.len(), + num_slots + ) + .into(), + )); + } + if offsets.first().copied() != Some(0) { + return Err(Error::invalid_input_source( + "Sparse structural list offsets must start at zero".into(), + )); + } + + let mut non_empty_positions = Vec::new(); + let mut counts = Vec::new(); + for slot in 0..num_slots { + let start = offsets[slot]; + let end = offsets[slot + 1]; + let count = end.checked_sub(start).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural list offsets decrease at slot {slot}: {start}..{end}" + ) + .into(), + ) + })?; + let is_valid = validity.is_none_or(|validity| validity.value(slot)); + if !is_valid && count != 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural null list slot {slot} has {count} child slots" + ) + .into(), + )); + } + if is_valid && count > 0 { + non_empty_positions.push(usize_to_u64(slot, "list position")?); + counts.push(u64::try_from(count).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural list count {count} exceeds u64::MAX") + .into(), + ) + })?); + } + } + + let num_slots_u64 = usize_to_u64(num_slots, "list slot count")?; + let num_non_empty = usize_to_u64(non_empty_positions.len(), "list position count")?; + num_items = num_items + .checked_add(num_slots_u64 - num_non_empty) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural item count overflows u64".into(), + ) + })?; + let num_child_slots = offsets.last().copied().ok_or_else(|| { + Error::invalid_input_source("Sparse structural list has no offsets".into()) + })?; + let num_child_slots = u64::try_from(num_child_slots).map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse structural list child slot count {num_child_slots} is negative" + ) + .into(), + ) + })?; + layers.push(SparseStructuralLayerPlan::List { + num_slots: num_slots_u64, + num_child_slots, + non_empty_positions: SparsePositionSet::from_positions( + non_empty_positions, + num_slots_u64, + "list non-empty", + )?, + counts: SparseCountSet::from_counts(counts)?, + validity: validity_set(validity, num_slots, "list validity")?, + }); + } + } + } + + let row_domain = match layers.first() { + Some(SparseStructuralLayerPlan::Validity { num_slots, .. }) + | Some(SparseStructuralLayerPlan::List { num_slots, .. }) + | Some(SparseStructuralLayerPlan::FixedSizeList { num_slots, .. }) => *num_slots, + None => { + return Err(Error::invalid_input_source( + "Sparse structural encoding requires at least one Arrow structural layer".into(), + )); + } + }; + let plan = SparseStructuralPlan { + layers, + num_items, + num_visible_items, + }; + plan.validate(row_domain)?; + Ok(plan) +} + +/// ConstantLayout remains the canonical representation when no value payload exists. +pub(in crate::encodings::logical::primitive) fn uses_constant_layout( + plan: &SparseStructuralPlan, + field: &Field, +) -> bool { + if plan.num_visible_items == 0 { + return true; + } + if matches!(field.data_type(), arrow_schema::DataType::Struct(fields) if fields.is_empty()) { + return true; + } + + let Some(layer) = plan.layers.last() else { + return false; + }; + let (num_slots, validity) = match layer { + SparseStructuralLayerPlan::Validity { + num_slots, + validity, + } + | SparseStructuralLayerPlan::List { + num_slots, + validity, + .. + } + | SparseStructuralLayerPlan::FixedSizeList { + num_slots, + validity, + .. + } => (*num_slots, validity), + }; + match validity.meaning { + SparseValidityMeaning::NullPositions => validity.positions.len() == num_slots, + SparseValidityMeaning::ValidPositions => validity.positions.is_empty(), + } +} + +struct SparseMiniBlockChunk { + buffer_sizes: Vec, + num_values: u32, +} + +struct SparseMiniBlockCompressed { + data: Vec, + chunks: Vec, +} + +struct SerializedValuePage { + num_buffers: u64, + data: LanceBuffer, + metadata: LanceBuffer, +} + +struct EncodedStructuralPlan { + layers: Vec, + buffers: Vec, +} + +fn with_explicit_value_counts( + compressed: MiniBlockCompressed, +) -> Result { + let mut values_in_previous_chunks = 0_u64; + let mut chunks = Vec::with_capacity(compressed.chunks.len()); + for chunk in compressed.chunks { + let num_values = chunk.num_values(values_in_previous_chunks, compressed.num_values); + values_in_previous_chunks = values_in_previous_chunks + .checked_add(num_values) + .ok_or_else(|| Error::internal("Sparse value count overflows u64".to_string()))?; + chunks.push(SparseMiniBlockChunk { + buffer_sizes: chunk.buffer_sizes, + num_values: u32::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse value chunk has {num_values} visible values, which exceeds the u32 metadata limit" + ) + .into(), + ) + })?, + }); + } + if values_in_previous_chunks != compressed.num_values { + return Err(Error::internal(format!( + "Sparse value chunks describe {values_in_previous_chunks} values, expected {}", + compressed.num_values + ))); + } + Ok(SparseMiniBlockCompressed { + data: compressed.data, + chunks, + }) +} + +fn serialize_value_chunks( + compressed: SparseMiniBlockCompressed, + support_large_chunk: bool, +) -> Result { + let bytes_data = compressed.data.iter().map(LanceBuffer::len).sum::(); + let num_buffers = compressed.data.len(); + let mut data_buffer = Vec::with_capacity(bytes_data + 9 * num_buffers); + let mut metadata = Vec::with_capacity(compressed.chunks.len() * 8); + let mut buffer_offsets = vec![0_usize; num_buffers]; + + for chunk in compressed.chunks { + if chunk.buffer_sizes.len() != num_buffers { + return Err(Error::internal(format!( + "Sparse chunk has {} value buffer sizes, expected {num_buffers}", + chunk.buffer_sizes.len() + ))); + } + + let chunk_start = data_buffer.len(); + debug_assert_eq!(chunk_start % MINIBLOCK_ALIGNMENT, 0); + data_buffer.extend_from_slice(&0_u16.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 ({buffer_size} bytes) exceeds 16-bit metadata" + )) + })?; + data_buffer.extend_from_slice(&buffer_size.to_le_bytes()); + } + } + let add_padding = |buffer: &mut Vec| { + let padding = pad_bytes::(buffer.len()); + buffer.extend(iter::repeat_n(FILL_BYTE, padding)); + }; + add_padding(&mut data_buffer); + + for (buffer_size, (buffer, buffer_offset)) in chunk + .buffer_sizes + .iter() + .zip(compressed.data.iter().zip(buffer_offsets.iter_mut())) + { + let start = *buffer_offset; + let end = start.checked_add(*buffer_size as usize).ok_or_else(|| { + Error::internal("Sparse value buffer range overflows".to_string()) + })?; + let bytes = buffer.as_ref().get(start..end).ok_or_else(|| { + Error::internal(format!( + "Sparse value chunk requests bytes {start}..{end} from a {}-byte buffer", + buffer.len() + )) + })?; + *buffer_offset = end; + data_buffer.extend_from_slice(bytes); + add_padding(&mut data_buffer); + } + + let chunk_bytes = data_buffer.len() - chunk_start; + if chunk_bytes == 0 || !chunk_bytes.is_multiple_of(MINIBLOCK_ALIGNMENT) { + return Err(Error::internal(format!( + "Sparse value chunk size {chunk_bytes} is not a positive multiple of {MINIBLOCK_ALIGNMENT}" + ))); + } + let words_minus_one = chunk_bytes / MINIBLOCK_ALIGNMENT - 1; + metadata.extend_from_slice( + &u32::try_from(words_minus_one) + .map_err(|_| { + Error::internal(format!( + "Sparse value chunk size {chunk_bytes} exceeds the metadata limit" + )) + })? + .to_le_bytes(), + ); + metadata.extend_from_slice(&chunk.num_values.to_le_bytes()); + } + + for (index, (consumed, buffer)) in buffer_offsets + .iter() + .zip(compressed.data.iter()) + .enumerate() + { + if *consumed != buffer.len() { + return Err(Error::internal(format!( + "Sparse value buffer {index} consumed {consumed} bytes, expected {}", + buffer.len() + ))); + } + } + + Ok(SerializedValuePage { + num_buffers: usize_to_u64(num_buffers, "value buffer count")?, + data: LanceBuffer::from(data_buffer), + metadata: LanceBuffer::from(metadata), + }) +} + +fn encode_u64_values( + values: Vec, + compression_strategy: &dyn CompressionStrategy, +) -> Result<(LanceBuffer, CompressiveEncoding)> { + let num_values = usize_to_u64(values.len(), "u64 value count")?; + 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("", arrow_schema::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 = 0_u64; + positions + .iter() + .copied() + .enumerate() + .map(|(index, position)| { + if index > 0 && position <= previous { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} positions must be strictly increasing") + .into(), + )); + } + let delta = if index == 0 { + position + } else { + position - previous + }; + previous = position; + Ok(delta) + }) + .collect() +} + +fn encode_position_set( + positions: &SparsePositionSet, + compression_strategy: &dyn CompressionStrategy, + label: &str, +) -> Result<(Option, pb21::SparsePositionSet)> { + let (buffer, positions_pb) = 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) => { + if positions.is_empty() { + return Err(Error::internal(format!( + "Sparse structural {label} explicit set is empty" + ))); + } + let (buffer, encoding) = + encode_u64_values(positions_to_deltas(positions, label)?, compression_strategy)?; + ( + Some(buffer), + pb21::sparse_position_set::Positions::Explicit(encoding), + ) + } + }; + Ok(( + buffer, + pb21::SparsePositionSet { + positions: Some(positions_pb), + num_positions: positions.len(), + }, + )) +} + +fn encode_count_set( + counts: &SparseCountSet, + compression_strategy: &dyn CompressionStrategy, +) -> Result<(Option, pb21::SparseCountSet)> { + let (buffer, counts_pb) = 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) = encode_u64_values(counts.to_vec(), compression_strategy)?; + ( + Some(buffer), + pb21::sparse_count_set::Counts::Explicit(encoding), + ) + } + }; + Ok(( + buffer, + pb21::SparseCountSet { + counts: Some(counts_pb), + }, + )) +} + +fn encode_validity_set( + validity: &SparseValiditySet, + compression_strategy: &dyn CompressionStrategy, + label: &str, +) -> Result<(Option, pb21::SparseValiditySet)> { + let (buffer, positions) = + encode_position_set(&validity.positions, compression_strategy, label)?; + let meaning = match validity.meaning { + SparseValidityMeaning::NullPositions => { + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions + } + SparseValidityMeaning::ValidPositions => { + pb21::sparse_validity_set::Meaning::SparseValidityValidPositions + } + }; + Ok(( + buffer, + pb21::SparseValiditySet { + meaning: meaning as i32, + positions: Some(positions), + }, + )) +} + +fn encode_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, + validity, + } => { + let (validity_buffer, validity) = + encode_validity_set(validity, compression_strategy, "validity")?; + buffers.extend(validity_buffer); + layers.push(pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::Validity( + pb21::SparseValidityLayer { + num_slots: *num_slots, + validity: Some(validity), + }, + )), + }); + } + SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + } => { + 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 (position_buffer, non_empty_positions) = encode_position_set( + non_empty_positions, + compression_strategy, + "list non-empty", + )?; + buffers.extend(position_buffer); + let (count_buffer, counts) = encode_count_set(counts, compression_strategy)?; + buffers.extend(count_buffer); + let (validity_buffer, validity) = + encode_validity_set(validity, compression_strategy, "list validity")?; + buffers.extend(validity_buffer); + layers.push(pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::List( + pb21::SparseListLayer { + num_slots: *num_slots, + num_child_slots: *num_child_slots, + non_empty_positions: Some(non_empty_positions), + counts: Some(counts), + validity: Some(validity), + }, + )), + }); + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots, + dimension, + validity, + } => { + let (validity_buffer, validity) = encode_validity_set( + validity, + compression_strategy, + "fixed-size-list validity", + )?; + buffers.extend(validity_buffer); + num_slots.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={num_slots}, dimension={dimension}" + ) + .into(), + ) + })?; + layers.push(pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::FixedSizeList( + pb21::SparseFixedSizeListLayer { + num_slots: *num_slots, + dimension: *dimension, + validity: Some(validity), + }, + )), + }); + } + } + } + + Ok(EncodedStructuralPlan { layers, buffers }) +} + +#[allow(clippy::too_many_arguments)] +pub(in crate::encodings::logical::primitive) fn encode_page( + column_idx: u32, + field: &Field, + compression_strategy: &dyn CompressionStrategy, + data: DataBlock, + plan: SparseStructuralPlan, + row_number: u64, + num_rows: u64, + support_large_chunk: bool, +) -> Result { + if plan.num_visible_items != data.num_values() { + return Err(Error::internal(format!( + "Sparse structural plan has {} visible items but data has {} values", + plan.num_visible_items, + data.num_values() + ))); + } + match &data { + DataBlock::AllNull(_) => { + return Err(Error::internal( + "All-null values must use ConstantLayout".to_string(), + )); + } + DataBlock::Dictionary(_) => { + return Err(Error::not_supported_source( + "Sparse layout does not support dictionary data blocks".into(), + )); + } + DataBlock::Struct(data) if data.has_variable_width_child() => { + return Err(Error::not_supported_source( + "Sparse layout does not support variable-width packed struct data blocks".into(), + )); + } + _ => {} + } + + let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; + let (compressed, value_compression) = compressor.compress(data)?; + let values = + serialize_value_chunks(with_explicit_value_counts(compressed)?, support_large_chunk)?; + let structural = encode_structural_plan(&plan, compression_strategy)?; + let description = pb21::PageLayout { + layout: Some(pb21::page_layout::Layout::SparseLayout( + pb21::SparseLayout { + value_compression: Some(value_compression), + num_buffers: values.num_buffers, + num_items: plan.num_items, + num_visible_items: plan.num_visible_items, + has_large_chunk: support_large_chunk, + structural_layers: structural.layers, + }, + )), + }; + + let mut page_data = Vec::with_capacity(2 + structural.buffers.len()); + page_data.push(values.metadata); + page_data.push(values.data); + page_data.extend(structural.buffers); + Ok(EncodedPage { + data: page_data, + description: PageEncoding::Structural(description), + num_rows, + row_number, + column_idx, + }) +} + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::Arc}; + + use arrow_array::{ + Array, ArrayRef, FixedSizeListArray, Int32Array, LargeListArray, ListArray, StructArray, + builder::{Int32Builder, MapBuilder, StringBuilder}, + }; + use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field as ArrowField, Fields}; + + use crate::{ + constants::{ + STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, + STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_SPARSE, + }, + encoder::{ + ColumnIndexSequence, EncodingOptions, FieldEncoder, MIN_PAGE_BUFFER_ALIGNMENT, + OutOfLineBuffers, default_encoding_strategy, + }, + testing::{TestCases, check_round_trip_encoding_of_data}, + version::LanceFileVersion, + }; + + use super::*; + + fn sparse_metadata() -> HashMap { + HashMap::from([( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_SPARSE.to_string(), + )]) + } + + fn structural_metadata(value: &str) -> HashMap { + HashMap::from([(STRUCTURAL_ENCODING_META_KEY.to_string(), value.to_string())]) + } + + fn null_buffer(validity: impl IntoIterator) -> NullBuffer { + NullBuffer::new(BooleanBuffer::from_iter(validity)) + } + + fn list_i32(offsets: Vec, validity: Option>) -> ArrayRef { + let num_values = offsets.last().copied().unwrap_or_default(); + let values = Arc::new(Int32Array::from_iter_values(0..num_values)) as ArrayRef; + Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + values, + validity.map(null_buffer), + ) + .unwrap(), + ) + } + + fn large_list_i32(offsets: Vec, validity: Option>) -> ArrayRef { + let num_values = offsets.last().copied().unwrap_or_default(); + let values = Arc::new(Int32Array::from_iter_values(0..num_values as i32)) as ArrayRef; + Arc::new( + LargeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + values, + validity.map(null_buffer), + ) + .unwrap(), + ) + } + + fn map_i32() -> ArrayRef { + let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + builder.keys().append_value("a"); + builder.values().append_value(1); + builder.append(true).unwrap(); + builder.append(false).unwrap(); + builder.append(true).unwrap(); + builder.keys().append_value("b"); + builder.values().append_null(); + builder.keys().append_value("c"); + builder.values().append_value(3); + builder.append(true).unwrap(); + Arc::new(builder.finish()) + } + + fn fixed_size_list_struct() -> ArrayRef { + let fields = Fields::from(vec![ArrowField::new("value", DataType::Int32, true)]); + let child = Arc::new(StructArray::new( + fields.clone(), + vec![Arc::new(Int32Array::from(vec![ + Some(0), + Some(1), + None, + Some(3), + Some(4), + Some(5), + None, + Some(7), + Some(8), + Some(9), + Some(10), + None, + ]))], + Some(null_buffer([ + true, true, false, true, true, true, true, false, true, true, true, true, + ])), + )) as ArrayRef; + Arc::new( + FixedSizeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Struct(fields), true)), + 2, + child, + Some(null_buffer([true, false, true, true, false, true])), + ) + .unwrap(), + ) + } + + fn list_fixed_size_list_struct() -> ArrayRef { + let struct_fields = Fields::from(vec![ArrowField::new("value", DataType::Int32, true)]); + let structs = Arc::new(StructArray::new( + struct_fields.clone(), + vec![Arc::new(Int32Array::from(vec![ + Some(0), + None, + Some(2), + Some(3), + None, + Some(5), + ]))], + Some(null_buffer([true, true, false, true, true, true])), + )) as ArrayRef; + let fixed_size_list = Arc::new( + FixedSizeListArray::try_new( + Arc::new(ArrowField::new( + "item", + DataType::Struct(struct_fields), + true, + )), + 2, + structs, + Some(null_buffer([true, false, true])), + ) + .unwrap(), + ) as ArrayRef; + Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new( + "item", + fixed_size_list.data_type().clone(), + true, + )), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1, 1, 3, 3])), + fixed_size_list, + Some(null_buffer([true, false, true, true])), + ) + .unwrap(), + ) + } + + fn nullable_struct() -> ArrayRef { + let fields = Fields::from(vec![ArrowField::new("value", DataType::Int32, true)]); + Arc::new(StructArray::new( + fields, + vec![Arc::new(Int32Array::from(vec![ + Some(10), + Some(20), + None, + Some(40), + Some(50), + ]))], + Some(null_buffer([true, false, true, true, false])), + )) + } + + fn deeply_nested() -> ArrayRef { + let leaf = Arc::new(Int32Array::from(vec![ + Some(0), + None, + Some(2), + Some(3), + None, + Some(5), + Some(6), + Some(7), + ])) as ArrayRef; + let inner = Arc::new( + LargeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i64, 2, 2, 3, 5, 5, 8])), + leaf, + Some(null_buffer([true, false, true, true, true, true])), + ) + .unwrap(), + ) as ArrayRef; + let struct_fields = Fields::from(vec![ArrowField::new( + "inner", + inner.data_type().clone(), + true, + )]); + let structs = Arc::new(StructArray::new( + struct_fields.clone(), + vec![inner], + Some(null_buffer([true, true, false, true, true, true])), + )) as ArrayRef; + Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new( + "item", + DataType::Struct(struct_fields), + true, + )), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 2, 2, 4, 6])), + structs, + Some(null_buffer([true, false, true, true, true])), + ) + .unwrap(), + ) + } + + fn page_layout(page: &EncodedPage) -> &pb21::page_layout::Layout { + let PageEncoding::Structural(layout) = &page.description else { + panic!("expected structural page encoding"); + }; + layout.layout.as_ref().expect("page layout must be present") + } + + fn create_encoder( + array: &ArrayRef, + version: LanceFileVersion, + metadata: HashMap, + ) -> Result> { + let arrow_field = + ArrowField::new("values", array.data_type().clone(), true).with_metadata(metadata); + let field = Field::try_from(&arrow_field)?; + let strategy = default_encoding_strategy(version); + let options = EncodingOptions { + cache_bytes_per_column: 1, + version, + ..Default::default() + }; + strategy.create_field_encoder( + strategy.as_ref(), + &field, + &mut ColumnIndexSequence::default(), + &options, + ) + } + + async fn encode_pages( + array: ArrayRef, + version: LanceFileVersion, + metadata: HashMap, + ) -> Result> { + encode_chunks(vec![array], version, metadata).await + } + + async fn encode_chunks( + arrays: Vec, + version: LanceFileVersion, + metadata: HashMap, + ) -> Result> { + let first = arrays + .first() + .ok_or_else(|| Error::internal("test input has no arrays".to_string()))?; + let mut encoder = create_encoder(first, version, metadata)?; + let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT); + let mut pages = Vec::new(); + let mut row_number = 0_u64; + for array in arrays { + let num_rows = array.len() as u64; + for task in encoder.maybe_encode( + array, + &mut external_buffers, + crate::repdef::RepDefBuilder::default(), + row_number, + num_rows, + )? { + pages.push(task.await?); + } + row_number += num_rows; + } + for task in encoder.flush(&mut external_buffers)? { + pages.push(task.await?); + } + for column in encoder.finish(&mut external_buffers).await? { + pages.extend(column.final_pages); + } + Ok(pages) + } + + fn sparse_layout(page: &EncodedPage) -> &pb21::SparseLayout { + let pb21::page_layout::Layout::SparseLayout(sparse) = page_layout(page) else { + panic!("expected SparseLayout, got {:?}", page_layout(page)); + }; + sparse + } + + fn list_layer(sparse: &pb21::SparseLayout) -> &pb21::SparseListLayer { + sparse + .structural_layers + .iter() + .find_map(|layer| match layer.layer.as_ref() { + Some(pb21::sparse_structural_layer::Layer::List(layer)) => Some(layer), + _ => None, + }) + .expect("expected sparse list layer") + } + + fn validity_layer(layer: &pb21::SparseStructuralLayer) -> Option<&pb21::SparseValidityLayer> { + match layer.layer.as_ref() { + Some(pb21::sparse_structural_layer::Layer::Validity(layer)) => Some(layer), + _ => None, + } + } + + fn layer_num_slots(layer: &pb21::SparseStructuralLayer) -> u64 { + match layer.layer.as_ref().expect("expected sparse layer variant") { + pb21::sparse_structural_layer::Layer::Validity(layer) => layer.num_slots, + pb21::sparse_structural_layer::Layer::List(layer) => layer.num_slots, + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => layer.num_slots, + } + } + + fn fixed_size_list_dimension(layer: &pb21::SparseStructuralLayer) -> Option { + match layer.layer.as_ref() { + Some(pb21::sparse_structural_layer::Layer::FixedSizeList(layer)) => { + Some(layer.dimension) + } + _ => None, + } + } + + fn planned_list( + offsets: Vec, + list_validity: Option>, + leaf_validity: Option>, + ) -> SparseStructuralPlan { + let num_values = u64::try_from(*offsets.last().unwrap()).unwrap(); + let mut builder = crate::repdef::RepDefBuilder::default(); + assert!(!builder.add_offsets( + OffsetBuffer::new(ScalarBuffer::from(offsets)), + list_validity.map(null_buffer), + )); + if let Some(leaf_validity) = leaf_validity { + builder.add_validity_bitmap(null_buffer(leaf_validity)); + } else { + builder.add_no_null(num_values as usize); + } + let normalized = crate::repdef::RepDefBuilder::normalize(vec![builder]); + plan(&normalized, num_values).unwrap() + } + + fn planned_list_layer(plan: &SparseStructuralPlan) -> &SparseStructuralLayerPlan { + plan.layers + .iter() + .find(|layer| matches!(layer, SparseStructuralLayerPlan::List { .. })) + .expect("expected planned list layer") + } + + #[test] + fn test_semantic_position_and_count_forms() { + let empty = planned_list(vec![0, 0, 0, 0], None, None); + assert_eq!(empty.num_items, 3); + assert_eq!(empty.num_visible_items, 0); + assert!(matches!( + planned_list_layer(&empty), + SparseStructuralLayerPlan::List { + non_empty_positions: SparsePositionSet::Empty, + counts: SparseCountSet::Empty, + .. + } + )); + + let all = planned_list(vec![0, 2, 4, 6], None, None); + assert_eq!(all.num_items, 6); + assert!(matches!( + planned_list_layer(&all), + SparseStructuralLayerPlan::List { + non_empty_positions: SparsePositionSet::All { len: 3 }, + counts: SparseCountSet::Constant { value: 2, len: 3 }, + .. + } + )); + + let range = planned_list(vec![0, 2, 4, 4, 4], None, None); + assert_eq!(range.num_items, 6); + assert!(matches!( + planned_list_layer(&range), + SparseStructuralLayerPlan::List { + non_empty_positions: SparsePositionSet::Range { start: 0, len: 2 }, + counts: SparseCountSet::Constant { value: 2, len: 2 }, + .. + } + )); + + let explicit = planned_list(vec![0, 1, 1, 4, 4, 6], None, None); + assert_eq!(explicit.num_items, 8); + assert!(matches!( + planned_list_layer(&explicit), + SparseStructuralLayerPlan::List { + non_empty_positions: SparsePositionSet::Explicit(positions), + counts: SparseCountSet::Explicit { counts, .. }, + .. + } if positions == &vec![0, 2, 4] && counts.as_ref() == [1, 3, 2] + )); + } + + #[test] + fn test_validity_polarity_uses_semantic_encoded_cost() { + let mostly_valid = BooleanBuffer::from_iter([true, false, true, true, false, true]); + let validity = validity_set(Some(&mostly_valid), mostly_valid.len(), "test").unwrap(); + assert_eq!(validity.meaning, SparseValidityMeaning::NullPositions); + assert!(matches!(validity.positions, SparsePositionSet::Explicit(ref p) if p == &[1, 4])); + + let mostly_null = BooleanBuffer::from_iter([false, true, false, false, true, false]); + let validity = validity_set(Some(&mostly_null), mostly_null.len(), "test").unwrap(); + assert_eq!(validity.meaning, SparseValidityMeaning::ValidPositions); + assert!(matches!(validity.positions, SparsePositionSet::Explicit(ref p) if p == &[1, 4])); + + let valid_island = BooleanBuffer::from_iter([false, false, true, true, false]); + let validity = validity_set(Some(&valid_island), valid_island.len(), "test").unwrap(); + assert_eq!(validity.meaning, SparseValidityMeaning::ValidPositions); + assert!(matches!( + validity.positions, + SparsePositionSet::Range { start: 2, len: 2 } + )); + + let all_valid = BooleanBuffer::from_iter([true, true, true]); + let validity = validity_set(Some(&all_valid), all_valid.len(), "test").unwrap(); + assert_eq!(validity.meaning, SparseValidityMeaning::NullPositions); + assert!(matches!(validity.positions, SparsePositionSet::Empty)); + + let all_null = BooleanBuffer::from_iter([false, false, false]); + let validity = validity_set(Some(&all_null), all_null.len(), "test").unwrap(); + assert_eq!(validity.meaning, SparseValidityMeaning::ValidPositions); + assert!(matches!(validity.positions, SparsePositionSet::Empty)); + } + + #[tokio::test] + async fn test_explicit_sparse_nullable_primitive_roundtrip() { + let array = Arc::new(Int32Array::from(vec![ + Some(10), + None, + Some(20), + Some(30), + None, + Some(40), + ])) as ArrayRef; + let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert_eq!(pages.len(), 1); + let sparse = sparse_layout(&pages[0]); + assert_eq!(sparse.num_items, 6); + assert_eq!(sparse.num_visible_items, 6); + assert_eq!(sparse.structural_layers.len(), 1); + let validity = validity_layer(&sparse.structural_layers[0]) + .unwrap() + .validity + .as_ref() + .unwrap(); + assert_eq!( + validity.meaning, + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions as i32 + ); + assert!(matches!( + validity.positions.as_ref().unwrap().positions, + Some(pb21::sparse_position_set::Positions::Explicit(_)) + )); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(1..5) + .with_indices(vec![0, 2, 5]); + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_nullable_struct_roundtrip() { + let array = nullable_struct(); + let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert_eq!(pages.len(), 1); + let sparse = sparse_layout(&pages[0]); + assert_eq!(sparse.structural_layers.len(), 2); + assert!( + sparse + .structural_layers + .iter() + .all(|layer| validity_layer(layer).is_some()) + ); + let struct_validity = validity_layer(&sparse.structural_layers[0]) + .unwrap() + .validity + .as_ref() + .unwrap(); + assert_eq!( + struct_validity.meaning, + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions as i32 + ); + assert!(matches!( + struct_validity.positions.as_ref().unwrap().positions, + Some(pb21::sparse_position_set::Positions::Explicit(_)) + )); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(1..5) + .with_indices(vec![0, 2, 4]); + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_struct_with_constant_and_sparse_children() { + let fields = Fields::from(vec![ + ArrowField::new("constant", DataType::Int32, true), + ArrowField::new("sparse", DataType::Int32, true), + ]); + let array = Arc::new(StructArray::new( + fields, + vec![ + Arc::new(Int32Array::from(vec![None::; 5])), + Arc::new(Int32Array::from(vec![ + Some(10), + Some(20), + None, + Some(40), + Some(50), + ])), + ], + Some(null_buffer([true, false, true, true, false])), + )) as ArrayRef; + let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert!(pages.iter().any(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::ConstantLayout(_) + ))); + assert!(pages.iter().any(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::SparseLayout(_) + ))); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(1..5) + .with_indices(vec![0, 2, 4]); + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_emits_both_validity_polarities() { + let mostly_valid = Arc::new(Int32Array::from(vec![ + Some(0), + None, + Some(2), + Some(3), + None, + Some(5), + ])) as ArrayRef; + let mostly_null = Arc::new(Int32Array::from(vec![ + None, + Some(1), + None, + None, + Some(4), + None, + ])) as ArrayRef; + + let null_positions = encode_pages( + mostly_valid.clone(), + LanceFileVersion::V2_3, + sparse_metadata(), + ) + .await + .unwrap(); + let validity = validity_layer(&sparse_layout(&null_positions[0]).structural_layers[0]) + .unwrap() + .validity + .as_ref() + .unwrap(); + assert_eq!( + validity.meaning, + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions as i32 + ); + + let valid_positions = encode_pages( + mostly_null.clone(), + LanceFileVersion::V2_3, + sparse_metadata(), + ) + .await + .unwrap(); + let validity = validity_layer(&sparse_layout(&valid_positions[0]).structural_layers[0]) + .unwrap() + .validity + .as_ref() + .unwrap(); + assert_eq!( + validity.meaning, + pb21::sparse_validity_set::Meaning::SparseValidityValidPositions as i32 + ); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(1..5) + .with_indices(vec![0, 2, 5]); + for array in [mostly_valid, mostly_null] { + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + } + + #[tokio::test] + async fn test_explicit_sparse_nested_page_boundaries_range_and_take() { + let nested = deeply_nested(); + let chunks = vec![nested.slice(0, 2), nested.slice(2, 3)]; + let pages = encode_chunks(chunks.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert!(pages.len() >= 2, "expected multiple sparse pages"); + let sparse = pages + .iter() + .map(sparse_layout) + .collect::>(); + assert!(sparse.iter().any(|layout| { + layout + .structural_layers + .iter() + .filter(|layer| { + matches!( + layer.layer.as_ref(), + Some(pb21::sparse_structural_layer::Layer::List(_)) + ) + }) + .count() + >= 2 + })); + assert!(sparse.iter().any(|layout| { + layout + .structural_layers + .iter() + .any(|layer| validity_layer(layer).is_some()) + })); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_batch_size(2) + .with_range(1..5) + .with_range(2..4) + .with_indices(vec![0, 2, 4]) + .with_indices(vec![1, 3]); + check_round_trip_encoding_of_data(chunks, &cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_list_and_large_list_null_empty_roundtrip() { + let list = list_i32( + vec![0, 0, 0, 2, 3, 3], + Some(vec![false, true, true, true, true]), + ); + let large_list = large_list_i32( + vec![0, 0, 0, 2, 3, 3], + Some(vec![false, true, true, true, true]), + ); + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(0..5) + .with_range(1..4) + .with_indices(vec![0, 1, 4]) + .with_indices(vec![2, 3]); + + for array in [list, large_list] { + let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert!(pages.iter().map(sparse_layout).all(|layout| { + layout.structural_layers.iter().any(|layer| { + matches!( + layer.layer.as_ref(), + Some(pb21::sparse_structural_layer::Layer::List(_)) + ) + }) + })); + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + } + + #[tokio::test] + async fn test_explicit_sparse_map_and_fixed_size_list_roundtrip() { + let map = map_i32(); + let map_pages = encode_pages(map.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert!(map_pages.iter().map(sparse_layout).any(|layout| { + layout.structural_layers.iter().any(|layer| { + matches!( + layer.layer.as_ref(), + Some(pb21::sparse_structural_layer::Layer::List(_)) + ) + }) + })); + let map_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(1..4) + .with_indices(vec![0, 2, 3]); + check_round_trip_encoding_of_data(vec![map], &map_cases, sparse_metadata()).await; + + let fsl = fixed_size_list_struct(); + let fsl_pages = encode_pages(fsl.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + for page in &fsl_pages { + let layout = sparse_layout(page); + let outer_slots = layer_num_slots(layout.structural_layers.first().unwrap()); + let fixed_size_scale = layout + .structural_layers + .iter() + .filter_map(fixed_size_list_dimension) + .product::(); + assert_eq!(page.num_rows, outer_slots * fixed_size_scale); + } + assert!(fsl_pages.iter().map(sparse_layout).any(|layout| { + layout + .structural_layers + .iter() + .any(|layer| fixed_size_list_dimension(layer) == Some(2)) + })); + let fsl_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(1..6) + .with_indices(vec![0, 3, 5]); + check_round_trip_encoding_of_data(vec![fsl], &fsl_cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_list_fixed_size_list_struct_roundtrip() { + let array = list_fixed_size_list_struct(); + let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert!(pages.iter().map(sparse_layout).any(|layout| { + let kinds = layout + .structural_layers + .iter() + .map(|layer| match layer.layer.as_ref().unwrap() { + pb21::sparse_structural_layer::Layer::Validity(_) => "validity", + pb21::sparse_structural_layer::Layer::List(_) => "list", + pb21::sparse_structural_layer::Layer::FixedSizeList(_) => "fixed-size-list", + }) + .collect::>(); + kinds.starts_with(&["list", "fixed-size-list"]) + })); + for page in &pages { + let layout = sparse_layout(page); + let outer_slots = layer_num_slots(layout.structural_layers.first().unwrap()); + assert_eq!(page.num_rows, outer_slots * 2); + } + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(1..4) + .with_indices(vec![0, 2, 3]) + .with_indices(vec![1, 3]); + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_serializes_semantic_list_forms() { + let all_array = list_i32(vec![0, 2, 4, 6], None); + let all = encode_pages(all_array.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + let all_layer = list_layer(sparse_layout(&all[0])); + assert!(matches!( + all_layer.non_empty_positions.as_ref().unwrap().positions, + Some(pb21::sparse_position_set::Positions::All(_)) + )); + assert!(matches!( + all_layer.counts.as_ref().unwrap().counts, + Some(pb21::sparse_count_set::Counts::Constant( + pb21::SparseCountConstant { value: 2 } + )) + )); + assert_eq!(all[0].data.len(), 2); + + let range_array = list_i32(vec![0, 2, 4, 4, 4], None); + let range = encode_pages( + range_array.clone(), + LanceFileVersion::V2_3, + sparse_metadata(), + ) + .await + .unwrap(); + let range_layer = list_layer(sparse_layout(&range[0])); + assert!(matches!( + range_layer.non_empty_positions.as_ref().unwrap().positions, + Some(pb21::sparse_position_set::Positions::Range( + pb21::SparsePositionRange { + start: 0, + length: 2 + } + )) + )); + assert!(matches!( + range_layer.counts.as_ref().unwrap().counts, + Some(pb21::sparse_count_set::Counts::Constant( + pb21::SparseCountConstant { value: 2 } + )) + )); + assert_eq!(range[0].data.len(), 2); + + let explicit_array = list_i32(vec![0, 1, 1, 4, 4, 6], None); + let explicit = encode_pages( + explicit_array.clone(), + LanceFileVersion::V2_3, + sparse_metadata(), + ) + .await + .unwrap(); + let explicit_layer = list_layer(sparse_layout(&explicit[0])); + assert!(matches!( + explicit_layer + .non_empty_positions + .as_ref() + .unwrap() + .positions, + Some(pb21::sparse_position_set::Positions::Explicit(_)) + )); + assert!(matches!( + explicit_layer.counts.as_ref().unwrap().counts, + Some(pb21::sparse_count_set::Counts::Explicit(_)) + )); + assert_eq!(explicit[0].data.len(), 4); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(1..3) + .with_indices(vec![0, 2]); + for array in [all_array, range_array, explicit_array] { + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + } + + #[tokio::test] + async fn test_constant_layout_boundary_is_explicit() { + let structural_only = encode_pages( + list_i32(vec![0, 0, 0, 0], None), + LanceFileVersion::V2_3, + sparse_metadata(), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&structural_only[0]), + pb21::page_layout::Layout::ConstantLayout(_) + )); + + let empty_struct = Arc::new(StructArray::new_empty_fields(3, None)) as ArrayRef; + let empty_struct_pages = + encode_pages(empty_struct, LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert!(matches!( + page_layout(&empty_struct_pages[0]), + pb21::page_layout::Layout::ConstantLayout(_) + )); + + let all_null = Arc::new(Int32Array::from(vec![None, None, None])) as ArrayRef; + let all_null_pages = encode_pages(all_null, LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert!(matches!( + page_layout(&all_null_pages[0]), + pb21::page_layout::Layout::ConstantLayout(_) + )); + + let constant = Arc::new(Int32Array::from(vec![7, 7, 7])) as ArrayRef; + let constant_pages = encode_pages(constant, LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert!(matches!( + page_layout(&constant_pages[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + } + + #[tokio::test] + async fn test_default_and_explicit_dense_layouts_are_unchanged() { + let array = Arc::new(Int32Array::from_iter_values(0..16)) as ArrayRef; + let v2_2_default = encode_pages(array.clone(), LanceFileVersion::V2_2, HashMap::new()) + .await + .unwrap(); + let v2_2_miniblock = encode_pages( + array.clone(), + LanceFileVersion::V2_2, + structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), + ) + .await + .unwrap(); + assert_eq!(v2_2_default.len(), v2_2_miniblock.len()); + for (default, explicit) in v2_2_default.iter().zip(v2_2_miniblock.iter()) { + assert!(matches!( + page_layout(default), + pb21::page_layout::Layout::MiniBlockLayout(_) + )); + assert_eq!(page_layout(default), page_layout(explicit)); + assert_eq!(default.data, explicit.data); + } + + let v2_3_default = encode_pages(array.clone(), LanceFileVersion::V2_3, HashMap::new()) + .await + .unwrap(); + assert!(matches!( + page_layout(&v2_3_default[0]), + pb21::page_layout::Layout::MiniBlockLayout(_) + )); + + let v2_3_miniblock = encode_pages( + array.clone(), + LanceFileVersion::V2_3, + structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&v2_3_miniblock[0]), + pb21::page_layout::Layout::MiniBlockLayout(_) + )); + + let v2_3_fullzip = encode_pages( + array.clone(), + LanceFileVersion::V2_3, + structural_metadata(STRUCTURAL_ENCODING_FULLZIP), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&v2_3_fullzip[0]), + pb21::page_layout::Layout::FullZipLayout(_) + )); + + let v2_3_sparse = encode_pages(array, LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert!(matches!( + page_layout(&v2_3_sparse[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + } + + #[test] + fn test_explicit_sparse_rejects_lance_2_2() { + let array = Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef; + let Err(error) = create_encoder(&array, LanceFileVersion::V2_2, sparse_metadata()) else { + panic!("expected Lance 2.2 to reject explicit sparse encoding"); + }; + assert!( + error + .to_string() + .contains("requires Lance file format 2.3+") + ); + + let structural_only = list_i32(vec![0, 0, 0], None); + let Err(error) = + create_encoder(&structural_only, LanceFileVersion::V2_2, sparse_metadata()) + else { + panic!("expected Lance 2.2 structural-only input to reject explicit sparse encoding"); + }; + assert!( + error + .to_string() + .contains("requires Lance file format 2.3+") + ); + } +} diff --git a/rust/lance-encoding/src/encodings/logical/struct.rs b/rust/lance-encoding/src/encodings/logical/struct.rs index e8c6289275a..dc734da72b5 100644 --- a/rust/lance-encoding/src/encodings/logical/struct.rs +++ b/rust/lance-encoding/src/encodings/logical/struct.rs @@ -368,26 +368,57 @@ impl StructuralDecodeArrayTask for RepDefStructDecodeTask { .map(|task| task.decode()) .collect::>>()?; let mut children = Vec::with_capacity(arrays.len()); + let mut repdefs = Vec::with_capacity(arrays.len()); let mut data_size = 0u64; let mut arrays_iter = arrays.into_iter(); - let first_array = arrays_iter.next().unwrap(); + let first_array = arrays_iter.next().ok_or_else(|| { + Error::internal("Struct decoder unexpectedly has no child arrays".to_string()) + })?; let length = first_array.array.len(); // The repdef should be identical across all children at this point - let mut repdef = first_array.repdef; + repdefs.push(first_array.repdef); data_size += first_array.data_size; children.push(first_array.array); for array in arrays_iter { - debug_assert_eq!(length, array.array.len()); + if length != array.array.len() { + return Err(Error::invalid_input_source( + format!( + "Struct child array length {} does not match sibling length {}", + array.array.len(), + length + ) + .into(), + )); + } data_size += array.data_size; children.push(array.array); + repdefs.push(array.repdef); + } + + // Dense rep/def state can retain child-specific repetition information after a child + // decoder finishes, so comparing dense siblings is not meaningful. If any child carries + // sparse state, keep a sparse child as the canonical structural plan and compare it with + // every other sparse sibling so sparse metadata is never silently discarded. + let primary_repdef = repdefs + .iter() + .position(CompositeRepDefUnraveler::has_sparse) + .unwrap_or(0); + let mut repdef = repdefs.swap_remove(primary_repdef); + if repdef.has_sparse() { + for sibling in repdefs { + if sibling.has_sparse() { + repdef.add_compatibility_check(sibling); + } + } } let validity = if self.is_root { + repdef.ensure_exhausted()?; None } else { - repdef.unravel_validity(length) + repdef.unravel_validity(length)? }; let array = StructArray::try_new(self.child_fields, children, validity) diff --git a/rust/lance-encoding/src/repdef.rs b/rust/lance-encoding/src/repdef.rs index b418b906de8..5aa79726917 100644 --- a/rust/lance-encoding/src/repdef.rs +++ b/rust/lance-encoding/src/repdef.rs @@ -118,7 +118,10 @@ use arrow_buffer::{ }; use lance_core::{Error, Result, utils::bit::log_2_ceil}; -use crate::buffer::LanceBuffer; +use crate::{ + buffer::LanceBuffer, + encodings::logical::primitive::sparse::{SparseStructuralPlan, SparseStructuralUnraveler}, +}; pub type LevelBuffer = Vec; @@ -199,6 +202,143 @@ enum RawRepDef { Fsl(FslDesc), } +/// A normalized Arrow structural layer shared by dense and sparse serializers. +#[derive(Clone, Copy, Debug)] +pub(crate) enum NormalizedStructuralLayer<'a> { + List { + offsets: &'a [i64], + validity: Option<&'a BooleanBuffer>, + num_slots: usize, + }, + Validity { + validity: Option<&'a BooleanBuffer>, + num_slots: usize, + }, + FixedSizeList { + validity: Option<&'a BooleanBuffer>, + dimension: usize, + num_slots: usize, + }, +} + +/// Structural layers concatenated across input batches exactly once. +/// +/// Dense rep/def serialization and sparse metadata planning both consume this +/// representation so the Arrow nesting is not independently reconstructed. +#[derive(Debug)] +pub(crate) struct NormalizedStructuralPlan { + layers: Vec, + dense_all_valid: bool, +} + +impl NormalizedStructuralPlan { + pub(crate) fn layers(&self) -> impl ExactSizeIterator> { + self.layers.iter().map(|layer| match layer { + RawRepDef::Offsets(OffsetDesc { + offsets, + validity, + num_values, + .. + }) => NormalizedStructuralLayer::List { + offsets, + validity: validity.as_ref(), + num_slots: *num_values, + }, + RawRepDef::Validity(ValidityDesc { + validity, + num_values, + }) => NormalizedStructuralLayer::Validity { + validity: validity.as_ref(), + num_slots: *num_values, + }, + RawRepDef::Fsl(FslDesc { + validity, + dimension, + num_values, + }) => NormalizedStructuralLayer::FixedSizeList { + validity: validity.as_ref(), + dimension: *dimension, + num_slots: *num_values, + }, + }) + } + + fn into_serializer(self) -> (SerializerContext, Option) { + if self.dense_all_valid { + let def_meaning = self + .layers + .iter() + .map(|_| DefinitionInterpretation::AllValidItem) + .collect::>(); + return ( + SerializerContext { + def_meaning, + rep_levels: LevelBuffer::default(), + spare_rep: LevelBuffer::default(), + def_levels: LevelBuffer::default(), + spare_def: LevelBuffer::default(), + current_rep: 0, + current_def: 0, + current_len: 0, + current_num_specials: 0, + has_fsl: false, + }, + None, + ); + } + + let total_len = self.layers.last().map_or(0, RawRepDef::num_values) + + self + .layers + .iter() + .map(RawRepDef::num_specials) + .sum::(); + let max_rep = self.layers.iter().map(RawRepDef::max_rep).sum::(); + let max_def = self.layers.iter().map(RawRepDef::max_def).sum::(); + let bits_per_rep = if max_rep > 0 { + u64::from(u16::BITS - max_rep.leading_zeros()) + } else { + 0 + }; + let bits_per_def = if max_def > 0 { + u64::from(u16::BITS - max_def.leading_zeros()) + } else { + 0 + }; + let bits_per_level = + (bits_per_rep + bits_per_def > 0).then_some(bits_per_rep + bits_per_def); + + let num_layers = self.layers.len(); + let mut context = SerializerContext::new(total_len, num_layers, max_rep, max_def); + for layer in self.layers { + match layer { + RawRepDef::Validity(def) => context.record_validity(&def), + RawRepDef::Offsets(rep) => context.record_offsets(&rep), + RawRepDef::Fsl(fsl) => context.record_fsl(&fsl), + } + } + (context, bits_per_level) + } + + pub(crate) fn serialize(self) -> SerializedRepDefs { + self.into_serializer().0.build() + } + + pub(crate) fn serialize_with_structural_plan( + self, + max_levels_for_bits: impl FnOnce(u64) -> u64, + num_rows: u64, + num_values: u64, + ) -> Result<(SerializedRepDefs, StructuralPagePlan)> { + let (context, bits_per_level) = self.into_serializer(); + context.build_with_structural_plan( + bits_per_level.map(max_levels_for_bits), + num_rows, + num_values, + ) + } +} + impl RawRepDef { // Are there any nulls in this layer fn has_nulls(&self) -> bool { @@ -1409,54 +1549,18 @@ impl RepDefBuilder { /// Converts the validity / offsets buffers that have been gathered so far /// into repetition and definition levels pub fn serialize(builders: Vec) -> SerializedRepDefs { - Self::serialize_builders(builders).0.build() + Self::normalize(builders).serialize() } - /// Converts gathered structural buffers into rep/def levels and an encode-time plan. - pub(crate) fn serialize_with_structural_plan( - builders: Vec, - max_levels_for_bits: impl FnOnce(u64) -> u64, - num_rows: u64, - num_values: u64, - ) -> Result<(SerializedRepDefs, StructuralPagePlan)> { - let (context, bits_per_level) = Self::serialize_builders(builders); - context.build_with_structural_plan( - bits_per_level.map(max_levels_for_bits), - num_rows, - num_values, - ) - } - - fn serialize_builders(builders: Vec) -> (SerializerContext, Option) { + pub(crate) fn normalize(builders: Vec) -> NormalizedStructuralPlan { assert!(!builders.is_empty()); - if builders.iter().all(|b| b.is_empty()) { - // No repetition, all-valid - let def_meaning = builders - .first() - .unwrap() - .repdefs - .iter() - .map(|_| DefinitionInterpretation::AllValidItem) - .collect::>(); - return ( - SerializerContext { - def_meaning, - rep_levels: LevelBuffer::default(), - spare_rep: LevelBuffer::default(), - def_levels: LevelBuffer::default(), - spare_def: LevelBuffer::default(), - current_rep: 0, - current_def: 0, - current_len: 0, - current_num_specials: 0, - has_fsl: false, - }, - None, - ); - } - let num_layers = builders[0].num_layers(); - let combined_layers = (0..num_layers) + debug_assert!( + builders + .iter() + .all(|builder| builder.num_layers() == num_layers) + ); + let layers = (0..num_layers) .map(|layer_index| { Self::concat_layers( builders.iter().map(|b| &b.repdefs[layer_index]), @@ -1464,47 +1568,10 @@ impl RepDefBuilder { ) }) .collect::>(); - debug_assert!( - builders - .iter() - .all(|b| b.num_layers() == builders[0].num_layers()) - ); - - let total_len = combined_layers.last().unwrap().num_values() - + combined_layers - .iter() - .map(|l| l.num_specials()) - .sum::(); - let max_rep = combined_layers.iter().map(|l| l.max_rep()).sum::(); - let max_def = combined_layers.iter().map(|l| l.max_def()).sum::(); - let bits_per_rep = if max_rep > 0 { - u64::from(u16::BITS - max_rep.leading_zeros()) - } else { - 0 - }; - let bits_per_def = if max_def > 0 { - u64::from(u16::BITS - max_def.leading_zeros()) - } else { - 0 - }; - let bits_per_level = - (bits_per_rep + bits_per_def > 0).then_some(bits_per_rep + bits_per_def); - - let mut context = SerializerContext::new(total_len, num_layers, max_rep, max_def); - for layer in combined_layers.into_iter() { - match layer { - RawRepDef::Validity(def) => { - context.record_validity(&def); - } - RawRepDef::Offsets(rep) => { - context.record_offsets(&rep); - } - RawRepDef::Fsl(fsl) => { - context.record_fsl(&fsl); - } - } + NormalizedStructuralPlan { + layers, + dense_all_valid: builders.iter().all(Self::is_empty), } - (context, bits_per_level) } } @@ -1514,6 +1581,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 +1635,7 @@ impl RepDefUnraveler { } } Self { + sparse: None, rep_levels, def_levels, current_def_cmp: 0, @@ -1578,7 +1647,35 @@ impl RepDefUnraveler { } } + pub(crate) fn new_sparse(plan: SparseStructuralPlan) -> Self { + Self { + sparse: Some(SparseStructuralUnraveler::new(plan)), + rep_levels: None, + def_levels: None, + levels_to_rep: Vec::new(), + def_meaning: Arc::new([]), + current_def_cmp: 0, + current_rep_cmp: 0, + current_layer: 0, + num_items: 0, + } + } + + fn ensure_exhausted(&self) -> Result<()> { + if let Some(sparse) = &self.sparse { + sparse.ensure_exhausted()?; + } + Ok(()) + } + + fn is_sparse(&self) -> bool { + self.sparse.is_some() + } + 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() } @@ -1587,15 +1684,19 @@ 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 { + pub fn max_lists(&self) -> Result { + if let Some(sparse) = &self.sparse { + return sparse.max_lists(); + } debug_assert!( self.def_meaning[self.current_layer] != DefinitionInterpretation::NullableItem ); - self.rep_levels + Ok(self + .rep_levels .as_ref() // Worst case every rep item is max_rep and a new list .map(|levels| levels.len()) - .unwrap_or(0) + .unwrap_or(0)) } /// Unravels a layer of offsets from the unraveler into the given offset width @@ -1607,6 +1708,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() @@ -1757,18 +1861,25 @@ impl RepDefUnraveler { } } - pub fn skip_validity(&mut self) { + pub fn skip_validity(&mut self) -> Result<()> { + if let Some(sparse) = self.sparse.as_mut() { + return sparse.skip_validity(); + } debug_assert!(self.is_all_valid()); self.current_layer += 1; + Ok(()) } /// Unravels a layer of validity from the definition levels - pub fn unravel_validity(&mut self, validity: &mut BooleanBufferBuilder) { + pub fn unravel_validity(&mut self, validity: &mut BooleanBufferBuilder) -> Result<()> { + if let Some(sparse) = self.sparse.as_mut() { + return sparse.unravel_validity(validity); + } let meaning = self.def_meaning[self.current_layer]; if meaning == DefinitionInterpretation::AllValidItem || self.def_levels.is_none() { self.current_layer += 1; validity.append_n(self.num_items as usize, true); - return; + return Ok(()); } self.current_layer += 1; @@ -1786,9 +1897,13 @@ impl RepDefUnraveler { }) { validity.append(is_valid); } + Ok(()) } - pub fn decimate(&mut self, dimension: usize) { + pub fn decimate(&mut self, dimension: usize) -> Result<()> { + if let Some(sparse) = self.sparse.as_mut() { + return sparse.decimate(dimension); + } 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: @@ -1804,7 +1919,7 @@ impl RepDefUnraveler { todo!("Not yet supported FSL<...List<...>>"); } let Some(def_levels) = self.def_levels.as_mut() else { - return; + return Ok(()); }; let mut read_idx = 0; let mut write_idx = 0; @@ -1816,6 +1931,7 @@ impl RepDefUnraveler { read_idx += dimension; } def_levels.truncate(write_idx); + Ok(()) } } @@ -1835,44 +1951,104 @@ impl RepDefUnraveler { #[derive(Debug)] pub struct CompositeRepDefUnraveler { unravelers: Vec, + comparisons: Vec, } impl CompositeRepDefUnraveler { pub fn new(unravelers: Vec) -> Self { - Self { unravelers } + Self { + unravelers, + comparisons: Vec::new(), + } + } + + pub(crate) fn add_compatibility_check(&mut self, other: Self) { + self.comparisons.push(other); + } + + pub(crate) fn has_sparse(&self) -> bool { + self.unravelers.iter().any(RepDefUnraveler::is_sparse) + || self.comparisons.iter().any(Self::has_sparse) + } + + pub(crate) fn ensure_exhausted(&self) -> Result<()> { + for unraveler in &self.unravelers { + unraveler.ensure_exhausted()?; + } + for comparison in &self.comparisons { + comparison.ensure_exhausted()?; + } + Ok(()) + } + + fn null_buffers_equal( + left: &Option, + right: &Option, + expected_len: usize, + ) -> bool { + match (left, right) { + (None, None) => true, + (Some(left), Some(right)) => { + left.len() == expected_len + && right.len() == expected_len + && left.iter().eq(right.iter()) + } + (None, Some(right)) => right.len() == expected_len && right.null_count() == 0, + (Some(left), None) => left.len() == expected_len && left.null_count() == 0, + } + } + + fn decimate(&mut self, dimension: usize) -> Result<()> { + for unraveler in &mut self.unravelers { + unraveler.decimate(dimension)?; + } + for comparison in &mut self.comparisons { + comparison.decimate(dimension)?; + } + Ok(()) } /// Unravels a layer of validity /// /// Returns None if there are no null items in this layer - pub fn unravel_validity(&mut self, num_values: usize) -> Option { + pub fn unravel_validity(&mut self, num_values: usize) -> Result> { let is_all_valid = self .unravelers .iter() .all(|unraveler| unraveler.is_all_valid()); - if is_all_valid { + let validity = if is_all_valid { for unraveler in self.unravelers.iter_mut() { - unraveler.skip_validity(); + unraveler.skip_validity()?; } None } else { let mut validity = BooleanBufferBuilder::new(num_values); for unraveler in self.unravelers.iter_mut() { - unraveler.unravel_validity(&mut validity); + unraveler.unravel_validity(&mut validity)?; } Some(NullBuffer::new(validity.finish())) + }; + for comparison in &mut self.comparisons { + let other = comparison.unravel_validity(num_values)?; + if !Self::null_buffers_equal(&validity, &other, num_values) { + return Err(Error::invalid_input_source( + format!( + "Structural sibling fields have incompatible validity metadata for {num_values} values" + ) + .into(), + )); + } } + Ok(validity) } pub fn unravel_fsl_validity( &mut self, num_values: usize, dimension: usize, - ) -> Option { - for unraveler in self.unravelers.iter_mut() { - unraveler.decimate(dimension); - } + ) -> Result> { + self.decimate(dimension)?; self.unravel_validity(num_values) } @@ -1881,10 +2057,16 @@ impl CompositeRepDefUnraveler { &mut self, ) -> Result<(OffsetBuffer, Option)> { let mut is_all_valid = true; - let mut max_num_lists = 0; + let mut max_num_lists: usize = 0; for unraveler in self.unravelers.iter() { is_all_valid &= unraveler.is_all_valid(); - max_num_lists += unraveler.max_lists(); + max_num_lists = max_num_lists + .checked_add(unraveler.max_lists()?) + .ok_or_else(|| { + Error::invalid_input_source( + "Combined repetition/definition list count exceeds usize::MAX".into(), + ) + })?; } let mut validity = if is_all_valid { @@ -1901,10 +2083,28 @@ impl CompositeRepDefUnraveler { unraveler.unravel_offsets(&mut offsets, validity.as_mut())?; } - Ok(( - OffsetBuffer::new(ScalarBuffer::from(offsets)), - validity.map(|mut v| NullBuffer::new(v.finish())), - )) + let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets)); + let validity = validity.map(|mut v| NullBuffer::new(v.finish())); + for comparison in &mut self.comparisons { + let (other_offsets, other_validity) = comparison.unravel_offsets::()?; + if offsets.as_ref() != other_offsets.as_ref() + || !Self::null_buffers_equal( + &validity, + &other_validity, + offsets.len().saturating_sub(1), + ) + { + return Err(Error::invalid_input_source( + format!( + "Structural sibling fields have incompatible list metadata for {} slots", + offsets.len().saturating_sub(1) + ) + .into(), + )); + } + } + + Ok((offsets, validity)) } } @@ -2603,6 +2803,10 @@ impl ControlWordParser { mod tests { use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; + use crate::encodings::logical::primitive::sparse::{ + SparsePositionSet, SparseStructuralLayerPlan, SparseStructuralPlan, SparseValidityMeaning, + SparseValiditySet, + }; use crate::repdef::{ CompositeRepDefUnraveler, DefinitionInterpretation, RepDefUnraveler, SerializedRepDefs, }; @@ -2621,6 +2825,31 @@ mod tests { OffsetBuffer::::new(ScalarBuffer::from_iter(values.iter().copied())) } + #[test] + fn sparse_sibling_validity_mismatch_is_invalid_input() { + let sparse = |positions| { + RepDefUnraveler::new_sparse(SparseStructuralPlan { + layers: vec![SparseStructuralLayerPlan::Validity { + num_slots: 2, + validity: SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions, + }, + }], + num_items: 2, + num_visible_items: 2, + }) + }; + let mut repdef = CompositeRepDefUnraveler::new(vec![sparse(SparsePositionSet::Empty)]); + repdef.add_compatibility_check(CompositeRepDefUnraveler::new(vec![sparse( + SparsePositionSet::Explicit(vec![0]), + )])); + + let err = repdef.unravel_validity(2).unwrap_err(); + assert!(matches!(err, lance_core::Error::InvalidInput { .. })); + assert!(err.to_string().contains("incompatible validity metadata")); + } + #[test] fn test_repdef_empty_offsets() { // Empty offsets should serialize without panicking. @@ -2666,7 +2895,7 @@ mod tests { // Note: validity doesn't exactly round-trip because repdef normalizes some of the // redundant validity values assert_eq!( - unraveler.unravel_validity(9), + unraveler.unravel_validity(9).unwrap(), Some(validity(&[ true, true, true, false, false, false, true, true, false ])) @@ -2804,14 +3033,14 @@ mod tests { )]); assert_eq!( - unraveler.unravel_validity(8), + unraveler.unravel_validity(8).unwrap(), Some(validity(&[ true, false, true, false, false, false, false, false ])) ); - assert_eq!(unraveler.unravel_fsl_validity(4, 2), None); + assert_eq!(unraveler.unravel_fsl_validity(4, 2).unwrap(), None); assert_eq!( - unraveler.unravel_fsl_validity(2, 2), + unraveler.unravel_fsl_validity(2, 2).unwrap(), Some(validity(&[true, false])) ); } @@ -2847,10 +3076,10 @@ mod tests { 8, )]); - assert_eq!(unraveler.unravel_validity(8), None); - assert_eq!(unraveler.unravel_fsl_validity(4, 2), None); + assert_eq!(unraveler.unravel_validity(8).unwrap(), None); + assert_eq!(unraveler.unravel_fsl_validity(4, 2).unwrap(), None); assert_eq!( - unraveler.unravel_fsl_validity(2, 2), + unraveler.unravel_fsl_validity(2, 2).unwrap(), Some(validity(&[true, false])) ); } @@ -2928,7 +3157,7 @@ mod tests { 8, )]); - assert_eq!(unraveler.unravel_validity(6), None); + assert_eq!(unraveler.unravel_validity(6).unwrap(), None); let (off, val) = unraveler.unravel_offsets::().unwrap(); assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner()); assert_eq!(val, None); @@ -2954,7 +3183,7 @@ mod tests { 9, )]); - assert_eq!(unraveler.unravel_validity(9), None); + assert_eq!(unraveler.unravel_validity(9).unwrap(), None); let (off, val) = unraveler.unravel_offsets::().unwrap(); assert_eq!(off.inner(), offsets_32(&[0, 1, 3, 5, 7, 9]).inner()); assert_eq!(val, None); @@ -3018,7 +3247,7 @@ mod tests { 8, )]); - assert_eq!(unraveler.unravel_validity(6), None); + assert_eq!(unraveler.unravel_validity(6).unwrap(), None); let (off, val) = unraveler.unravel_offsets::().unwrap(); assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner()); assert_eq!(val, None); @@ -3048,7 +3277,7 @@ mod tests { 8, )]); - assert_eq!(unraveler.unravel_validity(6), None); + assert_eq!(unraveler.unravel_validity(6).unwrap(), None); let (off, val) = unraveler.unravel_offsets::().unwrap(); assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner()); assert_eq!(val, Some(validity(&[true, false, false, true]))); @@ -3078,7 +3307,7 @@ mod tests { 8, )]); - assert_eq!(unraveler.unravel_validity(6), None); + assert_eq!(unraveler.unravel_validity(6).unwrap(), None); let (off, val) = unraveler.unravel_offsets::().unwrap(); assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner()); assert_eq!(val, Some(validity(&[true, false, true, true]))); @@ -3106,11 +3335,11 @@ mod tests { )]); assert_eq!( - unraveler.unravel_validity(4), + unraveler.unravel_validity(4).unwrap(), Some(validity(&[false, true, false, false])) ); assert_eq!( - unraveler.unravel_validity(4), + unraveler.unravel_validity(4).unwrap(), Some(validity(&[false, true, false, false])) ); let (off, val) = unraveler.unravel_offsets::().unwrap(); @@ -3139,14 +3368,14 @@ mod tests { )]); assert_eq!( - unraveler.unravel_validity(5), + unraveler.unravel_validity(5).unwrap(), Some(validity(&[false, false, true, true, false])) ); assert_eq!( - unraveler.unravel_validity(5), + unraveler.unravel_validity(5).unwrap(), Some(validity(&[false, false, true, true, true])) ); - assert_eq!(unraveler.unravel_validity(5), None); + assert_eq!(unraveler.unravel_validity(5).unwrap(), None); } #[test] @@ -3188,7 +3417,7 @@ mod tests { let mut unraveler = CompositeRepDefUnraveler::new(vec![unravel1, unravel2]); - assert!(unraveler.unravel_validity(9).is_none()); + assert!(unraveler.unravel_validity(9).unwrap().is_none()); let (off, val) = unraveler.unravel_offsets::().unwrap(); assert_eq!( off.inner(), @@ -3484,11 +3713,11 @@ mod tests { 0, )]); - assert_eq!(unraveler.unravel_validity(0), None); + assert_eq!(unraveler.unravel_validity(0).unwrap(), None); let (off, val) = unraveler.unravel_offsets::().unwrap(); assert_eq!(off.inner(), offsets_32(&[0, 0, 0, 0]).inner()); assert_eq!(val, Some(validity(&[false, false, false]))); - let val = unraveler.unravel_validity(3).unwrap(); + let val = unraveler.unravel_validity(3).unwrap().unwrap(); assert_eq!(val.inner(), validity(&[true, false, true]).inner()); } @@ -3516,7 +3745,7 @@ mod tests { 1, )]); - assert_eq!(unraveler.unravel_validity(1), None); + assert_eq!(unraveler.unravel_validity(1).unwrap(), None); let (off, val) = unraveler.unravel_offsets::().unwrap(); assert_eq!(off.inner(), offsets_32(&[0, 1, 1]).inner()); assert_eq!(val, Some(validity(&[true, false]))); @@ -3547,7 +3776,7 @@ mod tests { ]); assert_eq!( - unraveler.unravel_validity(8), + unraveler.unravel_validity(8).unwrap(), Some(validity(&[ true, false, true, false, true, true, true, true ])) @@ -3584,7 +3813,7 @@ mod tests { ]); assert_eq!( - unraveler.unravel_validity(4), + unraveler.unravel_validity(4).unwrap(), Some(validity(&[true, false, true, true])) ); assert_eq!( @@ -3616,7 +3845,7 @@ mod tests { ]); assert_eq!( - unraveler.unravel_validity(8), + unraveler.unravel_validity(8).unwrap(), Some(validity(&[ true, false, true, false, true, true, true, true ])) diff --git a/rust/lance-encoding/src/testing.rs b/rust/lance-encoding/src/testing.rs index 176083d6d64..9825cb5949d 100644 --- a/rust/lance-encoding/src/testing.rs +++ b/rust/lance-encoding/src/testing.rs @@ -666,6 +666,11 @@ fn collect_page_encoding(layout: &PageLayout, actual_chain: &mut Vec) -> collect_page_encoding(inner_layout.as_ref(), actual_chain)? } } + Layout::SparseLayout(sparse) => { + if let Some(value_compression) = &sparse.value_compression { + actual_chain.extend(extract_array_encoding_chain(value_compression)); + } + } } } diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index cea89235d73..4c70e27c562 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -919,11 +919,17 @@ impl FileReader { let mut global_bufs_cursor = Cursor::new(gbo_bytes); let mut global_buffers = Vec::with_capacity(footer.num_global_buffers as usize); - for _ in 0..footer.num_global_buffers { + for buffer_index in 0..footer.num_global_buffers { let buf_pos = global_bufs_cursor.read_u64::()?; - assert!( - version < LanceFileVersion::V2_1 || buf_pos % PAGE_BUFFER_ALIGNMENT as u64 == 0 - ); + if version >= LanceFileVersion::V2_1 && buf_pos % PAGE_BUFFER_ALIGNMENT as u64 != 0 { + return Err(Error::invalid_input_source( + format!( + "Global buffer {} position {} is not aligned to {} bytes", + buffer_index, buf_pos, PAGE_BUFFER_ALIGNMENT + ) + .into(), + )); + } let buf_size = global_bufs_cursor.read_u64::()?; global_buffers.push(BufferDescriptor { position: buf_pos, @@ -1021,7 +1027,7 @@ impl FileReader { let num_data_bytes = footer.column_meta_start - num_global_buffer_bytes; let num_column_metadata_bytes = footer.global_buff_offsets_start - footer.column_meta_start; - let column_infos = Self::meta_to_col_infos(column_metadatas.as_slice(), file_version); + let column_infos = Self::meta_to_col_infos(column_metadatas.as_slice(), file_version)?; // The tail read above already pulled in any global buffer that lives within // the captured window. Copy those user buffers (index >= 1; the schema at 0 @@ -1129,74 +1135,139 @@ impl FileReader { Self::read_metadata_index_with_known_schema(scheduler, Some((file_schema, num_rows))).await } - fn fetch_encoding(encoding: &pbfile::Encoding) -> M { + fn fetch_encoding(encoding: &pbfile::Encoding) -> Result { match &encoding.location { - Some(pbfile::encoding::Location::Indirect(_)) => todo!(), + Some(pbfile::encoding::Location::Indirect(_)) => Err(Error::invalid_input_source( + "Indirect file encodings are not supported".into(), + )), Some(pbfile::encoding::Location::Direct(encoding)) => { let encoding_buf = Bytes::from(encoding.encoding.clone()); - let encoding_any = prost_types::Any::decode(encoding_buf).unwrap(); - encoding_any.to_msg::().unwrap() + let encoding_any = prost_types::Any::decode(encoding_buf).map_err(|error| { + Error::invalid_input_source( + format!("Invalid direct {} encoding envelope: {error}", M::NAME).into(), + ) + })?; + encoding_any.to_msg::().map_err(|error| { + Error::invalid_input_source( + format!("Invalid direct {} encoding: {error}", M::NAME).into(), + ) + }) } - Some(pbfile::encoding::Location::None(_)) => panic!(), - None => panic!(), + Some(pbfile::encoding::Location::None(_)) => Err(Error::invalid_input_source( + format!("Missing {} encoding description", M::NAME).into(), + )), + None => Err(Error::invalid_input_source( + format!("Missing {} encoding location", M::NAME).into(), + )), } } fn meta_to_col_infos( column_metadatas: &[pbfile::ColumnMetadata], file_version: LanceFileVersion, - ) -> Vec> { + ) -> Result>> { column_metadatas .iter() .enumerate() .map(|(col_idx, col_meta)| { - Self::meta_to_col_info(col_idx as u32, col_meta, file_version) + let col_idx = u32::try_from(col_idx).map_err(|_| { + Error::invalid_input_source("File has more than u32::MAX columns".into()) + })?; + Self::meta_to_col_info(col_idx, col_meta, file_version) }) - .collect::>() + .collect() } fn meta_to_col_info( col_idx: u32, col_meta: &pbfile::ColumnMetadata, file_version: LanceFileVersion, - ) -> Arc { + ) -> Result> { let page_infos = col_meta .pages .iter() - .map(|page| { + .enumerate() + .map(|(page_idx, page)| { let num_rows = page.length; let encoding = match file_version { LanceFileVersion::V2_0 => { PageEncoding::Legacy(Self::fetch_encoding::( - page.encoding.as_ref().unwrap(), - )) + page.encoding.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!( + "Column {} page {} is missing its encoding", + col_idx, page_idx + ) + .into(), + ) + })?, + )?) + } + _ => { + PageEncoding::Structural(Self::fetch_encoding::( + page.encoding.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!( + "Column {} page {} is missing its encoding", + col_idx, page_idx + ) + .into(), + ) + })?, + )?) } - _ => PageEncoding::Structural(Self::fetch_encoding::( - page.encoding.as_ref().unwrap(), - )), }; + if page.buffer_offsets.len() != page.buffer_sizes.len() { + return Err(Error::invalid_input_source( + format!( + "Column {} page {} has {} buffer offsets but {} buffer sizes", + col_idx, + page_idx, + page.buffer_offsets.len(), + page.buffer_sizes.len() + ) + .into(), + )); + } let buffer_offsets_and_sizes = Arc::from( page.buffer_offsets .iter() .zip(page.buffer_sizes.iter()) - .map(|(offset, size)| { - // Starting with version 2.1 we can assert that page buffers are aligned - assert!( - file_version < LanceFileVersion::V2_1 - || offset % PAGE_BUFFER_ALIGNMENT as u64 == 0 - ); - (*offset, *size) + .map(|(offset, size)| -> Result<_> { + if file_version >= LanceFileVersion::V2_1 + && offset % PAGE_BUFFER_ALIGNMENT as u64 != 0 + { + return Err(Error::invalid_input_source( + format!( + "Column {} page {} buffer offset {} is not aligned to {} bytes", + col_idx, page_idx, offset, PAGE_BUFFER_ALIGNMENT + ) + .into(), + )); + } + Ok((*offset, *size)) }) - .collect::>(), + .collect::>>()?, ); - PageInfo { + Ok(PageInfo { buffer_offsets_and_sizes, encoding, num_rows, priority: page.priority, - } + }) }) - .collect::>(); + .collect::>>()?; + if col_meta.buffer_offsets.len() != col_meta.buffer_sizes.len() { + return Err(Error::invalid_input_source( + format!( + "Column {} has {} buffer offsets but {} buffer sizes", + col_idx, + col_meta.buffer_offsets.len(), + col_meta.buffer_sizes.len() + ) + .into(), + )); + } let buffer_offsets_and_sizes = Arc::from( col_meta .buffer_offsets @@ -1205,12 +1276,16 @@ impl FileReader { .map(|(offset, size)| (*offset, *size)) .collect::>(), ); - Arc::new(ColumnInfo { + Ok(Arc::new(ColumnInfo { index: col_idx, page_infos: Arc::from(page_infos), buffer_offsets_and_sizes, - encoding: Self::fetch_encoding(col_meta.encoding.as_ref().unwrap()), - }) + encoding: Self::fetch_encoding(col_meta.encoding.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Column {} is missing its encoding", col_idx).into(), + ) + })?)?, + })) } fn validate_projection( @@ -1946,7 +2021,7 @@ impl FileMetadataProvider { column_index, &column_metadata, metadata_index.version, - ); + )?; let cached = Arc::new(CachedColumnMetadata { column_metadata, column_info: column_info.clone(), @@ -2108,7 +2183,32 @@ impl FileReadCore { column_infos.len() )) })?; - Ok(info.page_infos.iter().map(|page| page.num_rows).sum()) + info.page_infos.iter().try_fold(0_u64, |rows, page| { + let page_rows = match &page.encoding { + PageEncoding::Structural(layout) => match &layout.layout { + Some(pbenc21::page_layout::Layout::SparseLayout(sparse)) => sparse + .structural_layers + .first() + .and_then(|layer| layer.layer.as_ref()) + .map_or(page.num_rows, |layer| match layer { + pbenc21::sparse_structural_layer::Layer::Validity(layer) => { + layer.num_slots + } + pbenc21::sparse_structural_layer::Layer::List(layer) => { + layer.num_slots + } + pbenc21::sparse_structural_layer::Layer::FixedSizeList(layer) => { + layer.num_slots + } + }), + _ => page.num_rows, + }, + _ => page.num_rows, + }; + rows.checked_add(page_rows).ok_or_else(|| { + Error::invalid_input_source("Column row count overflows u64".into()) + }) + }) }; let column_indices = &prepared.decoder_projection.column_indices; let fields = &prepared.decoder_projection.schema.fields; @@ -2512,7 +2612,7 @@ impl EncodedBatchReaderExt for EncodedBatch { footer.minor_version as u32, )?; - let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version); + let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version)?; Ok(Self { data: bytes, @@ -2561,7 +2661,7 @@ impl EncodedBatchReaderExt for EncodedBatch { let column_metadatas = FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?; - let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version); + let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version)?; Ok(Self { data: bytes, @@ -2585,9 +2685,10 @@ mod tests { }; use arrow_array::{ - RecordBatch, UInt32Array, + Int32Array, ListArray, RecordBatch, RecordBatchIterator, UInt32Array, types::{Float64Type, Int32Type}, }; + use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema}; use bytes::Bytes; use futures::{StreamExt, prelude::stream::TryStreamExt}; @@ -2595,10 +2696,13 @@ mod tests { use lance_core::{ArrowResult, datatypes::Schema}; use lance_datagen::{BatchCount, ByteCount, RowCount, array, gen_batch}; use lance_encoding::{ + constants::{STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_SPARSE}, decoder::{ - DecodeBatchScheduler, DecoderPlugins, FilterExpression, ReadBatchTask, decode_batch, + DecodeBatchScheduler, DecoderPlugins, FilterExpression, PageEncoding, ReadBatchTask, + decode_batch, }, encoder::{EncodedBatch, EncodingOptions, default_encoding_strategy, encode_batch}, + format::pb21, version::LanceFileVersion, }; use lance_io::{stream::RecordBatchStream, utils::CachedFileSize}; @@ -2614,6 +2718,135 @@ mod tests { use crate::writer::{EncodedBatchWriteExt, FileWriter, FileWriterOptions}; use lance_encoding::decoder::DecoderConfig; + #[tokio::test] + async fn sparse_file_writer_reader_scan_range_and_take_roundtrip() { + let fs = FsFixture::default(); + let sparse_metadata = HashMap::from([( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_SPARSE.to_string(), + )]); + let value_field = + Field::new("values", DataType::Int32, true).with_metadata(sparse_metadata.clone()); + let item_field = Arc::new(Field::new("item", DataType::Int32, true)); + let list_field = Field::new("items", DataType::List(item_field.clone()), true) + .with_metadata(sparse_metadata); + let arrow_schema = Arc::new(ArrowSchema::new(vec![value_field, list_field])); + let list = ListArray::try_new( + item_field, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 2, 2, 3, 3, 5])), + Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + Some(4), + Some(5), + ])), + Some(NullBuffer::from(vec![true, false, true, true, true, true])), + ) + .unwrap(); + let batch = RecordBatch::try_new( + arrow_schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![ + Some(10), + None, + Some(30), + Some(40), + None, + Some(60), + ])), + Arc::new(list), + ], + ) + .unwrap(); + let input = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema); + write_lance_file( + input, + &fs, + FileWriterOptions { + format_version: Some(LanceFileVersion::V2_3), + ..Default::default() + }, + ) + .await; + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let file_reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &test_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + assert_eq!(file_reader.metadata.column_infos.len(), 2); + assert!( + file_reader + .metadata + .column_infos + .iter() + .flat_map(|column| column.page_infos.iter()) + .all(|page| { + matches!( + &page.encoding, + PageEncoding::Structural(layout) + if matches!( + layout.layout, + Some(pb21::page_layout::Layout::SparseLayout(_)) + ) + ) + }) + ); + + let scan = file_reader + .read_stream( + lance_io::ReadBatchParams::RangeFull, + 1024, + 1, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(scan, vec![batch.clone()]); + + let range = file_reader + .read_stream( + lance_io::ReadBatchParams::Range(1..5), + 1024, + 1, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(range, vec![batch.slice(1, 4)]); + + let indices = UInt32Array::from(vec![0, 3, 5]); + let take = file_reader + .read_stream( + lance_io::ReadBatchParams::Indices(indices.clone()), + 1024, + 1, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(take, vec![batch.take(&indices).unwrap()]); + } + async fn create_some_file(fs: &FsFixture, version: LanceFileVersion) -> WrittenFile { let location_type = DataType::Struct(Fields::from(vec![ Field::new("x", DataType::Float64, true), diff --git a/rust/lance/benches/s3_file_reader_diagnostics.rs b/rust/lance/benches/s3_file_reader_diagnostics.rs index 5c661f447e3..1762048d495 100644 --- a/rust/lance/benches/s3_file_reader_diagnostics.rs +++ b/rust/lance/benches/s3_file_reader_diagnostics.rs @@ -580,6 +580,7 @@ fn page_layout_kind(encoding: &PageEncoding) -> &'static str { Some(pb21::page_layout::Layout::ConstantLayout(_)) => "constant", Some(pb21::page_layout::Layout::FullZipLayout(_)) => "fullzip", Some(pb21::page_layout::Layout::BlobLayout(_)) => "blob", + Some(pb21::page_layout::Layout::SparseLayout(_)) => "sparse", None => "missing", }, }