diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index d10137309c4..f1d14a3123b 100644 --- a/docs/src/format/file/encoding.md +++ b/docs/src/format/file/encoding.md @@ -587,6 +587,7 @@ on a per-value basis. We use ☑️ to mark a technique that is applied on a per | Constant | ✅ (2.1) | ❓ | ❓ | | Range | ✅ (2.3) | ❌ | ❓ | | Delta | ✅ (2.3) | ❌ | ❓ | +| Dictionary | ✅ (2.3) | ❌ | ❓ | | Bitpacking | ✅ (2.1) | ❓ | ✅ (2.1) | | Fsst | ❓ | ✅ (2.1) | ✅ (2.1) | | Rle | ✅ (2.2) | ❌ | ✅ (2.1) | @@ -616,7 +617,19 @@ with checked prefix sums. %%% proto.message.Delta %%% ``` -Lance 2.0 through 2.2 writers do not emit either encoding. +Block Dictionary stores `u32` indices and typed dictionary items as child codec trees. When at least one child +has a payload, both children are combined into one outer payload: + +```text +Dictionary payload: + u64 indices_payload_bytes + u64 items_payload_bytes + indices payload + dictionary-items payload +``` + +The framed length for a metadata-only child must be zero. Readers validate the item count, frame boundaries, +child cardinalities, and every index. Lance 2.0 through 2.2 writers do not emit Range, Delta, or block Dictionary. ### Flat diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index ef832726ae8..133aa3f5df8 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -1209,7 +1209,7 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { Compression::OutOfLineBitpacking(_) => Err(Error::not_supported_source( "this runtime was not built with bitpacking support".into(), )), - Compression::Range(_) | Compression::Delta(_) => { + Compression::Dictionary(_) | Compression::Range(_) | Compression::Delta(_) => { let value_type = block::infer_block_value_type(description)?; block::create_block_decompressor(description, value_type) .map(|(decompressor, _)| decompressor) diff --git a/rust/lance-encoding/src/compression/block.rs b/rust/lance-encoding/src/compression/block.rs index 792ca41937c..f2029c5f86a 100644 --- a/rust/lance-encoding/src/compression/block.rs +++ b/rust/lance-encoding/src/compression/block.rs @@ -8,6 +8,7 @@ use lance_core::{Error, Result}; +pub(crate) const MAX_DICTIONARY_ITEMS: usize = 4096; #[cfg(feature = "bitpacking")] pub(crate) const BITPACK_CHUNK_VALUES: u64 = 1024; diff --git a/rust/lance-encoding/src/compression/block/factory.rs b/rust/lance-encoding/src/compression/block/factory.rs index 525049cdce0..f92ee8f6bef 100644 --- a/rust/lance-encoding/src/compression/block/factory.rs +++ b/rust/lance-encoding/src/compression/block/factory.rs @@ -14,6 +14,7 @@ use crate::{ block::{CompressionConfig, CompressionScheme}, constant::ConstantBlockDecompressor, delta::DeltaDecompressor, + dictionary::DictionaryBlockDecompressor, general::GenericGeneralBlockDecompressor, range::RangeDecompressor, rle::{BlockRleDecompressor, BlockRunCount, MetadataRunLengths}, @@ -436,6 +437,47 @@ fn create_inner( values_have_payload || run_lengths_have_payload, )) } + Compression::Dictionary(dictionary) => { + if position != Position::Root { + return Err(Error::invalid_input( + "Dictionary is not supported as a block codec child", + )); + } + if dictionary.num_dictionary_items == 0 + || dictionary.num_dictionary_items as usize > MAX_DICTIONARY_ITEMS + { + return Err(Error::invalid_input(format!( + "Dictionary item count {} is outside 1..={MAX_DICTIONARY_ITEMS}", + dictionary.num_dictionary_items + ))); + } + let indices_encoding = dictionary.indices.as_deref().ok_or_else(|| { + Error::invalid_input("Dictionary is missing its indices encoding") + })?; + let items_encoding = dictionary + .items + .as_deref() + .ok_or_else(|| Error::invalid_input("Dictionary is missing its items encoding"))?; + let (indices, indices_have_payload) = create_inner( + indices_encoding, + BlockValueType::UInt32, + Position::Child, + false, + )?; + let (items, items_have_payload) = + create_inner(items_encoding, expected_type, Position::Child, false)?; + Ok(( + Box::new(DictionaryBlockDecompressor::new( + expected_type, + dictionary.num_dictionary_items, + indices, + items, + indices_have_payload, + items_have_payload, + )), + indices_have_payload || items_have_payload, + )) + } other => Err(Error::invalid_input(format!( "Unsupported block sequence encoding: {}", compression_name(other) @@ -559,6 +601,13 @@ fn infer_inner(encoding: &CompressiveEncoding, position: Position) -> Result infer_inner( + dictionary + .items + .as_deref() + .ok_or_else(|| Error::invalid_input("Dictionary is missing its items encoding"))?, + Position::Child, + ), other => Err(Error::invalid_input(format!( "Cannot infer bounded block value type from {} at {position:?}", compression_name(other) diff --git a/rust/lance-encoding/src/compression/block/tests.rs b/rust/lance-encoding/src/compression/block/tests.rs index bd3b566396c..eef76c4d570 100644 --- a/rust/lance-encoding/src/compression/block/tests.rs +++ b/rust/lance-encoding/src/compression/block/tests.rs @@ -3,13 +3,15 @@ use super::*; +use std::sync::Arc; + use crate::{ buffer::LanceBuffer, compression::BlockCompressor, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, encodings::physical::{ - constant::ConstantBlockCompressor, range::RangeEncoder, rle::BlockRleCompressor, - value::FixedWidthBlockCompressor, + constant::ConstantBlockCompressor, dictionary::DictionaryBlockCompressor, + range::RangeEncoder, rle::BlockRleCompressor, value::FixedWidthBlockCompressor, }, format::{ProtobufUtils21, pb21::CompressiveEncoding}, }; @@ -145,6 +147,30 @@ fn rle_compressor_owns_and_reuses_children() { round_trip_u64(&values, compressor, encoding); } +#[test] +fn dictionary_materializes_once_when_built() { + use crate::encodings::physical::dictionary::BLOCK_MATERIALIZATION_COUNT; + + let values = (0..4096_u64) + .map(|index| [u64::MAX - 9, 17, 1_u64 << 50, 991][index as usize % 4]) + .collect::>(); + let dictionary_items = Arc::from([17_u64, 991, 1_u64 << 50, u64::MAX - 9].as_slice()); + let compressor = Box::new(DictionaryBlockCompressor::new( + BlockValueType::UInt64, + dictionary_items, + Box::new(FixedWidthBlockCompressor::new(BlockValueType::UInt32)), + Box::new(FixedWidthBlockCompressor::new(BlockValueType::UInt64)), + )); + let encoding = ProtobufUtils21::dictionary( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::flat(64, None), + 4, + ); + BLOCK_MATERIALIZATION_COUNT.with(|count| count.set(0)); + round_trip_u64(&values, compressor, encoding); + BLOCK_MATERIALIZATION_COUNT.with(|count| assert_eq!(count.get(), 1)); +} + #[cfg(feature = "bitpacking")] #[test] fn out_of_line_bitpacking_round_trip() { @@ -206,6 +232,18 @@ fn factory_rejects_unbounded_or_mistyped_trees() { .to_string() .contains("expected 64") ); + + let oversized = ProtobufUtils21::dictionary( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::flat(64, None), + (MAX_DICTIONARY_ITEMS + 1) as u32, + ); + assert!( + create_block_decompressor(&oversized, BlockValueType::UInt64) + .unwrap_err() + .to_string() + .contains("outside") + ); } #[test] @@ -328,6 +366,60 @@ fn rle_framing_is_fallible() { assert!(error.to_string().contains("Metadata-only RLE run-length")); } +#[test] +fn metadata_dictionary_checks_index_bounds() { + let encoding = ProtobufUtils21::dictionary( + ProtobufUtils21::range(32, 0, 1), + ProtobufUtils21::range(64, 10, 1), + 2, + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(!has_payload); + let error = decoder.decompress(None, 3).unwrap_err(); + assert!(error.to_string().contains("out of bounds")); +} + +#[test] +fn dictionary_framing_rejects_inconsistent_lengths() { + let encoding = ProtobufUtils21::dictionary( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::flat(64, None), + 2, + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(has_payload); + let mut payload = 4_u64.to_le_bytes().to_vec(); + payload.extend_from_slice(&16_u64.to_le_bytes()); + payload.extend_from_slice(&[0; 8]); + assert!( + decoder + .decompress(Some(LanceBuffer::from(payload)), 1) + .is_err() + ); +} + +#[test] +fn dictionary_framing_rejects_payload_for_metadata_child() { + let encoding = ProtobufUtils21::dictionary( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::range(64, 10, 1), + 2, + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(has_payload); + let mut payload = 4_u64.to_le_bytes().to_vec(); + payload.extend_from_slice(&1_u64.to_le_bytes()); + payload.extend_from_slice(&0_u32.to_le_bytes()); + payload.push(0); + let error = decoder + .decompress(Some(LanceBuffer::from(payload)), 1) + .unwrap_err(); + assert!(error.to_string().contains("Metadata-only Dictionary items")); +} + #[test] fn decoder_allocation_overflow_is_fallible() { let (decoder, has_payload) = create_block_decompressor( diff --git a/rust/lance-encoding/src/encodings/physical.rs b/rust/lance-encoding/src/encodings/physical.rs index 849b720e4d0..46c77d271e2 100644 --- a/rust/lance-encoding/src/encodings/physical.rs +++ b/rust/lance-encoding/src/encodings/physical.rs @@ -10,6 +10,7 @@ pub mod block; pub mod byte_stream_split; pub mod constant; pub mod delta; +pub mod dictionary; pub mod fsst; pub mod general; pub mod packed; diff --git a/rust/lance-encoding/src/encodings/physical/dictionary.rs b/rust/lance-encoding/src/encodings/physical/dictionary.rs new file mode 100644 index 00000000000..5713936fd69 --- /dev/null +++ b/rust/lance-encoding/src/encodings/physical/dictionary.rs @@ -0,0 +1,329 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Dictionary codec for bounded unsigned block sequences. + +#[cfg(test)] +use std::cell::Cell; +#[cfg(test)] +use std::{collections::BTreeMap, sync::Arc}; + +use crate::{ + buffer::LanceBuffer, + compression::{ + BlockDecompressor, BlockValueType, + block::{fixed_block, read_unsigned_values, validate_fixed_payload_len}, + }, + data::DataBlock, + encodings::physical::try_vec_with_capacity, +}; +#[cfg(test)] +use crate::{ + compression::{ + BlockCompressor, + block::{fixed_from_u64_values, visit_unsigned_values}, + }, + data::{BlockInfo, FixedWidthDataBlock}, +}; +use lance_core::{Error, Result}; + +pub(crate) const BLOCK_FRAME_BYTES: u64 = 16; + +#[cfg(test)] +thread_local! { + pub(crate) static BLOCK_MATERIALIZATION_COUNT: Cell = const { Cell::new(0) }; +} + +/// Dictionary compressor that owns its indices and items compressors. +#[cfg(test)] +#[derive(Debug)] +pub(crate) struct DictionaryBlockCompressor { + value_type: BlockValueType, + dictionary_items: Arc<[u64]>, + indices: Box, + items: Box, +} + +#[cfg(test)] +impl DictionaryBlockCompressor { + pub(crate) fn new( + value_type: BlockValueType, + dictionary_items: Arc<[u64]>, + indices: Box, + items: Box, + ) -> Self { + Self { + value_type, + dictionary_items, + indices, + items, + } + } +} + +#[cfg(test)] +impl BlockCompressor for DictionaryBlockCompressor { + fn compress(&self, data: DataBlock) -> Result> { + #[cfg(test)] + BLOCK_MATERIALIZATION_COUNT.with(|count| count.set(count.get().saturating_add(1))); + + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Dictionary block compression requires fixed-width data", + )); + }; + if data.bits_per_value != self.value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Dictionary compressor expects {}-bit values, got {}", + self.value_type.bits_per_value(), + data.bits_per_value + ))); + } + let dictionary_index = self + .dictionary_items + .iter() + .enumerate() + .map(|(index, value)| (*value, index as u32)) + .collect::>(); + let mut encoded_indices = + try_vec_with_capacity::(data.num_values, "Dictionary indices")?; + let mut position = 0_u64; + visit_unsigned_values(&data, self.value_type, |value| { + encoded_indices.push(*dictionary_index.get(&value).ok_or_else(|| { + Error::invalid_input(format!( + "Dictionary compressor does not contain input value {value} at position {position}" + )) + })?); + position += 1; + Ok(()) + })?; + + let indices_block = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(encoded_indices), + num_values: data.num_values, + block_info: BlockInfo::default(), + }; + let items_block = + fixed_from_u64_values(&self.dictionary_items, self.value_type, "Dictionary items")?; + let indices_payload = self + .indices + .compress(DataBlock::FixedWidth(indices_block))?; + let items_payload = self.items.compress(DataBlock::FixedWidth(items_block))?; + if indices_payload.is_none() && items_payload.is_none() { + return Ok(None); + } + let indices_payload = indices_payload.unwrap_or_else(LanceBuffer::empty); + let items_payload = items_payload.unwrap_or_else(LanceBuffer::empty); + + let mut output = try_frame(indices_payload.len(), items_payload.len())?; + output.extend_from_slice(&(indices_payload.len() as u64).to_le_bytes()); + output.extend_from_slice(&(items_payload.len() as u64).to_le_bytes()); + output.extend_from_slice(&indices_payload); + output.extend_from_slice(&items_payload); + Ok(Some(LanceBuffer::from(output))) + } +} + +/// Dictionary decompressor that owns its indices and items decompressors. +#[derive(Debug)] +pub(crate) struct DictionaryBlockDecompressor { + value_type: BlockValueType, + num_dictionary_items: u32, + indices: Box, + items: Box, + indices_have_payload: bool, + items_have_payload: bool, +} + +impl DictionaryBlockDecompressor { + pub(crate) fn new( + value_type: BlockValueType, + num_dictionary_items: u32, + indices: Box, + items: Box, + indices_have_payload: bool, + items_have_payload: bool, + ) -> Self { + Self { + value_type, + num_dictionary_items, + indices, + items, + indices_have_payload, + items_have_payload, + } + } +} + +impl BlockDecompressor for DictionaryBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let has_payload = self.indices_have_payload || self.items_have_payload; + let (indices_payload, items_payload) = if let Some(data) = data { + if !has_payload { + return Err(Error::invalid_input( + "Metadata-only Dictionary expects no payload", + )); + } + if data.len() < BLOCK_FRAME_BYTES as usize { + return Err(Error::invalid_input(format!( + "Dictionary payload has {} bytes, shorter than its {BLOCK_FRAME_BYTES}-byte header", + data.len() + ))); + } + let indices_size = + u64::from_le_bytes(data[..8].try_into().expect("header length was checked")); + let items_size = + u64::from_le_bytes(data[8..16].try_into().expect("header length was checked")); + let indices_size = usize::try_from(indices_size).map_err(|_| { + Error::invalid_input("Dictionary indices payload length does not fit usize") + })?; + let items_size = usize::try_from(items_size).map_err(|_| { + Error::invalid_input("Dictionary items payload length does not fit usize") + })?; + let indices_start = BLOCK_FRAME_BYTES as usize; + let items_start = indices_start + .checked_add(indices_size) + .ok_or_else(|| Error::invalid_input("Dictionary indices payload end overflows"))?; + let end = items_start + .checked_add(items_size) + .ok_or_else(|| Error::invalid_input("Dictionary items payload end overflows"))?; + if end != data.len() { + return Err(Error::invalid_input(format!( + "Dictionary framing describes {end} bytes, payload has {}", + data.len() + ))); + } + if !self.indices_have_payload && indices_size != 0 { + return Err(Error::invalid_input(format!( + "Metadata-only Dictionary indices child has {indices_size} framed payload bytes" + ))); + } + if !self.items_have_payload && items_size != 0 { + return Err(Error::invalid_input(format!( + "Metadata-only Dictionary items child has {items_size} framed payload bytes" + ))); + } + ( + self.indices_have_payload + .then(|| data.slice_with_length(indices_start, indices_size)), + self.items_have_payload + .then(|| data.slice_with_length(items_start, items_size)), + ) + } else { + if has_payload { + return Err(Error::invalid_input("Dictionary requires one payload")); + } + (None, None) + }; + + let indices = self.indices.decompress(indices_payload, num_values)?; + let items = self + .items + .decompress(items_payload, u64::from(self.num_dictionary_items))?; + let DataBlock::FixedWidth(indices) = indices else { + return Err(Error::invalid_input( + "Dictionary indices decoded to a non fixed-width block", + )); + }; + let DataBlock::FixedWidth(items) = items else { + return Err(Error::invalid_input( + "Dictionary items decoded to a non fixed-width block", + )); + }; + validate_fixed_payload_len( + &indices.data, + BlockValueType::UInt32, + num_values, + "Dictionary indices", + )?; + validate_fixed_payload_len( + &items.data, + self.value_type, + u64::from(self.num_dictionary_items), + "Dictionary items", + )?; + let indices = indices.data.borrow_to_typed_slice::(); + let items = read_unsigned_values(&items, self.value_type)?; + let output = match self.value_type { + BlockValueType::UInt8 => { + let mut output = try_vec_with_capacity::(num_values, "Dictionary output")?; + append_items( + &mut output, + &indices, + &items, + self.num_dictionary_items, + |value| value as u8, + )?; + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt16 => { + let mut output = try_vec_with_capacity::(num_values, "Dictionary output")?; + append_items( + &mut output, + &indices, + &items, + self.num_dictionary_items, + |value| value as u16, + )?; + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt32 => { + let mut output = try_vec_with_capacity::(num_values, "Dictionary output")?; + append_items( + &mut output, + &indices, + &items, + self.num_dictionary_items, + |value| value as u32, + )?; + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt64 => { + let mut output = try_vec_with_capacity::(num_values, "Dictionary output")?; + append_items( + &mut output, + &indices, + &items, + self.num_dictionary_items, + |value| value, + )?; + LanceBuffer::reinterpret_vec(output) + } + }; + Ok(fixed_block(self.value_type, num_values, output)) + } +} + +fn append_items( + output: &mut Vec, + indices: &[u32], + items: &[u64], + num_dictionary_items: u32, + convert: impl Fn(u64) -> T, +) -> Result<()> { + for (position, index) in indices.iter().enumerate() { + let value = *items.get(*index as usize).ok_or_else(|| { + Error::invalid_input(format!( + "Dictionary index {index} at position {position} is out of bounds for {num_dictionary_items} items" + )) + })?; + output.push(convert(value)); + } + Ok(()) +} + +#[cfg(test)] +fn try_frame(indices_payload_bytes: usize, items_payload_bytes: usize) -> Result> { + let capacity = (BLOCK_FRAME_BYTES as usize) + .checked_add(indices_payload_bytes) + .and_then(|capacity| capacity.checked_add(items_payload_bytes)) + .ok_or_else(|| Error::invalid_input("Dictionary frame length overflows usize"))?; + let mut output = Vec::new(); + output.try_reserve_exact(capacity).map_err(|error| { + Error::invalid_input(format!( + "Dictionary could not reserve {capacity} frame bytes: {error}" + )) + })?; + Ok(output) +} diff --git a/rust/lance-encoding/src/format.rs b/rust/lance-encoding/src/format.rs index 74f1eafb831..f0c2ca575c6 100644 --- a/rust/lance-encoding/src/format.rs +++ b/rust/lance-encoding/src/format.rs @@ -706,6 +706,24 @@ impl ProtobufUtils21 { } } + pub fn dictionary( + indices: crate::format::pb21::CompressiveEncoding, + items: crate::format::pb21::CompressiveEncoding, + num_dictionary_items: u32, + ) -> crate::format::pb21::CompressiveEncoding { + crate::format::pb21::CompressiveEncoding { + compression: Some( + crate::format::pb21::compressive_encoding::Compression::Dictionary(Box::new( + crate::format::pb21::Dictionary { + indices: Some(Box::new(indices)), + items: Some(Box::new(items)), + num_dictionary_items, + }, + )), + ), + } + } + pub fn constant_layout( def_meaning: &[DefinitionInterpretation], inline_value: Option>,