parquet_index/
parquet-index.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//! Binary that prints the [page index] of a parquet file
19//!
20//! # Install
21//!
22//! `parquet-layout` can be installed using `cargo`:
23//! ```
24//! cargo install parquet --features=cli
25//! ```
26//! After this `parquet-index` should be available:
27//! ```
28//! parquet-index XYZ.parquet COLUMN_PATH
29//! ```
30//!
31//! The binary can also be built from the source code and run as follows:
32//! ```
33//! cargo run --features=cli --bin parquet-index XYZ.parquet COLUMN_PATH
34//! ```
35//!
36//! [page index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
37
38mod cli;
39
40use clap::Parser;
41use parquet::data_type::ByteArray;
42use parquet::errors::{ParquetError, Result};
43use parquet::file::page_index::column_index::{
44    ByteArrayColumnIndex, ColumnIndexMetaData, PrimitiveColumnIndex,
45};
46use parquet::file::page_index::offset_index::{OffsetIndexMetaData, PageLocation};
47use parquet::file::reader::{FileReader, SerializedFileReader};
48use parquet::file::serialized_reader::ReadOptionsBuilder;
49use std::fs::File;
50
51#[derive(Debug, Parser)]
52#[clap(author, version, about("Prints the page index of a parquet file"), long_about = None)]
53struct Args {
54    #[clap(help("Path to a parquet file"))]
55    file: String,
56
57    #[clap(help(
58        "Dot-separated column path to get index for. Literal dots can be escaped as `\\.`"
59    ))]
60    column: String,
61}
62
63impl Args {
64    fn run(&self) -> Result<()> {
65        let file = File::open(&self.file)?;
66        let options = ReadOptionsBuilder::new().with_page_index().build();
67        let reader = SerializedFileReader::new_with_options(file, options)?;
68
69        let schema = reader.metadata().file_metadata().schema_descr();
70        let column_path = cli::parse_column_path(&self.column).map_err(ParquetError::General)?;
71        let column_idx = schema
72            .columns()
73            .iter()
74            .position(|x| x.path().parts() == column_path.as_slice())
75            .ok_or_else(|| {
76                ParquetError::General(format!("Failed to find column {}", self.column))
77            })?;
78
79        let page_index = reader
80            .metadata()
81            .page_index()
82            .ok_or_else(|| ParquetError::General("Page index not found".to_string()))?;
83
84        // Iterate through each row group
85        for row_group_idx in 0..reader.metadata().num_row_groups() {
86            let row_group = reader.metadata().row_group(row_group_idx);
87            println!("Row Group: {row_group_idx}");
88            let offset_index = page_index
89                .offset_index(row_group_idx, column_idx)
90                .ok_or_else(|| {
91                    ParquetError::General(format!(
92                        "No offset index for row group {row_group_idx} column chunk {column_idx}"
93                    ))
94                })?;
95
96            let row_counts =
97                compute_row_counts(offset_index.page_locations(), row_group.num_rows());
98            match page_index.column_index(row_group_idx, column_idx) {
99                None => println!("NO INDEX"),
100                Some(ColumnIndexMetaData::BOOLEAN(v)) => {
101                    print_index::<bool>(v, offset_index, &row_counts)?
102                }
103                Some(ColumnIndexMetaData::INT32(v)) => print_index(v, offset_index, &row_counts)?,
104                Some(ColumnIndexMetaData::INT64(v)) => print_index(v, offset_index, &row_counts)?,
105                Some(ColumnIndexMetaData::INT96(v)) => print_index(v, offset_index, &row_counts)?,
106                Some(ColumnIndexMetaData::FLOAT(v)) => print_index(v, offset_index, &row_counts)?,
107                Some(ColumnIndexMetaData::DOUBLE(v)) => print_index(v, offset_index, &row_counts)?,
108                Some(ColumnIndexMetaData::BYTE_ARRAY(v)) => {
109                    print_bytes_index(v, offset_index, &row_counts)?
110                }
111                Some(ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(v)) => {
112                    print_bytes_index(v, offset_index, &row_counts)?
113                }
114            }
115        }
116        Ok(())
117    }
118}
119
120/// Computes the number of rows in each page within a column chunk
121fn compute_row_counts(offset_index: &[PageLocation], rows: i64) -> Vec<i64> {
122    if offset_index.is_empty() {
123        return vec![];
124    }
125
126    let mut last = offset_index[0].first_row_index;
127    let mut out = Vec::with_capacity(offset_index.len());
128    for o in offset_index.iter().skip(1) {
129        out.push(o.first_row_index - last);
130        last = o.first_row_index;
131    }
132    out.push(rows - last);
133    out
134}
135
136/// Prints index information for a single column chunk
137fn print_index<T: std::fmt::Display>(
138    column_index: &PrimitiveColumnIndex<T>,
139    offset_index: &OffsetIndexMetaData,
140    row_counts: &[i64],
141) -> Result<()> {
142    if column_index.num_pages() as usize != offset_index.page_locations().len() {
143        return Err(ParquetError::General(format!(
144            "Index length mismatch, got {} and {}",
145            column_index.num_pages(),
146            offset_index.page_locations().len()
147        )));
148    }
149
150    for (idx, (((min, max), o), row_count)) in column_index
151        .min_values_iter()
152        .zip(column_index.max_values_iter())
153        .zip(offset_index.page_locations())
154        .zip(row_counts)
155        .enumerate()
156    {
157        print!(
158            "Page {:>5} at offset {:>10} with length {:>10} and row count {:>10}",
159            idx, o.offset, o.compressed_page_size, row_count
160        );
161        match min {
162            Some(m) => print!(", min {m:>10}"),
163            None => print!(", min {:>10}", "NONE"),
164        }
165
166        match max {
167            Some(m) => print!(", max {m:>10}"),
168            None => print!(", max {:>10}", "NONE"),
169        }
170        println!()
171    }
172
173    Ok(())
174}
175
176fn print_bytes_index(
177    column_index: &ByteArrayColumnIndex,
178    offset_index: &OffsetIndexMetaData,
179    row_counts: &[i64],
180) -> Result<()> {
181    if column_index.num_pages() as usize != offset_index.page_locations().len() {
182        return Err(ParquetError::General(format!(
183            "Index length mismatch, got {} and {}",
184            column_index.num_pages(),
185            offset_index.page_locations().len()
186        )));
187    }
188
189    for (idx, (((min, max), o), row_count)) in column_index
190        .min_values_iter()
191        .zip(column_index.max_values_iter())
192        .zip(offset_index.page_locations())
193        .zip(row_counts)
194        .enumerate()
195    {
196        print!(
197            "Page {:>5} at offset {:>10} with length {:>10} and row count {:>10}",
198            idx, o.offset, o.compressed_page_size, row_count
199        );
200        match min {
201            Some(m) => match String::from_utf8(m.to_vec()) {
202                Ok(s) => print!(", min {s:>10}"),
203                Err(_) => print!(", min {:>10}", ByteArray::from(m)),
204            },
205            None => print!(", min {:>10}", "NONE"),
206        }
207
208        match max {
209            Some(m) => match String::from_utf8(m.to_vec()) {
210                Ok(s) => print!(", max {s:>10}"),
211                Err(_) => print!(", min {:>10}", ByteArray::from(m)),
212            },
213            None => print!(", max {:>10}", "NONE"),
214        }
215        println!()
216    }
217
218    Ok(())
219}
220
221fn main() -> Result<()> {
222    Args::parse().run()
223}
parquet_index/parquet-index.rs

parquet_index/
parquet-index.rs
