Skip to main content

parquet/file/page_index/
column_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//! [`ColumnIndexMetaData`] structures holding decoded [`ColumnIndex`] information
19//!
20//! [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
21//!
22
23use crate::{
24    data_type::{ByteArray, FixedLenByteArray},
25    errors::{ParquetError, Result},
26    parquet_thrift::{
27        ElementType, FieldType, ThriftCompactOutputProtocol, WriteThrift, WriteThriftField,
28    },
29};
30use std::ops::Deref;
31
32use crate::{
33    basic::BoundaryOrder,
34    data_type::{Int96, private::ParquetValueType},
35    file::page_index::index_reader::ThriftColumnIndex,
36};
37
38/// Common bits of the column index
39#[derive(Debug, Clone, PartialEq)]
40pub struct ColumnIndex {
41    pub(crate) null_pages: Vec<bool>,
42    pub(crate) boundary_order: BoundaryOrder,
43    pub(crate) null_counts: Option<Vec<i64>>,
44    pub(crate) repetition_level_histograms: Option<Vec<i64>>,
45    pub(crate) definition_level_histograms: Option<Vec<i64>>,
46    pub(crate) nan_counts: Option<Vec<i64>>,
47}
48
49impl ColumnIndex {
50    /// Returns the number of pages
51    pub fn num_pages(&self) -> u64 {
52        self.null_pages.len() as u64
53    }
54
55    /// Returns the number of null values in the page indexed by `idx`
56    ///
57    /// Returns `None` if no null counts have been set in the index
58    pub fn null_count(&self, idx: usize) -> Option<i64> {
59        self.null_counts.as_ref().map(|nc| nc[idx])
60    }
61
62    /// Returns the number of NaN values in the page indexed by `idx`
63    ///
64    /// Returns `None` if no NaN counts have been set in the index
65    pub fn nan_count(&self, idx: usize) -> Option<i64> {
66        self.nan_counts.as_ref().map(|nc| nc[idx])
67    }
68
69    /// Returns the repetition level histogram for the page indexed by `idx`
70    pub fn repetition_level_histogram(&self, idx: usize) -> Option<&[i64]> {
71        if let Some(rep_hists) = self.repetition_level_histograms.as_ref() {
72            let num_lvls = rep_hists.len() / self.num_pages() as usize;
73            let start = num_lvls * idx;
74            Some(&rep_hists[start..start + num_lvls])
75        } else {
76            None
77        }
78    }
79
80    /// Returns the definition level histogram for the page indexed by `idx`
81    pub fn definition_level_histogram(&self, idx: usize) -> Option<&[i64]> {
82        if let Some(def_hists) = self.definition_level_histograms.as_ref() {
83            let num_lvls = def_hists.len() / self.num_pages() as usize;
84            let start = num_lvls * idx;
85            Some(&def_hists[start..start + num_lvls])
86        } else {
87            None
88        }
89    }
90
91    /// Returns whether the page indexed by `idx` consists of all null values
92    pub fn is_null_page(&self, idx: usize) -> bool {
93        self.null_pages[idx]
94    }
95}
96
97/// Column index for primitive types
98#[derive(Debug, Clone, PartialEq)]
99pub struct PrimitiveColumnIndex<T> {
100    pub(crate) column_index: ColumnIndex,
101    pub(crate) min_values: Vec<T>,
102    pub(crate) max_values: Vec<T>,
103}
104
105impl<T: ParquetValueType> PrimitiveColumnIndex<T> {
106    #[allow(clippy::too_many_arguments)]
107    pub(crate) fn try_new(
108        null_pages: Vec<bool>,
109        boundary_order: BoundaryOrder,
110        null_counts: Option<Vec<i64>>,
111        nan_counts: Option<Vec<i64>>,
112        repetition_level_histograms: Option<Vec<i64>>,
113        definition_level_histograms: Option<Vec<i64>>,
114        min_bytes: Vec<&[u8]>,
115        max_bytes: Vec<&[u8]>,
116    ) -> Result<Self> {
117        let len = null_pages.len();
118
119        if min_bytes.len() != len || max_bytes.len() != len {
120            return Err(ParquetError::General(format!(
121                "ColumnIndex min/max length mismatch: expected {len}, got min={} max={}",
122                min_bytes.len(),
123                max_bytes.len()
124            )));
125        }
126        if let Some(ref nc) = null_counts
127            && nc.len() != len
128        {
129            return Err(ParquetError::General(format!(
130                "ColumnIndex null_counts length mismatch: expected {len}, got {}",
131                nc.len()
132            )));
133        }
134        if let Some(ref rep) = repetition_level_histograms
135            && len != 0
136            && rep.len() % len != 0
137        {
138            return Err(ParquetError::General(
139                "Invalid repetition_level_histograms length".to_string(),
140            ));
141        }
142        if let Some(ref def) = definition_level_histograms
143            && len != 0
144            && def.len() % len != 0
145        {
146            return Err(ParquetError::General(
147                "Invalid definition_level_histograms length".to_string(),
148            ));
149        }
150
151        let mut min_values = Vec::with_capacity(len);
152        let mut max_values = Vec::with_capacity(len);
153
154        for (i, is_null) in null_pages.iter().enumerate().take(len) {
155            if !is_null {
156                let min = min_bytes[i];
157                min_values.push(T::try_from_le_slice(min)?);
158
159                let max = max_bytes[i];
160                max_values.push(T::try_from_le_slice(max)?);
161            } else {
162                // need placeholders
163                min_values.push(Default::default());
164                max_values.push(Default::default());
165            }
166        }
167
168        Ok(Self {
169            column_index: ColumnIndex {
170                null_pages,
171                boundary_order,
172                null_counts,
173                repetition_level_histograms,
174                definition_level_histograms,
175                nan_counts,
176            },
177            min_values,
178            max_values,
179        })
180    }
181
182    pub(super) fn try_from_thrift(index: ThriftColumnIndex) -> Result<Self> {
183        Self::try_new(
184            index.null_pages,
185            index.boundary_order,
186            index.null_counts,
187            index.nan_counts,
188            index.repetition_level_histograms,
189            index.definition_level_histograms,
190            index.min_values,
191            index.max_values,
192        )
193    }
194}
195
196impl<T> PrimitiveColumnIndex<T> {
197    /// Returns an array containing the min values for each page.
198    ///
199    /// Values in the returned slice are only valid if [`ColumnIndex::is_null_page()`]
200    /// is `false` for the same index.
201    pub fn min_values(&self) -> &[T] {
202        &self.min_values
203    }
204
205    /// Returns an array containing the max values for each page.
206    ///
207    /// Values in the returned slice are only valid if [`ColumnIndex::is_null_page()`]
208    /// is `false` for the same index.
209    pub fn max_values(&self) -> &[T] {
210        &self.max_values
211    }
212
213    /// Returns an iterator over the min values.
214    ///
215    /// Values may be `None` when [`ColumnIndex::is_null_page()`] is `true`.
216    pub fn min_values_iter(&self) -> impl Iterator<Item = Option<&T>> {
217        self.min_values.iter().enumerate().map(|(i, min)| {
218            if self.is_null_page(i) {
219                None
220            } else {
221                Some(min)
222            }
223        })
224    }
225
226    /// Returns an iterator over the max values.
227    ///
228    /// Values may be `None` when [`ColumnIndex::is_null_page()`] is `true`.
229    pub fn max_values_iter(&self) -> impl Iterator<Item = Option<&T>> {
230        self.max_values.iter().enumerate().map(|(i, min)| {
231            if self.is_null_page(i) {
232                None
233            } else {
234                Some(min)
235            }
236        })
237    }
238
239    /// Returns the min value for the page indexed by `idx`
240    ///
241    /// It is `None` when all values are null
242    #[inline]
243    pub fn min_value(&self, idx: usize) -> Option<&T> {
244        if self.null_pages[idx] {
245            None
246        } else {
247            Some(&self.min_values[idx])
248        }
249    }
250
251    /// Returns the max value for the page indexed by `idx`
252    ///
253    /// It is `None` when all values are null
254    #[inline]
255    pub fn max_value(&self, idx: usize) -> Option<&T> {
256        if self.null_pages[idx] {
257            None
258        } else {
259            Some(&self.max_values[idx])
260        }
261    }
262}
263
264impl<T> Deref for PrimitiveColumnIndex<T> {
265    type Target = ColumnIndex;
266
267    fn deref(&self) -> &Self::Target {
268        &self.column_index
269    }
270}
271
272impl<T: ParquetValueType> WriteThrift for PrimitiveColumnIndex<T> {
273    const ELEMENT_TYPE: ElementType = ElementType::Struct;
274    fn write_thrift<W: std::io::Write>(
275        &self,
276        writer: &mut ThriftCompactOutputProtocol<W>,
277    ) -> Result<()> {
278        self.null_pages.write_thrift_field(writer, 1, 0)?;
279
280        // need to handle min/max manually
281        let len = self.null_pages.len();
282        writer.write_field_begin(FieldType::List, 2, 1)?;
283        writer.write_list_begin(ElementType::Binary, len)?;
284        for i in 0..len {
285            let min = self.min_value(i).map(|m| m.as_bytes()).unwrap_or(&[]);
286            min.write_thrift(writer)?;
287        }
288        writer.write_field_begin(FieldType::List, 3, 2)?;
289        writer.write_list_begin(ElementType::Binary, len)?;
290        for i in 0..len {
291            let max = self.max_value(i).map(|m| m.as_bytes()).unwrap_or(&[]);
292            max.write_thrift(writer)?;
293        }
294        let mut last_field_id = self.boundary_order.write_thrift_field(writer, 4, 3)?;
295        if let Some(null_counts) = &self.null_counts {
296            last_field_id = null_counts.write_thrift_field(writer, 5, last_field_id)?;
297        }
298        if let Some(repetition_level_histograms) = &self.repetition_level_histograms {
299            last_field_id =
300                repetition_level_histograms.write_thrift_field(writer, 6, last_field_id)?;
301        }
302        if let Some(definition_level_histograms) = &self.definition_level_histograms {
303            last_field_id =
304                definition_level_histograms.write_thrift_field(writer, 7, last_field_id)?;
305        }
306        if let Some(nan_counts) = &self.nan_counts {
307            nan_counts.write_thrift_field(writer, 8, last_field_id)?;
308        }
309        writer.write_struct_end()
310    }
311}
312
313/// Column index for byte arrays (fixed length and variable)
314#[derive(Debug, Clone, PartialEq)]
315pub struct ByteArrayColumnIndex {
316    pub(crate) column_index: ColumnIndex,
317    // raw bytes for min and max values
318    pub(crate) min_bytes: Vec<u8>,
319    pub(crate) min_offsets: Vec<usize>,
320    pub(crate) max_bytes: Vec<u8>,
321    pub(crate) max_offsets: Vec<usize>,
322}
323
324impl ByteArrayColumnIndex {
325    #[allow(clippy::too_many_arguments)]
326    pub(crate) fn try_new(
327        null_pages: Vec<bool>,
328        boundary_order: BoundaryOrder,
329        null_counts: Option<Vec<i64>>,
330        nan_counts: Option<Vec<i64>>,
331        repetition_level_histograms: Option<Vec<i64>>,
332        definition_level_histograms: Option<Vec<i64>>,
333        min_values: Vec<&[u8]>,
334        max_values: Vec<&[u8]>,
335    ) -> Result<Self> {
336        let len = null_pages.len();
337
338        if min_values.len() != len || max_values.len() != len {
339            return Err(ParquetError::General(format!(
340                "ColumnIndex min/max length mismatch: expected {len}, got min={} max={}",
341                min_values.len(),
342                max_values.len()
343            )));
344        }
345        if let Some(ref nc) = null_counts
346            && nc.len() != len
347        {
348            return Err(ParquetError::General(format!(
349                "ColumnIndex null_counts length mismatch: expected {len}, got {}",
350                nc.len()
351            )));
352        }
353        if let Some(ref rep) = repetition_level_histograms
354            && len != 0
355            && rep.len() % len != 0
356        {
357            return Err(ParquetError::General(
358                "Invalid repetition_level_histograms length".to_string(),
359            ));
360        }
361        if let Some(ref def) = definition_level_histograms
362            && len != 0
363            && def.len() % len != 0
364        {
365            return Err(ParquetError::General(
366                "Invalid definition_level_histograms length".to_string(),
367            ));
368        }
369
370        let min_len = min_values.iter().map(|&v| v.len()).sum();
371        let max_len = max_values.iter().map(|&v| v.len()).sum();
372        let mut min_bytes = vec![0u8; min_len];
373        let mut max_bytes = vec![0u8; max_len];
374
375        let mut min_offsets = vec![0usize; len + 1];
376        let mut max_offsets = vec![0usize; len + 1];
377
378        let mut min_pos = 0;
379        let mut max_pos = 0;
380
381        for (i, is_null) in null_pages.iter().enumerate().take(len) {
382            if !is_null {
383                let min = min_values[i];
384                let dst = &mut min_bytes[min_pos..min_pos + min.len()];
385                dst.copy_from_slice(min);
386                min_offsets[i] = min_pos;
387                min_pos += min.len();
388
389                let max = max_values[i];
390                let dst = &mut max_bytes[max_pos..max_pos + max.len()];
391                dst.copy_from_slice(max);
392                max_offsets[i] = max_pos;
393                max_pos += max.len();
394            } else {
395                min_offsets[i] = min_pos;
396                max_offsets[i] = max_pos;
397            }
398        }
399
400        min_offsets[len] = min_pos;
401        max_offsets[len] = max_pos;
402
403        Ok(Self {
404            column_index: ColumnIndex {
405                null_pages,
406                boundary_order,
407                null_counts,
408                repetition_level_histograms,
409                definition_level_histograms,
410                nan_counts,
411            },
412            min_bytes,
413            min_offsets,
414            max_bytes,
415            max_offsets,
416        })
417    }
418
419    pub(super) fn try_from_thrift(index: ThriftColumnIndex) -> Result<Self> {
420        Self::try_new(
421            index.null_pages,
422            index.boundary_order,
423            index.null_counts,
424            index.nan_counts,
425            index.repetition_level_histograms,
426            index.definition_level_histograms,
427            index.min_values,
428            index.max_values,
429        )
430    }
431
432    /// Returns the min value for the page indexed by `idx`
433    ///
434    /// It is `None` when all values are null
435    pub fn min_value(&self, idx: usize) -> Option<&[u8]> {
436        if self.null_pages[idx] {
437            None
438        } else {
439            let start = self.min_offsets[idx];
440            let end = self.min_offsets[idx + 1];
441            Some(&self.min_bytes[start..end])
442        }
443    }
444
445    /// Returns the max value for the page indexed by `idx`
446    ///
447    /// It is `None` when all values are null
448    pub fn max_value(&self, idx: usize) -> Option<&[u8]> {
449        if self.null_pages[idx] {
450            None
451        } else {
452            let start = self.max_offsets[idx];
453            let end = self.max_offsets[idx + 1];
454            Some(&self.max_bytes[start..end])
455        }
456    }
457
458    /// Returns an iterator over the min values.
459    ///
460    /// Values may be `None` when [`ColumnIndex::is_null_page()`] is `true`.
461    pub fn min_values_iter(&self) -> impl Iterator<Item = Option<&[u8]>> {
462        (0..self.num_pages() as usize).map(|i| self.min_value(i))
463    }
464
465    /// Returns an iterator over the max values.
466    ///
467    /// Values may be `None` when [`ColumnIndex::is_null_page()`] is `true`.
468    pub fn max_values_iter(&self) -> impl Iterator<Item = Option<&[u8]>> {
469        (0..self.num_pages() as usize).map(|i| self.max_value(i))
470    }
471}
472
473impl Deref for ByteArrayColumnIndex {
474    type Target = ColumnIndex;
475
476    fn deref(&self) -> &Self::Target {
477        &self.column_index
478    }
479}
480
481impl WriteThrift for ByteArrayColumnIndex {
482    const ELEMENT_TYPE: ElementType = ElementType::Struct;
483    fn write_thrift<W: std::io::Write>(
484        &self,
485        writer: &mut ThriftCompactOutputProtocol<W>,
486    ) -> Result<()> {
487        self.null_pages.write_thrift_field(writer, 1, 0)?;
488
489        // need to handle min/max manually
490        let len = self.null_pages.len();
491        writer.write_field_begin(FieldType::List, 2, 1)?;
492        writer.write_list_begin(ElementType::Binary, len)?;
493        for i in 0..len {
494            let min = self.min_value(i).unwrap_or(&[]);
495            min.write_thrift(writer)?;
496        }
497        writer.write_field_begin(FieldType::List, 3, 2)?;
498        writer.write_list_begin(ElementType::Binary, len)?;
499        for i in 0..len {
500            let max = self.max_value(i).unwrap_or(&[]);
501            max.write_thrift(writer)?;
502        }
503        let mut last_field_id = self.boundary_order.write_thrift_field(writer, 4, 3)?;
504        if let Some(null_counts) = &self.null_counts {
505            last_field_id = null_counts.write_thrift_field(writer, 5, last_field_id)?;
506        }
507        if let Some(repetition_level_histograms) = &self.repetition_level_histograms {
508            last_field_id =
509                repetition_level_histograms.write_thrift_field(writer, 6, last_field_id)?;
510        }
511        if let Some(definition_level_histograms) = &self.definition_level_histograms {
512            last_field_id =
513                definition_level_histograms.write_thrift_field(writer, 7, last_field_id)?;
514        }
515        if let Some(nan_counts) = &self.nan_counts {
516            nan_counts.write_thrift_field(writer, 8, last_field_id)?;
517        }
518        writer.write_struct_end()
519    }
520}
521
522// Macro to generate getter functions for ColumnIndexMetaData.
523macro_rules! colidx_enum_func {
524    ($self:ident, $func:ident, $arg:ident) => {{
525        match *$self {
526            Self::BOOLEAN(ref typed) => typed.$func($arg),
527            Self::INT32(ref typed) => typed.$func($arg),
528            Self::INT64(ref typed) => typed.$func($arg),
529            Self::INT96(ref typed) => typed.$func($arg),
530            Self::FLOAT(ref typed) => typed.$func($arg),
531            Self::DOUBLE(ref typed) => typed.$func($arg),
532            Self::BYTE_ARRAY(ref typed) => typed.$func($arg),
533            Self::FIXED_LEN_BYTE_ARRAY(ref typed) => typed.$func($arg),
534            _ => panic!(concat!(
535                "Cannot call ",
536                stringify!($func),
537                " on ColumnIndexMetaData::NONE"
538            )),
539        }
540    }};
541    ($self:ident, $func:ident) => {{
542        match *$self {
543            Self::BOOLEAN(ref typed) => typed.$func(),
544            Self::INT32(ref typed) => typed.$func(),
545            Self::INT64(ref typed) => typed.$func(),
546            Self::INT96(ref typed) => typed.$func(),
547            Self::FLOAT(ref typed) => typed.$func(),
548            Self::DOUBLE(ref typed) => typed.$func(),
549            Self::BYTE_ARRAY(ref typed) => typed.$func(),
550            Self::FIXED_LEN_BYTE_ARRAY(ref typed) => typed.$func(),
551            _ => panic!(concat!(
552                "Cannot call ",
553                stringify!($func),
554                " on ColumnIndexMetaData::NONE"
555            )),
556        }
557    }};
558}
559
560/// Parsed [`ColumnIndex`] information for a Parquet file.
561///
562/// See [`ParquetColumnIndex`] for more information.
563///
564/// [`ParquetColumnIndex`]: crate::file::metadata::ParquetColumnIndex
565/// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
566#[derive(Debug, Clone, PartialEq)]
567#[allow(non_camel_case_types)]
568pub enum ColumnIndexMetaData {
569    /// Sometimes reading page index from parquet file
570    /// will only return pageLocations without min_max index,
571    /// `NONE` represents this lack of index information
572    NONE,
573    /// Boolean type index
574    BOOLEAN(PrimitiveColumnIndex<bool>),
575    /// 32-bit integer type index
576    INT32(PrimitiveColumnIndex<i32>),
577    /// 64-bit integer type index
578    INT64(PrimitiveColumnIndex<i64>),
579    /// 96-bit integer type (timestamp) index
580    INT96(PrimitiveColumnIndex<Int96>),
581    /// 32-bit floating point type index
582    FLOAT(PrimitiveColumnIndex<f32>),
583    /// 64-bit floating point type index
584    DOUBLE(PrimitiveColumnIndex<f64>),
585    /// Byte array type index
586    BYTE_ARRAY(ByteArrayColumnIndex),
587    /// Fixed length byte array type index
588    FIXED_LEN_BYTE_ARRAY(ByteArrayColumnIndex),
589}
590
591impl ColumnIndexMetaData {
592    /// Return min/max elements inside ColumnIndex are ordered or not.
593    pub fn is_sorted(&self) -> bool {
594        // 0:UNORDERED, 1:ASCENDING ,2:DESCENDING,
595        if let Some(order) = self.get_boundary_order() {
596            order != BoundaryOrder::UNORDERED
597        } else {
598            false
599        }
600    }
601
602    /// Get boundary_order of this page index.
603    pub fn get_boundary_order(&self) -> Option<BoundaryOrder> {
604        match self {
605            Self::NONE => None,
606            Self::BOOLEAN(index) => Some(index.boundary_order),
607            Self::INT32(index) => Some(index.boundary_order),
608            Self::INT64(index) => Some(index.boundary_order),
609            Self::INT96(index) => Some(index.boundary_order),
610            Self::FLOAT(index) => Some(index.boundary_order),
611            Self::DOUBLE(index) => Some(index.boundary_order),
612            Self::BYTE_ARRAY(index) => Some(index.boundary_order),
613            Self::FIXED_LEN_BYTE_ARRAY(index) => Some(index.boundary_order),
614        }
615    }
616
617    /// Returns array of null counts, one per page.
618    ///
619    /// Returns `None` if no null counts have been set in the index
620    pub fn null_counts(&self) -> Option<&Vec<i64>> {
621        match self {
622            Self::NONE => None,
623            Self::BOOLEAN(index) => index.null_counts.as_ref(),
624            Self::INT32(index) => index.null_counts.as_ref(),
625            Self::INT64(index) => index.null_counts.as_ref(),
626            Self::INT96(index) => index.null_counts.as_ref(),
627            Self::FLOAT(index) => index.null_counts.as_ref(),
628            Self::DOUBLE(index) => index.null_counts.as_ref(),
629            Self::BYTE_ARRAY(index) => index.null_counts.as_ref(),
630            Self::FIXED_LEN_BYTE_ARRAY(index) => index.null_counts.as_ref(),
631        }
632    }
633
634    /// Returns array of NaN counts, one per page.
635    ///
636    /// Returns `None` if no NaN counts have been set in the index
637    pub fn nan_counts(&self) -> Option<&Vec<i64>> {
638        match self {
639            Self::NONE => None,
640            Self::BOOLEAN(index) => index.nan_counts.as_ref(),
641            Self::INT32(index) => index.nan_counts.as_ref(),
642            Self::INT64(index) => index.nan_counts.as_ref(),
643            Self::INT96(index) => index.nan_counts.as_ref(),
644            Self::FLOAT(index) => index.nan_counts.as_ref(),
645            Self::DOUBLE(index) => index.nan_counts.as_ref(),
646            Self::BYTE_ARRAY(index) => index.nan_counts.as_ref(),
647            Self::FIXED_LEN_BYTE_ARRAY(index) => index.nan_counts.as_ref(),
648        }
649    }
650
651    /// Returns the number of pages
652    pub fn num_pages(&self) -> u64 {
653        colidx_enum_func!(self, num_pages)
654    }
655
656    /// Returns the number of null values in the page indexed by `idx`
657    ///
658    /// Returns `None` if no null counts have been set in the index
659    pub fn null_count(&self, idx: usize) -> Option<i64> {
660        colidx_enum_func!(self, null_count, idx)
661    }
662
663    /// Returns the number of NaN values in the page indexed by `idx`
664    ///
665    /// Returns `None` if no NaN counts have been set in the index
666    pub fn nan_count(&self, idx: usize) -> Option<i64> {
667        colidx_enum_func!(self, nan_count, idx)
668    }
669
670    /// Returns the repetition level histogram for the page indexed by `idx`
671    pub fn repetition_level_histogram(&self, idx: usize) -> Option<&[i64]> {
672        colidx_enum_func!(self, repetition_level_histogram, idx)
673    }
674
675    /// Returns the definition level histogram for the page indexed by `idx`
676    pub fn definition_level_histogram(&self, idx: usize) -> Option<&[i64]> {
677        colidx_enum_func!(self, definition_level_histogram, idx)
678    }
679
680    /// Returns whether the page indexed by `idx` consists of all null values
681    #[inline]
682    pub fn is_null_page(&self, idx: usize) -> bool {
683        colidx_enum_func!(self, is_null_page, idx)
684    }
685}
686
687/// Provides iterators over min and max values of a [`ColumnIndexMetaData`]
688pub trait ColumnIndexIterators {
689    /// Can be one of `bool`, `i32`, `i64`, `Int96`, `f32`, `f64`, [`ByteArray`],
690    /// or [`FixedLenByteArray`]
691    type Item;
692
693    /// Return iterator over the min values for the index
694    fn min_values_iter(colidx: &ColumnIndexMetaData) -> impl Iterator<Item = Option<Self::Item>>;
695
696    /// Return iterator over the max values for the index
697    fn max_values_iter(colidx: &ColumnIndexMetaData) -> impl Iterator<Item = Option<Self::Item>>;
698}
699
700macro_rules! column_index_iters {
701    ($item: ident, $variant: ident, $conv:expr) => {
702        impl ColumnIndexIterators for $item {
703            type Item = $item;
704
705            fn min_values_iter(
706                colidx: &ColumnIndexMetaData,
707            ) -> impl Iterator<Item = Option<Self::Item>> {
708                if let ColumnIndexMetaData::$variant(index) = colidx {
709                    index.min_values_iter().map($conv)
710                } else {
711                    panic!(concat!("Wrong type for ", stringify!($item), " iterator"))
712                }
713            }
714
715            fn max_values_iter(
716                colidx: &ColumnIndexMetaData,
717            ) -> impl Iterator<Item = Option<Self::Item>> {
718                if let ColumnIndexMetaData::$variant(index) = colidx {
719                    index.max_values_iter().map($conv)
720                } else {
721                    panic!(concat!("Wrong type for ", stringify!($item), " iterator"))
722                }
723            }
724        }
725    };
726}
727
728column_index_iters!(bool, BOOLEAN, |v| v.copied());
729column_index_iters!(i32, INT32, |v| v.copied());
730column_index_iters!(i64, INT64, |v| v.copied());
731column_index_iters!(Int96, INT96, |v| v.copied());
732column_index_iters!(f32, FLOAT, |v| v.copied());
733column_index_iters!(f64, DOUBLE, |v| v.copied());
734column_index_iters!(ByteArray, BYTE_ARRAY, |v| v
735    .map(|v| ByteArray::from(v.to_owned())));
736column_index_iters!(FixedLenByteArray, FIXED_LEN_BYTE_ARRAY, |v| v
737    .map(|v| FixedLenByteArray::from(v.to_owned())));
738
739impl WriteThrift for ColumnIndexMetaData {
740    const ELEMENT_TYPE: ElementType = ElementType::Struct;
741
742    fn write_thrift<W: std::io::Write>(
743        &self,
744        writer: &mut ThriftCompactOutputProtocol<W>,
745    ) -> Result<()> {
746        match self {
747            ColumnIndexMetaData::BOOLEAN(index) => index.write_thrift(writer),
748            ColumnIndexMetaData::INT32(index) => index.write_thrift(writer),
749            ColumnIndexMetaData::INT64(index) => index.write_thrift(writer),
750            ColumnIndexMetaData::INT96(index) => index.write_thrift(writer),
751            ColumnIndexMetaData::FLOAT(index) => index.write_thrift(writer),
752            ColumnIndexMetaData::DOUBLE(index) => index.write_thrift(writer),
753            ColumnIndexMetaData::BYTE_ARRAY(index) => index.write_thrift(writer),
754            ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index) => index.write_thrift(writer),
755            _ => Err(general_err!("Cannot serialize NONE index")),
756        }
757    }
758}
759
760#[cfg(test)]
761mod tests {
762    use super::*;
763
764    #[test]
765    fn test_page_index_min_max_null() {
766        let column_index = PrimitiveColumnIndex {
767            column_index: ColumnIndex {
768                null_pages: vec![false],
769                boundary_order: BoundaryOrder::ASCENDING,
770                null_counts: Some(vec![0]),
771                nan_counts: None,
772                repetition_level_histograms: Some(vec![1, 2]),
773                definition_level_histograms: Some(vec![1, 2, 3]),
774            },
775            min_values: vec![-123],
776            max_values: vec![234],
777        };
778
779        assert_eq!(column_index.min_value(0), Some(&-123));
780        assert_eq!(column_index.max_value(0), Some(&234));
781        assert_eq!(column_index.null_count(0), Some(0));
782        assert_eq!(column_index.repetition_level_histogram(0).unwrap(), &[1, 2]);
783        assert_eq!(
784            column_index.definition_level_histogram(0).unwrap(),
785            &[1, 2, 3]
786        );
787    }
788
789    #[test]
790    fn test_page_index_min_max_null_none() {
791        let column_index: PrimitiveColumnIndex<i32> = PrimitiveColumnIndex::<i32> {
792            column_index: ColumnIndex {
793                null_pages: vec![true],
794                boundary_order: BoundaryOrder::ASCENDING,
795                null_counts: Some(vec![1]),
796                nan_counts: None,
797                repetition_level_histograms: None,
798                definition_level_histograms: Some(vec![1, 0]),
799            },
800            min_values: vec![Default::default()],
801            max_values: vec![Default::default()],
802        };
803
804        assert_eq!(column_index.min_value(0), None);
805        assert_eq!(column_index.max_value(0), None);
806        assert_eq!(column_index.null_count(0), Some(1));
807        assert_eq!(column_index.repetition_level_histogram(0), None);
808        assert_eq!(column_index.definition_level_histogram(0).unwrap(), &[1, 0]);
809    }
810
811    #[test]
812    fn test_invalid_column_index() {
813        let column_index = ThriftColumnIndex {
814            null_pages: vec![true, false],
815            min_values: vec![
816                &[],
817                &[], // this shouldn't be empty as null_pages[1] is false
818            ],
819            max_values: vec![
820                &[],
821                &[], // this shouldn't be empty as null_pages[1] is false
822            ],
823            null_counts: None,
824            nan_counts: None,
825            repetition_level_histograms: None,
826            definition_level_histograms: None,
827            boundary_order: BoundaryOrder::UNORDERED,
828        };
829
830        let err = PrimitiveColumnIndex::<i32>::try_from_thrift(column_index).unwrap_err();
831        assert_eq!(
832            err.to_string(),
833            "Parquet error: error converting value, expected 4 bytes got 0"
834        );
835    }
836
837    #[test]
838    fn test_column_index_rejects_mismatched_min_max_lengths() {
839        // Two pages, but only one min/max entry. The entry itself is valid i32 bytes,
840        // so this specifically checks that lengths must match the number of pages.
841        let column_index = ThriftColumnIndex {
842            null_pages: vec![false, false],
843            min_values: vec![&[1u8, 0, 0, 0]],
844            max_values: vec![&[10u8, 0, 0, 0]],
845            null_counts: None,
846            repetition_level_histograms: None,
847            definition_level_histograms: None,
848            boundary_order: BoundaryOrder::UNORDERED,
849            nan_counts: None,
850        };
851
852        // ColumnIndex arrays must align with the number of pages (null_pages.len()).
853        let err = PrimitiveColumnIndex::<i32>::try_from_thrift(column_index).unwrap_err();
854        // Should fail because min/max lengths don’t match null_pages
855        assert!(err.to_string().contains("length mismatch"));
856    }
857}