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

parquet_index/
parquet-index.rs
