lance_encoding/encodings/physical/
dictionary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Dictionary encoding for bounded unsigned block sequences.
5
6use std::{collections::BTreeMap, sync::Arc};
7
8use crate::{
9    buffer::LanceBuffer,
10    compression::{BlockCompressor, BlockDecompressor},
11    data::{BlockInfo, DataBlock, FixedWidthDataBlock},
12    encodings::physical::{checked_fixed_values, try_vec_with_capacity},
13    format::{EncodingUtils, metadata::CompressiveEncoding},
14};
15use lance_core::{Error, Result};
16
17pub(crate) const MAX_BLOCK_DICTIONARY_ITEMS: usize = 4096;
18const BLOCK_FRAME_BYTES: usize = 16;
19
20/// Encodes unsigned values through concrete index and dictionary-item codecs.
21#[derive(Debug)]
22pub struct BlockDictionaryEncoder {
23    bits_per_value: u64,
24    dictionary_items: Arc<[u64]>,
25    indices: Box<dyn BlockCompressor>,
26    items: Box<dyn BlockCompressor>,
27}
28
29impl BlockDictionaryEncoder {
30    /// Creates a bounded dictionary codec for `u32` or `u64` values.
31    pub fn try_new(
32        bits_per_value: u64,
33        dictionary_items: Arc<[u64]>,
34        indices: Box<dyn BlockCompressor>,
35        items: Box<dyn BlockCompressor>,
36    ) -> Result<Self> {
37        validate_dictionary_header(bits_per_value, dictionary_items.len())?;
38        if bits_per_value == 32
39            && dictionary_items
40                .iter()
41                .any(|value| *value > u32::MAX as u64)
42        {
43            return Err(Error::invalid_input(
44                "Dictionary contains an item that exceeds u32::MAX",
45            ));
46        }
47        Ok(Self {
48            bits_per_value,
49            dictionary_items,
50            indices,
51            items,
52        })
53    }
54}
55
56impl BlockCompressor for BlockDictionaryEncoder {
57    fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)> {
58        let DataBlock::FixedWidth(data) = data else {
59            return Err(Error::invalid_input(
60                "Dictionary block compression requires fixed-width data",
61            ));
62        };
63        if data.bits_per_value != self.bits_per_value {
64            return Err(Error::invalid_input(format!(
65                "Dictionary codec expects {}-bit values, got {}",
66                self.bits_per_value, data.bits_per_value
67            )));
68        }
69
70        let dictionary_index = self
71            .dictionary_items
72            .iter()
73            .enumerate()
74            .map(|(index, value)| (*value, index as u32))
75            .collect::<BTreeMap<_, _>>();
76        let values = read_unsigned_values(&data, "Dictionary input")?;
77        let mut encoded_indices =
78            try_vec_with_capacity::<u32>(data.num_values, "Dictionary indices")?;
79        for (position, value) in values.iter().enumerate() {
80            encoded_indices.push(*dictionary_index.get(value).ok_or_else(|| {
81                Error::invalid_input(format!(
82                    "Dictionary does not contain input value {value} at position {position}"
83                ))
84            })?);
85        }
86
87        let indices_block = FixedWidthDataBlock {
88            bits_per_value: 32,
89            data: LanceBuffer::reinterpret_vec(encoded_indices),
90            num_values: data.num_values,
91            block_info: BlockInfo::default(),
92        };
93        let items_block = fixed_from_u64_values(
94            &self.dictionary_items,
95            self.bits_per_value,
96            "Dictionary items",
97        )?;
98        let (indices_payload, indices_encoding) = self
99            .indices
100            .compress(DataBlock::FixedWidth(indices_block))?;
101        let (items_payload, items_encoding) =
102            self.items.compress(DataBlock::FixedWidth(items_block))?;
103        let encoding = EncodingUtils::dictionary(
104            indices_encoding,
105            items_encoding,
106            self.dictionary_items.len() as u32,
107        );
108        if indices_payload.is_none() && items_payload.is_none() {
109            return Ok((None, encoding));
110        }
111
112        let indices_payload = indices_payload.unwrap_or_else(LanceBuffer::empty);
113        let items_payload = items_payload.unwrap_or_else(LanceBuffer::empty);
114        let mut output = try_frame(indices_payload.len(), items_payload.len())?;
115        output.extend_from_slice(&(indices_payload.len() as u64).to_le_bytes());
116        output.extend_from_slice(&(items_payload.len() as u64).to_le_bytes());
117        output.extend_from_slice(&indices_payload);
118        output.extend_from_slice(&items_payload);
119        Ok((Some(LanceBuffer::from(output)), encoding))
120    }
121}
122
123/// Decodes a framed block dictionary through concrete child codecs.
124#[derive(Debug)]
125pub struct BlockDictionaryDecompressor {
126    bits_per_value: u64,
127    num_dictionary_items: u32,
128    indices: Box<dyn BlockDecompressor>,
129    items: Box<dyn BlockDecompressor>,
130}
131
132impl BlockDictionaryDecompressor {
133    /// Creates a bounded dictionary decoder for `u32` or `u64` values.
134    pub fn try_new(
135        bits_per_value: u64,
136        num_dictionary_items: u32,
137        indices: Box<dyn BlockDecompressor>,
138        items: Box<dyn BlockDecompressor>,
139    ) -> Result<Self> {
140        validate_dictionary_header(bits_per_value, num_dictionary_items as usize)?;
141        Ok(Self {
142            bits_per_value,
143            num_dictionary_items,
144            indices,
145            items,
146        })
147    }
148}
149
150impl BlockDecompressor for BlockDictionaryDecompressor {
151    fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
152        let indices_have_payload = self.indices.requires_payload();
153        let items_have_payload = self.items.requires_payload();
154        let has_payload = indices_have_payload || items_have_payload;
155        let (indices_payload, items_payload) =
156            split_payload(data, indices_have_payload, items_have_payload, has_payload)?;
157
158        let indices = self.indices.decompress(indices_payload, num_values)?;
159        let items = self
160            .items
161            .decompress(items_payload, u64::from(self.num_dictionary_items))?;
162        let DataBlock::FixedWidth(indices) = indices else {
163            return Err(Error::invalid_input(
164                "Dictionary indices decoded to a non fixed-width block",
165            ));
166        };
167        let DataBlock::FixedWidth(items) = items else {
168            return Err(Error::invalid_input(
169                "Dictionary items decoded to a non fixed-width block",
170            ));
171        };
172        if indices.bits_per_value != 32 {
173            return Err(Error::invalid_input(format!(
174                "Dictionary indices decoded as {}-bit values, expected 32",
175                indices.bits_per_value
176            )));
177        }
178        if items.bits_per_value != self.bits_per_value {
179            return Err(Error::invalid_input(format!(
180                "Dictionary items decoded as {}-bit values, expected {}",
181                items.bits_per_value, self.bits_per_value
182            )));
183        }
184
185        let indices = checked_fixed_values::<u32>(&indices, "Dictionary indices")?;
186        let items = read_unsigned_values(&items, "Dictionary items")?;
187        let output = match self.bits_per_value {
188            32 => {
189                let mut output = try_vec_with_capacity::<u32>(num_values, "Dictionary output")?;
190                append_items(
191                    &mut output,
192                    &indices,
193                    &items,
194                    self.num_dictionary_items,
195                    |value| value as u32,
196                )?;
197                LanceBuffer::reinterpret_vec(output)
198            }
199            64 => {
200                let mut output = try_vec_with_capacity::<u64>(num_values, "Dictionary output")?;
201                append_items(
202                    &mut output,
203                    &indices,
204                    &items,
205                    self.num_dictionary_items,
206                    |value| value,
207                )?;
208                LanceBuffer::reinterpret_vec(output)
209            }
210            _ => unreachable!("dictionary width was validated at construction"),
211        };
212        Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
213            bits_per_value: self.bits_per_value,
214            data: output,
215            num_values,
216            block_info: BlockInfo::default(),
217        }))
218    }
219
220    fn requires_payload(&self) -> bool {
221        self.indices.requires_payload() || self.items.requires_payload()
222    }
223}
224
225fn validate_dictionary_header(bits_per_value: u64, num_dictionary_items: usize) -> Result<()> {
226    if !matches!(bits_per_value, 32 | 64) {
227        return Err(Error::invalid_input(format!(
228            "Dictionary only supports 32 or 64-bit values, got {bits_per_value}"
229        )));
230    }
231    if !(1..=MAX_BLOCK_DICTIONARY_ITEMS).contains(&num_dictionary_items) {
232        return Err(Error::invalid_input(format!(
233            "Dictionary item count {num_dictionary_items} is outside 1..={MAX_BLOCK_DICTIONARY_ITEMS}"
234        )));
235    }
236    Ok(())
237}
238
239fn split_payload(
240    data: Option<LanceBuffer>,
241    indices_have_payload: bool,
242    items_have_payload: bool,
243    has_payload: bool,
244) -> Result<(Option<LanceBuffer>, Option<LanceBuffer>)> {
245    let Some(data) = data else {
246        if has_payload {
247            return Err(Error::invalid_input("Dictionary requires one payload"));
248        }
249        return Ok((None, None));
250    };
251    if !has_payload {
252        return Err(Error::invalid_input(
253            "Metadata-only Dictionary expects no payload",
254        ));
255    }
256    if data.len() < BLOCK_FRAME_BYTES {
257        return Err(Error::invalid_input(format!(
258            "Dictionary payload has {} bytes, shorter than its {BLOCK_FRAME_BYTES}-byte header",
259            data.len()
260        )));
261    }
262
263    let indices_size = u64::from_le_bytes(data[..8].try_into().unwrap());
264    let items_size = u64::from_le_bytes(data[8..16].try_into().unwrap());
265    let indices_size = usize::try_from(indices_size).map_err(|_| {
266        Error::invalid_input("Dictionary indices payload length does not fit usize")
267    })?;
268    let items_size = usize::try_from(items_size)
269        .map_err(|_| Error::invalid_input("Dictionary items payload length does not fit usize"))?;
270    let items_start = BLOCK_FRAME_BYTES
271        .checked_add(indices_size)
272        .ok_or_else(|| Error::invalid_input("Dictionary indices payload end overflows"))?;
273    let end = items_start
274        .checked_add(items_size)
275        .ok_or_else(|| Error::invalid_input("Dictionary items payload end overflows"))?;
276    if end != data.len() {
277        return Err(Error::invalid_input(format!(
278            "Dictionary framing describes {end} bytes, payload has {}",
279            data.len()
280        )));
281    }
282    if !indices_have_payload && indices_size != 0 {
283        return Err(Error::invalid_input(format!(
284            "Metadata-only Dictionary indices child has {indices_size} framed payload bytes"
285        )));
286    }
287    if !items_have_payload && items_size != 0 {
288        return Err(Error::invalid_input(format!(
289            "Metadata-only Dictionary items child has {items_size} framed payload bytes"
290        )));
291    }
292    Ok((
293        indices_have_payload.then(|| data.slice_with_length(BLOCK_FRAME_BYTES, indices_size)),
294        items_have_payload.then(|| data.slice_with_length(items_start, items_size)),
295    ))
296}
297
298fn append_items<T>(
299    output: &mut Vec<T>,
300    indices: &[u32],
301    items: &[u64],
302    num_dictionary_items: u32,
303    convert: impl Fn(u64) -> T,
304) -> Result<()> {
305    for (position, index) in indices.iter().enumerate() {
306        let value = *items.get(*index as usize).ok_or_else(|| {
307            Error::invalid_input(format!(
308                "Dictionary index {index} at position {position} is out of bounds for {num_dictionary_items} items"
309            ))
310        })?;
311        output.push(convert(value));
312    }
313    Ok(())
314}
315
316fn read_unsigned_values(data: &FixedWidthDataBlock, label: &str) -> Result<Vec<u64>> {
317    match data.bits_per_value {
318        32 => Ok(checked_fixed_values::<u32>(data, label)?
319            .iter()
320            .map(|value| u64::from(*value))
321            .collect()),
322        64 => Ok(checked_fixed_values::<u64>(data, label)?.to_vec()),
323        bits_per_value => Err(Error::invalid_input(format!(
324            "{label} uses unsupported {bits_per_value}-bit values"
325        ))),
326    }
327}
328
329fn fixed_from_u64_values(
330    values: &[u64],
331    bits_per_value: u64,
332    label: &str,
333) -> Result<FixedWidthDataBlock> {
334    let data = match bits_per_value {
335        32 => LanceBuffer::reinterpret_vec(
336            values
337                .iter()
338                .map(|value| {
339                    u32::try_from(*value).map_err(|_| {
340                        Error::invalid_input(format!("{label} value {value} exceeds u32::MAX"))
341                    })
342                })
343                .collect::<Result<Vec<_>>>()?,
344        ),
345        64 => LanceBuffer::reinterpret_vec(values.to_vec()),
346        _ => unreachable!("dictionary width was validated at construction"),
347    };
348    Ok(FixedWidthDataBlock {
349        bits_per_value,
350        data,
351        num_values: values.len() as u64,
352        block_info: BlockInfo::default(),
353    })
354}
355
356fn try_frame(indices_payload_bytes: usize, items_payload_bytes: usize) -> Result<Vec<u8>> {
357    let capacity = BLOCK_FRAME_BYTES
358        .checked_add(indices_payload_bytes)
359        .and_then(|capacity| capacity.checked_add(items_payload_bytes))
360        .ok_or_else(|| Error::invalid_input("Dictionary frame length overflows usize"))?;
361    let mut output = Vec::new();
362    output.try_reserve_exact(capacity).map_err(|error| {
363        Error::invalid_input(format!(
364            "Dictionary could not reserve {capacity} frame bytes: {error}"
365        ))
366    })?;
367    Ok(output)
368}
369
370#[cfg(test)]
371mod tests {
372    use crate::encodings::physical::{
373        range::{RangeDecompressor, RangeEncoder},
374        value::{ValueDecompressor, ValueEncoder},
375    };
376    use crate::format::metadata::Flat;
377
378    use super::*;
379
380    fn input_u64(values: Vec<u64>) -> DataBlock {
381        DataBlock::FixedWidth(FixedWidthDataBlock {
382            num_values: values.len() as u64,
383            bits_per_value: 64,
384            data: LanceBuffer::reinterpret_vec(values),
385            block_info: BlockInfo::default(),
386        })
387    }
388
389    fn flat_decoder(bits_per_value: u64) -> Box<dyn BlockDecompressor> {
390        Box::new(ValueDecompressor::from_flat(&Flat {
391            bits_per_value,
392            data: None,
393        }))
394    }
395
396    fn decoded_u64(block: DataBlock) -> Vec<u64> {
397        block
398            .as_fixed_width()
399            .unwrap()
400            .data
401            .borrow_to_typed_slice::<u64>()
402            .to_vec()
403    }
404
405    #[test]
406    fn dictionary_round_trip_with_payload_children() {
407        let items: Arc<[u64]> = Arc::from([10, 20]);
408        let encoder = BlockDictionaryEncoder::try_new(
409            64,
410            items,
411            Box::new(ValueEncoder::default()),
412            Box::new(ValueEncoder::default()),
413        )
414        .unwrap();
415        let (payload, encoding) = encoder
416            .compress(input_u64(vec![10, 20, 10, 20, 20]))
417            .unwrap();
418        let payload = payload.unwrap();
419        assert_eq!(u64::from_le_bytes(payload[..8].try_into().unwrap()), 20);
420        assert_eq!(u64::from_le_bytes(payload[8..16].try_into().unwrap()), 16);
421        assert!(matches!(
422            encoding.compression,
423            Some(crate::format::metadata::compressive_encoding::Compression::Dictionary(_))
424        ));
425
426        let decoder =
427            BlockDictionaryDecompressor::try_new(64, 2, flat_decoder(32), flat_decoder(64))
428                .unwrap();
429        assert_eq!(
430            decoded_u64(decoder.decompress(Some(payload), 5).unwrap()),
431            vec![10, 20, 10, 20, 20]
432        );
433    }
434
435    #[test]
436    fn dictionary_round_trip_without_payload() {
437        let encoder = BlockDictionaryEncoder::try_new(
438            64,
439            Arc::from([10, 20, 30]),
440            Box::new(RangeEncoder::new(32, 0, 1)),
441            Box::new(RangeEncoder::new(64, 10, 10)),
442        )
443        .unwrap();
444        let (payload, _) = encoder.compress(input_u64(vec![10, 20, 30])).unwrap();
445        assert!(payload.is_none());
446
447        let decoder = BlockDictionaryDecompressor::try_new(
448            64,
449            3,
450            Box::new(RangeDecompressor::new(32, 0, 1)),
451            Box::new(RangeDecompressor::new(64, 10, 10)),
452        )
453        .unwrap();
454        assert!(!decoder.requires_payload());
455        assert_eq!(
456            decoded_u64(decoder.decompress(None, 3).unwrap()),
457            vec![10, 20, 30]
458        );
459    }
460
461    #[test]
462    fn dictionary_frames_a_mixed_payload() {
463        let encoder = BlockDictionaryEncoder::try_new(
464            64,
465            Arc::from([10, 20]),
466            Box::new(ValueEncoder::default()),
467            Box::new(RangeEncoder::new(64, 10, 10)),
468        )
469        .unwrap();
470        let (payload, _) = encoder.compress(input_u64(vec![10, 20, 10])).unwrap();
471        let payload = payload.unwrap();
472        assert_eq!(u64::from_le_bytes(payload[..8].try_into().unwrap()), 12);
473        assert_eq!(u64::from_le_bytes(payload[8..16].try_into().unwrap()), 0);
474
475        let decoder = BlockDictionaryDecompressor::try_new(
476            64,
477            2,
478            flat_decoder(32),
479            Box::new(RangeDecompressor::new(64, 10, 10)),
480        )
481        .unwrap();
482        assert_eq!(
483            decoded_u64(decoder.decompress(Some(payload), 3).unwrap()),
484            vec![10, 20, 10]
485        );
486    }
487
488    #[test]
489    fn dictionary_rejects_out_of_bounds_index() {
490        let indices = LanceBuffer::reinterpret_vec(vec![0_u32, 2]);
491        let items = LanceBuffer::reinterpret_vec(vec![9_u64]);
492        let mut frame = try_frame(indices.len(), items.len()).unwrap();
493        frame.extend_from_slice(&(indices.len() as u64).to_le_bytes());
494        frame.extend_from_slice(&(items.len() as u64).to_le_bytes());
495        frame.extend_from_slice(&indices);
496        frame.extend_from_slice(&items);
497
498        let decoder =
499            BlockDictionaryDecompressor::try_new(64, 1, flat_decoder(32), flat_decoder(64))
500                .unwrap();
501        let error = decoder
502            .decompress(Some(LanceBuffer::from(frame)), 2)
503            .unwrap_err();
504        assert!(error.to_string().contains("out of bounds"));
505    }
506
507    #[test]
508    fn dictionary_rejects_invalid_frame_lengths() {
509        let mut frame = Vec::from(16_u64.to_le_bytes());
510        frame.extend_from_slice(&0_u64.to_le_bytes());
511
512        let decoder =
513            BlockDictionaryDecompressor::try_new(64, 1, flat_decoder(32), flat_decoder(64))
514                .unwrap();
515        let error = decoder
516            .decompress(Some(LanceBuffer::from(frame)), 1)
517            .unwrap_err();
518        assert!(error.to_string().contains("framing describes"));
519    }
520}
lance_encoding/encodings/physical/dictionary.rs

lance_encoding/encodings/physical/
dictionary.rs
