parquet/file/metadata/
dictionary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Decoding a column chunk's dictionary page directly into an Arrow array,
6//! independent of the row-by-row [`ArrayReader`] machinery.
7//!
8//! This is useful for callers that want the *set* of distinct values stored
9//! in a dictionary-encoded column chunk without reading any data pages, e.g.
10//! to prune a row group when the query predicate's literals are known not to
11//! be in the dictionary.
12//!
13//! [`ArrayReader`]: crate::arrow::array_reader::ArrayReader
14
15use crate::arrow::{ByteArrayDecoderPlain, OffsetBuffer};
16use crate::basic::{Encoding, PageType, Type as PhysicalType};
17use crate::column::page::Page;
18use crate::compression::{CodecOptions, create_codec};
19#[cfg(feature = "encryption")]
20use crate::encryption::decrypt::CryptoContext;
21use crate::errors::{ParquetError, Result};
22#[cfg(feature = "encryption")]
23use crate::file::metadata::ColumnChunkMetaData;
24use crate::file::metadata::ParquetMetaData;
25use crate::file::serialized_reader::{
26    SerializedPageReaderContext, decode_page, read_page_header_len_from_bytes, verify_page_size,
27};
28use arrow_array::ArrayRef;
29use arrow_schema::DataType as ArrowType;
30use bytes::Bytes;
31#[cfg(feature = "encryption")]
32use std::sync::Arc;
33
34/// Decodes the dictionary page of a column chunk into an [`ArrayRef`].
35///
36/// `buffer` must contain the entire dictionary page, byte-for-byte, i.e. the
37/// range `[dictionary_page_offset, data_page_offset)` of the column chunk.
38///
39/// Only `BYTE_ARRAY` columns are currently supported; other physical types
40/// return an error. The returned array never contains nulls: dictionary
41/// pages only store the distinct non-null values, with nulls represented via
42/// definition levels in the data pages.
43/// The returned array has `Binary` values, regardless of the column's logical type.
44///
45/// Note this only decodes whatever dictionary page is present -- it does
46/// **not** verify that the entire column chunk is dictionary-encoded (i.e.
47/// that every value in the chunk is drawn from this dictionary). Callers
48/// that need that guarantee (for example, to use the dictionary as an exact
49/// membership index) must check that themselves, e.g. via
50/// [`crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask`].
51pub(crate) fn decode_dictionary_page(
52    buffer: Bytes,
53    parquet_meta_data: &ParquetMetaData,
54    row_group_idx: usize,
55    column_idx: usize,
56) -> Result<ArrayRef> {
57    let column_metadata = parquet_meta_data
58        .row_group(row_group_idx)
59        .column(column_idx);
60    let column_descriptor = column_metadata.column_descr();
61
62    if column_descriptor.physical_type() != PhysicalType::BYTE_ARRAY {
63        return Err(ParquetError::General(format!(
64            "decode_dictionary_page only supports BYTE_ARRAY columns, got {}",
65            column_descriptor.physical_type()
66        )));
67    }
68
69    // Dictionary pages are subject to the same modular encryption as data
70    // pages: both the page header and the page body may be ciphertext, so
71    // we must route through the same crypto-aware header/data path that
72    // `SerializedPageReader` uses rather than parsing the header directly.
73    let page_context = SerializedPageReaderContext {
74        read_stats: true,
75        #[cfg(feature = "encryption")]
76        crypto_context: dictionary_page_crypto_context(
77            parquet_meta_data,
78            column_metadata,
79            row_group_idx,
80            column_idx,
81        )?,
82    };
83
84    let (consumed, header) =
85        read_page_header_len_from_bytes(&page_context, buffer.as_ref(), 0, true)?;
86    if header.r#type != PageType::DICTIONARY_PAGE {
87        return Err(ParquetError::General(format!(
88            "Expected a dictionary page, found {:?}",
89            header.r#type
90        )));
91    }
92
93    // `compressed_page_size` comes from the (possibly maliciously crafted)
94    // file header; `verify_page_size` bounds-checks it against what we
95    // actually fetched before we slice, instead of trusting it blindly.
96    let remaining = (buffer.len() - consumed) as u64;
97    verify_page_size(
98        header.compressed_page_size,
99        header.uncompressed_page_size,
100        remaining,
101    )?;
102    let compressed_size = header.compressed_page_size as usize;
103    let page_buf = buffer.slice(consumed..consumed + compressed_size);
104    let page_buf = page_context.decrypt_page_data(page_buf, 0, true)?;
105
106    let mut decompressor = create_codec(column_metadata.compression(), &CodecOptions::default())?;
107    let page = decode_page(
108        header,
109        page_buf,
110        column_descriptor.physical_type(),
111        decompressor.as_mut(),
112    )?;
113    let Page::DictionaryPage {
114        buf,
115        num_values,
116        encoding,
117        ..
118    } = page
119    else {
120        return Err(ParquetError::General(
121            "Expected a dictionary page".to_string(),
122        ));
123    };
124    let num_values = num_values as usize;
125
126    // The dictionary page is always PLAIN-encoded, regardless of what the
127    // data pages' encoding is (RLE_DICTIONARY/PLAIN_DICTIONARY only describe
128    // how *data* pages reference the dictionary by index).
129    if encoding != Encoding::PLAIN {
130        return Err(ParquetError::General(format!(
131            "Dictionary page encoding must be PLAIN, got {encoding:?}"
132        )));
133    }
134    let mut decoder = ByteArrayDecoderPlain::new(buf, num_values, Some(num_values), false);
135    let mut offsets = OffsetBuffer::<i32>::with_capacity(num_values);
136    decoder.read(&mut offsets, usize::MAX)?;
137    if offsets.len() != num_values {
138        return Err(ParquetError::General(format!(
139            "Expected {num_values} dictionary values, decoded {}",
140            offsets.len()
141        )));
142    }
143
144    Ok(offsets.into_array(None, ArrowType::Binary))
145}
146
147/// Builds the crypto context needed to decrypt the dictionary page of
148/// `column_metadata`, or `None` if the file (or this column) isn't encrypted.
149#[cfg(feature = "encryption")]
150fn dictionary_page_crypto_context(
151    parquet_meta_data: &ParquetMetaData,
152    column_metadata: &ColumnChunkMetaData,
153    row_group_idx: usize,
154    column_idx: usize,
155) -> Result<Option<Arc<CryptoContext>>> {
156    let Some(file_decryptor) = parquet_meta_data.file_decryptor() else {
157        return Ok(None);
158    };
159    let Some(crypto_metadata) = column_metadata.crypto_metadata() else {
160        return Ok(None);
161    };
162    let ordinal = parquet_meta_data
163        .row_group(row_group_idx)
164        .ordinal()
165        .ok_or_else(|| {
166            ParquetError::General("Encrypted row group is missing its file ordinal".to_string())
167        })?;
168    let ordinal = usize::try_from(ordinal).map_err(|_| {
169        ParquetError::General("Encrypted row group has an invalid file ordinal".to_string())
170    })?;
171    let crypto_context =
172        CryptoContext::for_column(file_decryptor, crypto_metadata, ordinal, column_idx)?
173            .for_dictionary_page();
174    Ok(Some(Arc::new(crypto_context)))
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use crate::arrow::ArrowWriter;
181    use crate::basic::Encoding;
182    use crate::file::metadata::ParquetMetaDataReader;
183    use crate::file::properties::WriterProperties;
184    use crate::file::reader::{ChunkReader, FileReader, SerializedFileReader};
185    use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift};
186    use arrow_array::{Array, BinaryArray, RecordBatch, StringArray};
187    use arrow_schema::{Field, Schema};
188    use std::sync::Arc;
189
190    fn write_dictionary_encoded_strings(values: &[&str]) -> Bytes {
191        let schema = Arc::new(Schema::new(vec![Field::new("s", ArrowType::Utf8, false)]));
192        let array = Arc::new(StringArray::from_iter_values(values.iter().copied()));
193        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
194
195        let props = WriterProperties::builder()
196            .set_dictionary_enabled(true)
197            .build();
198        let mut buf = Vec::new();
199        {
200            let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
201            writer.write(&batch).unwrap();
202            writer.close().unwrap();
203        }
204        Bytes::from(buf)
205    }
206
207    #[test]
208    fn decode_dictionary_page_round_trips_strings() {
209        let distinct_values = ["alpha", "beta", "gamma"];
210        // Repeat so the column is worth dictionary-encoding but the
211        // dictionary itself only contains the distinct values.
212        let values: Vec<&str> = distinct_values.iter().copied().cycle().take(30).collect();
213        let data = write_dictionary_encoded_strings(&values);
214
215        let reader = SerializedFileReader::new(data.clone()).unwrap();
216        let metadata = reader.metadata();
217        let column_metadata = metadata.row_group(0).column(0);
218
219        assert!(
220            column_metadata.dictionary_page_offset().is_some(),
221            "expected the column chunk to be dictionary-encoded"
222        );
223
224        let start = column_metadata.dictionary_page_offset().unwrap() as u64;
225        let end = column_metadata.data_page_offset() as u64;
226        let buffer = data.get_bytes(start, (end - start) as usize).unwrap();
227
228        let array = decode_dictionary_page(buffer, metadata, 0, 0).unwrap();
229        let array = array.as_any().downcast_ref::<BinaryArray>().unwrap();
230        let decoded: Vec<&[u8]> = array.iter().map(|v| v.unwrap()).collect();
231        assert_eq!(decoded, distinct_values.map(str::as_bytes));
232    }
233
234    #[test]
235    fn decode_dictionary_page_errors_on_truncated_buffer() {
236        let distinct_values = ["alpha", "beta", "gamma"];
237        let values: Vec<&str> = distinct_values.iter().copied().cycle().take(30).collect();
238        let data = write_dictionary_encoded_strings(&values);
239
240        let reader = SerializedFileReader::new(data.clone()).unwrap();
241        let metadata = reader.metadata();
242        let column_metadata = metadata.row_group(0).column(0);
243
244        let start = column_metadata.dictionary_page_offset().unwrap() as u64;
245        let end = column_metadata.data_page_offset() as u64;
246        let buffer = data.get_bytes(start, (end - start) as usize).unwrap();
247
248        // Simulate a truncated/malformed file: the page header's declared
249        // `compressed_page_size` no longer fits in what was actually
250        // fetched. This must return an error rather than panic while
251        // slicing (`Bytes::slice` panics on out-of-bounds ranges).
252        let truncated = buffer.slice(..buffer.len() - 1);
253        let err = decode_dictionary_page(truncated, metadata, 0, 0).unwrap_err();
254        assert!(
255            matches!(err, ParquetError::EOF(_)),
256            "unexpected error: {err}"
257        );
258    }
259
260    fn dictionary_page_with_header_change(
261        data: &Bytes,
262        change: impl FnOnce(&mut crate::file::metadata::thrift::PageHeader),
263    ) -> (Bytes, ParquetMetaData) {
264        let reader = SerializedFileReader::new(data.clone()).unwrap();
265        let metadata = reader.metadata().clone();
266        let column = metadata.row_group(0).column(0);
267        let start = column.dictionary_page_offset().unwrap() as u64;
268        let end = column.data_page_offset() as u64;
269        let buffer = data.get_bytes(start, (end - start) as usize).unwrap();
270        let context = SerializedPageReaderContext {
271            read_stats: true,
272            #[cfg(feature = "encryption")]
273            crypto_context: None,
274        };
275        let (header_len, mut header) =
276            read_page_header_len_from_bytes(&context, &buffer, 0, true).unwrap();
277        change(&mut header);
278        let mut changed = Vec::new();
279        header
280            .write_thrift(&mut ThriftCompactOutputProtocol::new(&mut changed))
281            .unwrap();
282        changed.extend_from_slice(&buffer[header_len..]);
283        (Bytes::from(changed), metadata)
284    }
285
286    #[test]
287    fn decode_dictionary_page_rejects_missing_values() {
288        let data = write_dictionary_encoded_strings(&["alpha", "beta", "alpha"]);
289        let (buffer, metadata) = dictionary_page_with_header_change(&data, |header| {
290            header.dictionary_page_header.as_mut().unwrap().num_values += 1;
291        });
292        let err = decode_dictionary_page(buffer, &metadata, 0, 0).unwrap_err();
293        assert!(err.to_string().contains("dictionary values"), "{err}");
294    }
295
296    #[test]
297    fn decode_dictionary_page_rejects_non_plain_encoding() {
298        let data = write_dictionary_encoded_strings(&["alpha", "beta", "alpha"]);
299        let (buffer, metadata) = dictionary_page_with_header_change(&data, |header| {
300            header.dictionary_page_header.as_mut().unwrap().encoding = Encoding::RLE_DICTIONARY;
301        });
302        let err = decode_dictionary_page(buffer, &metadata, 0, 0).unwrap_err();
303        assert!(err.to_string().contains("PLAIN"), "{err}");
304    }
305
306    #[test]
307    fn decode_dictionary_page_rejects_non_byte_array() {
308        let schema = Arc::new(Schema::new(vec![Field::new("i", ArrowType::Int32, false)]));
309        let array = Arc::new(arrow_array::Int32Array::from(vec![1, 2, 3]));
310        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
311        let mut buf = Vec::new();
312        {
313            let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
314            writer.write(&batch).unwrap();
315            writer.close().unwrap();
316        }
317        let data = Bytes::from(buf);
318
319        let reader = SerializedFileReader::new(data).unwrap();
320        let metadata = reader.metadata();
321
322        let err = decode_dictionary_page(Bytes::new(), metadata, 0, 0).unwrap_err();
323        assert!(err.to_string().contains("BYTE_ARRAY"));
324    }
325
326    #[test]
327    fn read_column_dictionary_round_trips_via_metadata_reader() {
328        let distinct_values = ["alpha", "beta", "gamma"];
329        let values: Vec<&str> = distinct_values.iter().copied().cycle().take(30).collect();
330        let data = write_dictionary_encoded_strings(&values);
331
332        let reader = SerializedFileReader::new(data.clone()).unwrap();
333        let metadata = reader.metadata();
334
335        let array = ParquetMetaDataReader::read_column_dictionary(&data, metadata, 0, 0)
336            .unwrap()
337            .unwrap();
338        let array = array.as_any().downcast_ref::<BinaryArray>().unwrap();
339        let decoded: Vec<&[u8]> = array.iter().map(|v| v.unwrap()).collect();
340        assert_eq!(decoded, distinct_values.map(str::as_bytes));
341    }
342
343    #[cfg(feature = "encryption")]
344    #[test]
345    fn read_column_dictionary_round_trips_with_encryption() {
346        use crate::encryption::decrypt::FileDecryptionProperties;
347        use crate::encryption::encrypt::FileEncryptionProperties;
348        const FOOTER_KEY: &[u8] = b"0123456789012345";
349
350        let distinct_values = ["alpha", "beta", "gamma"];
351        let values: Vec<&str> = distinct_values.iter().copied().cycle().take(30).collect();
352
353        let schema = Arc::new(Schema::new(vec![Field::new("s", ArrowType::Utf8, false)]));
354        let array = Arc::new(StringArray::from_iter_values(values.iter().copied()));
355        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
356
357        let encryption_properties = FileEncryptionProperties::builder(FOOTER_KEY.to_vec())
358            .build()
359            .unwrap();
360        let props = WriterProperties::builder()
361            .set_dictionary_enabled(true)
362            .with_file_encryption_properties(encryption_properties)
363            .build();
364        let mut buf = Vec::new();
365        {
366            let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
367            writer.write(&batch).unwrap();
368            writer.close().unwrap();
369        }
370        let data = Bytes::from(buf);
371
372        let decryption_properties = FileDecryptionProperties::builder(FOOTER_KEY.to_vec())
373            .build()
374            .unwrap();
375        let metadata = ParquetMetaDataReader::new()
376            .with_decryption_properties(Some(decryption_properties))
377            .parse_and_finish(&data)
378            .unwrap();
379
380        let array = ParquetMetaDataReader::read_column_dictionary(&data, &metadata, 0, 0)
381            .unwrap()
382            .unwrap();
383        let array = array.as_any().downcast_ref::<BinaryArray>().unwrap();
384        let decoded: Vec<&[u8]> = array.iter().map(|v| v.unwrap()).collect();
385        assert_eq!(decoded, distinct_values.map(str::as_bytes));
386    }
387
388    #[cfg(feature = "encryption")]
389    #[test]
390    fn read_column_dictionary_uses_file_ordinal_after_filtering() {
391        use crate::encryption::decrypt::FileDecryptionProperties;
392        use crate::encryption::encrypt::FileEncryptionProperties;
393        use crate::file::metadata::ParquetMetaDataBuilder;
394
395        const FOOTER_KEY: &[u8] = b"0123456789012345";
396        let schema = Arc::new(Schema::new(vec![Field::new("s", ArrowType::Utf8, false)]));
397        let props = WriterProperties::builder()
398            .set_dictionary_enabled(true)
399            .with_file_encryption_properties(
400                FileEncryptionProperties::builder(FOOTER_KEY.to_vec())
401                    .build()
402                    .unwrap(),
403            )
404            .build();
405        let mut buf = Vec::new();
406        {
407            let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap();
408            for value in ["first", "second"] {
409                let array = Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
410                    value, 30,
411                )));
412                let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
413                writer.write(&batch).unwrap();
414                writer.flush().unwrap();
415            }
416            writer.close().unwrap();
417        }
418        let data = Bytes::from(buf);
419        let metadata = ParquetMetaDataReader::new()
420            .with_decryption_properties(Some(
421                FileDecryptionProperties::builder(FOOTER_KEY.to_vec())
422                    .build()
423                    .unwrap(),
424            ))
425            .parse_and_finish(&data)
426            .unwrap();
427        assert_eq!(metadata.row_group(1).ordinal(), Some(1));
428        let second = metadata.row_group(1).clone();
429        let filtered = ParquetMetaDataBuilder::new_from_metadata(metadata)
430            .set_row_groups(vec![second])
431            .build();
432        let array = ParquetMetaDataReader::read_column_dictionary(&data, &filtered, 0, 0)
433            .unwrap()
434            .unwrap();
435        let array = array.as_any().downcast_ref::<BinaryArray>().unwrap();
436        assert_eq!(array.value(0), b"second");
437
438        let invalid = filtered
439            .row_group(0)
440            .clone()
441            .into_builder()
442            .set_ordinal(-1)
443            .build()
444            .unwrap();
445        let invalid_metadata = ParquetMetaDataBuilder::new_from_metadata(filtered)
446            .set_row_groups(vec![invalid])
447            .build();
448        let err = ParquetMetaDataReader::read_column_dictionary(&data, &invalid_metadata, 0, 0)
449            .unwrap_err();
450        assert!(err.to_string().contains("invalid file ordinal"), "{err}");
451    }
452
453    #[test]
454    fn read_column_dictionary_returns_none_without_dictionary_page() {
455        use crate::file::metadata::ParquetMetaDataReader;
456
457        let schema = Arc::new(Schema::new(vec![Field::new("s", ArrowType::Utf8, false)]));
458        let array = Arc::new(StringArray::from_iter_values(["a", "b", "c"]));
459        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
460
461        let props = WriterProperties::builder()
462            .set_dictionary_enabled(false)
463            .build();
464        let mut buf = Vec::new();
465        {
466            let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
467            writer.write(&batch).unwrap();
468            writer.close().unwrap();
469        }
470        let data = Bytes::from(buf);
471
472        let reader = SerializedFileReader::new(data.clone()).unwrap();
473        let metadata = reader.metadata();
474        assert!(
475            metadata
476                .row_group(0)
477                .column(0)
478                .dictionary_page_offset()
479                .is_none()
480        );
481
482        let result = ParquetMetaDataReader::read_column_dictionary(&data, metadata, 0, 0).unwrap();
483        assert!(result.is_none());
484    }
485}
parquet/file/metadata/dictionary.rs

parquet/file/metadata/
dictionary.rs
