tustvold commented on code in PR #8294: URL: https://github.com/apache/arrow-datafusion/pull/8294#discussion_r1402772669
########## datafusion/core/src/datasource/physical_plan/parquet/statistics.rs: ########## @@ -0,0 +1,807 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`RowGroupStatisticsConverter`] tp converts parquet RowGroup statistics to arrow [`ArrayRef`]. + +use arrow::{array::ArrayRef, datatypes::DataType}; +use arrow_array::new_empty_array; +use arrow_schema::Field; +use datafusion_common::{Result, ScalarValue}; +use parquet::file::{ + metadata::RowGroupMetaData, statistics::Statistics as ParquetStatistics, +}; + +// Convert the bytes array to i128. +// The endian of the input bytes array must be big-endian. +pub(crate) fn from_bytes_to_i128(b: &[u8]) -> i128 { + // The bytes array are from parquet file and must be the big-endian. + // The endian is defined by parquet format, and the reference document + // https://github.com/apache/parquet-format/blob/54e53e5d7794d383529dd30746378f19a12afd58/src/main/thrift/parquet.thrift#L66 + i128::from_be_bytes(sign_extend_be(b)) +} + +// Copy from arrow-rs +// https://github.com/apache/arrow-rs/blob/733b7e7fd1e8c43a404c3ce40ecf741d493c21b4/parquet/src/arrow/buffer/bit_util.rs#L55 +// Convert the byte slice to fixed length byte array with the length of 16 +fn sign_extend_be(b: &[u8]) -> [u8; 16] { + assert!(b.len() <= 16, "Array too large, expected less than 16"); + let is_negative = (b[0] & 128u8) == 128u8; + let mut result = if is_negative { [255u8; 16] } else { [0u8; 16] }; + for (d, s) in result.iter_mut().skip(16 - b.len()).zip(b) { + *d = *s; + } + result +} + +/// Converts parquet RowGroup statistics (stored in +/// [`RowGroupMetaData`]) into an arrow [`ArrayRef`] +/// +/// For example, given a parquet file with 3 Row Groups, when asked for +/// statistics for column "A" it will return a single array with 3 elements, +/// +pub(crate) struct RowGroupStatisticsConverter<'a> { + field: &'a Field, +} + +/// Extract a single min/max statistics from a [`ParquetStatistics`] object +/// +/// * `$column_statistics` is the `ParquetStatistics` object +/// * `$func is the function` (`min`/`max`) to call to get the value +/// * `$bytes_func` is the function (`min_bytes`/`max_bytes`) to call to get the value as bytes +/// * `$target_arrow_type` is the [`DataType`] of the target statistics +macro_rules! get_statistic { + ($column_statistics:expr, $func:ident, $bytes_func:ident, $target_arrow_type:expr) => {{ + if !$column_statistics.has_min_max_set() { + return None; + } + match $column_statistics { + ParquetStatistics::Boolean(s) => Some(ScalarValue::Boolean(Some(*s.$func()))), + ParquetStatistics::Int32(s) => { + match $target_arrow_type { + // int32 to decimal with the precision and scale + Some(DataType::Decimal128(precision, scale)) => { + Some(ScalarValue::Decimal128( + Some(*s.$func() as i128), + precision, + scale, + )) + } + _ => Some(ScalarValue::Int32(Some(*s.$func()))), + } + } + ParquetStatistics::Int64(s) => { + match $target_arrow_type { + // int64 to decimal with the precision and scale + Some(DataType::Decimal128(precision, scale)) => { + Some(ScalarValue::Decimal128( + Some(*s.$func() as i128), + precision, + scale, + )) + } + _ => Some(ScalarValue::Int64(Some(*s.$func()))), + } + } + // 96 bit ints not supported + ParquetStatistics::Int96(_) => None, + ParquetStatistics::Float(s) => Some(ScalarValue::Float32(Some(*s.$func()))), + ParquetStatistics::Double(s) => Some(ScalarValue::Float64(Some(*s.$func()))), + ParquetStatistics::ByteArray(s) => { + match $target_arrow_type { + // decimal data type + Some(DataType::Decimal128(precision, scale)) => { + Some(ScalarValue::Decimal128( + Some(from_bytes_to_i128(s.$bytes_func())), + precision, + scale, + )) + } + _ => { + let s = std::str::from_utf8(s.$bytes_func()) + .map(|s| s.to_string()) + .ok(); + Some(ScalarValue::Utf8(s)) + } + } + } + // type not supported yet + ParquetStatistics::FixedLenByteArray(s) => { + match $target_arrow_type { + // just support the decimal data type + Some(DataType::Decimal128(precision, scale)) => { + Some(ScalarValue::Decimal128( + Some(from_bytes_to_i128(s.$bytes_func())), + precision, + scale, + )) + } + _ => None, + } + } + } + }}; +} + +#[derive(Debug, Clone, Copy)] +enum MinMax { + Min, + Max, +} + +impl<'a> RowGroupStatisticsConverter<'a> { + /// Create a new RowGoupStatisticsConverter suitable that can extract + /// statistics for the specified field + pub fn new(field: &'a Field) -> Self { + Self { field } + } + + /// Returns the min value for the column into an array ref. + pub fn min<'b>( + &self, + row_group_meta_data: impl IntoIterator<Item = &'b RowGroupMetaData>, + ) -> Result<ArrayRef> { + self.min_max_impl(MinMax::Min, row_group_meta_data) + } + + /// Returns the max value for the column into an array ref. + pub fn max<'b>( + &self, + row_group_meta_data: impl IntoIterator<Item = &'b RowGroupMetaData>, + ) -> Result<ArrayRef> { + self.min_max_impl(MinMax::Max, row_group_meta_data) + } + + /// Extracts all min/max values for the column into an array ref. + fn min_max_impl<'b>( + &self, + mm: MinMax, + row_group_meta_data: impl IntoIterator<Item = &'b RowGroupMetaData>, + ) -> Result<ArrayRef> { + let mut row_group_meta_data = row_group_meta_data.into_iter().peekable(); + + // if it is empty, return empty array + if row_group_meta_data.peek().is_none() { + return Ok(new_empty_array(self.field.data_type())); + } + + let maybe_index = row_group_meta_data.peek().and_then(|rg_meta| { + rg_meta + .columns() + .iter() + .enumerate() + .find(|(_idx, c)| c.column_descr().name() == self.field.name()) + .map(|(idx, _c)| idx) + }); + + // don't have this column, return an array of all NULLs + let Some(column_index) = maybe_index else { + let num_row_groups = row_group_meta_data.count(); + let sv = ScalarValue::try_from(self.field.data_type())?; + return sv.to_array_of_size(num_row_groups); + }; + + let stats_iter = row_group_meta_data.map(move |row_group_meta_data| { + row_group_meta_data.column(column_index).statistics() + }); + + // this is the value to use when the statistics are not set + let null_value = ScalarValue::try_from(self.field.data_type())?; + match mm { + MinMax::Min => { + let values = stats_iter.map(|column_statistics| { + column_statistics + .and_then(|column_statistics| { + get_statistic!( + column_statistics, + min, + min_bytes, + Some(self.field.data_type().clone()) + ) + }) + .unwrap_or_else(|| null_value.clone()) + }); + ScalarValue::iter_to_array(values) + } + MinMax::Max => { + let values = stats_iter.map(|column_statistics| { + column_statistics + .and_then(|column_statistics| { + get_statistic!( + column_statistics, + max, + max_bytes, + Some(self.field.data_type().clone()) + ) + }) + .unwrap_or_else(|| null_value.clone()) + }); + ScalarValue::iter_to_array(values) + } + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use arrow_array::{ + BinaryArray, BooleanArray, Decimal128Array, Float32Array, Float64Array, + Int32Array, Int64Array, RecordBatch, StringArray, TimestampNanosecondArray, + }; + use arrow_schema::SchemaRef; + use bytes::Bytes; + use datafusion_common::test_util::parquet_test_data; + use parquet::arrow::arrow_reader::ArrowReaderBuilder; + use parquet::arrow::arrow_writer::ArrowWriter; + use parquet::file::metadata::ParquetMetaData; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; + use std::path::PathBuf; + use std::sync::Arc; + + // TODO error cases (with parquet statistics that are mismatched in expected type) + Review Comment: Whilst I'm not entirely sure how repeated fields could be handled, I think a test of nested groups, i.e. StructArray, would be worthwhile to demonstrate the mismatch between arrow and parquet schema representations, and I think is a useful test. We could also have a test that repeated fields are correctly ignored or something -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
