use crate::{error::Error, metadata::get_sort_order};
use super::{column_order::ColumnOrder, schema_descriptor::SchemaDescriptor, RowGroupMetaData};
use parquet_format_safe::ColumnOrder as TColumnOrder;
pub use crate::thrift_format::KeyValue;
#[derive(Debug, Clone)]
pub struct FileMetaData {
pub version: i32,
pub num_rows: usize,
pub created_by: Option<String>,
pub row_groups: Vec<RowGroupMetaData>,
pub key_value_metadata: Option<Vec<KeyValue>>,
pub schema_descr: SchemaDescriptor,
pub column_orders: Option<Vec<ColumnOrder>>,
}
impl FileMetaData {
pub fn schema(&self) -> &SchemaDescriptor {
&self.schema_descr
}
pub fn key_value_metadata(&self) -> &Option<Vec<KeyValue>> {
&self.key_value_metadata
}
pub fn column_order(&self, i: usize) -> ColumnOrder {
self.column_orders
.as_ref()
.map(|data| data[i])
.unwrap_or(ColumnOrder::Undefined)
}
pub fn try_from_thrift(metadata: parquet_format_safe::FileMetaData) -> Result<Self, Error> {
let schema_descr = SchemaDescriptor::try_from_thrift(&metadata.schema)?;
let row_groups = metadata
.row_groups
.into_iter()
.map(|rg| RowGroupMetaData::try_from_thrift(&schema_descr, rg))
.collect::<Result<_, Error>>()?;
let column_orders = metadata
.column_orders
.map(|orders| parse_column_orders(&orders, &schema_descr));
Ok(FileMetaData {
version: metadata.version,
num_rows: metadata.num_rows.try_into()?,
created_by: metadata.created_by,
row_groups,
key_value_metadata: metadata.key_value_metadata,
schema_descr,
column_orders,
})
}
pub fn into_thrift(self) -> parquet_format_safe::FileMetaData {
parquet_format_safe::FileMetaData {
version: self.version,
schema: self.schema_descr.into_thrift(),
num_rows: self.num_rows as i64,
row_groups: self
.row_groups
.into_iter()
.map(|v| v.into_thrift())
.collect(),
key_value_metadata: self.key_value_metadata,
created_by: self.created_by,
column_orders: None, encryption_algorithm: None,
footer_signing_key_metadata: None,
}
}
}
fn parse_column_orders(
orders: &[TColumnOrder],
schema_descr: &SchemaDescriptor,
) -> Vec<ColumnOrder> {
schema_descr
.columns()
.iter()
.zip(orders.iter())
.map(|(column, order)| match order {
TColumnOrder::TYPEORDER(_) => {
let sort_order = get_sort_order(
&column.descriptor.primitive_type.logical_type,
&column.descriptor.primitive_type.converted_type,
&column.descriptor.primitive_type.physical_type,
);
ColumnOrder::TypeDefinedOrder(sort_order)
}
})
.collect()
}