1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

//! APIs exposing `parquet2`'s statistics as arrow's statistics.
use std::any::Any;

use parquet2::metadata::ColumnChunkMetaData;
use parquet2::schema::types::PhysicalType;
use parquet2::statistics::PrimitiveStatistics as ParquetPrimitiveStatistics;
use parquet2::statistics::Statistics as ParquetStatistics;

use crate::datatypes::DataType;
use crate::datatypes::Field;
use crate::error::ArrowError;
use crate::error::Result;

mod primitive;
pub use primitive::*;
mod binary;
pub use binary::*;
mod boolean;
pub use boolean::*;
mod fixlen;
pub use fixlen::*;

use super::get_field_columns;

/// Trait representing a deserialized parquet statistics into arrow.
pub trait Statistics: std::fmt::Debug {
    /// returns the [`DataType`] of the statistics.
    fn data_type(&self) -> &DataType;

    /// Returns `dyn Any` can used to downcast to a physical type.
    fn as_any(&self) -> &dyn Any;

    /// Return the null count statistic
    fn null_count(&self) -> Option<i64>;
}

impl PartialEq for &dyn Statistics {
    fn eq(&self, other: &Self) -> bool {
        self.data_type() == other.data_type()
    }
}

impl PartialEq for Box<dyn Statistics> {
    fn eq(&self, other: &Self) -> bool {
        self.data_type() == other.data_type()
    }
}

/// Deserializes [`ParquetStatistics`] into [`Statistics`] based on `data_type`.
/// This takes into account the Arrow schema declared in Parquet's schema
fn _deserialize_statistics(
    stats: &dyn ParquetStatistics,
    data_type: DataType,
) -> Result<Box<dyn Statistics>> {
    match stats.physical_type() {
        PhysicalType::Int32 => {
            let stats = stats.as_any().downcast_ref().unwrap();
            primitive::statistics_from_i32(stats, data_type)
        }
        PhysicalType::Int64 => {
            let stats = stats.as_any().downcast_ref().unwrap();
            primitive::statistics_from_i64(stats, data_type)
        }
        PhysicalType::ByteArray => {
            let stats = stats.as_any().downcast_ref().unwrap();
            binary::statistics_from_byte_array(stats, data_type)
        }
        PhysicalType::Boolean => {
            let stats = stats.as_any().downcast_ref().unwrap();
            Ok(Box::new(BooleanStatistics::from(stats)))
        }
        PhysicalType::Float => {
            let stats = stats
                .as_any()
                .downcast_ref::<ParquetPrimitiveStatistics<f32>>()
                .unwrap();
            Ok(Box::new(PrimitiveStatistics::<f32>::from((
                stats, data_type,
            ))))
        }
        PhysicalType::Double => {
            let stats = stats
                .as_any()
                .downcast_ref::<ParquetPrimitiveStatistics<f64>>()
                .unwrap();
            Ok(Box::new(PrimitiveStatistics::<f64>::from((
                stats, data_type,
            ))))
        }
        PhysicalType::FixedLenByteArray(_) => {
            let stats = stats.as_any().downcast_ref().unwrap();
            fixlen::statistics_from_fix_len(stats, data_type)
        }
        _ => Err(ArrowError::NotYetImplemented(
            "Reading Fixed-len array statistics is not yet supported".to_string(),
        )),
    }
}

fn get_fields(field: &Field) -> Vec<&Field> {
    match field.data_type.to_logical_type() {
        DataType::List(inner) => get_fields(inner),
        DataType::LargeList(inner) => get_fields(inner),
        DataType::Struct(fields) => fields.iter().flat_map(get_fields).collect(),
        _ => vec![field],
    }
}

/// Deserializes [`ParquetStatistics`] into [`Statistics`] associated to `field`
///
/// For non-nested types, it returns a single column.
/// For nested types, it returns one column per parquet primitive column.
pub fn deserialize_statistics(
    field: &Field,
    columns: &[ColumnChunkMetaData],
) -> Result<Vec<Option<Box<dyn Statistics>>>> {
    let columns = get_field_columns(columns, field.name.as_ref());

    let fields = get_fields(field);

    columns
        .into_iter()
        .zip(fields.into_iter())
        .map(|(column, field)| {
            column
                .statistics()
                .map(|x| _deserialize_statistics(x?.as_ref(), field.data_type.clone()))
                .transpose()
        })
        .collect()
}