parquet/file/metadata/
dictionary.rs

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

parquet/file/metadata/
dictionary.rs
