diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index d9adacff4101..1e74d6bfbb45 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -21,7 +21,8 @@ use half::f16; 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, + compare_greater_with, fallback_encoding, has_dictionary_support, is_nan_with, update_max, + update_min, }; use crate::data_type::DataType; use crate::data_type::private::ParquetValueType; @@ -400,9 +401,12 @@ where T: ParquetValueType + 'a, I: Iterator, { + let logical_type = descr.logical_type_ref(); + let converted_type = descr.converted_type(); + let first = loop { let next = iter.next()?; - if !is_nan(descr, next) { + if !is_nan_with(logical_type, next) { break next; } }; @@ -410,13 +414,13 @@ where let mut min = first; let mut max = first; for val in iter { - if is_nan(descr, val) { + if is_nan_with(logical_type, val) { continue; } - if compare_greater(descr, min, val) { + if compare_greater_with(logical_type, converted_type, min, val) { min = val; } - if compare_greater(descr, val, max) { + if compare_greater_with(logical_type, converted_type, val, max) { max = val; } } @@ -429,14 +433,18 @@ where // // For max, it has similar logic but will be written as 0.0 // (positive zero) - let min = replace_zero(min, descr, -0.0); - let max = replace_zero(max, descr, 0.0); + let min = replace_zero(min, logical_type, -0.0); + let max = replace_zero(max, logical_type, 0.0); Some((min, max)) } -#[inline] -fn replace_zero(val: &T, descr: &ColumnDescriptor, replace: f32) -> T { +#[inline(always)] +fn replace_zero( + val: &T, + logical_type_ref: Option<&LogicalType>, + replace: f32, +) -> T { match T::PHYSICAL_TYPE { Type::FLOAT if f32::from_le_bytes(val.as_bytes().try_into().unwrap()) == 0.0 => { T::try_from_le_slice(&f32::to_le_bytes(replace)).unwrap() @@ -445,7 +453,7 @@ fn replace_zero(val: &T, descr: &ColumnDescriptor, replace: T::try_from_le_slice(&f64::to_le_bytes(replace as f64)).unwrap() } Type::FIXED_LEN_BYTE_ARRAY - if descr.logical_type_ref() == Some(LogicalType::Float16).as_ref() + if logical_type_ref == Some(LogicalType::Float16).as_ref() && f16::from_le_bytes(val.as_bytes().try_into().unwrap()) == f16::NEG_ZERO => { T::try_from_le_slice(&f16::to_le_bytes(f16::from_f32(replace))).unwrap() diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index aa9cef16c5ad..1c649548c775 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -1649,12 +1649,16 @@ fn update_max(descr: &ColumnDescriptor, val: &T, max: &mut update_stat::(descr, val, max, |cur| compare_greater(descr, val, cur)) } -#[inline] -#[allow(clippy::eq_op)] fn is_nan(descr: &ColumnDescriptor, val: &T) -> bool { + is_nan_with(descr.logical_type_ref(), val) +} + +#[inline(always)] +#[allow(clippy::eq_op)] +fn is_nan_with(logical_type_ref: Option<&LogicalType>, val: &T) -> bool { match T::PHYSICAL_TYPE { Type::FLOAT | Type::DOUBLE => val != val, - Type::FIXED_LEN_BYTE_ARRAY if descr.logical_type_ref() == Some(&LogicalType::Float16) => { + Type::FIXED_LEN_BYTE_ARRAY if logical_type_ref == Some(&LogicalType::Float16) => { let val = val.as_bytes(); let val = f16::from_le_bytes([val[0], val[1]]); val.is_nan() @@ -1685,18 +1689,29 @@ fn update_stat( } /// Evaluate `a > b` according to underlying logical type. +#[inline] fn compare_greater(descr: &ColumnDescriptor, a: &T, b: &T) -> bool { + compare_greater_with(descr.logical_type_ref(), descr.converted_type(), a, b) +} + +#[inline(always)] +fn compare_greater_with( + logical_type_ref: Option<&LogicalType>, + converted_type: ConvertedType, + a: &T, + b: &T, +) -> bool { match T::PHYSICAL_TYPE { Type::INT32 | Type::INT64 => { if let Some(LogicalType::Integer(IntType { is_signed: false, .. - })) = descr.logical_type_ref() + })) = logical_type_ref { // need to compare unsigned return compare_greater_unsigned_int(a, b); } - match descr.converted_type() { + match converted_type { ConvertedType::UINT_8 | ConvertedType::UINT_16 | ConvertedType::UINT_32 @@ -1707,13 +1722,13 @@ fn compare_greater(descr: &ColumnDescriptor, a: &T, b: &T) }; } Type::FIXED_LEN_BYTE_ARRAY | Type::BYTE_ARRAY => { - if let Some(LogicalType::Decimal(_)) = descr.logical_type_ref() { + if let Some(LogicalType::Decimal(_)) = logical_type_ref { return compare_greater_byte_array_decimals(a.as_bytes(), b.as_bytes()); } - if let ConvertedType::DECIMAL = descr.converted_type() { + if let ConvertedType::DECIMAL = converted_type { return compare_greater_byte_array_decimals(a.as_bytes(), b.as_bytes()); } - if let Some(LogicalType::Float16) = descr.logical_type_ref() { + if let Some(LogicalType::Float16) = logical_type_ref { return compare_greater_f16(a.as_bytes(), b.as_bytes()); } }