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