diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index 8c775b95bb1..e2cde116d80 100644 --- a/docs/src/format/file/encoding.md +++ b/docs/src/format/file/encoding.md @@ -355,15 +355,18 @@ 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. +Writers may emit this layout only for Lance 2.3+ fields. A field can request it explicitly with +`lance-encoding:structural-encoding=sparse`; the same request is an input error for earlier file versions. Without an +explicit structural encoding, the Lance 2.3 writer selects sparse only when the dense mini-block repetition/definition +budget would split the page or one top-level row exceeds that budget, and only when the value path is supported by the +sparse writer. Explicit `miniblock`, `fullzip`, and `sparse` requests are not changed by this automatic policy. Lance +2.2 and earlier writers never select sparse. + +Unsupported sparse value paths, including dictionary values and variable-width packed structs, retain their dense +behavior. Writers normalize Arrow validity and list structure once. Within-budget dense pages do not build sparse +position/count plans. When sparse is selected, writers choose the validity polarity with the lower semantic encoded +cost; ties use null positions. Field 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. 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 @@ -643,7 +646,7 @@ options. However, they can also be set in the field metadata in the schema. | `lance-encoding:dict-values-compression-level` | Integers (scheme dependent) | Varies by scheme | Compression level for dictionary values general compression | | `lance-encoding:general` | `off`, `on` | `off` | Whether to apply general compression. | | `lance-encoding:packed` | Any string | Not set | Whether to apply packed struct encoding (see above). | -| `lance-encoding:structural-encoding` | `miniblock`, `fullzip` | Not set | Force a particular structural encoding to be applied (only useful for testing purposes) | +| `lance-encoding:structural-encoding` | `miniblock`, `fullzip`, `sparse` | Not set | Force a particular structural encoding to be applied (only useful for testing purposes) | ### Configuration Details diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 5429067efe5..90e99bab69b 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -92,6 +92,7 @@ pub mod blob; pub mod constant; pub mod dict; pub mod fullzip; +mod layout; pub mod miniblock; pub(crate) mod sparse; @@ -3813,7 +3814,10 @@ enum PrimitivePageStructure { repdef: SerializedRepDefs, unsplittable_miniblock_levels: Option, }, - Sparse(sparse::SparseStructuralPlan), + Sparse { + plan: sparse::SparseStructuralPlan, + prepared_values: Option, + }, } // A primitive page after optional structural splitting. @@ -5407,7 +5411,10 @@ impl PrimitiveStructuralEncoder { repdef, unsplittable_miniblock_levels, } => (repdef, unsplittable_miniblock_levels), - PrimitivePageStructure::Sparse(plan) => { + PrimitivePageStructure::Sparse { + plan, + prepared_values, + } => { log::debug!( "Encoding column {} with {} visible items ({} rows) using sparse layout", column_idx, @@ -5418,7 +5425,14 @@ impl PrimitiveStructuralEncoder { column_idx, &field, compression_strategy.as_ref(), - DataBlock::from_arrays(&arrays, num_values), + prepared_values.map_or_else( + || { + sparse::writer::SparseValueInput::Unprepared(DataBlock::from_arrays( + &arrays, num_values, + )) + }, + sparse::writer::SparseValueInput::Prepared, + ), plan, row_number, num_rows, @@ -5734,35 +5748,78 @@ impl PrimitiveStructuralEncoder { let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity()); let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty()); let normalized = RepDefBuilder::normalize(repdefs); - let requests_sparse = self - .encoding_metadata - .get(STRUCTURAL_ENCODING_META_KEY) + let requested_encoding = self.encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY); + let requests_sparse = requested_encoding .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 { + let pages = if let Some(plan) = sparse_plan + && !sparse::writer::uses_constant_layout(&plan, &self.field) + { + vec![PrimitivePageData { + arrays, + structure: PrimitivePageStructure::Sparse { + plan, + prepared_values: None, + }, + row_number, + num_rows, + }] + } else { + let (repdef, structural_plan) = normalized.serialize_with_structural_plan( + miniblock::max_repdef_levels_per_chunk, + num_rows, + num_values, + )?; + let automatic_sparse = layout::select_automatic_sparse( + self.version.resolve(), + requested_encoding.map(String::as_str), + &structural_plan, + || { + let data = DataBlock::from_arrays(&arrays, num_values); + if !sparse::writer::supports_value_block(&data) { + return Ok(None); + } + let prepared_values = match sparse::writer::prepare_values( + &self.field, + self.compression_strategy.as_ref(), + data, + self.support_large_chunk, + ) { + Ok(prepared_values) => prepared_values, + Err(error) => { + trace!( + "Keeping column {} on its dense structural path because sparse value preparation is unavailable: {}", + self.column_index, error + ); + return Ok(None); + } + }; + let plan = sparse::writer::plan(&normalized, num_values)?; + if sparse::writer::uses_constant_layout(&plan, &self.field) { + return Ok(None); + } + Ok(Some((plan, prepared_values))) + }, + )?; + match automatic_sparse { + Some((plan, prepared_values)) => vec![PrimitivePageData { arrays, - structure: PrimitivePageStructure::Sparse(plan), + structure: PrimitivePageStructure::Sparse { + plan, + prepared_values: Some(prepared_values), + }, 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( + }], + None => Self::split_structural_pages_for_miniblock_budget( arrays, repdef, structural_plan, row_number, num_rows, - )? + )?, } }; diff --git a/rust/lance-encoding/src/encodings/logical/primitive/layout.rs b/rust/lance-encoding/src/encodings/logical/primitive/layout.rs new file mode 100644 index 00000000000..7135e1e527d --- /dev/null +++ b/rust/lance-encoding/src/encodings/logical/primitive/layout.rs @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use lance_core::Result; + +use crate::{repdef::StructuralPagePlan, version::LanceFileVersion}; + +/// Runs automatic sparse planning only after the dense mini-block budget makes it useful. +pub(super) fn select_automatic_sparse( + version: LanceFileVersion, + requested_encoding: Option<&str>, + dense_plan: &StructuralPagePlan, + candidate: impl FnOnce() -> Result>, +) -> Result> { + if version < LanceFileVersion::V2_3 + || requested_encoding.is_some() + || !matches!( + dense_plan, + StructuralPagePlan::Split(_) | StructuralPagePlan::UnsplittableOverBudget(_) + ) + { + return Ok(None); + } + candidate() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::constants::{ + STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_SPARSE, + }; + + #[test] + fn within_budget_does_not_construct_sparse_candidate() { + let selected = select_automatic_sparse::<()>( + LanceFileVersion::V2_3, + None, + &StructuralPagePlan::Fits, + || panic!("within-budget pages must not construct sparse candidates"), + ) + .unwrap(); + assert!(selected.is_none()); + } + + #[test] + fn over_budget_selects_only_eligible_candidates() { + let split = StructuralPagePlan::Split(Vec::new()); + let selected = + select_automatic_sparse(LanceFileVersion::V2_3, None, &split, || Ok(Some(42))).unwrap(); + assert_eq!(selected, Some(42)); + + let ineligible = + select_automatic_sparse::<()>(LanceFileVersion::V2_3, None, &split, || Ok(None)) + .unwrap(); + assert!(ineligible.is_none()); + + let unsplittable = StructuralPagePlan::UnsplittableOverBudget(70_000); + let selected = + select_automatic_sparse(LanceFileVersion::V2_3, None, &unsplittable, || Ok(Some(7))) + .unwrap(); + assert_eq!(selected, Some(7)); + } + + #[test] + fn explicit_modes_and_lance_2_2_do_not_auto_select() { + let split = StructuralPagePlan::Split(Vec::new()); + for requested in [ + STRUCTURAL_ENCODING_MINIBLOCK, + STRUCTURAL_ENCODING_FULLZIP, + STRUCTURAL_ENCODING_SPARSE, + ] { + let selected = select_automatic_sparse::<()>( + LanceFileVersion::V2_3, + Some(requested), + &split, + || panic!("explicit modes must not invoke automatic sparse planning"), + ) + .unwrap(); + assert!(selected.is_none()); + } + + let selected = select_automatic_sparse::<()>(LanceFileVersion::V2_2, None, &split, || { + panic!("Lance 2.2 must not invoke automatic sparse planning") + }) + .unwrap(); + assert!(selected.is_none()); + } +} diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs index a1eeace6574..971e6e99948 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -346,6 +346,30 @@ pub(in crate::encodings::logical::primitive) fn uses_constant_layout( } } +fn supports_fixed_size_list_values(data: &DataBlock) -> bool { + match data { + DataBlock::FixedWidth(_) => true, + DataBlock::FixedSizeList(list) => supports_fixed_size_list_values(list.child.as_ref()), + DataBlock::Nullable(nullable) => supports_fixed_size_list_values(nullable.data.as_ref()), + _ => false, + } +} + +/// Whether the sparse writer can encode this value block without changing its value path. +pub(in crate::encodings::logical::primitive) fn supports_value_block(data: &DataBlock) -> bool { + match data { + DataBlock::FixedWidth(_) | DataBlock::VariableWidth(_) => true, + DataBlock::Struct(data) => !data.has_variable_width_child(), + DataBlock::FixedSizeList(data) => supports_fixed_size_list_values(data.child.as_ref()), + DataBlock::Empty() + | DataBlock::Constant(_) + | DataBlock::AllNull(_) + | DataBlock::Nullable(_) + | DataBlock::Opaque(_) + | DataBlock::Dictionary(_) => false, + } +} + struct SparseMiniBlockChunk { buffer_sizes: Vec, num_values: u32, @@ -362,6 +386,17 @@ struct SerializedValuePage { metadata: LanceBuffer, } +pub(in crate::encodings::logical::primitive) struct PreparedSparseValues { + num_values: u64, + value_compression: CompressiveEncoding, + values: SerializedValuePage, +} + +pub(in crate::encodings::logical::primitive) enum SparseValueInput { + Unprepared(DataBlock), + Prepared(PreparedSparseValues), +} + struct EncodedStructuralPlan { layers: Vec, buffers: Vec, @@ -501,6 +536,43 @@ fn serialize_value_chunks( }) } +pub(in crate::encodings::logical::primitive) fn prepare_values( + field: &Field, + compression_strategy: &dyn CompressionStrategy, + data: DataBlock, + support_large_chunk: bool, +) -> Result { + 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 num_values = data.num_values(); + 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)?; + Ok(PreparedSparseValues { + num_values, + value_compression, + values, + }) +} + fn encode_u64_values( values: Vec, compression_strategy: &dyn CompressionStrategy, @@ -748,42 +820,28 @@ pub(in crate::encodings::logical::primitive) fn encode_page( column_idx: u32, field: &Field, compression_strategy: &dyn CompressionStrategy, - data: DataBlock, + values: SparseValueInput, plan: SparseStructuralPlan, row_number: u64, num_rows: u64, support_large_chunk: bool, ) -> Result { - if plan.num_visible_items != data.num_values() { + let PreparedSparseValues { + num_values, + value_compression, + values, + } = match values { + SparseValueInput::Unprepared(data) => { + prepare_values(field, compression_strategy, data, support_large_chunk)? + } + SparseValueInput::Prepared(prepared) => prepared, + }; + if plan.num_visible_items != num_values { return Err(Error::internal(format!( "Sparse structural plan has {} visible items but data has {} values", - plan.num_visible_items, - data.num_values() + plan.num_visible_items, 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( @@ -816,17 +874,20 @@ mod tests { use std::{collections::HashMap, sync::Arc}; use arrow_array::{ - Array, ArrayRef, FixedSizeListArray, Int32Array, LargeListArray, ListArray, StructArray, + Array, ArrayRef, DictionaryArray, FixedSizeBinaryArray, FixedSizeListArray, Int8Array, + Int32Array, LargeListArray, ListArray, StringArray, StructArray, builder::{Int32Builder, MapBuilder, StringBuilder}, + types::Int8Type, }; 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, + PACKED_STRUCT_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_SPARSE, }, + data::FixedSizeListBlock, encoder::{ ColumnIndexSequence, EncodingOptions, FieldEncoder, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers, default_encoding_strategy, @@ -866,6 +927,101 @@ mod tests { ) } + fn sparse_list_values( + num_rows: usize, + stride: usize, + values: ArrayRef, + item_field: Arc, + ) -> ArrayRef { + let mut offsets = Vec::with_capacity(num_rows + 1); + let mut num_values = 0_i32; + offsets.push(num_values); + for row in 0..num_rows { + if (row + 1).is_multiple_of(stride) || row + 1 == num_rows { + num_values += 1; + } + offsets.push(num_values); + } + assert_eq!(values.len(), num_values as usize); + Arc::new( + ListArray::try_new( + item_field, + OffsetBuffer::new(ScalarBuffer::from(offsets)), + values, + None, + ) + .unwrap(), + ) + } + + fn sparse_i32_list(num_rows: usize, stride: usize) -> ArrayRef { + let num_values = num_rows.div_ceil(stride); + sparse_list_values( + num_rows, + stride, + Arc::new(Int32Array::from_iter_values(0..num_values as i32)), + Arc::new(ArrowField::new("item", DataType::Int32, true)), + ) + } + + fn unsplittable_nested_list(values: ArrayRef, item_field: Arc) -> ArrayRef { + const NUM_INNER_LISTS: usize = 70_000; + + let mut inner_offsets = vec![0_i32; NUM_INNER_LISTS + 1]; + assert!(!values.is_empty()); + assert!(values.len() <= NUM_INNER_LISTS); + for value_index in 1..=values.len() { + inner_offsets[NUM_INNER_LISTS - values.len() + value_index] = value_index as i32; + } + let inner = Arc::new( + ListArray::try_new( + item_field, + OffsetBuffer::new(ScalarBuffer::from(inner_offsets)), + values, + None, + ) + .unwrap(), + ) as ArrayRef; + Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new("item", inner.data_type().clone(), true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, NUM_INNER_LISTS as i32])), + inner, + None, + ) + .unwrap(), + ) + } + + fn variable_packed_struct_values(num_values: usize) -> (ArrayRef, Arc) { + let fields = Fields::from(vec![ArrowField::new("value", DataType::Utf8, false)]); + let values = Arc::new(StructArray::new( + fields.clone(), + vec![Arc::new(StringArray::from_iter_values( + (0..num_values).map(|index| format!("value-{index}")), + ))], + None, + )) as ArrayRef; + let item_field = Arc::new( + ArrowField::new("item", DataType::Struct(fields), true).with_metadata(HashMap::from([ + (PACKED_STRUCT_META_KEY.to_string(), "true".to_string()), + ])), + ); + (values, item_field) + } + + fn dictionary_values(num_values: usize) -> (ArrayRef, Arc) { + let keys = Int8Array::from_iter_values((0..num_values).map(|index| (index % 2) as i8)); + let values = Arc::new(StringArray::from(vec!["value-0", "value-1"])); + let dictionary = Arc::new(DictionaryArray::::try_new(keys, values).unwrap()); + let item_field = Arc::new(ArrowField::new( + "item", + dictionary.data_type().clone(), + true, + )); + (dictionary, item_field) + } + 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; @@ -1074,6 +1230,64 @@ mod tests { .expect("expected sparse list layer") } + #[test] + fn test_sparse_value_block_eligibility() { + let fixed = DataBlock::from_array(Int32Array::from_iter_values(0..4)); + assert!(supports_value_block(&fixed)); + + let variable = DataBlock::from_array(StringArray::from(vec!["a", "b"])); + assert!(supports_value_block(&variable)); + + let nullable = DataBlock::from_array(Int32Array::from(vec![Some(1), None])); + assert!(!supports_value_block(&nullable)); + + let all_null = DataBlock::from_array(Int32Array::from(vec![None, None])); + assert!(!supports_value_block(&all_null)); + + let (dictionary, _) = dictionary_values(4); + assert!(!supports_value_block(&DataBlock::from_arrays( + &[dictionary], + 4 + ))); + + let fixed_fields = Fields::from(vec![ArrowField::new("value", DataType::Int32, false)]); + let fixed_struct = StructArray::new( + fixed_fields, + vec![Arc::new(Int32Array::from_iter_values(0..4))], + None, + ); + assert!(supports_value_block(&DataBlock::from_array(fixed_struct))); + + let variable_fields = Fields::from(vec![ArrowField::new("value", DataType::Utf8, false)]); + let variable_struct = StructArray::new( + variable_fields, + vec![Arc::new(StringArray::from(vec!["a", "b"]))], + None, + ); + assert!(!supports_value_block(&DataBlock::from_array( + variable_struct + ))); + + let fixed_size_list = FixedSizeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, false)), + 2, + Arc::new(Int32Array::from_iter_values(0..4)), + None, + ) + .unwrap(); + assert!(supports_value_block(&DataBlock::from_array( + fixed_size_list + ))); + + let unsupported_fixed_size_list = DataBlock::FixedSizeList(FixedSizeListBlock { + child: Box::new(DataBlock::from_array(StringArray::from(vec![ + "a", "b", "c", "d", + ]))), + dimension: 2, + }); + assert!(!supports_value_block(&unsupported_fixed_size_list)); + } + fn planned_list( offsets: Vec, list_validity: Option>, @@ -1579,6 +1793,325 @@ mod tests { )); } + #[tokio::test] + async fn test_auto_sparse_for_split_required_page() { + const NUM_ROWS: usize = 70_000; + let array = sparse_i32_list(NUM_ROWS, 2_000); + + let within_budget = encode_pages( + sparse_i32_list(4_096, 1_024), + LanceFileVersion::V2_3, + HashMap::new(), + ) + .await + .unwrap(); + assert_eq!(within_budget.len(), 1); + assert!(matches!( + page_layout(&within_budget[0]), + pb21::page_layout::Layout::MiniBlockLayout(_) + )); + + let automatic = encode_pages(array.clone(), LanceFileVersion::V2_3, HashMap::new()) + .await + .unwrap(); + assert_eq!(automatic.len(), 1); + assert!(matches!( + page_layout(&automatic[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + + let miniblock = encode_pages( + array.clone(), + LanceFileVersion::V2_3, + structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), + ) + .await + .unwrap(); + assert!(miniblock.len() > 1); + assert!(miniblock.iter().all(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::MiniBlockLayout(_) + ))); + + let fullzip = encode_pages( + array.clone(), + LanceFileVersion::V2_3, + structural_metadata(STRUCTURAL_ENCODING_FULLZIP), + ) + .await + .unwrap(); + assert_eq!(fullzip.len(), miniblock.len()); + assert!(fullzip.iter().all(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::FullZipLayout(_) + ))); + + let sparse = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert_eq!(sparse.len(), 1); + assert!(matches!( + page_layout(&sparse[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + + 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_eq!(page_layout(default), page_layout(explicit)); + assert_eq!(default.data, explicit.data); + } + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1024 * 1024]) + .with_batch_size(NUM_ROWS as u32) + .with_range(1_999..2_002) + .with_indices(vec![0, 1_999, 2_000, NUM_ROWS as u64 - 1]); + check_round_trip_encoding_of_data(vec![array], &cases, HashMap::new()).await; + } + + #[tokio::test] + async fn test_auto_sparse_for_single_row_over_budget() { + let array = unsplittable_nested_list( + Arc::new(Int32Array::from(vec![42, 43])), + Arc::new(ArrowField::new("item", DataType::Int32, true)), + ); + + let automatic = encode_pages(array.clone(), LanceFileVersion::V2_3, HashMap::new()) + .await + .unwrap(); + assert_eq!(automatic.len(), 1); + assert!(matches!( + page_layout(&automatic[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + + let v2_2 = encode_pages(array.clone(), LanceFileVersion::V2_2, HashMap::new()) + .await + .unwrap(); + assert_eq!(v2_2.len(), 1); + assert!(matches!( + page_layout(&v2_2[0]), + pb21::page_layout::Layout::FullZipLayout(_) + )); + + let miniblock_error = encode_pages( + array.clone(), + LanceFileVersion::V2_3, + structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), + ) + .await + .unwrap_err(); + assert!( + miniblock_error + .to_string() + .contains("Mini-block cannot encode 70000 rep/def levels") + ); + + let fullzip = encode_pages( + array.clone(), + LanceFileVersion::V2_3, + structural_metadata(STRUCTURAL_ENCODING_FULLZIP), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&fullzip[0]), + pb21::page_layout::Layout::FullZipLayout(_) + )); + + let explicit_sparse = + encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap(); + assert!(matches!( + page_layout(&explicit_sparse[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_range(0..1) + .with_indices(vec![0]); + check_round_trip_encoding_of_data(vec![array], &cases, HashMap::new()).await; + } + + #[tokio::test] + async fn test_auto_sparse_keeps_unsupported_values_dense() { + const NUM_ROWS: usize = 70_000; + let num_values = NUM_ROWS.div_ceil(2_000); + + let (dictionary, dictionary_field) = dictionary_values(num_values); + let dictionary_array = sparse_list_values(NUM_ROWS, 2_000, dictionary, dictionary_field); + let dictionary_pages = encode_pages( + dictionary_array.clone(), + LanceFileVersion::V2_3, + HashMap::new(), + ) + .await + .unwrap(); + assert!(dictionary_pages.len() > 1); + assert!(dictionary_pages.iter().all(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::MiniBlockLayout(_) + ))); + + let (packed_values, packed_field) = variable_packed_struct_values(num_values); + let packed_array = sparse_list_values(NUM_ROWS, 2_000, packed_values, packed_field); + let packed_pages = + encode_pages(packed_array.clone(), LanceFileVersion::V2_3, HashMap::new()) + .await + .unwrap(); + assert!(packed_pages.len() > 1); + assert!(packed_pages.iter().all(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::FullZipLayout(_) + ))); + + let (packed_values, packed_field) = variable_packed_struct_values(2); + let unsplittable_packed = unsplittable_nested_list(packed_values, packed_field); + let packed_fallback = encode_pages( + unsplittable_packed.clone(), + LanceFileVersion::V2_3, + HashMap::new(), + ) + .await + .unwrap(); + assert_eq!(packed_fallback.len(), 1); + assert!(matches!( + page_layout(&packed_fallback[0]), + pb21::page_layout::Layout::FullZipLayout(_) + )); + + let (dictionary, dictionary_field) = dictionary_values(2); + let unsplittable_dictionary = unsplittable_nested_list(dictionary, dictionary_field); + let v2_2_error = encode_pages( + unsplittable_dictionary.clone(), + LanceFileVersion::V2_2, + HashMap::new(), + ) + .await + .unwrap_err(); + let v2_3_error = encode_pages( + unsplittable_dictionary, + LanceFileVersion::V2_3, + HashMap::new(), + ) + .await + .unwrap_err(); + assert_eq!(v2_3_error.to_string(), v2_2_error.to_string()); + assert!( + v2_3_error + .to_string() + .contains("Mini-block cannot encode 70000 rep/def levels") + ); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1024 * 1024]) + .with_batch_size(NUM_ROWS as u32) + .with_range(1_999..2_002) + .with_indices(vec![0, 2_000, NUM_ROWS as u64 - 1]); + check_round_trip_encoding_of_data(vec![dictionary_array], &cases, HashMap::new()).await; + let packed_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1024 * 1024]) + .with_batch_size(NUM_ROWS as u32); + check_round_trip_encoding_of_data(vec![packed_array], &packed_cases, HashMap::new()).await; + + let single_row_cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1024 * 1024]) + .with_range(0..1) + .with_indices(vec![0]); + check_round_trip_encoding_of_data( + vec![unsplittable_packed], + &single_row_cases, + HashMap::new(), + ) + .await; + } + + #[tokio::test] + async fn test_auto_sparse_wide_values_keep_dense_fallback() { + let first = vec![0xAB_u8; 5_000]; + let second = vec![0xCD_u8; 5_000]; + let fixed_size_binary = Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + [Some(first.as_slice()), Some(second.as_slice())].into_iter(), + 5_000, + ) + .unwrap(), + ) as ArrayRef; + + const FSL_DIMENSION: i32 = 2_048; + let fixed_size_list = Arc::new( + FixedSizeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + FSL_DIMENSION, + Arc::new(Int32Array::from_iter_values(0..(FSL_DIMENSION * 2))), + None, + ) + .unwrap(), + ) as ArrayRef; + + for (label, values) in [ + ("fixed-size binary", fixed_size_binary), + ("fixed-size list", fixed_size_list), + ] { + let item_field = Arc::new(ArrowField::new("item", values.data_type().clone(), true)); + let array = unsplittable_nested_list(values, item_field); + + let v2_2 = encode_pages(array.clone(), LanceFileVersion::V2_2, HashMap::new()) + .await + .unwrap(); + let v2_3 = encode_pages(array.clone(), LanceFileVersion::V2_3, HashMap::new()) + .await + .unwrap(); + for pages in [&v2_2, &v2_3] { + assert_eq!(pages.len(), 1, "unexpected {label} page count"); + assert!( + matches!( + page_layout(&pages[0]), + pb21::page_layout::Layout::FullZipLayout(_) + ), + "{label} should retain the dense full-zip fallback" + ); + } + + let explicit_error = + encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) + .await + .unwrap_err(); + assert!( + explicit_error.to_string().contains("too wide"), + "explicit sparse should preserve the {label} value error: {explicit_error}" + ); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_2) + .with_max_file_version(LanceFileVersion::V2_3) + .with_range(0..1) + .with_indices(vec![0]); + check_round_trip_encoding_of_data(vec![array], &cases, HashMap::new()).await; + } + } + #[test] fn test_explicit_sparse_rejects_lance_2_2() { let array = Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef; diff --git a/rust/lance-encoding/src/repdef.rs b/rust/lance-encoding/src/repdef.rs index 2bff898a273..afd51b41683 100644 --- a/rust/lance-encoding/src/repdef.rs +++ b/rust/lance-encoding/src/repdef.rs @@ -263,7 +263,7 @@ impl NormalizedStructuralPlan { }) } - fn into_serializer(self) -> (SerializerContext, Option) { + fn serializer(&self) -> (SerializerContext, Option) { if self.dense_all_valid { let def_meaning = self .layers @@ -310,27 +310,27 @@ impl NormalizedStructuralPlan { let num_layers = self.layers.len(); let mut context = SerializerContext::new(total_len, num_layers, max_rep, max_def); - for layer in self.layers { + 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), + 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(&self) -> SerializedRepDefs { + self.serializer().0.build() } pub(crate) fn serialize_with_structural_plan( - self, + &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(); + let (context, bits_per_level) = self.serializer(); context.build_with_structural_plan( bits_per_level.map(max_levels_for_bits), num_rows,