diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index 4a908e87ed9..8c775b95bb1 100644 --- a/docs/src/format/file/encoding.md +++ b/docs/src/format/file/encoding.md @@ -353,6 +353,24 @@ buffer. Every layer has a `SparseValiditySet` whose meaning is explicit: The unspecified validity meaning is invalid. Both polarities are part of the wire contract and have identical Arrow semantics after normalization. +#### Writer Selection + +Writers may emit this layout only for Lance 2.3+ fields that explicitly set +`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. + +Sparse selection is opt-in. 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. They should choose the +validity polarity with the lower semantic encoded cost; ties use 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: diff --git a/rust/lance-encoding/src/constants.rs b/rust/lance-encoding/src/constants.rs index c95b587a532..0bd31d676cc 100644 --- a/rust/lance-encoding/src/constants.rs +++ b/rust/lance-encoding/src/constants.rs @@ -55,6 +55,8 @@ pub const STRUCTURAL_ENCODING_META_KEY: &str = "lance-encoding:structural-encodi pub const STRUCTURAL_ENCODING_MINIBLOCK: &str = "miniblock"; /// Value for fullzip structural encoding pub const STRUCTURAL_ENCODING_FULLZIP: &str = "fullzip"; +/// Value for sparse structural encoding +pub const STRUCTURAL_ENCODING_SPARSE: &str = "sparse"; // Byte stream split metadata keys /// Metadata key for byte stream split encoding configuration diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index d55f1c576b3..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}, @@ -3807,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. @@ -3853,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, @@ -5322,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), }]); } @@ -5348,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) @@ -5370,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. @@ -5683,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/sparse.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs index 17337c616bf..ed39325961c 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs @@ -5,6 +5,8 @@ use super::*; use arrow_array::new_empty_array; use arrow_buffer::ArrowNativeType; +pub(super) mod writer; + fn usize_from_u64(value: u64, label: &str) -> Result { usize::try_from(value).map_err(|_| { Error::invalid_input_source( 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..a1eeace6574 --- /dev/null +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -0,0 +1,1606 @@ +// 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.clone(), 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 { + kind: pb21::sparse_structural_layer::Kind::SparseLayerValidity as i32, + num_slots: *num_slots, + num_child_slots: *num_slots, + non_empty_positions: None, + counts: None, + fixed_size_list_dimension: 0, + 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 { + kind: pb21::sparse_structural_layer::Kind::SparseLayerList as i32, + num_slots: *num_slots, + num_child_slots: *num_child_slots, + non_empty_positions: Some(non_empty_positions), + counts: Some(counts), + fixed_size_list_dimension: 0, + 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); + 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={num_slots}, dimension={dimension}" + ) + .into(), + ) + })?; + layers.push(pb21::SparseStructuralLayer { + kind: pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList as i32, + num_slots: *num_slots, + num_child_slots, + non_empty_positions: None, + counts: None, + fixed_size_list_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 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::SparseStructuralLayer { + sparse + .structural_layers + .iter() + .find(|layer| layer.kind == pb21::sparse_structural_layer::Kind::SparseLayerList as i32) + .expect("expected sparse list layer") + } + + 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 == &vec![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 = sparse.structural_layers[0].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| { + layer.kind == pb21::sparse_structural_layer::Kind::SparseLayerValidity as i32 + })); + let struct_validity = sparse.structural_layers[0].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_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, LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + let validity = sparse_layout(&null_positions[0]).structural_layers[0] + .validity + .as_ref() + .unwrap(); + assert_eq!( + validity.meaning, + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions as i32 + ); + + let valid_positions = encode_pages(mostly_null, LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + let validity = sparse_layout(&valid_positions[0]).structural_layers[0] + .validity + .as_ref() + .unwrap(); + assert_eq!( + validity.meaning, + pb21::sparse_validity_set::Meaning::SparseValidityValidPositions as i32 + ); + } + + #[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| { + layer.kind == pb21::sparse_structural_layer::Kind::SparseLayerList as i32 + }) + .count() + >= 2 + })); + assert!(sparse.iter().any(|layout| { + layout.structural_layers.iter().any(|layer| { + layer.kind == pb21::sparse_structural_layer::Kind::SparseLayerValidity 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_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| { + layer.kind == pb21::sparse_structural_layer::Kind::SparseLayerList as i32 + }) + })); + 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| { + layer.kind == pb21::sparse_structural_layer::Kind::SparseLayerList as i32 + }) + })); + 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(); + assert!(fsl_pages.iter().map(sparse_layout).any(|layout| { + layout.structural_layers.iter().any(|layer| { + layer.kind == pb21::sparse_structural_layer::Kind::SparseLayerFixedSizeList as i32 + && layer.fixed_size_list_dimension == 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_serializes_semantic_list_forms() { + let all = encode_pages( + list_i32(vec![0, 2, 4, 6], None), + 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 = encode_pages( + list_i32(vec![0, 2, 4, 4, 4], None), + 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 = encode_pages( + list_i32(vec![0, 1, 1, 4, 4, 6], None), + 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); + } + + #[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/repdef.rs b/rust/lance-encoding/src/repdef.rs index 0ac9f07f558..2bff898a273 100644 --- a/rust/lance-encoding/src/repdef.rs +++ b/rust/lance-encoding/src/repdef.rs @@ -202,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 { @@ -1412,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() - } - - /// 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, - ) + Self::normalize(builders).serialize() } - 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]), @@ -1467,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) } }