Skip to main content

parquet/file/page_index/
index_reader.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//! Support for reading [`ColumnIndexMetaData`] and [`OffsetIndexMetaData`] from parquet metadata.
19
20use crate::basic::{BoundaryOrder, Type};
21use crate::data_type::Int96;
22use crate::errors::{ParquetError, Result};
23use crate::file::page_index::column_index::{
24    ByteArrayColumnIndex, ColumnIndexMetaData, PrimitiveColumnIndex,
25};
26use crate::file::page_index::offset_index::OffsetIndexMetaData;
27use crate::parquet_thrift::{
28    ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
29    ThriftSliceInputProtocol, WriteThrift, WriteThriftField, read_thrift_vec,
30};
31use crate::thrift_struct;
32use std::io::Write;
33
34/// Decode a Thrift [`OffsetIndex`] from the provided bytes.
35///
36/// The passed in bytes contain a serialized Thrift `OffsetIndex` struct as
37/// read from a Parquet file.
38///
39/// Returns an [`OffsetIndexMetaData`] containing page location information.
40///
41/// # Example
42///
43/// ```
44/// # use parquet::file::reader::{FileReader, SerializedFileReader};
45/// # use parquet::file::page_index::index_reader::decode_offset_index;
46/// # use std::fs::File;
47/// # use std::io::{Read, Seek};
48/// # use parquet::errors::Result;
49/// #
50/// # fn read_offset_index() -> Result<()> {
51/// // Open the Parquet file
52/// let mut file = File::open("data.parquet")?;
53/// let reader = SerializedFileReader::new(file.try_clone()?)?;
54/// let metadata = reader.metadata();
55///
56/// // Select a row group and column to read
57/// let row_group_idx = 0;
58/// let column_idx = 0;
59///
60/// // Get the column chunk metadata
61/// let row_group = metadata.row_group(row_group_idx);
62/// let column_chunk = row_group.column(column_idx);
63///
64/// // Get the offset index byte range from the column metadata
65/// if let Some(range) = column_chunk.offset_index_range() {
66///     // Read the offset index bytes from the file
67///     let mut buffer = vec![0u8; (range.end - range.start) as usize];
68///     file.seek(std::io::SeekFrom::Start(range.start))?;
69///     file.read_exact(&mut buffer)?;
70///
71///     // Decode the offset index
72///     let offset_index = decode_offset_index(&buffer)?;
73///
74///     // Access page location information
75///     for (i, page_location) in offset_index.page_locations().iter().enumerate() {
76///         println!("Page {}: offset={}, size={}, first_row={}",
77///             i,
78///             page_location.offset,
79///             page_location.compressed_page_size,
80///             page_location.first_row_index
81///         );
82///     }
83/// }
84/// # Ok(())
85/// # }
86/// ```
87///
88/// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/e94a5d090b324a0c0ee1adbb8ea6b099852dc3cc/src/main/thrift/parquet.thrift#L1253-L1273
89pub fn decode_offset_index(data: &[u8]) -> Result<OffsetIndexMetaData, ParquetError> {
90    let mut prot = ThriftSliceInputProtocol::new(data);
91
92    // Try to read fast-path first. If that fails, fall back to slower but more robust
93    // decoder.
94    match OffsetIndexMetaData::try_from_fast(&mut prot) {
95        Ok(offset_index) => Ok(offset_index),
96        Err(_) => {
97            prot = ThriftSliceInputProtocol::new(data);
98            OffsetIndexMetaData::read_thrift(&mut prot)
99        }
100    }
101}
102
103// private struct only used for decoding then discarded
104thrift_struct!(
105pub(super) struct ThriftColumnIndex<'a> {
106  1: required list<bool> null_pages
107  2: required list<'a><binary> min_values
108  3: required list<'a><binary> max_values
109  4: required BoundaryOrder boundary_order
110  5: optional list<i64> null_counts
111  6: optional list<i64> repetition_level_histograms;
112  7: optional list<i64> definition_level_histograms;
113  8: optional list<i64> nan_counts
114}
115);
116
117/// Decode a Thrift [`ColumnIndex`] from the provided bytes.
118///
119/// The passed in bytes contain a serialized Thrift `OffsetIndex` struct as
120/// read from a Parquet file. The `column_type` can be obtained via
121/// [`ColumnChunkMetaData::column_type`].
122///
123/// Returns a [`ColumnIndexMetaData`] containing per-page statistics.
124///
125/// # Example
126///
127/// ```
128/// # use parquet::file::reader::{FileReader, SerializedFileReader};
129/// # use parquet::file::page_index::index_reader::decode_column_index;
130/// # use std::fs::File;
131/// # use std::io::{Read, Seek};
132/// # use parquet::errors::Result;
133/// #
134/// # fn read_column_index() -> Result<()> {
135/// // Open the Parquet file
136/// let mut file = File::open("data.parquet")?;
137/// let reader = SerializedFileReader::new(file.try_clone()?)?;
138/// let metadata = reader.metadata();
139///
140/// // Select a row group and column to read
141/// let row_group_idx = 0;
142/// let column_idx = 0;
143///
144/// // Get the column chunk metadata
145/// let row_group = metadata.row_group(row_group_idx);
146/// let column_chunk = row_group.column(column_idx);
147///
148/// // Get the column index byte range from the column metadata
149/// if let Some(range) = column_chunk.column_index_range() {
150///     // Get the column type for proper deserialization
151///     let column_type = column_chunk.column_type();
152///
153///     // Read the column index bytes from the file
154///     let mut buffer = vec![0u8; (range.end - range.start) as usize];
155///     file.seek(std::io::SeekFrom::Start(range.start))?;
156///     file.read_exact(&mut buffer)?;
157///
158///     // Decode the column index
159///     let column_index = decode_column_index(&buffer, column_type)?;
160///
161///     // Access per-page statistics (example for INT32 column)
162///     use parquet::file::page_index::column_index::ColumnIndexMetaData;
163///     match column_index {
164///         ColumnIndexMetaData::INT32(index) => {
165///             for (i, (min, max)) in index.min_values().iter()
166///                 .zip(index.max_values().iter())
167///                 .enumerate() {
168///                 println!("Page {}: min={}, max={}", i, min, max);
169///             }
170///         }
171///         _ => println!("Column is not INT32 type"),
172///     }
173/// }
174/// # Ok(())
175/// # }
176/// ```
177///
178/// [`ColumnChunkMetaData::column_type`]: crate::file::metadata::ColumnChunkMetaData::column_type
179/// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/e94a5d090b324a0c0ee1adbb8ea6b099852dc3cc/src/main/thrift/parquet.thrift#L1275-1373
180pub fn decode_column_index(
181    data: &[u8],
182    column_type: Type,
183) -> Result<ColumnIndexMetaData, ParquetError> {
184    let mut prot = ThriftSliceInputProtocol::new(data);
185    let index = ThriftColumnIndex::read_thrift(&mut prot)?;
186
187    let index = match column_type {
188        Type::BOOLEAN => {
189            ColumnIndexMetaData::BOOLEAN(PrimitiveColumnIndex::<bool>::try_from_thrift(index)?)
190        }
191        Type::INT32 => {
192            ColumnIndexMetaData::INT32(PrimitiveColumnIndex::<i32>::try_from_thrift(index)?)
193        }
194        Type::INT64 => {
195            ColumnIndexMetaData::INT64(PrimitiveColumnIndex::<i64>::try_from_thrift(index)?)
196        }
197        Type::INT96 => {
198            ColumnIndexMetaData::INT96(PrimitiveColumnIndex::<Int96>::try_from_thrift(index)?)
199        }
200        Type::FLOAT => {
201            ColumnIndexMetaData::FLOAT(PrimitiveColumnIndex::<f32>::try_from_thrift(index)?)
202        }
203        Type::DOUBLE => {
204            ColumnIndexMetaData::DOUBLE(PrimitiveColumnIndex::<f64>::try_from_thrift(index)?)
205        }
206        Type::BYTE_ARRAY => {
207            ColumnIndexMetaData::BYTE_ARRAY(ByteArrayColumnIndex::try_from_thrift(index)?)
208        }
209        Type::FIXED_LEN_BYTE_ARRAY => {
210            ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(ByteArrayColumnIndex::try_from_thrift(index)?)
211        }
212    };
213
214    Ok(index)
215}