From 7de28e7a63cceedf80517f934d242ee22e9d10b0 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 5 Aug 2026 20:28:56 +0800 Subject: [PATCH] feat(encoding): add bounded block dictionary --- docs/src/format/file/encoding.md | 16 + rust/lance-encoding/src/encodings/physical.rs | 1 + .../src/encodings/physical/dictionary.rs | 520 ++++++++++++++++++ rust/lance-encoding/src/format.rs | 18 + 4 files changed, 555 insertions(+) create mode 100644 rust/lance-encoding/src/encodings/physical/dictionary.rs diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index 0b92eabcb90..b701f63f55e 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) | @@ -639,6 +640,21 @@ has one, and readers reconstruct values with checked prefix sums. %%% proto.message.Delta %%% ``` +### Block Dictionary + +A block dictionary represents unsigned `u32` or `u64` values with a `u32` indices child and a dictionary-items +child. Both children are concrete block codecs. When either child has a payload, the outer payload is framed as: + +```text +u64 indices_payload_bytes +u64 items_payload_bytes +indices payload +dictionary-items payload +``` + +A metadata-only child has a zero framed length. If both children are metadata-only, the dictionary itself has no +payload. Readers validate the frame boundaries, child widths and cardinalities, and every dictionary index. + ### Bitpacking Bitpacking is a compression technique that removes the unused bits from a set of values. For example, if we have diff --git a/rust/lance-encoding/src/encodings/physical.rs b/rust/lance-encoding/src/encodings/physical.rs index 119de66e53b..0c4ea45816e 100644 --- a/rust/lance-encoding/src/encodings/physical.rs +++ b/rust/lance-encoding/src/encodings/physical.rs @@ -13,6 +13,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..567cd161a96 --- /dev/null +++ b/rust/lance-encoding/src/encodings/physical/dictionary.rs @@ -0,0 +1,520 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Dictionary encoding for bounded unsigned block sequences. + +use std::{collections::BTreeMap, sync::Arc}; + +use crate::{ + buffer::LanceBuffer, + compression::{BlockCompressor, BlockDecompressor}, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + encodings::physical::{checked_fixed_values, try_vec_with_capacity}, + format::{ProtobufUtils21, pb21::CompressiveEncoding}, +}; +use lance_core::{Error, Result}; + +pub(crate) const MAX_BLOCK_DICTIONARY_ITEMS: usize = 4096; +const BLOCK_FRAME_BYTES: usize = 16; + +/// Encodes unsigned values through concrete index and dictionary-item codecs. +#[derive(Debug)] +pub struct BlockDictionaryEncoder { + bits_per_value: u64, + dictionary_items: Arc<[u64]>, + indices: Box, + items: Box, +} + +impl BlockDictionaryEncoder { + /// Creates a bounded dictionary codec for `u32` or `u64` values. + pub fn try_new( + bits_per_value: u64, + dictionary_items: Arc<[u64]>, + indices: Box, + items: Box, + ) -> Result { + validate_dictionary_header(bits_per_value, dictionary_items.len())?; + if bits_per_value == 32 + && dictionary_items + .iter() + .any(|value| *value > u32::MAX as u64) + { + return Err(Error::invalid_input( + "Dictionary contains an item that exceeds u32::MAX", + )); + } + Ok(Self { + bits_per_value, + dictionary_items, + indices, + items, + }) + } +} + +impl BlockCompressor for BlockDictionaryEncoder { + fn compress(&self, data: DataBlock) -> Result<(Option, CompressiveEncoding)> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Dictionary block compression requires fixed-width data", + )); + }; + if data.bits_per_value != self.bits_per_value { + return Err(Error::invalid_input(format!( + "Dictionary codec expects {}-bit values, got {}", + self.bits_per_value, data.bits_per_value + ))); + } + + let dictionary_index = self + .dictionary_items + .iter() + .enumerate() + .map(|(index, value)| (*value, index as u32)) + .collect::>(); + let values = read_unsigned_values(&data, "Dictionary input")?; + let mut encoded_indices = + try_vec_with_capacity::(data.num_values, "Dictionary indices")?; + for (position, value) in values.iter().enumerate() { + encoded_indices.push(*dictionary_index.get(value).ok_or_else(|| { + Error::invalid_input(format!( + "Dictionary does not contain input value {value} at position {position}" + )) + })?); + } + + 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.bits_per_value, + "Dictionary items", + )?; + let (indices_payload, indices_encoding) = self + .indices + .compress(DataBlock::FixedWidth(indices_block))?; + let (items_payload, items_encoding) = + self.items.compress(DataBlock::FixedWidth(items_block))?; + let encoding = ProtobufUtils21::dictionary( + indices_encoding, + items_encoding, + self.dictionary_items.len() as u32, + ); + if indices_payload.is_none() && items_payload.is_none() { + return Ok((None, encoding)); + } + + 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)), encoding)) + } +} + +/// Decodes a framed block dictionary through concrete child codecs. +#[derive(Debug)] +pub struct BlockDictionaryDecompressor { + bits_per_value: u64, + num_dictionary_items: u32, + indices: Box, + items: Box, +} + +impl BlockDictionaryDecompressor { + /// Creates a bounded dictionary decoder for `u32` or `u64` values. + pub fn try_new( + bits_per_value: u64, + num_dictionary_items: u32, + indices: Box, + items: Box, + ) -> Result { + validate_dictionary_header(bits_per_value, num_dictionary_items as usize)?; + Ok(Self { + bits_per_value, + num_dictionary_items, + indices, + items, + }) + } +} + +impl BlockDecompressor for BlockDictionaryDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let indices_have_payload = self.indices.requires_payload(); + let items_have_payload = self.items.requires_payload(); + let has_payload = indices_have_payload || items_have_payload; + let (indices_payload, items_payload) = + split_payload(data, indices_have_payload, items_have_payload, has_payload)?; + + 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", + )); + }; + if indices.bits_per_value != 32 { + return Err(Error::invalid_input(format!( + "Dictionary indices decoded as {}-bit values, expected 32", + indices.bits_per_value + ))); + } + if items.bits_per_value != self.bits_per_value { + return Err(Error::invalid_input(format!( + "Dictionary items decoded as {}-bit values, expected {}", + items.bits_per_value, self.bits_per_value + ))); + } + + let indices = checked_fixed_values::(&indices, "Dictionary indices")?; + let items = read_unsigned_values(&items, "Dictionary items")?; + let output = match self.bits_per_value { + 32 => { + 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) + } + 64 => { + 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) + } + _ => unreachable!("dictionary width was validated at construction"), + }; + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: self.bits_per_value, + data: output, + num_values, + block_info: BlockInfo::default(), + })) + } + + fn requires_payload(&self) -> bool { + self.indices.requires_payload() || self.items.requires_payload() + } +} + +fn validate_dictionary_header(bits_per_value: u64, num_dictionary_items: usize) -> Result<()> { + if !matches!(bits_per_value, 32 | 64) { + return Err(Error::invalid_input(format!( + "Dictionary only supports 32 or 64-bit values, got {bits_per_value}" + ))); + } + if !(1..=MAX_BLOCK_DICTIONARY_ITEMS).contains(&num_dictionary_items) { + return Err(Error::invalid_input(format!( + "Dictionary item count {num_dictionary_items} is outside 1..={MAX_BLOCK_DICTIONARY_ITEMS}" + ))); + } + Ok(()) +} + +fn split_payload( + data: Option, + indices_have_payload: bool, + items_have_payload: bool, + has_payload: bool, +) -> Result<(Option, Option)> { + let Some(data) = data else { + if has_payload { + return Err(Error::invalid_input("Dictionary requires one payload")); + } + return Ok((None, None)); + }; + if !has_payload { + return Err(Error::invalid_input( + "Metadata-only Dictionary expects no payload", + )); + } + if data.len() < BLOCK_FRAME_BYTES { + 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().unwrap()); + let items_size = u64::from_le_bytes(data[8..16].try_into().unwrap()); + 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 items_start = BLOCK_FRAME_BYTES + .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 !indices_have_payload && indices_size != 0 { + return Err(Error::invalid_input(format!( + "Metadata-only Dictionary indices child has {indices_size} framed payload bytes" + ))); + } + if !items_have_payload && items_size != 0 { + return Err(Error::invalid_input(format!( + "Metadata-only Dictionary items child has {items_size} framed payload bytes" + ))); + } + Ok(( + indices_have_payload.then(|| data.slice_with_length(BLOCK_FRAME_BYTES, indices_size)), + items_have_payload.then(|| data.slice_with_length(items_start, items_size)), + )) +} + +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(()) +} + +fn read_unsigned_values(data: &FixedWidthDataBlock, label: &str) -> Result> { + match data.bits_per_value { + 32 => Ok(checked_fixed_values::(data, label)? + .iter() + .map(|value| u64::from(*value)) + .collect()), + 64 => Ok(checked_fixed_values::(data, label)?.to_vec()), + bits_per_value => Err(Error::invalid_input(format!( + "{label} uses unsupported {bits_per_value}-bit values" + ))), + } +} + +fn fixed_from_u64_values( + values: &[u64], + bits_per_value: u64, + label: &str, +) -> Result { + let data = match bits_per_value { + 32 => LanceBuffer::reinterpret_vec( + values + .iter() + .map(|value| { + u32::try_from(*value).map_err(|_| { + Error::invalid_input(format!("{label} value {value} exceeds u32::MAX")) + }) + }) + .collect::>>()?, + ), + 64 => LanceBuffer::reinterpret_vec(values.to_vec()), + _ => unreachable!("dictionary width was validated at construction"), + }; + Ok(FixedWidthDataBlock { + bits_per_value, + data, + num_values: values.len() as u64, + block_info: BlockInfo::default(), + }) +} + +fn try_frame(indices_payload_bytes: usize, items_payload_bytes: usize) -> Result> { + let capacity = BLOCK_FRAME_BYTES + .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) +} + +#[cfg(test)] +mod tests { + use crate::encodings::physical::{ + range::{RangeDecompressor, RangeEncoder}, + value::{ValueDecompressor, ValueEncoder}, + }; + use crate::format::pb21::Flat; + + use super::*; + + fn input_u64(values: Vec) -> DataBlock { + DataBlock::FixedWidth(FixedWidthDataBlock { + num_values: values.len() as u64, + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(values), + block_info: BlockInfo::default(), + }) + } + + fn flat_decoder(bits_per_value: u64) -> Box { + Box::new(ValueDecompressor::from_flat(&Flat { + bits_per_value, + data: None, + })) + } + + fn decoded_u64(block: DataBlock) -> Vec { + block + .as_fixed_width() + .unwrap() + .data + .borrow_to_typed_slice::() + .to_vec() + } + + #[test] + fn dictionary_round_trip_with_payload_children() { + let items: Arc<[u64]> = Arc::from([10, 20]); + let encoder = BlockDictionaryEncoder::try_new( + 64, + items, + Box::new(ValueEncoder::default()), + Box::new(ValueEncoder::default()), + ) + .unwrap(); + let (payload, encoding) = encoder + .compress(input_u64(vec![10, 20, 10, 20, 20])) + .unwrap(); + let payload = payload.unwrap(); + assert_eq!(u64::from_le_bytes(payload[..8].try_into().unwrap()), 20); + assert_eq!(u64::from_le_bytes(payload[8..16].try_into().unwrap()), 16); + assert!(matches!( + encoding.compression, + Some(crate::format::pb21::compressive_encoding::Compression::Dictionary(_)) + )); + + let decoder = + BlockDictionaryDecompressor::try_new(64, 2, flat_decoder(32), flat_decoder(64)) + .unwrap(); + assert_eq!( + decoded_u64(decoder.decompress(Some(payload), 5).unwrap()), + vec![10, 20, 10, 20, 20] + ); + } + + #[test] + fn dictionary_round_trip_without_payload() { + let encoder = BlockDictionaryEncoder::try_new( + 64, + Arc::from([10, 20, 30]), + Box::new(RangeEncoder::new(32, 0, 1)), + Box::new(RangeEncoder::new(64, 10, 10)), + ) + .unwrap(); + let (payload, _) = encoder.compress(input_u64(vec![10, 20, 30])).unwrap(); + assert!(payload.is_none()); + + let decoder = BlockDictionaryDecompressor::try_new( + 64, + 3, + Box::new(RangeDecompressor::new(32, 0, 1)), + Box::new(RangeDecompressor::new(64, 10, 10)), + ) + .unwrap(); + assert!(!decoder.requires_payload()); + assert_eq!( + decoded_u64(decoder.decompress(None, 3).unwrap()), + vec![10, 20, 30] + ); + } + + #[test] + fn dictionary_frames_a_mixed_payload() { + let encoder = BlockDictionaryEncoder::try_new( + 64, + Arc::from([10, 20]), + Box::new(ValueEncoder::default()), + Box::new(RangeEncoder::new(64, 10, 10)), + ) + .unwrap(); + let (payload, _) = encoder.compress(input_u64(vec![10, 20, 10])).unwrap(); + let payload = payload.unwrap(); + assert_eq!(u64::from_le_bytes(payload[..8].try_into().unwrap()), 12); + assert_eq!(u64::from_le_bytes(payload[8..16].try_into().unwrap()), 0); + + let decoder = BlockDictionaryDecompressor::try_new( + 64, + 2, + flat_decoder(32), + Box::new(RangeDecompressor::new(64, 10, 10)), + ) + .unwrap(); + assert_eq!( + decoded_u64(decoder.decompress(Some(payload), 3).unwrap()), + vec![10, 20, 10] + ); + } + + #[test] + fn dictionary_rejects_out_of_bounds_index() { + let indices = LanceBuffer::reinterpret_vec(vec![0_u32, 2]); + let items = LanceBuffer::reinterpret_vec(vec![9_u64]); + let mut frame = try_frame(indices.len(), items.len()).unwrap(); + frame.extend_from_slice(&(indices.len() as u64).to_le_bytes()); + frame.extend_from_slice(&(items.len() as u64).to_le_bytes()); + frame.extend_from_slice(&indices); + frame.extend_from_slice(&items); + + let decoder = + BlockDictionaryDecompressor::try_new(64, 1, flat_decoder(32), flat_decoder(64)) + .unwrap(); + let error = decoder + .decompress(Some(LanceBuffer::from(frame)), 2) + .unwrap_err(); + assert!(error.to_string().contains("out of bounds")); + } + + #[test] + fn dictionary_rejects_invalid_frame_lengths() { + let mut frame = Vec::from(16_u64.to_le_bytes()); + frame.extend_from_slice(&0_u64.to_le_bytes()); + + let decoder = + BlockDictionaryDecompressor::try_new(64, 1, flat_decoder(32), flat_decoder(64)) + .unwrap(); + let error = decoder + .decompress(Some(LanceBuffer::from(frame)), 1) + .unwrap_err(); + assert!(error.to_string().contains("framing describes")); + } +} 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>,