This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11523-8f4bb0caa723c9536e59669418d475a6b544d2e9 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit aa3d58378d61812f047a9cc80b60156edca8053f Author: Raz Luvaton <[email protected]> AuthorDate: Fri Oct 16 20:11:31 2026 +0000 perf(parquet): optimize `i32`, `u32`, `i64`, `u64`, `f32`, `f64` min_max calculation (#11523) # Which issue does this PR close? None # Rationale for this change calculating min and max should be super fast # What changes are included in this PR? implemented custom highly vectorized min_max for `i32`, `u32`, `i64`, `u64`, `f32`, `f64` in parquet writer # Are these changes tested? Yes + existing tests # Are there any significant user-facing changes? No --- parquet/src/column/writer/encoder.rs | 54 +---- parquet/src/column/writer/min_max.rs | 389 +++++++++++++++++++++++++++++++++++ parquet/src/column/writer/mod.rs | 1 + parquet/src/data_type.rs | 38 +++- 4 files changed, 426 insertions(+), 56 deletions(-) diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 4939e19ce4..c0d10aeb95 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -19,9 +19,7 @@ use bytes::Bytes; use crate::basic::{ConvertedType, Encoding, LogicalType, Type}; use crate::bloom_filter::Sbbf; -use crate::column::writer::{ - compare_greater, fallback_encoding, has_dictionary_support, is_nan, update_max, update_min, -}; +use crate::column::writer::{fallback_encoding, has_dictionary_support, update_max, update_min}; use crate::data_type::DataType; use crate::data_type::private::ParquetValueType; use crate::encodings::encoding::{DictEncoder, Encoder, get_encoder_with_properties}; @@ -29,7 +27,7 @@ use crate::errors::{ParquetError, Result}; use crate::file::properties::{EnabledStatistics, ResolvedColumnProperties, WriterProperties}; use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accumulator}; use crate::geospatial::statistics::GeospatialStatistics; -use crate::schema::types::{BasicTypeInfo, ColumnDescPtr}; +use crate::schema::types::ColumnDescPtr; /// A collection of [`ParquetValueType`] encoded by a [`ColumnValueEncoder`] pub trait ColumnValues { @@ -226,7 +224,7 @@ impl<T: DataType> ColumnValueEncoderImpl<T> { if let Some(accumulator) = self.geo_stats_accumulator.as_deref_mut() { update_geo_stats_accumulator(accumulator, slice.iter()); } else if let Some((min, max, nan_count)) = - get_min_max(self.descr.get_basic_info(), slice.iter()) + T::T::min_max(self.descr.get_basic_info(), slice) { update_min(&self.descr, &min, &mut self.min_value); update_max(&self.descr, &max, &mut self.max_value); @@ -456,52 +454,6 @@ impl<T: DataType> ColumnValueEncoder for ColumnValueEncoderImpl<T> { } } -// Get min and max values for all values in `iter`. -// -// For floating point we need to compare NaN values until we encounter a non-NaN -// value which then becomes the new min/max. After this, only non-NaN values are -// evaluated. If all values are NaN, then the min/max NaNs as determined by -// IEEE 754 total order are returned. -fn get_min_max<'a, T, I>(basic_type_info: &BasicTypeInfo, mut iter: I) -> Option<(T, T, u64)> -where - T: ParquetValueType + 'a, - I: Iterator<Item = &'a T>, -{ - let first = iter.next()?; - let mut min_max_nan = is_nan(basic_type_info, first); - let mut nan_count = min_max_nan as u64; - - let mut min = first; - let mut max = first; - for val in iter { - match (min_max_nan, is_nan(basic_type_info, val)) { - // skip NaNs if we've encounter non-NaN - (false, true) => { - nan_count += 1; - } - // if min/max are NaN, check for non-NaN and reset - (true, false) => { - min = val; - max = val; - min_max_nan = false; - } - // both are NaN or non-NaN, so do the comparison - (_, val_is_nan) => { - nan_count += val_is_nan as u64; - // we've already initialized min and max, so a single value can't be both - // extremes - if compare_greater(basic_type_info, min, val) { - min = val; - } else if compare_greater(basic_type_info, val, max) { - max = val; - } - } - } - } - - Some((min.clone(), max.clone(), nan_count)) -} - /// Creates a bloom filter sized for the column's configured NDV, returning the filter /// and the target FPP for folding. pub(crate) fn create_bloom_filter( diff --git a/parquet/src/column/writer/min_max.rs b/parquet/src/column/writer/min_max.rs new file mode 100644 index 0000000000..bcaa13ccf4 --- /dev/null +++ b/parquet/src/column/writer/min_max.rs @@ -0,0 +1,389 @@ +// 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. + +use crate::basic::SortOrder; +use crate::column::writer::{compare_greater, is_nan}; +use crate::data_type::private::ParquetValueType; +use crate::schema::types::BasicTypeInfo; + +/// Accumulators for 64 bytes of values, i.e. a cache line and one AVX-512 +/// register, or four NEON / two AVX2 registers. +const LANES_32: usize = 16; +const LANES_64: usize = 8; + +/// Returns the minimum and maximum of `values` after mapping them with `key`, +/// or `None` if `values` is empty. +#[inline(always)] +fn min_max_by_key<T: Copy, K: Copy + Ord, const LANES: usize>( + values: &[T], + key: impl Fn(T) -> K, + key_min: K, + key_max: K, +) -> Option<(K, K)> { + if values.is_empty() { + return None; + } + + let mut mins = [key_max; LANES]; + let mut maxs = [key_min; LANES]; + + let (chunks, remainder) = values.as_chunks::<LANES>(); + for chunk in chunks { + for i in 0..LANES { + let k = key(chunk[i]); + mins[i] = mins[i].min(k); + maxs[i] = maxs[i].max(k); + } + } + for (i, v) in remainder.iter().enumerate() { + let k = key(*v); + mins[i] = mins[i].min(k); + maxs[i] = maxs[i].max(k); + } + + let min = mins.into_iter().min().unwrap(); + let max = maxs.into_iter().max().unwrap(); + Some((min, max)) +} + +#[inline] +fn is_unsigned(basic_type_info: &BasicTypeInfo) -> bool { + matches!(basic_type_info.sort_order(), SortOrder::UNSIGNED) +} + +/// Returns the min and max of an `INT32` column, or `None` if `values` is empty. +pub(crate) fn min_max_i32(basic_type_info: &BasicTypeInfo, values: &[i32]) -> Option<(i32, i32)> { + if is_unsigned(basic_type_info) { + min_max_by_key::<_, _, LANES_32>(values, |v| v as u32, u32::MIN, u32::MAX) + .map(|(min, max)| (min as i32, max as i32)) + } else { + min_max_by_key::<_, _, LANES_32>(values, |v| v, i32::MIN, i32::MAX) + } +} + +/// Returns the min and max of an `INT64` column, or `None` if `values` is empty. +pub(crate) fn min_max_i64(basic_type_info: &BasicTypeInfo, values: &[i64]) -> Option<(i64, i64)> { + if is_unsigned(basic_type_info) { + min_max_by_key::<_, _, LANES_64>(values, |v| v as u64, u64::MIN, u64::MAX) + .map(|(min, max)| (min as i64, max as i64)) + } else { + min_max_by_key::<_, _, LANES_64>(values, |v| v, i64::MIN, i64::MAX) + } +} + +macro_rules! float_min_max { + ($name:ident, $float:ty, $int:ty, $uint:ty, $lanes:expr) => { + /// Returns the min, max and NaN count of a floating point column, + /// comparing by IEEE 754 total order and skipping NaN. + /// + /// Returns `None` if `values` is empty or contains only NaN, in which + /// case the min and max are NaN themselves and the caller falls back + /// to the scalar implementation. + pub(crate) fn $name(values: &[$float]) -> Option<($float, $float, u64)> { + const SIGN_SHIFT: u32 = <$int>::BITS - 1; + + // Maps the bits of a float to an integer with the same ordering as + // `total_cmp`: negative values have their magnitude bits flipped + // so that they order in reverse. The mapping is its own inverse. + // + // To avoid calling `{f32,f64}::total_cmp` for both min and max values and every other comparison + // we transform the input once (like what rust is doing internally before comparing) and reuse it, this avoids the overhead of calling `total_cmp` twice for each value. + // and doing the same calculation. + // + // This is the same transformation used in Rust `f32::total_cmp` and `f64::total_cmp` + // but we avoid calling + // f32: https://github.com/rust-lang/rust/blob/b5427502dab50cb27d9327d2893ddf2337dd4388/library/core/src/num/f32.rs#L1626 + // f64: https://github.com/rust-lang/rust/blob/b5427502dab50cb27d9327d2893ddf2337dd4388/library/core/src/num/f64.rs#L1604 + #[inline(always)] + fn total_order_key(bits: $int) -> $int { + bits ^ ((((bits >> SIGN_SHIFT) as $uint) >> 1) as $int) + } + + // Bounds the per lane NaN counters well below overflow on every + // platform, including 32-bit. + const BLOCK: usize = 1 << 30; + + let mut min = <$int>::MAX; + let mut max = <$int>::MIN; + let mut nan_count = 0u64; + + for block in values.chunks(BLOCK) { + let mut mins = [<$int>::MAX; $lanes]; + let mut maxs = [<$int>::MIN; $lanes]; + let mut nans = [<$uint>::MIN; $lanes]; + + // NaN lanes are replaced by the identity of each accumulator, so + // they never win, without a branch. + let mut accumulate = |i: usize, v: $float| { + let bits = v.to_bits(); + let is_nan = v.is_nan(); + let key = total_order_key(bits as $int); + mins[i] = mins[i].min(if is_nan { <$int>::MAX } else { key }); + maxs[i] = maxs[i].max(if is_nan { <$int>::MIN } else { key }); + nans[i] += is_nan as $uint; + }; + + let (chunks, remainder) = block.as_chunks::<$lanes>(); + for chunk in chunks { + for i in 0..$lanes { + accumulate(i, chunk[i]); + } + } + for (i, v) in remainder.iter().enumerate() { + accumulate(i, *v); + } + + min = min.min(mins.into_iter().min().unwrap()); + max = max.max(maxs.into_iter().max().unwrap()); + nan_count += nans.into_iter().map(u64::from).sum::<u64>(); + } + + if nan_count as usize == values.len() { + return None; + } + + let min = <$float>::from_bits(total_order_key(min) as $uint); + let max = <$float>::from_bits(total_order_key(max) as $uint); + Some((min, max, nan_count)) + } + }; +} + +float_min_max!(min_max_f32, f32, i32, u32, LANES_32); +float_min_max!(min_max_f64, f64, i64, u64, LANES_64); + +// Get min and max values for all values in `iter`. +// +// For floating point we need to compare NaN values until we encounter a non-NaN +// value which then becomes the new min/max. After this, only non-NaN values are +// evaluated. If all values are NaN, then the min/max NaNs as determined by +// IEEE 754 total order are returned. +pub(crate) fn fallback_min_max<'a, T, I>( + basic_type_info: &BasicTypeInfo, + mut iter: I, +) -> Option<(T, T, u64)> +where + T: ParquetValueType + 'a, + I: Iterator<Item = &'a T>, +{ + let first = iter.next()?; + let mut min_max_nan = is_nan(basic_type_info, first); + let mut nan_count = min_max_nan as u64; + + let mut min = first; + let mut max = first; + for val in iter { + match (min_max_nan, is_nan(basic_type_info, val)) { + // skip NaNs if we've encounter non-NaN + (false, true) => { + nan_count += 1; + } + // if min/max are NaN, check for non-NaN and reset + (true, false) => { + min = val; + max = val; + min_max_nan = false; + } + // both are NaN or non-NaN, so do the comparison + (_, val_is_nan) => { + nan_count += val_is_nan as u64; + // we've already initialized min and max, so a single value can't be both + // extremes + if compare_greater(basic_type_info, min, val) { + min = val; + } else if compare_greater(basic_type_info, val, max) { + max = val; + } + } + } + } + + Some((min.clone(), max.clone(), nan_count)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::basic::{ConvertedType, IntType, LogicalType, Type as PhysicalType}; + use crate::schema::types::Type; + use rand::prelude::*; + + fn type_info(physical: PhysicalType, logical: Option<LogicalType>) -> BasicTypeInfo { + Type::primitive_type_builder("col", physical) + .with_logical_type(logical) + .build() + .unwrap() + .get_basic_info() + .clone() + } + + fn signed_int(bits: u8) -> Option<LogicalType> { + Some(LogicalType::Integer(IntType { + bit_width: bits as i8, + is_signed: true, + })) + } + + fn unsigned_int(bits: u8) -> Option<LogicalType> { + Some(LogicalType::Integer(IntType { + bit_width: bits as i8, + is_signed: false, + })) + } + + /// Lengths around every lane count and remainder boundary + fn lengths() -> impl Iterator<Item = usize> { + (0..=70).chain([127, 128, 129, 1000, 4097]) + } + + fn check_ints<T, F>( + rng: &mut StdRng, + info: &BasicTypeInfo, + kernel: F, + gen_value: impl Fn(&mut StdRng) -> T, + ) where + T: crate::data_type::private::ParquetValueType + Copy, + F: Fn(&BasicTypeInfo, &[T]) -> Option<(T, T)>, + { + for len in lengths() { + let values: Vec<T> = (0..len).map(|_| gen_value(rng)).collect(); + let expected = fallback_min_max(info, values.iter()); + let actual = kernel(info, &values).map(|(min, max)| (min, max, 0)); + assert_eq!(actual, expected, "len {len}: {values:?}"); + } + } + + #[test] + fn test_min_max_i32() { + let mut rng = StdRng::seed_from_u64(42); + let edges = [i32::MIN, i32::MIN + 1, -1, 0, 1, i32::MAX - 1, i32::MAX]; + for info in [ + type_info(PhysicalType::INT32, None), + type_info(PhysicalType::INT32, signed_int(32)), + type_info(PhysicalType::INT32, unsigned_int(32)), + type_info(PhysicalType::INT32, unsigned_int(8)), + ] { + check_ints(&mut rng, &info, min_max_i32, |r| r.random()); + check_ints(&mut rng, &info, min_max_i32, |r| *edges.choose(r).unwrap()); + check_ints(&mut rng, &info, min_max_i32, |r| r.random_range(-3..3)); + } + } + + #[test] + fn test_min_max_i64() { + let mut rng = StdRng::seed_from_u64(42); + let edges = [i64::MIN, i64::MIN + 1, -1, 0, 1, i64::MAX - 1, i64::MAX]; + for info in [ + type_info(PhysicalType::INT64, None), + type_info(PhysicalType::INT64, signed_int(64)), + type_info(PhysicalType::INT64, unsigned_int(64)), + ] { + check_ints(&mut rng, &info, min_max_i64, |r| r.random()); + check_ints(&mut rng, &info, min_max_i64, |r| *edges.choose(r).unwrap()); + check_ints(&mut rng, &info, min_max_i64, |r| r.random_range(-3..3)); + } + } + + #[test] + fn test_unsigned_int32_converted_type() { + let info = Type::primitive_type_builder("col", PhysicalType::INT32) + .with_converted_type(ConvertedType::UINT_32) + .build() + .unwrap() + .get_basic_info() + .clone(); + assert_eq!(min_max_i32(&info, &[-1, 0, 5]), Some((0, -1))); + } + + macro_rules! float_tests { + ($test:ident, $float:ty, $uint:ty, $kernel:ident, $physical:expr) => { + #[test] + fn $test() { + let info = type_info($physical, None); + let mut rng = StdRng::seed_from_u64(42); + + let quiet_nan = <$float>::NAN; + let neg_nan = -<$float>::NAN; + let payload_nan = <$float>::from_bits(<$float>::NAN.to_bits() | 1); + let max_nan = <$float>::from_bits(<$uint>::MAX >> 1); + let edges = [ + <$float>::NEG_INFINITY, + <$float>::MIN, + -1.0, + -<$float>::MIN_POSITIVE, + -0.0, + 0.0, + <$float>::MIN_POSITIVE, + <$float>::from_bits(1), + 1.0, + <$float>::MAX, + <$float>::INFINITY, + quiet_nan, + neg_nan, + payload_nan, + max_nan, + ]; + let nans = [quiet_nan, neg_nan, payload_nan, max_nan]; + + let generators: [&dyn Fn(&mut StdRng) -> $float; 5] = [ + &|r| r.random::<$float>() * 2.0 - 1.0, + &|r| <$float>::from_bits(r.random()), + &|r| *edges.choose(r).unwrap(), + &|r| *nans.choose(r).unwrap(), + &|r| { + if r.random_bool(0.95) { + quiet_nan + } else { + r.random() + } + }, + ]; + + for generator in generators { + for len in lengths() { + let values: Vec<$float> = (0..len).map(|_| generator(&mut rng)).collect(); + let expected = fallback_min_max(&info, values.iter()); + match $kernel(&values) { + Some(actual) => { + // Compare bits to tell apart -0.0 / 0.0 + let to_bits = |(min, max, nans): ($float, $float, u64)| { + (min.to_bits(), max.to_bits(), nans) + }; + assert_eq!( + to_bits(actual), + to_bits(expected.unwrap()), + "len {len}: {values:?}" + ); + } + // Only all NaN (or empty) input defers to the scalar path + None => assert!(values.iter().all(|v| v.is_nan()), "{values:?}"), + } + } + } + } + }; + } + + float_tests!(test_min_max_f32, f32, u32, min_max_f32, PhysicalType::FLOAT); + float_tests!( + test_min_max_f64, + f64, + u64, + min_max_f64, + PhysicalType::DOUBLE + ); +} diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 432c97bfe8..4f8c04b4ac 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -53,6 +53,7 @@ use crate::schema::types::{BasicTypeInfo, ColumnDescPtr, ColumnDescriptor}; mod byte_budget_chunker; pub(crate) mod encoder; +pub(crate) mod min_max; use byte_budget_chunker::{ByteBudgetChunker, SubBatchStrategy}; diff --git a/parquet/src/data_type.rs b/parquet/src/data_type.rs index 3af6d426f3..4987e669fa 100644 --- a/parquet/src/data_type.rs +++ b/parquet/src/data_type.rs @@ -698,7 +698,11 @@ pub(crate) mod private { use super::{ParquetError, Result, SliceAsBytes}; use crate::basic::Type; + use crate::column::writer::min_max::{ + fallback_min_max, min_max_f32, min_max_f64, min_max_i32, min_max_i64, + }; use crate::file::metadata::HeapSize; + use crate::schema::types::BasicTypeInfo; /// Sealed trait to start to remove specialisation from implementations /// @@ -768,6 +772,12 @@ pub(crate) mod private { .map(|x| x as u64) } + /// Return the min, max and NaN count of `values` for column statistics, + /// or `None` if `values` is empty + fn min_max(basic_type_info: &BasicTypeInfo, values: &[Self]) -> Option<(Self, Self, u64)> { + fallback_min_max(basic_type_info, values.iter()) + } + /// Return the value as an Any to allow for downcasts without transmutation fn as_any(&self) -> &dyn std::any::Any; @@ -835,7 +845,7 @@ pub(crate) mod private { } macro_rules! impl_from_raw { - ($ty: ty, $physical_ty: expr, $self: ident => $as_i64: block) => { + ($ty: ty, $physical_ty: expr, $self: ident => $as_i64: block, min_max: ($info: ident, $values: ident) => $min_max: block) => { impl ParquetValueType for $ty { const PHYSICAL_TYPE: Type = $physical_ty; @@ -907,6 +917,11 @@ pub(crate) mod private { $as_i64 } + #[inline] + fn min_max($info: &BasicTypeInfo, $values: &[Self]) -> Option<(Self, Self, u64)> { + $min_max + } + #[inline] fn as_any(&self) -> &dyn std::any::Any { self @@ -920,10 +935,23 @@ pub(crate) mod private { } } - impl_from_raw!(i32, Type::INT32, self => { Ok(*self as i64) }); - impl_from_raw!(i64, Type::INT64, self => { Ok(*self) }); - impl_from_raw!(f32, Type::FLOAT, self => { Err(general_err!("Type cannot be converted to i64")) }); - impl_from_raw!(f64, Type::DOUBLE, self => { Err(general_err!("Type cannot be converted to i64")) }); + impl_from_raw!( + i32, Type::INT32, self => { Ok(*self as i64) }, + min_max: (info, values) => { min_max_i32(info, values).map(|(min, max)| (min, max, 0)) } + ); + impl_from_raw!( + i64, Type::INT64, self => { Ok(*self) }, + min_max: (info, values) => { min_max_i64(info, values).map(|(min, max)| (min, max, 0)) } + ); + // The float kernels return `None` for all NaN input, which the scalar path handles + impl_from_raw!( + f32, Type::FLOAT, self => { Err(general_err!("Type cannot be converted to i64")) }, + min_max: (info, values) => { min_max_f32(values).or_else(|| fallback_min_max(info, values.iter())) } + ); + impl_from_raw!( + f64, Type::DOUBLE, self => { Err(general_err!("Type cannot be converted to i64")) }, + min_max: (info, values) => { min_max_f64(values).or_else(|| fallback_min_max(info, values.iter())) } + ); impl ParquetValueType for super::Int96 { const PHYSICAL_TYPE: Type = Type::INT96;
