From ae27070e71ececd9daefccc99b92394f361c87d4 Mon Sep 17 00:00:00 2001 From: Congxian Qiu Date: Tue, 19 May 2026 02:55:20 +0800 Subject: [PATCH 01/51] [Variant] Align cast logic for from/to_decimal for variant (#9689) # Which issue does this PR close? - Closes #9688 . # What changes are included in this PR? - Extract some logic in arrow-cast - Reuse the extracted logic in arrow-cast and parquet-variant # Are these changes tested? Reuse the existing tests in arrow-test # Are there any user-facing changes? Yes, changed the docs --- arrow-cast/src/cast/decimal.rs | 17 +- arrow-cast/src/cast/mod.rs | 23 +- .../src/type_conversion.rs | 20 +- parquet-variant/src/variant.rs | 231 ++++++++++++++---- 4 files changed, 230 insertions(+), 61 deletions(-) diff --git a/arrow-cast/src/cast/decimal.rs b/arrow-cast/src/cast/decimal.rs index 3553f2b6a76f..789dcea89ed1 100644 --- a/arrow-cast/src/cast/decimal.rs +++ b/arrow-cast/src/cast/decimal.rs @@ -531,7 +531,7 @@ where /// Parses given string to specified decimal native (i128/i256) based on given /// scale. Returns an `Err` if it cannot parse given string. -pub(crate) fn parse_string_to_decimal_native( +pub fn parse_string_to_decimal_native( value_str: &str, scale: usize, ) -> Result @@ -777,7 +777,7 @@ where if cast_options.safe { array .unary_opt::<_, D>(|v| { - D::Native::from_f64((mul * v.as_()).round()) + single_float_to_decimal::(v.as_(), mul) .filter(|v| D::is_valid_decimal_precision(*v, precision)) }) .with_precision_and_scale(precision, scale) @@ -785,7 +785,7 @@ where } else { array .try_unary::<_, D, _>(|v| { - D::Native::from_f64((mul * v.as_()).round()) + single_float_to_decimal::(v.as_(), mul) .ok_or_else(|| { ArrowError::CastError(format!( "Cannot cast to {}({}, {}). Overflowing on {:?}", @@ -802,6 +802,17 @@ where } } +/// Cast a single floating point value to a decimal native with the given multiple. +/// Returns `None` if the value cannot be represented with the requested precision. +#[inline(always)] +pub fn single_float_to_decimal(input: f64, mul: f64) -> Option +where + D: DecimalType + ArrowPrimitiveType, + ::Native: DecimalCast, +{ + D::Native::from_f64((mul * input).round()) +} + pub(crate) fn cast_decimal_to_integer( array: &dyn Array, base: D::Native, diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs index 2d9923c9a299..e100b22d728f 100644 --- a/arrow-cast/src/cast/mod.rs +++ b/arrow-cast/src/cast/mod.rs @@ -72,9 +72,25 @@ use arrow_schema::*; use arrow_select::take::take; use num_traits::{NumCast, ToPrimitive, cast::AsPrimitive}; -pub use decimal::{DecimalCast, rescale_decimal}; +pub use decimal::{ + DecimalCast, parse_string_to_decimal_native, rescale_decimal, single_float_to_decimal, +}; pub use string::cast_single_string_to_boolean_default; +/// Lossy conversion from decimal to float. +/// +/// Conversion is lossy and follows standard floating point semantics. Values +/// that exceed the representable range become `INFINITY` or `-INFINITY` without +/// returning an error. +#[inline(always)] +pub fn single_decimal_to_float_lossy(f: &F, x: D::Native, scale: i32) -> f64 +where + D: DecimalType, + F: Fn(D::Native) -> f64, +{ + f(x) / 10_f64.powi(scale) +} + /// CastOptions provides a way to override the default cast behaviors #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct CastOptions<'a> { @@ -2314,10 +2330,11 @@ where Int32 => cast_decimal_to_integer::(array, base, *scale, cast_options), Int64 => cast_decimal_to_integer::(array, base, *scale, cast_options), Float32 => cast_decimal_to_float::(array, |x| { - (as_float(x) / 10_f64.powi(*scale as i32)) as f32 + single_decimal_to_float_lossy::(&as_float, x, >::from(*scale)) + as f32 }), Float64 => cast_decimal_to_float::(array, |x| { - as_float(x) / 10_f64.powi(*scale as i32) + single_decimal_to_float_lossy::(&as_float, x, >::from(*scale)) }), Utf8View => value_to_string_view(array, cast_options), Utf8 => value_to_string::(array, cast_options), diff --git a/parquet-variant-compute/src/type_conversion.rs b/parquet-variant-compute/src/type_conversion.rs index 7b9eb67d1a95..2255d4316b25 100644 --- a/parquet-variant-compute/src/type_conversion.rs +++ b/parquet-variant-compute/src/type_conversion.rs @@ -17,7 +17,10 @@ //! Module for transforming a typed arrow `Array` to `VariantArray`. -use arrow::compute::{CastOptions, DecimalCast, rescale_decimal}; +use arrow::compute::{ + CastOptions, DecimalCast, parse_string_to_decimal_native, rescale_decimal, + single_float_to_decimal, +}; use arrow::datatypes::{ self, ArrowPrimitiveType, ArrowTimestampType, Decimal32Type, Decimal64Type, Decimal128Type, DecimalType, @@ -204,9 +207,12 @@ impl_timestamp_from_variant!( /// /// - `precision` and `scale` specify the target Arrow decimal parameters /// - Integer variants (`Int8/16/32/64`) are treated as decimals with scale 0 +/// - Floating point variants (`Float/Double`) are converted to decimals with the given scale +/// - String variants (`String/ShortString`) are parsed as decimals with the given scale /// - Decimal variants (`Decimal4/8/16`) use their embedded precision and scale /// -/// The value is rescaled to (`precision`, `scale`) using `rescale_decimal` and +/// The value is rescaled to (`precision`, `scale`) using `rescale_decimal` for integers, +/// `single_float_to_decimal` for floats, and `parse_string_to_decimal_native` for strings. /// returns `None` if it cannot fit the requested precision. pub(crate) fn variant_to_unscaled_decimal( variant: &Variant<'_, '_>, @@ -217,6 +223,8 @@ where O: DecimalType, O::Native: DecimalCast, { + let mul = 10_f64.powi(scale as i32); + match variant { Variant::Int8(i) => rescale_decimal::( *i as i32, @@ -246,6 +254,14 @@ where precision, scale, ), + Variant::Float(f) => single_float_to_decimal::(f64::from(*f), mul), + Variant::Double(f) => single_float_to_decimal::(*f, mul), + // arrow-cast only support cast string to decimal with scale >=0 for now + // Please see `cast_string_to_decimal` in arrow-cast/src/cast/decimal.rs for more detail + Variant::String(v) if scale >= 0 => parse_string_to_decimal_native::(v, scale as _).ok(), + Variant::ShortString(v) if scale >= 0 => { + parse_string_to_decimal_native::(v, scale as _).ok() + } Variant::Decimal4(d) => rescale_decimal::( d.integer(), VariantDecimal4::MAX_PRECISION, diff --git a/parquet-variant/src/variant.rs b/parquet-variant/src/variant.rs index accff009045a..3b5a04ddfd6c 100644 --- a/parquet-variant/src/variant.rs +++ b/parquet-variant/src/variant.rs @@ -29,9 +29,14 @@ use crate::decoder::{ }; use crate::path::{VariantPath, VariantPathElement}; use crate::utils::{first_byte_from_slice, slice_from_slice}; +use arrow::array::ArrowNativeTypeOp; use arrow::compute::{ - cast_num_to_bool, cast_single_string_to_boolean_default, num_cast, single_bool_to_numeric, + DecimalCast, cast_num_to_bool, cast_single_string_to_boolean_default, num_cast, + parse_string_to_decimal_native, single_bool_to_numeric, single_decimal_to_float_lossy, + single_float_to_decimal, }; +use arrow::datatypes::{Decimal32Type, Decimal64Type, Decimal128Type, DecimalType}; + use arrow_schema::ArrowError; use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Timelike, Utc}; use num_traits::NumCast; @@ -166,10 +171,11 @@ impl Deref for ShortString<'_> { /// Arrow UTF8-to-boolean cast rules. /// - Numeric accessors such as [`Self::as_int8`], [`Self::as_int64`], [`Self::as_u8`], /// [`Self::as_u64`], [`Self::as_f16`], [`Self::as_f32`], and [`Self::as_f64`] accept -/// boolean and numeric variants (integers, floating-point, and decimals with scale `0`). +/// boolean and numeric variants (integers, floating-point, and decimals). /// They return `None` when conversion is not possible. /// - Decimal accessors such as [`Self::as_decimal4`], [`Self::as_decimal8`], and -/// [`Self::as_decimal16`] accept compatible decimal variants and integer variants. +/// [`Self::as_decimal16`] accept compatible decimal variants, integer variants, +/// float variants and string variants. /// They return `None` when conversion is not possible. /// /// # Examples: @@ -294,6 +300,35 @@ pub enum Variant<'m, 'v> { // We don't want this to grow because it could hurt performance of a frequently-created type. const _: () = crate::utils::expect_size_of::(80); +enum NumericKind { + Integer, + Float, +} + +trait DecimalCastTarget: NumCast + Default { + const KIND: NumericKind; +} + +macro_rules! impl_decimal_cast_target { + ($raw_type: ident, $target_kind:expr) => { + impl DecimalCastTarget for $raw_type { + const KIND: NumericKind = $target_kind; + } + }; +} + +impl_decimal_cast_target!(i8, NumericKind::Integer); +impl_decimal_cast_target!(i16, NumericKind::Integer); +impl_decimal_cast_target!(i32, NumericKind::Integer); +impl_decimal_cast_target!(i64, NumericKind::Integer); +impl_decimal_cast_target!(u8, NumericKind::Integer); +impl_decimal_cast_target!(u16, NumericKind::Integer); +impl_decimal_cast_target!(u32, NumericKind::Integer); +impl_decimal_cast_target!(u64, NumericKind::Integer); +impl_decimal_cast_target!(f16, NumericKind::Float); +impl_decimal_cast_target!(f32, NumericKind::Float); +impl_decimal_cast_target!(f64, NumericKind::Float); + impl<'m, 'v> Variant<'m, 'v> { /// Attempts to interpret a metadata and value buffer pair as a new `Variant`. /// @@ -797,14 +832,37 @@ impl<'m, 'v> Variant<'m, 'v> { } } - /// Converts a boolean or numeric variant(integers, floating-point, and decimals with scale 0) + fn cast_decimal_to_num(raw: D::Native, scale: u8, as_float: F) -> Option + where + D: DecimalType, + D::Native: NumCast + ArrowNativeTypeOp, + T: DecimalCastTarget, + F: Fn(D::Native) -> f64, + { + let base: D::Native = NumCast::from(10)?; + + let div = base.pow_checked(>::from(scale)).ok()?; + match T::KIND { + NumericKind::Integer => raw + .div_checked(div) + .ok() + .and_then(::from::), + NumericKind::Float => T::from(single_decimal_to_float_lossy::( + &as_float, + raw, + >::from(scale), + )), + } + } + + /// Converts a boolean or numeric variant(integers, floating-point, and decimals) /// to the specified numeric type `T`. /// /// Uses Arrow's casting logic to perform the conversion. Returns `Some(T)` if /// the conversion succeeds, `None` if the variant can't be casted to type `T`. fn as_num(&self) -> Option where - T: NumCast + Default, + T: DecimalCastTarget, { match *self { Variant::BooleanFalse => single_bool_to_numeric(false), @@ -815,9 +873,21 @@ impl<'m, 'v> Variant<'m, 'v> { Variant::Int64(i) => num_cast(i), Variant::Float(f) => num_cast(f), Variant::Double(d) => num_cast(d), - Variant::Decimal4(d) if d.scale() == 0 => num_cast(d.integer()), - Variant::Decimal8(d) if d.scale() == 0 => num_cast(d.integer()), - Variant::Decimal16(d) if d.scale() == 0 => num_cast(d.integer()), + Variant::Decimal4(d) => { + Self::cast_decimal_to_num::(d.integer(), d.scale(), |x| { + x as f64 + }) + } + Variant::Decimal8(d) => { + Self::cast_decimal_to_num::(d.integer(), d.scale(), |x| { + x as f64 + }) + } + Variant::Decimal16(d) => { + Self::cast_decimal_to_num::(d.integer(), d.scale(), |x| { + x as f64 + }) + } _ => None, } } @@ -962,17 +1032,17 @@ impl<'m, 'v> Variant<'m, 'v> { /// let v2 = Variant::from(d); /// assert_eq!(v2.as_u8(), Some(26u8)); /// + /// // or a variant that decimal with scale not equal to zero + /// let d = VariantDecimal4::try_new(123, 2).unwrap(); + /// let v3 = Variant::from(d); + /// assert_eq!(v3.as_u8(), Some(1)); + /// /// // or from boolean variant - /// let v3 = Variant::BooleanFalse; - /// assert_eq!(v3.as_u8(), Some(0)); + /// let v4 = Variant::BooleanFalse; + /// assert_eq!(v4.as_u8(), Some(0)); /// /// // but not a variant that can't fit into the range - /// let v4 = Variant::from(-1); - /// assert_eq!(v4.as_u8(), None); - /// - /// // not a variant that decimal with scale not equal to zero - /// let d = VariantDecimal4::try_new(1, 2).unwrap(); - /// let v5 = Variant::from(d); + /// let v5 = Variant::from(-1); /// assert_eq!(v5.as_u8(), None); /// /// // or not a variant that cannot be cast into an integer @@ -1003,17 +1073,17 @@ impl<'m, 'v> Variant<'m, 'v> { /// let v2 = Variant::from(d); /// assert_eq!(v2.as_u16(), Some(u16::MAX)); /// + /// // or a variant that decimal with scale not equal to zero + /// let d = VariantDecimal4::try_new(123, 2).unwrap(); + /// let v3 = Variant::from(d); + /// assert_eq!(v3.as_u16(), Some(1)); + /// /// // or from boolean variant - /// let v3= Variant::BooleanFalse; - /// assert_eq!(v3.as_u16(), Some(0)); + /// let v4= Variant::BooleanFalse; + /// assert_eq!(v4.as_u16(), Some(0)); /// /// // but not a variant that can't fit into the range - /// let v4 = Variant::from(-1); - /// assert_eq!(v4.as_u16(), None); - /// - /// // not a variant that decimal with scale not equal to zero - /// let d = VariantDecimal4::try_new(1, 2).unwrap(); - /// let v5 = Variant::from(d); + /// let v5 = Variant::from(-1); /// assert_eq!(v5.as_u16(), None); /// /// // or not a variant that cannot be cast into an integer @@ -1044,17 +1114,17 @@ impl<'m, 'v> Variant<'m, 'v> { /// let v2 = Variant::from(d); /// assert_eq!(v2.as_u32(), Some(u32::MAX)); /// + /// // or a variant that decimal with scale not equal to zero + /// let d = VariantDecimal8::try_new(123, 2).unwrap(); + /// let v3 = Variant::from(d); + /// assert_eq!(v3.as_u32(), Some(1)); + /// /// // or from boolean variant - /// let v3 = Variant::BooleanFalse; - /// assert_eq!(v3.as_u32(), Some(0)); + /// let v4 = Variant::BooleanFalse; + /// assert_eq!(v4.as_u32(), Some(0)); /// /// // but not a variant that can't fit into the range - /// let v4 = Variant::from(-1); - /// assert_eq!(v4.as_u32(), None); - /// - /// // not a variant that decimal with scale not equal to zero - /// let d = VariantDecimal8::try_new(1, 2).unwrap(); - /// let v5 = Variant::from(d); + /// let v5 = Variant::from(-1); /// assert_eq!(v5.as_u32(), None); /// /// // or not a variant that cannot be cast into an integer @@ -1085,17 +1155,17 @@ impl<'m, 'v> Variant<'m, 'v> { /// let v2 = Variant::from(d); /// assert_eq!(v2.as_u64(), Some(u64::MAX)); /// + /// // or a variant that decimal with scale not equal to zero + /// let d = VariantDecimal16::try_new(123, 2).unwrap(); + /// let v3 = Variant::from(d); + /// assert_eq!(v3.as_u64(), Some(1)); + /// /// // or from boolean variant - /// let v3 = Variant::BooleanFalse; - /// assert_eq!(v3.as_u64(), Some(0)); + /// let v4 = Variant::BooleanFalse; + /// assert_eq!(v4.as_u64(), Some(0)); /// /// // but not a variant that can't fit into the range - /// let v4 = Variant::from(-1); - /// assert_eq!(v4.as_u64(), None); - /// - /// // not a variant that decimal with scale not equal to zero - /// let d = VariantDecimal16::try_new(1, 2).unwrap(); - /// let v5 = Variant::from(d); + /// let v5 = Variant::from(-1); /// assert_eq!(v5.as_u64(), None); /// /// // or not a variant that cannot be cast into an integer @@ -1106,6 +1176,21 @@ impl<'m, 'v> Variant<'m, 'v> { self.as_num() } + fn convert_string_to_decimal(input: &str) -> Option + where + D: DecimalType, + VD: VariantDecimalType, + D::Native: NumCast + DecimalCast, + { + // find the last '.' + let scale_usize = input.rsplit_once('.').map_or(0, |(_, frac)| frac.len()); + + let scale = u8::try_from(scale_usize).ok()?; + + let raw = parse_string_to_decimal_native::(input, scale_usize).ok()?; + VD::try_new(raw, scale).ok() + } + /// Converts this variant to tuple with a 4-byte unscaled value if possible. /// /// Returns `Some((i32, u8))` for decimal variants where the unscaled value @@ -1125,19 +1210,31 @@ impl<'m, 'v> Variant<'m, 'v> { /// let v2 = Variant::from(VariantDecimal8::try_new(1234_i64, 2).unwrap()); /// assert_eq!(v2.as_decimal4(), VariantDecimal4::try_new(1234_i32, 2).ok()); /// + /// // or from string variants if they can be parsed as decimals + /// let v3 = Variant::from("123.45"); + /// assert_eq!(v3.as_decimal4(), VariantDecimal4::try_new(12345, 2).ok()); + /// /// // but not if the value would overflow i32 - /// let v3 = Variant::from(VariantDecimal8::try_new(12345678901i64, 2).unwrap()); - /// assert_eq!(v3.as_decimal4(), None); + /// let v4 = Variant::from(VariantDecimal8::try_new(12345678901i64, 2).unwrap()); + /// assert_eq!(v4.as_decimal4(), None); /// /// // or if the variant is not a decimal - /// let v4 = Variant::from("hello!"); - /// assert_eq!(v4.as_decimal4(), None); + /// let v5 = Variant::from("hello!"); + /// assert_eq!(v5.as_decimal4(), None); /// ``` pub fn as_decimal4(&self) -> Option { match *self { Variant::Int8(_) | Variant::Int16(_) | Variant::Int32(_) | Variant::Int64(_) => { self.as_num::().and_then(|x| x.try_into().ok()) } + Variant::Float(f) => single_float_to_decimal::(f as _, 1f64) + .and_then(|x: i32| x.try_into().ok()), + Variant::Double(f) => single_float_to_decimal::(f, 1f64) + .and_then(|x: i32| x.try_into().ok()), + Variant::String(v) => Self::convert_string_to_decimal::(v), + Variant::ShortString(v) => { + Self::convert_string_to_decimal::(v.as_str()) + } Variant::Decimal4(decimal4) => Some(decimal4), Variant::Decimal8(decimal8) => decimal8.try_into().ok(), Variant::Decimal16(decimal16) => decimal16.try_into().ok(), @@ -1148,7 +1245,7 @@ impl<'m, 'v> Variant<'m, 'v> { /// Converts this variant to tuple with an 8-byte unscaled value if possible. /// /// Returns `Some((i64, u8))` for decimal variants where the unscaled value - /// fits in `i64` range, + /// fits in `i64` range, the scale will be 0 if the input is string variants. /// `None` for non-decimal variants or decimal values that would overflow. /// /// # Examples @@ -1164,19 +1261,31 @@ impl<'m, 'v> Variant<'m, 'v> { /// let v2 = Variant::from(VariantDecimal16::try_new(1234_i128, 2).unwrap()); /// assert_eq!(v2.as_decimal8(), VariantDecimal8::try_new(1234_i64, 2).ok()); /// + /// // or from string variants if they can be parsed as decimals + /// let v3 = Variant::from("123.45"); + /// assert_eq!(v3.as_decimal8(), VariantDecimal8::try_new(12345, 2).ok()); + /// /// // but not if the value would overflow i64 - /// let v3 = Variant::from(VariantDecimal16::try_new(2e19 as i128, 2).unwrap()); - /// assert_eq!(v3.as_decimal8(), None); + /// let v4 = Variant::from(VariantDecimal16::try_new(2e19 as i128, 2).unwrap()); + /// assert_eq!(v4.as_decimal8(), None); /// /// // or if the variant is not a decimal - /// let v4 = Variant::from("hello!"); - /// assert_eq!(v4.as_decimal8(), None); + /// let v5 = Variant::from("hello!"); + /// assert_eq!(v5.as_decimal8(), None); /// ``` pub fn as_decimal8(&self) -> Option { match *self { Variant::Int8(_) | Variant::Int16(_) | Variant::Int32(_) | Variant::Int64(_) => { self.as_num::().and_then(|x| x.try_into().ok()) } + Variant::Float(f) => single_float_to_decimal::(f as _, 1f64) + .and_then(|x: i64| x.try_into().ok()), + Variant::Double(f) => single_float_to_decimal::(f, 1f64) + .and_then(|x: i64| x.try_into().ok()), + Variant::String(v) => Self::convert_string_to_decimal::(v), + Variant::ShortString(v) => { + Self::convert_string_to_decimal::(v.as_str()) + } Variant::Decimal4(decimal4) => Some(decimal4.into()), Variant::Decimal8(decimal8) => Some(decimal8), Variant::Decimal16(decimal16) => decimal16.try_into().ok(), @@ -1187,7 +1296,7 @@ impl<'m, 'v> Variant<'m, 'v> { /// Converts this variant to tuple with a 16-byte unscaled value if possible. /// /// Returns `Some((i128, u8))` for decimal variants where the unscaled value - /// fits in `i128` range, + /// fits in `i128` range, the scale will be 0 if the input is string variants. /// `None` for non-decimal variants or decimal values that would overflow. /// /// # Examples @@ -1199,14 +1308,30 @@ impl<'m, 'v> Variant<'m, 'v> { /// let v1 = Variant::from(VariantDecimal4::try_new(1234_i32, 2).unwrap()); /// assert_eq!(v1.as_decimal16(), VariantDecimal16::try_new(1234_i128, 2).ok()); /// + /// // or from a string variant if it can be parsed as decimal + /// let v2 = Variant::from("123.45"); + /// assert_eq!(v2.as_decimal16(), VariantDecimal16::try_new(12345, 2).ok()); + /// /// // but not if the variant is not a decimal - /// let v2 = Variant::from("hello!"); - /// assert_eq!(v2.as_decimal16(), None); + /// let v3 = Variant::from("hello!"); + /// assert_eq!(v3.as_decimal16(), None); /// ``` pub fn as_decimal16(&self) -> Option { match *self { Variant::Int8(_) | Variant::Int16(_) | Variant::Int32(_) | Variant::Int64(_) => { - self.as_num::().and_then(|x| x.try_into().ok()) + let x = self.as_num::()?; + >::from(x).try_into().ok() + } + Variant::Float(f) => { + single_float_to_decimal::(>::from(f), 1f64) + .and_then(|x| x.try_into().ok()) + } + Variant::Double(f) => { + single_float_to_decimal::(f, 1f64).and_then(|x| x.try_into().ok()) + } + Variant::String(v) => Self::convert_string_to_decimal::(v), + Variant::ShortString(v) => { + Self::convert_string_to_decimal::(v.as_str()) } Variant::Decimal4(decimal4) => Some(decimal4.into()), Variant::Decimal8(decimal8) => Some(decimal8.into()), From da09872193bb3805a3177678bf9005ddeea2333d Mon Sep 17 00:00:00 2001 From: Vegard Stikbakke Date: Wed, 20 May 2026 03:44:03 +0200 Subject: [PATCH 02/51] Implement native interleave for ListView (#9558) This PR adds a native implementation of interleave for the ListView type which uses a good heuristic thanks to @asubiotto, either 1. copy each row's elements and put them all into a new flat array, or 2. concatenate all source value array (and adjust offsets). The latter is best when there is sharing of elements. Closes #9342. --------- Signed-off-by: Alfonso Subiotto Marques Co-authored-by: Alfonso Subiotto Marques --- arrow-select/src/interleave.rs | 232 +++++++++++++++++++++++++++++++++ 1 file changed, 232 insertions(+) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 1bd34a7f3a58..13d28747dfcf 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -114,6 +114,8 @@ pub fn interleave( DataType::Int64 => interleave_run_end::(values, indices), t => unreachable!("illegal run-end type {t}"), }, + DataType::ListView(field) => interleave_list_view::(values, indices, field), + DataType::LargeListView(field) => interleave_list_view::(values, indices, field), _ => interleave_fallback(values, indices) } } @@ -482,6 +484,114 @@ fn interleave_run_end( )?)) } +fn interleave_list_view( + values: &[&dyn Array], + indices: &[(usize, usize)], + field: &FieldRef, +) -> Result { + let interleaved = Interleave::<'_, GenericListViewArray>::new(values, indices); + + // Pick whichever strategy produces fewer child elements: + // - Per-row copy: total = sum of selected sizes. Better for sparse selections. + // - Concat + offset adjustment: total = sum of source backing array lengths. + // Better when rows share backing elements via overlapping offset/size ranges. + let concat_cost: usize = interleaved.arrays.iter().map(|lv| lv.values().len()).sum(); + let per_row_cost: usize = indices + .iter() + .map(|&(a, r)| interleaved.arrays[a].sizes()[r].as_usize()) + .sum(); + + if per_row_cost <= concat_cost { + interleave_list_view_copy::(&interleaved, indices, field) + } else { + interleave_list_view_concat::(&interleaved, indices, field) + } +} + +/// Per-row copy: copies each selected row's child elements into a new flat array. +fn interleave_list_view_copy( + interleaved: &Interleave<'_, GenericListViewArray>, + indices: &[(usize, usize)], + field: &FieldRef, +) -> Result { + let mut capacity = 0usize; + let mut offsets = Vec::with_capacity(indices.len()); + let mut sizes = Vec::with_capacity(indices.len()); + for &(array_idx, row_idx) in indices { + let list = interleaved.arrays[array_idx]; + let size = list.sizes()[row_idx].as_usize(); + offsets.push( + O::from_usize(capacity).ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?, + ); + sizes.push(O::from_usize(size).ok_or_else(|| ArrowError::OffsetOverflowError(size))?); + capacity += size; + } + + let child_data: Vec<_> = interleaved + .arrays + .iter() + .map(|list| list.values().to_data()) + .collect(); + let child_data_refs: Vec<_> = child_data.iter().collect(); + let mut mutable_child = MutableArrayData::new(child_data_refs, false, capacity); + for &(array_idx, row_idx) in indices { + let list = interleaved.arrays[array_idx]; + let start = list.offsets()[row_idx].as_usize(); + let size = list.sizes()[row_idx].as_usize(); + if size > 0 { + mutable_child.extend(array_idx, start, start + size); + } + } + + Ok(Arc::new(GenericListViewArray::::new( + field.clone(), + offsets.into(), + sizes.into(), + make_array(mutable_child.freeze()), + interleaved.nulls.clone(), + ))) +} + +/// Concat backing arrays: concatenates all source value arrays and adjusts offsets. +/// Preserves within-source element sharing. +fn interleave_list_view_concat( + interleaved: &Interleave<'_, GenericListViewArray>, + indices: &[(usize, usize)], + field: &FieldRef, +) -> Result { + let child_arrays: Vec<&dyn Array> = interleaved + .arrays + .iter() + .map(|lv| lv.values().as_ref()) + .collect(); + let mut base_offsets = Vec::with_capacity(interleaved.arrays.len()); + let mut running = 0usize; + for lv in &interleaved.arrays { + base_offsets.push(running); + running += lv.values().len(); + } + let combined_values = concat(&child_arrays)?; + + let mut new_offsets = Vec::with_capacity(indices.len()); + let mut new_sizes = Vec::with_capacity(indices.len()); + for &(array_idx, row_idx) in indices { + let lv = interleaved.arrays[array_idx]; + let adjusted = lv.offsets()[row_idx].as_usize() + base_offsets[array_idx]; + new_offsets.push( + O::from_usize(adjusted).ok_or_else(|| ArrowError::OffsetOverflowError(adjusted))?, + ); + new_sizes.push(lv.sizes()[row_idx]); + } + + Ok(Arc::new(GenericListViewArray::::new( + field.clone(), + new_offsets.into(), + new_sizes.into(), + combined_values, + interleaved.nulls.clone(), + ))) +} + /// Fallback implementation of interleave using [`MutableArrayData`] fn interleave_fallback( values: &[&dyn Array], @@ -841,6 +951,128 @@ mod tests { test_interleave_lists::(); } + fn test_interleave_list_views() { + // [[1, 2], null, [3]] + let mut a = GenericListBuilder::::new(Int32Builder::new()); + a.values().append_value(1); + a.values().append_value(2); + a.append(true); + a.append(false); + a.values().append_value(3); + a.append(true); + let a: GenericListViewArray = a.finish().into(); + + // [[4], null, [5, 6, null]] + let mut b = GenericListBuilder::::new(Int32Builder::new()); + b.values().append_value(4); + b.append(true); + b.append(false); + b.values().append_value(5); + b.values().append_value(6); + b.values().append_null(); + b.append(true); + let b: GenericListViewArray = b.finish().into(); + + let values = interleave(&[&a, &b], &[(0, 2), (0, 1), (1, 0), (1, 2), (1, 1)]).unwrap(); + let v = values + .as_any() + .downcast_ref::>() + .unwrap(); + + // [[3], null, [4], [5, 6, null], null] + let mut expected = GenericListBuilder::::new(Int32Builder::new()); + expected.values().append_value(3); + expected.append(true); + expected.append(false); + expected.values().append_value(4); + expected.append(true); + expected.values().append_value(5); + expected.values().append_value(6); + expected.values().append_null(); + expected.append(true); + expected.append(false); + let expected: GenericListViewArray = expected.finish().into(); + + assert_eq!(v, &expected); + } + + #[test] + fn test_list_views() { + test_interleave_list_views::(); + } + + #[test] + fn test_large_list_views() { + test_interleave_list_views::(); + } + + #[test] + fn test_interleave_list_view_overlapping() { + let field = Arc::new(Field::new_list_field(DataType::Int64, false)); + + // lv_a: 10 rows, two groups of 5 sharing the same backing elements. + // rows 0-4 → offset 0, size 5 → [0,1,2,3,4] + // rows 5-9 → offset 5, size 5 → [5,6,7,8,9] + let lv_a = ListViewArray::new( + Arc::clone(&field), + ScalarBuffer::from(vec![0i32, 0, 0, 0, 0, 5, 5, 5, 5, 5]), + ScalarBuffer::from(vec![5i32; 10]), + Arc::new(Int64Array::from_iter_values(0..10)), + None, + ); + + // lv_b: 8 rows, two groups of 4 sharing the same backing elements. + // rows 0-3 → offset 0, size 3 → [100,101,102] + // rows 4-7 → offset 3, size 3 → [103,104,105] + let lv_b = ListViewArray::new( + Arc::clone(&field), + ScalarBuffer::from(vec![0i32, 0, 0, 0, 3, 3, 3, 3]), + ScalarBuffer::from(vec![3i32; 8]), + Arc::new(Int64Array::from_iter_values(100..106)), + None, + ); + + let indices: Vec<(usize, usize)> = vec![ + (0, 0), + (1, 0), + (0, 5), + (1, 4), + (0, 1), + (1, 1), + (0, 6), + (1, 5), + ]; + let result = interleave(&[&lv_a as &dyn Array, &lv_b as &dyn Array], &indices).unwrap(); + result + .to_data() + .validate_full() + .expect("result must be valid"); + + let result_lv = result.as_list_view::(); + assert_eq!(result_lv.len(), 8); + assert_eq!( + result_lv.value(0).as_primitive::().values(), + &[0, 1, 2, 3, 4] + ); + assert_eq!( + result_lv.value(1).as_primitive::().values(), + &[100, 101, 102] + ); + assert_eq!( + result_lv.value(2).as_primitive::().values(), + &[5, 6, 7, 8, 9] + ); + assert_eq!( + result_lv.value(3).as_primitive::().values(), + &[103, 104, 105] + ); + + // Backing elements = sum of source arrays (10 + 6 = 16), not per-row + // expansion (8 rows × avg ~4 = 32). Overlapping sharing is preserved. + let total_input_elements = lv_a.values().len() + lv_b.values().len(); + assert_eq!(result_lv.values().len(), total_input_elements); + } + #[test] fn test_struct_without_nulls() { let fields = Fields::from(vec![ From d48c3057e1e8113463fe5d3cee3ede6db092eaf2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alfonso=20Subiotto=20Marqu=C3=A9s?= Date: Wed, 20 May 2026 14:56:04 +0200 Subject: [PATCH 03/51] fix(arrow-schema): allow empty metadata value for UUID extension type (#10001) # Which issue does this PR close? - Closes #10000 # Rationale for this change Some ipc writers such as arrow-go unconditionally write an empty metadata value, which breaks deserialization. # What changes are included in this PR? Empty check # Are these changes tested? Yes # Are there any user-facing changes? Signed-off-by: Alfonso Subiotto Marques --- arrow-schema/src/extension/canonical/uuid.rs | 31 ++++++++++++++++---- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/arrow-schema/src/extension/canonical/uuid.rs b/arrow-schema/src/extension/canonical/uuid.rs index affea04acb3c..7a9f6e6e538e 100644 --- a/arrow-schema/src/extension/canonical/uuid.rs +++ b/arrow-schema/src/extension/canonical/uuid.rs @@ -53,10 +53,14 @@ impl ExtensionType for Uuid { fn deserialize_metadata(metadata: Option<&str>) -> Result { metadata.map_or_else( || Ok(()), - |_| { - Err(ArrowError::InvalidArgumentError( - "Uuid extension type expects no metadata".to_owned(), - )) + |v| { + if !v.is_empty() { + Err(ArrowError::InvalidArgumentError( + "Uuid extension type expects no metadata".to_owned(), + )) + } else { + Ok(()) + } }, ) } @@ -122,11 +126,28 @@ mod tests { let field = Field::new("", DataType::FixedSizeBinary(16), false).with_metadata( [ (EXTENSION_TYPE_NAME_KEY.to_owned(), Uuid::NAME.to_owned()), - (EXTENSION_TYPE_METADATA_KEY.to_owned(), "".to_owned()), + ( + EXTENSION_TYPE_METADATA_KEY.to_owned(), + "unexpected".to_owned(), + ), ] .into_iter() .collect(), ); field.extension_type::(); } + + #[test] + fn empty_metadata_string_is_treated_as_none() -> Result<(), ArrowError> { + let field = Field::new("", DataType::FixedSizeBinary(16), false).with_metadata( + [ + (EXTENSION_TYPE_NAME_KEY.to_owned(), Uuid::NAME.to_owned()), + (EXTENSION_TYPE_METADATA_KEY.to_owned(), "".to_owned()), + ] + .into_iter() + .collect(), + ); + field.try_extension_type::()?; + Ok(()) + } } From accb1cf12a33878f179f8bb3c78c1a740fb97216 Mon Sep 17 00:00:00 2001 From: hsiang-c <137842490+hsiang-c@users.noreply.github.com> Date: Wed, 20 May 2026 10:51:53 -0700 Subject: [PATCH 04/51] [parquet] Allow more encryption algorithms (#9203) # Which issue does this PR close? - Closes #9202. # Rationale for this change - Iceberg [spec](https://iceberg.apache.org/gcm-stream-spec/#encryption-algorithm) supports AES key sizes of 128, 192 and 256 bits. Iceberg Rust depends on `arrow-rs` for Parquet I/O, I'd like to start supporting AES 256 with this PR. # What changes are included in this PR? - `RingGcmBlockEncryptor` and `RingGcmBlockDecryptor` will pick AES-128 or AES-256 based on key size - Refactor `encryption_async.rs` and `encryption.rs` to test both AES-128 and AES-256 encrypted parquet files # Are these changes tested? Yes, unit test and on AES-256 encrypted Parquet files defined in https://github.com/apache/parquet-testing/tree/master/data/aes256 # Are there any user-facing changes? No --- parquet-testing | 2 +- parquet/src/encryption/ciphers.rs | 35 +- parquet/tests/encryption/encryption.rs | 892 ++++++++++++------- parquet/tests/encryption/encryption_async.rs | 542 +++++++---- parquet/tests/encryption/encryption_util.rs | 80 ++ 5 files changed, 1015 insertions(+), 536 deletions(-) diff --git a/parquet-testing b/parquet-testing index a3d96a65e11e..5a6cf84678df 160000 --- a/parquet-testing +++ b/parquet-testing @@ -1 +1 @@ -Subproject commit a3d96a65e11e2bbca7d22a894e8313ede90a33a3 +Subproject commit 5a6cf84678df3af65c231109b78b86ba9bf495df diff --git a/parquet/src/encryption/ciphers.rs b/parquet/src/encryption/ciphers.rs index a94c72dcd5ec..beb3acda268d 100644 --- a/parquet/src/encryption/ciphers.rs +++ b/parquet/src/encryption/ciphers.rs @@ -19,7 +19,7 @@ use crate::errors::ParquetError; use crate::errors::ParquetError::General; use crate::errors::Result; use crate::file::metadata::HeapSize; -use ring::aead::{AES_128_GCM, Aad, LessSafeKey, NonceSequence, UnboundKey}; +use ring::aead::{AES_128_GCM, AES_256_GCM, Aad, LessSafeKey, NonceSequence, UnboundKey}; use ring::rand::{SecureRandom, SystemRandom}; use std::fmt::Debug; @@ -40,10 +40,20 @@ pub(crate) struct RingGcmBlockDecryptor { } impl RingGcmBlockDecryptor { + /// Create a new `RingGcmBlockDecryptor` with a given key. pub(crate) fn new(key_bytes: &[u8]) -> Result { - // todo support other key sizes - let key = UnboundKey::new(&AES_128_GCM, key_bytes) - .map_err(|_| General("Failed to create AES key".to_string()))?; + let algorithm = if key_bytes.len() == AES_128_GCM.key_len() { + &AES_128_GCM + } else if key_bytes.len() == AES_256_GCM.key_len() { + &AES_256_GCM + } else { + return Err(general_err!( + "Error creating RingGcmBlockDecryptor with unsupported key length: {}", + key_bytes.len() + )); + }; + let key = UnboundKey::new(algorithm, key_bytes) + .map_err(|_| general_err!("Failed to create {:?} key", algorithm))?; Ok(Self { key: LessSafeKey::new(key), @@ -144,10 +154,19 @@ impl RingGcmBlockEncryptor { /// return an error if it wraps around. pub(crate) fn new(key_bytes: &[u8]) -> Result { let rng = SystemRandom::new(); - - // todo support other key sizes - let key = UnboundKey::new(&AES_128_GCM, key_bytes) - .map_err(|e| general_err!("Error creating AES key: {}", e))?; + let algorithm = if key_bytes.len() == AES_128_GCM.key_len() { + &AES_128_GCM + } else if key_bytes.len() == AES_256_GCM.key_len() { + &AES_256_GCM + } else { + return Err(general_err!( + "Error creating RingGcmBlockEncryptor with unsupported key length: {}", + key_bytes.len() + )); + }; + + let key = UnboundKey::new(algorithm, key_bytes) + .map_err(|e| general_err!("Error creating {:?} key: {}", algorithm, e))?; let nonce = CounterNonce::new(&rng)?; Ok(Self { diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index c26a0edee682..edd26f29619c 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -17,7 +17,12 @@ //! This module contains tests for reading encrypted Parquet files with the Arrow API +use crate::encryption_util; use crate::encryption_util::{ + AES_128_COLUMN_KEYS, AES_128_COLUMN_NAME_KEYS, AES_128_COLUMN_NAMES, AES_128_FOOTER_KEY, + AES_128_FOOTER_KEY_NAME, AES_128_KEY_NAME_KEY, AES_128_KEY_NAMES, AES_256_COLUMN_KEYS, + AES_256_COLUMN_NAME_KEYS, AES_256_COLUMN_NAMES, AES_256_FOOTER_KEY, AES_256_FOOTER_KEY_NAME, + AES_256_KEY_NAME_KEY, AES_256_KEY_NAMES, BAD_AES_128_FOOTER_KEY, BAD_AES_256_FOOTER_KEY, TestKeyRetriever, read_and_roundtrip_to_encrypted_file, verify_column_indexes, verify_encryption_test_file_read, }; @@ -43,120 +48,131 @@ use std::sync::Arc; #[test] fn test_non_uniform_encryption_plaintext_footer() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn non_uniform_encryption_plaintext_footer(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + verify_encryption_test_file_read(file, decryption_properties); + } - // There is always a footer key even with a plaintext footer, + // AES-128: there is always a footer key even with a plaintext footer, // but this is used for signing the footer. - let footer_key = "0123456789012345".as_bytes(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes(); - let column_2_key = "1234567890123451".as_bytes(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key.to_vec()) - .with_column_key("float_field", column_2_key.to_vec()) - .build() - .unwrap(); + non_uniform_encryption_plaintext_footer(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); - verify_encryption_test_file_read(file, decryption_properties); + // AES-256 + non_uniform_encryption_plaintext_footer(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); } #[test] fn test_plaintext_footer_signature_verification() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path.clone()).unwrap(); - - let footer_key = "0000000000000000".as_bytes(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes(); - let column_2_key = "1234567890123451".as_bytes(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .disable_footer_signature_verification() - .with_column_key("double_field", column_1_key.to_vec()) - .with_column_key("float_field", column_2_key.to_vec()) - .build() - .unwrap(); - - verify_encryption_test_file_read(file, decryption_properties); + fn plaintext_footer_signature_verification(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path.clone()).unwrap(); + let mut disable_footer_signature_verification_builder = + FileDecryptionProperties::builder(footer_key.to_vec()) + .disable_footer_signature_verification(); + for (column_name, key) in column_keys { + disable_footer_signature_verification_builder = + disable_footer_signature_verification_builder + .with_column_key(column_name, key.to_vec()); + } + let decryption_properties = disable_footer_signature_verification_builder + .build() + .unwrap(); + verify_encryption_test_file_read(file, decryption_properties); - let file = File::open(path.clone()).unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key.to_vec()) - .with_column_key("float_field", column_2_key.to_vec()) - .build() - .unwrap(); + let file = File::open(path.clone()).unwrap(); + let decryption_properties = builder.build().unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .starts_with("Parquet error: Footer signature verification failed. Computed: [") + ); + } - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .starts_with("Parquet error: Footer signature verification failed. Computed: [") - ); + plaintext_footer_signature_verification(BAD_AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + plaintext_footer_signature_verification(BAD_AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS) } #[test] fn test_non_uniform_encryption_disabled_aad_storage() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = - format!("{test_data}/encrypt_columns_and_footer_disable_aad_storage.parquet.encrypted"); - let file = File::open(path.clone()).unwrap(); - - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let column_1_key = b"1234567890123450".to_vec(); - let column_2_key = b"1234567890123451".to_vec(); - - // Can read successfully when providing the correct AAD prefix - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .with_aad_prefix(b"tester".to_vec()) - .build() - .unwrap(); - - verify_encryption_test_file_read(file, decryption_properties); - - // Using wrong AAD prefix should fail - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .with_aad_prefix(b"wrong_aad_prefix".to_vec()) - .build() - .unwrap(); + fn non_uniform_encryption_disabled_aad_storage( + footer_key: &[u8], + column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer_disable_aad_storage.parquet.encrypted", + ); - let file = File::open(path.clone()).unwrap(); - let options = ArrowReaderOptions::default() - .with_file_decryption_properties(decryption_properties.clone()); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer" - ); + // Can read successfully when providing the correct AAD prefix + let file = File::open(path.clone()).unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_aad_prefix(b"tester".to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + verify_encryption_test_file_read(file, decryption_properties); + + // Using wrong AAD prefix should fail + let file = File::open(path.clone()).unwrap(); + let mut wrong_aad_builder = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_aad_prefix(b"wrong_aad_prefix".to_vec()); + for (column_name, key) in column_keys { + wrong_aad_builder = wrong_aad_builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = wrong_aad_builder.build().unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + std::assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer" + ); - // Not providing any AAD prefix should fail as it isn't stored in the file - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + // Not providing any AAD prefix should fail as it isn't stored in the file + let mut no_aad_builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + no_aad_builder = no_aad_builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = no_aad_builder.build().unwrap(); - let file = File::open(path).unwrap(); - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file was encrypted with an AAD prefix that is not stored in the file, \ + let file = File::open(path).unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file was encrypted with an AAD prefix that is not stored in the file, \ but no AAD prefix was provided in the file decryption properties" - ); + ); + } + + non_uniform_encryption_disabled_aad_storage(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + non_uniform_encryption_disabled_aad_storage(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); } #[test] @@ -167,33 +183,44 @@ fn test_plaintext_footer_read_without_decryption() { #[test] fn test_non_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn non_uniform_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let column_1_key = b"1234567890123450".to_vec(); - let column_2_key = b"1234567890123451".to_vec(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + verify_encryption_test_file_read(file, decryption_properties); + } - verify_encryption_test_file_read(file, decryption_properties); + non_uniform_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + non_uniform_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); } #[test] fn test_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn uniform_encryption(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let key_code = b"0123456789012345".to_vec(); - let decryption_properties = FileDecryptionProperties::builder(key_code).build().unwrap(); + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } - verify_encryption_test_file_read(file, decryption_properties); + uniform_encryption(AES_128_FOOTER_KEY); + uniform_encryption(AES_256_FOOTER_KEY); } #[test] @@ -213,173 +240,298 @@ fn test_decrypting_without_decryption_properties_fails() { #[test] fn test_aes_ctr_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer_ctr.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn aes_ctr_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer_ctr.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let footer_key = b"0123456789012345".to_vec(); - let column_1_key = b"1234567890123450".to_vec(); - let column_2_key = b"1234567890123451".to_vec(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + let options = + ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + let metadata = ArrowReaderMetadata::load(&file, options); - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let metadata = ArrowReaderMetadata::load(&file, options); + match metadata { + Err(ParquetError::NYI(s)) => { + assert!(s.contains("AES_GCM_CTR_V1")); + } + _ => { + panic!("Expected ParquetError::NYI"); + } + }; + } - match metadata { - Err(parquet::errors::ParquetError::NYI(s)) => { - assert!(s.contains("AES_GCM_CTR_V1")); - } - _ => { - panic!("Expected ParquetError::NYI"); - } - }; + aes_ctr_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + aes_ctr_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); } #[test] fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn non_uniform_encryption_plaintext_footer_with_key_retriever( + footer_key: &[u8], + keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let key_retriever = TestKeyRetriever::new() - .with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()) - .with_key("kc1".to_owned(), "1234567890123450".as_bytes().to_vec()) - .with_key("kc2".to_owned(), "1234567890123451".as_bytes().to_vec()); + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() - .unwrap(); + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); - verify_encryption_test_file_read(file, decryption_properties); + verify_encryption_test_file_read(file, decryption_properties); + } + + non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_128_FOOTER_KEY, + AES_128_KEY_NAME_KEY, + ); + + non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_256_FOOTER_KEY, + AES_256_KEY_NAME_KEY, + ); } #[test] -fn test_uniform_encryption_plaintext_footer_with_key_retriever() { - let test_data = arrow::util::test_util::parquet_test_data(); +fn test_roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever() { + fn roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( + footer_key: &[u8], + footer_key_metadata: &str, + wrong_footer_key: &[u8], + dec_column_keys: &[(&str, &[u8])], + enc_column_keys: &[(&str, &[u8], &str)], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - // Read example data with key retriever - let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + let mut key_retriever = + TestKeyRetriever::new().with_key(footer_key_metadata.to_owned(), footer_key.to_vec()); - let key_retriever = Arc::new( - TestKeyRetriever::new() - .with_key("kf".to_owned(), b"0123456789012345".to_vec()) - .with_key("kc1".to_owned(), b"1234567890123450".to_vec()) - .with_key("kc2".to_owned(), b"1234567890123451".to_vec()), - ); + for (key_name, key) in dec_column_keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } - let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever.clone()) - .build() - .unwrap(); + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + + // Write data into temporary file with plaintext footer and footer key metadata + let temp_file = tempfile::tempfile().unwrap(); + let mut encryption_properties_builder = + FileEncryptionProperties::builder(footer_key.to_vec()) + .with_footer_key_metadata(footer_key_metadata.into()); + for (column_name, key, metadata) in enc_column_keys { + encryption_properties_builder = encryption_properties_builder + .with_column_key_and_metadata(column_name, key.to_vec(), (*metadata).into()); + } + let encryption_properties = encryption_properties_builder + .with_plaintext_footer(true) + .build() + .unwrap(); - // Write data into temporary file with plaintext footer and footer key metadata - let temp_file = tempfile::tempfile().unwrap(); - let encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".to_vec()) - .with_footer_key_metadata("kf".into()) - .with_column_key_and_metadata("double_field", b"1234567890123450".to_vec(), b"kc1".into()) - .with_column_key_and_metadata("float_field", b"1234567890123451".to_vec(), b"kc2".into()) - .with_plaintext_footer(true) - .build() - .unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + let batch_reader = builder.build().unwrap(); + let batches = batch_reader + .collect::, _>>() + .unwrap(); - let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); - let batch_reader = builder.build().unwrap(); - let batches = batch_reader - .collect::, _>>() + let props = WriterProperties::builder() + .with_file_encryption_properties(encryption_properties) + .build(); + + let mut writer = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ) .unwrap(); + for batch in batches { + writer.write(&batch).unwrap(); + } - let props = WriterProperties::builder() - .with_file_encryption_properties(encryption_properties) - .build(); + writer.close().unwrap(); - let mut writer = ArrowWriter::try_new( - temp_file.try_clone().unwrap(), - metadata.schema().clone(), - Some(props), - ) - .unwrap(); - for batch in batches { - writer.write(&batch).unwrap(); - } + // Read temporary file with plaintext metadata using key retriever + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let _ = ArrowReaderMetadata::load(&temp_file, options.clone()).unwrap(); - writer.close().unwrap(); + // Read temporary file with plaintext metadata using key retriever with invalid key + let mut key_retriever = TestKeyRetriever::new() + .with_key(footer_key_metadata.to_owned(), wrong_footer_key.to_vec()); - // Read temporary file with plaintext metadata using key retriever - let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever) - .build() - .unwrap(); + for (key_name, key) in dec_column_keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&temp_file, options.clone()); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .starts_with("Parquet error: Footer signature verification failed. Computed: [") + ); + } - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let _ = ArrowReaderMetadata::load(&temp_file, options.clone()).unwrap(); - - // Read temporary file with plaintext metadata using key retriever with invalid key - let key_retriever = Arc::new( - TestKeyRetriever::new() - .with_key("kf".to_owned(), b"0133756789012345".to_vec()) - .with_key("kc1".to_owned(), b"1234567890123450".to_vec()) - .with_key("kc2".to_owned(), b"1234567890123451".to_vec()), + roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_128_FOOTER_KEY, + AES_128_FOOTER_KEY_NAME, + BAD_AES_128_FOOTER_KEY, + &[ + (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), + ], + &[ + ( + AES_128_COLUMN_NAMES[0], + AES_128_COLUMN_KEYS[0], + AES_128_KEY_NAMES[0], + ), + ( + AES_128_COLUMN_NAMES[1], + AES_128_COLUMN_KEYS[1], + AES_128_KEY_NAMES[1], + ), + ], ); - let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever) - .build() - .unwrap(); - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let result = ArrowReaderMetadata::load(&temp_file, options.clone()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .starts_with("Parquet error: Footer signature verification failed. Computed: [") + + roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_256_FOOTER_KEY, + AES_256_FOOTER_KEY_NAME, + BAD_AES_256_FOOTER_KEY, + &[ + (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), + ], + &[ + ( + AES_256_COLUMN_NAMES[0], + AES_256_COLUMN_KEYS[0], + AES_256_KEY_NAMES[0], + ), + ( + AES_256_COLUMN_NAMES[1], + AES_256_COLUMN_KEYS[1], + AES_256_KEY_NAMES[1], + ), + ( + AES_256_COLUMN_NAMES[2], + AES_256_COLUMN_KEYS[2], + AES_256_KEY_NAMES[2], + ), + ( + AES_256_COLUMN_NAMES[3], + AES_256_COLUMN_KEYS[3], + AES_256_KEY_NAMES[3], + ), + ( + AES_256_COLUMN_NAMES[4], + AES_256_COLUMN_KEYS[4], + AES_256_KEY_NAMES[4], + ), + ( + AES_256_COLUMN_NAMES[5], + AES_256_COLUMN_KEYS[5], + AES_256_KEY_NAMES[5], + ), + ( + AES_256_COLUMN_NAMES[6], + AES_256_COLUMN_KEYS[6], + AES_256_KEY_NAMES[6], + ), + ( + AES_256_COLUMN_NAMES[7], + AES_256_COLUMN_KEYS[7], + AES_256_KEY_NAMES[7], + ), + ], ); } #[test] fn test_non_uniform_encryption_with_key_retriever() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn non_uniform_encryption_with_key_retriever(footer_key: &[u8], keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let key_retriever = TestKeyRetriever::new() - .with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()) - .with_key("kc1".to_owned(), "1234567890123450".as_bytes().to_vec()) - .with_key("kc2".to_owned(), "1234567890123451".as_bytes().to_vec()); + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() - .unwrap(); + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } + + non_uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY, AES_128_KEY_NAME_KEY); - verify_encryption_test_file_read(file, decryption_properties); + non_uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY, AES_256_KEY_NAME_KEY); } #[test] fn test_uniform_encryption_with_key_retriever() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn uniform_encryption_with_key_retriever(key_name: &str, footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let key_retriever = - TestKeyRetriever::new().with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()); + let key_retriever = + TestKeyRetriever::new().with_key(key_name.to_owned(), footer_key.to_vec()); - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() - .unwrap(); + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } - verify_encryption_test_file_read(file, decryption_properties); + uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY); + uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY); } fn row_group_sizes(metadata: &ParquetMetaData) -> Vec { @@ -426,7 +578,7 @@ fn uniform_encryption_roundtrip( let file = tempfile::tempfile()?; - let footer_key = b"0123456789012345"; + let footer_key = AES_128_FOOTER_KEY; let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()).build()?; @@ -531,7 +683,7 @@ fn uniform_encryption_page_skipping(page_index: bool) -> parquet::errors::Result let file = tempfile::tempfile()?; - let footer_key = b"0123456789012345"; + let footer_key = AES_128_FOOTER_KEY; let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()).build()?; @@ -618,27 +770,50 @@ fn uniform_encryption_page_skipping(page_index: bool) -> parquet::errors::Result #[test] fn test_write_non_uniform_encryption() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn write_non_uniform_encryption( + footer_key: &[u8], + column_names: Vec<&str>, + column_keys: Vec>, + encryption_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let column_names = vec!["double_field", "float_field"]; - let column_keys = vec![b"1234567890123450".to_vec(), b"1234567890123451".to_vec()]; + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_column_keys(column_names.to_vec(), column_keys.clone()) + .unwrap() + .build() + .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_keys(column_names.clone(), column_keys.clone()) - .unwrap() - .build() - .unwrap(); + let mut builder = FileEncryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in encryption_column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let file_encryption_properties = builder.build().unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .with_column_keys(column_names, column_keys) - .unwrap() - .build() - .unwrap(); + read_and_roundtrip_to_encrypted_file( + &file, + decryption_properties, + file_encryption_properties, + ); + } - read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); + write_non_uniform_encryption( + AES_128_FOOTER_KEY, + AES_128_COLUMN_NAMES.to_vec(), + AES_128_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), + AES_128_COLUMN_NAME_KEYS, + ); + + write_non_uniform_encryption( + AES_256_FOOTER_KEY, + AES_256_COLUMN_NAMES.to_vec(), + AES_256_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), + AES_256_COLUMN_NAME_KEYS, + ); } #[test] @@ -647,20 +822,20 @@ fn test_write_uniform_encryption_plaintext_footer() { let path = format!("{testdata}/encrypt_columns_plaintext_footer.parquet.encrypted"); let file = File::open(path).unwrap(); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let wrong_footer_key = b"0000000000000000".to_vec(); // 128bit/16 - let column_1_key = b"1234567890123450".to_vec(); - let column_2_key = b"1234567890123451".to_vec(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let wrong_footer_key = BAD_AES_128_FOOTER_KEY.to_vec(); + let column_1_key = AES_128_COLUMN_KEYS[0].to_vec(); + let column_2_key = AES_128_COLUMN_KEYS[1].to_vec(); let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) + .with_column_key(AES_128_COLUMN_NAMES[0], column_1_key.clone()) + .with_column_key(AES_128_COLUMN_NAMES[1], column_2_key.clone()) .build() .unwrap(); let wrong_decryption_properties = FileDecryptionProperties::builder(wrong_footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) + .with_column_key(AES_128_COLUMN_NAMES[0], column_1_key) + .with_column_key(AES_128_COLUMN_NAMES[1], column_2_key) .build() .unwrap(); @@ -721,8 +896,8 @@ fn test_write_uniform_encryption_plaintext_footer() { #[test] pub fn test_column_statistics_with_plaintext_footer() { - let footer_key = b"0123456789012345".to_vec(); - let column_key = b"1234567890123450".to_vec(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key = AES_128_COLUMN_KEYS[0].to_vec(); // Encrypt with a plaintext footer and column-specific keys let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) @@ -755,7 +930,7 @@ pub fn test_column_statistics_with_plaintext_footer() { #[test] pub fn test_column_statistics_with_plaintext_footer_and_uniform_encryption() { - let footer_key = b"0123456789012345".to_vec(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); // Write with uniform encryption and a plaintext footer. let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) @@ -780,8 +955,8 @@ pub fn test_column_statistics_with_plaintext_footer_and_uniform_encryption() { #[test] pub fn test_column_statistics_with_encrypted_footer() { - let footer_key = b"0123456789012345".to_vec(); - let column_key = b"1234567890123450".to_vec(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key = AES_128_COLUMN_KEYS[0].to_vec(); // Encrypt with an encrypted footer and column-specific keys let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) @@ -810,7 +985,7 @@ pub fn test_column_statistics_with_encrypted_footer() { #[test] pub fn test_column_statistics_with_encrypted_footer_and_uniform_encryption() { - let footer_key = b"0123456789012345".to_vec(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); // Encrypt with an encrypted footer and uniform encryption let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) @@ -934,64 +1109,94 @@ fn write_and_read_stats( #[test] fn test_write_uniform_encryption() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/uniform_encryption.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn write_uniform_encryption(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .build() - .unwrap(); + let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .build() - .unwrap(); + read_and_roundtrip_to_encrypted_file( + &file, + decryption_properties, + file_encryption_properties, + ); + } - read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); + write_uniform_encryption(AES_128_FOOTER_KEY); + write_uniform_encryption(AES_256_FOOTER_KEY); } #[test] fn test_write_non_uniform_encryption_column_missmatch() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); + fn write_non_uniform_encryption_column_missmatch( + footer_key: &[u8], + column_keys: &[(&str, &[u8])], + encryption_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let column_1_key = b"1234567890123450".to_vec(); - let column_2_key = b"1234567890123451".to_vec(); + let mut decryption_builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + decryption_builder = decryption_builder.with_column_key(column_name, key.to_vec()); + } - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .build() - .unwrap(); + let decryption_properties = decryption_builder.build().unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("other_field", column_1_key) - .with_column_key("yet_another_field", column_2_key) - .build() - .unwrap(); + let mut file_encryption_builder = FileEncryptionProperties::builder(footer_key.to_vec()) + .with_column_key("other_field", encryption_column_keys[0].1.to_vec()) + .with_column_key("yet_another_field", encryption_column_keys[1].1.to_vec()); - let temp_file = tempfile::tempfile().unwrap(); + for (column_name, key) in encryption_column_keys { + file_encryption_builder = + file_encryption_builder.with_column_key(column_name, key.to_vec()); + } - // read example data - let file = File::open(path).unwrap(); - let options = ArrowReaderOptions::default() - .with_file_decryption_properties(decryption_properties.clone()); - let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); - let props = WriterProperties::builder() - .with_file_encryption_properties(file_encryption_properties) - .build(); + let file_encryption_properties = file_encryption_builder.build().unwrap(); - let result = ArrowWriter::try_new( - temp_file.try_clone().unwrap(), - metadata.schema().clone(), - Some(props), + let temp_file = tempfile::tempfile().unwrap(); + + // read example data + let file = File::open(path).unwrap(); + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + let props = WriterProperties::builder() + .with_file_encryption_properties(file_encryption_properties) + .build(); + + let result = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: The following columns with encryption keys specified were not found in the schema: other_field, yet_another_field" + ); + } + + write_non_uniform_encryption_column_missmatch( + AES_128_FOOTER_KEY, + AES_128_COLUMN_NAME_KEYS, + AES_128_COLUMN_NAME_KEYS, ); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: The following columns with encryption keys specified were not found in the schema: other_field, yet_another_field" + + write_non_uniform_encryption_column_missmatch( + AES_256_FOOTER_KEY, + AES_256_COLUMN_NAME_KEYS, + AES_256_COLUMN_NAME_KEYS, ); } @@ -1014,7 +1219,7 @@ fn test_write_encrypted_column() { let file: File = tempfile::tempfile().unwrap(); let builder = WriterProperties::builder(); - let footer_key: &[u8] = "0123456789012345".as_bytes(); + let footer_key: &[u8] = AES_128_FOOTER_KEY; let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()) .build() .unwrap(); @@ -1117,9 +1322,9 @@ fn test_write_encrypted_struct_field() { // keys need to be specified for each leaf-level Parquet column using the full "." separated // column path. let builder = WriterProperties::builder(); - let footer_key = b"0123456789012345".to_vec(); - let column_key_1 = b"1234567890123450".to_vec(); - let column_key_2 = b"1234567890123451".to_vec(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key_1 = AES_128_COLUMN_KEYS[0].to_vec(); + let column_key_2 = AES_128_COLUMN_KEYS[1].to_vec(); let file_encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) .with_column_key("struct_col.int64_col", column_key_1.clone()) .with_column_key("struct_col.float64_col", column_key_2.clone()) @@ -1184,8 +1389,8 @@ pub fn test_retrieve_row_group_statistics_after_encrypted_write() { let temp_file = tempfile::tempfile().unwrap(); - let footer_key = b"0123456789012345".to_vec(); - let column_key = b"1234567890123450".to_vec(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key = AES_128_COLUMN_KEYS[0].to_vec(); let file_encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) .with_column_key("x", column_key.clone()) .build() @@ -1218,33 +1423,40 @@ pub fn test_retrieve_row_group_statistics_after_encrypted_write() { #[test] fn test_decrypt_page_index_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); + fn decrypt_page_index_uniform(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); - let key_code: &[u8] = "0123456789012345".as_bytes(); - let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) - .build() - .unwrap(); + test_decrypt_page_index(&path, decryption_properties).unwrap(); + } - test_decrypt_page_index(&path, decryption_properties).unwrap(); + decrypt_page_index_uniform(AES_128_FOOTER_KEY); + decrypt_page_index_uniform(AES_256_FOOTER_KEY); } #[test] fn test_decrypt_page_index_non_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - - let footer_key = "0123456789012345".as_bytes().to_vec(); - let column_1_key = "1234567890123450".as_bytes().to_vec(); - let column_2_key = "1234567890123451".as_bytes().to_vec(); + fn decrypt_page_index_non_uniform(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + test_decrypt_page_index(&path, decryption_properties).unwrap(); + } - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + decrypt_page_index_non_uniform(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); - test_decrypt_page_index(&path, decryption_properties).unwrap(); + decrypt_page_index_non_uniform(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); } fn test_decrypt_page_index( diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index 48c844afb99e..f86ab59bf755 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -17,9 +17,13 @@ //! This module contains tests for reading encrypted Parquet files with the async Arrow API +use crate::encryption_util; use crate::encryption_util::{ - TestKeyRetriever, read_encrypted_file, verify_column_indexes, - verify_encryption_double_test_data, verify_encryption_test_data, + AES_128_COLUMN_KEYS, AES_128_COLUMN_NAME_KEYS, AES_128_COLUMN_NAMES, AES_128_FOOTER_KEY, + AES_128_FOOTER_KEY_NAME, AES_128_KEY_NAME_KEY, AES_256_COLUMN_KEYS, AES_256_COLUMN_NAME_KEYS, + AES_256_COLUMN_NAMES, AES_256_FOOTER_KEY, AES_256_FOOTER_KEY_NAME, AES_256_KEY_NAME_KEY, + BAD_AES_128_FOOTER_KEY, BAD_AES_256_FOOTER_KEY, TestKeyRetriever, read_encrypted_file, + verify_column_indexes, verify_encryption_double_test_data, verify_encryption_test_data, }; use arrow_array::RecordBatch; use arrow_schema::Schema; @@ -47,61 +51,64 @@ use tokio::task::JoinHandle; #[tokio::test] async fn test_non_uniform_encryption_plaintext_footer() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); + async fn non_uniform_encryption_plaintext_footer( + footer_key: &[u8], + column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } - // There is always a footer key even with a plaintext footer, + // AES-128: there is always a footer key even with a plaintext footer, // but this is used for signing the footer. - let footer_key = "0123456789012345".as_bytes().to_vec(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes().to_vec(); - let column_2_key = "1234567890123451".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + non_uniform_encryption_plaintext_footer(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + // AES-256 + non_uniform_encryption_plaintext_footer(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; } #[tokio::test] async fn test_misspecified_encryption_keys() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - - // There is always a footer key even with a plaintext footer, - // but this is used for signing the footer. - let footer_key = "0123456789012345".as_bytes(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes(); - let column_2_key = "1234567890123451".as_bytes(); - // read file with keys and check for expected error message async fn check_for_error( expected_message: &str, - path: &String, footer_key: &[u8], column_1_key: &[u8], column_2_key: &[u8], + additional_column_keys: &[(&str, &[u8])], ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); let mut file = File::open(&path).await.unwrap(); - let mut decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); if !column_1_key.is_empty() { - decryption_properties = - decryption_properties.with_column_key("double_field", column_1_key.to_vec()); + builder = builder.with_column_key(AES_128_COLUMN_NAMES[0], column_1_key.to_vec()); } if !column_2_key.is_empty() { - decryption_properties = - decryption_properties.with_column_key("float_field", column_2_key.to_vec()); + builder = builder.with_column_key(AES_128_COLUMN_NAMES[1], column_2_key.to_vec()); + } + + for (column_name, key) in additional_column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); } - let decryption_properties = decryption_properties.build().unwrap(); + let decryption_properties = builder.build().unwrap(); match verify_encryption_test_file_read_async(&mut file, decryption_properties).await { Ok(_) => { @@ -113,63 +120,149 @@ async fn test_misspecified_encryption_keys() { } } + // There is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + let footer_key = AES_128_FOOTER_KEY; + let column_1_key = AES_128_COLUMN_KEYS[0]; + let column_2_key = AES_128_COLUMN_KEYS[1]; + let empty_column_key = &[]; + // Too short footer key check_for_error( - "Parquet error: Invalid footer key. Failed to create AES key", - &path, - "bad_pwd".as_bytes(), + format!("Parquet error: Invalid footer key. Error creating RingGcmBlockDecryptor with unsupported key length: {}", "bad_pwd".len()).as_str(), + b"bad_pwd", column_1_key, column_2_key, + empty_column_key ) .await; // Wrong footer key check_for_error( "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer", - &path, - "1123456789012345".as_bytes(), + BAD_AES_128_FOOTER_KEY, column_1_key, column_2_key, + empty_column_key, ) .await; // Missing column key check_for_error( "Parquet error: No column decryption key set for encrypted column 'double_field'", - &path, footer_key, "".as_bytes(), column_2_key, + empty_column_key, ) .await; // Too short column key check_for_error( - "Parquet error: Failed to create AES key", - &path, + format!( + "Parquet error: Error creating RingGcmBlockDecryptor with unsupported key length: {}", + "abc".len() + ) + .as_str(), footer_key, "abc".as_bytes(), column_2_key, + empty_column_key, ) .await; // Wrong column key check_for_error( "Parquet error: Unable to decrypt column 'double_field', perhaps the column key is wrong?", - &path, footer_key, "1123456789012345".as_bytes(), column_2_key, + empty_column_key, ) .await; // Mixed up keys check_for_error( "Parquet error: Unable to decrypt column 'float_field', perhaps the column key is wrong?", - &path, footer_key, column_2_key, column_1_key, + empty_column_key, + ) + .await; + + let aes256_footer_key = AES_256_FOOTER_KEY; + let aes256_column_1_key = AES_256_COLUMN_KEYS[0]; + let aes256_column_2_key = AES_256_COLUMN_KEYS[1]; + let additional_column_keys = &[ + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), + ]; + + // Too short footer key + check_for_error( + format!("Parquet error: Invalid footer key. Error creating RingGcmBlockDecryptor with unsupported key length: {}", "bad_pwd".len()).as_str(), + b"bad_pwd", + aes256_column_1_key, + aes256_column_2_key, + additional_column_keys + ).await; + + // Wrong footer key + check_for_error( + "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer", + BAD_AES_256_FOOTER_KEY, + aes256_column_1_key, + aes256_column_2_key, + additional_column_keys, + ) + .await; + + // Missing column key + check_for_error( + "Parquet error: No column decryption key set for encrypted column 'double_field'", + aes256_footer_key, + "".as_bytes(), + aes256_column_2_key, + additional_column_keys, + ) + .await; + + // Too short column key + check_for_error( + format!( + "Parquet error: Error creating RingGcmBlockDecryptor with unsupported key length: {}", + "abc".len() + ) + .as_str(), + aes256_footer_key, + "abc".as_bytes(), + aes256_column_2_key, + additional_column_keys, + ) + .await; + + // Wrong column key + check_for_error( + "Parquet error: Unable to decrypt column 'double_field', perhaps the column key is wrong?", + aes256_footer_key, + "22345678901234567890123456789012".as_bytes(), + aes256_column_2_key, + additional_column_keys, + ) + .await; + + // Mixed up keys + check_for_error( + "Parquet error: Unable to decrypt column 'float_field', perhaps the column key is wrong?", + aes256_footer_key, + aes256_column_2_key, + aes256_column_1_key, + additional_column_keys, ) .await; } @@ -183,68 +276,81 @@ async fn test_plaintext_footer_read_without_decryption() { #[tokio::test] async fn test_non_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); + async fn non_uniform_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - let footer_key = "0123456789012345".as_bytes().to_vec(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes().to_vec(); - let column_2_key = "1234567890123451".as_bytes().to_vec(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + non_uniform_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; + non_uniform_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; } #[tokio::test] async fn test_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); + async fn uniform_encryption(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - let key_code: &[u8] = "0123456789012345".as_bytes(); - let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) - .build() - .unwrap(); + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } + + uniform_encryption(AES_128_FOOTER_KEY).await; + uniform_encryption(AES_256_FOOTER_KEY).await; } #[tokio::test] async fn test_aes_ctr_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer_ctr.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); + async fn aes_ctr_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer_ctr.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - let footer_key = "0123456789012345".as_bytes().to_vec(); - let column_1_key = "1234567890123450".as_bytes().to_vec(); - //let column_2_key = "1234567890123451".as_bytes().to_vec(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_1_key) - .build() - .unwrap(); + let options = + ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + let metadata = ArrowReaderMetadata::load_async(&mut file, options).await; - let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); - let metadata = ArrowReaderMetadata::load_async(&mut file, options).await; + match metadata { + Err(ParquetError::NYI(s)) => { + assert!(s.contains("AES_GCM_CTR_V1")); + } + _ => { + panic!("Expected ParquetError::NYI"); + } + }; + } - match metadata { - Err(ParquetError::NYI(s)) => { - assert!(s.contains("AES_GCM_CTR_V1")); - } - _ => { - panic!("Expected ParquetError::NYI"); - } - }; + aes_ctr_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; + aes_ctr_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; } #[tokio::test] @@ -264,32 +370,54 @@ async fn test_decrypting_without_decryption_properties_fails() { #[tokio::test] async fn test_write_non_uniform_encryption() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); + async fn write_non_uniform_encryption( + footer_key: &[u8], + column_names: Vec<&str>, + column_keys: Vec>, + encryption_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_column_keys(column_names.to_vec(), column_keys.clone()) + .unwrap() + .build() + .unwrap(); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let column_names = vec!["double_field", "float_field"]; - let column_keys = vec![b"1234567890123450".to_vec(), b"1234567890123451".to_vec()]; + let mut builder = FileEncryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in encryption_column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let file_encryption_properties = builder.build().unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_keys(column_names.clone(), column_keys.clone()) - .unwrap() - .build() + read_and_roundtrip_to_encrypted_file_async( + &path, + decryption_properties, + file_encryption_properties, + ) + .await .unwrap(); + } - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .with_column_keys(column_names, column_keys) - .unwrap() - .build() - .unwrap(); + write_non_uniform_encryption( + AES_128_FOOTER_KEY, + AES_128_COLUMN_NAMES.to_vec(), + AES_128_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), + AES_128_COLUMN_NAME_KEYS, + ) + .await; - read_and_roundtrip_to_encrypted_file_async( - &path, - decryption_properties, - file_encryption_properties, + // AES-256 + write_non_uniform_encryption( + AES_256_FOOTER_KEY, + AES_256_COLUMN_NAMES.to_vec(), + AES_256_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), + AES_256_COLUMN_NAME_KEYS, ) - .await - .unwrap(); + .await; } #[cfg(feature = "object_store")] @@ -338,98 +466,138 @@ async fn test_read_encrypted_file_from_object_store() { #[tokio::test] async fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); + async fn non_uniform_encryption_plaintext_footer_with_key_retriever( + footer_key: &[u8], + keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - let key_retriever = TestKeyRetriever::new() - .with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()) - .with_key("kc1".to_owned(), "1234567890123450".as_bytes().to_vec()) - .with_key("kc2".to_owned(), "1234567890123451".as_bytes().to_vec()); + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await .unwrap(); + } - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_128_FOOTER_KEY, + AES_128_KEY_NAME_KEY, + ) + .await; + + non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_256_FOOTER_KEY, + AES_256_KEY_NAME_KEY, + ) + .await; } #[tokio::test] async fn test_non_uniform_encryption_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); + async fn non_uniform_encryption_with_key_retriever(footer_key: &[u8], keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - let key_retriever = TestKeyRetriever::new() - .with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()) - .with_key("kc1".to_owned(), "1234567890123450".as_bytes().to_vec()) - .with_key("kc2".to_owned(), "1234567890123451".as_bytes().to_vec()); + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await .unwrap(); + } - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + non_uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY, AES_128_KEY_NAME_KEY).await; + non_uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY, AES_256_KEY_NAME_KEY).await; } #[tokio::test] async fn test_uniform_encryption_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/uniform_encryption.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); + async fn uniform_encryption_with_key_retriever(key_name: &str, footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - let key_retriever = - TestKeyRetriever::new().with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()); + let key_retriever = + TestKeyRetriever::new().with_key(key_name.to_owned(), footer_key.to_vec()); - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await .unwrap(); + } - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY).await; + uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY).await; } #[tokio::test] async fn test_decrypt_page_index_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); + async fn decrypt_page_index_uniform(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); - let key_code: &[u8] = "0123456789012345".as_bytes(); - let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) - .build() - .unwrap(); + test_decrypt_page_index(&path, decryption_properties) + .await + .unwrap(); + } - test_decrypt_page_index(&path, decryption_properties) - .await - .unwrap(); + decrypt_page_index_uniform(AES_128_FOOTER_KEY).await; + decrypt_page_index_uniform(AES_256_FOOTER_KEY).await; } #[tokio::test] async fn test_decrypt_page_index_non_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - - let footer_key = "0123456789012345".as_bytes().to_vec(); - let column_1_key = "1234567890123450".as_bytes().to_vec(); - let column_2_key = "1234567890123451".as_bytes().to_vec(); + async fn decrypt_page_index_non_uniform(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + test_decrypt_page_index(&path, decryption_properties) + .await + .unwrap(); + } - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + decrypt_page_index_non_uniform(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; - test_decrypt_page_index(&path, decryption_properties) - .await - .unwrap(); + decrypt_page_index_non_uniform(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; } async fn test_decrypt_page_index( @@ -675,14 +843,14 @@ async fn test_concurrent_encrypted_writing_over_multiple_row_groups() { let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); let file = std::fs::File::open(path).unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let file_encryption_properties = FileEncryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let decryption_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); @@ -748,14 +916,14 @@ async fn test_multi_threaded_encrypted_writing() { let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); let file = std::fs::File::open(path).unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let file_encryption_properties = FileEncryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let decryption_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); @@ -849,14 +1017,14 @@ async fn test_multi_threaded_encrypted_writing_deprecated() { let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); let file = std::fs::File::open(path).unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let file_encryption_properties = FileEncryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let decryption_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); diff --git a/parquet/tests/encryption/encryption_util.rs b/parquet/tests/encryption/encryption_util.rs index 7f4cc5a9da45..f0fe66651d11 100644 --- a/parquet/tests/encryption/encryption_util.rs +++ b/parquet/tests/encryption/encryption_util.rs @@ -26,10 +26,79 @@ use parquet::encryption::encrypt::FileEncryptionProperties; use parquet::errors::{ParquetError, Result}; use parquet::file::metadata::ParquetMetaData; use parquet::file::properties::WriterProperties; +use ring::aead::AES_256_GCM; use std::collections::HashMap; use std::fs::File; use std::sync::{Arc, Mutex}; +pub(crate) const AES_128_FOOTER_KEY: &[u8; 16] = b"0123456789012345"; // 128bit/16 +pub(crate) const BAD_AES_128_FOOTER_KEY: &[u8; 16] = b"0000000000000000"; +pub(crate) const AES_128_FOOTER_KEY_NAME: &str = "kf"; +pub(crate) const AES_128_KEY_NAMES: [&str; 2] = ["kc1", "kc2"]; +pub(crate) const AES_128_COLUMN_NAMES: [&str; 2] = ["double_field", "float_field"]; +pub(crate) const AES_128_COLUMN_KEYS: [&[u8; 16]; 2] = [b"1234567890123450", b"1234567890123451"]; + +pub(crate) const AES_128_COLUMN_NAME_KEYS: &[(&str, &[u8]); 2] = &[ + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), +]; + +pub(crate) const AES_128_KEY_NAME_KEY: &[(&str, &[u8]); 3] = &[ + (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), + (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), +]; + +pub(crate) const AES_256_FOOTER_KEY: &[u8; 32] = b"01234567890123456789012345678901"; // 256bit/32 +pub(crate) const BAD_AES_256_FOOTER_KEY: &[u8; 32] = b"00000000000000000000000000000000"; +pub(crate) const AES_256_FOOTER_KEY_NAME: &str = "kf"; +pub(crate) const AES_256_KEY_NAMES: [&str; 8] = + ["kc1", "kc2", "kc3", "kc4", "kc5", "kc6", "kc7", "kc8"]; + +pub(crate) const AES_256_COLUMN_NAMES: [&str; 8] = [ + "double_field", + "float_field", + "boolean_field", + "int32_field", + "ba_field", + "flba_field", + "int64_field.list.element", + "int96_field", +]; +pub(crate) const AES_256_COLUMN_KEYS: [&[u8]; 8] = [ + b"12345678901234567890123456789012", + b"12345678901234567890123456789013", + b"12345678901234567890123456789014", + b"12345678901234567890123456789015", + b"12345678901234567890123456789016", + b"12345678901234567890123456789017", + b"12345678901234567890123456789018", + b"12345678901234567890123456789019", +]; + +pub(crate) const AES_256_COLUMN_NAME_KEYS: &[(&str, &[u8]); 8] = &[ + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), +]; + +pub(crate) const AES_256_KEY_NAME_KEY: &[(&str, &[u8]); 9] = &[ + (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), + (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), +]; + pub(crate) fn verify_encryption_double_test_data( record_batches: Vec, metadata: &ParquetMetaData, @@ -317,3 +386,14 @@ impl KeyRetriever for TestKeyRetriever { } } } + +pub fn encrypted_data_path(footer_key: &[u8], file_name: &str) -> String { + let test_data = arrow::util::test_util::parquet_test_data(); + let subpath = if AES_256_GCM.key_len() == footer_key.len() { + "aes256/" + } else { + "" + }; + let path = format!("{test_data}/{subpath}/{file_name}"); + path +} From 7c6eb2cbd958369fc2d44b5947bcfba480b6dbf8 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 20 May 2026 11:11:36 -0700 Subject: [PATCH 05/51] feat(parquet): Add `ParquetPushDecoder::into_builder` to allow swapping projections / row filters at row group boundaries (#9968) This is the decoder piece of the work presented at the NYC DataFusion meetup. The idea is that we'll be able to adaptively promote and demote filters into row filters based on runtime selectivity stats. --------- Co-authored-by: Claude Opus 4.7 (1M context) --- parquet/src/arrow/arrow_reader/mod.rs | 7 +- parquet/src/arrow/async_reader/mod.rs | 4 +- parquet/src/arrow/push_decoder/mod.rs | 590 +++++++++++++++++- .../arrow/push_decoder/reader_builder/mod.rs | 67 ++ parquet/src/arrow/push_decoder/remaining.rs | 73 ++- parquet/src/util/push_buffers.rs | 29 +- 6 files changed, 733 insertions(+), 37 deletions(-) diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 70d3ce7cf9a9..12c3e192cdfc 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -110,11 +110,12 @@ pub const DEFAULT_BATCH_SIZE: usize = 1024; pub struct ArrowReaderBuilder { /// The "input" to read parquet data from. /// - /// Note in the case of the [`ParquetPushDecoderBuilder`], there - /// is no underlying input, which is indicated by a type parameter of [`NoInput`] + /// Note in the case of the [`ParquetPushDecoderBuilder`] there is no + /// underlying reader; the input is instead [`PushDecoderInput`], the buffer that + /// caller-pushed bytes accumulate in. /// /// [`ParquetPushDecoderBuilder`]: crate::arrow::push_decoder::ParquetPushDecoderBuilder - /// [`NoInput`]: crate::arrow::push_decoder::NoInput + /// [`PushDecoderInput`]: crate::arrow::push_decoder::PushDecoderInput pub(crate) input: T, pub(crate) metadata: Arc, diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index 6c7890b75c45..3bba746e740e 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -54,7 +54,7 @@ pub use metadata::*; mod store; use crate::DecodeResult; -use crate::arrow::push_decoder::{NoInput, ParquetPushDecoder, ParquetPushDecoderBuilder}; +use crate::arrow::push_decoder::{ParquetPushDecoder, ParquetPushDecoderBuilder, PushDecoderInput}; #[cfg(feature = "object_store")] pub use store::*; @@ -600,7 +600,7 @@ impl ParquetRecordBatchStreamBuilder { let projected_schema = Arc::new(Schema::new(projected_fields)); let decoder = ParquetPushDecoderBuilder { - input: NoInput, + input: PushDecoderInput::default(), metadata, schema, fields, diff --git a/parquet/src/arrow/push_decoder/mod.rs b/parquet/src/arrow/push_decoder/mod.rs index f905d6fb2ccb..6dc5520bb975 100644 --- a/parquet/src/arrow/push_decoder/mod.rs +++ b/parquet/src/arrow/push_decoder/mod.rs @@ -27,11 +27,11 @@ use crate::arrow::arrow_reader::{ }; use crate::errors::ParquetError; use crate::file::metadata::ParquetMetaData; -use crate::util::push_buffers::PushBuffers; +pub use crate::util::push_buffers::PushBuffers; use arrow_array::RecordBatch; use bytes::Bytes; -use reader_builder::{RowBudget, RowGroupReaderBuilder}; -use remaining::RemainingRowGroups; +use reader_builder::{RowBudget, RowGroupReaderBuilder, RowGroupReaderBuilderParts}; +use remaining::{RemainingRowGroups, RemainingRowGroupsParts}; use std::ops::Range; use std::sync::Arc; @@ -109,19 +109,100 @@ use std::sync::Arc; /// } /// } /// ``` -pub type ParquetPushDecoderBuilder = ArrowReaderBuilder; - -/// Type that represents "No input" for the [`ParquetPushDecoderBuilder`] /// -/// There is no "input" for the push decoder by design (the idea is that -/// the caller pushes data to the decoder as needed).. +/// # Adaptive scans +/// +/// The scan strategy is not fixed once [`build`](Self::build) is called: it +/// can be changed *while decoding*, at row-group boundaries. +/// +/// The important API for this is [`ParquetPushDecoder::try_next_reader`]. +/// Unlike [`try_decode`](ParquetPushDecoder::try_decode), which barrels +/// straight through row-group boundaries, `try_next_reader` returns once per +/// row group — leaving a clean window *between* row groups. At any such +/// boundary, [`ParquetPushDecoder::into_builder`] hands back a +/// `ParquetPushDecoderBuilder` for the row groups not yet decoded. Change any +/// option on it (projection, row filter, row selection policy, …) and +/// [`build`](Self::build) a fresh decoder that resumes from the next row +/// group. This is how a query engine promotes or demotes filters — for +/// example turning a row filter on or off — based on the selectivity observed +/// in the row groups decoded so far. /// -/// However, [`ArrowReaderBuilder`] is shared with the sync and async readers, -/// which DO have an `input`. To support reusing the same builder code for -/// all three types of decoders, we define this `NoInput` for the push decoder to -/// denote in the type system there is no type. -#[derive(Debug, Clone, Copy)] -pub struct NoInput; +/// ``` +/// # use std::ops::Range; +/// # use std::sync::Arc; +/// # use bytes::Bytes; +/// # use arrow_array::record_batch; +/// # use parquet::DecodeResult; +/// # use parquet::arrow::ProjectionMask; +/// # use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; +/// # use parquet::arrow::ArrowWriter; +/// # use parquet::file::metadata::ParquetMetaDataPushDecoder; +/// # use parquet::file::properties::WriterProperties; +/// # let file_bytes = { +/// # let batch = record_batch!( +/// # ("a", Int32, [1, 2, 3, 4, 5, 6]), +/// # ("b", Int32, [6, 5, 4, 3, 2, 1]) +/// # ).unwrap(); +/// # // Small row groups so the test file has two of them. +/// # let props = WriterProperties::builder().set_max_row_group_row_count(Some(3)).build(); +/// # let mut buffer = vec![]; +/// # let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); +/// # writer.write(&batch).unwrap(); +/// # writer.close().unwrap(); +/// # Bytes::from(buffer) +/// # }; +/// # let get_range = |r: &Range| file_bytes.slice(r.start as usize..r.end as usize); +/// # let file_length = file_bytes.len() as u64; +/// # let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_length).unwrap(); +/// # metadata_decoder.push_ranges(vec![0..file_length], vec![file_bytes.clone()]).unwrap(); +/// # let DecodeResult::Data(parquet_metadata) = metadata_decoder.try_decode().unwrap() else { panic!() }; +/// # let parquet_metadata = Arc::new(parquet_metadata); +/// let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(parquet_metadata) +/// .unwrap() +/// .build() +/// .unwrap(); +/// +/// // Drive the decoder one row group at a time with `try_next_reader`. +/// loop { +/// match decoder.try_next_reader().unwrap() { +/// DecodeResult::NeedsData(ranges) => { +/// // Fetch and hand over the bytes the decoder asked for. +/// let data = ranges.iter().map(|r| get_range(r)).collect(); +/// decoder.push_ranges(ranges, data).unwrap(); +/// } +/// DecodeResult::Data(reader) => { +/// // Decode this row group's batches. +/// for batch in reader { +/// assert!(batch.unwrap().num_rows() > 0); +/// } +/// // We are now at a row-group boundary. Based on whatever stats +/// // were gathered, optionally change strategy for the row groups +/// // still to come: drop or promote a row filter, narrow or widen +/// // the projection, etc. +/// if decoder.is_at_row_group_boundary() && decoder.row_groups_remaining() > 0 { +/// let builder = decoder.into_builder().unwrap(); +/// // e.g. column "b" turned out not to be needed. +/// let projection = ProjectionMask::columns(builder.parquet_schema(), ["a"]); +/// decoder = builder.with_projection(projection).build().unwrap(); +/// } +/// } +/// DecodeResult::Finished => break, +/// } +/// } +/// ``` +pub type ParquetPushDecoderBuilder = ArrowReaderBuilder; + +/// The `input` of a [`ParquetPushDecoderBuilder`]. +/// +/// The shared [`ArrowReaderBuilder`] is generic over an `input`. The sync and +/// async builders read from a file or async reader; the push decoder has no +/// reader, so its input is the [`PushBuffers`] that caller-pushed bytes +/// accumulate in (empty for a fresh builder). +#[derive(Debug, Default)] +pub struct PushDecoderInput { + /// Bytes pushed into the decoder, awaiting decode. + buffers: PushBuffers, +} /// Methods for building a ParquetDecoder. See the base [`ArrowReaderBuilder`] for /// more options that can be configured. @@ -156,15 +237,24 @@ impl ParquetPushDecoderBuilder { /// See [`ArrowReaderMetadata::try_new`] for how to create the metadata from /// the Parquet metadata and reader options. pub fn new_with_metadata(arrow_reader_metadata: ArrowReaderMetadata) -> Self { - Self::new_builder(NoInput, arrow_reader_metadata) + Self::new_builder(PushDecoderInput::default(), arrow_reader_metadata) + } + + /// Provide a preexisting [`PushBuffers`] for the built decoder to read + /// from, so bytes already fetched are not requested again. + pub fn with_buffers(self, buffers: PushBuffers) -> Self { + Self { + input: PushDecoderInput { buffers }, + ..self + } } /// Create a [`ParquetPushDecoder`] with the configured options pub fn build(self) -> Result { let Self { - input: NoInput, + input: PushDecoderInput { buffers }, metadata: parquet_metadata, - schema: _, + schema, fields, batch_size, row_groups, @@ -185,9 +275,9 @@ impl ParquetPushDecoderBuilder { .as_ref() .is_some_and(|filter| !filter.predicates.is_empty()); - // Prepare to build RowGroup readers - let file_len = 0; // not used in push decoder - let buffers = PushBuffers::new(file_len); + // Prepare to build RowGroup readers. `buffers` carries any bytes the + // caller already pushed (preserved across `into_builder`); a fresh + // builder supplies an empty `PushBuffers`. let row_group_reader_builder = RowGroupReaderBuilder::new( batch_size, projection, @@ -202,6 +292,7 @@ impl ParquetPushDecoderBuilder { // Initialize the decoder with the configured options let remaining_row_groups = RemainingRowGroups::new( + schema, parquet_metadata, row_groups, selection, @@ -218,6 +309,55 @@ impl ParquetPushDecoderBuilder { } } +/// Reassemble a [`ParquetPushDecoderBuilder`] from a decoder's not-yet-decoded +/// state — the inverse of [`ParquetPushDecoderBuilder::build`]. The rebuilt +/// builder pins the remaining row groups and carries the remaining row +/// selection, offset/limit budget, and buffered bytes. +fn builder_from_remaining(parts: RemainingRowGroupsParts) -> ParquetPushDecoderBuilder { + let RemainingRowGroupsParts { + metadata, + schema, + row_groups, + selection, + offset, + limit, + reader_builder, + } = parts; + let RowGroupReaderBuilderParts { + batch_size, + projection, + fields, + filter, + max_predicate_cache_size, + metrics, + row_selection_policy, + buffers, + } = reader_builder; + + ArrowReaderBuilder { + input: PushDecoderInput::default(), + metadata, + schema, + fields, + batch_size, + // The frontier tracks remaining row groups explicitly, so the rebuilt + // builder always pins them (even if the original left `row_groups` as + // `None` meaning "all"). + row_groups: Some(row_groups), + projection, + filter, + selection, + row_selection_policy, + limit, + offset, + metrics, + max_predicate_cache_size, + } + // Carry the decoder's already-fetched bytes across the rebuild so the new + // decoder does not re-request them. + .with_buffers(buffers) +} + /// A push based Parquet Decoder /// /// See [`ParquetPushDecoderBuilder`] for an example of how to build and use the decoder. @@ -377,6 +517,80 @@ impl ParquetPushDecoder { pub fn clear_all_ranges(&mut self) { self.state.clear_all_ranges(); } + + /// True iff the decoder is at a row-group boundary, where + /// [`Self::into_builder`] can reconfigure the scan. + /// + /// A boundary is "between row groups": the previous row group's + /// [`ParquetRecordBatchReader`] has been fully extracted (via + /// [`Self::try_next_reader`]) or fully drained (via [`Self::try_decode`]), + /// and the next row group has not yet been planned. While + /// [`Self::try_decode`] is iterating an active row group's reader this + /// returns `false`; with [`Self::try_next_reader`] there is a clean + /// window between two consecutive returns where this is `true`. + pub fn is_at_row_group_boundary(&self) -> bool { + self.state.is_at_row_group_boundary() + } + + /// Number of row groups left to decode after the one currently in flight. + /// Useful as a "should I bother reconfiguring the scan?" signal. + pub fn row_groups_remaining(&self) -> usize { + self.state.row_groups_remaining() + } + + /// Decompose this decoder back into a [`ParquetPushDecoderBuilder`] for the + /// row groups that have *not* yet been decoded. + /// + /// This is the API for *adaptive* scans. Drive the decoder with + /// [`Self::try_next_reader`]; at any row-group boundary, call + /// `into_builder` to recover a builder, adjust it with the usual + /// [`ParquetPushDecoderBuilder`] setters, and + /// [`build`](ParquetPushDecoderBuilder::build) a fresh decoder that resumes + /// from the next row group: + /// + /// ```no_run + /// # use parquet::arrow::push_decoder::ParquetPushDecoder; + /// # use parquet::arrow::arrow_reader::RowFilter; + /// # fn get_decoder() -> ParquetPushDecoder { unimplemented!() } + /// # fn new_filter() -> RowFilter { unimplemented!() } + /// let mut decoder = get_decoder(); + /// // ... drive `decoder.try_next_reader()` for a few row groups ... + /// if decoder.is_at_row_group_boundary() && decoder.row_groups_remaining() > 0 { + /// decoder = decoder + /// .into_builder() + /// .unwrap() + /// // any builder option can be changed here, e.g. promote a + /// // filter into a row filter based on observed selectivity + /// .with_row_filter(new_filter()) + /// .build() + /// .unwrap(); + /// } + /// ``` + /// + /// The returned builder pins the not-yet-decoded row groups (via + /// [`with_row_groups`](ArrowReaderBuilder::with_row_groups)) and carries the + /// not-yet-consumed row selection and offset/limit budget, so rows from + /// already-decoded row groups are not produced again. Every other option — + /// projection, row filter, row selection policy, batch size, metrics, + /// predicate-cache size — is left exactly as the decoder had it and can be + /// overridden before [`build`](ParquetPushDecoderBuilder::build). + /// + /// # Errors + /// + /// Returns `Err(ParquetError::General)` when the decoder is not at a + /// row-group boundary (check [`Self::is_at_row_group_boundary`] first) or + /// has already finished. The decoder is consumed either way. + /// + /// # Buffered bytes + /// + /// The decoder's buffered bytes are carried across the rebuild: bytes + /// already fetched for row groups the new configuration still reads are + /// not re-requested. Bytes the new configuration no longer needs stay + /// buffered until [`clear_all_ranges`](Self::clear_all_ranges) is called + /// or the rebuilt decoder is dropped. + pub fn into_builder(self) -> Result { + self.state.into_builder() + } } /// Internal state machine for the [`ParquetPushDecoder`] @@ -599,6 +813,59 @@ impl ParquetDecoderState { ParquetDecoderState::Finished => {} } } + + fn is_at_row_group_boundary(&self) -> bool { + match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups.is_at_row_group_boundary(), + // Mid-row-group: the active reader holds an `ArrayReader` and + // `ReadPlan` keyed to the *current* projection/filter; rebuilding + // would require throwing that work away. + ParquetDecoderState::DecodingRowGroup { .. } => false, + ParquetDecoderState::Finished => false, + } + } + + fn row_groups_remaining(&self) -> usize { + match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups.row_groups_remaining(), + ParquetDecoderState::DecodingRowGroup { + remaining_row_groups, + .. + } => remaining_row_groups.row_groups_remaining(), + ParquetDecoderState::Finished => 0, + } + } + + fn into_builder(self) -> Result { + let remaining_row_groups = match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups, + ParquetDecoderState::DecodingRowGroup { .. } => { + return Err(ParquetError::General( + "into_builder called while a row group is being decoded; \ + check is_at_row_group_boundary() first" + .to_string(), + )); + } + ParquetDecoderState::Finished => { + return Err(ParquetError::General( + "into_builder called on a finished decoder".to_string(), + )); + } + }; + if !remaining_row_groups.is_at_row_group_boundary() { + return Err(ParquetError::General( + "into_builder called mid-row-group; check is_at_row_group_boundary() first" + .to_string(), + )); + } + Ok(builder_from_remaining(remaining_row_groups.into_parts())) + } } #[cfg(test)] @@ -1476,6 +1743,270 @@ mod test { expect_finished(decoder.try_decode()); } + /// `into_builder` between row groups recovers a builder for the + /// not-yet-decoded row groups; rebuilding it with a new row filter + /// applies that filter to the subsequent row groups while leaving the + /// already-decoded row group's results untouched. + /// + /// See the "Adaptive scans" section of [`ParquetPushDecoderBuilder`] for + /// the high-level overview. + #[test] + fn test_into_builder_installs_filter_between_row_groups() { + let schema_descr = test_file_parquet_metadata() + .file_metadata() + .schema_descr_ptr(); + let mut decoder = prefetched_decoder(1024); + + // Reader for row group 0 — no filter. + let reader0 = expect_data(decoder.try_next_reader()); + let batches0: Vec<_> = reader0.collect::>().unwrap(); + let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap(); + assert_eq!(batch0, TEST_BATCH.slice(0, 200)); + + // We're between row groups now. Rebuild with a filter on column "a". + assert!(decoder.is_at_row_group_boundary()); + assert_eq!(decoder.row_groups_remaining(), 1); + let filter = + ArrowPredicateFn::new(ProjectionMask::columns(&schema_descr, ["a"]), |batch| { + gt(batch.column(0), &Int64Array::new_scalar(250)) + }); + let mut decoder = decoder + .into_builder() + .unwrap() + .with_row_filter(RowFilter::new(vec![Box::new(filter)])) + .build() + .unwrap(); + + // Reader for row group 1 — filter applied. The rebuilt decoder kept + // the buffered bytes (see `test_into_builder_preserves_buffered_bytes`) + // so no data needs to be re-supplied. Column "a" in RG1 has values + // 200..399; `a > 250` keeps 251..399 = 149 rows. + let reader1 = expect_data(decoder.try_next_reader()); + let batches1: Vec<_> = reader1.collect::>().unwrap(); + let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap(); + assert_eq!(batch1, TEST_BATCH.slice(251, 149)); + expect_finished(decoder.try_next_reader()); + } + + /// `into_builder` is rejected while a row group's reader is being + /// drained (`DecodingRowGroup`); the error points at + /// `is_at_row_group_boundary`. + #[test] + fn test_into_builder_rejected_mid_row_group() { + let mut decoder = prefetched_decoder(50); + + // Decode one batch to land mid-row-group, inside `DecodingRowGroup` + // with an active reader — not a boundary. + expect_data(decoder.try_decode()); + assert!(!decoder.is_at_row_group_boundary()); + + let err = decoder.into_builder().unwrap_err(); + let err_msg = format!("{err}"); + assert!( + err_msg.contains("is_at_row_group_boundary"), + "unexpected error: {err_msg}" + ); + } + + /// `into_builder` is rejected once the decoder has finished. + #[test] + fn test_into_builder_rejected_on_finished_decoder() { + let mut decoder = prefetched_decoder(1024); + expect_data(decoder.try_decode()); + expect_data(decoder.try_decode()); + expect_finished(decoder.try_decode()); + assert!(!decoder.is_at_row_group_boundary()); + + let err = decoder.into_builder().unwrap_err(); + assert!( + format!("{err}").contains("finished"), + "unexpected error: {err}" + ); + } + + /// `try_next_reader` hands the active reader off to the caller and + /// transitions the decoder back to `ReadingRowGroup` — so the caller + /// can call `into_builder` even while still holding the returned + /// reader. (The handed-off reader has no link back to the decoder's + /// projection/filter; it has its own `ArrayReader` and `ReadPlan`.) + #[test] + fn test_into_builder_allowed_while_iterating_handed_off_reader() { + let mut decoder = prefetched_decoder(1024); + + let reader0 = expect_data(decoder.try_next_reader()); + // Decoder no longer owns the reader, so it considers itself + // "between row groups". + assert!(decoder.is_at_row_group_boundary()); + // Recovering the builder consumes the decoder but leaves `reader0` + // valid: iterating it is independent of the decoder's state. + let _builder = decoder.into_builder().unwrap(); + let batches: Vec<_> = reader0.collect::>().unwrap(); + let batch0 = concat_batches(&TEST_BATCH.schema(), &batches).unwrap(); + assert_eq!(batch0, TEST_BATCH.slice(0, 200)); + } + + /// `into_builder` recovers a builder for the *remaining* row groups and + /// carries the not-yet-consumed offset/limit budget, so a rebuilt + /// decoder resumes where the original left off rather than restarting. + #[test] + fn test_into_builder_resumes_remaining_budget() { + // limit = 250 spans both 200-row row groups: all 200 rows of RG0 + // plus the first 50 rows of RG1. + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_batch_size(1024) + .with_limit(250) + .build() + .unwrap(); + prefetch_test_file(&mut decoder); + + // RG0 contributes all 200 of its rows. + let reader0 = expect_data(decoder.try_next_reader()); + let batches0: Vec<_> = reader0.collect::>().unwrap(); + let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap(); + assert_eq!(batch0, TEST_BATCH.slice(0, 200)); + + // Rebuild without changing anything: the remaining 50-row limit and + // the not-yet-decoded RG1 must carry through (as do the buffers, so + // no data needs re-supplying). + assert!(decoder.is_at_row_group_boundary()); + let mut decoder = decoder.into_builder().unwrap().build().unwrap(); + + let reader1 = expect_data(decoder.try_next_reader()); + let batches1: Vec<_> = reader1.collect::>().unwrap(); + let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap(); + // Only the first 50 rows of RG1 (200..249) — the rest of the limit. + assert_eq!(batch1, TEST_BATCH.slice(200, 50)); + expect_finished(decoder.try_next_reader()); + } + + /// `into_builder` carries the decoder's buffered bytes across the + /// rebuild: the rebuilt decoder keeps them and does not re-request data + /// it already holds. + #[test] + fn test_into_builder_preserves_buffered_bytes() { + let mut decoder = prefetched_decoder(1024); + assert_eq!(decoder.buffered_bytes(), test_file_len()); + + // Drain RG0. + let reader0 = expect_data(decoder.try_next_reader()); + let _: Vec<_> = reader0.collect::>().unwrap(); + // RG1's bytes are still staged inside the decoder. + let buffered = decoder.buffered_bytes(); + assert!(buffered > 0); + + // Rebuilding via into_builder keeps the staged bytes. + let mut decoder = decoder.into_builder().unwrap().build().unwrap(); + assert_eq!(decoder.buffered_bytes(), buffered); + + // RG1's bytes are already buffered, so it decodes without a + // `NeedsData` round-trip. + let reader1 = expect_data(decoder.try_next_reader()); + let batches1: Vec<_> = reader1.collect::>().unwrap(); + let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap(); + assert_eq!(batch1, TEST_BATCH.slice(200, 200)); + expect_finished(decoder.try_next_reader()); + } + + /// Drive the decoder incrementally. Start with a narrow projection, + /// drain RG0, then `into_builder` and widen the projection to all three + /// columns. The rebuilt decoder's `NeedsData` for RG1 must request + /// bytes for *all three* columns, not just the originally-projected + /// "a". The expected ranges are hardcoded because `TEST_BATCH` and the + /// writer settings are static; this pins the layout cleanly without a + /// parallel reference decoder. + #[test] + fn test_into_builder_expand_projection_requests_new_bytes() { + let metadata = test_file_parquet_metadata(); + let schema_descr = metadata.file_metadata().schema_descr_ptr(); + + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(metadata) + .unwrap() + .with_batch_size(1024) + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .build() + .unwrap(); + + // RG0: incrementally satisfy the narrow request — a single + // contiguous range for the "a"-only projection. + let ranges_rg0 = expect_needs_data(decoder.try_next_reader()); + assert_eq!(ranges_rg0, vec![4..1860]); + push_ranges_to_decoder(&mut decoder, ranges_rg0); + + let reader0 = expect_data(decoder.try_next_reader()); + let batches0: Vec<_> = reader0.collect::>().unwrap(); + let batch0 = concat_batches(&batches0[0].schema(), &batches0).unwrap(); + assert_eq!(batch0, TEST_BATCH.slice(0, 200).project(&[0]).unwrap()); + + // Widen the projection at the boundary. + assert!(decoder.is_at_row_group_boundary()); + let mut decoder = decoder + .into_builder() + .unwrap() + .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b", "c"])) + .build() + .unwrap(); + + // RG1 now requests "a", "b", and "c" column chunks: ~1.8KiB each + // for "a" and "b", ~7.5KiB for the StringView column "c". + let ranges_rg1 = expect_needs_data(decoder.try_next_reader()); + assert_eq!(ranges_rg1, vec![11062..12918, 12918..14774, 14774..22230]); + push_ranges_to_decoder(&mut decoder, ranges_rg1); + + let reader1 = expect_data(decoder.try_next_reader()); + let batches1: Vec<_> = reader1.collect::>().unwrap(); + let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap(); + assert_eq!(batch1, TEST_BATCH.slice(200, 200)); + expect_finished(decoder.try_next_reader()); + } + + /// Mirror of [`test_into_builder_expand_projection_requests_new_bytes`]: + /// start with the full projection, drain RG0, then `into_builder` and + /// narrow the projection to just column "a". RG1's `NeedsData` must + /// request only the single "a" column-chunk range, not the three a + /// wide projection would. + #[test] + fn test_into_builder_narrow_projection_requests_fewer_bytes() { + let metadata = test_file_parquet_metadata(); + let schema_descr = metadata.file_metadata().schema_descr_ptr(); + + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(metadata) + .unwrap() + .with_batch_size(1024) + .build() + .unwrap(); + + // RG0 with the default (full) projection — three column ranges. + let ranges_rg0 = expect_needs_data(decoder.try_next_reader()); + assert_eq!(ranges_rg0, vec![4..1860, 1860..3716, 3716..11062]); + push_ranges_to_decoder(&mut decoder, ranges_rg0); + + let reader0 = expect_data(decoder.try_next_reader()); + let batches0: Vec<_> = reader0.collect::>().unwrap(); + let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap(); + assert_eq!(batch0, TEST_BATCH.slice(0, 200)); + + // Narrow the projection at the boundary. + assert!(decoder.is_at_row_group_boundary()); + let mut decoder = decoder + .into_builder() + .unwrap() + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .build() + .unwrap(); + + // RG1 now requests column "a" only — a single 1856-byte range. + let ranges_rg1 = expect_needs_data(decoder.try_next_reader()); + assert_eq!(ranges_rg1, vec![11062..12918]); + push_ranges_to_decoder(&mut decoder, ranges_rg1); + + let reader1 = expect_data(decoder.try_next_reader()); + let batches1: Vec<_> = reader1.collect::>().unwrap(); + let batch1 = concat_batches(&batches1[0].schema(), &batches1).unwrap(); + assert_eq!(batch1, TEST_BATCH.slice(200, 200).project(&[0]).unwrap()); + expect_finished(decoder.try_next_reader()); + } + /// Returns a batch with 400 rows, with 3 columns: "a", "b", "c" /// /// Note c is a different types (so the data page sizes will be different) @@ -1582,6 +2113,25 @@ mod test { metadata_decoder.push_ranges(ranges, data).unwrap(); } + /// Push the entire test file into `decoder`. + fn prefetch_test_file(decoder: &mut ParquetPushDecoder) { + decoder + .push_range(test_file_range(), TEST_FILE_DATA.clone()) + .unwrap(); + } + + /// Build a decoder over the test file with the given batch size and + /// prefetch the whole file into it. + fn prefetched_decoder(batch_size: usize) -> ParquetPushDecoder { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_batch_size(batch_size) + .build() + .unwrap(); + prefetch_test_file(&mut decoder); + decoder + } + fn push_ranges_to_decoder(decoder: &mut ParquetPushDecoder, ranges: Vec>) { let data = ranges .iter() diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs index 0452cea436a3..dacf1a2caad9 100644 --- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs +++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs @@ -104,6 +104,16 @@ impl RowBudget { matches!(self.limit, Some(0)) } + /// The offset still to be skipped before the next readable row group. + pub(crate) fn offset(self) -> Option { + self.offset + } + + /// The number of output rows still permitted across the remaining row groups. + pub(crate) fn limit(self) -> Option { + self.limit + } + /// Returns how many selected rows remain after applying this budget. pub(crate) fn rows_after(self, rows_before_budget: usize) -> usize { let rows_after_offset = rows_before_budget.saturating_sub(self.offset.unwrap_or(0)); @@ -263,6 +273,25 @@ pub(crate) struct RowGroupReaderBuilder { buffers: PushBuffers, } +/// The parts of a [`RowGroupReaderBuilder`] needed to rebuild it, recovered by +/// [`RowGroupReaderBuilder::into_parts`]. +/// +/// `metadata` is not included: it is a whole-file property carried alongside +/// `schema` in `RemainingRowGroupsParts`. +#[derive(Debug)] +pub(crate) struct RowGroupReaderBuilderParts { + pub batch_size: usize, + pub projection: ProjectionMask, + pub fields: Option>, + pub filter: Option, + pub max_predicate_cache_size: usize, + pub metrics: ArrowReaderMetrics, + pub row_selection_policy: RowSelectionPolicy, + /// Bytes already pushed into the decoder, carried across a rebuild so they + /// are not re-requested. + pub buffers: PushBuffers, +} + impl RowGroupReaderBuilder { /// Create a new RowGroupReaderBuilder #[expect(clippy::too_many_arguments)] @@ -291,11 +320,49 @@ impl RowGroupReaderBuilder { } } + /// Decompose into [`RowGroupReaderBuilderParts`] so the builder can be + /// reconstructed. The runtime decode `state` is discarded; `metadata` is + /// recovered from the frontier instead (see `RemainingRowGroups::into_parts`). + pub(crate) fn into_parts(self) -> RowGroupReaderBuilderParts { + // If a new field is added to `RowGroupReaderBuilder`, it must be added here and in `RowGroupReaderBuilderParts`, + // or at least evaluate how it should be handled in the decomposition and reconstruction of the builder. + let Self { + batch_size, + projection, + metadata: _, + fields, + filter, + max_predicate_cache_size, + metrics, + row_selection_policy, + state: _, + buffers, + } = self; + RowGroupReaderBuilderParts { + batch_size, + projection, + fields, + filter, + max_predicate_cache_size, + metrics, + row_selection_policy, + buffers, + } + } + /// Push new data buffers that can be used to satisfy pending requests pub fn push_data(&mut self, ranges: Vec>, buffers: Vec) { self.buffers.push_ranges(ranges, buffers); } + /// True iff the inner state is `Finished`. This is the only state in + /// which it is safe to decompose the builder via [`Self::into_parts`], + /// because no `RowGroupInfo`, `FilterInfo`, or in-flight `DataRequest` + /// is referencing the row-group-scoped decode state. + pub(crate) fn is_finished(&self) -> bool { + matches!(self.state, Some(RowGroupDecoderState::Finished)) + } + /// Returns the total number of buffered bytes available pub fn buffered_bytes(&self) -> u64 { self.buffers.buffered_bytes() diff --git a/parquet/src/arrow/push_decoder/remaining.rs b/parquet/src/arrow/push_decoder/remaining.rs index 33e13abf9c12..d1070d2aa69f 100644 --- a/parquet/src/arrow/push_decoder/remaining.rs +++ b/parquet/src/arrow/push_decoder/remaining.rs @@ -18,10 +18,11 @@ use crate::DecodeResult; use crate::arrow::arrow_reader::{ParquetRecordBatchReader, RowSelection}; use crate::arrow::push_decoder::reader_builder::{ - RowBudget, RowGroupBuildResult, RowGroupReaderBuilder, + RowBudget, RowGroupBuildResult, RowGroupReaderBuilder, RowGroupReaderBuilderParts, }; use crate::errors::ParquetError; use crate::file::metadata::ParquetMetaData; +use arrow_schema::SchemaRef; use bytes::Bytes; use std::collections::VecDeque; use std::ops::Range; @@ -185,6 +186,11 @@ impl RowGroupFrontier { /// work item. [`RowGroupReaderBuilder`] owns decoding for the active row group. #[derive(Debug)] pub(crate) struct RemainingRowGroups { + /// The arrow schema of the decoded output. Carried only so + /// [`Self::into_parts`] can hand it to a rebuilt builder; unused while + /// decoding. + schema: SchemaRef, + /// Cross-row-group scan state for queued work. frontier: RowGroupFrontier, @@ -192,8 +198,30 @@ pub(crate) struct RemainingRowGroups { row_group_reader_builder: RowGroupReaderBuilder, } +/// The state recovered from a [`RemainingRowGroups`] by +/// [`RemainingRowGroups::into_parts`], describing the row groups *not* yet +/// decoded so a builder reconstructed from it resumes where the decoder left off. +#[derive(Debug)] +pub(crate) struct RemainingRowGroupsParts { + /// The arrow schema of the decoded output. + pub schema: SchemaRef, + /// The Parquet file metadata. + pub metadata: Arc, + /// Row groups not yet handed to the reader builder. + pub row_groups: Vec, + /// The not-yet-consumed slice of the global row selection. + pub selection: Option, + /// Offset still to be skipped before the next readable row group. + pub offset: Option, + /// Output rows still permitted across the remaining row groups. + pub limit: Option, + /// Builder-configurable parts of the inner row-group reader builder. + pub reader_builder: RowGroupReaderBuilderParts, +} + impl RemainingRowGroups { pub fn new( + schema: SchemaRef, parquet_metadata: Arc, row_groups: Vec, selection: Option, @@ -202,6 +230,7 @@ impl RemainingRowGroups { row_group_reader_builder: RowGroupReaderBuilder, ) -> Self { Self { + schema, frontier: RowGroupFrontier::new( parquet_metadata, row_groups, @@ -213,6 +242,36 @@ impl RemainingRowGroups { } } + /// Decompose into [`RemainingRowGroupsParts`]. + /// + /// Must be called at a row-group boundary (see + /// [`Self::is_at_row_group_boundary`]). The inner reader builder's runtime + /// decode state is discarded; its buffered bytes are carried through. + pub(crate) fn into_parts(self) -> RemainingRowGroupsParts { + let Self { + schema, + frontier, + row_group_reader_builder, + } = self; + // `has_predicates` is recomputed by `build()` from the filter. + let RowGroupFrontier { + parquet_metadata, + row_groups, + selection, + budget, + has_predicates: _, + } = frontier; + RemainingRowGroupsParts { + schema, + metadata: parquet_metadata, + row_groups: Vec::from(row_groups), + selection, + offset: budget.offset(), + limit: budget.limit(), + reader_builder: row_group_reader_builder.into_parts(), + } + } + /// Push new data buffers that can be used to satisfy pending requests pub fn push_data(&mut self, ranges: Vec>, buffers: Vec) { self.row_group_reader_builder.push_data(ranges, buffers); @@ -228,6 +287,18 @@ impl RemainingRowGroups { self.row_group_reader_builder.clear_all_ranges(); } + /// True iff the inner row-group reader is between row groups (state + /// `Finished`). Forward to [`RowGroupReaderBuilder::is_finished`]. + pub fn is_at_row_group_boundary(&self) -> bool { + self.row_group_reader_builder.is_finished() + } + + /// Number of row groups remaining (not including the one currently + /// being decoded). + pub fn row_groups_remaining(&self) -> usize { + self.frontier.row_groups.len() + } + /// returns [`ParquetRecordBatchReader`] suitable for reading the next /// group of rows from the Parquet data, or the list of data ranges still /// needed to proceed diff --git a/parquet/src/util/push_buffers.rs b/parquet/src/util/push_buffers.rs index b8225ab3a1db..be245b3359ab 100644 --- a/parquet/src/util/push_buffers.rs +++ b/parquet/src/util/push_buffers.rs @@ -21,9 +21,12 @@ use bytes::Bytes; use std::fmt::Display; use std::ops::Range; -/// Holds multiple buffers of data +/// Holds multiple non-contiguous, caller-provided buffers of file data. /// -/// This is the in-memory buffer for the ParquetDecoder and ParquetMetadataDecoders +/// This is the in-memory buffer used by the push-based Parquet decoders +/// (`ParquetPushDecoder` and `ParquetMetaDataPushDecoder`). It can be +/// constructed up front and handed to a builder so the decoder reuses bytes +/// that have already been fetched. /// /// Features: /// 1. Zero copy @@ -38,8 +41,8 @@ use std::ops::Range; /// /// Thus, the implementation defers to the caller to coalesce subsequent requests /// if desired. -#[derive(Debug, Clone)] -pub(crate) struct PushBuffers { +#[derive(Debug, Clone, Default)] +pub struct PushBuffers { /// the virtual "offset" of this buffers (added to any request) offset: u64, /// The total length of the file being decoded @@ -75,7 +78,11 @@ impl Display for PushBuffers { } impl PushBuffers { - /// Create a new Buffers instance with the given file length + /// Create a new, empty `PushBuffers` for a file of the given length. + /// + /// Use [`PushBuffers::default`] when the file length is unknown or + /// irrelevant (e.g. the push decoder, which tracks ranges by absolute + /// offset and never consults `file_len`). pub fn new(file_len: u64) -> Self { Self { offset: 0, @@ -109,7 +116,7 @@ impl PushBuffers { } /// Returns true if the Buffers contains data for the given range - pub fn has_range(&self, range: &Range) -> bool { + pub(crate) fn has_range(&self, range: &Range) -> bool { self.ranges .iter() .any(|r| r.start <= range.start && r.end >= range.end) @@ -120,25 +127,25 @@ impl PushBuffers { } /// return the file length of the Parquet file being read - pub fn file_len(&self) -> u64 { + pub(crate) fn file_len(&self) -> u64 { self.file_len } /// Specify a new offset - pub fn with_offset(mut self, offset: u64) -> Self { + fn with_offset(mut self, offset: u64) -> Self { self.offset = offset; self } /// Return the total of all buffered ranges #[cfg(feature = "arrow")] - pub fn buffered_bytes(&self) -> u64 { + pub(crate) fn buffered_bytes(&self) -> u64 { self.ranges.iter().map(|r| r.end - r.start).sum() } /// Clear any range and corresponding buffer that is exactly in the ranges_to_clear #[cfg(feature = "arrow")] - pub fn clear_ranges(&mut self, ranges_to_clear: &[Range]) { + pub(crate) fn clear_ranges(&mut self, ranges_to_clear: &[Range]) { let mut new_ranges = Vec::new(); let mut new_buffers = Vec::new(); @@ -156,7 +163,7 @@ impl PushBuffers { } /// Clear all buffered ranges and their corresponding data - pub fn clear_all_ranges(&mut self) { + pub(crate) fn clear_all_ranges(&mut self) { self.ranges.clear(); self.buffers.clear(); } From e7c37ded17298172ebfbe349d9044ef100f82e40 Mon Sep 17 00:00:00 2001 From: Ed Seidl Date: Wed, 20 May 2026 11:44:15 -0700 Subject: [PATCH 06/51] Add helper functions to create `LogicalType` struct variants (#9996) # Which issue does this PR close? - Part of #9995. # Rationale for this change Before switching `LogicalType` from struct variants to tuple variants, add some helper functions that will hide some of the increase in complexity. # What changes are included in this PR? Adds functions to the `LogicalType` impl for creating instances of the non-unit variants (`Integer`, `Decimal`, `Time`, `Timestamp`, `Variant`, `Geometry`, `Geography`). # Are these changes tested? Should be covered by existing tests. # Are there any user-facing changes? Adds to the `LogicalType` API --- parquet/src/arrow/schema/extension.rs | 18 +- parquet/src/arrow/schema/mod.rs | 59 ++-- parquet/src/basic.rs | 382 +++++++++----------------- parquet/src/column/writer/mod.rs | 5 +- parquet/src/file/writer.rs | 5 +- parquet/src/geospatial/accumulator.rs | 7 +- parquet/src/schema/parser.rs | 57 +--- parquet/src/schema/printer.rs | 71 ++--- parquet/src/schema/types.rs | 20 +- parquet/src/variant.rs | 4 +- parquet_derive/src/parquet_field.rs | 40 +-- 11 files changed, 209 insertions(+), 459 deletions(-) diff --git a/parquet/src/arrow/schema/extension.rs b/parquet/src/arrow/schema/extension.rs index a2e9c32dee8c..0244c1b6bb99 100644 --- a/parquet/src/arrow/schema/extension.rs +++ b/parquet/src/arrow/schema/extension.rs @@ -112,9 +112,7 @@ pub(crate) fn has_extension_type(parquet_type: &Type) -> bool { pub(crate) fn logical_type_for_struct(field: &Field) -> Option { use parquet_variant_compute::VariantType; if field.has_valid_extension_type::() { - Some(LogicalType::Variant { - specification_version: None, - }) + Some(LogicalType::variant(None)) } else { None } @@ -167,13 +165,13 @@ pub(crate) fn logical_type_for_binary(field: &Field) -> Option { match field.extension_type_name() { Some(n) if n == WkbType::NAME => match field.try_extension_type::() { Ok(wkb_type) => match wkb_type.metadata().type_hint() { - WkbTypeHint::Geometry => Some(LogicalType::Geometry { - crs: wkb_type.metadata().crs.as_ref().map(|c| c.to_string()), - }), - WkbTypeHint::Geography => Some(LogicalType::Geography { - crs: wkb_type.metadata().crs.as_ref().map(|c| c.to_string()), - algorithm: wkb_type.metadata().algorithm.map(|a| a.into()), - }), + WkbTypeHint::Geometry => Some(LogicalType::geometry( + wkb_type.metadata().crs.as_ref().map(|c| c.to_string()), + )), + WkbTypeHint::Geography => Some(LogicalType::geography( + wkb_type.metadata().crs.as_ref().map(|c| c.to_string()), + wkb_type.metadata().algorithm.map(|a| a.into()), + )), }, Err(_e) => None, }, diff --git a/parquet/src/arrow/schema/mod.rs b/parquet/src/arrow/schema/mod.rs index b2b93687ba89..7fe6fbc9d93d 100644 --- a/parquet/src/arrow/schema/mod.rs +++ b/parquet/src/arrow/schema/mod.rs @@ -556,18 +556,12 @@ fn arrow_to_parquet_type(field: &Field, coerce_types: bool) -> Result { .with_id(id) .build(), DataType::Int8 => Type::primitive_type_builder(name, PhysicalType::INT32) - .with_logical_type(Some(LogicalType::Integer { - bit_width: 8, - is_signed: true, - })) + .with_logical_type(Some(LogicalType::integer(8, true))) .with_repetition(repetition) .with_id(id) .build(), DataType::Int16 => Type::primitive_type_builder(name, PhysicalType::INT32) - .with_logical_type(Some(LogicalType::Integer { - bit_width: 16, - is_signed: true, - })) + .with_logical_type(Some(LogicalType::integer(16, true))) .with_repetition(repetition) .with_id(id) .build(), @@ -580,34 +574,22 @@ fn arrow_to_parquet_type(field: &Field, coerce_types: bool) -> Result { .with_id(id) .build(), DataType::UInt8 => Type::primitive_type_builder(name, PhysicalType::INT32) - .with_logical_type(Some(LogicalType::Integer { - bit_width: 8, - is_signed: false, - })) + .with_logical_type(Some(LogicalType::integer(8, false))) .with_repetition(repetition) .with_id(id) .build(), DataType::UInt16 => Type::primitive_type_builder(name, PhysicalType::INT32) - .with_logical_type(Some(LogicalType::Integer { - bit_width: 16, - is_signed: false, - })) + .with_logical_type(Some(LogicalType::integer(16, false))) .with_repetition(repetition) .with_id(id) .build(), DataType::UInt32 => Type::primitive_type_builder(name, PhysicalType::INT32) - .with_logical_type(Some(LogicalType::Integer { - bit_width: 32, - is_signed: false, - })) + .with_logical_type(Some(LogicalType::integer(32, false))) .with_repetition(repetition) .with_id(id) .build(), DataType::UInt64 => Type::primitive_type_builder(name, PhysicalType::INT64) - .with_logical_type(Some(LogicalType::Integer { - bit_width: 64, - is_signed: false, - })) + .with_logical_type(Some(LogicalType::integer(64, false))) .with_repetition(repetition) .with_id(id) .build(), @@ -634,16 +616,16 @@ fn arrow_to_parquet_type(field: &Field, coerce_types: bool) -> Result { } DataType::Timestamp(time_unit, tz) => { Type::primitive_type_builder(name, PhysicalType::INT64) - .with_logical_type(Some(LogicalType::Timestamp { + .with_logical_type(Some(LogicalType::timestamp( // If timezone set, values are normalized to UTC timezone - is_adjusted_to_u_t_c: matches!(tz, Some(z) if !z.as_ref().is_empty()), - unit: match time_unit { + matches!(tz, Some(z) if !z.as_ref().is_empty()), + match time_unit { TimeUnit::Second => unreachable!(), TimeUnit::Millisecond => ParquetTimeUnit::MILLIS, TimeUnit::Microsecond => ParquetTimeUnit::MICROS, TimeUnit::Nanosecond => ParquetTimeUnit::NANOS, }, - })) + ))) .with_repetition(repetition) .with_id(id) .build() @@ -675,25 +657,25 @@ fn arrow_to_parquet_type(field: &Field, coerce_types: bool) -> Result { .build() } DataType::Time32(unit) => Type::primitive_type_builder(name, PhysicalType::INT32) - .with_logical_type(Some(LogicalType::Time { - is_adjusted_to_u_t_c: field.metadata().contains_key("adjusted_to_utc"), - unit: match unit { + .with_logical_type(Some(LogicalType::time( + field.metadata().contains_key("adjusted_to_utc"), + match unit { TimeUnit::Millisecond => ParquetTimeUnit::MILLIS, u => unreachable!("Invalid unit for Time32: {:?}", u), }, - })) + ))) .with_repetition(repetition) .with_id(id) .build(), DataType::Time64(unit) => Type::primitive_type_builder(name, PhysicalType::INT64) - .with_logical_type(Some(LogicalType::Time { - is_adjusted_to_u_t_c: field.metadata().contains_key("adjusted_to_utc"), - unit: match unit { + .with_logical_type(Some(LogicalType::time( + field.metadata().contains_key("adjusted_to_utc"), + match unit { TimeUnit::Microsecond => ParquetTimeUnit::MICROS, TimeUnit::Nanosecond => ParquetTimeUnit::NANOS, u => unreachable!("Invalid unit for Time64: {:?}", u), }, - })) + ))) .with_repetition(repetition) .with_id(id) .build(), @@ -749,10 +731,7 @@ fn arrow_to_parquet_type(field: &Field, coerce_types: bool) -> Result { .with_repetition(repetition) .with_id(id) .with_length(length) - .with_logical_type(Some(LogicalType::Decimal { - scale: *scale as i32, - precision: *precision as i32, - })) + .with_logical_type(Some(LogicalType::decimal(*scale as i32, *precision as i32))) .with_precision(*precision as i32) .with_scale(*scale as i32) .build() diff --git a/parquet/src/basic.rs b/parquet/src/basic.rs index 17e5b9c321bb..796779358cd2 100644 --- a/parquet/src/basic.rs +++ b/parquet/src/basic.rs @@ -314,6 +314,54 @@ pub enum LogicalType { }, } +impl LogicalType { + /// Create a [`LogicalType::Integer`] variant with the given `bit_width` and `is_signed` + pub fn integer(bit_width: i8, is_signed: bool) -> Self { + Self::Integer { + bit_width, + is_signed, + } + } + + /// Create a [`LogicalType::Decimal`] variant with the given `scale` and `precision` + pub fn decimal(scale: i32, precision: i32) -> Self { + Self::Decimal { scale, precision } + } + + /// Create a [`LogicalType::Time`] variant with the given `is_adjusted_to_u_t_c` and `unit` + pub fn time(is_adjusted_to_u_t_c: bool, unit: TimeUnit) -> Self { + Self::Time { + is_adjusted_to_u_t_c, + unit, + } + } + + /// Create a [`LogicalType::Timestamp`] variant with the given `is_adjusted_to_u_t_c` and `unit` + pub fn timestamp(is_adjusted_to_u_t_c: bool, unit: TimeUnit) -> Self { + Self::Timestamp { + is_adjusted_to_u_t_c, + unit, + } + } + + /// Create a [`LogicalType::Variant`] variant with the given `specification_version` + pub fn variant(specification_version: Option) -> Self { + Self::Variant { + specification_version, + } + } + + /// Create a [`LogicalType::Geometry`] variant with the given `crs` + pub fn geometry(crs: Option) -> Self { + Self::Geometry { crs } + } + + /// Create a [`LogicalType::Geography`] variant with the given `crs` and `algorithm` + pub fn geography(crs: Option, algorithm: Option) -> Self { + Self::Geography { crs, algorithm } + } +} + impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for LogicalType { fn read_thrift(prot: &mut R) -> Result { let field_ident = prot.read_field_begin(0)?; @@ -339,10 +387,7 @@ impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for LogicalType { } 5 => { let val = DecimalType::read_thrift(&mut *prot)?; - Self::Decimal { - scale: val.scale, - precision: val.precision, - } + Self::decimal(val.scale, val.precision) } 6 => { prot.skip_empty_struct()?; @@ -350,24 +395,15 @@ impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for LogicalType { } 7 => { let val = TimeType::read_thrift(&mut *prot)?; - Self::Time { - is_adjusted_to_u_t_c: val.is_adjusted_to_u_t_c, - unit: val.unit, - } + Self::time(val.is_adjusted_to_u_t_c, val.unit) } 8 => { let val = TimestampType::read_thrift(&mut *prot)?; - Self::Timestamp { - is_adjusted_to_u_t_c: val.is_adjusted_to_u_t_c, - unit: val.unit, - } + Self::timestamp(val.is_adjusted_to_u_t_c, val.unit) } 10 => { let val = IntType::read_thrift(&mut *prot)?; - Self::Integer { - is_signed: val.is_signed, - bit_width: val.bit_width, - } + Self::integer(val.bit_width, val.is_signed) } 11 => { prot.skip_empty_struct()?; @@ -391,15 +427,11 @@ impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for LogicalType { } 16 => { let val = VariantType::read_thrift(&mut *prot)?; - Self::Variant { - specification_version: val.specification_version, - } + Self::variant(val.specification_version) } 17 => { let val = GeometryType::read_thrift(&mut *prot)?; - Self::Geometry { - crs: val.crs.map(|s| s.to_owned()), - } + Self::geometry(val.crs.map(|s| s.to_owned())) } 18 => { let val = GeographyType::read_thrift(&mut *prot)?; @@ -408,10 +440,7 @@ impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for LogicalType { let algorithm = val .algorithm .unwrap_or(EdgeInterpolationAlgorithm::SPHERICAL); - Self::Geography { - crs: val.crs.map(|s| s.to_owned()), - algorithm: Some(algorithm), - } + Self::geography(val.crs.map(|s| s.to_owned()), Some(algorithm)) } _ => { prot.skip(field_ident.field_type)?; @@ -1511,26 +1540,14 @@ impl str::FromStr for LogicalType { fn from_str(s: &str) -> Result { match s { // The type is a placeholder that gets updated elsewhere - "INTEGER" => Ok(LogicalType::Integer { - bit_width: 8, - is_signed: false, - }), + "INTEGER" => Ok(LogicalType::integer(8, false)), "MAP" => Ok(LogicalType::Map), "LIST" => Ok(LogicalType::List), "ENUM" => Ok(LogicalType::Enum), - "DECIMAL" => Ok(LogicalType::Decimal { - precision: -1, - scale: -1, - }), + "DECIMAL" => Ok(LogicalType::decimal(-1, -1)), "DATE" => Ok(LogicalType::Date), - "TIME" => Ok(LogicalType::Time { - is_adjusted_to_u_t_c: false, - unit: TimeUnit::MILLIS, - }), - "TIMESTAMP" => Ok(LogicalType::Timestamp { - is_adjusted_to_u_t_c: false, - unit: TimeUnit::MILLIS, - }), + "TIME" => Ok(LogicalType::time(false, TimeUnit::MILLIS)), + "TIMESTAMP" => Ok(LogicalType::timestamp(false, TimeUnit::MILLIS)), "STRING" => Ok(LogicalType::String), "JSON" => Ok(LogicalType::Json), "BSON" => Ok(LogicalType::Bson), @@ -1540,11 +1557,12 @@ impl str::FromStr for LogicalType { "Interval parquet logical type not yet supported" )), "FLOAT16" => Ok(LogicalType::Float16), - "GEOMETRY" => Ok(LogicalType::Geometry { crs: None }), - "GEOGRAPHY" => Ok(LogicalType::Geography { - crs: None, - algorithm: Some(EdgeInterpolationAlgorithm::SPHERICAL), - }), + "VARIANT" => Ok(LogicalType::variant(None)), + "GEOMETRY" => Ok(LogicalType::geometry(None)), + "GEOGRAPHY" => Ok(LogicalType::geography( + None, + Some(EdgeInterpolationAlgorithm::SPHERICAL), + )), other => Err(general_err!("Invalid parquet logical type {}", other)), } } @@ -1860,10 +1878,7 @@ mod tests { let logical_none: Option = None; assert_eq!(ConvertedType::from(logical_none), ConvertedType::NONE); assert_eq!( - ConvertedType::from(Some(LogicalType::Decimal { - precision: 20, - scale: 5 - })), + ConvertedType::from(Some(LogicalType::decimal(5, 20))), ConvertedType::DECIMAL ); assert_eq!( @@ -1883,101 +1898,59 @@ mod tests { ConvertedType::DATE ); assert_eq!( - ConvertedType::from(Some(LogicalType::Time { - unit: TimeUnit::MILLIS, - is_adjusted_to_u_t_c: true, - })), + ConvertedType::from(Some(LogicalType::time(true, TimeUnit::MILLIS))), ConvertedType::TIME_MILLIS ); assert_eq!( - ConvertedType::from(Some(LogicalType::Time { - unit: TimeUnit::MICROS, - is_adjusted_to_u_t_c: true, - })), + ConvertedType::from(Some(LogicalType::time(true, TimeUnit::MICROS))), ConvertedType::TIME_MICROS ); assert_eq!( - ConvertedType::from(Some(LogicalType::Time { - unit: TimeUnit::NANOS, - is_adjusted_to_u_t_c: false, - })), + ConvertedType::from(Some(LogicalType::time(false, TimeUnit::NANOS))), ConvertedType::NONE ); assert_eq!( - ConvertedType::from(Some(LogicalType::Timestamp { - unit: TimeUnit::MILLIS, - is_adjusted_to_u_t_c: true, - })), + ConvertedType::from(Some(LogicalType::timestamp(true, TimeUnit::MILLIS))), ConvertedType::TIMESTAMP_MILLIS ); assert_eq!( - ConvertedType::from(Some(LogicalType::Timestamp { - unit: TimeUnit::MICROS, - is_adjusted_to_u_t_c: false, - })), + ConvertedType::from(Some(LogicalType::timestamp(false, TimeUnit::MICROS))), ConvertedType::TIMESTAMP_MICROS ); assert_eq!( - ConvertedType::from(Some(LogicalType::Timestamp { - unit: TimeUnit::NANOS, - is_adjusted_to_u_t_c: false, - })), + ConvertedType::from(Some(LogicalType::timestamp(false, TimeUnit::NANOS))), ConvertedType::NONE ); assert_eq!( - ConvertedType::from(Some(LogicalType::Integer { - bit_width: 8, - is_signed: false - })), + ConvertedType::from(Some(LogicalType::integer(8, false))), ConvertedType::UINT_8 ); assert_eq!( - ConvertedType::from(Some(LogicalType::Integer { - bit_width: 8, - is_signed: true - })), + ConvertedType::from(Some(LogicalType::integer(8, true))), ConvertedType::INT_8 ); assert_eq!( - ConvertedType::from(Some(LogicalType::Integer { - bit_width: 16, - is_signed: false - })), + ConvertedType::from(Some(LogicalType::integer(16, false))), ConvertedType::UINT_16 ); assert_eq!( - ConvertedType::from(Some(LogicalType::Integer { - bit_width: 16, - is_signed: true - })), + ConvertedType::from(Some(LogicalType::integer(16, true))), ConvertedType::INT_16 ); assert_eq!( - ConvertedType::from(Some(LogicalType::Integer { - bit_width: 32, - is_signed: false - })), + ConvertedType::from(Some(LogicalType::integer(32, false))), ConvertedType::UINT_32 ); assert_eq!( - ConvertedType::from(Some(LogicalType::Integer { - bit_width: 32, - is_signed: true - })), + ConvertedType::from(Some(LogicalType::integer(32, true))), ConvertedType::INT_32 ); assert_eq!( - ConvertedType::from(Some(LogicalType::Integer { - bit_width: 64, - is_signed: false - })), + ConvertedType::from(Some(LogicalType::integer(64, false))), ConvertedType::UINT_64 ); assert_eq!( - ConvertedType::from(Some(LogicalType::Integer { - bit_width: 64, - is_signed: true - })), + ConvertedType::from(Some(LogicalType::integer(64, true))), ConvertedType::INT_64 ); assert_eq!( @@ -2001,14 +1974,15 @@ mod tests { ConvertedType::NONE ); assert_eq!( - ConvertedType::from(Some(LogicalType::Geometry { crs: None })), + ConvertedType::from(Some(LogicalType::variant(None))), ConvertedType::NONE ); assert_eq!( - ConvertedType::from(Some(LogicalType::Geography { - crs: None, - algorithm: Some(EdgeInterpolationAlgorithm::default()), - })), + ConvertedType::from(Some(LogicalType::geometry(None))), + ConvertedType::NONE + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::geography(None, Some(Default::default())))), ConvertedType::NONE ); assert_eq!( @@ -2023,81 +1997,42 @@ mod tests { test_roundtrip(LogicalType::Map); test_roundtrip(LogicalType::List); test_roundtrip(LogicalType::Enum); - test_roundtrip(LogicalType::Decimal { - scale: 0, - precision: 20, - }); + test_roundtrip(LogicalType::decimal(0, 20)); test_roundtrip(LogicalType::Date); - test_roundtrip(LogicalType::Time { - is_adjusted_to_u_t_c: true, - unit: TimeUnit::MICROS, - }); - test_roundtrip(LogicalType::Time { - is_adjusted_to_u_t_c: false, - unit: TimeUnit::MILLIS, - }); - test_roundtrip(LogicalType::Time { - is_adjusted_to_u_t_c: false, - unit: TimeUnit::NANOS, - }); - test_roundtrip(LogicalType::Timestamp { - is_adjusted_to_u_t_c: false, - unit: TimeUnit::MICROS, - }); - test_roundtrip(LogicalType::Timestamp { - is_adjusted_to_u_t_c: true, - unit: TimeUnit::MILLIS, - }); - test_roundtrip(LogicalType::Timestamp { - is_adjusted_to_u_t_c: true, - unit: TimeUnit::NANOS, - }); - test_roundtrip(LogicalType::Integer { - bit_width: 8, - is_signed: true, - }); - test_roundtrip(LogicalType::Integer { - bit_width: 16, - is_signed: false, - }); - test_roundtrip(LogicalType::Integer { - bit_width: 32, - is_signed: true, - }); - test_roundtrip(LogicalType::Integer { - bit_width: 64, - is_signed: false, - }); + test_roundtrip(LogicalType::time(true, TimeUnit::MICROS)); + test_roundtrip(LogicalType::time(false, TimeUnit::MILLIS)); + test_roundtrip(LogicalType::time(false, TimeUnit::NANOS)); + test_roundtrip(LogicalType::timestamp(false, TimeUnit::MICROS)); + test_roundtrip(LogicalType::timestamp(true, TimeUnit::MILLIS)); + test_roundtrip(LogicalType::timestamp(true, TimeUnit::NANOS)); + test_roundtrip(LogicalType::integer(8, true)); + test_roundtrip(LogicalType::integer(16, false)); + test_roundtrip(LogicalType::integer(32, true)); + test_roundtrip(LogicalType::integer(64, false)); test_roundtrip(LogicalType::Json); test_roundtrip(LogicalType::Bson); test_roundtrip(LogicalType::Uuid); test_roundtrip(LogicalType::Float16); - test_roundtrip(LogicalType::Variant { - specification_version: Some(1), - }); - test_roundtrip(LogicalType::Variant { - specification_version: None, - }); - test_roundtrip(LogicalType::Geometry { - crs: Some("foo".to_owned()), - }); - test_roundtrip(LogicalType::Geometry { crs: None }); - test_roundtrip(LogicalType::Geography { - crs: Some("foo".to_owned()), - algorithm: Some(EdgeInterpolationAlgorithm::ANDOYER), - }); - test_roundtrip(LogicalType::Geography { - crs: None, - algorithm: Some(EdgeInterpolationAlgorithm::KARNEY), - }); - test_roundtrip(LogicalType::Geography { - crs: Some("foo".to_owned()), - algorithm: Some(EdgeInterpolationAlgorithm::SPHERICAL), - }); - test_roundtrip(LogicalType::Geography { - crs: None, - algorithm: Some(EdgeInterpolationAlgorithm::SPHERICAL), - }); + test_roundtrip(LogicalType::variant(Some(1))); + test_roundtrip(LogicalType::variant(None)); + test_roundtrip(LogicalType::geometry(Some("foo".to_owned()))); + test_roundtrip(LogicalType::geometry(None)); + test_roundtrip(LogicalType::geography( + Some("foo".to_owned()), + Some(EdgeInterpolationAlgorithm::ANDOYER), + )); + test_roundtrip(LogicalType::geography( + None, + Some(EdgeInterpolationAlgorithm::KARNEY), + )); + test_roundtrip(LogicalType::geography( + Some("foo".to_owned()), + Some(EdgeInterpolationAlgorithm::SPHERICAL), + )); + test_roundtrip(LogicalType::geography( + None, + Some(EdgeInterpolationAlgorithm::SPHERICAL), + )); } #[test] @@ -2291,72 +2226,27 @@ mod tests { LogicalType::Bson, LogicalType::Enum, LogicalType::Uuid, - LogicalType::Integer { - bit_width: 8, - is_signed: false, - }, - LogicalType::Integer { - bit_width: 16, - is_signed: false, - }, - LogicalType::Integer { - bit_width: 32, - is_signed: false, - }, - LogicalType::Integer { - bit_width: 64, - is_signed: false, - }, + LogicalType::integer(8, false), + LogicalType::integer(16, false), + LogicalType::integer(32, false), + LogicalType::integer(64, false), ]; check_sort_order(unsigned, SortOrder::UNSIGNED); // Signed comparison (physical type does not matter) let signed = vec![ - LogicalType::Integer { - bit_width: 8, - is_signed: true, - }, - LogicalType::Integer { - bit_width: 8, - is_signed: true, - }, - LogicalType::Integer { - bit_width: 8, - is_signed: true, - }, - LogicalType::Integer { - bit_width: 8, - is_signed: true, - }, - LogicalType::Decimal { - scale: 20, - precision: 4, - }, + LogicalType::integer(8, true), + LogicalType::integer(16, true), + LogicalType::integer(32, true), + LogicalType::integer(64, true), + LogicalType::decimal(20, 4), LogicalType::Date, - LogicalType::Time { - is_adjusted_to_u_t_c: false, - unit: TimeUnit::MILLIS, - }, - LogicalType::Time { - is_adjusted_to_u_t_c: false, - unit: TimeUnit::MICROS, - }, - LogicalType::Time { - is_adjusted_to_u_t_c: true, - unit: TimeUnit::NANOS, - }, - LogicalType::Timestamp { - is_adjusted_to_u_t_c: false, - unit: TimeUnit::MILLIS, - }, - LogicalType::Timestamp { - is_adjusted_to_u_t_c: false, - unit: TimeUnit::MICROS, - }, - LogicalType::Timestamp { - is_adjusted_to_u_t_c: true, - unit: TimeUnit::NANOS, - }, + LogicalType::time(false, TimeUnit::MILLIS), + LogicalType::time(false, TimeUnit::MICROS), + LogicalType::time(true, TimeUnit::NANOS), + LogicalType::timestamp(false, TimeUnit::MILLIS), + LogicalType::timestamp(false, TimeUnit::MICROS), + LogicalType::timestamp(true, TimeUnit::NANOS), LogicalType::Float16, ]; check_sort_order(signed, SortOrder::SIGNED); @@ -2365,11 +2255,9 @@ mod tests { let undefined = vec![ LogicalType::List, LogicalType::Map, - LogicalType::Geometry { crs: None }, - LogicalType::Geography { - crs: None, - algorithm: Some(EdgeInterpolationAlgorithm::default()), - }, + LogicalType::variant(None), + LogicalType::geometry(None), + LogicalType::geography(None, Some(Default::default())), ]; check_sort_order(undefined, SortOrder::UNDEFINED); } diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 5d14ac6856f9..595eadbc90f2 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -4449,10 +4449,7 @@ mod tests { let path = ColumnPath::from("col"); let tpe = SchemaType::primitive_type_builder("col", T::get_physical_type()) .with_length(16) - .with_logical_type(Some(LogicalType::Decimal { - scale: 2, - precision: 3, - })) + .with_logical_type(Some(LogicalType::decimal(2, 3))) .with_scale(2) .with_precision(3) .build() diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs index ec4d02a47d79..942013ea6238 100644 --- a/parquet/src/file/writer.rs +++ b/parquet/src/file/writer.rs @@ -1280,10 +1280,7 @@ mod tests { #[test] fn test_file_writer_v2_with_metadata() { let file = tempfile::tempfile().unwrap(); - let field_logical_type = Some(LogicalType::Integer { - bit_width: 8, - is_signed: false, - }); + let field_logical_type = Some(LogicalType::integer(8, false)); let field = Arc::new( types::Type::primitive_type_builder("col1", Type::INT32) .with_logical_type(field_logical_type.clone()) diff --git a/parquet/src/geospatial/accumulator.rs b/parquet/src/geospatial/accumulator.rs index d25a47930f6e..3aad1060e05c 100644 --- a/parquet/src/geospatial/accumulator.rs +++ b/parquet/src/geospatial/accumulator.rs @@ -253,7 +253,7 @@ mod test { // Check that we have a working accumulator for Geometry let parquet_type = Type::primitive_type_builder("geom", crate::basic::Type::BYTE_ARRAY) - .with_logical_type(Some(LogicalType::Geometry { crs: None })) + .with_logical_type(Some(LogicalType::geometry(None))) .build() .unwrap(); let column_descr = @@ -271,10 +271,7 @@ mod test { // Check that we have a void accumulator for Geography let parquet_type = Type::primitive_type_builder("geom", crate::basic::Type::BYTE_ARRAY) - .with_logical_type(Some(LogicalType::Geography { - crs: None, - algorithm: None, - })) + .with_logical_type(Some(LogicalType::geography(None, None))) .build() .unwrap(); let column_descr = diff --git a/parquet/src/schema/parser.rs b/parquet/src/schema/parser.rs index 36cf2dc5175e..071962aa4ed8 100644 --- a/parquet/src/schema/parser.rs +++ b/parquet/src/schema/parser.rs @@ -353,7 +353,7 @@ impl Parser<'_> { } else { scale = 0 } - logical = Some(LogicalType::Decimal { scale, precision }); + logical = Some(LogicalType::decimal(scale, precision)); converted = ConvertedType::from(logical.clone()); } } @@ -371,10 +371,7 @@ impl Parser<'_> { "Failed to parse timezone info for TIME type", )?; assert_token(self.tokenizer.next(), ")")?; - logical = Some(LogicalType::Time { - is_adjusted_to_u_t_c, - unit, - }); + logical = Some(LogicalType::time(is_adjusted_to_u_t_c, unit)); converted = ConvertedType::from(logical.clone()); } else { // Invalid token for unit @@ -396,10 +393,7 @@ impl Parser<'_> { "Failed to parse timezone info for TIMESTAMP type", )?; assert_token(self.tokenizer.next(), ")")?; - logical = Some(LogicalType::Timestamp { - is_adjusted_to_u_t_c, - unit, - }); + logical = Some(LogicalType::timestamp(is_adjusted_to_u_t_c, unit)); converted = ConvertedType::from(logical.clone()); } else { // Invalid token for unit @@ -446,10 +440,7 @@ impl Parser<'_> { "Failed to parse is_signed for INTEGER type", )?; assert_token(self.tokenizer.next(), ")")?; - logical = Some(LogicalType::Integer { - bit_width, - is_signed, - }); + logical = Some(LogicalType::integer(bit_width, is_signed)); converted = ConvertedType::from(logical.clone()); } else { // Invalid token for unit @@ -833,10 +824,7 @@ mod tests { .with_fields(vec![ Arc::new( Type::primitive_type_builder("f1", PhysicalType::FIXED_LEN_BYTE_ARRAY) - .with_logical_type(Some(LogicalType::Decimal { - precision: 9, - scale: 3, - })) + .with_logical_type(Some(LogicalType::decimal(3, 9))) .with_converted_type(ConvertedType::DECIMAL) .with_length(5) .with_precision(9) @@ -846,10 +834,7 @@ mod tests { ), Arc::new( Type::primitive_type_builder("f2", PhysicalType::FIXED_LEN_BYTE_ARRAY) - .with_logical_type(Some(LogicalType::Decimal { - precision: 38, - scale: 18, - })) + .with_logical_type(Some(LogicalType::decimal(18, 38))) .with_converted_type(ConvertedType::DECIMAL) .with_length(16) .with_precision(38) @@ -1038,20 +1023,14 @@ mod tests { Arc::new( Type::primitive_type_builder("_1", PhysicalType::INT32) .with_repetition(Repetition::REQUIRED) - .with_logical_type(Some(LogicalType::Integer { - bit_width: 8, - is_signed: true, - })) + .with_logical_type(Some(LogicalType::integer(8, true))) .build() .unwrap(), ), Arc::new( Type::primitive_type_builder("_2", PhysicalType::INT32) .with_repetition(Repetition::REQUIRED) - .with_logical_type(Some(LogicalType::Integer { - bit_width: 16, - is_signed: false, - })) + .with_logical_type(Some(LogicalType::integer(16, false))) .build() .unwrap(), ), @@ -1075,37 +1054,25 @@ mod tests { ), Arc::new( Type::primitive_type_builder("_6", PhysicalType::INT32) - .with_logical_type(Some(LogicalType::Time { - unit: TimeUnit::MILLIS, - is_adjusted_to_u_t_c: false, - })) + .with_logical_type(Some(LogicalType::time(false, TimeUnit::MILLIS))) .build() .unwrap(), ), Arc::new( Type::primitive_type_builder("_7", PhysicalType::INT64) - .with_logical_type(Some(LogicalType::Time { - unit: TimeUnit::MICROS, - is_adjusted_to_u_t_c: true, - })) + .with_logical_type(Some(LogicalType::time(true, TimeUnit::MICROS))) .build() .unwrap(), ), Arc::new( Type::primitive_type_builder("_8", PhysicalType::INT64) - .with_logical_type(Some(LogicalType::Timestamp { - unit: TimeUnit::MILLIS, - is_adjusted_to_u_t_c: true, - })) + .with_logical_type(Some(LogicalType::timestamp(true, TimeUnit::MILLIS))) .build() .unwrap(), ), Arc::new( Type::primitive_type_builder("_9", PhysicalType::INT64) - .with_logical_type(Some(LogicalType::Timestamp { - unit: TimeUnit::NANOS, - is_adjusted_to_u_t_c: false, - })) + .with_logical_type(Some(LogicalType::timestamp(false, TimeUnit::NANOS))) .build() .unwrap(), ), diff --git a/parquet/src/schema/printer.rs b/parquet/src/schema/printer.rs index 68398005b6a5..dbeddcfc128c 100644 --- a/parquet/src/schema/printer.rs +++ b/parquet/src/schema/printer.rs @@ -457,7 +457,7 @@ mod tests { use std::sync::Arc; - use crate::basic::{EdgeInterpolationAlgorithm, Repetition, Type as PhysicalType}; + use crate::basic::{Repetition, Type as PhysicalType}; use crate::errors::Result; use crate::schema::parser::parse_message_type; @@ -543,10 +543,7 @@ mod tests { "field", None, PhysicalType::INT32, - Some(LogicalType::Integer { - bit_width: 32, - is_signed: true, - }), + Some(LogicalType::integer(32, true)), ConvertedType::NONE, Repetition::REQUIRED, ) @@ -558,10 +555,7 @@ mod tests { "field", None, PhysicalType::INT32, - Some(LogicalType::Integer { - bit_width: 8, - is_signed: false, - }), + Some(LogicalType::integer(8, false)), ConvertedType::NONE, Repetition::OPTIONAL, ) @@ -573,10 +567,7 @@ mod tests { "field", None, PhysicalType::INT32, - Some(LogicalType::Integer { - bit_width: 16, - is_signed: true, - }), + Some(LogicalType::integer(16, true)), ConvertedType::INT_16, Repetition::REPEATED, ) @@ -588,10 +579,7 @@ mod tests { "field", Some(42), PhysicalType::INT32, - Some(LogicalType::Integer { - bit_width: 16, - is_signed: true, - }), + Some(LogicalType::integer(16, true)), ConvertedType::INT_16, Repetition::REPEATED, ) @@ -651,10 +639,7 @@ mod tests { "field", None, PhysicalType::INT64, - Some(LogicalType::Timestamp { - is_adjusted_to_u_t_c: true, - unit: TimeUnit::MILLIS, - }), + Some(LogicalType::timestamp(true, TimeUnit::MILLIS)), ConvertedType::NONE, Repetition::REQUIRED, ) @@ -678,10 +663,7 @@ mod tests { "field", None, PhysicalType::INT32, - Some(LogicalType::Time { - unit: TimeUnit::MILLIS, - is_adjusted_to_u_t_c: false, - }), + Some(LogicalType::time(false, TimeUnit::MILLIS)), ConvertedType::TIME_MILLIS, Repetition::REQUIRED, ) @@ -693,10 +675,7 @@ mod tests { "field", Some(42), PhysicalType::INT32, - Some(LogicalType::Time { - unit: TimeUnit::MILLIS, - is_adjusted_to_u_t_c: false, - }), + Some(LogicalType::time(false, TimeUnit::MILLIS)), ConvertedType::TIME_MILLIS, Repetition::REQUIRED, ) @@ -792,7 +771,7 @@ mod tests { "field", None, PhysicalType::BYTE_ARRAY, - Some(LogicalType::Geometry { crs: None }), + Some(LogicalType::geometry(None)), ConvertedType::NONE, Repetition::REQUIRED, ) @@ -804,9 +783,7 @@ mod tests { "field", None, PhysicalType::BYTE_ARRAY, - Some(LogicalType::Geometry { - crs: Some("non-missing CRS".to_string()), - }), + Some(LogicalType::geometry(Some("non-missing CRS".to_string()))), ConvertedType::NONE, Repetition::REQUIRED, ) @@ -818,10 +795,7 @@ mod tests { "field", None, PhysicalType::BYTE_ARRAY, - Some(LogicalType::Geography { - crs: None, - algorithm: Some(EdgeInterpolationAlgorithm::default()), - }), + Some(LogicalType::geography(None, Some(Default::default()))), ConvertedType::NONE, Repetition::REQUIRED, ) @@ -833,10 +807,10 @@ mod tests { "field", None, PhysicalType::BYTE_ARRAY, - Some(LogicalType::Geography { - crs: Some("non-missing CRS".to_string()), - algorithm: Some(EdgeInterpolationAlgorithm::default()), - }), + Some(LogicalType::geography( + Some("non-missing CRS".to_string()), + Some(Default::default()), + )), ConvertedType::NONE, Repetition::REQUIRED, ) @@ -887,10 +861,7 @@ mod tests { ), ( Type::primitive_type_builder("decimal", PhysicalType::FIXED_LEN_BYTE_ARRAY) - .with_logical_type(Some(LogicalType::Decimal { - precision: 32, - scale: 20, - })) + .with_logical_type(Some(LogicalType::decimal(20, 32))) .with_precision(32) .with_scale(20) .with_length(decimal_length_from_precision(32)) @@ -1178,10 +1149,7 @@ mod tests { fn test_print_and_parse_decimal() { let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32) .with_repetition(Repetition::OPTIONAL) - .with_logical_type(Some(LogicalType::Decimal { - precision: 9, - scale: 2, - })) + .with_logical_type(Some(LogicalType::decimal(2, 9))) .with_converted_type(ConvertedType::DECIMAL) .with_precision(9) .with_scale(2) @@ -1190,10 +1158,7 @@ mod tests { let f2 = Type::primitive_type_builder("f2", PhysicalType::INT32) .with_repetition(Repetition::OPTIONAL) - .with_logical_type(Some(LogicalType::Decimal { - precision: 9, - scale: 0, - })) + .with_logical_type(Some(LogicalType::decimal(0, 9))) .with_converted_type(ConvertedType::DECIMAL) .with_precision(9) .with_scale(0) diff --git a/parquet/src/schema/types.rs b/parquet/src/schema/types.rs index 0d504e16fc28..d8b3456d4723 100644 --- a/parquet/src/schema/types.rs +++ b/parquet/src/schema/types.rs @@ -1449,10 +1449,7 @@ mod tests { #[test] fn test_primitive_type() { let mut result = Type::primitive_type_builder("foo", PhysicalType::INT32) - .with_logical_type(Some(LogicalType::Integer { - bit_width: 32, - is_signed: true, - })) + .with_logical_type(Some(LogicalType::integer(32, true))) .with_id(Some(0)) .build(); assert!(result.is_ok()); @@ -1464,10 +1461,7 @@ mod tests { assert_eq!(basic_info.repetition(), Repetition::OPTIONAL); assert_eq!( basic_info.logical_type_ref(), - Some(&LogicalType::Integer { - bit_width: 32, - is_signed: true - }) + Some(&LogicalType::integer(32, true)) ); assert_eq!(basic_info.converted_type(), ConvertedType::INT_32); assert_eq!(basic_info.id(), 0); @@ -1482,10 +1476,7 @@ mod tests { // Test illegal inputs with logical type result = Type::primitive_type_builder("foo", PhysicalType::INT64) .with_repetition(Repetition::REPEATED) - .with_logical_type(Some(LogicalType::Integer { - is_signed: true, - bit_width: 8, - })) + .with_logical_type(Some(LogicalType::integer(8, true))) .build(); assert!(result.is_err()); if let Err(e) = result { @@ -1524,10 +1515,7 @@ mod tests { result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY) .with_repetition(Repetition::REQUIRED) - .with_logical_type(Some(LogicalType::Decimal { - scale: 32, - precision: 12, - })) + .with_logical_type(Some(LogicalType::decimal(32, 12))) .with_precision(-1) .with_scale(-1) .build(); diff --git a/parquet/src/variant.rs b/parquet/src/variant.rs index cdbdd849683d..55df08673611 100644 --- a/parquet/src/variant.rs +++ b/parquet/src/variant.rs @@ -199,9 +199,7 @@ mod tests { // data should have been written with the Variant logical type assert_eq!( field.get_basic_info().logical_type_ref(), - Some(&crate::basic::LogicalType::Variant { - specification_version: None - }) + Some(&crate::basic::LogicalType::variant(None)) ); } diff --git a/parquet_derive/src/parquet_field.rs b/parquet_derive/src/parquet_field.rs index 7473f2305517..17b8d8543725 100644 --- a/parquet_derive/src/parquet_field.rs +++ b/parquet_derive/src/parquet_field.rs @@ -693,42 +693,18 @@ impl Type { match last_part.trim() { "bool" => quote! { None }, - "u8" => quote! { Some(LogicalType::Integer { - bit_width: 8, - is_signed: false, - }) }, - "u16" => quote! { Some(LogicalType::Integer { - bit_width: 16, - is_signed: false, - }) }, - "u32" => quote! { Some(LogicalType::Integer { - bit_width: 32, - is_signed: false, - }) }, - "u64" => quote! { Some(LogicalType::Integer { - bit_width: 64, - is_signed: false, - }) }, - "i8" => quote! { Some(LogicalType::Integer { - bit_width: 8, - is_signed: true, - }) }, - "i16" => quote! { Some(LogicalType::Integer { - bit_width: 16, - is_signed: true, - }) }, + "u8" => quote! { Some(LogicalType::integer(8, false)) }, + "u16" => quote! { Some(LogicalType::integer(16, false)) }, + "u32" => quote! { Some(LogicalType::integer(32, false)) }, + "u64" => quote! { Some(LogicalType::integer(64, false)) }, + "i8" => quote! { Some(LogicalType::integer(8, true)) }, + "i16" => quote! { Some(LogicalType::integer(16, true)) }, "i32" | "i64" => quote! { None }, "usize" => { - quote! { Some(LogicalType::Integer { - bit_width: usize::BITS as i8, - is_signed: false - }) } + quote! { Some(LogicalType::integer(usize::BITS as i8, false)) } } "isize" => { - quote! { Some(LogicalType::Integer { - bit_width: usize::BITS as i8, - is_signed: true - }) } + quote! { Some(LogicalType::integer(usize::BITS as i8, true)) } } "NaiveDate" => quote! { Some(LogicalType::Date) }, "NaiveDateTime" => quote! { None }, From 39c8814f04150f4b8f1c3b30d7914022e7c2a292 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 20 May 2026 21:44:42 +0300 Subject: [PATCH 07/51] feat: extract `has_false` and `has_true` from BooleanArray to `BooleanBuffer` and reuse for no nulls (#9987) # Which issue does this PR close? N/A # Rationale for this change So we can use the useful helpers without creating temp `BooleanArray` # What changes are included in this PR? Extract non null variants for computing `has_false` and `has_true` from `BooleanArray` to `BooleanBuffer` and call them instead, and copied the tests for not nullable # Are these changes tested? Yes # Are there any user-facing changes? 2 new functions ----- Cc @alamb as talked in: - https://github.com/apache/datafusion/pull/22158#discussion_r3241321436 --------- Co-authored-by: Andrew Lamb --- arrow-array/src/array/boolean_array.rs | 52 +------- arrow-buffer/src/buffer/boolean.rs | 165 ++++++++++++++++++++++++- 2 files changed, 166 insertions(+), 51 deletions(-) diff --git a/arrow-array/src/array/boolean_array.rs b/arrow-array/src/array/boolean_array.rs index 22a1ba7653ac..0fa5b5d91209 100644 --- a/arrow-array/src/array/boolean_array.rs +++ b/arrow-array/src/array/boolean_array.rs @@ -19,7 +19,6 @@ use crate::array::print_long_array; use crate::builder::BooleanBuilder; use crate::iterator::BooleanIter; use crate::{Array, ArrayAccessor, ArrayRef, Scalar}; -use arrow_buffer::bit_chunk_iterator::UnalignedBitChunk; use arrow_buffer::{BooleanBuffer, Buffer, MutableBuffer, NullBuffer, bit_util}; use arrow_data::{ArrayData, ArrayDataBuilder}; use arrow_schema::DataType; @@ -157,16 +156,6 @@ impl BooleanArray { &self.values } - /// Block size for chunked fold operations in [`Self::has_true`] and [`Self::has_false`]. - /// Using `chunks_exact` with this size lets the compiler fully unroll the inner - /// fold (no inner branch/loop), enabling short-circuit exits every N chunks. - const CHUNK_FOLD_BLOCK_SIZE: usize = 16; - - /// Returns an [`UnalignedBitChunk`] over this array's values. - fn unaligned_bit_chunks(&self) -> UnalignedBitChunk<'_> { - UnalignedBitChunk::new(self.values().values(), self.values().offset(), self.len()) - } - /// Returns the number of non null, true values within this array. /// If you only need to check if there is at least one true value, consider using `has_true()` which can short-circuit and be more efficient. pub fn true_count(&self) -> usize { @@ -202,16 +191,7 @@ impl BooleanArray { let value_chunks = self.values().bit_chunks().iter_padded(); null_chunks.zip(value_chunks).any(|(n, v)| (n & v) != 0) } - None => { - let bit_chunks = self.unaligned_bit_chunks(); - let chunks = bit_chunks.chunks(); - let mut exact = chunks.chunks_exact(Self::CHUNK_FOLD_BLOCK_SIZE); - let found = bit_chunks.prefix().unwrap_or(0) != 0 - || exact.any(|block| block.iter().fold(0u64, |acc, &c| acc | c) != 0); - found - || exact.remainder().iter().any(|&c| c != 0) - || bit_chunks.suffix().unwrap_or(0) != 0 - } + None => self.values().has_true(), } } @@ -228,35 +208,7 @@ impl BooleanArray { let value_chunks = self.values().bit_chunks().iter_padded(); null_chunks.zip(value_chunks).any(|(n, v)| (n & !v) != 0) } - None => { - let bit_chunks = self.unaligned_bit_chunks(); - // UnalignedBitChunk zeros padding bits; fill them with 1s so - // they don't appear as false values. - let lead_mask = !((1u64 << bit_chunks.lead_padding()) - 1); - let trail_mask = if bit_chunks.trailing_padding() == 0 { - u64::MAX - } else { - (1u64 << (64 - bit_chunks.trailing_padding())) - 1 - }; - let (prefix_fill, suffix_fill) = match (bit_chunks.prefix(), bit_chunks.suffix()) { - (Some(_), Some(_)) => (!lead_mask, !trail_mask), - (Some(_), None) => (!lead_mask | !trail_mask, 0), - (None, Some(_)) => (0, !trail_mask), - (None, None) => (0, 0), - }; - let chunks = bit_chunks.chunks(); - let mut exact = chunks.chunks_exact(Self::CHUNK_FOLD_BLOCK_SIZE); - let found = bit_chunks - .prefix() - .is_some_and(|v| (v | prefix_fill) != u64::MAX) - || exact - .any(|block| block.iter().fold(u64::MAX, |acc, &c| acc & c) != u64::MAX); - found - || exact.remainder().iter().any(|&c| c != u64::MAX) - || bit_chunks - .suffix() - .is_some_and(|v| (v | suffix_fill) != u64::MAX) - } + None => self.values().has_false(), } } diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index 420bbf59f3be..52df67080a3c 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::bit_chunk_iterator::BitChunks; +use crate::bit_chunk_iterator::{BitChunks, UnalignedBitChunk}; use crate::bit_iterator::{BitIndexIterator, BitIndexU32Iterator, BitIterator, BitSliceIterator}; use crate::bit_util::read_u64; use crate::{ @@ -611,6 +611,11 @@ impl BooleanBuffer { self.into_iter() } + /// Returns an [`UnalignedBitChunk`] over this buffer's values. + fn unaligned_bit_chunks(&self) -> UnalignedBitChunk<'_> { + UnalignedBitChunk::new(self.values(), self.offset(), self.len()) + } + /// Returns an iterator over the set bit positions in this [`BooleanBuffer`] pub fn set_indices(&self) -> BitIndexIterator<'_> { BitIndexIterator::new(self.values(), self.bit_offset, self.bit_len) @@ -625,6 +630,61 @@ impl BooleanBuffer { pub fn set_slices(&self) -> BitSliceIterator<'_> { BitSliceIterator::new(self.values(), self.bit_offset, self.bit_len) } + + /// Block size for chunked fold operations in [`Self::has_true`] and [`Self::has_false`]. + /// Using `chunks_exact` with this size lets the compiler fully unroll the inner + /// fold (no inner branch/loop), enabling short-circuit exits every N chunks. + const CHUNK_FOLD_BLOCK_SIZE: usize = 16; + + /// Returns whether there is at least one `true` value in this buffer. + /// + /// This is more efficient than `count_set_bits() > 0` because it can short-circuit + /// as soon as a `true` value is found, without counting all set bits. + /// + /// Returns `false` for empty buffer. + pub fn has_true(&self) -> bool { + let bit_chunks = self.unaligned_bit_chunks(); + let chunks = bit_chunks.chunks(); + let mut exact = chunks.chunks_exact(Self::CHUNK_FOLD_BLOCK_SIZE); + let found = bit_chunks.prefix().unwrap_or(0) != 0 + || exact.any(|block| block.iter().fold(0u64, |acc, &c| acc | c) != 0); + found || exact.remainder().iter().any(|&c| c != 0) || bit_chunks.suffix().unwrap_or(0) != 0 + } + + /// Returns whether there is at least one `false` value in this buffer. + /// + /// This is more efficient than `len() > count_set_bits()` because it can short-circuit + /// as soon as a `false` value is found, without counting all set bits. + /// + /// Returns `false` for empty buffer. + pub fn has_false(&self) -> bool { + let bit_chunks = self.unaligned_bit_chunks(); + // UnalignedBitChunk zeros padding bits; fill them with 1s so + // they don't appear as false values. + let lead_mask = !((1u64 << bit_chunks.lead_padding()) - 1); + let trail_mask = if bit_chunks.trailing_padding() == 0 { + u64::MAX + } else { + (1u64 << (64 - bit_chunks.trailing_padding())) - 1 + }; + let (prefix_fill, suffix_fill) = match (bit_chunks.prefix(), bit_chunks.suffix()) { + (Some(_), Some(_)) => (!lead_mask, !trail_mask), + (Some(_), None) => (!lead_mask | !trail_mask, 0), + (None, Some(_)) => (0, !trail_mask), + (None, None) => (0, 0), + }; + let chunks = bit_chunks.chunks(); + let mut exact = chunks.chunks_exact(Self::CHUNK_FOLD_BLOCK_SIZE); + let found = bit_chunks + .prefix() + .is_some_and(|v| (v | prefix_fill) != u64::MAX) + || exact.any(|block| block.iter().fold(u64::MAX, |acc, &c| acc & c) != u64::MAX); + found + || exact.remainder().iter().any(|&c| c != u64::MAX) + || bit_chunks + .suffix() + .is_some_and(|v| (v | suffix_fill) != u64::MAX) + } } impl Not for &BooleanBuffer { @@ -1245,4 +1305,107 @@ mod tests { let buffer = BooleanBuffer::new_unset(100); assert_eq!(buffer.clone().find_nth_set_bit_position(0, 1), 100); } + + #[test] + fn test_has_true_has_false_all_true() { + let arr = BooleanBuffer::from(vec![true, true, true]); + assert!(arr.has_true()); + assert!(!arr.has_false()); + } + + #[test] + fn test_has_true_has_false_all_false() { + let arr = BooleanBuffer::from(vec![false, false, false]); + assert!(!arr.has_true()); + assert!(arr.has_false()); + } + + #[test] + fn test_has_true_has_false_mixed() { + let arr = BooleanBuffer::from(vec![true, false, true]); + assert!(arr.has_true()); + assert!(arr.has_false()); + } + + #[test] + fn test_has_true_has_false_empty() { + let arr = BooleanBuffer::from(Vec::::new()); + assert!(!arr.has_true()); + assert!(!arr.has_false()); + } + + #[test] + fn test_has_false_aligned_suffix_all_true() { + let arr = BooleanBuffer::from(vec![true; 129]); + assert!(arr.has_true()); + assert!(!arr.has_false()); + } + + #[test] + fn test_has_false_non_aligned_all_true() { + // 65 elements: exercises the remainder path in has_false + let arr = BooleanBuffer::from(vec![true; 65]); + assert!(arr.has_true()); + assert!(!arr.has_false()); + } + + #[test] + fn test_has_false_non_aligned_last_false() { + // 64 trues + 1 false: remainder path should find the false + let mut values = vec![true; 64]; + values.push(false); + let arr = BooleanBuffer::from(values); + assert!(arr.has_true()); + assert!(arr.has_false()); + } + + #[test] + fn test_has_false_exact_64_all_true() { + // Exactly 64 elements, no remainder + let arr = BooleanBuffer::from(vec![true; 64]); + assert!(arr.has_true()); + assert!(!arr.has_false()); + } + + #[test] + fn test_has_true_has_false_unaligned_slices() { + let cases = [ + (1, 129, true, false), + (3, 130, true, false), + (5, 65, true, false), + (7, 64, true, false), + ]; + + let base = BooleanBuffer::from(vec![true; 300]); + + for (offset, len, expected_has_true, expected_has_false) in cases { + let arr = base.slice(offset, len); + assert_eq!( + arr.has_true(), + expected_has_true, + "offset={offset} len={len}" + ); + assert_eq!( + arr.has_false(), + expected_has_false, + "offset={offset} len={len}" + ); + } + } + + #[test] + fn test_has_true_has_false_exact_multiples_of_64() { + let cases = [ + (64, true, false), + (128, true, false), + (192, true, false), + (256, true, false), + ]; + + for (len, expected_has_true, expected_has_false) in cases { + let arr = BooleanBuffer::from(vec![true; len]); + assert_eq!(arr.has_true(), expected_has_true, "len={len}"); + assert_eq!(arr.has_false(), expected_has_false, "len={len}"); + } + } } From 73c513afd36b0c6a3b41ef9dc135aee568184255 Mon Sep 17 00:00:00 2001 From: RyanStewart <47729789+RyanJamesStewart@users.noreply.github.com> Date: Wed, 20 May 2026 11:49:47 -0700 Subject: [PATCH 08/51] Bulk-fill definition levels for majority-null leaf columns (#9967) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Contributes to #9731. ## AI assistance Implementation drafted with AI assistance and iterated against the benchmarks below. I've reviewed and own the code, including the gate threshold which I picked from the sweep in [Threshold (`BULK_FILL_MIN_LEN`)](#threshold-bulk_fill_min_len). Per the project's [CONTRIBUTING guidance on AI-generated submissions](https://github.com/apache/arrow-rs/blob/main/CONTRIBUTING.md#ai-generated-submissions). ## Rationale for this change When writing a nullable leaf (primitive) Arrow array, `write_leaf` builds the definition-level buffer one element at a time, mapping each null bit to a level. For columns that are mostly null this does ~`num_rows` of branchy work and allocates a `num_rows`-element level buffer even though almost every produced level is the same value. #9954 adds an O(1) fast path for the *entirely* null case; this PR covers the *sparse* (mostly-but-not-entirely null) case it doesn't handle, the literal subject of #9731 ("a column that is 99% null … ~100x more work than necessary"). ## What changes are included in this PR? A single popcount pass over the null mask (`Buffer::count_set_bits_offset`, O(`num_rows`/64)) counts the valid values in the range. When the slice is majority-null, the definition-level buffer is bulk-filled with the null level (a vectorized `Vec::resize` memset) and only the non-null positions (from `NullBuffer::valid_indices()`) are overwritten. The existing per-row path is kept for non-majority-null slices, so balanced and null-light columns are unaffected. Both branches share the same `let range_nulls = nulls.slice(range.start, len)` slicing idiom; the slow path uses `range_nulls.iter()` for the def-level map and `range_nulls.valid_indices().map(|i| i + range.start)` for `non_null_indices`, with no `unsafe`. Output is byte-identical: the level *values* are unchanged, just produced via memset+scatter (fast path) or via the high-level `NullBuffer` iterators (slow path) instead of a manual `BitIndexIterator` walk. ## Threshold (`BULK_FILL_MIN_LEN`) The bulk-fill fast path is gated on two conditions: - `len >= BULK_FILL_MIN_LEN` (currently 64). Per-call slice/popcount/iterator overhead only amortizes on sizable sub-ranges. List/struct paths call `write_leaf` many times with tiny ranges (avg list length 1-5); paying any per-call popcount there would regress them. A threshold sweep at T = {0, 16, 32, 64, 128, 256} on Ryzen 9 9950X shows the regression floor settles by T=32, and the choice of 64 gives ~12x margin over the average list length without losing the flat-primitive wins. - `nulls.null_count() * 2 >= nulls.len()`. The cached `null_count()` is O(1), so this check is free. We use the buffer-wide density as a heuristic for the sub-range; for full-array writes (the primary target, flat primitive columns) it's exact. Even when the gate skips the fast path, evaluating it across high-frequency call sites (~10K calls in some list benchmarks) is a small structural cost (~1-2% on list-sparse cases). The wins on the targeted shapes (-35% sparse-primitive, -66% all-null primitive) far outweigh that. Reducing the cost further would require hoisting the decision into the caller. ## Are these changes tested? Existing tests cover this path: `cargo test -p parquet --features arrow --lib arrow_writer` is green (136 tests, full of nulls and roundtrips); full `cargo test -p parquet --features arrow` green modulo the pre-existing `PARQUET_TEST_DATA` submodule failures (unrelated, same on `main`). `cargo clippy -p parquet --features arrow --lib` and `cargo fmt --check` clean. The `unsafe get_unchecked_mut` flagged in the original revision was replaced via `NullBuffer::valid_indices()`; the slow-path also dropped its `unsafe value_unchecked` for the same reason. ## Are there any user-facing changes? None. ## Benchmarks `cargo bench -p parquet --bench arrow_writer`, 1M rows × 7 nullable primitive columns, local Ryzen 9 9950X: ``` primitive_sparse_99pct_null/default 11.88 ms -> 9.13 ms (-23%) <- the case #9731 calls out primitive_all_null/default 5.65 ms -> 2.33 ms (-59%) (subsumed by #9954's O(1) path if that lands first) struct_sparse_99pct_null/default 5.67 ms -> 5.32 ms (-6%) struct_all_null/default 1.52 ms -> 1.31 ms (-14%) list_primitive_sparse_99pct_null, primitive (25% null), primitive_non_null, bool, string: within noise (no regression) ``` The CI benchmark bot (GKE `c4a-highmem-16`, Neoverse-V2) on the post-fixup revision shows the same shape with stronger relative wins on the targeted cases: ``` primitive_all_null/default 2.47x (11.0ms -> 4.4ms) primitive_sparse_99pct_null/default 1.60x (16.8ms -> 10.5ms) primitive_all_null/{bloom_filter,cdc,parquet_2,zstd,zstd_parquet_2} 1.38x to 2.48x primitive_sparse_99pct_null/{...} 1.28x to 1.59x list_primitive*, list_primitive_sparse_99pct_null*: 1.00x to 1.01x (within noise) ``` Microbench of the definition-level fill in isolation: 10.3x @ 100%-null, 8.6x @ 99%, 5.2x @ 90%, 1.9x @ 50%, 0.93x @ 10%, 0.81x @ 0%. Crossover ≈ 12-15% null, clean win above ~25%; the `>= 50% null` guard is conservative. This is the *materialization*-cost half of #9731 (~30% of the 99%-null write); the *walk*-cost half, a run-length input to the level encoder so the column writer doesn't even iterate all `num_rows` levels, is the larger structural change #9653 is heading toward. This PR is deliberately small and isolated so it lands independently of and rebases cleanly under that work. --------- Co-authored-by: Ryan Stewart --- parquet/src/arrow/arrow_writer/levels.rs | 52 +++++++++++++++++++----- 1 file changed, 41 insertions(+), 11 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/levels.rs b/parquet/src/arrow/arrow_writer/levels.rs index 73f922137620..10f90f707c08 100644 --- a/parquet/src/arrow/arrow_writer/levels.rs +++ b/parquet/src/arrow/arrow_writer/levels.rs @@ -153,6 +153,13 @@ enum LevelInfoBuilder { Struct(Vec, LevelContext, Option), } +/// Minimum sub-range length before the bulk-fill fast path in `write_leaf` +/// becomes profitable for null-heavy leaf columns. Below this, per-call +/// slice + popcount overhead regresses list/struct paths that call +/// `write_leaf` many times with tiny ranges. Picked via threshold sweep; +/// see for the rationale. +const BULK_FILL_MIN_LEN: usize = 64; + impl LevelInfoBuilder { /// Create a new [`LevelInfoBuilder`] for the given [`Field`] and parent [`LevelContext`] fn try_new(field: &Field, parent_ctx: LevelContext, array: &ArrayRef) -> Result { @@ -676,20 +683,43 @@ impl LevelInfoBuilder { match &info.logical_nulls { Some(nulls) => { assert!(range.end <= nulls.len()); - let nulls = nulls.inner(); - info.def_levels.extend_from_iter(range.clone().map(|i| { - // Safety: range.end was asserted to be in bounds earlier - let valid = unsafe { nulls.value_unchecked(i) }; - max_def_level - (!valid as i16) - })); - info.non_null_indices.reserve(len); - info.non_null_indices.extend( - BitIndexIterator::new(nulls.inner(), nulls.offset() + range.start, len) - .map(|i| i + range.start), - ); + // Bulk-fill is profitable only on null-heavy ranges long enough to + // amortize the slice/popcount cost; see `BULK_FILL_MIN_LEN` and the + // PR description for the threshold sweep. The gate uses the cached + // buffer-wide `null_count` (O(1)) to stay cheap on the cold path. + if len >= BULK_FILL_MIN_LEN && nulls.null_count() * 2 >= nulls.len() { + let range_nulls = nulls.slice(range.start, len); + let valid_in_range = len - range_nulls.null_count(); + let null_def_level = max_def_level - 1; + let buf = info + .def_levels + .materialize_mut() + .expect("definition levels present"); + let base = buf.len(); + buf.resize(base + len, null_def_level); + for i in range_nulls.valid_indices() { + buf[base + i] = max_def_level; + } + info.non_null_indices.reserve(valid_in_range); + info.non_null_indices + .extend(range_nulls.valid_indices().map(|i| i + range.start)); + } else { + let bits = nulls.inner(); + info.def_levels.extend_from_iter(range.clone().map(|i| { + // Safety: range.end was asserted to be in bounds earlier + let valid = unsafe { bits.value_unchecked(i) }; + max_def_level - (!valid as i16) + })); + info.non_null_indices.reserve(len); + info.non_null_indices.extend( + BitIndexIterator::new(bits.inner(), bits.offset() + range.start, len) + .map(|i| i + range.start), + ); + } } None => { info.append_def_level_run(max_def_level, len); + info.non_null_indices.reserve(len); info.non_null_indices.extend(range.clone()); } } From dc821a93e9eb0b49b77cbea12915076f7142bedf Mon Sep 17 00:00:00 2001 From: Rishab Joshi <8187657+rishvin@users.noreply.github.com> Date: Wed, 20 May 2026 13:44:00 -0700 Subject: [PATCH 09/51] Add support for FixedSizeList to variant_to_arrow (#9663) - Closes #9531 # Rationale for this change Add support for `FixedSizeList` when invoking `variant_to_arrow`. # What changes are included in this PR? - Introduces a new builder `VariantToFixedSizeListArrowRowBuilder`. - Adds test cases for shredding and getting variant by `FixedSizeList`. # Are these changes tested? By adding few test cases. # Are there any user-facing changes? N/A. --------- Co-authored-by: Konstantin Tarasov <33369833+sdf-jkl@users.noreply.github.com> --- parquet-variant-compute/src/shred_variant.rs | 88 +++++++++++-- parquet-variant-compute/src/variant_get.rs | 87 ++++++++---- .../src/variant_to_arrow.rs | 124 +++++++++++++++++- 3 files changed, 259 insertions(+), 40 deletions(-) diff --git a/parquet-variant-compute/src/shred_variant.rs b/parquet-variant-compute/src/shred_variant.rs index 7b919660d69e..440f4b716521 100644 --- a/parquet-variant-compute/src/shred_variant.rs +++ b/parquet-variant-compute/src/shred_variant.rs @@ -68,6 +68,14 @@ use std::sync::Arc; /// See [`ShreddedSchemaBuilder`] for a convenient way to build the `as_type` /// value passed to this function. pub fn shred_variant(array: &VariantArray, as_type: &DataType) -> Result { + shred_variant_with_options(array, as_type, &CastOptions::default()) +} + +pub(crate) fn shred_variant_with_options( + array: &VariantArray, + as_type: &DataType, + cast_options: &CastOptions, +) -> Result { if array.typed_value_field().is_some() { return Err(ArrowError::InvalidArgumentError( "Input is already shredded".to_string(), @@ -79,10 +87,9 @@ pub fn shred_variant(array: &VariantArray, as_type: &DataType) -> Result VariantToShreddedArrayVariantRowBuilder<'a> { Variant::List(list) => { self.nulls.append_non_null(); self.value_builder.append_null(); + + // NOTE: A `FixedSizeList` with incorrect size will hard fail during shredding. self.typed_value_builder .append_value(&Variant::List(list))?; Ok(true) @@ -694,9 +703,9 @@ mod tests { use crate::VariantArrayBuilder; use crate::variant_array::{binary_array_value, variant_from_arrays_at}; use arrow::array::{ - Array, BinaryViewArray, FixedSizeBinaryArray, Float64Array, GenericListArray, - GenericListViewArray, Int64Array, LargeBinaryArray, LargeStringArray, ListArray, - ListLikeArray, OffsetSizeTrait, PrimitiveArray, StringArray, + Array, BinaryViewArray, FixedSizeBinaryArray, FixedSizeListArray, Float64Array, + GenericListArray, GenericListViewArray, Int64Array, LargeBinaryArray, LargeStringArray, + ListArray, ListLikeArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray, }; use arrow::datatypes::{ ArrowPrimitiveType, DataType, Field, Fields, Int64Type, TimeUnit, UnionFields, UnionMode, @@ -1621,6 +1630,67 @@ mod tests { #[test] fn test_array_shredding_as_fixed_size_list() { + let input = build_variant_array(vec![ + VariantRow::List(vec![VariantValue::from(1i64), VariantValue::from(2i64)]), + VariantRow::Value(VariantValue::from("This should not be shredded")), + VariantRow::List(vec![VariantValue::from(3i64), VariantValue::from(4i64)]), + ]); + + let list_schema = + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 2); + let result = shred_variant(&input, &list_schema).unwrap(); + assert_eq!(result.len(), 3); + + // The first row should be shredded, so the `value` field should be null and the + // `typed_value` field should contain the list + assert!(result.is_valid(0)); + assert!(result.value_field().unwrap().is_null(0)); + assert!(result.typed_value_field().unwrap().is_valid(0)); + + // The second row should not be shredded because the provided schema for shredding did not + // match. Hence, the `value` field should contain the raw value and the `typed_value` field + // should be null. + assert!(result.is_valid(1)); + assert!(result.value_field().unwrap().is_valid(1)); + assert!(result.typed_value_field().unwrap().is_null(1)); + + // The third row should be shredded, so the `value` field should be null and the + // `typed_value` field should contain the list + assert!(result.is_valid(2)); + assert!(result.value_field().unwrap().is_null(2)); + assert!(result.typed_value_field().unwrap().is_valid(2)); + + let typed_value = result.typed_value_field().unwrap(); + let fixed_size_list = typed_value + .as_any() + .downcast_ref::() + .expect("Expected FixedSizeListArray"); + + // Verify that typed value is `FixedSizeList`. + assert_eq!(fixed_size_list.len(), 3); + assert_eq!(fixed_size_list.value_length(), 2); + + // Verify that the first entry in the `FixedSizeList` contains the expected value. + let val0 = fixed_size_list.value(0); + let val0_struct = val0.as_any().downcast_ref::().unwrap(); + let val0_typed = val0_struct.column_by_name("typed_value").unwrap(); + let val0_ints = val0_typed.as_any().downcast_ref::().unwrap(); + assert_eq!(val0_ints.values(), &[1i64, 2i64]); + + // Verify that second entry in the `FixedSizeList` cannot be shredded hence the value is + // invalid. + assert!(fixed_size_list.is_null(1)); + + // Verify that the third entry in the `FixedSizeList` contains the expected value. + let val2 = fixed_size_list.value(2); + let val2_struct = val2.as_any().downcast_ref::().unwrap(); + let val2_typed = val2_struct.column_by_name("typed_value").unwrap(); + let val2_ints = val2_typed.as_any().downcast_ref::().unwrap(); + assert_eq!(val2_ints.values(), &[3i64, 4i64]); + } + + #[test] + fn test_array_shredding_as_fixed_size_list_wrong_size() { let input = build_variant_array(vec![VariantRow::List(vec![ VariantValue::from(1i64), VariantValue::from(2i64), @@ -1628,10 +1698,12 @@ mod tests { ])]); let list_schema = DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 2); + let err = shred_variant(&input, &list_schema).unwrap_err(); - assert_eq!( - err.to_string(), - "Not yet implemented: Converting unshredded variant arrays to arrow fixed-size lists" + assert!( + err.to_string() + .contains("Expected fixed size list of size 2, got size 3"), + "got: {err}", ); } diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index e95c774f2926..f76dc4e44600 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -361,10 +361,11 @@ mod test { use arrow::array::{ Array, ArrayRef, AsArray, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, - Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, - LargeBinaryArray, LargeListArray, LargeListViewArray, LargeStringArray, ListArray, - ListViewArray, NullArray, NullBuilder, StringArray, StringViewArray, StructArray, - Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray, + FixedSizeListArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, + Int64Array, LargeBinaryArray, LargeListArray, LargeListViewArray, LargeStringArray, + ListArray, ListViewArray, NullArray, NullBuilder, StringArray, StringViewArray, + StructArray, Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, + Time64NanosecondArray, }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::compute::CastOptions; @@ -4138,13 +4139,29 @@ mod test { ( DataType::LargeListView(field.clone()), Arc::new(LargeListViewArray::new( - field, + field.clone(), ScalarBuffer::from(vec![0, 3]), ScalarBuffer::from(vec![3, 0]), element_array, Some(NullBuffer::from(vec![true, false])), )) as ArrayRef, ), + ( + DataType::FixedSizeList(field.clone(), 3), + Arc::new(FixedSizeListArray::new( + field, + 3, + Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(3), + None, + None, + None, + ])), + Some(NullBuffer::from(vec![true, false])), + )) as ArrayRef, + ), ]; for (request_type, expected) in expectations { @@ -4307,7 +4324,8 @@ mod test { DataType::List(item_field.clone()), DataType::LargeList(item_field.clone()), DataType::ListView(item_field.clone()), - DataType::LargeListView(item_field), + DataType::LargeListView(item_field.clone()), + DataType::FixedSizeList(item_field, 2), ]; for data_type in data_types { @@ -4324,28 +4342,47 @@ mod test { } #[test] - fn test_variant_get_fixed_size_list_not_implemented() { - let string_array: ArrayRef = Arc::new(StringArray::from(vec!["[1, 2]", "\"not a list\""])); + fn test_variant_get_fixed_size_list_wrong_size() { + let string_array: ArrayRef = Arc::new(StringArray::from(vec!["[1, 2, 3]"])); let variant_array = ArrayRef::from(json_to_variant(&string_array).unwrap()); let item_field = Arc::new(Field::new("item", Int64, true)); - for safe in [true, false] { - let options = GetOptions::new() - .with_as_type(Some(FieldRef::from(Field::new( - "result", - DataType::FixedSizeList(item_field.clone(), 2), - true, - )))) - .with_cast_options(CastOptions { - safe, - ..Default::default() - }); - let err = variant_get(&variant_array, options).unwrap_err(); - assert!( - err.to_string() - .contains("Converting unshredded variant arrays to arrow fixed-size lists") - ); - } + // With `safe` set to true, size mismatch should return Null. + let options = GetOptions::new() + .with_as_type(Some(FieldRef::from(Field::new( + "result", + DataType::FixedSizeList(item_field.clone(), 2), + true, + )))) + .with_cast_options(CastOptions { + safe: true, + ..Default::default() + }); + let result = variant_get(&variant_array, options).unwrap(); + let fixed_size_list = result + .as_any() + .downcast_ref::() + .expect("Expected FixedSizeListArray"); + assert_eq!(fixed_size_list.len(), 1); + assert!(fixed_size_list.is_null(0)); + + // With `safe` set to false, error should be raised on wrong sized fixed list. + let options = GetOptions::new() + .with_as_type(Some(FieldRef::from(Field::new( + "result", + DataType::FixedSizeList(item_field.clone(), 2), + true, + )))) + .with_cast_options(CastOptions { + safe: false, + ..Default::default() + }); + let err = variant_get(&variant_array, options).unwrap_err(); + assert!( + err.to_string() + .contains("Expected fixed size list of size 2, got size 3"), + "got: {err}", + ); } macro_rules! perfectly_shredded_preserves_top_level_nulls_test { diff --git a/parquet-variant-compute/src/variant_to_arrow.rs b/parquet-variant-compute/src/variant_to_arrow.rs index 6d1626640c10..ee6f1049ed65 100644 --- a/parquet-variant-compute/src/variant_to_arrow.rs +++ b/parquet-variant-compute/src/variant_to_arrow.rs @@ -27,9 +27,10 @@ use crate::variant_array::ShreddedVariantFieldArray; use crate::{VariantArray, VariantValueArrayBuilder}; use arrow::array::{ ArrayRef, ArrowNativeTypeOp, BinaryBuilder, BinaryLikeArrayBuilder, BinaryViewBuilder, - BooleanBuilder, FixedSizeBinaryBuilder, GenericListArray, GenericListViewArray, - LargeBinaryBuilder, LargeStringBuilder, NullArray, NullBufferBuilder, OffsetSizeTrait, - PrimitiveBuilder, StringBuilder, StringLikeArrayBuilder, StringViewBuilder, StructArray, + BooleanBuilder, FixedSizeBinaryBuilder, FixedSizeListArray, GenericListArray, + GenericListViewArray, LargeBinaryBuilder, LargeStringBuilder, NullArray, NullBufferBuilder, + OffsetSizeTrait, PrimitiveBuilder, StringBuilder, StringLikeArrayBuilder, StringViewBuilder, + StructArray, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::compute::{CastOptions, DecimalCast}; @@ -506,6 +507,7 @@ pub(crate) enum ArrayVariantToArrowRowBuilder<'a> { LargeList(VariantToListArrowRowBuilder<'a, i64, false>), ListView(VariantToListArrowRowBuilder<'a, i32, true>), LargeListView(VariantToListArrowRowBuilder<'a, i64, true>), + FixedSizeList(VariantToFixedSizeListArrowRowBuilder<'a>), } pub(crate) struct StructVariantToArrowRowBuilder<'a> { @@ -618,10 +620,15 @@ impl<'a> ArrayVariantToArrowRowBuilder<'a> { DataType::LargeList(field) => make_list_builder!(LargeList, i64, false, field), DataType::ListView(field) => make_list_builder!(ListView, i32, true, field), DataType::LargeListView(field) => make_list_builder!(LargeListView, i64, true, field), - DataType::FixedSizeList(..) => { - return Err(ArrowError::NotYetImplemented( - "Converting unshredded variant arrays to arrow fixed-size lists".to_string(), - )); + DataType::FixedSizeList(field, size) => { + FixedSizeList(VariantToFixedSizeListArrowRowBuilder::try_new( + field.clone(), + field.data_type(), + *size, + cast_options, + capacity, + shredded, + )?) } other => { return Err(ArrowError::InvalidArgumentError(format!( @@ -638,6 +645,7 @@ impl<'a> ArrayVariantToArrowRowBuilder<'a> { Self::LargeList(builder) => builder.append_null(), Self::ListView(builder) => builder.append_null(), Self::LargeListView(builder) => builder.append_null(), + Self::FixedSizeList(builder) => builder.append_null(), } } @@ -647,6 +655,7 @@ impl<'a> ArrayVariantToArrowRowBuilder<'a> { Self::LargeList(builder) => builder.append_value(value), Self::ListView(builder) => builder.append_value(value), Self::LargeListView(builder) => builder.append_value(value), + Self::FixedSizeList(builder) => builder.append_value(value), } } @@ -656,6 +665,7 @@ impl<'a> ArrayVariantToArrowRowBuilder<'a> { Self::LargeList(builder) => builder.finish(), Self::ListView(builder) => builder.finish(), Self::LargeListView(builder) => builder.finish(), + Self::FixedSizeList(builder) => builder.finish(), } } } @@ -910,6 +920,13 @@ enum ListElementBuilder<'a> { } impl<'a> ListElementBuilder<'a> { + fn append_null(&mut self) -> Result<()> { + match self { + Self::Typed(b) => b.append_null(), + Self::Shredded(b) => b.append_null(), + } + } + fn append_value(&mut self, value: Variant<'_, '_>) -> Result { match self { Self::Typed(b) => b.append_value(value), @@ -1049,6 +1066,99 @@ where } } +pub(crate) struct VariantToFixedSizeListArrowRowBuilder<'a> { + field: FieldRef, + list_size: i32, + element_builder: ListElementBuilder<'a>, + nulls: NullBufferBuilder, + cast_options: &'a CastOptions<'a>, + shredded: bool, +} + +impl<'a> VariantToFixedSizeListArrowRowBuilder<'a> { + fn try_new( + field: FieldRef, + element_data_type: &'a DataType, + list_size: i32, + cast_options: &'a CastOptions, + capacity: usize, + shredded: bool, + ) -> Result { + let element_builder = if shredded { + let builder = make_variant_to_shredded_variant_arrow_row_builder( + element_data_type, + cast_options, + capacity, + NullValue::ArrayElement, + )?; + ListElementBuilder::Shredded(Box::new(builder)) + } else { + let builder = + make_typed_variant_to_arrow_row_builder(element_data_type, cast_options, capacity)?; + ListElementBuilder::Typed(Box::new(builder)) + }; + Ok(Self { + field, + list_size, + element_builder, + nulls: NullBufferBuilder::new(capacity), + cast_options, + shredded, + }) + } + + fn append_null(&mut self) -> Result<()> { + for _ in 0..self.list_size { + self.element_builder.append_null()?; + } + self.nulls.append_null(); + Ok(()) + } + + fn append_value(&mut self, value: &Variant<'_, '_>) -> Result { + match variant_cast_with_options(value, self.cast_options, Variant::as_list) { + Ok(Some(list)) => { + let len = list.len(); + if len != self.list_size as usize { + if self.cast_options.safe && !self.shredded { + self.append_null()?; + return Ok(false); + } + return Err(ArrowError::CastError(format!( + "Expected fixed size list of size {}, got size {}", + self.list_size, len + ))); + } + for element in list.iter() { + self.element_builder.append_value(element)?; + } + self.nulls.append_non_null(); + Ok(true) + } + Ok(None) => { + self.append_null()?; + Ok(false) + } + Err(_) => Err(ArrowError::CastError(format!( + "Failed to extract list from variant {value:?}" + ))), + } + } + + fn finish(mut self) -> Result { + let element_array: ArrayRef = self.element_builder.finish()?; + let field = Arc::new( + self.field + .as_ref() + .clone() + .with_data_type(element_array.data_type().clone()), + ); + let fixed_size_list_array = + FixedSizeListArray::try_new(field, self.list_size, element_array, self.nulls.finish())?; + Ok(Arc::new(fixed_size_list_array)) + } +} + /// Builder for creating VariantArray output (for path extraction without type conversion) pub(crate) struct VariantToBinaryVariantArrowRowBuilder { metadata: ArrayRef, From 4b80f0e1587b003aa01082fc3f8b15873800f219 Mon Sep 17 00:00:00 2001 From: Hippolyte Barraud Date: Wed, 20 May 2026 17:08:20 -0400 Subject: [PATCH 10/51] feat(parquet): generalize value encoder inputs (#9955) # Which issue does this PR close? - Spawn off from #9653 - Contributes to #9731 # Rationale for this change See #9731 # What changes are included in this PR? Changes `byte_array` encoder methods (`FallbackEncoder::encode`, `DictEncoder::encode`, etc) and all `get_*_array_slice` functions from `&[usize]` to `impl ExactSizeIterator`. # Are these changes tested? All tests passing. # Are there any user-facing changes? None. Signed-off-by: Hippolyte Barraud --- parquet/src/arrow/arrow_writer/byte_array.rs | 30 ++++++----- parquet/src/arrow/arrow_writer/mod.rs | 57 ++++++++++---------- 2 files changed, 48 insertions(+), 39 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index f56f9570adfb..9cb0718b4d84 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -165,7 +165,7 @@ impl FallbackEncoder { } /// Encode `values` to the in-progress page - fn encode(&mut self, values: T, indices: &[usize]) + fn encode(&mut self, values: T, indices: impl ExactSizeIterator) where T: ArrayAccessor + Copy, T::Item: AsRef<[u8]>, @@ -174,7 +174,7 @@ impl FallbackEncoder { match &mut self.encoder { FallbackEncoderImpl::Plain { buffer } => { for idx in indices { - let value = values.value(*idx); + let value = values.value(idx); let value = value.as_ref(); buffer.extend_from_slice((value.len() as u32).as_bytes()); buffer.extend_from_slice(value); @@ -183,7 +183,7 @@ impl FallbackEncoder { } FallbackEncoderImpl::DeltaLength { buffer, lengths } => { for idx in indices { - let value = values.value(*idx); + let value = values.value(idx); let value = value.as_ref(); lengths.put(&[value.len() as i32]).unwrap(); buffer.extend_from_slice(value); @@ -197,7 +197,7 @@ impl FallbackEncoder { suffix_lengths, } => { for idx in indices { - let value = values.value(*idx); + let value = values.value(idx); let value = value.as_ref(); let mut prefix_length = 0; @@ -343,7 +343,7 @@ struct DictEncoder { impl DictEncoder { /// Encode `values` to the in-progress page - fn encode(&mut self, values: T, indices: &[usize]) + fn encode(&mut self, values: T, indices: impl ExactSizeIterator) where T: ArrayAccessor + Copy, T::Item: AsRef<[u8]>, @@ -351,7 +351,7 @@ impl DictEncoder { self.indices.reserve(indices.len()); for idx in indices { - let value = values.value(*idx); + let value = values.value(idx); let interned = self.interner.intern(value.as_ref()); self.indices.push(interned); self.variable_length_bytes += value.as_ref().len() as i64; @@ -471,7 +471,13 @@ impl ColumnValueEncoder for ByteArrayEncoder { } fn write_gather(&mut self, values: &Self::Values, indices: &[usize]) -> Result<()> { - downcast_op!(values.data_type(), values, encode, indices, self); + downcast_op!( + values.data_type(), + values, + encode, + indices.iter().copied(), + self + ); Ok(()) } @@ -554,15 +560,16 @@ impl ColumnValueEncoder for ByteArrayEncoder { /// Encodes the provided `values` and `indices` to `encoder` /// /// This is a free function so it can be used with `downcast_op!` -fn encode(values: T, indices: &[usize], encoder: &mut ByteArrayEncoder) +fn encode(values: T, indices: I, encoder: &mut ByteArrayEncoder) where T: ArrayAccessor + Copy, T::Item: Copy + Ord + AsRef<[u8]>, + I: ExactSizeIterator + Clone, { if encoder.statistics_enabled != EnabledStatistics::None { if let Some(accumulator) = encoder.geo_stats_accumulator.as_mut() { - update_geo_stats_accumulator(accumulator.as_mut(), values, indices.iter().cloned()); - } else if let Some((min, max)) = compute_min_max(values, indices.iter().cloned()) { + update_geo_stats_accumulator(accumulator.as_mut(), values, indices.clone()); + } else if let Some((min, max)) = compute_min_max(values, indices.clone()) { if encoder.min_value.as_ref().is_none_or(|m| m > &min) { encoder.min_value = Some(min); } @@ -575,8 +582,7 @@ where // encode the values into bloom filter if enabled if let Some(bloom_filter) = &mut encoder.bloom_filter { - let valid = indices.iter().cloned(); - for idx in valid { + for idx in indices.clone() { bloom_filter.insert(values.value(idx).as_ref()); } } diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 04a798391429..79542caed9b7 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -1387,7 +1387,7 @@ fn write_leaf( } ColumnWriter::BoolColumnWriter(typed) => { let array = column.as_boolean(); - let values = get_bool_array_slice(array, indices); + let values = get_bool_array_slice(array, indices.iter().copied()); typed.write_batch_internal( values.as_slice(), None, @@ -1504,11 +1504,11 @@ fn write_leaf( ArrowDataType::Interval(interval_unit) => match interval_unit { IntervalUnit::YearMonth => { let array = column.as_primitive::(); - get_interval_ym_array_slice(array, indices) + get_interval_ym_array_slice(array, indices.iter().copied()) } IntervalUnit::DayTime => { let array = column.as_primitive::(); - get_interval_dt_array_slice(array, indices) + get_interval_dt_array_slice(array, indices.iter().copied()) } _ => { return Err(ParquetError::NYI(format!( @@ -1518,27 +1518,27 @@ fn write_leaf( }, ArrowDataType::FixedSizeBinary(_) => { let array = column.as_fixed_size_binary(); - get_fsb_array_slice(array, indices) + get_fsb_array_slice(array, indices.iter().copied()) } ArrowDataType::Decimal32(_, _) => { let array = column.as_primitive::(); - get_decimal_32_array_slice(array, indices) + get_decimal_32_array_slice(array, indices.iter().copied()) } ArrowDataType::Decimal64(_, _) => { let array = column.as_primitive::(); - get_decimal_64_array_slice(array, indices) + get_decimal_64_array_slice(array, indices.iter().copied()) } ArrowDataType::Decimal128(_, _) => { let array = column.as_primitive::(); - get_decimal_128_array_slice(array, indices) + get_decimal_128_array_slice(array, indices.iter().copied()) } ArrowDataType::Decimal256(_, _) => { let array = column.as_primitive::(); - get_decimal_256_array_slice(array, indices) + get_decimal_256_array_slice(array, indices.iter().copied()) } ArrowDataType::Float16 => { let array = column.as_primitive::(); - get_float_16_array_slice(array, indices) + get_float_16_array_slice(array, indices.iter().copied()) } _ => { return Err(ParquetError::NYI( @@ -1575,10 +1575,13 @@ fn write_primitive( ) } -fn get_bool_array_slice(array: &arrow_array::BooleanArray, indices: &[usize]) -> Vec { +fn get_bool_array_slice( + array: &arrow_array::BooleanArray, + indices: impl ExactSizeIterator, +) -> Vec { let mut values = Vec::with_capacity(indices.len()); for i in indices { - values.push(array.value(*i)) + values.push(array.value(i)) } values } @@ -1587,11 +1590,11 @@ fn get_bool_array_slice(array: &arrow_array::BooleanArray, indices: &[usize]) -> /// An Arrow YearMonth interval only stores months, thus only the first 4 bytes are populated. fn get_interval_ym_array_slice( array: &arrow_array::IntervalYearMonthArray, - indices: &[usize], + indices: impl ExactSizeIterator, ) -> Vec { let mut values = Vec::with_capacity(indices.len()); for i in indices { - let mut value = array.value(*i).to_le_bytes().to_vec(); + let mut value = array.value(i).to_le_bytes().to_vec(); let mut suffix = vec![0; 8]; value.append(&mut suffix); values.push(FixedLenByteArray::from(ByteArray::from(value))) @@ -1603,12 +1606,12 @@ fn get_interval_ym_array_slice( /// An Arrow DayTime interval only stores days and millis, thus the first 4 bytes are not populated. fn get_interval_dt_array_slice( array: &arrow_array::IntervalDayTimeArray, - indices: &[usize], + indices: impl ExactSizeIterator, ) -> Vec { let mut values = Vec::with_capacity(indices.len()); for i in indices { let mut out = [0; 12]; - let value = array.value(*i); + let value = array.value(i); out[4..8].copy_from_slice(&value.days.to_le_bytes()); out[8..12].copy_from_slice(&value.milliseconds.to_le_bytes()); values.push(FixedLenByteArray::from(ByteArray::from(out.to_vec()))); @@ -1618,12 +1621,12 @@ fn get_interval_dt_array_slice( fn get_decimal_32_array_slice( array: &arrow_array::Decimal32Array, - indices: &[usize], + indices: impl ExactSizeIterator, ) -> Vec { let mut values = Vec::with_capacity(indices.len()); let size = decimal_length_from_precision(array.precision()); for i in indices { - let as_be_bytes = array.value(*i).to_be_bytes(); + let as_be_bytes = array.value(i).to_be_bytes(); let resized_value = as_be_bytes[(4 - size)..].to_vec(); values.push(FixedLenByteArray::from(ByteArray::from(resized_value))); } @@ -1632,12 +1635,12 @@ fn get_decimal_32_array_slice( fn get_decimal_64_array_slice( array: &arrow_array::Decimal64Array, - indices: &[usize], + indices: impl ExactSizeIterator, ) -> Vec { let mut values = Vec::with_capacity(indices.len()); let size = decimal_length_from_precision(array.precision()); for i in indices { - let as_be_bytes = array.value(*i).to_be_bytes(); + let as_be_bytes = array.value(i).to_be_bytes(); let resized_value = as_be_bytes[(8 - size)..].to_vec(); values.push(FixedLenByteArray::from(ByteArray::from(resized_value))); } @@ -1646,12 +1649,12 @@ fn get_decimal_64_array_slice( fn get_decimal_128_array_slice( array: &arrow_array::Decimal128Array, - indices: &[usize], + indices: impl ExactSizeIterator, ) -> Vec { let mut values = Vec::with_capacity(indices.len()); let size = decimal_length_from_precision(array.precision()); for i in indices { - let as_be_bytes = array.value(*i).to_be_bytes(); + let as_be_bytes = array.value(i).to_be_bytes(); let resized_value = as_be_bytes[(16 - size)..].to_vec(); values.push(FixedLenByteArray::from(ByteArray::from(resized_value))); } @@ -1660,12 +1663,12 @@ fn get_decimal_128_array_slice( fn get_decimal_256_array_slice( array: &arrow_array::Decimal256Array, - indices: &[usize], + indices: impl ExactSizeIterator, ) -> Vec { let mut values = Vec::with_capacity(indices.len()); let size = decimal_length_from_precision(array.precision()); for i in indices { - let as_be_bytes = array.value(*i).to_be_bytes(); + let as_be_bytes = array.value(i).to_be_bytes(); let resized_value = as_be_bytes[(32 - size)..].to_vec(); values.push(FixedLenByteArray::from(ByteArray::from(resized_value))); } @@ -1674,11 +1677,11 @@ fn get_decimal_256_array_slice( fn get_float_16_array_slice( array: &arrow_array::Float16Array, - indices: &[usize], + indices: impl ExactSizeIterator, ) -> Vec { let mut values = Vec::with_capacity(indices.len()); for i in indices { - let value = array.value(*i).to_le_bytes().to_vec(); + let value = array.value(i).to_le_bytes().to_vec(); values.push(FixedLenByteArray::from(ByteArray::from(value))); } values @@ -1686,11 +1689,11 @@ fn get_float_16_array_slice( fn get_fsb_array_slice( array: &arrow_array::FixedSizeBinaryArray, - indices: &[usize], + indices: impl ExactSizeIterator, ) -> Vec { let mut values = Vec::with_capacity(indices.len()); for i in indices { - let value = array.value(*i).to_vec(); + let value = array.value(i).to_vec(); values.push(FixedLenByteArray::from(ByteArray::from(value))) } values From f7907211873fef5f80ae22b1c3779dd24041e940 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Thu, 21 May 2026 18:03:54 +0200 Subject: [PATCH 11/51] fix(parquet): validate INT96 column metadata statistics (#10003) # Which issue does this PR close? Closes #10002. # Rationale for this change Malformed Parquet footer metadata can contain INT96 statistics whose encoded min or max value is longer than 12 bytes. The footer metadata conversion path checked that INT96 statistics were at least 12 bytes, but then asserted they were exactly 12 bytes. That allowed malformed input to panic instead of returning an error. The page-statistics path already returns an error for non-12-byte INT96 statistics, so this change makes the footer metadata path behave consistently. # What changes are included in this PR? This PR replaces the INT96 min/max length assertions in footer metadata statistics conversion with explicit `ParquetError` returns. It also adds a regression test covering overlong INT96 min and max values in column metadata statistics. # Are these changes tested? Yes. I ran: - `cargo fmt --all` - `cargo +stable fmt --all -- --check` - `cargo fmt -p parquet -- --check --config skip_children=true $(find ./parquet -name "*.rs" ! -name format.rs)` - `cargo test -p parquet --lib file::metadata::thrift::tests::test_convert_stats_returns_error_for_overlong_int96_statistics` - `cargo test -p parquet --lib file::metadata::thrift::tests` - `cargo test -p parquet` - `cargo check -p parquet --all-targets` - `cargo clippy -p parquet --all-targets --all-features -- -D warnings` # Are there any user-facing changes? Malformed INT96 column metadata statistics now return an error instead of panicking. --- parquet/src/file/metadata/thrift/mod.rs | 48 +++++++++++++++++++++++-- 1 file changed, 46 insertions(+), 2 deletions(-) diff --git a/parquet/src/file/metadata/thrift/mod.rs b/parquet/src/file/metadata/thrift/mod.rs index 9be697c0fe43..d5a0112a5e1a 100644 --- a/parquet/src/file/metadata/thrift/mod.rs +++ b/parquet/src/file/metadata/thrift/mod.rs @@ -274,13 +274,17 @@ fn convert_stats( Type::INT96 => { // INT96 statistics may not be correct, because comparison is signed let min = if let Some(data) = min { - assert_eq!(data.len(), 12); + if data.len() != 12 { + return Err(general_err!("Incorrect Int96 min statistics")); + } Some(Int96::try_from_le_slice(data)?) } else { None }; let max = if let Some(data) = max { - assert_eq!(data.len(), 12); + if data.len() != 12 { + return Err(general_err!("Incorrect Int96 max statistics")); + } Some(Int96::try_from_le_slice(data)?) } else { None @@ -1916,6 +1920,46 @@ pub(crate) mod tests { assert_eq!(decoded_zero.null_count_opt(), Some(0)); } + #[test] + fn test_convert_stats_returns_error_for_overlong_int96_statistics() { + let primitive = + crate::schema::types::Type::primitive_type_builder("col", PhysicalType::INT96) + .build() + .unwrap(); + let column_descr = Arc::new(ColumnDescriptor::new( + Arc::new(primitive), + 0, + 0, + ColumnPath::new(vec![]), + )); + let invalid = (0..13).collect::>(); + + let make_stats = |min, max| super::Statistics { + max, + min, + null_count: Some(0), + distinct_count: None, + max_value: None, + min_value: None, + is_max_value_exact: None, + is_min_value_exact: None, + }; + + let err = super::convert_stats(&column_descr, Some(make_stats(Some(&invalid), None))) + .unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Incorrect Int96 min statistics" + ); + + let err = super::convert_stats(&column_descr, Some(make_stats(None, Some(&invalid)))) + .unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Incorrect Int96 max statistics" + ); + } + #[test] fn malformed_bool_field_returns_error_not_panic() { let page_header = PageHeader { From c46f419b4f87319637869021d46823abeb33ce1f Mon Sep 17 00:00:00 2001 From: mwish Date: Fri, 22 May 2026 02:17:11 +0800 Subject: [PATCH 12/51] fix(cast): Trying to fix cast losting schema problem (#10005) # Which issue does this PR close? - Closes #10004 . # Rationale for this change Previously, just `data_type` is considered. Now the field is taking into account. # What changes are included in this PR? Previously, just `data_type` is considered. Now the field is taking into account. # Are these changes tested? Yes # Are there any user-facing changes? Maybe cast would be a bit more strict --- arrow-cast/src/cast/list.rs | 2 +- arrow-cast/src/cast/mod.rs | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/arrow-cast/src/cast/list.rs b/arrow-cast/src/cast/list.rs index 5d7209ee1111..591afbea5ca4 100644 --- a/arrow-cast/src/cast/list.rs +++ b/arrow-cast/src/cast/list.rs @@ -102,7 +102,7 @@ fn cast_fixed_size_list_to_list_inner = + HashMap::from([("PARQUET:field_id".to_string(), "89".to_string())]); + + let src = Arc::new( + FixedSizeListArray::from_iter_primitive::( + [[1.0_f32, 2.0].map(Some), [3.0, 4.0].map(Some)].map(Some), + 2, + ), + ) as ArrayRef; + + let target_field = Arc::new( + Field::new("element", DataType::Float32, true).with_metadata(metadata.clone()), + ); + + let target_types = [ + DataType::List(target_field.clone()), + DataType::LargeList(target_field.clone()), + DataType::ListView(target_field.clone()), + DataType::LargeListView(target_field.clone()), + ]; + + for target_type in &target_types { + let result = cast(&src, target_type).unwrap(); + assert_eq!( + result.data_type(), + target_type, + "Cast to {target_type:?} should preserve field metadata" + ); + } + } + #[test] fn test_cast_utf8_to_list() { // DataType::List From 2f923f72989ab9df0cb02c749891c5ab3093f743 Mon Sep 17 00:00:00 2001 From: Swanand Mulay <73115739+swanandx@users.noreply.github.com> Date: Fri, 22 May 2026 12:28:39 +0530 Subject: [PATCH 13/51] fix(arrow-cast): support full Date32 range when parsing extended-year dates (#9961) `Date32Type::parse` previously used `chrono::NaiveDate`, which caps at roughly +-262,143 years and rejected valid ISO 8601 extended-year inputs like `+2739877-01-03` As Gregorian repeats in 400-year era (146,097 days), we find the current era and then calculate & validate the date in current era. We recover the absolute day count by adding era * 146,097. Claude code's help was taken to come up with this. # Which issue does this PR close? - Closes #9960 # Rationale for this change Supporting full range of date's allows other dependents like delta-rs to parse data written/managed by other engines like Spark which support full Date32 changing `parse_date()` signature is also other option but would need changes with Date64 as well. # What changes are included in this PR? calculating number of days without converting it full extended year to NaiveDate. And tests for it. # Are these changes tested? added tests and relying on existing tests for verification # Are there any user-facing changes? Maybe as we parse some data successfully which would have previously been None / error Signed-off-by: Swanand Mulay <73115739+swanandx@users.noreply.github.com> --- arrow-cast/src/parse.rs | 92 +++++++++++++++++++++++++++++++++-------- 1 file changed, 75 insertions(+), 17 deletions(-) diff --git a/arrow-cast/src/parse.rs b/arrow-cast/src/parse.rs index a23f421c34db..e7c6f90e75a1 100644 --- a/arrow-cast/src/parse.rs +++ b/arrow-cast/src/parse.rs @@ -585,6 +585,32 @@ const EPOCH_DAYS_FROM_CE: i32 = 719_163; /// Error message if nanosecond conversion request beyond supported interval const ERR_NANOSECONDS_NOT_SUPPORTED: &str = "The dates that can be represented as nanoseconds have to be between 1677-09-21T00:12:44.0 and 2262-04-11T23:47:16.854775804"; +/// Parse the ISO 8601 signed extended-year form (`±YYYY[Y...]-MM-DD`) into +/// raw `(year, month, day)` components, without validating the calendar date. +/// +/// The caller must have already verified that `string` begins with `+` or `-`; +/// the year must have at least 4 digits. Returns `None` if the shape is +/// malformed or any component fails to parse numerically. +fn parse_extended_ymd(string: &str) -> Option<(i32, u32, u32)> { + debug_assert!(string.starts_with('+') || string.starts_with('-')); + // Skip the sign and look for the hyphen that terminates the year digits. + // Per ISO 8601 the unsigned year part must be at least 4 digits. + let rest = &string[1..]; + let hyphen = rest.find('-')?; + if hyphen < 4 { + return None; + } + // The year substring is the sign and the digits (but not the separator), + // e.g. for "+10999-12-31", hyphen is 5 and s[..6] is "+10999". + let year: i32 = string[..hyphen + 1].parse().ok()?; + // The remainder should begin with a '-' which we strip off, leaving the month-day part. + let remainder = string[hyphen + 1..].strip_prefix('-')?; + let mut parts = remainder.splitn(2, '-'); + let month: u32 = parts.next()?.parse().ok()?; + let day: u32 = parts.next()?.parse().ok()?; + Some((year, month, day)) +} + fn parse_date(string: &str) -> Option { // If the date has an extended (signed) year such as "+10999-12-31" or "-0012-05-06" // @@ -594,21 +620,7 @@ fn parse_date(string: &str) -> Option { // // [ISO 8601]: https://docs.oracle.com/en/java/javase/17/docs/api/java.base/java/time/format/DateTimeFormatter.html#ISO_LOCAL_DATE if string.starts_with('+') || string.starts_with('-') { - // Skip the sign and look for the hyphen that terminates the year digits. - // According to ISO 8601 the unsigned part must be at least 4 digits. - let rest = &string[1..]; - let hyphen = rest.find('-')?; - if hyphen < 4 { - return None; - } - // The year substring is the sign and the digits (but not the separator) - // e.g. for "+10999-12-31", hyphen is 5 and s[..6] is "+10999" - let year: i32 = string[..hyphen + 1].parse().ok()?; - // The remainder should begin with a '-' which we strip off, leaving the month-day part. - let remainder = string[hyphen + 1..].strip_prefix('-')?; - let mut parts = remainder.splitn(2, '-'); - let month: u32 = parts.next()?.parse().ok()?; - let day: u32 = parts.next()?.parse().ok()?; + let (year, month, day) = parse_extended_ymd(string)?; return NaiveDate::from_ymd_opt(year, month, day); } @@ -679,10 +691,30 @@ fn parse_date(string: &str) -> Option { NaiveDate::from_ymd_opt(year as _, month as _, day as _) } +/// Parse a date string into days since 1970-01-01, covering the full +/// `Date32` range (years ≈ ±5,881,580) for the signed extended-year form. +/// +/// The Gregorian calendar repeats exactly every 400 years (146,097 days), so +/// we fold the year into `[0, 400)`, validate the folded date, and add +/// `era * 146_097` to recover the absolute day count. +/// +/// For all other inputs, behavior matches [`parse_date`]. +fn parse_date_to_days(string: &str) -> Option { + if string.starts_with('+') || string.starts_with('-') { + let (year, month, day) = parse_extended_ymd(string)?; + let y = year as i64; + let era = y.div_euclid(400); + let yoe = y.rem_euclid(400) as i32; + let nd = NaiveDate::from_ymd_opt(yoe, month, day)?; + let in_era = (nd.num_days_from_ce() - EPOCH_DAYS_FROM_CE) as i64; + return i32::try_from(era * 146_097 + in_era).ok(); + } + parse_date(string).map(|nd| nd.num_days_from_ce() - EPOCH_DAYS_FROM_CE) +} + impl Parser for Date32Type { fn parse(string: &str) -> Option { - let date = parse_date(string)?; - Some(date.num_days_from_ce() - EPOCH_DAYS_FROM_CE) + parse_date_to_days(string) } fn parse_formatted(string: &str, format: &str) -> Option { @@ -1797,6 +1829,32 @@ mod tests { } } + #[test] + fn parse_date32_extended_year() { + // `Date32` covers any i32 days-from-epoch, verify we can parse it + let cases: &[(&str, i32)] = &[ + ("+1970-01-01", 0), + ("+2024-01-01", 19_723), + ("-0001-01-01", -719_893), + ("+29349-01-26", 10_000_000), + ("+2739877-01-03", 1_000_000_000), + // Extremes of the Date32 representable range. + ("+5881580-07-11", i32::MAX), + ("-5877641-06-23", i32::MIN), + ]; + for (input, expected) in cases { + assert_eq!(Date32Type::parse(input), Some(*expected), "input: {input}"); + } + + // One past Date32::MAX / MIN overflows i32 days-from-epoch. + assert_eq!(Date32Type::parse("+5881580-07-12"), None); + assert_eq!(Date32Type::parse("-5877641-06-22"), None); + // Invalid calendar dates still rejected regardless of year magnitude. + assert_eq!(Date32Type::parse("+2739877-02-30"), None); + assert_eq!(Date32Type::parse("+2739877-13-01"), None); + assert_eq!(Date32Type::parse("-2739877-02-30"), None); + } + #[test] fn parse_time64_nanos() { assert_eq!( From edfb9aba45e3fba6fdc52d525c8d4b4132a0d857 Mon Sep 17 00:00:00 2001 From: Ed Seidl Date: Fri, 22 May 2026 08:33:47 -0700 Subject: [PATCH 14/51] Use Thrift macro to generate Parquet `LogicalType` serialization code (#9997) # Which issue does this PR close? - Closes #9995. # Rationale for this change See issue. Improve code maintainability by using thrift macro to generate `LogicalType` serialization code. # What changes are included in this PR? Adds a new macro to generate code for a Thrift `union` that needs to be forward compatible. Does this by adding a catchall `_Unknown` variant for unknown field ids. # Are there any user-facing changes? Yes this is a breaking API change because the `LogicalType` enum will now use tuple variants rather than struct. This also makes public some structs that were previously private. --------- Co-authored-by: Andrew Lamb --- parquet/src/arrow/schema/extension.rs | 15 +- parquet/src/arrow/schema/primitive.rs | 58 ++-- parquet/src/basic.rs | 415 ++++++-------------------- parquet/src/column/writer/mod.rs | 9 +- parquet/src/parquet_macros.rs | 80 +++++ parquet/src/schema/printer.rs | 39 +-- parquet/src/schema/types.rs | 39 ++- parquet/tests/geospatial.rs | 23 +- 8 files changed, 259 insertions(+), 419 deletions(-) diff --git a/parquet/src/arrow/schema/extension.rs b/parquet/src/arrow/schema/extension.rs index 0244c1b6bb99..353770ddbdbb 100644 --- a/parquet/src/arrow/schema/extension.rs +++ b/parquet/src/arrow/schema/extension.rs @@ -48,7 +48,7 @@ pub(crate) fn try_add_extension_type( }; Ok(match parquet_logical_type { #[cfg(feature = "variant_experimental")] - LogicalType::Variant { .. } => { + LogicalType::Variant(_) => { let mut arrow_field = arrow_field; arrow_field.try_with_extension_type(parquet_variant_compute::VariantType)?; arrow_field @@ -66,16 +66,19 @@ pub(crate) fn try_add_extension_type( arrow_field } #[cfg(feature = "geospatial")] - LogicalType::Geometry { crs } => { - let md = parquet_geospatial::WkbMetadata::new(crs.as_deref(), None); + LogicalType::Geometry(geometry) => { + let md = parquet_geospatial::WkbMetadata::new(geometry.crs.as_deref(), None); let mut arrow_field = arrow_field; arrow_field.try_with_extension_type(parquet_geospatial::WkbType::new(Some(md)))?; arrow_field } #[cfg(feature = "geospatial")] - LogicalType::Geography { crs, algorithm } => { - let algorithm = algorithm.map(|a| a.try_as_edges()).transpose()?; - let md = parquet_geospatial::WkbMetadata::new(crs.as_deref(), algorithm); + LogicalType::Geography(geography) => { + let algorithm = geography + .algorithm() + .map(|a| a.try_as_edges()) + .transpose()?; + let md = parquet_geospatial::WkbMetadata::new(geography.crs.as_deref(), algorithm); let mut arrow_field = arrow_field; arrow_field.try_with_extension_type(parquet_geospatial::WkbType::new(Some(md)))?; arrow_field diff --git a/parquet/src/arrow/schema/primitive.rs b/parquet/src/arrow/schema/primitive.rs index b440753cc83b..2272014a9361 100644 --- a/parquet/src/arrow/schema/primitive.rs +++ b/parquet/src/arrow/schema/primitive.rs @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. -use crate::basic::{ConvertedType, LogicalType, TimeUnit as ParquetTimeUnit, Type as PhysicalType}; +use crate::basic::{ + ConvertedType, IntType, LogicalType, TimeUnit as ParquetTimeUnit, Type as PhysicalType, +}; use crate::errors::{ParquetError, Result}; use crate::schema::types::{BasicTypeInfo, Type}; use arrow_schema::{DECIMAL128_MAX_PRECISION, DataType, IntervalUnit, TimeUnit}; @@ -171,15 +173,7 @@ fn decimal_256_type(scale: i32, precision: i32) -> Result { fn from_int32(info: &BasicTypeInfo, scale: i32, precision: i32) -> Result { match (info.logical_type_ref(), info.converted_type()) { (None, ConvertedType::NONE) => Ok(DataType::Int32), - ( - Some( - ref t @ LogicalType::Integer { - bit_width, - is_signed, - }, - ), - _, - ) => match (bit_width, is_signed) { + (Some(ref t @ LogicalType::Integer(int)), _) => match (int.bit_width, int.is_signed) { (8, true) => Ok(DataType::Int8), (16, true) => Ok(DataType::Int16), (32, true) => Ok(DataType::Int32), @@ -188,15 +182,15 @@ fn from_int32(info: &BasicTypeInfo, scale: i32, precision: i32) -> Result Ok(DataType::UInt32), _ => Err(arrow_err!("Cannot create INT32 physical type from {:?}", t)), }, - (Some(LogicalType::Decimal { scale, precision }), _) => { - decimal_128_type(*scale, *precision) + (Some(LogicalType::Decimal(decimal)), _) => { + decimal_128_type(decimal.scale, decimal.precision) } (Some(LogicalType::Date), _) => Ok(DataType::Date32), - (Some(LogicalType::Time { unit, .. }), _) => match unit { + (Some(LogicalType::Time(time)), _) => match time.unit { ParquetTimeUnit::MILLIS => Ok(DataType::Time32(TimeUnit::Millisecond)), _ => Err(arrow_err!( "Cannot create INT32 physical type from {:?}", - unit + time.unit )), }, (None, ConvertedType::UINT_8) => Ok(DataType::UInt8), @@ -220,35 +214,29 @@ fn from_int64(info: &BasicTypeInfo, scale: i32, precision: i32) -> Result Ok(DataType::Int64), ( - Some(LogicalType::Integer { + Some(LogicalType::Integer(IntType { bit_width: 64, is_signed, - }), + })), _, ) => match is_signed { true => Ok(DataType::Int64), false => Ok(DataType::UInt64), }, - (Some(LogicalType::Time { unit, .. }), _) => match unit { + (Some(LogicalType::Time(time)), _) => match time.unit { ParquetTimeUnit::MILLIS => { Err(arrow_err!("Cannot create INT64 from MILLIS time unit",)) } ParquetTimeUnit::MICROS => Ok(DataType::Time64(TimeUnit::Microsecond)), ParquetTimeUnit::NANOS => Ok(DataType::Time64(TimeUnit::Nanosecond)), }, - ( - Some(LogicalType::Timestamp { - is_adjusted_to_u_t_c, - unit, - }), - _, - ) => Ok(DataType::Timestamp( - match unit { + (Some(LogicalType::Timestamp(timestamp)), _) => Ok(DataType::Timestamp( + match timestamp.unit { ParquetTimeUnit::MILLIS => TimeUnit::Millisecond, ParquetTimeUnit::MICROS => TimeUnit::Microsecond, ParquetTimeUnit::NANOS => TimeUnit::Nanosecond, }, - if *is_adjusted_to_u_t_c { + if timestamp.is_adjusted_to_u_t_c { Some("UTC".into()) } else { None @@ -265,9 +253,7 @@ fn from_int64(info: &BasicTypeInfo, scale: i32, precision: i32) -> Result { - decimal_128_type(*scale, *precision) - } + (Some(LogicalType::Decimal(dec)), _) => decimal_128_type(dec.scale, dec.precision), (None, ConvertedType::DECIMAL) => decimal_128_type(scale, precision), (logical, converted) => Err(arrow_err!( "Unable to convert parquet INT64 logical type {:?} or converted type {}", @@ -291,13 +277,7 @@ fn from_byte_array(info: &BasicTypeInfo, precision: i32, scale: i32) -> Result Ok(DataType::Binary), (None, ConvertedType::ENUM) => Ok(DataType::Binary), (None, ConvertedType::UTF8) => Ok(DataType::Utf8), - ( - Some(LogicalType::Decimal { - scale: s, - precision: p, - }), - _, - ) => decimal_type(*s, *p), + (Some(LogicalType::Decimal(decimal)), _) => decimal_type(decimal.scale, decimal.precision), (None, ConvertedType::DECIMAL) => decimal_type(scale, precision), (logical, converted) => Err(arrow_err!( "Unable to convert parquet BYTE_ARRAY logical type {:?} or converted type {}", @@ -315,11 +295,11 @@ fn from_fixed_len_byte_array( ) -> Result { // TODO: This should check the type length for the decimal and interval types match (info.logical_type_ref(), info.converted_type()) { - (Some(LogicalType::Decimal { scale, precision }), _) => { + (Some(LogicalType::Decimal(decimal)), _) => { if type_length <= 16 { - decimal_128_type(*scale, *precision) + decimal_128_type(decimal.scale, decimal.precision) } else { - decimal_256_type(*scale, *precision) + decimal_256_type(decimal.scale, decimal.precision) } } (None, ConvertedType::DECIMAL) => { diff --git a/parquet/src/basic.rs b/parquet/src/basic.rs index 796779358cd2..b4f18f311723 100644 --- a/parquet/src/basic.rs +++ b/parquet/src/basic.rs @@ -30,7 +30,10 @@ use crate::parquet_thrift::{ ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, WriteThrift, WriteThriftField, validate_list_type, }; -use crate::{thrift_enum, thrift_struct, thrift_union_all_empty, write_thrift_field}; +use crate::{ + thrift_enum, thrift_struct, thrift_union_all_empty, thrift_union_with_unknown, + write_thrift_field, +}; use crate::errors::{ParquetError, Result}; @@ -183,383 +186,166 @@ union TimeUnit { // ---------------------------------------------------------------------- // Mirrors thrift union `LogicalType` -// private structs for decoding logical type - thrift_struct!( -struct DecimalType { +pub struct DecimalType { + /// The number of digits in the decimal. 1: required i32 scale + /// The location of the decimal point. 2: required i32 precision } ); thrift_struct!( -struct TimestampType { +pub struct TimestampType { + /// Whether the timestamp is adjusted to UTC. 1: required bool is_adjusted_to_u_t_c + /// The unit of time. 2: required TimeUnit unit } ); -// they are identical -use TimestampType as TimeType; +/// Identical to [`TimestampType`] +pub use TimestampType as TimeType; thrift_struct!( -struct IntType { +pub struct IntType { + /// The number of bits in the integer. 1: required i8 bit_width + /// Whether the integer is signed. 2: required bool is_signed } ); thrift_struct!( -struct VariantType { - // The version of the variant specification that the variant was - // written with. +pub struct VariantType { + /// The version of the variant specification that the variant was + /// written with. 1: optional i8 specification_version } ); thrift_struct!( -struct GeometryType<'a> { - 1: optional string<'a> crs; +pub struct GeometryType { + /// A custom CRS. If unset the CRS `OGC:CRS84` should be used, which means that the geometries + /// must be stored in longitude, latitude based on the WGS84 datum. + 1: optional string crs; } ); thrift_struct!( -struct GeographyType<'a> { - 1: optional string<'a> crs; +pub struct GeographyType { + /// A custom CRS. If unset the CRS `OGC:CRS84` should be used. + 1: optional string crs; + /// An optional algorithm can be set to correctly interpret edges interpolation + /// of the geometries. If unset, the `SPHERICAL` algorithm should be used. 2: optional EdgeInterpolationAlgorithm algorithm; } ); -// TODO(ets): should we switch to tuple variants so we can use -// the thrift macros? +impl GeographyType { + /// Accessor for the `GeographyType::algorithm` field. If this field is not set, this + /// function returns the default value (currently [`EdgeInterpolationAlgorithm::SPHERICAL`] + /// per the Parquet [specification]). + /// + /// [specification]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#geography + pub fn algorithm(&self) -> Option { + self.algorithm.or(Some(Default::default())) + } +} +thrift_union_with_unknown!( /// Logical types used by version 2.4.0+ of the Parquet format. /// /// This is an *entirely new* struct as of version /// 4.0.0. The struct previously named `LogicalType` was renamed to /// [`ConvertedType`]. Please see the README.md for more details. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum LogicalType { - /// A UTF8 encoded string. - String, - /// A map of key-value pairs. - Map, - /// A list of elements. - List, - /// A set of predefined values. - Enum, - /// A decimal value with a specified scale and precision. - Decimal { - /// The number of digits in the decimal. - scale: i32, - /// The location of the decimal point. - precision: i32, - }, - /// A date stored as days since Unix epoch. - Date, - /// A time stored as [`TimeUnit`] since midnight. - Time { - /// Whether the time is adjusted to UTC. - is_adjusted_to_u_t_c: bool, - /// The unit of time. - unit: TimeUnit, - }, - /// A timestamp stored as [`TimeUnit`] since Unix epoch. - Timestamp { - /// Whether the timestamp is adjusted to UTC. - is_adjusted_to_u_t_c: bool, - /// The unit of time. - unit: TimeUnit, - }, - /// An integer with a specified bit width and signedness. - Integer { - /// The number of bits in the integer. - bit_width: i8, - /// Whether the integer is signed. - is_signed: bool, - }, - /// An unknown logical type. - Unknown, - /// A JSON document. - Json, - /// A BSON document. - Bson, - /// A UUID. - Uuid, - /// A 16-bit floating point number. - Float16, - /// A Variant value. - Variant { - /// The version of the variant specification that the variant was written with. - specification_version: Option, - }, - /// A geospatial feature in the Well-Known Binary (WKB) format with linear/planar edges interpolation. - Geometry { - /// A custom CRS. If unset the defaults to `OGC:CRS84`, which means that the geometries - /// must be stored in longitude, latitude based on the WGS84 datum. - crs: Option, - }, - /// A geospatial feature in the WKB format with an explicit (non-linear/non-planar) edges interpolation. - Geography { - /// A custom CRS. If unset the defaults to `OGC:CRS84`. - crs: Option, - /// An optional algorithm can be set to correctly interpret edges interpolation - /// of the geometries. If unset, the algorithm defaults to `SPHERICAL`. - algorithm: Option, - }, - /// For forward compatibility; used when an unknown union value is encountered. - _Unknown { - /// The field id encountered when parsing the unknown logical type. - field_id: i16, - }, +union LogicalType { + /// A UTF8 encoded string. + 1: String + /// A map of key-value pairs. + 2: Map + /// A list of elements. + 3: List + /// A set of predefined values. + 4: Enum + /// A decimal value with a specified scale and precision. + 5: (DecimalType) Decimal + /// A date stored as days since Unix epoch. + 6: Date + /// A time stored as [`TimeUnit`] since midnight. + 7: (TimeType) Time + /// A timestamp stored as [`TimeUnit`] since Unix epoch. + 8: (TimestampType) Timestamp + // 9: reserved for INTERVAL + /// An integer with a specified bit width and signedness. + 10: (IntType) Integer + /// An unknown logical type. + 11: Unknown + /// A JSON document. + 12: Json + /// A BSON document. + 13: Bson + /// A UUID. + 14: Uuid + /// A 16-bit floating point number. + 15: Float16 + /// A Variant value. + 16: (VariantType) Variant + /// A geospatial feature in the Well-Known Binary (WKB) format with linear/planar edges interpolation. + 17: (GeometryType) Geometry + /// A geospatial feature in the WKB format with an explicit (non-linear/non-planar) edges interpolation. + 18: (GeographyType) Geography } +); impl LogicalType { /// Create a [`LogicalType::Integer`] variant with the given `bit_width` and `is_signed` pub fn integer(bit_width: i8, is_signed: bool) -> Self { - Self::Integer { + Self::Integer(IntType { bit_width, is_signed, - } + }) } /// Create a [`LogicalType::Decimal`] variant with the given `scale` and `precision` pub fn decimal(scale: i32, precision: i32) -> Self { - Self::Decimal { scale, precision } + Self::Decimal(DecimalType { scale, precision }) } /// Create a [`LogicalType::Time`] variant with the given `is_adjusted_to_u_t_c` and `unit` pub fn time(is_adjusted_to_u_t_c: bool, unit: TimeUnit) -> Self { - Self::Time { + Self::Time(TimeType { is_adjusted_to_u_t_c, unit, - } + }) } /// Create a [`LogicalType::Timestamp`] variant with the given `is_adjusted_to_u_t_c` and `unit` pub fn timestamp(is_adjusted_to_u_t_c: bool, unit: TimeUnit) -> Self { - Self::Timestamp { + Self::Timestamp(TimestampType { is_adjusted_to_u_t_c, unit, - } + }) } /// Create a [`LogicalType::Variant`] variant with the given `specification_version` pub fn variant(specification_version: Option) -> Self { - Self::Variant { + Self::Variant(VariantType { specification_version, - } + }) } /// Create a [`LogicalType::Geometry`] variant with the given `crs` pub fn geometry(crs: Option) -> Self { - Self::Geometry { crs } + Self::Geometry(GeometryType { crs }) } /// Create a [`LogicalType::Geography`] variant with the given `crs` and `algorithm` pub fn geography(crs: Option, algorithm: Option) -> Self { - Self::Geography { crs, algorithm } - } -} - -impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for LogicalType { - fn read_thrift(prot: &mut R) -> Result { - let field_ident = prot.read_field_begin(0)?; - if field_ident.field_type == FieldType::Stop { - return Err(general_err!("received empty union from remote LogicalType")); - } - let ret = match field_ident.id { - 1 => { - prot.skip_empty_struct()?; - Self::String - } - 2 => { - prot.skip_empty_struct()?; - Self::Map - } - 3 => { - prot.skip_empty_struct()?; - Self::List - } - 4 => { - prot.skip_empty_struct()?; - Self::Enum - } - 5 => { - let val = DecimalType::read_thrift(&mut *prot)?; - Self::decimal(val.scale, val.precision) - } - 6 => { - prot.skip_empty_struct()?; - Self::Date - } - 7 => { - let val = TimeType::read_thrift(&mut *prot)?; - Self::time(val.is_adjusted_to_u_t_c, val.unit) - } - 8 => { - let val = TimestampType::read_thrift(&mut *prot)?; - Self::timestamp(val.is_adjusted_to_u_t_c, val.unit) - } - 10 => { - let val = IntType::read_thrift(&mut *prot)?; - Self::integer(val.bit_width, val.is_signed) - } - 11 => { - prot.skip_empty_struct()?; - Self::Unknown - } - 12 => { - prot.skip_empty_struct()?; - Self::Json - } - 13 => { - prot.skip_empty_struct()?; - Self::Bson - } - 14 => { - prot.skip_empty_struct()?; - Self::Uuid - } - 15 => { - prot.skip_empty_struct()?; - Self::Float16 - } - 16 => { - let val = VariantType::read_thrift(&mut *prot)?; - Self::variant(val.specification_version) - } - 17 => { - let val = GeometryType::read_thrift(&mut *prot)?; - Self::geometry(val.crs.map(|s| s.to_owned())) - } - 18 => { - let val = GeographyType::read_thrift(&mut *prot)?; - // unset algorithm means SPHERICAL, per the spec: - // https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#geography - let algorithm = val - .algorithm - .unwrap_or(EdgeInterpolationAlgorithm::SPHERICAL); - Self::geography(val.crs.map(|s| s.to_owned()), Some(algorithm)) - } - _ => { - prot.skip(field_ident.field_type)?; - Self::_Unknown { - field_id: field_ident.id, - } - } - }; - let field_ident = prot.read_field_begin(field_ident.id)?; - if field_ident.field_type != FieldType::Stop { - return Err(general_err!( - "Received multiple fields for union from remote LogicalType" - )); - } - Ok(ret) + Self::Geography(GeographyType { crs, algorithm }) } } -impl WriteThrift for LogicalType { - const ELEMENT_TYPE: ElementType = ElementType::Struct; - - fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { - match self { - Self::String => { - writer.write_empty_struct(1, 0)?; - } - Self::Map => { - writer.write_empty_struct(2, 0)?; - } - Self::List => { - writer.write_empty_struct(3, 0)?; - } - Self::Enum => { - writer.write_empty_struct(4, 0)?; - } - Self::Decimal { scale, precision } => { - DecimalType { - scale: *scale, - precision: *precision, - } - .write_thrift_field(writer, 5, 0)?; - } - Self::Date => { - writer.write_empty_struct(6, 0)?; - } - Self::Time { - is_adjusted_to_u_t_c, - unit, - } => { - TimeType { - is_adjusted_to_u_t_c: *is_adjusted_to_u_t_c, - unit: *unit, - } - .write_thrift_field(writer, 7, 0)?; - } - Self::Timestamp { - is_adjusted_to_u_t_c, - unit, - } => { - TimestampType { - is_adjusted_to_u_t_c: *is_adjusted_to_u_t_c, - unit: *unit, - } - .write_thrift_field(writer, 8, 0)?; - } - Self::Integer { - bit_width, - is_signed, - } => { - IntType { - bit_width: *bit_width, - is_signed: *is_signed, - } - .write_thrift_field(writer, 10, 0)?; - } - Self::Unknown => { - writer.write_empty_struct(11, 0)?; - } - Self::Json => { - writer.write_empty_struct(12, 0)?; - } - Self::Bson => { - writer.write_empty_struct(13, 0)?; - } - Self::Uuid => { - writer.write_empty_struct(14, 0)?; - } - Self::Float16 => { - writer.write_empty_struct(15, 0)?; - } - Self::Variant { - specification_version, - } => { - VariantType { - specification_version: *specification_version, - } - .write_thrift_field(writer, 16, 0)?; - } - Self::Geometry { crs } => { - GeometryType { - crs: crs.as_ref().map(|s| s.as_str()), - } - .write_thrift_field(writer, 17, 0)?; - } - Self::Geography { crs, algorithm } => { - GeographyType { - crs: crs.as_ref().map(|s| s.as_str()), - algorithm: *algorithm, - } - .write_thrift_field(writer, 18, 0)?; - } - _ => return Err(nyi_err!("logical type")), - } - writer.write_struct_end() - } -} - -write_thrift_field!(LogicalType, FieldType::Struct); - // ---------------------------------------------------------------------- // Mirrors thrift enum `FieldRepetitionType` // @@ -1253,21 +1039,21 @@ impl ColumnOrder { LogicalType::String | LogicalType::Enum | LogicalType::Json | LogicalType::Bson => { SortOrder::UNSIGNED } - LogicalType::Integer { is_signed, .. } => match is_signed { + LogicalType::Integer(int) => match int.is_signed { true => SortOrder::SIGNED, false => SortOrder::UNSIGNED, }, LogicalType::Map | LogicalType::List => SortOrder::UNDEFINED, - LogicalType::Decimal { .. } => SortOrder::SIGNED, + LogicalType::Decimal(_) => SortOrder::SIGNED, LogicalType::Date => SortOrder::SIGNED, - LogicalType::Time { .. } => SortOrder::SIGNED, - LogicalType::Timestamp { .. } => SortOrder::SIGNED, + LogicalType::Time(_) => SortOrder::SIGNED, + LogicalType::Timestamp(_) => SortOrder::SIGNED, LogicalType::Unknown => SortOrder::UNDEFINED, LogicalType::Uuid => SortOrder::UNSIGNED, LogicalType::Float16 => SortOrder::SIGNED, - LogicalType::Variant { .. } - | LogicalType::Geometry { .. } - | LogicalType::Geography { .. } + LogicalType::Variant(_) + | LogicalType::Geometry(_) + | LogicalType::Geography(_) | LogicalType::_Unknown { .. } => SortOrder::UNDEFINED, }, // Fall back to converted type @@ -1426,20 +1212,17 @@ impl From> for ConvertedType { LogicalType::Enum => ConvertedType::ENUM, LogicalType::Decimal { .. } => ConvertedType::DECIMAL, LogicalType::Date => ConvertedType::DATE, - LogicalType::Time { unit, .. } => match unit { + LogicalType::Time(time) => match time.unit { TimeUnit::MILLIS => ConvertedType::TIME_MILLIS, TimeUnit::MICROS => ConvertedType::TIME_MICROS, TimeUnit::NANOS => ConvertedType::NONE, }, - LogicalType::Timestamp { unit, .. } => match unit { + LogicalType::Timestamp(time) => match time.unit { TimeUnit::MILLIS => ConvertedType::TIMESTAMP_MILLIS, TimeUnit::MICROS => ConvertedType::TIMESTAMP_MICROS, TimeUnit::NANOS => ConvertedType::NONE, }, - LogicalType::Integer { - bit_width, - is_signed, - } => match (bit_width, is_signed) { + LogicalType::Integer(int_type) => match (int_type.bit_width, int_type.is_signed) { (8, true) => ConvertedType::INT_8, (16, true) => ConvertedType::INT_16, (32, true) => ConvertedType::INT_32, @@ -1456,9 +1239,9 @@ impl From> for ConvertedType { LogicalType::Bson => ConvertedType::BSON, LogicalType::Uuid | LogicalType::Float16 - | LogicalType::Variant { .. } - | LogicalType::Geometry { .. } - | LogicalType::Geography { .. } + | LogicalType::Variant(_) + | LogicalType::Geometry(_) + | LogicalType::Geography(_) | LogicalType::_Unknown { .. } | LogicalType::Unknown => ConvertedType::NONE, }, diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 595eadbc90f2..4e53230bbf89 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -27,7 +27,8 @@ use std::collections::{BTreeSet, VecDeque}; use std::str; use crate::basic::{ - BoundaryOrder, Compression, ConvertedType, Encoding, EncodingMask, LogicalType, PageType, Type, + BoundaryOrder, Compression, ConvertedType, Encoding, EncodingMask, IntType, LogicalType, + PageType, Type, }; use crate::column::page::{CompressedPage, Page, PageWriteSpec, PageWriter}; use crate::column::writer::encoder::{ColumnValueEncoder, ColumnValueEncoderImpl, ColumnValues}; @@ -1522,9 +1523,9 @@ fn update_stat( fn compare_greater(descr: &ColumnDescriptor, a: &T, b: &T) -> bool { match T::PHYSICAL_TYPE { Type::INT32 | Type::INT64 => { - if let Some(LogicalType::Integer { + if let Some(LogicalType::Integer(IntType { is_signed: false, .. - }) = descr.logical_type_ref() + })) = descr.logical_type_ref() { // need to compare unsigned return compare_greater_unsigned_int(a, b); @@ -1541,7 +1542,7 @@ 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(_)) = descr.logical_type_ref() { return compare_greater_byte_array_decimals(a.as_bytes(), b.as_bytes()); } if let ConvertedType::DECIMAL = descr.converted_type() { diff --git a/parquet/src/parquet_macros.rs b/parquet/src/parquet_macros.rs index 8bb2ad23b054..f7ddf57b1446 100644 --- a/parquet/src/parquet_macros.rs +++ b/parquet/src/parquet_macros.rs @@ -260,6 +260,86 @@ macro_rules! thrift_union { } } +/// Macro used to generate Rust enums for Thrift unions where variants are a mix of unit and +/// tuple types. This version allows for unknown variants for forwards compatibility. +/// +/// Use of this macro requires modifying the thrift IDL. For variants with empty structs as their +/// type, delete the typename (i.e. `1: EmptyStruct Var1;` becomes `1: Var1`). For variants with a +/// non-empty type, the typename must be contained within parens (e.g. `1: MyType Var1;` becomes +/// `1: (MyType) Var1;`). +/// +/// This macro allows for specifying lifetime annotations for the resulting `enum` and its fields. +/// +/// When utilizing this macro the Thrift serialization traits and structs need to be in scope. +#[doc(hidden)] +#[macro_export] +#[allow(clippy::crate_in_macro_def)] +macro_rules! thrift_union_with_unknown { + ($(#[$($def_attrs:tt)*])* union $identifier:ident $(< $lt:lifetime >)? { $($(#[$($field_attrs:tt)*])* $field_id:literal : $( ( $field_type:ident $(< $element_type:ident >)? $(< $field_lt:lifetime >)?) )? $field_name:ident $(;)?)* }) => { + $(#[cfg_attr(not(doctest), $($def_attrs)*)])* + #[derive(Clone, Debug, Eq, PartialEq)] + #[allow(non_camel_case_types)] + #[allow(non_snake_case)] + #[allow(missing_docs)] + pub enum $identifier $(<$lt>)? { + $($(#[cfg_attr(not(doctest), $($field_attrs)*)])* $field_name $( ( $crate::__thrift_union_type!{$field_type $($field_lt)? $($element_type)?} ) )?),*, + _Unknown { + /// The field id encountered when parsing the unknown variant. + field_id: i16, + }, + } + + impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for $identifier $(<$lt>)? { + fn read_thrift(prot: &mut R) -> Result { + let field_ident = prot.read_field_begin(0)?; + if field_ident.field_type == FieldType::Stop { + return Err(general_err!("Received empty union from remote {}", stringify!($identifier))); + } + let ret = match field_ident.id { + $($field_id => { + let val = $crate::__thrift_read_variant!(prot, $field_name $($field_type $($element_type)?)?); + val + })* + _ => { + prot.skip(field_ident.field_type)?; + Self::_Unknown { + field_id: field_ident.id, + } + } + }; + let field_ident = prot.read_field_begin(field_ident.id)?; + if field_ident.field_type != FieldType::Stop { + return Err(general_err!( + concat!("Received multiple fields for union from remote {}", stringify!($identifier)) + )); + } + Ok(ret) + } + } + + impl $(<$lt>)? WriteThrift for $identifier $(<$lt>)? { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + match self { + $($crate::__thrift_write_variant_lhs!($field_name $($field_type)?, variant_val) => + $crate::__thrift_write_variant_rhs!($field_id $($field_type)?, writer, variant_val),)* + Self::_Unknown{..} => return Err(general_err!("Trying to write unknown variant")), + }; + writer.write_struct_end() + } + } + + impl $(<$lt>)? WriteThriftField for $identifier $(<$lt>)? { + fn write_thrift_field(&self, writer: &mut ThriftCompactOutputProtocol, field_id: i16, last_field_id: i16) -> Result { + writer.write_field_begin(FieldType::Struct, field_id, last_field_id)?; + self.write_thrift(writer)?; + Ok(field_id) + } + } + } +} + /// Macro used to generate Rust structs from a Thrift `struct` definition. /// /// Note: diff --git a/parquet/src/schema/printer.rs b/parquet/src/schema/printer.rs index dbeddcfc128c..31d00fc23b3b 100644 --- a/parquet/src/schema/printer.rs +++ b/parquet/src/schema/printer.rs @@ -45,7 +45,10 @@ use std::{fmt, io}; -use crate::basic::{ConvertedType, LogicalType, TimeUnit, Type as PhysicalType}; +use crate::basic::{ + ConvertedType, DecimalType, GeographyType, GeometryType, IntType, LogicalType, TimeUnit, + Type as PhysicalType, VariantType, +}; use crate::file::metadata::{ColumnChunkMetaData, FileMetaData, ParquetMetaData, RowGroupMetaData}; use crate::schema::types::Type; @@ -284,30 +287,28 @@ fn print_logical_and_converted( ) -> String { match logical_type { Some(logical_type) => match logical_type { - LogicalType::Integer { + LogicalType::Integer(IntType { bit_width, is_signed, - } => { + }) => { format!("INTEGER({bit_width},{is_signed})") } - LogicalType::Decimal { scale, precision } => { + LogicalType::Decimal(DecimalType { scale, precision }) => { format!("DECIMAL({precision},{scale})") } - LogicalType::Timestamp { - is_adjusted_to_u_t_c, - unit, - } => { + LogicalType::Timestamp(timestamp) => { format!( "TIMESTAMP({},{})", - print_timeunit(unit), - is_adjusted_to_u_t_c + print_timeunit(×tamp.unit), + timestamp.is_adjusted_to_u_t_c ) } - LogicalType::Time { - is_adjusted_to_u_t_c, - unit, - } => { - format!("TIME({},{})", print_timeunit(unit), is_adjusted_to_u_t_c) + LogicalType::Time(time) => { + format!( + "TIME({},{})", + print_timeunit(&time.unit), + time.is_adjusted_to_u_t_c + ) } LogicalType::Date => "DATE".to_string(), LogicalType::Bson => "BSON".to_string(), @@ -318,17 +319,17 @@ fn print_logical_and_converted( LogicalType::List => "LIST".to_string(), LogicalType::Map => "MAP".to_string(), LogicalType::Float16 => "FLOAT16".to_string(), - LogicalType::Variant { + LogicalType::Variant(VariantType { specification_version, - } => format!("VARIANT({specification_version:?})"), - LogicalType::Geometry { crs } => { + }) => format!("VARIANT({specification_version:?})"), + LogicalType::Geometry(GeometryType { crs }) => { if let Some(crs) = crs { format!("GEOMETRY({crs})") } else { "GEOMETRY".to_string() } } - LogicalType::Geography { crs, algorithm } => { + LogicalType::Geography(GeographyType { crs, algorithm }) => { let algorithm = algorithm.unwrap_or_default(); if let Some(crs) = crs { format!("GEOGRAPHY({algorithm}, {crs})") diff --git a/parquet/src/schema/types.rs b/parquet/src/schema/types.rs index d8b3456d4723..1f9b8590fcf6 100644 --- a/parquet/src/schema/types.rs +++ b/parquet/src/schema/types.rs @@ -24,7 +24,8 @@ use crate::file::metadata::HeapSize; use crate::file::metadata::thrift::SchemaElement; use crate::basic::{ - ColumnOrder, ConvertedType, LogicalType, Repetition, SortOrder, TimeUnit, Type as PhysicalType, + ColumnOrder, ConvertedType, IntType, LogicalType, Repetition, SortOrder, TimeType, TimeUnit, + Type as PhysicalType, }; use crate::errors::{ParquetError, Result}; @@ -356,20 +357,20 @@ impl<'a> PrimitiveTypeBuilder<'a> { )); } (LogicalType::Enum, PhysicalType::BYTE_ARRAY) => {} - (LogicalType::Decimal { scale, precision }, _) => { + (LogicalType::Decimal(decimal), _) => { // Check that scale and precision are consistent with legacy values - if *scale != self.scale { + if decimal.scale != self.scale { return Err(general_err!( "DECIMAL logical type scale {} must match self.scale {} for field '{}'", - scale, + decimal.scale, self.scale, self.name )); } - if *precision != self.precision { + if decimal.precision != self.precision { return Err(general_err!( "DECIMAL logical type precision {} must match self.precision {} for field '{}'", - precision, + decimal.precision, self.precision, self.name )); @@ -378,32 +379,30 @@ impl<'a> PrimitiveTypeBuilder<'a> { } (LogicalType::Date, PhysicalType::INT32) => {} ( - LogicalType::Time { + LogicalType::Time(TimeType { unit: TimeUnit::MILLIS, .. - }, + }), PhysicalType::INT32, ) => {} - (LogicalType::Time { unit, .. }, PhysicalType::INT64) => { - if *unit == TimeUnit::MILLIS { + (LogicalType::Time(time), PhysicalType::INT64) => { + if time.unit == TimeUnit::MILLIS { return Err(general_err!( "Cannot use millisecond unit on INT64 type for field '{}'", self.name )); } } - (LogicalType::Timestamp { .. }, PhysicalType::INT64) => {} - (LogicalType::Integer { bit_width, .. }, PhysicalType::INT32) - if *bit_width <= 32 => {} - (LogicalType::Integer { bit_width, .. }, PhysicalType::INT64) - if *bit_width == 64 => {} + (LogicalType::Timestamp(_), PhysicalType::INT64) => {} + (LogicalType::Integer(int), PhysicalType::INT32) if int.bit_width <= 32 => {} + (LogicalType::Integer(int), PhysicalType::INT64) if int.bit_width == 64 => {} // Null type (LogicalType::Unknown, _) => {} (LogicalType::String, PhysicalType::BYTE_ARRAY) => {} (LogicalType::Json, PhysicalType::BYTE_ARRAY) => {} (LogicalType::Bson, PhysicalType::BYTE_ARRAY) => {} - (LogicalType::Geometry { .. }, PhysicalType::BYTE_ARRAY) => {} - (LogicalType::Geography { .. }, PhysicalType::BYTE_ARRAY) => {} + (LogicalType::Geometry(_), PhysicalType::BYTE_ARRAY) => {} + (LogicalType::Geography(_), PhysicalType::BYTE_ARRAY) => {} (LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) if self.length == 16 => {} (LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) => { return Err(general_err!( @@ -1280,8 +1279,8 @@ fn build_tree<'a>( /// Checks if the logical type is valid. fn check_logical_type(logical_type: &Option) -> Result<()> { - if let Some(LogicalType::Integer { bit_width, .. }) = *logical_type { - if bit_width != 8 && bit_width != 16 && bit_width != 32 && bit_width != 64 { + if let Some(LogicalType::Integer(IntType { bit_width, .. })) = logical_type { + if *bit_width != 8 && *bit_width != 16 && *bit_width != 32 && *bit_width != 64 { return Err(general_err!( "Bit width must be 8, 16, 32, or 64 for Integer logical type" )); @@ -1482,7 +1481,7 @@ mod tests { if let Err(e) = result { assert_eq!( format!("{e}"), - "Parquet error: Cannot annotate Integer { bit_width: 8, is_signed: true } from INT64 for field 'foo'" + "Parquet error: Cannot annotate Integer(IntType { bit_width: 8, is_signed: true }) from INT64 for field 'foo'" ); } diff --git a/parquet/tests/geospatial.rs b/parquet/tests/geospatial.rs index fcc93661ed97..bf34528d03e8 100644 --- a/parquet/tests/geospatial.rs +++ b/parquet/tests/geospatial.rs @@ -30,7 +30,7 @@ mod test { ArrowSchemaConverter, ArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder, arrow_writer::ArrowWriterOptions, }, - basic::{EdgeInterpolationAlgorithm, LogicalType}, + basic::LogicalType, column::reader::ColumnReader, data_type::{ByteArray, ByteArrayType}, file::{ @@ -62,29 +62,22 @@ mod test { let expected_metadata = [ ( "crs-default.parquet", - LogicalType::Geometry { crs: None }, + LogicalType::geometry(None), WkbMetadata::new(None, None), ), ( "crs-srid.parquet", - LogicalType::Geometry { - crs: Some("srid:5070".to_string()), - }, + LogicalType::geometry(Some("srid:5070".to_string())), WkbMetadata::new(Some("srid:5070"), None), ), ( "crs-projjson.parquet", - LogicalType::Geometry { - crs: Some("projjson:projjson_epsg_5070".to_string()), - }, + LogicalType::geometry(Some("projjson:projjson_epsg_5070".to_string())), WkbMetadata::new(Some("projjson:projjson_epsg_5070"), None), ), ( "crs-geography.parquet", - LogicalType::Geography { - crs: None, - algorithm: Some(EdgeInterpolationAlgorithm::SPHERICAL), - }, + LogicalType::geography(None, None), WkbMetadata::new(None, Some(WkbEdges::Spherical)), ), ]; @@ -109,8 +102,8 @@ mod test { let column_descr = metadata.file_metadata().schema_descr().column(1); let logical_type = column_descr.logical_type_ref().unwrap(); - if let LogicalType::Geometry { crs } = logical_type { - let crs = crs.as_ref(); + if let LogicalType::Geometry(geometry) = logical_type { + let crs = geometry.crs.as_ref(); let crs_parsed: Value = serde_json::from_str(crs.unwrap()).unwrap(); assert_eq!(crs_parsed.get("id").unwrap().get("code").unwrap(), 5070); } else { @@ -136,7 +129,7 @@ mod test { // optional binary field_id=-1 geometry (Geometry(crs=)); let fields = metadata.file_metadata().schema().get_fields(); let logical_type = fields[2].get_basic_info().logical_type_ref().unwrap(); - assert_eq!(logical_type, &LogicalType::Geometry { crs: None }); + assert_eq!(logical_type, &LogicalType::geometry(None)); let geo_statistics = metadata.row_group(0).column(2).geo_statistics(); assert!(geo_statistics.is_some()); From 8c88bd49b5c4c019b3473b652267b46464d2e58f Mon Sep 17 00:00:00 2001 From: Kevin Choubacha <2043529+choubacha@users.noreply.github.com> Date: Fri, 22 May 2026 11:22:26 -0700 Subject: [PATCH 15/51] Adds is_null function to RowAccessor (#9979) # Which issue does this PR close? - Closes #8076. # Rationale for this change When dealing with parquet files directly, having a built in option for checking for nulls is useful when reading the rows out of the file. This function allows gating calls to the other accessors so that they can be mapped to None instead of relying on errors or accessing the field vec itself. # Are these changes tested? Yes, a test was added. # Are there any user-facing changes? This does add to the Row api. The trait function has a documentation block. --- parquet/src/record/api.rs | 98 ++++++++++++++++++++++++++++++++++++--- 1 file changed, 91 insertions(+), 7 deletions(-) diff --git a/parquet/src/record/api.rs b/parquet/src/record/api.rs index 33a1464fa3d8..104a82e64d9d 100644 --- a/parquet/src/record/api.rs +++ b/parquet/src/record/api.rs @@ -140,6 +140,8 @@ impl<'a> Iterator for RowColumnIter<'a> { /// Trait for type-safe convenient access to fields within a Row. pub trait RowAccessor { + /// Check if the field at the index is null. + fn is_null(&self, i: usize) -> Result; /// Try to get a boolean value at the given index. fn get_bool(&self, i: usize) -> Result; /// Try to get a byte value at the given index. @@ -210,11 +212,12 @@ pub trait RowFormatter { macro_rules! row_primitive_accessor { ($METHOD:ident, $VARIANT:ident, $TY:ty) => { fn $METHOD(&self, i: usize) -> Result<$TY> { - match self.fields[i].1 { - Field::$VARIANT(v) => Ok(v), + match self.fields.get(i) { + Some((_, Field::$VARIANT(v))) => Ok(*v), + None => Err(ParquetError::IndexOutOfBound(i, self.fields.len())), _ => Err(general_err!( "Cannot access {} as {}", - self.fields[i].1.get_type_name(), + self.fields[i].1.get_type_name(), // Safe access as None is stringify!($VARIANT) )), } @@ -227,11 +230,13 @@ macro_rules! row_primitive_accessor { macro_rules! row_complex_accessor { ($METHOD:ident, $VARIANT:ident, $TY:ty) => { fn $METHOD(&self, i: usize) -> Result<&$TY> { - match self.fields[i].1 { - Field::$VARIANT(ref v) => Ok(v), + match self.fields.get(i) { + Some((_, Field::$VARIANT(v))) => Ok(v), + None => Err(ParquetError::IndexOutOfBound(i, self.fields.len())), _ => Err(general_err!( "Cannot access {} as {}", - self.fields[i].1.get_type_name(), + self.fields[i].1.get_type_name(), // Safe access as None is + // just checked. stringify!($VARIANT) )), } @@ -242,11 +247,23 @@ macro_rules! row_complex_accessor { impl RowFormatter for Row { /// Get Display reference for a given field. fn fmt(&self, i: usize) -> &dyn fmt::Display { - &self.fields[i].1 + if let Some((_, v)) = self.fields.get(i) { + v + } else { + &"" + } } } impl RowAccessor for Row { + fn is_null(&self, i: usize) -> Result { + match self.fields.get(i) { + Some((_, Field::Null)) => Ok(true), + None => Err(ParquetError::IndexOutOfBound(i, self.len())), + _ => Ok(false), + } + } + row_primitive_accessor!(get_bool, Bool, bool); row_primitive_accessor!(get_byte, Byte, i8); @@ -1651,6 +1668,8 @@ mod tests { ("p".to_string(), Field::Float16(f16::from_f32(9.1))), ]); + assert!(row.is_null(0).unwrap()); + assert!(!row.is_null(1).unwrap()); assert!(!row.get_bool(1).unwrap()); assert_eq!(3, row.get_byte(2).unwrap()); assert_eq!(4, row.get_short(3).unwrap()); @@ -1666,6 +1685,71 @@ mod tests { assert_eq!(5, row.get_bytes(13).unwrap().len()); assert_eq!(7, row.get_decimal(14).unwrap().precision()); assert!((f16::from_f32(9.1) - row.get_float16(15).unwrap()).abs() < f16::EPSILON); + + assert!(matches!( + row.is_null(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_bool(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_byte(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_short(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_int(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_long(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_ubyte(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_ushort(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_uint(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_ulong(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_float(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_double(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_string(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_bytes(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_decimal(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_float16(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); } #[test] From 5372e8acf70ba6a9455d8b98edddb7368e4c71f5 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 22 May 2026 19:23:31 +0100 Subject: [PATCH 16/51] Fix parquet-variant build on wasm targets (#9978) # Which issue does this PR close? - Closes #9977 # Rationale for this change Seems like a trivial fix to get it building on more targets. # What changes are included in this PR? 1. enables a feature for `uuid` that is required to build on WASM (only when building for WASM). 2. Change the const size assertions to take pointer_width into account # Are these changes tested? I've tested the change locally (for both WASM targets), not sure if its worth it to add to CI that currently only tests the top-level `arrow` crate on WASM # Are there any user-facing changes? None --- parquet-variant-compute/Cargo.toml | 10 ++++++---- parquet-variant-json/Cargo.toml | 4 +++- parquet-variant/Cargo.toml | 9 ++++++--- parquet-variant/src/variant.rs | 3 +++ parquet-variant/src/variant/list.rs | 4 ++++ parquet-variant/src/variant/metadata.rs | 4 ++++ parquet-variant/src/variant/object.rs | 4 ++++ 7 files changed, 30 insertions(+), 8 deletions(-) diff --git a/parquet-variant-compute/Cargo.toml b/parquet-variant-compute/Cargo.toml index 85d66a9cf706..bcfb36b8710c 100644 --- a/parquet-variant-compute/Cargo.toml +++ b/parquet-variant-compute/Cargo.toml @@ -27,18 +27,21 @@ keywords = ["arrow", "parquet", "variant"] edition = { workspace = true } rust-version = { workspace = true } - [dependencies] -arrow = { workspace = true , features = ["canonical_extension_types"]} +arrow = { workspace = true, features = ["canonical_extension_types"] } arrow-schema = { workspace = true } half = { version = "2.1", default-features = false } indexmap = "2.10.0" parquet-variant = { workspace = true } parquet-variant-json = { workspace = true } chrono = { workspace = true } -uuid = { version = "1.18.0", features = ["v4"]} +uuid = { version = "1.18.0", features = ["v4"] } serde_json = "1.0" +# uuid requires the `js` feature to run on wasm +[target.'cfg(target_arch = "wasm32")'.dependencies] +uuid = { version = "1.18.0", features = ["v4", "js"] } + [lib] name = "parquet_variant_compute" bench = false @@ -48,7 +51,6 @@ rand = "0.9.1" criterion = { workspace = true, default-features = false } arrow = { workspace = true, features = ["test_utils"] } - [[bench]] name = "variant_kernels" harness = false diff --git a/parquet-variant-json/Cargo.toml b/parquet-variant-json/Cargo.toml index f9550adc26af..76c639d2ee38 100644 --- a/parquet-variant-json/Cargo.toml +++ b/parquet-variant-json/Cargo.toml @@ -28,7 +28,6 @@ readme = "../parquet-variant/README.md" edition = { workspace = true } rust-version = { workspace = true } - [dependencies] arrow-schema = { workspace = true } parquet-variant = { workspace = true } @@ -37,6 +36,9 @@ serde_json = "1.0" base64 = "0.22" uuid = "1.18.0" +# uuid requires the `js` feature to run on wasm +[target.'cfg(target_arch = "wasm32")'.dependencies] +uuid = { version = "1.18.0", features = ["js"] } [lib] name = "parquet_variant_json" diff --git a/parquet-variant/Cargo.toml b/parquet-variant/Cargo.toml index 7d5064331e4c..ba186432d28f 100644 --- a/parquet-variant/Cargo.toml +++ b/parquet-variant/Cargo.toml @@ -29,15 +29,18 @@ edition = { workspace = true } rust-version = { workspace = true } [dependencies] -arrow = { workspace = true , features = ["canonical_extension_types"] } +arrow = { workspace = true, features = ["canonical_extension_types"] } arrow-schema = { workspace = true } chrono = { workspace = true } half = { version = "2.1", default-features = false } indexmap = "2.10.0" num-traits = { version = "0.2", default-features = false } -uuid = { version = "1.18.0", features = ["v4"]} +uuid = { version = "1.18.0", features = ["v4"] } +simdutf8 = { workspace = true, optional = true } -simdutf8 = { workspace = true , optional = true } +# uuid requires the `js` feature to run on wasm +[target.'cfg(target_arch = "wasm32")'.dependencies] +uuid = { version = "1.18.0", features = ["v4", "js"] } [lib] name = "parquet_variant" diff --git a/parquet-variant/src/variant.rs b/parquet-variant/src/variant.rs index 3b5a04ddfd6c..58a3f7eeb261 100644 --- a/parquet-variant/src/variant.rs +++ b/parquet-variant/src/variant.rs @@ -298,7 +298,10 @@ pub enum Variant<'m, 'v> { } // We don't want this to grow because it could hurt performance of a frequently-created type. +#[cfg(target_pointer_width = "64")] const _: () = crate::utils::expect_size_of::(80); +#[cfg(target_pointer_width = "32")] +const _: () = crate::utils::expect_size_of::(48); enum NumericKind { Integer, diff --git a/parquet-variant/src/variant/list.rs b/parquet-variant/src/variant/list.rs index fd71afba7342..7301d0570645 100644 --- a/parquet-variant/src/variant/list.rs +++ b/parquet-variant/src/variant/list.rs @@ -128,8 +128,12 @@ pub struct VariantList<'m, 'v> { } // We don't want this to grow because it could increase the size of `Variant` and hurt performance. +#[cfg(target_pointer_width = "64")] const _: () = crate::utils::expect_size_of::(64); +#[cfg(target_pointer_width = "32")] +const _: () = crate::utils::expect_size_of::(40); + impl<'m, 'v> VariantList<'m, 'v> { /// Attempts to interpret `value` as a variant array value. /// diff --git a/parquet-variant/src/variant/metadata.rs b/parquet-variant/src/variant/metadata.rs index 9f9688acd090..d5d08d204c83 100644 --- a/parquet-variant/src/variant/metadata.rs +++ b/parquet-variant/src/variant/metadata.rs @@ -140,8 +140,12 @@ pub struct VariantMetadata<'m> { // We don't want this to grow because it increases the size of VariantList and VariantObject, which // could increase the size of Variant. All those size increases could hurt performance. +#[cfg(target_pointer_width = "64")] const _: () = crate::utils::expect_size_of::(32); +#[cfg(target_pointer_width = "32")] +const _: () = crate::utils::expect_size_of::(20); + /// The canonical byte slice corresponding to an empty metadata dictionary. /// /// ``` diff --git a/parquet-variant/src/variant/object.rs b/parquet-variant/src/variant/object.rs index 52dc2ef42f98..bb91584cefa6 100644 --- a/parquet-variant/src/variant/object.rs +++ b/parquet-variant/src/variant/object.rs @@ -127,8 +127,12 @@ pub struct VariantObject<'m, 'v> { } // We don't want this to grow because it could increase the size of `Variant` and hurt performance. +#[cfg(target_pointer_width = "64")] const _: () = crate::utils::expect_size_of::(64); +#[cfg(target_pointer_width = "32")] +const _: () = crate::utils::expect_size_of::(44); + impl<'m, 'v> VariantObject<'m, 'v> { pub fn new(metadata: VariantMetadata<'m>, value: &'v [u8]) -> Self { Self::try_new_with_shallow_validation(metadata, value).expect("Invalid variant object") From df89537ea98dd2b8228fdcb94cfd6a6b92935081 Mon Sep 17 00:00:00 2001 From: Konstantin Tarasov <33369833+sdf-jkl@users.noreply.github.com> Date: Fri, 22 May 2026 14:24:20 -0400 Subject: [PATCH 17/51] [Variant] remove `BorrowedShreddingState` (#9791) # Which issue does this PR close? - Closes #9790. # Rationale for this change Check issue # What changes are included in this PR? - Drop `BorrowedShreddingState` - Replace it with `ShreddingState` - ~~Removed the lifetimes in `unshred_variant` as they required helpers to cover recursive `ShreddingState` handling.~~ - ~~Lifetimes removal introduces clone on `NullBuffer`. Extra 3 usize (24 bytes) per `Array`. Only used in `NullUnshredVariantBuilder`~~ Removed the only place where `NullBuffer` was stored. No regression. # Are these changes tested? Yes, unit tests. # Are there any user-facing changes? No. --- parquet-variant-compute/src/lib.rs | 2 +- .../src/unshred_variant.rs | 58 +++++++------- parquet-variant-compute/src/variant_array.rs | 75 ++----------------- parquet-variant-compute/src/variant_get.rs | 16 ++-- 4 files changed, 41 insertions(+), 110 deletions(-) diff --git a/parquet-variant-compute/src/lib.rs b/parquet-variant-compute/src/lib.rs index b05d0e023653..066fe15b7c8b 100644 --- a/parquet-variant-compute/src/lib.rs +++ b/parquet-variant-compute/src/lib.rs @@ -51,7 +51,7 @@ mod variant_array_builder; mod variant_get; mod variant_to_arrow; -pub use variant_array::{BorrowedShreddingState, ShreddingState, VariantArray, VariantType}; +pub use variant_array::{ShreddingState, VariantArray, VariantType}; pub use variant_array_builder::{VariantArrayBuilder, VariantValueArrayBuilder}; pub use cast_to_variant::{cast_to_variant, cast_to_variant_with_options}; diff --git a/parquet-variant-compute/src/unshred_variant.rs b/parquet-variant-compute/src/unshred_variant.rs index 7bf31e2256df..f4bfa73bded1 100644 --- a/parquet-variant-compute/src/unshred_variant.rs +++ b/parquet-variant-compute/src/unshred_variant.rs @@ -17,15 +17,14 @@ //! Module for unshredding VariantArray by folding typed_value columns back into the value column. -use crate::variant_array::binary_array_value; -use crate::{BorrowedShreddingState, VariantArray, VariantValueArrayBuilder}; +use crate::variant_array::{binary_array_value, validate_binary_array}; +use crate::{VariantArray, VariantValueArrayBuilder}; use arrow::array::{ Array, ArrayRef, AsArray as _, BinaryArray, BinaryViewArray, BooleanArray, FixedSizeBinaryArray, FixedSizeListArray, GenericListArray, GenericListViewArray, LargeBinaryArray, LargeStringArray, ListLikeArray, PrimitiveArray, StringArray, StringViewArray, StructArray, }; -use arrow::buffer::NullBuffer; use arrow::datatypes::{ ArrowPrimitiveType, DataType, Date32Type, Decimal32Type, Decimal64Type, Decimal128Type, DecimalType, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, @@ -67,8 +66,8 @@ pub fn unshred_variant(array: &VariantArray) -> Result { // NOTE: None/None at top-level is technically invalid, but the shredding spec requires us to // emit `Variant::Null` when a required value is missing. let nulls = array.nulls(); - let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.shredding_state().borrow())? - .unwrap_or_else(|| UnshredVariantRowBuilder::null(nulls)); + let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.inner())? + .unwrap_or_else(UnshredVariantRowBuilder::null); let metadata = array.metadata_field(); let mut value_builder = VariantValueArrayBuilder::new(array.len()); @@ -126,13 +125,13 @@ enum UnshredVariantRowBuilder<'a> { FixedSizeList(ListUnshredVariantBuilder<'a, FixedSizeListArray>), Struct(StructUnshredVariantBuilder<'a>), ValueOnly(ValueOnlyUnshredVariantBuilder<'a>), - Null(NullUnshredVariantBuilder<'a>), + Null(NullUnshredVariantBuilder), } impl<'a> UnshredVariantRowBuilder<'a> { /// Creates an all-null row builder. - fn null(nulls: Option<&'a NullBuffer>) -> Self { - Self::Null(NullUnshredVariantBuilder::new(nulls)) + fn null() -> Self { + Self::Null(NullUnshredVariantBuilder) } /// Appends a single row at the given value index to the supplied builder. @@ -175,12 +174,17 @@ impl<'a> UnshredVariantRowBuilder<'a> { } } - /// Creates a new UnshredVariantRowBuilder from shredding state - /// Returns None for None/None case - caller decides how to handle based on context - fn try_new_opt(shredding_state: BorrowedShreddingState<'a>) -> Result> { - let value = shredding_state.value_field(); - let typed_value = shredding_state.typed_value_field(); - let Some(typed_value) = typed_value else { + /// Creates a new UnshredVariantRowBuilder from the `(value, typed_value)` pair of a shredded + /// variant struct. Returns None for the None/None case - caller decides how to handle based on + /// context. + fn try_new_opt(inner_struct: &'a StructArray) -> Result> { + let value = if let Some(value_col) = inner_struct.column_by_name("value") { + validate_binary_array(value_col.as_ref(), "value")?; + Some(value_col) + } else { + None + }; + let Some(typed_value) = inner_struct.column_by_name("typed_value") else { // Copy the value across directly, if present. Else caller decides what to do. return Ok(value.map(|v| Self::ValueOnly(ValueOnlyUnshredVariantBuilder::new(v)))); }; @@ -289,27 +293,17 @@ impl<'a> UnshredVariantRowBuilder<'a> { } } -/// Builder for arrays with neither typed_value nor value (all NULL/Variant::Null) -struct NullUnshredVariantBuilder<'a> { - nulls: Option<&'a NullBuffer>, -} - -impl<'a> NullUnshredVariantBuilder<'a> { - fn new(nulls: Option<&'a NullBuffer>) -> Self { - Self { nulls } - } +/// Builder for arrays with neither typed_value nor value (all Variant::Null) +struct NullUnshredVariantBuilder; +impl NullUnshredVariantBuilder { fn append_row( &mut self, builder: &mut impl VariantBuilderExt, _metadata: &VariantMetadata, - index: usize, + _index: usize, ) -> Result<()> { - if self.nulls.is_some_and(|nulls| nulls.is_null(index)) { - builder.append_null(); - } else { - builder.append_value(Variant::Null); - } + builder.append_value(Variant::Null); Ok(()) } } @@ -590,7 +584,7 @@ impl<'a> StructUnshredVariantBuilder<'a> { field_array.data_type() ))); }; - let field_unshredder = UnshredVariantRowBuilder::try_new_opt(field_array.try_into()?)?; + let field_unshredder = UnshredVariantRowBuilder::try_new_opt(field_array)?; field_unshredders.insert(field.name().as_ref(), field_unshredder); } @@ -670,8 +664,8 @@ impl<'a, L: ListLikeArray> ListUnshredVariantBuilder<'a, L> { // // NOTE: A None/None array element is technically invalid, but the shredding spec // requires us to emit `Variant::Null` when a required value is missing. - let element_unshredder = UnshredVariantRowBuilder::try_new_opt(element_values.try_into()?)? - .unwrap_or_else(|| UnshredVariantRowBuilder::null(None)); + let element_unshredder = UnshredVariantRowBuilder::try_new_opt(element_values)? + .unwrap_or_else(UnshredVariantRowBuilder::null); Ok(Self { value, diff --git a/parquet-variant-compute/src/variant_array.rs b/parquet-variant-compute/src/variant_array.rs index 5f52939e712c..a6ee281002cc 100644 --- a/parquet-variant-compute/src/variant_array.rs +++ b/parquet-variant-compute/src/variant_array.rs @@ -64,7 +64,7 @@ pub(crate) fn variant_from_arrays_at<'m, 'v>( } /// Validates that an array has a binary-like data type. -fn validate_binary_array(array: &dyn Array, field_name: &str) -> Result<()> { +pub(crate) fn validate_binary_array(array: &dyn Array, field_name: &str) -> Result<()> { match array.data_type() { DataType::Binary | DataType::LargeBinary | DataType::BinaryView => Ok(()), _ => Err(ArrowError::InvalidArgumentError(format!( @@ -843,14 +843,6 @@ impl ShreddingState { self.typed_value.as_ref() } - /// Returns a borrowed version of this shredding state - pub fn borrow(&self) -> BorrowedShreddingState<'_> { - BorrowedShreddingState { - value: self.value_field(), - typed_value: self.typed_value_field(), - } - } - /// Slice all the underlying arrays pub fn slice(&self, offset: usize, length: usize) -> Self { Self { @@ -860,74 +852,19 @@ impl ShreddingState { } } -/// Similar to [`ShreddingState`] except it holds borrowed references of the target arrays. Useful -/// for avoiding clone operations when the caller does not need a self-standing shredding state. -#[derive(Clone, Debug)] -pub struct BorrowedShreddingState<'a> { - value: Option<&'a ArrayRef>, - typed_value: Option<&'a ArrayRef>, -} - -impl<'a> BorrowedShreddingState<'a> { - /// Create a new `BorrowedShreddingState` from the given `value` and `typed_value` fields - /// - /// Note you can create a `BorrowedShreddingState` from a &[`StructArray`] using - /// `BorrowedShreddingState::try_from(&struct_array)`, for example: - /// - /// ```no_run - /// # use arrow::array::StructArray; - /// # use parquet_variant_compute::BorrowedShreddingState; - /// # fn get_struct_array() -> StructArray { - /// # unimplemented!() - /// # } - /// let struct_array: StructArray = get_struct_array(); - /// let shredding_state = BorrowedShreddingState::try_from(&struct_array).unwrap(); - /// ``` - pub fn new(value: Option<&'a ArrayRef>, typed_value: Option<&'a ArrayRef>) -> Self { - Self { value, typed_value } - } - - /// Return a reference to the value field, if present - pub fn value_field(&self) -> Option<&'a ArrayRef> { - self.value - } - - /// Return a reference to the typed_value field, if present - pub fn typed_value_field(&self) -> Option<&'a ArrayRef> { - self.typed_value - } -} - -impl<'a> TryFrom<&'a StructArray> for BorrowedShreddingState<'a> { +impl TryFrom<&StructArray> for ShreddingState { type Error = ArrowError; - fn try_from(inner_struct: &'a StructArray) -> Result { + fn try_from(inner_struct: &StructArray) -> Result { // The `value` column need not exist, but if it does it must be a binary type. let value = if let Some(value_col) = inner_struct.column_by_name("value") { validate_binary_array(value_col.as_ref(), "value")?; - Some(value_col) + Some(value_col.clone()) } else { None }; - let typed_value = inner_struct.column_by_name("typed_value"); - Ok(BorrowedShreddingState::new(value, typed_value)) - } -} - -impl TryFrom<&StructArray> for ShreddingState { - type Error = ArrowError; - - fn try_from(inner_struct: &StructArray) -> Result { - Ok(BorrowedShreddingState::try_from(inner_struct)?.into()) - } -} - -impl From> for ShreddingState { - fn from(state: BorrowedShreddingState<'_>) -> Self { - ShreddingState { - value: state.value_field().cloned(), - typed_value: state.typed_value_field().cloned(), - } + let typed_value = inner_struct.column_by_name("typed_value").cloned(); + Ok(ShreddingState::new(value, typed_value)) } } diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index f76dc4e44600..774da0e72e8c 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -25,15 +25,15 @@ use arrow_schema::{ArrowError, DataType, FieldRef}; use parquet_variant::{VariantPath, VariantPathElement}; use crate::VariantArray; -use crate::variant_array::BorrowedShreddingState; +use crate::variant_array::ShreddingState; use crate::variant_to_arrow::make_variant_to_arrow_row_builder; use arrow::array::AsArray; use std::sync::Arc; -pub(crate) enum ShreddedPathStep<'a> { +pub(crate) enum ShreddedPathStep { /// Path step succeeded, return the new shredding state - Success(BorrowedShreddingState<'a>), + Success(ShreddingState), /// The path element is not present in the `typed_value` column and there is no `value` column, /// so we know it does not exist. It, and all paths under it, are all-NULL. Missing, @@ -49,11 +49,11 @@ pub(crate) enum ShreddedPathStep<'a> { /// a missing-path step (`Missing` or `NotShredded` depending on whether `value` exists). /// /// TODO: Support `VariantPathElement::Index`? It wouldn't be easy, and maybe not even possible. -pub(crate) fn follow_shredded_path_element<'a>( - shredding_state: &BorrowedShreddingState<'a>, +pub(crate) fn follow_shredded_path_element( + shredding_state: &ShreddingState, path_element: &VariantPathElement<'_>, _cast_options: &CastOptions, -) -> Result> { +) -> Result { // If the requested path element is not present in `typed_value`, and `value` is missing, then // we know it does not exist; it, and all paths under it, are all-NULL. let missing_path_step = || match shredding_state.value_field() { @@ -90,7 +90,7 @@ pub(crate) fn follow_shredded_path_element<'a>( )) })?; - let state = BorrowedShreddingState::try_from(struct_array)?; + let state = ShreddingState::try_from(struct_array)?; Ok(ShreddedPathStep::Success(state)) } VariantPathElement::Index { .. } => { @@ -154,7 +154,7 @@ fn shredded_get_path( // Peel away the prefix of path elements that traverses the shredded parts of this variant // column. Shredding will traverse the rest of the path on a per-row basis. - let mut shredding_state = input.shredding_state().borrow(); + let mut shredding_state = input.shredding_state().clone(); let mut accumulated_nulls = input.inner().nulls().cloned(); let mut path_index = 0; for path_element in path { From 7b335b7712089c1270ab1a989336908a240cc812 Mon Sep 17 00:00:00 2001 From: BoazC-MSFT Date: Fri, 22 May 2026 21:26:57 +0300 Subject: [PATCH 18/51] fix: prevent panic in record reader when row group metadata overcounts num_rows (#9993) # Which issue does this PR close? - Closes #9992. # Rationale for this change The record reader (`RowIter` / `get_row_iter`) panics with `index out of bounds` when a Parquet file's row group metadata declares more rows than a column chunk actually contains. This happens in production when reading third-party Parquet files with mismatched metadata. Instead of panicking, the reader should return an error. # What changes are included in this PR? Three layers of fix in `parquet/src/record/`: **triplet.rs - fix the inconsistent internal state:** - Reset `curr_triplet_index` to 0 in the exhaustion path of `read_next`, so the stale index from the previous batch never persists alongside empty buffers. - Return 0 from `current_def_level` and `current_rep_level` when `has_next` is false, as defense-in-depth against any caller that skips the `has_next` check. **reader.rs - return errors instead of panicking:** - Add `has_next()` guards before consuming column data in all `read_field` variants: `PrimitiveReader`, `OptionReader`, `RepeatedReader`, and `KeyValueReader`. When a column is exhausted mid-iteration, `read_field` now returns `Err("Unexpected end of column data")` which propagates through `ReaderIter::next` as `Some(Err(...))`. # Are these changes tested? Yes. Five new tests: - `test_current_def_level_safe_after_exhaustion` - drives a `TripletIter` to exhaustion on an optional column and asserts `current_def_level()` returns 0 instead of panicking. - `test_current_rep_level_safe_after_exhaustion` - same for `current_rep_level()` on a repeated column. - `test_reader_iter_returns_error_when_num_records_exceeds_data` - exercises the full `ReaderIter` stack with an optional field (via `nulls.snappy.parquet`). - `test_reader_iter_returns_error_for_repeated_field_when_num_records_exceeds_data` - same for a repeated primitive field (via `repeated_primitive_no_list.parquet`). - `test_reader_iter_returns_error_for_map_field_when_num_records_exceeds_data` - same for a map field projected alone (via `map_no_value.parquet`). Each integration test inflates `num_records` by 1 beyond actual data, asserts all real rows return `Ok`, and asserts the extra iteration returns `Err` containing "Unexpected end of column data". # Are there any user-facing changes? Callers of `get_row_iter` or `RowIter` that previously hit a panic on corrupt/truncated files will now receive an `Err` from the iterator instead. No API signature changes. --- parquet/src/record/reader.rs | 65 +++++++++++++++++++++++++++++++++++ parquet/src/record/triplet.rs | 44 ++++++++++++++++++++++++ 2 files changed, 109 insertions(+) diff --git a/parquet/src/record/reader.rs b/parquet/src/record/reader.rs index a6b8d2d54cd7..bee49c70d81d 100644 --- a/parquet/src/record/reader.rs +++ b/parquet/src/record/reader.rs @@ -437,11 +437,17 @@ impl Reader { fn read_field(&mut self) -> Result { let field = match *self { Reader::PrimitiveReader(_, ref mut column) => { + if !column.has_next() { + return Err(general_err!("Unexpected end of column data")); + } let value = column.current_value()?; column.read_next()?; value } Reader::OptionReader(def_level, ref mut reader) => { + if !reader.has_next() { + return Err(general_err!("Unexpected end of column data")); + } if reader.current_def_level() > def_level { reader.read_field()? } else { @@ -465,6 +471,9 @@ impl Reader { Field::Group(row) } Reader::RepeatedReader(_, def_level, rep_level, ref mut reader) => { + if !reader.has_next() { + return Err(general_err!("Unexpected end of column data")); + } let mut elements = Vec::new(); loop { if reader.current_def_level() > def_level { @@ -488,6 +497,9 @@ impl Reader { Field::ListInternal(make_list(elements)) } Reader::KeyValueReader(_, def_level, rep_level, ref mut keys, ref mut values) => { + if !keys.has_next() { + return Err(general_err!("Unexpected end of column data")); + } let mut pairs = Vec::new(); loop { if keys.current_def_level() > def_level { @@ -1912,4 +1924,57 @@ mod tests { ),]]; assert_eq!(rows, expected_rows); } + + fn assert_err_on_overcount(file_name: &str, proj_schema: Option) { + let file = get_test_file(file_name); + let file_reader = SerializedFileReader::new(file).unwrap(); + let metadata = file_reader.metadata(); + let row_group_reader = file_reader.get_row_group(0).unwrap(); + let actual_rows = row_group_reader.metadata().num_rows() as usize; + + let descr = match proj_schema { + Some(schema) => Arc::new(SchemaDescriptor::new(Arc::new(schema))), + None => metadata.file_metadata().schema_descr_ptr(), + }; + let reader = TreeBuilder::new().build(descr, &*row_group_reader).unwrap(); + let iter = ReaderIter::new(reader, actual_rows + 1).unwrap(); + + let rows: Vec> = iter.collect(); + assert_eq!(rows.len(), actual_rows + 1); + for row in &rows[..actual_rows] { + assert!(row.is_ok(), "Expected Ok row, got: {:?}", row); + } + let err = rows[actual_rows].as_ref().unwrap_err(); + assert!( + err.to_string().contains("Unexpected end of column data"), + "Unexpected error message: {}", + err + ); + } + + #[test] + fn test_reader_iter_returns_error_when_num_records_exceeds_data() { + assert_err_on_overcount("nulls.snappy.parquet", None); + } + + #[test] + fn test_reader_iter_returns_error_for_repeated_field_when_num_records_exceeds_data() { + assert_err_on_overcount("repeated_primitive_no_list.parquet", None); + } + + #[test] + fn test_reader_iter_returns_error_for_map_field_when_num_records_exceeds_data() { + let schema = parse_message_type( + "message schema { + REQUIRED group my_map (MAP) { + REPEATED group key_value { + REQUIRED INT32 key; + OPTIONAL INT32 value; + } + } + }", + ) + .unwrap(); + assert_err_on_overcount("map_no_value.parquet", Some(schema)); + } } diff --git a/parquet/src/record/triplet.rs b/parquet/src/record/triplet.rs index 8244dfb12823..b4d39bbbd9c9 100644 --- a/parquet/src/record/triplet.rs +++ b/parquet/src/record/triplet.rs @@ -263,6 +263,9 @@ impl TypedTripletIter { /// If field is required, then maximum definition level is returned. #[inline] fn current_def_level(&self) -> i16 { + if !self.has_next { + return 0; + } match self.def_levels { Some(ref vec) => vec[self.curr_triplet_index], None => self.max_def_level, @@ -273,6 +276,9 @@ impl TypedTripletIter { /// If field is required, then maximum repetition level is returned. #[inline] fn current_rep_level(&self) -> i16 { + if !self.has_next { + return 0; + } match self.rep_levels { Some(ref vec) => vec[self.curr_triplet_index], None => self.max_rep_level, @@ -315,6 +321,7 @@ impl TypedTripletIter { // No more values or levels to read if records_read == 0 && values_read == 0 && levels_read == 0 { + self.curr_triplet_index = 0; self.has_next = false; return Ok(false); } @@ -561,4 +568,41 @@ mod tests { assert_eq!(def_levels, expected_def_levels); assert_eq!(rep_levels, expected_rep_levels); } + + fn open_triplet_iter(file_name: &str, path: &[&str], batch_size: usize) -> TripletIter { + let column_path = ColumnPath::from(path.iter().map(|x| x.to_string()).collect::>()); + let file = get_test_file(file_name); + let file_reader = SerializedFileReader::new(file).unwrap(); + let metadata = file_reader.metadata(); + let schema = metadata.file_metadata().schema_descr(); + let row_group_reader = file_reader.get_row_group(0).unwrap(); + for i in 0..schema.num_columns() { + let descr = schema.column(i); + if descr.path() == &column_path { + let reader = row_group_reader.get_column_reader(i).unwrap(); + return TripletIter::new(descr.clone(), reader, batch_size); + } + } + panic!("Column {column_path:?} not found in {file_name}"); + } + + #[test] + fn test_current_def_level_safe_after_exhaustion() { + let mut iter = open_triplet_iter("nulls.snappy.parquet", &["b_struct", "b_c_int"], 256); + while let Ok(true) = iter.read_next() {} + assert!(!iter.has_next()); + assert_eq!(iter.current_def_level(), 0); + } + + #[test] + fn test_current_rep_level_safe_after_exhaustion() { + let mut iter = open_triplet_iter( + "nested_lists.snappy.parquet", + &["a", "list", "element", "list", "element", "list", "element"], + 256, + ); + while let Ok(true) = iter.read_next() {} + assert!(!iter.has_next()); + assert_eq!(iter.current_rep_level(), 0); + } } From 8acab7b5371470deff6c211899295d2bb3030dfc Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Sun, 24 May 2026 21:57:45 -0400 Subject: [PATCH 19/51] feat: Implement decimal <-> float16 casts (#10008) # Which issue does this PR close? - Closes #9123. # Rationale for this change Arrow supports casts between decimal and float64/float32; for consistency and completeness, we should also support casts between decimal and float16. In DataFusion, this will be particularly useful: once https://github.com/apache/datafusion/issues/14612 is fixed, `arrow_cast(0.0, 'Float16')` will no longer work, unless we first add support for decimal -> float16 casts in arrow-rs. # What changes are included in this PR? * Add support for decimal -> float16 cast * Add support for float16 -> decimal cast * Add unit tests for new behavior * Update docs/comment on supported casts # Are these changes tested? Yes; new tests added. # Are there any user-facing changes? Yes; new casts are now supported. Otherwise no changes. --- arrow-cast/src/cast/mod.rs | 282 ++++++++++++++++++++++++++++++++++++- 1 file changed, 278 insertions(+), 4 deletions(-) diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs index 6b4355ae3a8e..12442bd9fd30 100644 --- a/arrow-cast/src/cast/mod.rs +++ b/arrow-cast/src/cast/mod.rs @@ -188,7 +188,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool { ) => true, // signed numeric to decimal ( - Int8 | Int16 | Int32 | Int64 | Float32 | Float64, + Int8 | Int16 | Int32 | Int64 | Float16 | Float32 | Float64, Decimal32(_, _) | Decimal64(_, _) | Decimal128(_, _) | Decimal256(_, _), ) => true, // decimal to unsigned numeric @@ -199,7 +199,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool { // decimal to signed numeric ( Decimal32(_, _) | Decimal64(_, _) | Decimal128(_, _) | Decimal256(_, _), - Null | Int8 | Int16 | Int32 | Int64 | Float32 | Float64, + Null | Int8 | Int16 | Int32 | Int64 | Float16 | Float32 | Float64, ) => true, // decimal to string ( @@ -649,9 +649,9 @@ fn timestamp_to_date32( /// * `Time32 and `Time64`: precision lost when going to higher interval /// * `Timestamp` and `Date{32|64}`: precision lost when going to higher interval /// * Temporal to/from backing Primitive: zero-copy with data type change -/// * `Float32/Float64` to `Decimal(precision, scale)` rounds to the `scale` decimals +/// * `Float16/Float32/Float64` to `Decimal(precision, scale)` rounds to the `scale` decimals /// (i.e. casting `6.4999` to `Decimal(10, 1)` becomes `6.5`). -/// * `Decimal` to `Float32/Float64` is lossy and values outside the representable +/// * `Decimal` to `Float16/Float32/Float64` is lossy and values outside the representable /// range become `INFINITY` or `-INFINITY` without error. /// /// Unsupported Casts (check with `can_cast_types` before calling): @@ -2329,6 +2329,13 @@ where Int16 => cast_decimal_to_integer::(array, base, *scale, cast_options), Int32 => cast_decimal_to_integer::(array, base, *scale, cast_options), Int64 => cast_decimal_to_integer::(array, base, *scale, cast_options), + Float16 => cast_decimal_to_float::(array, |x| { + half::f16::from_f64(single_decimal_to_float_lossy::( + &as_float, + x, + >::from(*scale), + )) + }), Float32 => cast_decimal_to_float::(array, |x| { single_decimal_to_float_lossy::(&as_float, x, >::from(*scale)) as f32 @@ -2426,6 +2433,12 @@ where base, cast_options, ), + Float16 => cast_floating_point_to_decimal::<_, D>( + array.as_primitive::(), + *precision, + *scale, + cast_options, + ), Float32 => cast_floating_point_to_decimal::<_, D>( array.as_primitive::(), *precision, @@ -3366,6 +3379,84 @@ mod tests { } } + #[test] + fn test_cast_float16_to_decimals() { + let array = Float16Array::from(vec![ + Some(f16::from_f32(1.25)), + Some(f16::from_f32(-2.5)), + Some(f16::from_f32(1.125)), + Some(f16::from_f32(-1.125)), + Some(f16::from_f32(0.0)), + None, + ]); + + generate_cast_test_case!( + &array, + Decimal32Array, + &DataType::Decimal32(9, 2), + vec![ + Some(125_i32), + Some(-250_i32), + Some(113_i32), + Some(-113_i32), + Some(0_i32), + None + ] + ); + generate_cast_test_case!( + &array, + Decimal64Array, + &DataType::Decimal64(18, 2), + vec![ + Some(125_i64), + Some(-250_i64), + Some(113_i64), + Some(-113_i64), + Some(0_i64), + None + ] + ); + generate_cast_test_case!( + &array, + Decimal128Array, + &DataType::Decimal128(38, 2), + vec![ + Some(125_i128), + Some(-250_i128), + Some(113_i128), + Some(-113_i128), + Some(0_i128), + None + ] + ); + generate_cast_test_case!( + &array, + Decimal256Array, + &DataType::Decimal256(76, 2), + vec![ + Some(i256::from_i128(125_i128)), + Some(i256::from_i128(-250_i128)), + Some(i256::from_i128(113_i128)), + Some(i256::from_i128(-113_i128)), + Some(i256::from_i128(0_i128)), + None + ] + ); + + let array = Float16Array::from(vec![ + Some(f16::from_f32(1250.0)), + Some(f16::from_f32(-1250.0)), + Some(f16::from_f32(1249.0)), + None, + ]); + generate_cast_test_case!( + &array, + Decimal128Array, + &DataType::Decimal128(5, -2), + vec![Some(13_i128), Some(-13_i128), Some(12_i128), None] + ); + } + #[test] fn test_cast_decimal128_to_decimal128_overflow() { let input_type = DataType::Decimal128(38, 3); @@ -3620,6 +3711,19 @@ mod tests { &DataType::Int64, vec![Some(1_i64), Some(2_i64), Some(3_i64), None, Some(5_i64)] ); + // f16 + generate_cast_test_case!( + array, + Float16Array, + &DataType::Float16, + vec![ + Some(f16::from_f32(1.25)), + Some(f16::from_f32(2.25)), + Some(f16::from_f32(3.25)), + None, + Some(f16::from_f32(5.25)) + ] + ); // f32 generate_cast_test_case!( array, @@ -3847,6 +3951,19 @@ mod tests { &DataType::Int64, vec![Some(1_i64), Some(2_i64), Some(3_i64), None, Some(5_i64)] ); + // f16 + generate_cast_test_case!( + &array, + Float16Array, + &DataType::Float16, + vec![ + Some(f16::from_f32(1.25)), + Some(f16::from_f32(2.25)), + Some(f16::from_f32(3.25)), + None, + Some(f16::from_f32(5.25)) + ] + ); // f32 generate_cast_test_case!( &array, @@ -3957,6 +4074,60 @@ mod tests { ); } + #[test] + fn test_cast_decimal128_to_float16_overflow() { + let array = create_decimal128_array( + vec![ + Some(6_550_400_i128), + Some(100_000_000_i128), + Some(-100_000_000_i128), + None, + ], + 10, + 2, + ) + .unwrap(); + + generate_cast_test_case!( + &array, + Float16Array, + &DataType::Float16, + vec![ + Some(f16::from_f64(65504.0)), + Some(f16::INFINITY), + Some(f16::NEG_INFINITY), + None + ] + ); + } + + #[test] + fn test_cast_decimal256_to_float16_overflow() { + let array = create_decimal256_array( + vec![ + Some(i256::from_i128(6_550_400_i128)), + Some(i256::from_i128(100_000_000_i128)), + Some(i256::from_i128(-100_000_000_i128)), + None, + ], + 10, + 2, + ) + .unwrap(); + + generate_cast_test_case!( + &array, + Float16Array, + &DataType::Float16, + vec![ + Some(f16::from_f64(65504.0)), + Some(f16::INFINITY), + Some(f16::NEG_INFINITY), + None + ] + ); + } + #[test] fn test_cast_decimal_to_numeric_negative_scale() { let value_array: Vec> = vec![ @@ -3975,6 +4146,19 @@ mod tests { vec![Some(1_250), Some(2_250), Some(3_250), None, Some(5_250)] ); + let value_array: Vec> = vec![Some(12), Some(-12), None]; + let array = create_decimal128_array(value_array, 10, -2).unwrap(); + generate_cast_test_case!( + &array, + Float16Array, + &DataType::Float16, + vec![ + Some(f16::from_f32(1200.0)), + Some(f16::from_f32(-1200.0)), + None + ] + ); + let value_array: Vec> = vec![Some(125), Some(225), Some(325), None, Some(525)]; let array = create_decimal32_array(value_array, 8, -2).unwrap(); generate_cast_test_case!( @@ -10164,6 +10348,96 @@ mod tests { ); } + #[test] + fn test_cast_float16_to_decimal128_precision_overflow() { + let array = Float16Array::from(vec![f16::from_f32(1.1)]); + let array = Arc::new(array) as ArrayRef; + let casted_array = cast_with_options( + &array, + &DataType::Decimal128(2, 2), + &CastOptions { + safe: true, + format_options: FormatOptions::default(), + }, + ); + assert!(casted_array.is_ok()); + assert!(casted_array.unwrap().is_null(0)); + + let casted_array = cast_with_options( + &array, + &DataType::Decimal128(2, 2), + &CastOptions { + safe: false, + format_options: FormatOptions::default(), + }, + ); + let err = casted_array.unwrap_err().to_string(); + let expected_error = "Invalid argument error: 1.10 is too large to store in a Decimal128 of precision 2. Max is 0.99"; + assert_eq!(err, expected_error); + } + + #[test] + fn test_cast_float16_to_decimal256_precision_overflow() { + let array = Float16Array::from(vec![f16::from_f32(1.1)]); + let array = Arc::new(array) as ArrayRef; + let casted_array = cast_with_options( + &array, + &DataType::Decimal256(2, 2), + &CastOptions { + safe: true, + format_options: FormatOptions::default(), + }, + ); + assert!(casted_array.is_ok()); + assert!(casted_array.unwrap().is_null(0)); + + let casted_array = cast_with_options( + &array, + &DataType::Decimal256(2, 2), + &CastOptions { + safe: false, + format_options: FormatOptions::default(), + }, + ); + let err = casted_array.unwrap_err().to_string(); + let expected_error = "Invalid argument error: 1.10 is too large to store in a Decimal256 of precision 2. Max is 0.99"; + assert_eq!(err, expected_error); + } + + #[test] + fn test_cast_float16_to_decimal128_non_finite() { + let array = Float16Array::from(vec![f16::NAN, f16::INFINITY, f16::NEG_INFINITY]); + let array = Arc::new(array) as ArrayRef; + let casted_array = cast_with_options( + &array, + &DataType::Decimal128(38, 2), + &CastOptions { + safe: true, + format_options: FormatOptions::default(), + }, + ) + .unwrap(); + + assert!(casted_array.is_null(0)); + assert!(casted_array.is_null(1)); + assert!(casted_array.is_null(2)); + + let casted_array = cast_with_options( + &array, + &DataType::Decimal128(38, 2), + &CastOptions { + safe: false, + format_options: FormatOptions::default(), + }, + ); + let err = casted_array.unwrap_err().to_string(); + let expected_error = "Cannot cast to Decimal128(38, 2)"; + assert!( + err.contains(expected_error), + "did not find expected error '{expected_error}' in actual error '{err}'" + ); + } + #[test] fn test_cast_floating_point_to_decimal256_precision_overflow() { let array = Float64Array::from(vec![1.1]); From 1c679ef8a7d0481596a29736a0bdea10f7eb4fed Mon Sep 17 00:00:00 2001 From: Sergei Date: Mon, 25 May 2026 18:00:30 +0700 Subject: [PATCH 20/51] fix(parquet): exclude single-leaf struct roots from predicate cache (#9983) # Which issue does this PR close? - Closes #9982 . # Rationale for this change # What changes are included in this PR? ## Root cause `ProjectionMask::without_nested_types` (`parquet/src/arrow/mod.rs:427`) decides which leaves the predicate cache may cover. The check before this fix was: ```rust if root_leaf_counts[root_idx] == 1 && !root.is_list() { included_leaves.push(leaf_idx); } ``` PR #8866 added `!root.is_list()` to exclude lists, but a **struct** root with a single leaf still satisfies the condition and gets cached. ## Fix (1 line) `parquet/src/arrow/mod.rs:455`: ```diff - if root_leaf_counts[root_idx] == 1 && !root.is_list() { + if root_leaf_counts[root_idx] == 1 && root.is_primitive() { included_leaves.push(leaf_idx); } ``` # Are these changes tested? ## Tests added on the branch ### 1. Reproducer integration test **File:** `parquet/tests/arrow_reader/predicate_cache.rs` **Name:** `test_async_predicate_on_single_leaf_nullable_struct` Builds an in-memory Parquet file with `OPTIONAL group b { REQUIRED BYTE_ARRAY aa (UTF8); }`, writes two rows (parent NULL, parent non-NULL), then runs the same `IS NULL` row filter through the async reader twice: once with the default cache, once with `with_max_predicate_cache_size(0)`. It asserts that - the uncached control yields exactly 1 row (`address` NULL row matches); - the cached run yields the same row count as the uncached one. **Pre-fix:** panic at `struct_array.rs:142`. **Post-fix:** passes (1 row in both cases). ### 2. Unit test **File:** `parquet/src/arrow/mod.rs` (test module) **Name:** `test_projection_mask_without_nested_single_leaf_struct` Directly checks `ProjectionMask::without_nested_types` against a schema with `OPTIONAL group address { REQUIRED BYTE_ARRAY street; } REQUIRED INT32 id`, for three input masks (single nested leaf, mixed, all leaves). All three expected outputs reflect that the struct's leaf is now considered nested. **Pre-fix:** would return `Some([street_leaf])` for the single-leaf-only mask. **Post-fix:** returns `None` for the single-leaf-only mask; returns `Some([id])` for mixed. ## Verification matrix | Test | Pre-fix | Post-fix | |---|---|---| | `test_projection_mask_without_nested_single_leaf_struct` (new unit) | would FAIL | PASS | | `test_async_predicate_on_single_leaf_nullable_struct` (new integration) | PANIC | PASS | | `predicate_cache::test_default_read` | PASS | PASS | | `predicate_cache::test_async_cache_with_filters` | PASS | PASS | | `predicate_cache::test_sync_cache_with_filters` | PASS | PASS | | `predicate_cache::test_cache_disabled_with_filters` | PASS | PASS | | `predicate_cache::test_cache_projection_excludes_nested_columns` | PASS | PASS | | `test_projection_mask_without_nested_*` (5 existing) | PASS | PASS | # Are there any user-facing changes? --- parquet/src/arrow/mod.rs | 49 ++++++++---- parquet/tests/arrow_reader/predicate_cache.rs | 78 +++++++++++++++++++ 2 files changed, 114 insertions(+), 13 deletions(-) diff --git a/parquet/src/arrow/mod.rs b/parquet/src/arrow/mod.rs index 52152988166f..14f7b9b6b41b 100644 --- a/parquet/src/arrow/mod.rs +++ b/parquet/src/arrow/mod.rs @@ -435,24 +435,14 @@ impl ProjectionMask { root_leaf_counts[root_idx] += 1; } - // Keep only leaves whose root has exactly one leaf (non-nested) and is not a - // LIST. LIST is encoded as a wrapped logical type with a single leaf, e.g. - // https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#lists - // - // ```text - // // List (list non-null, elements nullable) - // required group my_list (LIST) { - // repeated group list { - // optional binary element (STRING); - // } - // } - // ``` + // Cache only top-level primitive columns. + // Even a one-leaf group is nested; caching it drops parent def levels. let mut included_leaves = Vec::new(); for leaf_idx in 0..num_leaves { if self.leaf_included(leaf_idx) { let root = schema.get_column_root(leaf_idx); let root_idx = schema.get_column_root_idx(leaf_idx); - if root_leaf_counts[root_idx] == 1 && !root.is_list() { + if root_leaf_counts[root_idx] == 1 && root.is_primitive() { included_leaves.push(leaf_idx); } } @@ -1042,6 +1032,39 @@ mod test { ); } + #[test] + fn test_projection_mask_without_nested_single_leaf_struct() { + // Regression: a single-leaf struct is still nested. + let schema = parse_schema( + " + message test_schema { + OPTIONAL group address { + REQUIRED BYTE_ARRAY street (UTF8); + } + REQUIRED INT32 id; + } + ", + ); + + // street -> empty; root is a struct + let mask = ProjectionMask::leaves(&schema, [0]); + assert_eq!(None, mask.without_nested_types(&schema)); + + // street, id --> id only + let mask = ProjectionMask::leaves(&schema, [0, 1]); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [1])), + mask.without_nested_types(&schema) + ); + + // all --> id only + let mask = ProjectionMask::all(); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [1])), + mask.without_nested_types(&schema) + ); + } + /// Converts a schema string into a `SchemaDescriptor` fn parse_schema(schema: &str) -> SchemaDescriptor { let parquet_group_type = parse_message_type(schema).unwrap(); diff --git a/parquet/tests/arrow_reader/predicate_cache.rs b/parquet/tests/arrow_reader/predicate_cache.rs index 85dba68c9c69..4029b4e19e20 100644 --- a/parquet/tests/arrow_reader/predicate_cache.rs +++ b/parquet/tests/arrow_reader/predicate_cache.rs @@ -18,6 +18,7 @@ //! Test for predicate cache in Parquet Arrow reader use super::io::TestReader; +use arrow::array::Array; use arrow::array::ArrayRef; use arrow::array::Int64Array; use arrow::compute::and; @@ -92,6 +93,83 @@ async fn test_cache_projection_excludes_nested_columns() { test.run_async(async_builder).await; } +/// Regression: cache must match no-cache for a nullable single-leaf struct. +#[tokio::test] +async fn test_async_predicate_on_single_leaf_nullable_struct() { + // Rows: b = NULL, then b.aa = "hello". + let aa: StringArray = StringArray::from(vec!["padding", "hello"]); + let nulls = arrow_buffer::NullBuffer::from(vec![false, true]); + let b = StructArray::new( + vec![Arc::new(Field::new("aa", DataType::Utf8, false))].into(), + vec![Arc::new(aa) as ArrayRef], + Some(nulls), + ); + let input_batch = RecordBatch::try_from_iter([("b", Arc::new(b) as ArrayRef)]).unwrap(); + + let mut output = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut output, input_batch.schema(), None).unwrap(); + writer.write(&input_batch).unwrap(); + writer.close().unwrap(); + let bytes = Bytes::from(output); + + // Since `aa` is required, `b.aa IS NULL` means `b` is NULL. + let build_is_null_filter = |schema_descr: &parquet::schema::types::SchemaDescPtr| -> RowFilter { + let mask = ProjectionMask::leaves(schema_descr, vec![0]); + let predicate = ArrowPredicateFn::new(mask.clone(), |batch: RecordBatch| { + let struct_arr = batch.column(0).as_struct(); + let leaf = struct_arr.column(0); + Ok((0..batch.num_rows()) + .map(|i| struct_arr.is_null(i) || leaf.is_null(i)) + .collect::()) + }); + RowFilter::new(vec![Box::new(predicate)]) + }; + + // Default cache. + let reader = TestReader::new(bytes.clone()); + let async_builder = + ParquetRecordBatchStreamBuilder::new_with_options(reader, ArrowReaderOptions::default()) + .await + .unwrap(); + let schema_descr = async_builder.metadata().file_metadata().schema_descr_ptr(); + let async_builder = async_builder + .with_projection(ProjectionMask::leaves(&schema_descr, vec![0])) + .with_row_filter(build_is_null_filter(&schema_descr)); + let mut stream = async_builder.build().unwrap(); + let mut row_count_cached = 0; + while let Some(batch) = stream.next().await { + row_count_cached += batch.unwrap().num_rows(); + } + + // Cache disabled. + let reader = TestReader::new(bytes.clone()); + let async_builder = + ParquetRecordBatchStreamBuilder::new_with_options(reader, ArrowReaderOptions::default()) + .await + .unwrap(); + let schema_descr = async_builder.metadata().file_metadata().schema_descr_ptr(); + let async_builder = async_builder + .with_projection(ProjectionMask::leaves(&schema_descr, vec![0])) + .with_row_filter(build_is_null_filter(&schema_descr)) + .with_max_predicate_cache_size(0); + let mut stream = async_builder.build().unwrap(); + let mut row_count_uncached = 0; + while let Some(batch) = stream.next().await { + row_count_uncached += batch.unwrap().num_rows(); + } + + assert_eq!( + row_count_uncached, 1, + "control: with cache disabled the predicate must match exactly one row (parent NULL)" + ); + assert_eq!( + row_count_cached, row_count_uncached, + "cached reader must match uncached reader; \ + got {row_count_cached} cached vs {row_count_uncached} uncached. \ + single-leaf struct roots must stay out of the cache." + ); +} + // -- Begin test infrastructure -- /// A test parquet file From aa61e0724612e8eb9e576de2903827e9d6c23a8d Mon Sep 17 00:00:00 2001 From: Liam Bao Date: Mon, 25 May 2026 07:18:05 -0400 Subject: [PATCH 21/51] [arrow-select] Replace `ArrayData` with direct `Array` construction in filter kernels (#9986) # Which issue does this PR close? - Part of #9298. # Rationale for this change # What changes are included in this PR? - Replaces several `ArrayDataBuilder` paths in `arrow-select/src/filter.rs` with direct typed array constructors. - Adds a small helper for filtered null buffers that reuses the already-computed null count. # Are these changes tested? Covered by exsiting tests # Are there any user-facing changes? No --------- Co-authored-by: Andrew Lamb --- arrow-select/src/filter.rs | 231 +++++++++++++------------------------ 1 file changed, 81 insertions(+), 150 deletions(-) diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index e95d01f2b592..aaab9d202013 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -26,9 +26,10 @@ use arrow_array::types::{ ArrowDictionaryKeyType, ArrowPrimitiveType, ByteArrayType, ByteViewType, RunEndIndexType, }; use arrow_array::*; -use arrow_buffer::{ArrowNativeType, BooleanBuffer, NullBuffer, RunEndBuffer, bit_util}; +use arrow_buffer::{ + ArrowNativeType, BooleanBuffer, NullBuffer, OffsetBuffer, RunEndBuffer, ScalarBuffer, bit_util, +}; use arrow_buffer::{Buffer, MutableBuffer}; -use arrow_data::ArrayDataBuilder; use arrow_data::bit_iterator::{BitIndexIterator, BitSliceIterator}; use arrow_data::transform::MutableArrayData; use arrow_schema::*; @@ -408,6 +409,22 @@ impl FilterPredicate { pub fn count(&self) -> usize { self.count } + + /// Filters the given `nulls` buffer using this predicate. + /// + /// Returns `None` when there is nothing to track in the output, either + /// because the input `nulls` was `None`, the input had no nulls, or the + /// filtered result has no nulls. Otherwise returns the filtered + /// [`NullBuffer`] with its precomputed null count. + pub fn filter_nulls(&self, nulls: Option<&NullBuffer>) -> Option { + let (null_count, nulls) = filter_null_mask(nulls, self)?; + let buffer = BooleanBuffer::new(nulls, 0, self.count); + + debug_assert_eq!(null_count, buffer.len() - buffer.count_set_bits()); + // SAFETY: `filter_null_mask` derived `null_count` from `buffer`, so it + // matches the number of unset bits as required by `new_unchecked`. + Some(unsafe { NullBuffer::new_unchecked(buffer, null_count) }) + } } fn filter_array(values: &dyn Array, predicate: &FilterPredicate) -> Result { @@ -624,18 +641,11 @@ fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { /// `filter` implementation for boolean buffers fn filter_boolean(array: &BooleanArray, predicate: &FilterPredicate) -> BooleanArray { - let values = filter_bits(array.values(), predicate); - - let mut builder = ArrayDataBuilder::new(DataType::Boolean) - .len(predicate.count) - .add_buffer(values); + let buffer = filter_bits(array.values(), predicate); + let values = BooleanBuffer::new(buffer, 0, predicate.count); + let nulls = predicate.filter_nulls(array.nulls()); - if let Some((null_count, nulls)) = filter_null_mask(array.nulls(), predicate) { - builder = builder.null_count(null_count).null_bit_buffer(Some(nulls)); - } - - let data = unsafe { builder.build_unchecked() }; - BooleanArray::from(data) + BooleanArray::new(values, nulls) } #[inline(never)] @@ -681,18 +691,17 @@ fn filter_primitive(array: &PrimitiveArray, predicate: &FilterPredicate) - where T: ArrowPrimitiveType, { - let values = array.values(); - let buffer = filter_native(values, predicate); - let mut builder = ArrayDataBuilder::new(array.data_type().clone()) - .len(predicate.count) - .add_buffer(buffer); - - if let Some((null_count, nulls)) = filter_null_mask(array.nulls(), predicate) { - builder = builder.null_count(null_count).null_bit_buffer(Some(nulls)); + let buffer = filter_native(array.values(), predicate); + let values = ScalarBuffer::new(buffer, 0, predicate.count); + let nulls = predicate.filter_nulls(array.nulls()); + let filtered = PrimitiveArray::new(values, nulls); + + // Avoid the compatibility check when the physical type already matches. + if array.data_type() == &T::DATA_TYPE { + filtered + } else { + filtered.with_data_type(array.data_type().clone()) } - - let data = unsafe { builder.build_unchecked() }; - PrimitiveArray::from(data) } /// [`FilterBytes`] is created from a source [`GenericByteArray`] and can be @@ -824,17 +833,15 @@ where IterationStrategy::All | IterationStrategy::None => unreachable!(), } - let mut builder = ArrayDataBuilder::new(T::DATA_TYPE) - .len(predicate.count) - .add_buffer(filter.dst_offsets.into()) - .add_buffer(filter.dst_values.into()); - - if let Some((null_count, nulls)) = filter_null_mask(array.nulls(), predicate) { - builder = builder.null_count(null_count).null_bit_buffer(Some(nulls)); - } + // SAFETY: `dst_offsets` starts at `[0]` and only grows by the running + // `cur_offset`, so it is monotonically non-decreasing. + let offsets = unsafe { OffsetBuffer::new_unchecked(filter.dst_offsets.into()) }; + let nulls = predicate.filter_nulls(array.nulls()); - let data = unsafe { builder.build_unchecked() }; - GenericByteArray::from(data) + // SAFETY: `offsets` index into `dst_values` by construction, and each slot + // is a byte-for-byte copy from `array`, so UTF-8 validity (if any) is preserved. + // Length invariant: `offsets.len() - 1 == predicate.count == nulls.len()`. + unsafe { GenericByteArray::new_unchecked(offsets, filter.dst_values.into(), nulls) } } /// `filter` implementation for byte view arrays. @@ -843,17 +850,14 @@ fn filter_byte_view( predicate: &FilterPredicate, ) -> GenericByteViewArray { let new_view_buffer = filter_native(array.views(), predicate); - - let mut builder = ArrayDataBuilder::new(T::DATA_TYPE) - .len(predicate.count) - .add_buffer(new_view_buffer) - .add_buffers(array.data_buffers().to_vec()); - - if let Some((null_count, nulls)) = filter_null_mask(array.nulls(), predicate) { - builder = builder.null_count(null_count).null_bit_buffer(Some(nulls)); - } - - GenericByteViewArray::from(unsafe { builder.build_unchecked() }) + let views = ScalarBuffer::new(new_view_buffer, 0, predicate.count); + let buffers = array.data_buffers().to_vec(); + let nulls = predicate.filter_nulls(array.nulls()); + + // SAFETY: each view is copied unchanged from `array.views()` and `buffers` + // is the same buffer list, so every view still points to an in-bounds + // (and, for strings, UTF-8 valid) range. + unsafe { GenericByteViewArray::new_unchecked(views, buffers, nulls) } } fn filter_fixed_size_binary( @@ -902,16 +906,10 @@ fn filter_fixed_size_binary( } IterationStrategy::All | IterationStrategy::None => unreachable!(), }; - let mut builder = ArrayDataBuilder::new(array.data_type().clone()) - .len(predicate.count) - .add_buffer(buffer.into()); - if let Some((null_count, nulls)) = filter_null_mask(array.nulls(), predicate) { - builder = builder.null_count(null_count).null_bit_buffer(Some(nulls)); - } + let nulls = predicate.filter_nulls(array.nulls()); - let data = unsafe { builder.build_unchecked() }; - FixedSizeBinaryArray::from(data) + FixedSizeBinaryArray::new(array.value_length(), buffer.into(), nulls) } /// `filter` implementation for dictionaries @@ -992,24 +990,19 @@ fn filter_list_view( let filtered_offsets = filter_native::(array.offsets(), predicate); let filtered_sizes = filter_native::(array.sizes(), predicate); - // Filter the nulls - let nulls = if let Some((null_count, nulls)) = filter_null_mask(array.nulls(), predicate) { - let buffer = BooleanBuffer::new(nulls, 0, predicate.count); - - Some(unsafe { NullBuffer::new_unchecked(buffer, null_count) }) - } else { - None + let field = match array.data_type() { + DataType::ListView(field) | DataType::LargeListView(field) => field.clone(), + _ => unreachable!(), }; - - let list_data = ArrayDataBuilder::new(array.data_type().clone()) - .nulls(nulls) - .buffers(vec![filtered_offsets, filtered_sizes]) - .child_data(vec![array.values().to_data()]) - .len(predicate.count); - - let list_data = unsafe { list_data.build_unchecked() }; - - GenericListViewArray::from(list_data) + let offsets = ScalarBuffer::new(filtered_offsets, 0, predicate.count); + let sizes = ScalarBuffer::new(filtered_sizes, 0, predicate.count); + let values = array.values().clone(); + let nulls = predicate.filter_nulls(array.nulls()); + + // SAFETY: each `(offset, size)` pair is copied unchanged from `array` and + // indexes into the same `values` child, so every range stays in-bounds. + // `field` and `values`' data type are unchanged from `array`. + unsafe { GenericListViewArray::new_unchecked(field, offsets, sizes, values, nulls) } } #[cfg(test)] @@ -1018,7 +1011,6 @@ mod tests { use arrow_array::builder::*; use arrow_array::cast::as_run_array; use arrow_array::types::*; - use arrow_data::ArrayData; use rand::distr::uniform::{UniformSampler, UniformUsize}; use rand::distr::{Alphanumeric, StandardUniform}; use rand::prelude::*; @@ -1494,49 +1486,22 @@ mod tests { #[test] fn test_filter_list_array() { - let value_data = ArrayData::builder(DataType::Int32) - .len(8) - .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7])) - .build() - .unwrap(); - - let value_offsets = Buffer::from_slice_ref([0i64, 3, 6, 8, 8]); - - let list_data_type = - DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int32, false))); - let list_data = ArrayData::builder(list_data_type) - .len(4) - .add_buffer(value_offsets) - .add_child_data(value_data) - .null_bit_buffer(Some(Buffer::from([0b00000111]))) - .build() - .unwrap(); - + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let offsets = OffsetBuffer::new(vec![0i64, 3, 6, 8, 8].into()); + let value_array = Arc::new(Int32Array::from_iter_values(0..8)); + let nulls = Some(NullBuffer::from(vec![true, true, true, false])); // a = [[0, 1, 2], [3, 4, 5], [6, 7], null] - let a = LargeListArray::from(list_data); + let a = LargeListArray::new(field.clone(), offsets, value_array, nulls); let b = BooleanArray::from(vec![false, true, false, true]); let result = filter(&a, &b).unwrap(); // expected: [[3, 4, 5], null] - let value_data = ArrayData::builder(DataType::Int32) - .len(3) - .add_buffer(Buffer::from_slice_ref([3, 4, 5])) - .build() - .unwrap(); - - let value_offsets = Buffer::from_slice_ref([0i64, 3, 3]); - - let list_data_type = - DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int32, false))); - let expected = ArrayData::builder(list_data_type) - .len(2) - .add_buffer(value_offsets) - .add_child_data(value_data) - .null_bit_buffer(Some(Buffer::from([0b00000001]))) - .build() - .unwrap(); + let offsets = OffsetBuffer::new(vec![0i64, 3, 3].into()); + let value_array = Arc::new(Int32Array::from_iter_values([3, 4, 5])); + let nulls = Some(NullBuffer::from(vec![true, false])); + let expected: ArrayRef = Arc::new(LargeListArray::new(field, offsets, value_array, nulls)); - assert_eq!(&make_array(expected), &result); + assert_eq!(&expected, &result); } fn test_case_filter_list_view() { @@ -1719,14 +1684,7 @@ mod tests { let truncated_length = mask_len - offset - truncate; - let data = ArrayDataBuilder::new(DataType::Boolean) - .len(truncated_length) - .offset(offset) - .add_buffer(buffer) - .build() - .unwrap(); - - let filter = BooleanArray::from(data); + let filter = BooleanArray::new(BooleanBuffer::new(buffer, offset, truncated_length), None); let slice_bits: Vec<_> = SlicesIterator::new(&filter) .flat_map(|(start, end)| start..end) @@ -1949,18 +1907,9 @@ mod tests { #[test] fn test_filter_fixed_size_list_arrays() { - let value_data = ArrayData::builder(DataType::Int32) - .len(9) - .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8])) - .build() - .unwrap(); - let list_data_type = DataType::new_fixed_size_list(DataType::Int32, 3, false); - let list_data = ArrayData::builder(list_data_type) - .len(3) - .add_child_data(value_data) - .build() - .unwrap(); - let array = FixedSizeListArray::from(list_data); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let value_array = Arc::new(Int32Array::from_iter_values(0..9)); + let array = FixedSizeListArray::new(field, 3, value_array, None); let filter_array = BooleanArray::from(vec![true, false, false]); @@ -1996,28 +1945,10 @@ mod tests { #[test] fn test_filter_fixed_size_list_arrays_with_null() { - let value_data = ArrayData::builder(DataType::Int32) - .len(10) - .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])) - .build() - .unwrap(); - - // Set null buts for the nested array: - // [[0, 1], null, null, [6, 7], [8, 9]] - // 01011001 00000001 - let mut null_bits: [u8; 1] = [0; 1]; - bit_util::set_bit(&mut null_bits, 0); - bit_util::set_bit(&mut null_bits, 3); - bit_util::set_bit(&mut null_bits, 4); - - let list_data_type = DataType::new_fixed_size_list(DataType::Int32, 2, false); - let list_data = ArrayData::builder(list_data_type) - .len(5) - .add_child_data(value_data) - .null_bit_buffer(Some(Buffer::from(null_bits))) - .build() - .unwrap(); - let array = FixedSizeListArray::from(list_data); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let value_array = Arc::new(Int32Array::from_iter_values(0..10)); + let nulls = Some(NullBuffer::from(vec![true, false, false, true, true])); + let array = FixedSizeListArray::new(field, 2, value_array, nulls); let filter_array = BooleanArray::from(vec![true, true, false, true, false]); From c4e154f38b487b3bdefd2d3c7c429bf12a2bcb81 Mon Sep 17 00:00:00 2001 From: Konstantin Tarasov <33369833+sdf-jkl@users.noreply.github.com> Date: Mon, 25 May 2026 07:18:33 -0400 Subject: [PATCH 22/51] Support Shredded Lists/Array in `variant_get` (#8354) # Which issue does this PR close? - closes #9443. # Rationale for this change We should be able to `variant_get` using Indices to path through `VariantArray`s # What changes are included in this PR? # Are these changes tested? Yes, unit tested. # Are there any user-facing changes? --------- Co-authored-by: Congxian Qiu Co-authored-by: Ryan Johnson --- parquet-variant-compute/src/variant_get.rs | 338 ++++++++++++++++++++- parquet-variant/src/path.rs | 9 + 2 files changed, 334 insertions(+), 13 deletions(-) diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index 774da0e72e8c..38c577564dea 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -15,17 +15,20 @@ // specific language governing permissions and limitations // under the License. use arrow::{ - array::{self, Array, ArrayRef, StructArray, make_array}, + array::{ + self, Array, ArrayRef, GenericListArray, GenericListViewArray, ListLikeArray, StructArray, + UInt64Array, make_array, + }, buffer::NullBuffer, - compute::CastOptions, + compute::{CastOptions, take}, datatypes::Field, error::Result, }; use arrow_schema::{ArrowError, DataType, FieldRef}; use parquet_variant::{VariantPath, VariantPathElement}; +use crate::ShreddingState; use crate::VariantArray; -use crate::variant_array::ShreddingState; use crate::variant_to_arrow::make_variant_to_arrow_row_builder; use arrow::array::AsArray; @@ -43,12 +46,70 @@ pub(crate) enum ShreddedPathStep { NotShredded, } +/// Build the next shredding state by taking one list-like element (at `index`) per input row. +/// +fn take_list_like_index_as_shredding_state( + typed_value: &dyn Array, + index: usize, +) -> Result> { + let list_array = typed_value.as_any().downcast_ref::().ok_or_else(|| { + ArrowError::ComputeError(format!( + "Expected array type '{}' while handling list-like path step, got '{}'", + std::any::type_name::(), + typed_value.data_type() + )) + })?; + + let values = list_array.values(); + + let Some(struct_array) = values.as_struct_opt() else { + return Ok(None); + }; + let shredding_state = ShreddingState::try_from(struct_array)?; + + let value_array = shredding_state.value_field(); + let typed_array = shredding_state.typed_value_field(); + + // If list elements have neither typed nor fallback value, this path step is missing. + if value_array.is_none() && typed_array.is_none() { + return Ok(None); + } + + let mut take_indices = Vec::with_capacity(list_array.len()); + for row in 0..list_array.len() { + let row_range = list_array.element_range(row); + let take_index = (index < row_range.len()).then(|| (row_range.start + index) as u64); + take_indices.push(take_index); + } + + let index_array = UInt64Array::from(take_indices); + + // Gather both typed and fallback values at the requested element index. + let taken_value = value_array + .map(|value| take(value, &index_array, None)) + .transpose()?; + let taken_typed = typed_array + .map(|typed| take(typed, &index_array, None)) + .transpose()?; + + Ok(Some(ShreddingState::new(taken_value, taken_typed))) +} + /// Given a shredded variant field -- a `(value?, typed_value?)` pair -- try to take one path step /// deeper. For a `VariantPathElement::Field`, if there is no `typed_value` at this level, if /// `typed_value` is not a struct, or if the requested field name does not exist, traversal returns /// a missing-path step (`Missing` or `NotShredded` depending on whether `value` exists). /// -/// TODO: Support `VariantPathElement::Index`? It wouldn't be easy, and maybe not even possible. +/// Safe-cast behavior (`cast_options.safe = true`): +/// - Type mismatch during path traversal (for example field access on non-struct, index access on +/// non-list) returns [`ShreddedPathStep::Missing`] or [`ShreddedPathStep::NotShredded`], allowing +/// the caller to continue with null/fallback semantics. +/// - List index out-of-bounds produces nulls for the corresponding rows. +/// +/// Unsafe-cast behavior (`cast_options.safe = false`): +/// - Field access on non-struct returns [`ArrowError::CastError`]. +/// - List index path steps follow JSONPath semantics and return missing/null for non-list or +/// out-of-bounds rows. pub(crate) fn follow_shredded_path_element( shredding_state: &ShreddingState, path_element: &VariantPathElement<'_>, @@ -69,7 +130,7 @@ pub(crate) fn follow_shredded_path_element( VariantPathElement::Field { name } => { // Try to step into the requested field name of a struct. // First, try to downcast to StructArray - let Some(struct_array) = typed_value.as_any().downcast_ref::() else { + let Some(struct_array) = typed_value.as_struct_opt() else { // Object field path step follows JSONPath semantics and returns missing path step (NotShredded/Missing) on non-struct path return Ok(missing_path_step()); }; @@ -93,12 +154,30 @@ pub(crate) fn follow_shredded_path_element( let state = ShreddingState::try_from(struct_array)?; Ok(ShreddedPathStep::Success(state)) } - VariantPathElement::Index { .. } => { - // TODO: Support array indexing. Among other things, it will require slicing not - // only the array we have here, but also the corresponding metadata and null masks. - Err(ArrowError::NotYetImplemented( - "Pathing into shredded variant array index".into(), - )) + VariantPathElement::Index { index } => { + let state = match typed_value.data_type() { + DataType::List(_) => take_list_like_index_as_shredding_state::< + GenericListArray, + >(typed_value.as_ref(), *index)?, + DataType::LargeList(_) => take_list_like_index_as_shredding_state::< + GenericListArray, + >(typed_value.as_ref(), *index)?, + DataType::ListView(_) => take_list_like_index_as_shredding_state::< + GenericListViewArray, + >(typed_value.as_ref(), *index)?, + DataType::LargeListView(_) => take_list_like_index_as_shredding_state::< + GenericListViewArray, + >(typed_value.as_ref(), *index)?, + _ => { + // JSONPath semantics: indexing a non-list yields no match. + return Ok(missing_path_step()); + } + }; + + match state { + Some(state) => Ok(ShreddedPathStep::Success(state)), + None => Ok(missing_path_step()), + } } } } @@ -356,7 +435,8 @@ mod test { use super::{GetOptions, variant_get}; use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder}; use crate::{ - VariantArray, VariantArrayBuilder, cast_to_variant, json_to_variant, shred_variant, + ShreddedSchemaBuilder, VariantArray, VariantArrayBuilder, cast_to_variant, json_to_variant, + shred_variant, }; use arrow::array::{ Array, ArrayRef, AsArray, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, @@ -448,7 +528,7 @@ mod test { fn get_primitive_variant_inside_object_of_list() { single_variant_get_test( r#"{"some_field": [1234]}"#, - VariantPath::try_from("some_field").unwrap().join(0), + VariantPath::try_from("some_field[0]").unwrap(), "1234", ); } @@ -1741,6 +1821,7 @@ mod test { Some(nulls), )) } + /// This test manually constructs a shredded variant array representing objects /// like {"x": 1, "y": "foo"} and {"x": 42} and tests extracting the "x" field /// as VariantArray using variant_get. @@ -1777,6 +1858,210 @@ mod test { assert_eq!(&result, &expected); } + type ShreddedListLikeArrayGen = fn() -> ArrayRef; + type ShreddedListLikeCase = (&'static str, ShreddedListLikeArrayGen); + + fn shredded_list_like_cases() -> [ShreddedListLikeCase; 4] { + [ + ("list", shredded_list_variant_array), + ("large_list", shredded_large_list_variant_array), + ("list_view", shredded_list_view_variant_array), + ("large_list_view", shredded_large_list_view_variant_array), + ] + } + + #[test] + fn test_shredded_list_like_index_access_from_value_field() { + let options = GetOptions::new_with_path(VariantPath::from(1)); + + for (case, array_gen) in shredded_list_like_cases() { + let array = array_gen(); + let result = variant_get(&array, options.clone()).unwrap(); + let result_variant = VariantArray::try_new(&result).unwrap(); + + assert_eq!(result_variant.value(0), Variant::from("drama"), "{case}"); + assert_eq!(result_variant.value(1).as_int64(), Some(123), "{case}"); + } + } + + #[test] + fn test_shredded_list_like_index_out_of_bounds_unsafe_cast_returns_null() { + let options = + GetOptions::new_with_path(VariantPath::from(10)).with_cast_options(CastOptions { + safe: false, + ..Default::default() + }); + + for (case, array_gen) in shredded_list_like_cases() { + let result = variant_get(&array_gen(), options.clone()).unwrap(); + let result_variant = VariantArray::try_new(&result).unwrap(); + assert_eq!(result_variant.value(0), Variant::Null, "{case}"); + assert_eq!(result_variant.value(1), Variant::Null, "{case}"); + } + } + + /// Test extracting shredded list-like field with type conversion. + #[test] + fn test_shredded_list_like_as_string() { + let field = Field::new("typed_value", DataType::Utf8, false); + let options = GetOptions::new_with_path(VariantPath::from(0)) + .with_as_type(Some(FieldRef::from(field))); + let expected: ArrayRef = Arc::new(StringArray::from(vec![Some("comedy"), Some("horror")])); + + for (case, array_gen) in shredded_list_like_cases() { + let result = variant_get(&array_gen(), options.clone()).unwrap(); + assert_eq!(&result, &expected, "{case}"); + } + } + + #[test] + fn test_shredded_list_like_index_access_from_value_field_as_int64() { + let field = Field::new("typed_value", DataType::Int64, true); + let options = GetOptions::new_with_path(VariantPath::from(1)) + .with_as_type(Some(FieldRef::from(field))); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(123)])); + + for (case, array_gen) in shredded_list_like_cases() { + let result = variant_get(&array_gen(), options.clone()).unwrap(); + // "drama" -> NULL, 123 -> 123. + assert_eq!(&result, &expected, "{case}"); + } + } + + #[test] + fn test_shredded_list_in_struct_index_access() { + let array = shredded_struct_with_list_variant_array(); + let options = GetOptions::new_with_path(VariantPath::try_from("a[1]").unwrap()); + let result = variant_get(&array, options).unwrap(); + let result_variant = VariantArray::try_new(&result).unwrap(); + + assert_eq!(result_variant.value(0), Variant::from("drama")); + assert_eq!(result_variant.value(1).as_int64(), Some(123)); + } + + #[test] + fn test_shredded_struct_in_list_field_access() { + let array = shredded_list_of_struct_variant_array(); + let field = Field::new("x", DataType::Int32, true); + let path = VariantPath::from(0).join("x"); + let options = GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field))); + let result = variant_get(&array, options).unwrap(); + + let expected: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(3)])); + assert_eq!(&result, &expected); + } + + #[test] + fn test_shredded_list_of_lists_index_access() { + let array = shredded_list_of_lists_variant_array(); + let path = VariantPath::from(0).join(1); + + let result = variant_get(&array, GetOptions::new_with_path(path.clone())).unwrap(); + let result_variant = VariantArray::try_new(&result).unwrap(); + assert_eq!(result_variant.value(0), Variant::from("b")); + assert_eq!(result_variant.value(1).as_int64(), Some(123)); + + let field = Field::new("typed_value", DataType::Int64, true); + let casted = variant_get( + &array, + GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field))), + ) + .unwrap(); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(123)])); + assert_eq!(&casted, &expected); + } + + /// Helper to create a shredded list-like variant array used by list index tests. + /// + /// Rows: + /// 1. `["comedy", "drama"]` (fully shred-able as `Utf8`) + /// 2. `["horror", 123]` (partially shredded, with fallback for the numeric element) + fn shredded_list_like_variant_array(list_schema: DataType) -> ArrayRef { + let json_rows: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"["comedy", "drama"]"#), + Some(r#"["horror", 123]"#), + ])); + let input = json_to_variant(&json_rows).unwrap(); + + let shredded = shred_variant(&input, &list_schema).unwrap(); + ArrayRef::from(shredded) + } + + fn shredded_list_of_lists_variant_array() -> ArrayRef { + let json_rows: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"[["a", "b"], ["c", "d"]]"#), + Some(r#"[["x", 123], ["y", "z"]]"#), + ])); + let input = json_to_variant(&json_rows).unwrap(); + + let inner_list = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))); + let outer_list = DataType::List(Arc::new(Field::new("item", inner_list, true))); + let shredded = shred_variant(&input, &outer_list).unwrap(); + ArrayRef::from(shredded) + } + + fn shredded_list_variant_array() -> ArrayRef { + shredded_list_like_variant_array(DataType::List(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))) + } + + fn shredded_large_list_variant_array() -> ArrayRef { + shredded_list_like_variant_array(DataType::LargeList(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))) + } + + fn shredded_list_view_variant_array() -> ArrayRef { + shredded_list_like_variant_array(DataType::ListView(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))) + } + + fn shredded_large_list_view_variant_array() -> ArrayRef { + shredded_list_like_variant_array(DataType::LargeListView(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))) + } + + fn shredded_struct_with_list_variant_array() -> ArrayRef { + let json_rows: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"{"a": ["comedy", "drama"]}"#), + Some(r#"{"a": ["horror", 123]}"#), + ])); + let input = json_to_variant(&json_rows).unwrap(); + + let list_schema = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))); + let shredding_schema = ShreddedSchemaBuilder::default() + .with_path("a", &list_schema) + .unwrap() + .build(); + let shredded = shred_variant(&input, &shredding_schema).unwrap(); + ArrayRef::from(shredded) + } + + fn shredded_list_of_struct_variant_array() -> ArrayRef { + let json_rows: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"[{"x": 1}, {"x": 2}]"#), + Some(r#"[{"x": 3}, {"y": 4}]"#), + ])); + let input = json_to_variant(&json_rows).unwrap(); + + let struct_type = + DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, true)])); + let list_schema = DataType::List(Arc::new(Field::new("item", struct_type, true))); + let shredded = shred_variant(&input, &list_schema).unwrap(); + ArrayRef::from(shredded) + } + /// Helper function to create a shredded variant array representing objects /// /// This creates an array that represents: @@ -2494,6 +2779,33 @@ mod test { } } + #[test] + fn test_strict_cast_options_index_on_non_list_returns_null() { + use arrow::compute::CastOptions; + use arrow::datatypes::{DataType, Field}; + use parquet_variant::VariantPath; + use std::sync::Arc; + + // Use existing test data that has Int32 typed_value at the top level. + let variant_array = perfectly_shredded_int32_variant_array(); + let options = GetOptions { + path: VariantPath::from(0), + as_type: Some(Arc::new(Field::new("result", DataType::Int32, true))), + cast_options: CastOptions { + safe: false, + ..Default::default() + }, + }; + + let variant_array_ref: Arc = variant_array.clone(); + let result = variant_get(&variant_array_ref, options).unwrap(); + + assert_eq!(result.len(), 3); + assert!(result.is_null(0)); + assert!(result.is_null(1)); + assert!(result.is_null(2)); + } + #[test] fn test_error_message_boolean_type_display() { let mut builder = VariantArrayBuilder::new(1); diff --git a/parquet-variant/src/path.rs b/parquet-variant/src/path.rs index 8e68d9efadf2..0837e49f8abf 100644 --- a/parquet-variant/src/path.rs +++ b/parquet-variant/src/path.rs @@ -346,5 +346,14 @@ mod tests { err.to_string(), "Parser error: Invalid token in bracket request: `abc`. Expected a quoted string or a number(e.g., `['field']` or `[123]`)" ); + + // Out-of-range integer indexes are invalid path tokens. + let too_large_index = (usize::MAX as u128) + 1; + let err = VariantPath::try_from(format!("[{too_large_index}]").as_str()).unwrap_err(); + assert!( + err.to_string() + .contains("Parser error: Invalid token in bracket request"), + "{err}" + ); } } From e28fd0d0f1f45e5608928eb52af3b12ffe9fee2e Mon Sep 17 00:00:00 2001 From: Konstantin Tarasov <33369833+sdf-jkl@users.noreply.github.com> Date: Mon, 25 May 2026 07:19:27 -0400 Subject: [PATCH 23/51] Add `DatePart` enum 1-indexed variants (#9965) # Which issue does this PR close? - Closes #9964. # Rationale for this change Remove extra computations done on the datafusion side for temporal functions. # What changes are included in this PR? - Added two new `DatePart` variant - Closed some tests lists drift # Are these changes tested? - Unit tests # Are there any user-facing changes? No --------- Co-authored-by: Claude Opus 4.7 (1M context) --- arrow-arith/src/temporal.rs | 88 +++++++++++++++++++++++++++++++++---- 1 file changed, 80 insertions(+), 8 deletions(-) diff --git a/arrow-arith/src/temporal.rs b/arrow-arith/src/temporal.rs index 96a3b43a3811..301ad172dabc 100644 --- a/arrow-arith/src/temporal.rs +++ b/arrow-arith/src/temporal.rs @@ -63,6 +63,10 @@ pub enum DatePart { DayOfWeekSunday0, /// Day of the week, in range `0..=6`, where Monday is `0` DayOfWeekMonday0, + /// Day of the week, in range `1..=7`, where Sunday is `1` + DayOfWeekSunday1, + /// ISO day of the week, in range `1..=7`, where Monday is `1` + DayOfWeekMonday1, /// Day of year, in range `1..=366` DayOfYear, /// Hour of the day, in range `0..=23` @@ -100,11 +104,6 @@ impl std::fmt::Display for DatePart { /// - `century`, `decade`, `millennium` — no matching [`DatePart`] variant. /// - `timezone`, `timezone_hour`, `timezone_minute` — not modeled by /// [`DatePart`]. -/// - `isodow` — PostgreSQL's `isodow` returns `1..=7` (Mon=1), but the -/// closest variant ([`DatePart::DayOfWeekMonday0`]) returns `0..=6`. -/// Accepting it here would silently shift the value range; callers that -/// need the PostgreSQL semantic should map the alias themselves and -/// add `1` to the extracted value. impl FromStr for DatePart { type Err = ArrowError; @@ -118,6 +117,9 @@ impl FromStr for DatePart { "isoweek" => Self::WeekISO, "d" | "day" | "days" => Self::Day, "dow" | "dayofweek" => Self::DayOfWeekSunday0, + "dow1" | "dayofweek1" => Self::DayOfWeekSunday1, + "isodow0" => Self::DayOfWeekMonday0, + "isodow" => Self::DayOfWeekMonday1, "doy" | "dayofyear" => Self::DayOfYear, "h" | "hr" | "hrs" | "hour" | "hours" => Self::Hour, "m" | "min" | "mins" | "minute" | "minutes" => Self::Minute, @@ -158,6 +160,8 @@ where DatePart::Day => |d| d.day() as i32, DatePart::DayOfWeekSunday0 => |d| d.num_days_from_sunday(), DatePart::DayOfWeekMonday0 => |d| d.num_days_from_monday(), + DatePart::DayOfWeekSunday1 => |d| d.num_days_from_sunday() + 1, + DatePart::DayOfWeekMonday1 => |d| d.num_days_from_monday() + 1, DatePart::DayOfYear => |d| d.ordinal() as i32, DatePart::Hour => |d| d.hour() as i32, DatePart::Minute => |d| d.minute() as i32, @@ -500,6 +504,8 @@ impl ExtractDatePartExt for PrimitiveArray { | DatePart::Day | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 + | DatePart::DayOfWeekSunday1 + | DatePart::DayOfWeekMonday1 | DatePart::DayOfYear | DatePart::Hour | DatePart::Minute @@ -536,6 +542,8 @@ impl ExtractDatePartExt for PrimitiveArray { | DatePart::Month | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 + | DatePart::DayOfWeekSunday1 + | DatePart::DayOfWeekMonday1 | DatePart::DayOfYear => { return_compute_error_with!(format!("{part} does not support"), self.data_type()) } @@ -578,6 +586,8 @@ impl ExtractDatePartExt for PrimitiveArray { | DatePart::YearISO | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 + | DatePart::DayOfWeekSunday1 + | DatePart::DayOfWeekMonday1 | DatePart::DayOfYear => { return_compute_error_with!(format!("{part} does not support"), self.data_type()) } @@ -610,6 +620,8 @@ impl ExtractDatePartExt for PrimitiveArray { | DatePart::Month | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 + | DatePart::DayOfWeekSunday1 + | DatePart::DayOfWeekMonday1 | DatePart::DayOfYear => { return_compute_error_with!(format!("{part} does not support"), self.data_type()) } @@ -642,6 +654,8 @@ impl ExtractDatePartExt for PrimitiveArray { | DatePart::Month | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 + | DatePart::DayOfWeekSunday1 + | DatePart::DayOfWeekMonday1 | DatePart::DayOfYear => { return_compute_error_with!(format!("{part} does not support"), self.data_type()) } @@ -674,6 +688,8 @@ impl ExtractDatePartExt for PrimitiveArray { | DatePart::Month | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 + | DatePart::DayOfWeekSunday1 + | DatePart::DayOfWeekMonday1 | DatePart::DayOfYear => { return_compute_error_with!(format!("{part} does not support"), self.data_type()) } @@ -706,6 +722,8 @@ impl ExtractDatePartExt for PrimitiveArray { | DatePart::Month | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 + | DatePart::DayOfWeekSunday1 + | DatePart::DayOfWeekMonday1 | DatePart::DayOfYear => { return_compute_error_with!(format!("{part} does not support"), self.data_type()) } @@ -985,6 +1003,46 @@ mod tests { assert_eq!(2, b.value(2)); } + #[test] + fn test_temporal_array_date64_dayofweek1() { + //1514764800000 -> 2018-01-01 (Monday) + //1550636625000 -> 2019-02-20 (Wednesday) + //1483228800000 -> 2017-01-01 (Sunday) + let a: PrimitiveArray = vec![ + Some(1514764800000), + None, + Some(1550636625000), + Some(1483228800000), + ] + .into(); + + let b = date_part_primitive(&a, DatePart::DayOfWeekSunday1).unwrap(); + assert_eq!(2, b.value(0)); + assert!(!b.is_valid(1)); + assert_eq!(4, b.value(2)); + assert_eq!(1, b.value(3)); + } + + #[test] + fn test_temporal_array_date64_isodow() { + //1514764800000 -> 2018-01-01 (Monday) + //1550636625000 -> 2019-02-20 (Wednesday) + //1483228800000 -> 2017-01-01 (Sunday) + let a: PrimitiveArray = vec![ + Some(1514764800000), + None, + Some(1550636625000), + Some(1483228800000), + ] + .into(); + + let b = date_part_primitive(&a, DatePart::DayOfWeekMonday1).unwrap(); + assert_eq!(1, b.value(0)); + assert!(!b.is_valid(1)); + assert_eq!(3, b.value(2)); + assert_eq!(7, b.value(3)); + } + #[test] fn test_temporal_array_date64_weekday0() { //1483228800000 -> 2017-01-01 (Sunday) @@ -1579,11 +1637,15 @@ mod tests { let invalid_parts = [ DatePart::Quarter, DatePart::Year, + DatePart::YearISO, DatePart::Month, DatePart::Week, + DatePart::WeekISO, DatePart::Day, DatePart::DayOfWeekSunday0, DatePart::DayOfWeekMonday0, + DatePart::DayOfWeekSunday1, + DatePart::DayOfWeekMonday1, DatePart::DayOfYear, ]; @@ -1813,8 +1875,12 @@ mod tests { fn ensure_returns_error(array: &dyn Array) { let invalid_parts = [ DatePart::Quarter, + DatePart::YearISO, + DatePart::WeekISO, DatePart::DayOfWeekSunday0, DatePart::DayOfWeekMonday0, + DatePart::DayOfWeekSunday1, + DatePart::DayOfWeekMonday1, DatePart::DayOfYear, ]; @@ -1956,10 +2022,14 @@ mod tests { fn ensure_returns_error(array: &dyn Array) { let invalid_parts = [ DatePart::Year, + DatePart::YearISO, DatePart::Quarter, DatePart::Month, + DatePart::WeekISO, DatePart::DayOfWeekSunday0, DatePart::DayOfWeekMonday0, + DatePart::DayOfWeekSunday1, + DatePart::DayOfWeekMonday1, DatePart::DayOfYear, ]; @@ -2157,6 +2227,11 @@ mod tests { ("day", DatePart::Day), ("dow", DatePart::DayOfWeekSunday0), ("DayOfWeek", DatePart::DayOfWeekSunday0), + ("dow1", DatePart::DayOfWeekSunday1), + ("dayofweek1", DatePart::DayOfWeekSunday1), + ("DAYOFWEEK1", DatePart::DayOfWeekSunday1), + ("isodow", DatePart::DayOfWeekMonday1), + ("ISODOW", DatePart::DayOfWeekMonday1), ("doy", DatePart::DayOfYear), ("DayOfYear", DatePart::DayOfYear), ("h", DatePart::Hour), @@ -2195,8 +2270,6 @@ mod tests { #[test] fn test_date_part_from_str_unknown() { // Names intentionally rejected — see the FromStr doc comment for why. - // `isodow` is here because mapping it to `DayOfWeekMonday0` would - // silently shift the value range vs. PostgreSQL's `isodow` (1..=7). let unknown = [ "epoch", "century", @@ -2205,7 +2278,6 @@ mod tests { "timezone", "timezone_hour", "timezone_minute", - "isodow", // Whitespace is not trimmed — pin this so the behavior doesn't // change silently. Callers should trim before parsing. " year ", From 2b2a95ac966ba73b3d020259d43111d2f76c557b Mon Sep 17 00:00:00 2001 From: ClSlaid Date: Wed, 27 May 2026 03:36:47 +0800 Subject: [PATCH 24/51] arrow: add oversized coalesce take benchmarks (#9799) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - add `coalesce_kernels` take benchmarks alongside the existing filter coverage for primitive, view, utf8, and dictionary schemas - add oversized repeated-index benchmark cases that stress the materialized `push_batch_with_indices` fallback with a configured `biggest_coalesce_batch_size` - extend the benchmark harness so take benchmarks can use explicit output lengths and optionally drain all completed batches while running ## Verification - cargo fmt -p arrow - cargo clippy -p arrow --bench coalesce_kernels --features test_utils -- -D warnings - cargo bench -p arrow --bench coalesce_kernels --features test_utils -- 'extra_large_repeat' Signed-off-by: 蔡略 --- arrow/benches/coalesce_kernels.rs | 394 +++++++++++++++++++++++++++++- 1 file changed, 393 insertions(+), 1 deletion(-) diff --git a/arrow/benches/coalesce_kernels.rs b/arrow/benches/coalesce_kernels.rs index b85c5cc532db..0816d1a2e8bb 100644 --- a/arrow/benches/coalesce_kernels.rs +++ b/arrow/benches/coalesce_kernels.rs @@ -25,6 +25,7 @@ use arrow_array::types::{Float64Type, Int32Type, TimestampNanosecondType}; use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; use arrow_select::coalesce::BatchCoalescer; use criterion::{Criterion, criterion_group, criterion_main}; +use rand::{Rng, SeedableRng, rngs::StdRng}; /// Benchmarks for generating evently sized output RecordBatches /// from a sequence of filtered source batches @@ -158,7 +159,155 @@ fn add_all_filter_benchmarks(c: &mut Criterion) { } } -criterion_group!(benches, add_all_filter_benchmarks); +fn add_all_take_benchmarks(c: &mut Criterion) { + let batch_size = 8192; + + let primitive_schema = SchemaRef::new(Schema::new(vec![ + Field::new("int32_val", DataType::Int32, true), + Field::new("float_val", DataType::Float64, true), + Field::new( + "timestamp_val", + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + true, + ), + ])); + + let single_schema = SchemaRef::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8View, + true, + )])); + + let single_binaryview_schema = SchemaRef::new(Schema::new(vec![Field::new( + "value", + DataType::BinaryView, + true, + )])); + + let mixed_utf8view_schema = SchemaRef::new(Schema::new(vec![ + Field::new("int32_val", DataType::Int32, true), + Field::new("float_val", DataType::Float64, true), + Field::new("utf8view_val", DataType::Utf8View, true), + ])); + + let mixed_binaryview_schema = SchemaRef::new(Schema::new(vec![ + Field::new("int32_val", DataType::Int32, true), + Field::new("float_val", DataType::Float64, true), + Field::new("binaryview_val", DataType::BinaryView, true), + ])); + + let mixed_utf8_schema = SchemaRef::new(Schema::new(vec![ + Field::new("int32_val", DataType::Int32, true), + Field::new("float_val", DataType::Float64, true), + Field::new("utf8", DataType::Utf8, true), + ])); + + let mixed_dict_schema = SchemaRef::new(Schema::new(vec![ + Field::new( + "string_dict", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + ), + Field::new("float_val1", DataType::Float64, true), + Field::new("float_val2", DataType::Float64, true), + ])); + + for null_density in [0.0, 0.1] { + for selectivity in [0.001, 0.01, 0.1, 0.8] { + for scenario in [ + TakeBenchmarkScenario { + name: "primitive", + num_output_batches: 50, + max_string_len: 30, + schema: &primitive_schema, + }, + TakeBenchmarkScenario { + name: "single_utf8view", + num_output_batches: 50, + max_string_len: 30, + schema: &single_schema, + }, + TakeBenchmarkScenario { + name: "single_binaryview", + num_output_batches: 50, + max_string_len: 30, + schema: &single_binaryview_schema, + }, + TakeBenchmarkScenario { + name: "mixed_utf8view (max_string_len=20)", + num_output_batches: 20, + max_string_len: 20, + schema: &mixed_utf8view_schema, + }, + TakeBenchmarkScenario { + name: "mixed_utf8view (max_string_len=128)", + num_output_batches: 20, + max_string_len: 128, + schema: &mixed_utf8view_schema, + }, + TakeBenchmarkScenario { + name: "mixed_binaryview (max_string_len=20)", + num_output_batches: 20, + max_string_len: 20, + schema: &mixed_binaryview_schema, + }, + TakeBenchmarkScenario { + name: "mixed_binaryview (max_string_len=128)", + num_output_batches: 20, + max_string_len: 128, + schema: &mixed_binaryview_schema, + }, + TakeBenchmarkScenario { + name: "mixed_utf8", + num_output_batches: 20, + max_string_len: 30, + schema: &mixed_utf8_schema, + }, + TakeBenchmarkScenario { + name: "mixed_dict", + num_output_batches: 10, + max_string_len: 30, + schema: &mixed_dict_schema, + }, + ] { + TakeBenchmarkBuilder::from_scenario( + c, + batch_size, + null_density, + selectivity, + scenario, + ) + .build(); + } + } + } + + // Repeated indices make the taken batch much larger than the source batch, + // which exercises the materialized fallback path for unsupported schemas. + for (name, schema) in [ + ("primitive extra_large_repeat", &primitive_schema), + ("mixed_utf8 extra_large_repeat", &mixed_utf8_schema), + ] { + TakeBenchmarkBuilder::from_scenario( + c, + batch_size, + 0.0, + 1.0, + TakeBenchmarkScenario { + name, + num_output_batches: 64, + max_string_len: 30, + schema, + }, + ) + .with_biggest_coalesce_batch_size(1024) + .with_index_output_len(131_072) + .with_drain_all_completed_batches() + .build(); + } +} + +criterion_group!(benches, add_all_filter_benchmarks, add_all_take_benchmarks); criterion_main!(benches); /// Run the filters with a batch_size, null_density, selectivity, and schema @@ -222,6 +371,129 @@ impl FilterBenchmarkBuilder<'_> { } } +struct TakeBenchmarkBuilder<'a> { + c: &'a mut Criterion, + name: &'a str, + batch_size: usize, + num_output_batches: usize, + null_density: f32, + selectivity: f32, + max_string_len: usize, + schema: &'a SchemaRef, + biggest_coalesce_batch_size: Option, + index_output_len: Option, + drain_all_completed_batches: bool, +} + +#[derive(Clone, Copy)] +struct TakeBenchmarkScenario<'a> { + name: &'a str, + num_output_batches: usize, + max_string_len: usize, + schema: &'a SchemaRef, +} + +impl<'a> TakeBenchmarkBuilder<'a> { + fn from_scenario( + c: &'a mut Criterion, + batch_size: usize, + null_density: f32, + selectivity: f32, + scenario: TakeBenchmarkScenario<'a>, + ) -> TakeBenchmarkBuilder<'a> { + let TakeBenchmarkScenario { + name, + num_output_batches, + max_string_len, + schema, + } = scenario; + + TakeBenchmarkBuilder { + c, + name, + batch_size, + num_output_batches, + null_density, + selectivity, + max_string_len, + schema, + biggest_coalesce_batch_size: None, + index_output_len: None, + drain_all_completed_batches: false, + } + } + + fn with_biggest_coalesce_batch_size(mut self, limit: usize) -> Self { + self.biggest_coalesce_batch_size = Some(limit); + self + } + + fn with_index_output_len(mut self, output_len: usize) -> Self { + self.index_output_len = Some(output_len); + self + } + + fn with_drain_all_completed_batches(mut self) -> Self { + self.drain_all_completed_batches = true; + self + } + + fn build(self) { + let Self { + c, + name, + batch_size, + num_output_batches, + null_density, + selectivity, + max_string_len, + schema, + biggest_coalesce_batch_size, + index_output_len, + drain_all_completed_batches, + } = self; + + let output_len = index_output_len + .unwrap_or_else(|| ((batch_size as f32) * selectivity).round().max(1.0) as usize); + + let indices = match index_output_len { + Some(_) => IndexStreamBuilder::new() + .with_batch_size(batch_size) + .with_output_len(output_len) + .build(), + None => IndexStreamBuilder::new() + .with_batch_size(batch_size) + .with_selectivity(selectivity) + .build(), + }; + + let data = DataStreamBuilder::new(Arc::clone(schema)) + .with_batch_size(batch_size) + .with_null_density(null_density) + .with_max_string_len(max_string_len) + .build(); + + let id = if index_output_len.is_some() || biggest_coalesce_batch_size.is_some() { + format!( + "take: {name}, input: {batch_size}, output: {output_len}, nulls: {null_density}, biggest: {biggest_coalesce_batch_size:?}" + ) + } else { + format!("take: {name}, {batch_size}, nulls: {null_density}, selectivity: {selectivity}") + }; + c.bench_function(&id, |b| { + b.iter(|| { + take_streams( + num_output_batches, + indices.clone(), + data.clone(), + biggest_coalesce_batch_size, + drain_all_completed_batches, + ); + }) + }); + } +} + /// Pull RecordBatches from a data stream and apply a sequence of /// filters from a filter stream until we have a specified number of output /// batches. @@ -247,6 +519,34 @@ fn filter_streams( } } +fn take_streams( + mut num_output_batches: usize, + mut index_stream: IndexStream, + mut data_stream: DataStream, + biggest_coalesce_batch_size: Option, + drain_all_completed_batches: bool, +) { + let schema = data_stream.schema(); + let batch_size = data_stream.batch_size(); + let mut coalescer = BatchCoalescer::new(Arc::clone(schema), batch_size) + .with_biggest_coalesce_batch_size(biggest_coalesce_batch_size); + + while num_output_batches > 0 { + let indices = index_stream.next_indices(); + let batch = data_stream.next_batch(); + coalescer + .push_batch_with_indices(batch.clone(), indices) + .unwrap(); + if drain_all_completed_batches { + while num_output_batches > 0 && coalescer.next_completed_batch().is_some() { + num_output_batches -= 1; + } + } else if coalescer.next_completed_batch().is_some() { + num_output_batches -= 1; + } + } +} + /// Stream of filters to apply to a sequence of input RecordBatches /// /// This pre-computes a sequence of filters and then repeats it forever. @@ -257,6 +557,25 @@ struct FilterStream { batches: Arc<[BooleanArray]>, } +#[derive(Debug, Clone)] +struct IndexStream { + index: usize, + batches: Arc<[UInt32Array]>, +} + +impl IndexStream { + fn next_indices(&mut self) -> &UInt32Array { + let current_index = self.index; + self.index += 1; + if self.index >= self.batches.len() { + self.index = 0; + } + self.batches + .get(current_index) + .expect("No more index batches available") + } +} + impl FilterStream { pub fn next_filter(&mut self) -> &BooleanArray { let current_index = self.index; @@ -324,6 +643,68 @@ impl FilterStreamBuilder { } } +#[derive(Debug)] +struct IndexStreamBuilder { + batch_size: usize, + num_batches: usize, + output_len: Option, + selectivity: f32, +} + +impl IndexStreamBuilder { + fn new() -> Self { + Self { + batch_size: 8192, + num_batches: 11, + output_len: None, + selectivity: 0.5, + } + } + + fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = batch_size; + self + } + + fn with_selectivity(mut self, selectivity: f32) -> Self { + assert!((0.0..=1.0).contains(&selectivity)); + self.selectivity = selectivity; + self + } + + fn with_output_len(mut self, output_len: usize) -> Self { + self.output_len = Some(output_len.max(1)); + self + } + + fn build(self) -> IndexStream { + let Self { + batch_size, + num_batches, + output_len, + selectivity, + } = self; + + let output_len = output_len + .unwrap_or_else(|| ((batch_size as f32) * selectivity).round().max(1.0) as usize); + let batches = (0..num_batches) + .map(|seed| create_index_array(batch_size, output_len, seed as u64)) + .collect::>(); + + IndexStream { + index: 0, + batches: Arc::from(batches), + } + } +} + +fn create_index_array(input_len: usize, output_len: usize, seed: u64) -> UInt32Array { + let mut rng = StdRng::seed_from_u64(seed); + UInt32Array::from_iter_values( + (0..output_len).map(|_| rng.random_range(0..u32::try_from(input_len).unwrap())), + ) +} + #[derive(Debug, Clone)] struct DataStream { schema: SchemaRef, @@ -455,6 +836,17 @@ impl DataStreamBuilder { self.max_string_len, )) // TODO seed } + DataType::BinaryView => Arc::new(BinaryViewArray::from_iter( + create_binary_array_with_len_range_and_prefix_and_seed::( + self.batch_size, + self.null_density, + 0, + self.max_string_len, + b"", + seed, + ) + .iter(), + )), DataType::Dictionary(key_type, value_type) if key_type.as_ref() == &DataType::Int32 && value_type.as_ref() == &DataType::Utf8 => From bbbe8a60b950d70b0f59991ee2099eb17f65ceb8 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 26 May 2026 18:40:23 -0500 Subject: [PATCH 25/51] bench(parquet): add short and large string `arrow_writer` benchmarks (#10021) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Which issue does this PR close? Split out of #9972 per [this review comment](https://github.com/apache/arrow-rs/pull/9972#discussion_r3307256819). # Rationale for this change #9972 makes the parquet writer's mini-batch sizing byte-budget aware so large variable-width values don't produce oversized data pages. To measure that change against a stable baseline — and in particular to see the difference in the large-string case — these benchmarks belong on `main` first. # What changes are included in this PR? Adds two BYTE_ARRAY write cases to the `arrow_writer` criterion bench: - **`short_string_non_null`** — 1M fixed-width 8-byte strings. The small-value hot path, where byte-budget-based sub-batch sizing should always resolve to the full chunk (no granular splitting, no regression). - **`large_string_non_null`** — 1024 × 256 KiB strings (256 MiB total). The large-value case: with the default 1 MiB page byte limit each value needs its own page, and a `write_batch_size` of 1024 would otherwise buffer all 256 MiB before the post-write size check runs. No library code changes — benchmarks only. # Are there any user-facing changes? No. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.7 (1M context) --- parquet/benches/arrow_writer.rs | 35 ++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/parquet/benches/arrow_writer.rs b/parquet/benches/arrow_writer.rs index 9b22bb04b3e7..3636cb040253 100644 --- a/parquet/benches/arrow_writer.rs +++ b/parquet/benches/arrow_writer.rs @@ -32,7 +32,7 @@ use std::sync::Arc; use arrow::datatypes::*; use arrow::util::bench_util::{create_f16_array, create_f32_array, create_f64_array}; use arrow::{record_batch::RecordBatch, util::data_gen::*}; -use arrow_array::RecordBatchOptions; +use arrow_array::{RecordBatchOptions, StringArray}; use parquet::errors::Result; use parquet::file::properties::{CdcOptions, WriterProperties, WriterVersion}; @@ -100,6 +100,29 @@ fn create_string_bench_batch( )?) } +/// 1 M short, fixed-width 8-byte strings. Exercises the BYTE_ARRAY hot path +/// for the case where individual values are small enough that the byte-budget +/// based sub-batch sizing in `write_batch_internal` should always resolve to +/// the full chunk (no granular splitting, no regression vs. current behavior). +fn create_short_string_bench_batch(size: usize) -> Result { + let array = Arc::new(StringArray::from_iter_values( + (0..size).map(|i| format!("{i:08}")), + )) as _; + Ok(RecordBatch::try_from_iter([("col", array)])?) +} + +/// `size` rows of `value_size`-byte strings. Exercises the BYTE_ARRAY path +/// where individual values are large enough that batching the default +/// `write_batch_size` of them would blow the page byte limit by orders of +/// magnitude — the case the page-size fix targets. +fn create_large_string_bench_batch(size: usize, value_size: usize) -> Result { + let value = "x".repeat(value_size); + let array = Arc::new(StringArray::from_iter_values( + (0..size).map(|_| value.as_str()), + )) as _; + Ok(RecordBatch::try_from_iter([("col", array)])?) +} + fn create_string_and_binary_view_bench_batch( size: usize, null_density: f32, @@ -392,6 +415,16 @@ fn create_batches() -> Vec<(&'static str, RecordBatch)> { let batch = create_string_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(); batches.push(("string", batch)); + let batch = create_short_string_bench_batch(BATCH_SIZE).unwrap(); + batches.push(("short_string_non_null", batch)); + + // 1024 rows × 256 KiB = 256 MiB total. With the default 1 MiB page byte + // limit, this is the case where the page-size fix kicks in: each value + // needs its own page, and `write_batch_size = 1024` would otherwise + // buffer all 256 MiB before the post-write check runs. + let batch = create_large_string_bench_batch(1024, 256 * 1024).unwrap(); + batches.push(("large_string_non_null", batch)); + let batch = create_string_and_binary_view_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(); batches.push(("string_and_binary_view", batch)); From 40ebcf03a12cd5934d4b9cf6ccd6e970eb658366 Mon Sep 17 00:00:00 2001 From: RIchard Baah <137434454+Rich-T-kid@users.noreply.github.com> Date: Wed, 27 May 2026 03:37:42 -0400 Subject: [PATCH 26/51] benchmarks for writing REE arrays to parquet (#9936) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Which issue does this PR close? - Closes #9935. # Rationale for this change there is no way to currently tell which approach to writing out REE columns to parquet is more performant. This PR aims to solve that. # What changes are included in this PR? Added a `create_string_ree_bench_batch()` function that builds record batches of REE data — it plugs into the existing benchmark structure. For controlling the shape of the generated REE arrays, I currently have two constants, `MIN_RUN` and `MAX_RUN`, that bound the run length. The intent is to let benchmarks cover long uniform runs as well as shorter / more sparse data, rather than only one shape. An alternative would be a small params struct with defaults that callers can override — happy to switch to that if it's preferred, but that would require changing other callsites # Are these changes tested? yes # Are there any user-facing changes? no --- arrow/src/util/data_gen.rs | 83 +++++++++++++++++++++++++++++++++ parquet/benches/arrow_writer.rs | 30 ++++++++++++ 2 files changed, 113 insertions(+) diff --git a/arrow/src/util/data_gen.rs b/arrow/src/util/data_gen.rs index e54ab3499472..a5a0647aa877 100644 --- a/arrow/src/util/data_gen.rs +++ b/arrow/src/util/data_gen.rs @@ -178,6 +178,9 @@ pub fn create_random_array( Map(_, _) => create_random_map_array(field, size, null_density, true_density)?, Decimal128(_, _) => create_random_decimal_array(field, size, null_density)?, Decimal256(_, _) => create_random_decimal_array(field, size, null_density)?, + RunEndEncoded(index, value) => { + create_random_run_end_encoded_array(index, value, size, null_density, true_density)? + } other => { return Err(ArrowError::NotYetImplemented(format!( "Generating random arrays not yet implemented for {other:?}" @@ -230,6 +233,62 @@ fn create_random_decimal_array(field: &Field, size: usize, null_density: f32) -> ))), } } +#[inline] +fn create_random_run_end_encoded_array( + index: &Field, + value: &Field, + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + const MIN_RUN: usize = 8; + const MAX_RUN: usize = 32; + + let mut rng = seedable_rng(); + let mut run_lengths: Vec = Vec::new(); + let mut remaining = size; + while remaining > 0 { + let len = rng.random_range(MIN_RUN..=MAX_RUN).min(remaining); + run_lengths.push(len); + remaining -= len; + } + let num_runs = run_lengths.len(); + + let mut cumulative: i64 = 0; + let run_ends_i64: Vec = run_lengths + .iter() + .map(|&l| { + cumulative += l as i64; + cumulative + }) + .collect(); + + let values = create_random_array(value, num_runs, null_density, true_density)?; + + match index.data_type() { + DataType::Int16 => { + let run_ends: Int16Array = run_ends_i64.iter().map(|&v| v as i16).collect(); + Ok(Arc::new(RunArray::::try_new( + &run_ends, &values, + )?)) + } + DataType::Int32 => { + let run_ends: Int32Array = run_ends_i64.iter().map(|&v| v as i32).collect(); + Ok(Arc::new(RunArray::::try_new( + &run_ends, &values, + )?)) + } + DataType::Int64 => { + let run_ends: Int64Array = run_ends_i64.iter().copied().collect(); + Ok(Arc::new(RunArray::::try_new( + &run_ends, &values, + )?)) + } + other => Err(ArrowError::InvalidArgumentError(format!( + "Unsupported run-ends type for REE: {other:?}" + ))), + } +} #[inline] fn create_random_list_array( @@ -648,6 +707,30 @@ mod tests { assert_eq!(col_d_y.null_count(), 0); } + #[test] + fn test_create_run_end_encoded_array() { + let size = 1000; + let ree_field = Field::new( + "ree", + DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", DataType::Utf8, true)), + ), + false, + ); + + let array = create_random_array(&ree_field, size, 0.25, 0.0).unwrap(); + assert_eq!(array.len(), size); + + let ree = array.as_run::(); + let run_ends = ree.run_ends().values(); + let num_runs = run_ends.len(); + + assert_eq!(*run_ends.last().unwrap() as usize, size); + + assert_eq!(ree.values().len(), num_runs); + } + #[test] fn test_create_list_array_nested_nullability() { let list_field = Field::new_list( diff --git a/parquet/benches/arrow_writer.rs b/parquet/benches/arrow_writer.rs index 3636cb040253..80d3e7144b19 100644 --- a/parquet/benches/arrow_writer.rs +++ b/parquet/benches/arrow_writer.rs @@ -159,6 +159,30 @@ fn create_string_dictionary_bench_batch( true_density, )?) } +// commenting out until implementation of RunEndEncoded is complete. See https://github.com/apache/arrow-rs/pull/9936#discussion_r3242936421 +#[allow(dead_code)] +fn create_ree_bench_batch( + value_dt: DataType, + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![Field::new( + "_1", + DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", value_dt, true)), + ), + true, + )]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} fn create_string_bench_batch_non_null( size: usize, @@ -434,6 +458,12 @@ fn create_batches() -> Vec<(&'static str, RecordBatch)> { let batch = create_string_bench_batch_non_null(BATCH_SIZE, 0.25, 0.75).unwrap(); batches.push(("string_non_null", batch)); + //let batch = create_ree_bench_batch(DataType::Utf8, BATCH_SIZE, 0.25, 0.75).unwrap(); + //batches.push(("string_ree", batch)); + + //let batch = create_ree_bench_batch(DataType::Int32, BATCH_SIZE, 0.25, 0.75).unwrap(); + //batches.push(("int32_ree", batch)); + let batch = create_float_bench_batch_with_nans(BATCH_SIZE, 0.5).unwrap(); batches.push(("float_with_nans", batch)); From 2eeb805b8a8b8ca67788917ec2f5220eb3e6f958 Mon Sep 17 00:00:00 2001 From: RIchard Baah <137434454+Rich-T-kid@users.noreply.github.com> Date: Wed, 27 May 2026 03:39:21 -0400 Subject: [PATCH 27/51] Implement AnyRee (#9959) # Which issue does this PR close? closes #9909. - Closes #9909. # Rationale for this change makes the API simpler to work with & less code duplication # What changes are included in this PR? Replace the per-key-type RunEndEncoded match arms in length/bit_length (arrow-string) and date_part (arrow-arith) with a single dispatch through the new `AsArray::as_any_ree_opt/as_any_ree` returning &dyn AnyRunEndArray, mirroring the existing dictionary handling. This removes the now-unused `ree_map!` macro, leaving one trait-object code path for all Int16/Int32/Int64 run-end types. # Are these changes tested? yes # Are there any user-facing changes? no --- arrow-arith/src/temporal.rs | 16 ++++----- arrow-array/src/array/run_array.rs | 57 +++++++++++++++++++++++++++--- arrow-array/src/cast.rs | 20 +++++++++++ arrow-string/src/length.rs | 27 +++++--------- 4 files changed, 87 insertions(+), 33 deletions(-) diff --git a/arrow-arith/src/temporal.rs b/arrow-arith/src/temporal.rs index 301ad172dabc..769d309f5812 100644 --- a/arrow-arith/src/temporal.rs +++ b/arrow-arith/src/temporal.rs @@ -24,7 +24,6 @@ use arrow_array::cast::AsArray; use cast::as_primitive_array; use chrono::{Datelike, TimeZone, Timelike, Utc}; -use arrow_array::ree_map; use arrow_array::temporal_conversions::{ MICROSECONDS, MICROSECONDS_IN_DAY, MILLISECONDS, MILLISECONDS_IN_DAY, NANOSECONDS, NANOSECONDS_IN_DAY, SECONDS_IN_DAY, date32_to_datetime, date64_to_datetime, @@ -253,15 +252,12 @@ pub fn date_part(array: &dyn Array, part: DatePart) -> Result match k.data_type() { - DataType::Int16 => ree_map!(array, Int16Type, |a| date_part(a, part)), - DataType::Int32 => ree_map!(array, Int32Type, |a| date_part(a, part)), - DataType::Int64 => ree_map!(array, Int64Type, |a| date_part(a, part)), - _ => Err(ArrowError::InvalidArgumentError(format!( - "Invalid run-end type: {:?}", - k.data_type() - ))), - }, + DataType::RunEndEncoded(_, _) => { + let array = array.as_any_ree(); + let values = date_part(array.values(), part)?; + let new_array = array.with_values(values); + Ok(new_array) + } t => return_compute_error_with!(format!("{part} does not support"), t), ) } diff --git a/arrow-array/src/array/run_array.rs b/arrow-array/src/array/run_array.rs index 02bc730b32df..09fb2998a22b 100644 --- a/arrow-array/src/array/run_array.rs +++ b/arrow-array/src/array/run_array.rs @@ -245,13 +245,20 @@ impl RunArray { /// assert_eq!(new_run_array.run_ends().values(), &[2, 3, 5]); /// ``` pub fn with_values(&self, values: ArrayRef) -> Self { - assert_eq!(values.len(), self.values().len()); + assert_eq!(values.len(), self.values.len()); let (run_ends_field, values_field) = match &self.data_type { - DataType::RunEndEncoded(r, v) => (r, v), + DataType::RunEndEncoded(r, v) => { + let new_v = Arc::new(Field::new( + v.name(), + values.data_type().clone(), + v.is_nullable(), + )); + (r, new_v) + } _ => unreachable!("RunArray should have type RunEndEncoded"), }; - let data_type = - DataType::RunEndEncoded(Arc::clone(run_ends_field), Arc::clone(values_field)); + let data_type = DataType::RunEndEncoded(Arc::clone(run_ends_field), values_field); + Self { data_type, run_ends: self.run_ends.clone(), @@ -781,6 +788,28 @@ where RunArrayIter::new(self) } } +/// An array that can be downcast to a [`RunArray`] of any run end type and any value type. +/// +/// This can be used to efficiently implement kernels for all possible run end +/// types without needing to create specialized implementations for each key type. +pub trait AnyRunEndArray: Array { + /// Returns the values of this array. + fn values(&self) -> &Arc; + + /// Returns a new run-end encoded array with the given values, preserving the + /// existing run ends. + fn with_values(&self, values: ArrayRef) -> ArrayRef; +} + +impl AnyRunEndArray for RunArray { + fn values(&self) -> &Arc { + &self.values + } + + fn with_values(&self, values: ArrayRef) -> ArrayRef { + Arc::new(RunArray::::with_values(self, values)) + } +} #[cfg(test)] mod tests { @@ -789,6 +818,7 @@ mod tests { use rand::seq::SliceRandom; use super::*; + use crate::Int64Array; use crate::builder::PrimitiveRunBuilder; use crate::cast::AsArray; use crate::new_empty_array; @@ -1055,6 +1085,25 @@ mod tests { let expected = ArrowError::InvalidArgumentError("The run_ends array length should be the same as values array length. Run_ends array length is 3, values array length is 4".to_string()); assert_eq!(expected.to_string(), actual.err().unwrap().to_string()); } + #[test] + fn test_run_array_with_values_changes_value_type() { + let values = StringArray::from(vec!["foo", "bar", "baz"]); + let run_ends: Int32Array = [Some(1), Some(2), Some(3)].into_iter().collect(); + let ree = RunArray::::try_new(&run_ends, &values).unwrap(); + + let new_values = Int64Array::from(vec![10, 20, 30]); + let result = ree.with_values(Arc::new(new_values)); + + match result.data_type() { + DataType::RunEndEncoded(_, v) => { + assert_eq!(v.data_type(), &DataType::Int64); + } + other => panic!("expected RunEndEncoded, got {other:?}"), + } + + assert_eq!(result.values().data_type(), &DataType::Int64); + assert_eq!(result.values().len(), 3); + } #[test] fn test_run_array_run_ends_with_null() { diff --git a/arrow-array/src/cast.rs b/arrow-array/src/cast.rs index d6cc242e0267..d30de5906f85 100644 --- a/arrow-array/src/cast.rs +++ b/arrow-array/src/cast.rs @@ -986,6 +986,14 @@ pub trait AsArray: private::Sealed { fn as_any_dictionary(&self) -> &dyn AnyDictionaryArray { self.as_any_dictionary_opt().expect("any dictionary array") } + + /// Downcasts this to a [`AnyRunEndArray`] returning `None` if not possible + fn as_any_ree_opt(&self) -> Option<&dyn AnyRunEndArray>; + + /// Downcasts this to a [`AnyRunEndArray`] panicking if not possible + fn as_any_ree(&self) -> &dyn AnyRunEndArray { + self.as_any_ree_opt().expect("any run end array") + } } impl private::Sealed for dyn Array + '_ {} @@ -1049,6 +1057,14 @@ impl AsArray for dyn Array + '_ { _ => None } } + + fn as_any_ree_opt(&self) -> Option<&dyn AnyRunEndArray> { + let array = self; + downcast_run_array! { + array => Some(array), + _ => None + } + } } impl private::Sealed for ArrayRef {} @@ -1105,6 +1121,10 @@ impl AsArray for ArrayRef { self.as_ref().as_any_dictionary_opt() } + fn as_any_ree_opt(&self) -> Option<&dyn AnyRunEndArray> { + self.as_ref().as_any_ree_opt() + } + fn as_run_opt(&self) -> Option<&RunArray> { self.as_ref().as_run_opt() } diff --git a/arrow-string/src/length.rs b/arrow-string/src/length.rs index feefe1247e2c..99a9bd69a6c3 100644 --- a/arrow-string/src/length.rs +++ b/arrow-string/src/length.rs @@ -17,7 +17,6 @@ //! Defines kernel for length of string arrays and binary arrays -use arrow_array::ree_map; use arrow_array::*; use arrow_array::{cast::AsArray, types::*}; use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer}; @@ -59,6 +58,10 @@ pub fn length(array: &dyn Array) -> Result { let lengths = length(d.values().as_ref())?; return Ok(d.with_values(lengths)); } + if let Some(ree) = array.as_any_ree_opt() { + let lengths = length(ree.values())?; + return Ok(ree.with_values(lengths)); + } match array.data_type() { DataType::List(_) => { let list = array.as_list::(); @@ -117,15 +120,6 @@ pub fn length(array: &dyn Array) -> Result { list.nulls().cloned(), )?)) } - DataType::RunEndEncoded(k, _) => match k.data_type() { - DataType::Int16 => ree_map!(array, Int16Type, length), - DataType::Int32 => ree_map!(array, Int32Type, length), - DataType::Int64 => ree_map!(array, Int64Type, length), - _ => Err(ArrowError::InvalidArgumentError(format!( - "Invalid run-end type: {:?}", - k.data_type() - ))), - }, other => Err(ArrowError::ComputeError(format!( "length not supported for {other:?}" ))), @@ -144,6 +138,10 @@ pub fn bit_length(array: &dyn Array) -> Result { let lengths = bit_length(d.values().as_ref())?; return Ok(d.with_values(lengths)); } + if let Some(ree) = array.as_any_ree_opt() { + let lengths = bit_length(ree.values())?; + return Ok(ree.with_values(lengths)); + } match array.data_type() { DataType::Utf8 => { @@ -190,15 +188,6 @@ pub fn bit_length(array: &dyn Array) -> Result { array.nulls().cloned(), )?)) } - DataType::RunEndEncoded(k, _) => match k.data_type() { - DataType::Int16 => ree_map!(array, Int16Type, bit_length), - DataType::Int32 => ree_map!(array, Int32Type, bit_length), - DataType::Int64 => ree_map!(array, Int64Type, bit_length), - _ => Err(ArrowError::InvalidArgumentError(format!( - "Invalid run-end type: {:?}", - k.data_type() - ))), - }, other => Err(ArrowError::ComputeError(format!( "bit_length not supported for {other:?}" ))), From e470187b93b13bf9821e67d5b2348a1d89612a39 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Wed, 27 May 2026 16:16:41 -0400 Subject: [PATCH 28/51] fix: Reject empty strings when casting strings to decimal (#10010) # Which issue does this PR close? - Closes #10009. # Rationale for this change When casting string values to decimal, `parse_string_to_decimal_native` treated empty strings and whitespace-only strings as valid input, resulting in a decimal with a value of 0. This is inconsistent with `parse_decimal` and how parsing and string -> numeric casts work for floating point types: in all of those cases, empty strings and whitespace-only strings are rejected. # What changes are included in this PR? * Change `parse_string_to_decimal_native` to reject empty strings and whitespace-only strings * Add test coverage # Are these changes tested? Yes, new tests added. # Are there any user-facing changes? Yes, this changes the behavior of string -> decimal casts. The previous behavior is (IMO) clearly incorrect but it is possible that some user code relies upon it. --- arrow-cast/src/cast/decimal.rs | 17 +++++++++++++++++ arrow-cast/src/cast/mod.rs | 21 +++++++++++++++------ 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/arrow-cast/src/cast/decimal.rs b/arrow-cast/src/cast/decimal.rs index 789dcea89ed1..9d1465567e0d 100644 --- a/arrow-cast/src/cast/decimal.rs +++ b/arrow-cast/src/cast/decimal.rs @@ -559,6 +559,12 @@ where let integers = first_part; let decimals = if parts.len() == 2 { parts[1] } else { "" }; + if integers.is_empty() && decimals.is_empty() { + return Err(ArrowError::InvalidArgumentError(format!( + "Invalid decimal format: {value_str:?}" + ))); + } + if !integers.is_empty() && !integers.as_bytes()[0].is_ascii_digit() { return Err(ArrowError::InvalidArgumentError(format!( "Invalid decimal format: {value_str:?}" @@ -973,6 +979,17 @@ mod tests { parse_string_to_decimal_native::("123.4567891", 5)?, 12345679_i128 ); + + for value in ["", " ", ".", "+", "-", "+.", "-."] { + assert!( + parse_string_to_decimal_native::(value, 2).is_err(), + "expected {value:?} to fail parsing as Decimal128" + ); + assert!( + parse_string_to_decimal_native::(value, 2).is_err(), + "expected {value:?} to fail parsing as Decimal256" + ); + } Ok(()) } diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs index 12442bd9fd30..67da85b8c1d6 100644 --- a/arrow-cast/src/cast/mod.rs +++ b/arrow-cast/src/cast/mod.rs @@ -10775,8 +10775,8 @@ mod tests { assert_eq!("0.12", decimal_arr.value_as_string(7)); assert_eq!("12.23", decimal_arr.value_as_string(8)); assert!(decimal_arr.is_null(9)); - assert_eq!("0.00", decimal_arr.value_as_string(10)); - assert_eq!("0.00", decimal_arr.value_as_string(11)); + assert!(decimal_arr.is_null(10)); + assert!(decimal_arr.is_null(11)); assert!(decimal_arr.is_null(12)); assert_eq!("-1.23", decimal_arr.value_as_string(13)); assert_eq!("-1.24", decimal_arr.value_as_string(14)); @@ -10815,8 +10815,8 @@ mod tests { assert_eq!("0.123", decimal_arr.value_as_string(7)); assert_eq!("12.234", decimal_arr.value_as_string(8)); assert!(decimal_arr.is_null(9)); - assert_eq!("0.000", decimal_arr.value_as_string(10)); - assert_eq!("0.000", decimal_arr.value_as_string(11)); + assert!(decimal_arr.is_null(10)); + assert!(decimal_arr.is_null(11)); assert!(decimal_arr.is_null(12)); assert_eq!("-1.235", decimal_arr.value_as_string(13)); assert_eq!("-1.236", decimal_arr.value_as_string(14)); @@ -10880,8 +10880,8 @@ mod tests { let test_cases = [ (None, None), - // (Some(""), None), - // (Some(" "), None), + (Some(""), None), + (Some(" "), None), (Some("0"), Some("0")), (Some("000.000"), Some("0")), (Some("12345"), Some("12345")), @@ -11029,6 +11029,15 @@ mod tests { .to_string() .contains("Cannot cast string '. 0.123' to value of Decimal128(38, 10) type") ); + + let str_array = StringArray::from(vec![""]); + let array = Arc::new(str_array) as ArrayRef; + let casted_err = cast_with_options(&array, &output_type, &option).unwrap_err(); + assert!( + casted_err + .to_string() + .contains("Cannot cast string '' to value of Decimal128(38, 10) type") + ); } fn test_cast_string_to_decimal128_overflow(overflow_array: ArrayRef) { From 1377761779afb1ce70a7a9a9038a308d2ff1ab88 Mon Sep 17 00:00:00 2001 From: eunsang Date: Fri, 29 May 2026 03:05:56 +0900 Subject: [PATCH 29/51] Validate FIXED_LEN_BYTE_ARRAY length for DECIMAL and INTERVAL types (#9985) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes #9984. ## Rationale for this change `from_fixed_len_byte_array` in `parquet/src/arrow/schema/primitive.rs` does not validate `type_length`. While `PrimitiveTypeBuilder::build()` enforces these constraints during schema construction (`schema/types.rs:477` for INTERVAL, `:565-580` for DECIMAL), schemas decoded directly from Thrift bypass that validation path entirely. As a result: - `DECIMAL` with a `type_length` outside `1..=32` was silently routed through `Decimal128` / `Decimal256` using invalid parameters. - `INTERVAL` with a `type_length != 12` silently returned `Interval(DayTime)` regardless. The same function already rejects `FLOAT16` when `type_length != 2`. This PR mirrors that pattern for DECIMAL and INTERVAL, closing the TODO introduced in #1682. ## What changes are included in this PR? - Added a `check_decimal_length` helper to reject `type_length` values outside `1..=32` for both `LogicalType::Decimal` and `ConvertedType::DECIMAL`. - Added an inline `type_length == 12` check for `ConvertedType::INTERVAL`. ## Are these changes tested? Yes. Added five new tests in `parquet/src/arrow/schema/primitive.rs::tests` covering: - Invalid lengths (`{-1, 0, 33}` for DECIMAL, `{0, 11, 13}` for INTERVAL) - Valid lengths (16 → `Decimal128`, 32 → `Decimal256`, 12 → `Interval(DayTime)`) To exercise the reader-side check, the tests construct a valid `Type::PrimitiveType` via the builder and directly modify the `type_length` on the resulting enum, simulating a malformed schema decoded from Thrift. ## Are there any user-facing changes? No public API changes. The only behavior change is on the reader side: schemas with an out-of-range `type_length` for DECIMAL or INTERVAL will now return a `ParquetError::General` instead of silently producing a mismatched Arrow type. --- parquet/src/arrow/schema/primitive.rs | 132 +++++++++++++++++++++++++- 1 file changed, 131 insertions(+), 1 deletion(-) diff --git a/parquet/src/arrow/schema/primitive.rs b/parquet/src/arrow/schema/primitive.rs index 2272014a9361..c70885355d74 100644 --- a/parquet/src/arrow/schema/primitive.rs +++ b/parquet/src/arrow/schema/primitive.rs @@ -170,6 +170,16 @@ fn decimal_256_type(scale: i32, precision: i32) -> Result { Ok(DataType::Decimal256(precision, scale)) } +#[allow(clippy::manual_range_contains)] +fn check_decimal_length(type_length: i32) -> Result<()> { + if type_length < 1 || type_length > 32 { + return Err(ParquetError::General(format!( + "DECIMAL must be a Fixed Length Byte Array with length 1 to 32, got {type_length}" + ))); + } + Ok(()) +} + fn from_int32(info: &BasicTypeInfo, scale: i32, precision: i32) -> Result { match (info.logical_type_ref(), info.converted_type()) { (None, ConvertedType::NONE) => Ok(DataType::Int32), @@ -293,9 +303,10 @@ fn from_fixed_len_byte_array( precision: i32, type_length: i32, ) -> Result { - // TODO: This should check the type length for the decimal and interval types match (info.logical_type_ref(), info.converted_type()) { (Some(LogicalType::Decimal(decimal)), _) => { + check_decimal_length(type_length)?; + // lengths 1..=16 map to Decimal128, 17..=32 to Decimal256 if type_length <= 16 { decimal_128_type(decimal.scale, decimal.precision) } else { @@ -303,6 +314,7 @@ fn from_fixed_len_byte_array( } } (None, ConvertedType::DECIMAL) => { + check_decimal_length(type_length)?; if type_length <= 16 { decimal_128_type(scale, precision) } else { @@ -310,6 +322,11 @@ fn from_fixed_len_byte_array( } } (None, ConvertedType::INTERVAL) => { + if type_length != 12 { + return Err(ParquetError::General(format!( + "INTERVAL must be a Fixed Length Byte Array with length 12, got {type_length}" + ))); + } // There is currently no reliable way of determining which IntervalUnit // to return. Thus without the original Arrow schema, the results // would be incorrect if all 12 bytes of the interval are populated @@ -328,3 +345,116 @@ fn from_fixed_len_byte_array( _ => Ok(DataType::FixedSizeBinary(type_length)), } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::basic::{DecimalType, Repetition}; + use crate::schema::types::Type; + + // The PrimitiveTypeBuilder rejects bad lengths at construction. To exercise + // the reader-side checks, build a valid type then overwrite its type_length, + // simulating a schema decoded from a file that wasn't produced via the builder. + fn with_type_length(ty: Type, type_length: i32) -> Type { + match ty { + Type::PrimitiveType { + basic_info, + physical_type, + precision, + scale, + .. + } => Type::PrimitiveType { + basic_info, + physical_type, + type_length, + precision, + scale, + }, + _ => unreachable!(), + } + } + + fn flba_decimal_logical(type_length: i32) -> Type { + let valid = Type::primitive_type_builder("c", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_logical_type(Some(LogicalType::Decimal(DecimalType { + precision: 5, + scale: 2, + }))) + .with_length(16) + .with_precision(5) + .with_scale(2) + .build() + .unwrap(); + with_type_length(valid, type_length) + } + + fn flba_decimal_converted(type_length: i32) -> Type { + let valid = Type::primitive_type_builder("c", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_length(16) + .with_precision(5) + .with_scale(2) + .build() + .unwrap(); + with_type_length(valid, type_length) + } + + fn flba_interval(type_length: i32) -> Type { + let valid = Type::primitive_type_builder("c", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INTERVAL) + .with_length(12) + .build() + .unwrap(); + with_type_length(valid, type_length) + } + + fn assert_err_contains(ty: &Type, needle: &str) { + let err = convert_primitive(ty, None).expect_err("expected an error"); + let msg = err.to_string(); + assert!(msg.contains(needle), "expected {needle:?} in error: {msg}"); + } + + #[test] + fn decimal_logical_rejects_invalid_length() { + for bad in [-1, 0, 33] { + assert_err_contains(&flba_decimal_logical(bad), "DECIMAL"); + } + } + + #[test] + fn decimal_converted_rejects_invalid_length() { + for bad in [-1, 0, 33] { + assert_err_contains(&flba_decimal_converted(bad), "DECIMAL"); + } + } + + #[test] + fn decimal_accepts_valid_lengths() { + assert!(matches!( + convert_primitive(&flba_decimal_logical(16), None).unwrap(), + DataType::Decimal128(_, _) + )); + assert!(matches!( + convert_primitive(&flba_decimal_logical(32), None).unwrap(), + DataType::Decimal256(_, _) + )); + } + + #[test] + fn interval_rejects_wrong_length() { + for bad in [0, 11, 13] { + assert_err_contains(&flba_interval(bad), "INTERVAL"); + } + } + + #[test] + fn interval_accepts_length_12() { + assert_eq!( + convert_primitive(&flba_interval(12), None).unwrap(), + DataType::Interval(IntervalUnit::DayTime) + ); + } +} From 9450f1e9cc17061049309e98b0cbccca91e7dfb7 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Fri, 29 May 2026 15:42:00 -0400 Subject: [PATCH 30/51] Call `align_buffers()` in `from_ffi`, remove redundant call from `arrow-pyarrow` (#10030) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Which issue does this PR close? - Closes #10028. # Rationale for this change `from_ffi` / `from_ffi_and_data_type` (and therefore `ArrowArrayStreamReader`) panic inside `ScalarBuffer::::from` when an FFI producer hands over a `Decimal128` buffer that is 8-byte aligned but not 16-byte aligned. The producer is spec-conformant — the C Data Interface only recommends 8-byte alignment — but `align_of::() == 16` since Rust 1.77 on x86 (always on ARM), so arrow-rs's typed arrays require 16. JVM producers like arrow-java's `NettyAllocationManager` hit this regularly. The IPC reader already handles this by calling `ArrayData::align_buffers()` on import (default of `IpcReadOptions::require_alignment`, see #5554), and `arrow-pyarrow` was patched the same way for #6471 / apache/arrow#43552. The C Data Interface entry points were the missing piece. # What changes are included in this PR? - `arrow::ffi::from_ffi` and `from_ffi_and_data_type`: call `data.align_buffers()` after `consume()`. No-op when buffers are already aligned; depends on #6462 making `align_buffers` recursive over child data. - `arrow-pyarrow`: drop the now-redundant `array_data.align_buffers()` call; it's covered by `from_ffi`. # Are these changes tested? Yes. New regression test `test_decimal128_under_aligned_round_trip` in `arrow-array/src/ffi.rs` constructs an 8-aligned-not-16-aligned `Decimal128` buffer via `Buffer::from_vec(...).slice(8)`, imports through `from_ffi`, and asserts the resulting `Decimal128Array` values are correct. The test panics without the fix with the exact error from #10028. # Are there any user-facing changes? No API changes. Behavior change: `from_ffi` / `from_ffi_and_data_type` (and `ArrowArrayStreamReader::next`) now silently realign under-aligned buffers instead of panicking. Already-aligned producers are unaffected; misaligned producers that previously panicked now succeed with a one-time copy of the offending buffer. --- arrow-array/src/ffi.rs | 52 ++++++++++++++++++++++++++++++++++++++-- arrow-pyarrow/src/lib.rs | 7 +----- 2 files changed, 51 insertions(+), 8 deletions(-) diff --git a/arrow-array/src/ffi.rs b/arrow-array/src/ffi.rs index f50dd3420baa..59b9f6b3b7b5 100644 --- a/arrow-array/src/ffi.rs +++ b/arrow-array/src/ffi.rs @@ -281,7 +281,14 @@ pub unsafe fn from_ffi(array: FFI_ArrowArray, schema: &FFI_ArrowSchema) -> Resul data_type: dt, owner: &array, }; - tmp.consume() + let mut data = tmp.consume()?; + // arrow-rs has stricter alignment requirements than the C Data Interface spec; + // a no-op when buffers are already aligned. Unreachable under + // `cfg(feature = "force_validate")`; tracked in #10034. + // See https://github.com/apache/arrow/issues/43552 and + // https://github.com/apache/arrow-rs/issues/10028 for context. + data.align_buffers(); + Ok(data) } /// Import [ArrayData] from the C Data Interface @@ -299,7 +306,14 @@ pub unsafe fn from_ffi_and_data_type( data_type, owner: &array, }; - tmp.consume() + let mut data = tmp.consume()?; + // arrow-rs has stricter alignment requirements than the C Data Interface spec; + // a no-op when buffers are already aligned. Unreachable under + // `cfg(feature = "force_validate")`; tracked in #10034. + // See https://github.com/apache/arrow/issues/43552 and + // https://github.com/apache/arrow-rs/issues/10028 for context. + data.align_buffers(); + Ok(data) } #[derive(Debug)] @@ -667,6 +681,40 @@ mod tests_to_then_from_ffi { } // case with nulls is tested in the docs, through the example on this module. + #[test] + #[cfg(not(feature = "force_validate"))] + fn test_decimal128_under_aligned_round_trip() -> Result<()> { + // Construct an 8-aligned-but-not-16-aligned i128 data buffer to model + // an FFI producer that only guarantees the C Data Interface's + // recommended 8-byte alignment (e.g. arrow-java). + let aligned = Buffer::from_vec(vec![0_i128, 1_i128, 2_i128]); + let under_aligned = aligned.slice(8); + assert_eq!(under_aligned.as_ptr().align_offset(8), 0); + assert_ne!(under_aligned.as_ptr().align_offset(16), 0); + + // SAFETY: buffer is large enough for 2 i128 elements; misaligned + // input is the condition under test. + let data = unsafe { + ArrayData::builder(DataType::Decimal128(10, 2)) + .len(2) + .add_buffer(under_aligned) + .build_unchecked() + }; + + let schema = FFI_ArrowSchema::try_from(data.data_type()).unwrap(); + let array = FFI_ArrowArray::new(&data); + + let imported = unsafe { from_ffi(array, &schema) }?; + let array = Decimal128Array::from(imported); + + // The little-endian byte layout of [0i128, 1, 2] sliced 8 bytes in + // yields elements `1 << 64` and `2 << 64`. + assert_eq!(array.len(), 2); + assert_eq!(array.value(0), 1_i128 << 64); + assert_eq!(array.value(1), 2_i128 << 64); + Ok(()) + } + #[test] fn test_null_count_handling() { let int32_data = ArrayData::builder(DataType::Int32) diff --git a/arrow-pyarrow/src/lib.rs b/arrow-pyarrow/src/lib.rs index d8f584e396d3..484324665cac 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/arrow-pyarrow/src/lib.rs @@ -353,7 +353,7 @@ impl FromPyArrow for RecordBatch { .pointer_checked(Some(ARROW_ARRAY_CAPSULE_NAME))? .cast::(); let ffi_array = unsafe { FFI_ArrowArray::from_raw(array_ptr.as_ptr()) }; - let mut array_data = + let array_data = unsafe { ffi::from_ffi(ffi_array, schema_ptr.as_ref()) }.map_err(to_py_err)?; if !matches!(array_data.data_type(), DataType::Struct(_)) { return Err(PyTypeError::new_err( @@ -361,11 +361,6 @@ impl FromPyArrow for RecordBatch { )); } let options = RecordBatchOptions::default().with_row_count(Some(array_data.len())); - // Ensure data is aligned (by potentially copying the buffers). - // This is needed because some python code (for example the - // python flight client) produces unaligned buffers - // See https://github.com/apache/arrow/issues/43552 for details - array_data.align_buffers(); let array = StructArray::from(array_data); // StructArray does not embed metadata from schema. We need to override // the output schema with the schema from the capsule. From 511ad068ae2b7f511e16215d9d3c96fb0f334f2c Mon Sep 17 00:00:00 2001 From: Dan Mattheiss Date: Sat, 30 May 2026 23:18:11 -0400 Subject: [PATCH 31/51] bench(parquet): add Sbbf check/insert benchmarks (#10041) Adds `bench_check` and `bench_insert` benchmarks for`Sbbf::{check,insert}`. Originally benchmarks were part of #10011 but were split out to follow Contributing guidelines # Are these changes tested? Benchmarks compiled and ran using `cargo bench -p parquet --bench bloom_filter`. # Are there any user-facing changes? No. --- parquet/benches/bloom_filter.rs | 123 +++++++++++++++++++++++++++++++- 1 file changed, 120 insertions(+), 3 deletions(-) diff --git a/parquet/benches/bloom_filter.rs b/parquet/benches/bloom_filter.rs index ca4f900067f8..cb2bfb6b4a2c 100644 --- a/parquet/benches/bloom_filter.rs +++ b/parquet/benches/bloom_filter.rs @@ -15,8 +15,12 @@ // specific language governing permissions and limitations // under the License. -use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use std::hint; + +use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; use parquet::bloom_filter::Sbbf; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; /// Build a bloom filter sized for `initial_ndv` at `fpp`, insert `num_values` distinct values, /// and return it ready for folding. @@ -46,7 +50,7 @@ fn bench_fold_to_target_fpp(c: &mut Criterion) { f.fold_to_target_fpp(fpp); f }, - criterion::BatchSize::SmallInput, + BatchSize::SmallInput, ); }); } @@ -104,10 +108,123 @@ fn bench_insert_only(c: &mut Criterion) { group.finish(); } +/// Benchmark `Sbbf::insert` across the same three cache regimes as +/// `bench_check`. The filter is allocated once per regime and reused +/// across iterations — bloom inserts are idempotent on identical +/// input, so this measures the pure insert kernel cost (hash + mask + +/// load + OR + store) without per-iteration allocation noise. +fn bench_insert(c: &mut Criterion) { + let regimes: [(&str, usize); 3] = [ + ("s_128KiB", 128 * 1024), + ("m_2MiB", 2 * 1024 * 1024), + ("l_32MiB", 32 * 1024 * 1024), + ]; + const NUM_INSERTS: usize = 50_000; + + let mut group = c.benchmark_group("insert"); + + for (label, num_bytes) in regimes { + let mut rng = StdRng::seed_from_u64(0xC0FFEE); + let keys: Vec<[u8; 16]> = (0..NUM_INSERTS) + .map(|_| { + let mut k = [0u8; 16]; + rng.fill(&mut k); + k + }) + .collect(); + + let mut filter = Sbbf::new_with_num_of_bytes(num_bytes); + + group.throughput(Throughput::Elements(NUM_INSERTS as u64)); + group.bench_with_input(BenchmarkId::new("ins", label), &keys, |b, keys| { + b.iter(|| { + for k in keys { + filter.insert(hint::black_box(k.as_slice())); + } + }); + }); + } + group.finish(); +} + +/// Benchmark `Sbbf::check` across three cache regimes and both miss-heavy +/// (the common case for Parquet row-group skipping) and hit-heavy workloads. +/// +/// The three filter sizes span the cache hierarchy on a typical server: +/// 128 KB fits in L2, 2 MB lives in L3, 32 MB spills out of L3. +fn bench_check(c: &mut Criterion) { + let regimes: [(&str, usize, usize); 3] = [ + ("s_128KiB", 128 * 1024, 10_000), + ("m_2MiB", 2 * 1024 * 1024, 200_000), + ("l_32MiB", 32 * 1024 * 1024, 3_000_000), + ]; + const NUM_QUERIES: usize = 50_000; + + let mut group = c.benchmark_group("check"); + + for (label, num_bytes, ndv) in regimes { + let mut rng = StdRng::seed_from_u64(0xC0FFEE); + + // Clear the high bit of byte 0 on inserted keys so the miss set below + // — which forces the same bit to 1 — is genuinely disjoint by + // construction (not just disjoint at birthday-paradox probability). + let mut filter = Sbbf::new_with_num_of_bytes(num_bytes); + let mut inserted: Vec<[u8; 16]> = Vec::with_capacity(ndv); + for _ in 0..ndv { + let mut k = [0u8; 16]; + rng.fill(&mut k); + k[0] &= 0x7f; + filter.insert(k.as_slice()); + inserted.push(k); + } + + // Disjoint miss set: high bit of byte 0 is 1 here and 0 in `inserted`, + // so no miss key can equal an inserted key. + let miss_keys: Vec<[u8; 16]> = (0..NUM_QUERIES) + .map(|_| { + let mut k = [0u8; 16]; + rng.fill(&mut k); + k[0] |= 0x80; + k + }) + .collect(); + let hit_keys: Vec<[u8; 16]> = (0..NUM_QUERIES) + .map(|_| inserted[rng.random_range(0..inserted.len())]) + .collect(); + + group.throughput(Throughput::Elements(NUM_QUERIES as u64)); + group.bench_with_input(BenchmarkId::new("miss", label), &miss_keys, |b, keys| { + b.iter(|| { + let mut found = 0u64; + for k in keys { + if hint::black_box(filter.check(k.as_slice())) { + found += 1; + } + } + hint::black_box(found) + }); + }); + group.bench_with_input(BenchmarkId::new("hit", label), &hit_keys, |b, keys| { + b.iter(|| { + let mut found = 0u64; + for k in keys { + if hint::black_box(filter.check(k.as_slice())) { + found += 1; + } + } + hint::black_box(found) + }); + }); + } + group.finish(); +} + criterion_group!( benches, bench_fold_to_target_fpp, bench_insert_and_fold, - bench_insert_only + bench_insert_only, + bench_insert, + bench_check, ); criterion_main!(benches); From c58f0b1f09fb455ec95d3d87a409c5cfc92940f1 Mon Sep 17 00:00:00 2001 From: Ed Seidl Date: Mon, 1 Jun 2026 07:32:15 -0700 Subject: [PATCH 32/51] chore: Remove some deprecated functions from parquet crate (#10035) # Which issue does this PR close? None # Rationale for this change Tidying up for 59.0.0. # What changes are included in this PR? Removes public APIs that were due to be removed per the [deprecation policy](https://github.com/apache/arrow-rs#deprecation-guidelines). # Are these changes tested? Should be covered by existing tests # Are there any user-facing changes? Yes, public functions have been removed, including `ParquetMetaDataReader::with_page_indexes`, `read_columns_indexes`, and `read_offset_indexes`. --- parquet/src/arrow/async_reader/mod.rs | 11 +-- parquet/src/arrow/mod.rs | 11 ++- parquet/src/file/metadata/reader.rs | 78 +++++------------ parquet/src/file/page_index/index_reader.rs | 93 --------------------- parquet/src/file/serialized_reader.rs | 32 +------ parquet/tests/arrow_reader/bad_data.rs | 9 +- 6 files changed, 35 insertions(+), 199 deletions(-) diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index 3bba746e740e..5a0083b7164d 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -1634,7 +1634,6 @@ mod tests { } #[tokio::test] - #[allow(deprecated)] async fn empty_offset_index_doesnt_panic_in_read_row_group() { use tokio::fs::File; let testdata = arrow::util::test_util::parquet_test_data(); @@ -1642,7 +1641,7 @@ mod tests { let mut file = File::open(&path).await.unwrap(); let file_size = file.metadata().await.unwrap().len(); let mut metadata = ParquetMetaDataReader::new() - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .load_and_finish(&mut file, file_size) .await .unwrap(); @@ -1660,7 +1659,6 @@ mod tests { } #[tokio::test] - #[allow(deprecated)] async fn non_empty_offset_index_doesnt_panic_in_read_row_group() { use tokio::fs::File; let testdata = arrow::util::test_util::parquet_test_data(); @@ -1668,7 +1666,7 @@ mod tests { let mut file = File::open(&path).await.unwrap(); let file_size = file.metadata().await.unwrap().len(); let metadata = ParquetMetaDataReader::new() - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .load_and_finish(&mut file, file_size) .await .unwrap(); @@ -1685,7 +1683,6 @@ mod tests { } #[tokio::test] - #[allow(deprecated)] async fn empty_offset_index_doesnt_panic_in_column_chunks() { use tempfile::TempDir; use tokio::fs::File; @@ -1705,7 +1702,7 @@ mod tests { use std::fs::File; let file = File::open(file).unwrap(); ParquetMetaDataReader::new() - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .parse_and_finish(&file) .unwrap() } @@ -1715,7 +1712,7 @@ mod tests { let mut file = File::open(&path).await.unwrap(); let file_size = file.metadata().await.unwrap().len(); let metadata = ParquetMetaDataReader::new() - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .load_and_finish(&mut file, file_size) .await .unwrap(); diff --git a/parquet/src/arrow/mod.rs b/parquet/src/arrow/mod.rs index 14f7b9b6b41b..b89db361eda1 100644 --- a/parquet/src/arrow/mod.rs +++ b/parquet/src/arrow/mod.rs @@ -485,7 +485,8 @@ pub fn parquet_column<'a>( mod test { use crate::arrow::ArrowWriter; use crate::file::metadata::{ - ParquetMetaData, ParquetMetaDataOptions, ParquetMetaDataReader, ParquetMetaDataWriter, + PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, ParquetMetaDataReader, + ParquetMetaDataWriter, }; use crate::file::properties::{EnabledStatistics, WriterProperties}; use crate::schema::parser::parse_message_type; @@ -497,7 +498,6 @@ mod test { use super::ProjectionMask; #[test] - #[allow(deprecated)] // Reproducer for https://github.com/apache/arrow-rs/issues/6464 fn test_metadata_read_write_partial_offset() { let parquet_bytes = create_parquet_file(); @@ -514,7 +514,7 @@ mod test { let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); let err = ParquetMetaDataReader::new() .with_metadata_options(Some(options)) - .with_page_indexes(true) // there are no page indexes in the metadata + .with_page_index_policy(PageIndexPolicy::Required) // there are no page indexes in the metadata .parse_and_finish(&metadata_bytes) .err() .unwrap(); @@ -553,7 +553,6 @@ mod test { } #[test] - #[allow(deprecated)] fn test_metadata_read_write_roundtrip_page_index() { let parquet_bytes = create_parquet_file(); @@ -562,7 +561,7 @@ mod test { let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); let original_metadata = ParquetMetaDataReader::new() .with_metadata_options(Some(options)) - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .parse_and_finish(&parquet_bytes) .unwrap(); @@ -571,7 +570,7 @@ mod test { let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); let roundtrip_metadata = ParquetMetaDataReader::new() .with_metadata_options(Some(options)) - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .parse_and_finish(&metadata_bytes) .unwrap(); diff --git a/parquet/src/file/metadata/reader.rs b/parquet/src/file/metadata/reader.rs index 60c57129fd24..844ae747c7c9 100644 --- a/parquet/src/file/metadata/reader.rs +++ b/parquet/src/file/metadata/reader.rs @@ -56,12 +56,12 @@ use crate::arrow::async_reader::{MetadataFetch, MetadataSuffixFetch}; /// /// # Example /// ```no_run -/// # use parquet::file::metadata::ParquetMetaDataReader; +/// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; /// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); } /// // read parquet metadata including page indexes from a file /// let file = open_parquet_file("some_path.parquet"); /// let mut reader = ParquetMetaDataReader::new() -/// .with_page_indexes(true); +/// .with_page_index_policy(PageIndexPolicy::Required); /// reader.try_parse(&file).unwrap(); /// let metadata = reader.finish().unwrap(); /// assert!(metadata.column_index().is_some()); @@ -117,33 +117,6 @@ impl ParquetMetaDataReader { } } - /// Enable or disable reading the page index structures described in - /// "[Parquet page index]: Layout to Support Page Skipping". - /// - /// [Parquet page index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md - #[deprecated(since = "56.1.0", note = "Use `with_page_index_policy` instead")] - pub fn with_page_indexes(self, val: bool) -> Self { - self.with_page_index_policy(PageIndexPolicy::from(val)) - } - - /// Enable or disable reading the Parquet [ColumnIndex] structure. - /// - /// [ColumnIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md - #[deprecated(since = "56.1.0", note = "Use `with_column_index_policy` instead")] - pub fn with_column_indexes(self, val: bool) -> Self { - let policy = PageIndexPolicy::from(val); - self.with_column_index_policy(policy) - } - - /// Enable or disable reading the Parquet [OffsetIndex] structure. - /// - /// [OffsetIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md - #[deprecated(since = "56.1.0", note = "Use `with_offset_index_policy` instead")] - pub fn with_offset_indexes(self, val: bool) -> Self { - let policy = PageIndexPolicy::from(val); - self.with_offset_index_policy(policy) - } - /// Sets the [`PageIndexPolicy`] for the column and offset indexes pub fn with_page_index_policy(self, policy: PageIndexPolicy) -> Self { self.with_column_index_policy(policy) @@ -218,12 +191,12 @@ impl ParquetMetaDataReader { /// /// # Example /// ```no_run - /// # use parquet::file::metadata::ParquetMetaDataReader; + /// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; /// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); } /// // read parquet metadata including page indexes /// let file = open_parquet_file("some_path.parquet"); /// let metadata = ParquetMetaDataReader::new() - /// .with_page_indexes(true) + /// .with_page_index_policy(PageIndexPolicy::Required) /// .parse_and_finish(&file).unwrap(); /// ``` pub fn parse_and_finish(mut self, reader: &R) -> Result { @@ -257,7 +230,7 @@ impl ParquetMetaDataReader { /// /// # Example /// ```no_run - /// # use parquet::file::metadata::ParquetMetaDataReader; + /// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; /// # use parquet::errors::ParquetError; /// # use crate::parquet::file::reader::Length; /// # fn get_bytes(file: &std::fs::File, range: std::ops::Range) -> bytes::Bytes { unimplemented!(); } @@ -266,7 +239,7 @@ impl ParquetMetaDataReader { /// let len = file.len(); /// // Speculatively read 1 kilobyte from the end of the file /// let bytes = get_bytes(&file, len - 1024..len); - /// let mut reader = ParquetMetaDataReader::new().with_page_indexes(true); + /// let mut reader = ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); /// match reader.try_parse_sized(&bytes, len) { /// Ok(_) => (), /// Err(ParquetError::NeedMoreData(needed)) => { @@ -284,7 +257,7 @@ impl ParquetMetaDataReader { /// to test for this. In the event the file metadata is present, re-parsing of the file /// metadata can be skipped by using [`Self::read_page_indexes_sized()`], as shown below. /// ```no_run - /// # use parquet::file::metadata::ParquetMetaDataReader; + /// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; /// # use parquet::errors::ParquetError; /// # use crate::parquet::file::reader::Length; /// # fn get_bytes(file: &std::fs::File, range: std::ops::Range) -> bytes::Bytes { unimplemented!(); } @@ -293,7 +266,7 @@ impl ParquetMetaDataReader { /// let len = file.len(); /// // Speculatively read 1 kilobyte from the end of the file /// let mut bytes = get_bytes(&file, len - 1024..len); - /// let mut reader = ParquetMetaDataReader::new().with_page_indexes(true); + /// let mut reader = ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); /// // Loop until `bytes` is large enough /// loop { /// match reader.try_parse_sized(&bytes, len) { @@ -720,18 +693,6 @@ impl ParquetMetaDataReader { } } - /// Decodes a [`FooterTail`] from the provided 8-byte slice. - #[deprecated(since = "57.0.0", note = "Use FooterTail::try_from instead")] - pub fn decode_footer_tail(slice: &[u8; FOOTER_SIZE]) -> Result { - FooterTail::try_new(slice) - } - - /// Decodes the Parquet footer, returning the metadata length in bytes - #[deprecated(since = "54.3.0", note = "Use decode_footer_tail instead")] - pub fn decode_footer(slice: &[u8; FOOTER_SIZE]) -> Result { - FooterTail::try_new(slice).map(|f| f.metadata_length()) - } - /// Decodes [`ParquetMetaData`] from the provided bytes. /// /// Typically, this is used to decode the metadata from the end of a parquet @@ -915,12 +876,12 @@ mod tests { } #[test] - #[allow(deprecated)] fn test_try_parse() { let file = get_test_file("alltypes_tiny_pages.parquet"); let len = file.len(); - let mut reader = ParquetMetaDataReader::new().with_page_indexes(true); + let mut reader = + ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); let bytes_for_range = |range: Range| { file.get_bytes(range.start, (range.end - range.start).try_into().unwrap()) @@ -964,7 +925,8 @@ mod tests { }; // not enough for file metadata, but keep trying until page indexes are read - let mut reader = ParquetMetaDataReader::new().with_page_indexes(true); + let mut reader = + ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); let mut bytes = bytes_for_range(452505..len); loop { match reader.try_parse_sized(&bytes, len) { @@ -1301,7 +1263,6 @@ mod async_tests { } #[tokio::test] - #[allow(deprecated)] async fn test_page_index() { let mut file = get_test_file("alltypes_tiny_pages.parquet"); let len = file.len(); @@ -1312,7 +1273,8 @@ mod async_tests { }; let f = MetadataFetchFn(&mut fetch); - let mut loader = ParquetMetaDataReader::new().with_page_indexes(true); + let mut loader = + ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); loader.try_load(f, len).await.unwrap(); assert_eq!(fetch_count.load(Ordering::SeqCst), 3); let metadata = loader.finish().unwrap(); @@ -1322,7 +1284,7 @@ mod async_tests { fetch_count.store(0, Ordering::SeqCst); let f = MetadataFetchFn(&mut fetch); let mut loader = ParquetMetaDataReader::new() - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .with_prefetch_hint(Some(1729)); loader.try_load(f, len).await.unwrap(); assert_eq!(fetch_count.load(Ordering::SeqCst), 2); @@ -1333,7 +1295,7 @@ mod async_tests { fetch_count.store(0, Ordering::SeqCst); let f = MetadataFetchFn(&mut fetch); let mut loader = ParquetMetaDataReader::new() - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .with_prefetch_hint(Some(130649)); loader.try_load(f, len).await.unwrap(); assert_eq!(fetch_count.load(Ordering::SeqCst), 2); @@ -1344,7 +1306,7 @@ mod async_tests { fetch_count.store(0, Ordering::SeqCst); let f = MetadataFetchFn(&mut fetch); let metadata = ParquetMetaDataReader::new() - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .with_prefetch_hint(Some(130650)) .load_and_finish(f, len) .await @@ -1356,7 +1318,7 @@ mod async_tests { fetch_count.store(0, Ordering::SeqCst); let f = MetadataFetchFn(&mut fetch); let metadata = ParquetMetaDataReader::new() - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .with_prefetch_hint(Some((len - 1000) as usize)) // prefetch entire file .load_and_finish(f, len) .await @@ -1368,7 +1330,7 @@ mod async_tests { fetch_count.store(0, Ordering::SeqCst); let f = MetadataFetchFn(&mut fetch); let metadata = ParquetMetaDataReader::new() - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .with_prefetch_hint(Some(len as usize)) // prefetch entire file .load_and_finish(f, len) .await @@ -1380,7 +1342,7 @@ mod async_tests { fetch_count.store(0, Ordering::SeqCst); let f = MetadataFetchFn(&mut fetch); let metadata = ParquetMetaDataReader::new() - .with_page_indexes(true) + .with_page_index_policy(PageIndexPolicy::Required) .with_prefetch_hint(Some((len + 1000) as usize)) // prefetch entire file .load_and_finish(f, len) .await diff --git a/parquet/src/file/page_index/index_reader.rs b/parquet/src/file/page_index/index_reader.rs index ce7fc5fbaac5..f0e40f7fdd1f 100644 --- a/parquet/src/file/page_index/index_reader.rs +++ b/parquet/src/file/page_index/index_reader.rs @@ -20,12 +20,10 @@ use crate::basic::{BoundaryOrder, Type}; use crate::data_type::Int96; use crate::errors::{ParquetError, Result}; -use crate::file::metadata::ColumnChunkMetaData; use crate::file::page_index::column_index::{ ByteArrayColumnIndex, ColumnIndexMetaData, PrimitiveColumnIndex, }; use crate::file::page_index::offset_index::OffsetIndexMetaData; -use crate::file::reader::ChunkReader; use crate::parquet_thrift::{ ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, ThriftSliceInputProtocol, WriteThrift, WriteThriftField, read_thrift_vec, @@ -44,97 +42,6 @@ pub(crate) fn acc_range(a: Option>, b: Option>) -> Option< } } -/// Reads per-column [`ColumnIndexMetaData`] for all columns of a row group by -/// decoding [`ColumnIndex`] . -/// -/// Returns a vector of `index[column_number]`. -/// -/// Returns `None` if this row group does not contain a [`ColumnIndex`]. -/// -/// See [Page Index Documentation] for more details. -/// -/// [Page Index Documentation]: https://github.com/apache/parquet-format/blob/master/PageIndex.md -/// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md -#[deprecated( - since = "55.2.0", - note = "Use ParquetMetaDataReader instead; will be removed in 58.0.0" -)] -pub fn read_columns_indexes( - reader: &R, - chunks: &[ColumnChunkMetaData], -) -> Result>, ParquetError> { - let fetch = chunks - .iter() - .fold(None, |range, c| acc_range(range, c.column_index_range())); - - let fetch = match fetch { - Some(r) => r, - None => return Ok(None), - }; - - let bytes = reader.get_bytes(fetch.start as _, (fetch.end - fetch.start).try_into()?)?; - - Some( - chunks - .iter() - .map(|c| match c.column_index_range() { - Some(r) => decode_column_index( - &bytes[usize::try_from(r.start - fetch.start)? - ..usize::try_from(r.end - fetch.start)?], - c.column_type(), - ), - None => Ok(ColumnIndexMetaData::NONE), - }) - .collect(), - ) - .transpose() -} - -/// Reads per-column [`OffsetIndexMetaData`] for all columns of a row group by -/// decoding [`OffsetIndex`] . -/// -/// Returns a vector of `offset_index[column_number]`. -/// -/// Returns `None` if this row group does not contain an [`OffsetIndex`]. -/// -/// See [Page Index Documentation] for more details. -/// -/// [Page Index Documentation]: https://github.com/apache/parquet-format/blob/master/PageIndex.md -/// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md -#[deprecated( - since = "55.2.0", - note = "Use ParquetMetaDataReader instead; will be removed in 58.0.0" -)] -pub fn read_offset_indexes( - reader: &R, - chunks: &[ColumnChunkMetaData], -) -> Result>, ParquetError> { - let fetch = chunks - .iter() - .fold(None, |range, c| acc_range(range, c.offset_index_range())); - - let fetch = match fetch { - Some(r) => r, - None => return Ok(None), - }; - - let bytes = reader.get_bytes(fetch.start as _, (fetch.end - fetch.start).try_into()?)?; - - Some( - chunks - .iter() - .map(|c| match c.offset_index_range() { - Some(r) => decode_offset_index( - &bytes[usize::try_from(r.start - fetch.start)? - ..usize::try_from(r.end - fetch.start)?], - ), - None => Err(general_err!("missing offset index")), - }) - .collect(), - ) - .transpose() -} - pub(crate) fn decode_offset_index(data: &[u8]) -> Result { let mut prot = ThriftSliceInputProtocol::new(data); diff --git a/parquet/src/file/serialized_reader.rs b/parquet/src/file/serialized_reader.rs index 4b71e3c14e6d..113e5203c94b 100644 --- a/parquet/src/file/serialized_reader.rs +++ b/parquet/src/file/serialized_reader.rs @@ -240,7 +240,6 @@ impl SerializedFileReader { /// Creates file reader from a Parquet file with read options. /// Returns an error if the Parquet file does not exist or is corrupt. - #[allow(deprecated)] pub fn new_with_options(chunk_reader: R, options: ReadOptions) -> Result { let mut metadata_builder = ParquetMetaDataReader::new() .with_metadata_options(Some(options.metadata_options.clone())) @@ -266,8 +265,8 @@ impl SerializedFileReader { // If page indexes are desired, build them with the filtered set of row groups if options.enable_page_index { - let mut reader = - ParquetMetaDataReader::new_with_metadata(metadata).with_page_indexes(true); + let mut reader = ParquetMetaDataReader::new_with_metadata(metadata) + .with_page_index_policy(PageIndexPolicy::Required); reader.read_page_indexes(&chunk_reader)?; metadata = reader.finish()?; } @@ -1185,8 +1184,6 @@ mod tests { use crate::data_type::private::ParquetValueType; use crate::data_type::{AsBytes, FixedLenByteArrayType, Int32Type}; use crate::file::metadata::thrift::DataPageHeaderV2; - #[allow(deprecated)] - use crate::file::page_index::index_reader::{read_columns_indexes, read_offset_indexes}; use crate::file::writer::SerializedFileWriter; use crate::record::RowAccessor; use crate::schema::parser::parse_message_type; @@ -2181,31 +2178,6 @@ mod tests { assert_eq!(0, page_offset.first_row_index); } - #[test] - #[allow(deprecated)] - fn test_page_index_reader_out_of_order() { - let test_file = get_test_file("alltypes_tiny_pages_plain.parquet"); - let options = ReadOptionsBuilder::new().with_page_index().build(); - let reader = SerializedFileReader::new_with_options(test_file, options).unwrap(); - let metadata = reader.metadata(); - - let test_file = get_test_file("alltypes_tiny_pages_plain.parquet"); - let columns = metadata.row_group(0).columns(); - let reversed: Vec<_> = columns.iter().cloned().rev().collect(); - - let a = read_columns_indexes(&test_file, columns).unwrap().unwrap(); - let mut b = read_columns_indexes(&test_file, &reversed) - .unwrap() - .unwrap(); - b.reverse(); - assert_eq!(a, b); - - let a = read_offset_indexes(&test_file, columns).unwrap().unwrap(); - let mut b = read_offset_indexes(&test_file, &reversed).unwrap().unwrap(); - b.reverse(); - assert_eq!(a, b); - } - #[test] fn test_page_index_reader_all_type() { let test_file = get_test_file("alltypes_tiny_pages_plain.parquet"); diff --git a/parquet/tests/arrow_reader/bad_data.rs b/parquet/tests/arrow_reader/bad_data.rs index 38fa69cdb1b2..9a110c136151 100644 --- a/parquet/tests/arrow_reader/bad_data.rs +++ b/parquet/tests/arrow_reader/bad_data.rs @@ -188,9 +188,8 @@ fn skip_unknown_types() { #[cfg(feature = "async")] #[tokio::test] -#[allow(deprecated)] async fn bad_metadata_err() { - use parquet::file::metadata::ParquetMetaDataReader; + use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; let metadata_buffer = Bytes::from_static(include_bytes!("bad_raw_metadata.bin")); @@ -199,13 +198,13 @@ async fn bad_metadata_err() { let mut reader = std::io::Cursor::new(&metadata_buffer); let mut loader = ParquetMetaDataReader::new(); loader.try_load(&mut reader, metadata_length).await.unwrap(); - loader = loader.with_page_indexes(false); + loader = loader.with_page_index_policy(PageIndexPolicy::Skip); loader.load_page_index(&mut reader).await.unwrap(); - loader = loader.with_offset_indexes(true); + loader = loader.with_offset_index_policy(PageIndexPolicy::Required); loader.load_page_index(&mut reader).await.unwrap(); - loader = loader.with_column_indexes(true); + loader = loader.with_column_index_policy(PageIndexPolicy::Required); let err = loader.load_page_index(&mut reader).await.unwrap_err(); assert_eq!( From 38778f0110725eb36f38f9fa2d43d2060b8cf928 Mon Sep 17 00:00:00 2001 From: quantumish Date: Mon, 1 Jun 2026 15:32:35 +0100 Subject: [PATCH 33/51] Replace `From>` impls with `TryFrom`s for `FixedSizeBinaryArray` (#10019) # Which issue does this PR close? - Closes #10018. # Rationale for this change There isn't a clear way to fix the `From>` implementations for `FixedSizeBinaryArray` that wouldn't be confusing, so making them `TryFrom` is a better fit since they are in genuine use across e.g. tests within the Arrow library as well as a terser way of calling `FixedSizeBinaryArray::try_from_iter` or `FixedSizeBinaryArray::try_from_sparse_iter`. # What changes are included in this PR? - Converts `From>`, `From>`, and `From>>` implementations for `FixedSizeBinaryArray` to `TryFrom` implementations. - Adds a `TryFrom>>` implementation for the missing combination of types. - Updates various test cases within the arrow/parquet libraries to use `try_from().unwrap()` instead of `from()`. # Are these changes tested? This is sort of a transparent change in that only the API for expressing failure cases has changed rather than the actual failure cases. All existing tests surrounding conversion failures have been updated to check whether a conversion has correctly failed. # Are there any user-facing changes? Yes, this is a breaking API change since user-facing trait implementations have been replaced with different trait implementations. --- .../src/array/fixed_size_binary_array.rs | 50 +++++++++++++------ arrow-cast/src/cast/mod.rs | 4 +- arrow-ord/src/comparison.rs | 6 ++- arrow-select/src/filter.rs | 2 +- arrow-string/src/concat_elements.rs | 28 +++++++---- arrow-string/src/substring.rs | 19 ++++++- parquet/tests/arrow_reader/mod.rs | 3 +- 7 files changed, 77 insertions(+), 35 deletions(-) diff --git a/arrow-array/src/array/fixed_size_binary_array.rs b/arrow-array/src/array/fixed_size_binary_array.rs index a44d21ad5b70..67af86322894 100644 --- a/arrow-array/src/array/fixed_size_binary_array.rs +++ b/arrow-array/src/array/fixed_size_binary_array.rs @@ -679,22 +679,42 @@ impl From for FixedSizeBinaryArray { } } -impl From>> for FixedSizeBinaryArray { - fn from(v: Vec>) -> Self { +impl TryFrom>> for FixedSizeBinaryArray { + type Error = ArrowError; + + fn try_from(v: Vec>) -> Result { #[allow(deprecated)] - Self::try_from_sparse_iter(v.into_iter()).unwrap() + Self::try_from_sparse_iter(v.into_iter()) + } +} + +impl TryFrom> for FixedSizeBinaryArray { + type Error = ArrowError; + + fn try_from(v: Vec<&[u8]>) -> Result { + Self::try_from_iter(v.into_iter()) } } -impl From> for FixedSizeBinaryArray { - fn from(v: Vec<&[u8]>) -> Self { - Self::try_from_iter(v.into_iter()).unwrap() +impl TryFrom>> for FixedSizeBinaryArray { + type Error = ArrowError; + + fn try_from(v: Vec>) -> Result { + N.try_into() + .map_err(|_| { + ArrowError::InvalidArgumentError(format!( + "FixedSizeBinaryArray value length exceeds i32, got {N}" + )) + }) + .and_then(|x| Self::try_from_sparse_iter_with_size(v.into_iter(), x)) } } -impl From> for FixedSizeBinaryArray { - fn from(v: Vec<&[u8; N]>) -> Self { - Self::try_from_iter(v.into_iter()).unwrap() +impl TryFrom> for FixedSizeBinaryArray { + type Error = ArrowError; + + fn try_from(v: Vec<&[u8; N]>) -> Result { + Self::try_from_iter(v.into_iter()) } } @@ -1009,7 +1029,7 @@ mod tests { #[test] fn test_fixed_size_binary_array_from_vec() { let values = vec!["one".as_bytes(), b"two", b"six", b"ten"]; - let array = FixedSizeBinaryArray::from(values); + let array = FixedSizeBinaryArray::try_from(values).unwrap(); assert_eq!(array.len(), 4); assert_eq!(array.null_count(), 0); assert_eq!(array.logical_null_count(), 0); @@ -1024,10 +1044,9 @@ mod tests { } #[test] - #[should_panic(expected = "Nested array size mismatch: one is 3, and the other is 5")] fn test_fixed_size_binary_array_from_vec_incorrect_length() { let values = vec!["one".as_bytes(), b"two", b"three", b"four"]; - let _ = FixedSizeBinaryArray::from(values); + assert!(FixedSizeBinaryArray::try_from(values).is_err()); } #[test] @@ -1039,7 +1058,7 @@ mod tests { Some(b"six"), Some(b"ten"), ]; - let array = FixedSizeBinaryArray::from(values); + let array = FixedSizeBinaryArray::try_from(values).unwrap(); assert_eq!(array.len(), 5); assert_eq!(array.value(0), b"one"); assert_eq!(array.value(1), b"two"); @@ -1053,7 +1072,6 @@ mod tests { } #[test] - #[should_panic(expected = "Nested array size mismatch: one is 3, and the other is 5")] fn test_fixed_size_binary_array_from_opt_vec_incorrect_length() { let values = vec![ Some("one".as_bytes()), @@ -1062,7 +1080,7 @@ mod tests { Some(b"three"), Some(b"four"), ]; - let _ = FixedSizeBinaryArray::from(values); + assert!(FixedSizeBinaryArray::try_from(values).is_err()); } #[test] @@ -1098,7 +1116,7 @@ mod tests { )] fn test_fixed_size_binary_array_get_value_index_out_of_bound() { let values = vec![Some("one".as_bytes()), Some(b"two"), None]; - let array = FixedSizeBinaryArray::from(values); + let array = FixedSizeBinaryArray::try_from(values).unwrap(); array.value(4); } diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs index 67da85b8c1d6..0367a5412108 100644 --- a/arrow-cast/src/cast/mod.rs +++ b/arrow-cast/src/cast/mod.rs @@ -5908,7 +5908,7 @@ mod tests { let bytes_2 = "Hello".as_bytes(); let binary_data = vec![Some(bytes_1), Some(bytes_2), None]; - let a1 = Arc::new(FixedSizeBinaryArray::from(binary_data.clone())) as ArrayRef; + let a1 = Arc::new(FixedSizeBinaryArray::try_from(binary_data.clone()).unwrap()) as ArrayRef; let array_ref = cast(&a1, &DataType::Binary).unwrap(); let down_cast = array_ref.as_binary::(); @@ -5935,7 +5935,7 @@ mod tests { let bytes_2 = "Hello".as_bytes(); let binary_data = vec![Some(bytes_1), Some(bytes_2), Some(bytes_1), None]; - let a1 = Arc::new(FixedSizeBinaryArray::from(binary_data.clone())) as ArrayRef; + let a1 = Arc::new(FixedSizeBinaryArray::try_from(binary_data.clone()).unwrap()) as ArrayRef; let cast_type = DataType::Dictionary( Box::new(DataType::Int8), diff --git a/arrow-ord/src/comparison.rs b/arrow-ord/src/comparison.rs index 3aff2c6234be..4d61eb875814 100644 --- a/arrow-ord/src/comparison.rs +++ b/arrow-ord/src/comparison.rs @@ -1801,7 +1801,8 @@ mod tests { expected ); - let fsb_array = FixedSizeBinaryArray::from(vec![&[0u8], &[0u8], &[0u8], &[1u8]]); + let fsb_array = + FixedSizeBinaryArray::try_from(vec![&[0u8], &[0u8], &[0u8], &[1u8]]).unwrap(); let scalar = FixedSizeBinaryArray::new_scalar([1u8]); let expected = BooleanArray::from(vec![Some(false), Some(false), Some(false), Some(true)]); assert_eq!(crate::cmp::eq(&fsb_array, &scalar).unwrap(), expected); @@ -1824,7 +1825,8 @@ mod tests { expected ); - let fsb_array = FixedSizeBinaryArray::from(vec![&[0u8], &[0u8], &[0u8], &[1u8]]); + let fsb_array = + FixedSizeBinaryArray::try_from(vec![&[0u8], &[0u8], &[0u8], &[1u8]]).unwrap(); let scalar = FixedSizeBinaryArray::new_scalar([1u8]); let expected = BooleanArray::from(vec![Some(true), Some(true), Some(true), Some(false)]); assert_eq!(crate::cmp::neq(&fsb_array, &scalar).unwrap(), expected); diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index aaab9d202013..fcbce82d5d9d 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -1273,7 +1273,7 @@ mod tests { let v2 = [3_u8, 4]; let v3 = [5_u8, 6]; let v = vec![&v1, &v2, &v3]; - let a = FixedSizeBinaryArray::from(v); + let a = FixedSizeBinaryArray::try_from(v).unwrap(); let b = BooleanArray::from(vec![true, false, true]); let c = filter(&a, &b).unwrap(); let d = c diff --git a/arrow-string/src/concat_elements.rs b/arrow-string/src/concat_elements.rs index cd4676d28752..8742b1739230 100644 --- a/arrow-string/src/concat_elements.rs +++ b/arrow-string/src/concat_elements.rs @@ -505,30 +505,33 @@ mod tests { #[test] fn test_fixed_size_binary_concat() { - let left = FixedSizeBinaryArray::from(vec![Some(b"foo" as &[u8]), Some(b"bar"), None]); - let right = FixedSizeBinaryArray::from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")]); + let left = FixedSizeBinaryArray::try_from(vec![Some(b"foo" as &[u8]), Some(b"bar"), None]) + .unwrap(); + let right = FixedSizeBinaryArray::try_from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")]) + .unwrap(); let output = concat_elements_fixed_size_binary(&left, &right).unwrap(); - let expected = FixedSizeBinaryArray::from(vec![None, Some(b"baryyy" as &[u8]), None]); + let expected = + FixedSizeBinaryArray::try_from(vec![None, Some(b"baryyy" as &[u8]), None]).unwrap(); assert_eq!(output, expected); } #[test] fn test_fixed_size_binary_concat_no_null() { - let left = FixedSizeBinaryArray::from(vec![b"ab" as &[u8], b"cd"]); - let right = FixedSizeBinaryArray::from(vec![b"12" as &[u8], b"34"]); + let left = FixedSizeBinaryArray::try_from(vec![b"ab" as &[u8], b"cd"]).unwrap(); + let right = FixedSizeBinaryArray::try_from(vec![b"12" as &[u8], b"34"]).unwrap(); let output = concat_elements_fixed_size_binary(&left, &right).unwrap(); - let expected = FixedSizeBinaryArray::from(vec![b"ab12" as &[u8], b"cd34"]); + let expected = FixedSizeBinaryArray::try_from(vec![b"ab12" as &[u8], b"cd34"]).unwrap(); assert_eq!(output, expected); } #[test] fn test_fixed_size_binary_concat_error() { - let left = FixedSizeBinaryArray::from(vec![b"ab" as &[u8], b"cd"]); - let right = FixedSizeBinaryArray::from(vec![b"12" as &[u8]]); + let left = FixedSizeBinaryArray::try_from(vec![b"ab" as &[u8], b"cd"]).unwrap(); + let right = FixedSizeBinaryArray::try_from(vec![b"12" as &[u8]]).unwrap(); let output = concat_elements_fixed_size_binary(&left, &right); assert_eq!( @@ -673,13 +676,16 @@ mod tests { assert_eq!(output, expected); // test for FixedSizeBinaryArray - let left = FixedSizeBinaryArray::from(vec![Some(b"foo" as &[u8]), Some(b"bar"), None]); - let right = FixedSizeBinaryArray::from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")]); + let left = FixedSizeBinaryArray::try_from(vec![Some(b"foo" as &[u8]), Some(b"bar"), None]) + .unwrap(); + let right = FixedSizeBinaryArray::try_from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")]) + .unwrap(); let output: FixedSizeBinaryArray = concat_elements_dyn(&left, &right) .unwrap() .into_data() .into(); - let expected = FixedSizeBinaryArray::from(vec![None, Some(b"baryyy" as &[u8]), None]); + let expected = + FixedSizeBinaryArray::try_from(vec![None, Some(b"baryyy" as &[u8]), None]).unwrap(); assert_eq!(output, expected); } diff --git a/arrow-string/src/substring.rs b/arrow-string/src/substring.rs index 123c77cdf9dc..563c66448069 100644 --- a/arrow-string/src/substring.rs +++ b/arrow-string/src/substring.rs @@ -429,6 +429,21 @@ mod tests { }; } + /// A helper macro to test the substring functions for array types only implementing TryFrom. + macro_rules! do_test_tryfrom { + ($cases:expr, $array_ty:ty, $substring_fn:ident) => { + $cases + .into_iter() + .for_each(|(array, start, length, expected)| { + let array = <$array_ty>::try_from(array).unwrap(); + let result = $substring_fn(&array, start, length).unwrap(); + let result = result.as_any().downcast_ref::<$array_ty>().unwrap(); + let expected = <$array_ty>::try_from(expected).unwrap(); + assert_eq!(&expected, result); + }) + }; + } + fn with_nulls_generic_binary() { let input = vec![ Some("hello".as_bytes()), @@ -591,7 +606,7 @@ mod tests { (-3, Some(4), input.clone()) ); - do_test!( + do_test_tryfrom!( [&base_case[..], &cases[..]].concat(), FixedSizeBinaryArray, substring @@ -630,7 +645,7 @@ mod tests { (-3, Some(4), input.clone()) ); - do_test!( + do_test_tryfrom!( [&base_case[..], &cases[..]].concat(), FixedSizeBinaryArray, substring diff --git a/parquet/tests/arrow_reader/mod.rs b/parquet/tests/arrow_reader/mod.rs index d9729fdc7703..404bebc05d0b 100644 --- a/parquet/tests/arrow_reader/mod.rs +++ b/parquet/tests/arrow_reader/mod.rs @@ -559,7 +559,8 @@ fn make_bytearray_batch( .iter() .map(|value| Some(value.as_slice())) .collect::>() - .into(); + .try_into() + .unwrap(); let service_large_binary: LargeBinaryArray = large_binary_values.iter().map(Some).collect(); let schema = Schema::new(vec![ From 51ffd8c873feb3fbc1659ebffcc2c72fca796e94 Mon Sep 17 00:00:00 2001 From: theirix Date: Tue, 2 Jun 2026 15:52:05 +0100 Subject: [PATCH 34/51] fix: better error handling for negative size of FixedSizeBinary (#10042) # Which issue does this PR close? - Closes #10033. # Rationale for this change Related to https://github.com/apache/datafusion/pull/22297, where using `FixedSizeBinary(-N)` caused failures. Actually, it will still panic, but with a proper error. It could be a good idea to introduce `try_new_null` to carry an error gracefully - thoughts? # What changes are included in this PR? - Return a `Result` on negative byte width when possible - Panic explicitly with a proper error message otherwise - Avoid silent overflow with a direct `len as usize` cast - Reject negative FSB when parsing from tokens # Are these changes tested? - Tests are passing # Are there any user-facing changes? No --- arrow-array/src/ffi.rs | 8 ++++++- arrow-data/src/data.rs | 30 +++++++++++++++++++++++- arrow-data/src/equal/fixed_binary.rs | 8 ++----- arrow-data/src/transform/fixed_binary.rs | 13 +++------- arrow-json/src/reader/binary_array.rs | 5 +++- arrow-row/src/fixed.rs | 6 +++++ arrow-schema/src/datatype_parse.rs | 13 ++++++++-- arrow-string/src/concat_elements.rs | 14 +++++++++-- arrow/src/util/data_gen.rs | 7 +++++- 9 files changed, 80 insertions(+), 24 deletions(-) diff --git a/arrow-array/src/ffi.rs b/arrow-array/src/ffi.rs index 59b9f6b3b7b5..110dfc9855bb 100644 --- a/arrow-array/src/ffi.rs +++ b/arrow-array/src/ffi.rs @@ -162,7 +162,13 @@ fn bit_width(data_type: &DataType, i: usize) -> Result { "The datatype \"{data_type}\" expects 2 buffers, but requested {i}. Please verify that the C data interface is correctly implemented." ))); } - (DataType::FixedSizeBinary(num_bytes), 1) => *num_bytes as usize * u8::BITS as usize, + (DataType::FixedSizeBinary(num_bytes), 1) => { + TryInto::::try_into(*num_bytes).map_err(|_| { + ArrowError::InvalidArgumentError(format!( + "cannot determine bit_width for FixedSizeBinary({num_bytes})" + )) + })? * u8::BITS as usize + } (DataType::FixedSizeList(f, num_elems), 1) => { let child_bit_width = bit_width(f.data_type(), 1)?; child_bit_width * (*num_elems as usize) diff --git a/arrow-data/src/data.rs b/arrow-data/src/data.rs index 34d5cb6d3aba..7b208b7af0be 100644 --- a/arrow-data/src/data.rs +++ b/arrow-data/src/data.rs @@ -134,6 +134,9 @@ pub(crate) fn new_buffers(data_type: &DataType, capacity: usize) -> [MutableBuff MutableBuffer::new(capacity * mem::size_of::()), ], DataType::FixedSizeBinary(size) => { + if *size < 0 { + panic!("cannot construct buffers from FixedSizeBinary({size})"); + } [MutableBuffer::new(capacity * *size as usize), empty_buffer] } DataType::Dictionary(k, _) => [ @@ -634,6 +637,10 @@ impl ArrayData { } /// Returns a new [`ArrayData`] valid for `data_type` containing `len` null values + /// + /// # Panics + /// This function panics if: + /// * the datatype `data_type` has incorrect layout pub fn new_null(data_type: &DataType, len: usize) -> Self { let bit_len = bit_util::ceil(len, 8); let zeroed = |len: usize| Buffer::from(MutableBuffer::from_len_zeroed(len)); @@ -650,7 +657,12 @@ impl ArrayData { DataType::LargeBinary | DataType::LargeUtf8 => { (vec![zeroed((len + 1) * 8), zeroed(0)], vec![], true) } - DataType::FixedSizeBinary(i) => (vec![zeroed(*i as usize * len)], vec![], true), + DataType::FixedSizeBinary(i) => { + if *i < 0 { + panic!("cannot construct null data from FixedSizeBinary({i})"); + } + (vec![zeroed(*i as usize * len)], vec![], true) + } DataType::List(f) | DataType::Map(f, _) => ( vec![zeroed((len + 1) * 4)], vec![ArrayData::new_empty(f.data_type())], @@ -2262,6 +2274,22 @@ impl From for ArrayDataBuilder { } } +/// Get byte width of FixedSizeBinary size +/// # Panics: +/// - Panics if the `data_type` is not FixedSizeBinary +/// - Panics if byte width is negative +pub(crate) fn get_fixed_size_binary_width(data_type: &DataType) -> usize { + match data_type { + DataType::FixedSizeBinary(i) => { + if *i < 0 { + panic!("cannot compare FixedSizeBinary({})", *i); + } + *i as usize + } + _ => unreachable!(), + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/arrow-data/src/equal/fixed_binary.rs b/arrow-data/src/equal/fixed_binary.rs index 0778d77e2fdd..8d29a2684db8 100644 --- a/arrow-data/src/equal/fixed_binary.rs +++ b/arrow-data/src/equal/fixed_binary.rs @@ -17,9 +17,8 @@ use crate::bit_iterator::BitSliceIterator; use crate::contains_nulls; -use crate::data::ArrayData; +use crate::data::{ArrayData, get_fixed_size_binary_width}; use crate::equal::primitive::NULL_SLICES_SELECTIVITY_THRESHOLD; -use arrow_schema::DataType; use super::utils::equal_len; @@ -30,10 +29,7 @@ pub(super) fn fixed_binary_equal( rhs_start: usize, len: usize, ) -> bool { - let size = match lhs.data_type() { - DataType::FixedSizeBinary(i) => *i as usize, - _ => unreachable!(), - }; + let size = get_fixed_size_binary_width(lhs.data_type()); let lhs_values = &lhs.buffers()[0].as_slice()[lhs.offset() * size..]; let rhs_values = &rhs.buffers()[0].as_slice()[rhs.offset() * size..]; diff --git a/arrow-data/src/transform/fixed_binary.rs b/arrow-data/src/transform/fixed_binary.rs index 626ecbee0261..7f4bde3de1dd 100644 --- a/arrow-data/src/transform/fixed_binary.rs +++ b/arrow-data/src/transform/fixed_binary.rs @@ -17,14 +17,10 @@ use super::{_MutableArrayData, Extend}; use crate::ArrayData; -use arrow_schema::DataType; +use crate::data::get_fixed_size_binary_width; pub(super) fn build_extend(array: &ArrayData) -> Extend<'_> { - let size = match array.data_type() { - DataType::FixedSizeBinary(i) => *i as usize, - _ => unreachable!(), - }; - + let size = get_fixed_size_binary_width(array.data_type()); let values = &array.buffers()[0].as_slice()[array.offset() * size..]; Box::new( move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| { @@ -35,10 +31,7 @@ pub(super) fn build_extend(array: &ArrayData) -> Extend<'_> { } pub(super) fn extend_nulls(mutable: &mut _MutableArrayData, len: usize) { - let size = match mutable.data_type { - DataType::FixedSizeBinary(i) => i as usize, - _ => unreachable!(), - }; + let size = get_fixed_size_binary_width(&mutable.data_type); let values_buffer = &mut mutable.buffer1; values_buffer.extend_zeros(len * size); diff --git a/arrow-json/src/reader/binary_array.rs b/arrow-json/src/reader/binary_array.rs index b1b736e83895..7f0d5219e290 100644 --- a/arrow-json/src/reader/binary_array.rs +++ b/arrow-json/src/reader/binary_array.rs @@ -133,7 +133,10 @@ impl ArrayDecoder for FixedSizeBinaryArrayDecoder { fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result { let mut builder = FixedSizeBinaryBuilder::with_capacity(pos.len(), self.len); // Preallocate for the decoded byte width (FixedSizeBinary len), not the hex string length. - let mut scratch = Vec::with_capacity(self.len as usize); + let capacity: usize = self.len.try_into().map_err(|_| { + ArrowError::InvalidArgumentError(format!("Cannot convert size '{}' to usize", self.len)) + })?; + let mut scratch = Vec::with_capacity(capacity); for p in pos { match tape.get(*p) { diff --git a/arrow-row/src/fixed.rs b/arrow-row/src/fixed.rs index 493e674018ab..422c0c1a32c8 100644 --- a/arrow-row/src/fixed.rs +++ b/arrow-row/src/fixed.rs @@ -475,6 +475,9 @@ where } /// Decodes a `FixedLengthBinary` from rows +/// +/// # Panics: +/// Panics if `size` is negative pub fn decode_fixed_size_binary( rows: &mut [&[u8]], size: i32, @@ -482,6 +485,9 @@ pub fn decode_fixed_size_binary( ) -> FixedSizeBinaryArray { let len = rows.len(); + if size < 0 { + panic!("cannot decode FixedSizeBinary({size})"); + } let mut values = MutableBuffer::new(size as usize * rows.len()); let (null_count, nulls) = decode_nulls(rows); diff --git a/arrow-schema/src/datatype_parse.rs b/arrow-schema/src/datatype_parse.rs index 9349635151aa..721bbda11a09 100644 --- a/arrow-schema/src/datatype_parse.rs +++ b/arrow-schema/src/datatype_parse.rs @@ -375,6 +375,12 @@ impl<'a> Parser<'a> { fn parse_fixed_size_binary(&mut self) -> ArrowResult { self.expect_token(Token::LParen)?; let length = self.parse_i32("FixedSizeBinary")?; + if length < 0 { + return Err(make_error( + self.val, + &format!("FixedSizeBinary length must be non-negative, got {length}"), + )); + } self.expect_token(Token::RParen)?; Ok(DataType::FixedSizeBinary(length)) } @@ -997,7 +1003,6 @@ mod test { DataType::BinaryView, DataType::FixedSizeBinary(0), DataType::FixedSizeBinary(1234), - DataType::FixedSizeBinary(-432), DataType::LargeBinary, DataType::Utf8, DataType::Utf8View, @@ -1312,7 +1317,6 @@ mod test { ("BinaryView", BinaryView), ("FixedSizeBinary(0)", FixedSizeBinary(0)), ("FixedSizeBinary(1234)", FixedSizeBinary(1234)), - ("FixedSizeBinary(-432)", FixedSizeBinary(-432)), ("LargeBinary", LargeBinary), ("Utf8", Utf8), ("Utf8View", Utf8View), @@ -1437,6 +1441,11 @@ mod test { "FixedSizeBinary(4000000000), ", "Error converting 4000000000 into i32 for FixedSizeBinary: out of range integral type conversion attempted", ), + // can't have negative width + ( + "FixedSizeBinary(-1), ", + "FixedSizeBinary length must be non-negative, got -1", + ), // can't have negative precision ( "Decimal32(-3, 5)", diff --git a/arrow-string/src/concat_elements.rs b/arrow-string/src/concat_elements.rs index 8742b1739230..4809cc5e108e 100644 --- a/arrow-string/src/concat_elements.rs +++ b/arrow-string/src/concat_elements.rs @@ -188,8 +188,18 @@ pub fn concat_elements_fixed_size_binary( ))); } - let left_size = left.value_length() as usize; - let right_size = right.value_length() as usize; + let left_size: usize = left.value_length().try_into().map_err(|_| { + ArrowError::InvalidArgumentError(format!( + "Invalid size of FixedSizeBinaryArray({})", + left.value_length() + )) + })?; + let right_size: usize = right.value_length().try_into().map_err(|_| { + ArrowError::InvalidArgumentError(format!( + "Invalid size of FixedSizeBinaryArray({})", + right.value_length() + )) + })?; let output_size = left_size + right_size; // Pre-compute combined null bitmap so the per-row NULL check is efficient diff --git a/arrow/src/util/data_gen.rs b/arrow/src/util/data_gen.rs index a5a0647aa877..481fb9e3a42d 100644 --- a/arrow/src/util/data_gen.rs +++ b/arrow/src/util/data_gen.rs @@ -159,7 +159,12 @@ pub fn create_random_array( )), Binary => Arc::new(create_binary_array::(size, null_density)), LargeBinary => Arc::new(create_binary_array::(size, null_density)), - FixedSizeBinary(len) => Arc::new(create_fsb_array(size, null_density, *len as usize)), + FixedSizeBinary(len) => { + let len = TryInto::::try_into(*len).map_err(|_| { + ArrowError::InvalidArgumentError(format!("cannot use FixedSizeBinary({len})")) + })?; + Arc::new(create_fsb_array(size, null_density, len)) + } BinaryView => Arc::new( create_string_view_array_with_len(size, null_density, 4, false).to_binary_view(), ), From 6881f73b0b3621f2263cb57608a8eff58efe58cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Franti=C5=A1ek=20Zatloukal?= Date: Tue, 2 Jun 2026 17:55:40 +0300 Subject: [PATCH 35/51] Adjust Variant size expectation for s390x architecture (#10027) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Which issue does this PR close? - Closes #10026 . # Rationale for this change The Variant enum has a different size on s390x (72 bytes) compared to other 64-bit architectures (80 bytes) due to architecture-specific alignment and padding requirements. # What changes are included in this PR? This change adds a conditional compilation check to expect 72 bytes on s390x while maintaining the existing 80-byte expectation for other 64-bit platforms. This ensures the size check passes on s390x without compromising the performance validation on other architectures. # Are these changes tested? Build-time test only # Are there any user-facing changes? N/A Signed-off-by: František Zatloukal --- parquet-variant/src/variant.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/parquet-variant/src/variant.rs b/parquet-variant/src/variant.rs index 58a3f7eeb261..c9f175c3a610 100644 --- a/parquet-variant/src/variant.rs +++ b/parquet-variant/src/variant.rs @@ -298,7 +298,9 @@ pub enum Variant<'m, 'v> { } // We don't want this to grow because it could hurt performance of a frequently-created type. -#[cfg(target_pointer_width = "64")] +#[cfg(all(target_pointer_width = "64", target_arch = "s390x"))] +const _: () = crate::utils::expect_size_of::(72); +#[cfg(all(target_pointer_width = "64", not(target_arch = "s390x")))] const _: () = crate::utils::expect_size_of::(80); #[cfg(target_pointer_width = "32")] const _: () = crate::utils::expect_size_of::(48); From 57eeb266af09f7fee47b6d4265ed2bdff6746929 Mon Sep 17 00:00:00 2001 From: Frederic Branczyk Date: Tue, 2 Jun 2026 14:55:53 +0000 Subject: [PATCH 36/51] arrow-cast: Add ability to cast plain struct to dictionary (#10039) # Which issue does this PR close? - Closes #10038 # What changes are included in this PR? A naive implementation of casting plain structs to dictionaries, that doesn't perform any deduplication. # Are these changes tested? Unit tests added. # Are there any user-facing changes? No, just a new feature. @alamb @Jefffrey --- arrow-cast/src/cast/dictionary.rs | 39 +++++++++ arrow-cast/src/cast/mod.rs | 126 ++++++++++++++++++++++++++++++ 2 files changed, 165 insertions(+) diff --git a/arrow-cast/src/cast/dictionary.rs b/arrow-cast/src/cast/dictionary.rs index 601f50a4d001..83aa691482e5 100644 --- a/arrow-cast/src/cast/dictionary.rs +++ b/arrow-cast/src/cast/dictionary.rs @@ -315,12 +315,51 @@ pub(crate) fn cast_to_dictionary( FixedSizeBinary(byte_size) => { pack_byte_to_fixed_size_dictionary::(array, cast_options, byte_size) } + Struct(_) => pack_struct_to_dictionary::(array, dict_value_type, cast_options), _ => Err(ArrowError::CastError(format!( "Unsupported output type for dictionary packing: {dict_value_type}" ))), } } +/// Wrap a struct-valued array as a `DictionaryArray` with identity +/// keys `[0, 1, ..., len-1]`. Unlike the primitive / byte packers above, no +/// deduplication is performed, since struct values have no general hash/equality +/// builder in arrow-rs. +/// +/// Each child field of the source is recursively cast to the matching field of +/// `dict_value_type` via `cast_with_options` before keys are emitted. If any +/// child cast fails, the whole pack fails, the same contract as the primitive +/// packers above. +fn pack_struct_to_dictionary( + array: &dyn Array, + dict_value_type: &DataType, + cast_options: &CastOptions, +) -> Result { + let cast_values = cast_with_options(array, dict_value_type, cast_options)?; + let len = cast_values.len(); + + // Identity keys `[0, 1, ..., len-1]`, with null entries wherever the + // source row is null so the dictionary's logical null mask matches. + let mut builder = PrimitiveBuilder::::with_capacity(len); + for i in 0..len { + if cast_values.is_null(i) { + builder.append_null(); + } else { + let key = K::Native::from_usize(i).ok_or_else(|| { + ArrowError::CastError(format!( + "Cannot fit {len} dictionary keys in {:?}", + K::DATA_TYPE, + )) + })?; + builder.append_value(key); + } + } + let keys = builder.finish(); + + Ok(Arc::new(DictionaryArray::::try_new(keys, cast_values)?)) +} + // Packs the data from the primitive array of type to a // DictionaryArray with keys of type K and values of value_type V pub(crate) fn pack_numeric_to_dictionary( diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs index 0367a5412108..4d67703ea6b8 100644 --- a/arrow-cast/src/cast/mod.rs +++ b/arrow-cast/src/cast/mod.rs @@ -6278,6 +6278,132 @@ mod tests { assert_ne!(keys.value(0), keys.value(1)); } + #[test] + fn test_cast_struct_array_to_dict_struct() { + // Cast a StructArray into Dictionary. The dictionary + // value type's child fields may differ from the source's (here: + // Utf8 source → Utf8View child for `name`), so the per-field cast + // must run before identity keys are emitted. This is the "as long as + // the struct can be cast to the dict value" contract. + let names = StringArray::from(vec![Some("alpha"), None, Some("gamma")]); + let ids = Int32Array::from(vec![Some(1), Some(2), Some(3)]); + let source = StructArray::from(vec![ + ( + Arc::new(Field::new("name", DataType::Utf8, true)), + Arc::new(names) as ArrayRef, + ), + ( + Arc::new(Field::new("id", DataType::Int32, false)), + Arc::new(ids) as ArrayRef, + ), + ]); + + let target_value_type = DataType::Struct( + vec![ + Field::new("name", DataType::Utf8View, true), + Field::new("id", DataType::Int64, false), + ] + .into(), + ); + let cast_type = DataType::Dictionary( + Box::new(DataType::UInt32), + Box::new(target_value_type.clone()), + ); + assert!(can_cast_types(source.data_type(), &cast_type)); + + let cast_array = cast(&source, &cast_type).unwrap(); + assert_eq!(cast_array.data_type(), &cast_type); + assert_eq!(cast_array.len(), 3); + + let dict = cast_array.as_dictionary::(); + assert_eq!(dict.values().data_type(), &target_value_type); + // No dedup is performed for struct values — one row, one key. + assert_eq!(dict.values().len(), 3); + + // Source row 1 was a `Utf8`-null in the `name` field but the whole + // struct row was valid (StructArray::from above takes per-field + // nulls only). The dictionary's logical null mask therefore mirrors + // the source struct's row-level null mask — all rows valid here. + let keys = dict.keys(); + assert_eq!(keys.values(), &[0u32, 1, 2]); + assert_eq!(keys.null_count(), 0); + + let struct_values = dict.values().as_struct(); + let names_out = struct_values + .column_by_name("name") + .unwrap() + .as_string_view(); + assert_eq!(names_out.value(0), "alpha"); + assert!(names_out.is_null(1)); + assert_eq!(names_out.value(2), "gamma"); + let ids_out = struct_values + .column_by_name("id") + .unwrap() + .as_primitive::(); + assert_eq!(ids_out.values(), &[1i64, 2, 3]); + } + + #[test] + fn test_cast_struct_array_to_dict_struct_row_nulls() { + // Row-level nulls on the source struct must surface as null keys on + // the dictionary, since the dictionary's logical null mask is + // determined by the keys. + let names = StringArray::from(vec![Some("alpha"), Some("beta"), Some("gamma")]); + let ids = Int32Array::from(vec![Some(1), Some(2), Some(3)]); + let source = StructArray::try_new( + vec![ + Field::new("name", DataType::Utf8, true), + Field::new("id", DataType::Int32, false), + ] + .into(), + vec![Arc::new(names) as ArrayRef, Arc::new(ids) as ArrayRef], + Some(NullBuffer::from(vec![true, false, true])), + ) + .unwrap(); + + let target_value_type = DataType::Struct( + vec![ + Field::new("name", DataType::Utf8, true), + Field::new("id", DataType::Int32, false), + ] + .into(), + ); + let cast_type = + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(target_value_type)); + + let cast_array = cast(&source, &cast_type).unwrap(); + let dict = cast_array.as_dictionary::(); + assert_eq!(dict.len(), 3); + let keys = dict.keys(); + assert!(!keys.is_null(0)); + assert!(keys.is_null(1)); + assert!(!keys.is_null(2)); + } + + #[test] + fn test_cast_struct_array_to_dict_struct_key_overflow() { + // Source has 300 rows but the dictionary key type is UInt8 (max 255). + // We must return a CastError instead of silently truncating. + let n = 300; + let names = StringArray::from((0..n).map(|i| Some(format!("v{i}"))).collect::>()); + let source = StructArray::from(vec![( + Arc::new(Field::new("name", DataType::Utf8, true)), + Arc::new(names) as ArrayRef, + )]); + + let cast_type = DataType::Dictionary( + Box::new(DataType::UInt8), + Box::new(DataType::Struct( + vec![Field::new("name", DataType::Utf8, true)].into(), + )), + ); + let err = cast(&source, &cast_type).unwrap_err().to_string(); + assert!( + err.contains("Cannot fit") && err.contains("dictionary keys"), + "expected key-overflow error, got: {err}" + ); + } + #[test] fn test_cast_empty_string_array_to_dict_utf8_view() { let array = StringArray::from(Vec::>::new()); From 870fb06f73c45d9d09565d7b83cd83139a74424f Mon Sep 17 00:00:00 2001 From: Soiman Vasile Date: Tue, 2 Jun 2026 18:36:07 +0300 Subject: [PATCH 37/51] test: add overflow tests for MutableBuffer (#9958) ### Description While analyzing the memory allocation logic in `MutableBuffer`, I identified that the `with_capacity` and `reserve` methods correctly use `.expect()` guards to prevent integer overflows. However, I couldnt find test cases for this `.expect()` guard in the current test suite. This PR adds 2 `#[should_panic]` tests in `mutable.rs` to verify that the API correctly panics **Changes:** * Added `test_mutable_new_capacity_overflow` to cover `MutableBuffer::new` * Added `test_mutable_reserve_overflow` to cover `MutableBuffer::reserve` ### Rationale Adding these tests ensures that the safety guards in `arrow-buffer` remain intact and provides explicit coverage for edge cases involving near-`usize::MAX` allocations. ### Tests - `test_mutable_new_capacity_overflow` - `test_mutable_reserve_overflow` Co-authored-by: Andrew Lamb --- arrow-buffer/src/buffer/mutable.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs index aff1f138c286..b6e6a70c6cba 100644 --- a/arrow-buffer/src/buffer/mutable.rs +++ b/arrow-buffer/src/buffer/mutable.rs @@ -1669,4 +1669,20 @@ mod tests { let data_1000: Vec = (0..1000).collect(); test_repeat_count(repeat_count, &data_1000); } + + #[test] + #[should_panic(expected = "failed to round upto multiple of 64")] + fn test_mutable_new_capacity_overflow() { + // Tests overflow during initial allocation + let _ = MutableBuffer::new(usize::MAX - 10); + } + + #[test] + #[should_panic(expected = "buffer length overflow")] + fn test_mutable_reserve_overflow() { + // Tests overflow during growth (checked_add) + let mut buf = MutableBuffer::new(1); + buf.push(1u8); + buf.reserve(usize::MAX); + } } From 1ae246942d09888633338bc623dcc53b07ba9d75 Mon Sep 17 00:00:00 2001 From: Ed Seidl Date: Tue, 2 Jun 2026 11:12:12 -0700 Subject: [PATCH 38/51] chore: Remove some deprecated Arrow functions from the public API (#10040) # Which issue does this PR close? None - related to #9110 # Rationale for this change Housecleaning for 59.0.0 # What changes are included in this PR? Remove some deprecated functions from public Arrow APIs. # Are these changes tested? Covered by existing tests # Are there any user-facing changes? Yes, deprecated public functions are removed --------- Co-authored-by: Andrew Lamb --- arrow-array/src/temporal_conversions.rs | 14 --- arrow-array/src/types.rs | 127 ------------------------ arrow-buffer/src/buffer/immutable.rs | 11 -- arrow-data/src/data.rs | 6 -- arrow-data/src/decimal.rs | 100 ------------------- arrow-ipc/src/reader.rs | 4 +- arrow-ipc/src/reader/stream.rs | 2 +- arrow-schema/src/fields.rs | 53 +--------- 8 files changed, 5 insertions(+), 312 deletions(-) diff --git a/arrow-array/src/temporal_conversions.rs b/arrow-array/src/temporal_conversions.rs index a5ec50da1fc6..375c19bc8722 100644 --- a/arrow-array/src/temporal_conversions.rs +++ b/arrow-array/src/temporal_conversions.rs @@ -215,26 +215,12 @@ pub(crate) fn split_second(v: i64, base: i64) -> (i64, u32) { (v.div_euclid(base), v.rem_euclid(base) as u32) } -/// converts a `i64` representing a `duration(s)` to [`Duration`] -#[inline] -#[deprecated(since = "55.2.0", note = "Use `try_duration_s_to_duration` instead")] -pub fn duration_s_to_duration(v: i64) -> Duration { - Duration::try_seconds(v).unwrap() -} - /// converts a `i64` representing a `duration(s)` to [`Option`] #[inline] pub fn try_duration_s_to_duration(v: i64) -> Option { Duration::try_seconds(v) } -/// converts a `i64` representing a `duration(ms)` to [`Duration`] -#[inline] -#[deprecated(since = "55.2.0", note = "Use `try_duration_ms_to_duration` instead")] -pub fn duration_ms_to_duration(v: i64) -> Duration { - Duration::try_seconds(v).unwrap() -} - /// converts a `i64` representing a `duration(ms)` to [`Option`] #[inline] pub fn try_duration_ms_to_duration(v: i64) -> Option { diff --git a/arrow-array/src/types.rs b/arrow-array/src/types.rs index 267011d8af80..a8de9267b3e6 100644 --- a/arrow-array/src/types.rs +++ b/arrow-array/src/types.rs @@ -1205,17 +1205,6 @@ impl Date32Type { } impl Date64Type { - /// Converts an arrow Date64Type into a chrono::NaiveDate - /// - /// # Arguments - /// - /// * `i` - The Date64Type to convert - #[deprecated(since = "56.0.0", note = "Use to_naive_date_opt instead.")] - pub fn to_naive_date(i: ::Native) -> NaiveDate { - Self::to_naive_date_opt(i) - .unwrap_or_else(|| panic!("Date64Type::to_naive_date overflowed for date: {i}",)) - } - /// Converts an arrow Date64Type into a chrono::NaiveDateTime if it fits in the range that chrono::NaiveDateTime can represent. /// Returns `None` if the calculation would overflow or underflow. /// @@ -1241,25 +1230,6 @@ impl Date64Type { d.sub(epoch).num_milliseconds() as ::Native } - /// Adds the given IntervalYearMonthType to an arrow Date64Type - /// - /// # Arguments - /// - /// * `date` - The date on which to perform the operation - /// * `delta` - The interval to add - #[deprecated( - since = "56.0.0", - note = "Use `add_year_months_opt` instead, which returns an Option to handle overflow." - )] - pub fn add_year_months( - date: ::Native, - delta: ::Native, - ) -> ::Native { - Self::add_year_months_opt(date, delta).unwrap_or_else(|| { - panic!("Date64Type::add_year_months overflowed for date: {date}, delta: {delta}",) - }) - } - /// Adds the given IntervalYearMonthType to an arrow Date64Type /// /// # Arguments @@ -1278,25 +1248,6 @@ impl Date64Type { Some(Date64Type::from_naive_date(posterior)) } - /// Adds the given IntervalDayTimeType to an arrow Date64Type - /// - /// # Arguments - /// - /// * `date` - The date on which to perform the operation - /// * `delta` - The interval to add - #[deprecated( - since = "56.0.0", - note = "Use `add_day_time_opt` instead, which returns an Option to handle overflow." - )] - pub fn add_day_time( - date: ::Native, - delta: ::Native, - ) -> ::Native { - Self::add_day_time_opt(date, delta).unwrap_or_else(|| { - panic!("Date64Type::add_day_time overflowed for date: {date}, delta: {delta:?}",) - }) - } - /// Adds the given IntervalDayTimeType to an arrow Date64Type /// /// # Arguments @@ -1316,25 +1267,6 @@ impl Date64Type { Some(Date64Type::from_naive_date(res)) } - /// Adds the given IntervalMonthDayNanoType to an arrow Date64Type - /// - /// # Arguments - /// - /// * `date` - The date on which to perform the operation - /// * `delta` - The interval to add - #[deprecated( - since = "56.0.0", - note = "Use `add_month_day_nano_opt` instead, which returns an Option to handle overflow." - )] - pub fn add_month_day_nano( - date: ::Native, - delta: ::Native, - ) -> ::Native { - Self::add_month_day_nano_opt(date, delta).unwrap_or_else(|| { - panic!("Date64Type::add_month_day_nano overflowed for date: {date}, delta: {delta:?}",) - }) - } - /// Adds the given IntervalMonthDayNanoType to an arrow Date64Type /// /// # Arguments @@ -1355,25 +1287,6 @@ impl Date64Type { Some(Date64Type::from_naive_date(res)) } - /// Subtract the given IntervalYearMonthType to an arrow Date64Type - /// - /// # Arguments - /// - /// * `date` - The date on which to perform the operation - /// * `delta` - The interval to subtract - #[deprecated( - since = "56.0.0", - note = "Use `subtract_year_months_opt` instead, which returns an Option to handle overflow." - )] - pub fn subtract_year_months( - date: ::Native, - delta: ::Native, - ) -> ::Native { - Self::subtract_year_months_opt(date, delta).unwrap_or_else(|| { - panic!("Date64Type::subtract_year_months overflowed for date: {date}, delta: {delta}",) - }) - } - /// Subtract the given IntervalYearMonthType to an arrow Date64Type /// /// # Arguments @@ -1392,25 +1305,6 @@ impl Date64Type { Some(Date64Type::from_naive_date(posterior)) } - /// Subtract the given IntervalDayTimeType to an arrow Date64Type - /// - /// # Arguments - /// - /// * `date` - The date on which to perform the operation - /// * `delta` - The interval to subtract - #[deprecated( - since = "56.0.0", - note = "Use `subtract_day_time_opt` instead, which returns an Option to handle overflow." - )] - pub fn subtract_day_time( - date: ::Native, - delta: ::Native, - ) -> ::Native { - Self::subtract_day_time_opt(date, delta).unwrap_or_else(|| { - panic!("Date64Type::subtract_day_time overflowed for date: {date}, delta: {delta:?}",) - }) - } - /// Subtract the given IntervalDayTimeType to an arrow Date64Type /// /// # Arguments @@ -1430,27 +1324,6 @@ impl Date64Type { Some(Date64Type::from_naive_date(res)) } - /// Subtract the given IntervalMonthDayNanoType to an arrow Date64Type - /// - /// # Arguments - /// - /// * `date` - The date on which to perform the operation - /// * `delta` - The interval to subtract - #[deprecated( - since = "56.0.0", - note = "Use `subtract_month_day_nano_opt` instead, which returns an Option to handle overflow." - )] - pub fn subtract_month_day_nano( - date: ::Native, - delta: ::Native, - ) -> ::Native { - Self::subtract_month_day_nano_opt(date, delta).unwrap_or_else(|| { - panic!( - "Date64Type::subtract_month_day_nano overflowed for date: {date}, delta: {delta:?}", - ) - }) - } - /// Subtract the given IntervalMonthDayNanoType to an arrow Date64Type /// /// # Arguments diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index a73cc55086ac..fc001afe5a07 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -103,17 +103,6 @@ unsafe impl Send for Buffer where Bytes: Send {} unsafe impl Sync for Buffer where Bytes: Sync {} impl Buffer { - /// Create a new Buffer from a (internal) `Bytes` - /// - /// NOTE despite the same name, `Bytes` is an internal struct in arrow-rs - /// and is different than [`bytes::Bytes`]. - /// - /// See examples on [`Buffer`] for ways to create a buffer from a [`bytes::Bytes`]. - #[deprecated(since = "54.1.0", note = "Use Buffer::from instead")] - pub fn from_bytes(bytes: Bytes) -> Self { - Self::from(bytes) - } - /// Returns the offset, in bytes, of `Self::ptr` to `Self::data` /// /// self.ptr and self.data can be different after slicing or advancing the buffer. diff --git a/arrow-data/src/data.rs b/arrow-data/src/data.rs index 7b208b7af0be..0ce98aa090df 100644 --- a/arrow-data/src/data.rs +++ b/arrow-data/src/data.rs @@ -2210,12 +2210,6 @@ impl ArrayDataBuilder { Ok(data) } - /// Creates an array data, validating all inputs, and aligning any buffers - #[deprecated(since = "54.1.0", note = "Use ArrayData::align_buffers instead")] - pub fn build_aligned(self) -> Result { - self.align_buffers(true).build() - } - /// Ensure that all buffers are aligned, copying data if necessary /// /// Rust requires that arrays are aligned to their corresponding primitive, diff --git a/arrow-data/src/decimal.rs b/arrow-data/src/decimal.rs index 2c997753bd5f..a5986755ad6e 100644 --- a/arrow-data/src/decimal.rs +++ b/arrow-data/src/decimal.rs @@ -688,106 +688,6 @@ pub const MIN_DECIMAL256_FOR_EACH_PRECISION: [i256; 77] = [ ]), ]; -/// `MAX_DECIMAL_FOR_EACH_PRECISION[p-1]` holds the maximum `i128` value that can -/// be stored in a [`Decimal128`] value of precision `p` -/// -/// [`Decimal128`]: arrow_schema::DataType::Decimal128 -#[deprecated( - since = "54.1.0", - note = "Use MAX_DECIMAL128_FOR_EACH_PRECISION (note indexes are different)" -)] -#[allow(dead_code)] // no longer used but is part of our public API -pub const MAX_DECIMAL_FOR_EACH_PRECISION: [i128; 38] = [ - 9, - 99, - 999, - 9999, - 99999, - 999999, - 9999999, - 99999999, - 999999999, - 9999999999, - 99999999999, - 999999999999, - 9999999999999, - 99999999999999, - 999999999999999, - 9999999999999999, - 99999999999999999, - 999999999999999999, - 9999999999999999999, - 99999999999999999999, - 999999999999999999999, - 9999999999999999999999, - 99999999999999999999999, - 999999999999999999999999, - 9999999999999999999999999, - 99999999999999999999999999, - 999999999999999999999999999, - 9999999999999999999999999999, - 99999999999999999999999999999, - 999999999999999999999999999999, - 9999999999999999999999999999999, - 99999999999999999999999999999999, - 999999999999999999999999999999999, - 9999999999999999999999999999999999, - 99999999999999999999999999999999999, - 999999999999999999999999999999999999, - 9999999999999999999999999999999999999, - 99999999999999999999999999999999999999, -]; - -/// `MIN_DECIMAL_FOR_EACH_PRECISION[p-1]` holds the minimum `i128` value that can -/// be stored in a [`Decimal128`] value of precision `p` -/// -/// [`Decimal128`]: arrow_schema::DataType::Decimal128 -#[allow(dead_code)] // no longer used but is part of our public API -#[deprecated( - since = "54.1.0", - note = "Use MIN_DECIMAL128_FOR_EACH_PRECISION (note indexes are different)" -)] -pub const MIN_DECIMAL_FOR_EACH_PRECISION: [i128; 38] = [ - -9, - -99, - -999, - -9999, - -99999, - -999999, - -9999999, - -99999999, - -999999999, - -9999999999, - -99999999999, - -999999999999, - -9999999999999, - -99999999999999, - -999999999999999, - -9999999999999999, - -99999999999999999, - -999999999999999999, - -9999999999999999999, - -99999999999999999999, - -999999999999999999999, - -9999999999999999999999, - -99999999999999999999999, - -999999999999999999999999, - -9999999999999999999999999, - -99999999999999999999999999, - -999999999999999999999999999, - -9999999999999999999999999999, - -99999999999999999999999999999, - -999999999999999999999999999999, - -9999999999999999999999999999999, - -99999999999999999999999999999999, - -999999999999999999999999999999999, - -9999999999999999999999999999999999, - -99999999999999999999999999999999999, - -999999999999999999999999999999999999, - -9999999999999999999999999999999999999, - -99999999999999999999999999999999999999, -]; - /// `MAX_DECIMAL128_FOR_EACH_PRECISION[p]` holds the maximum `i128` value that can /// be stored in [`Decimal128`] value of precision `p`. /// diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs index 1d5e06c6871c..6d1466febf0f 100644 --- a/arrow-ipc/src/reader.rs +++ b/arrow-ipc/src/reader.rs @@ -754,7 +754,7 @@ impl<'a> RecordBatchDecoder<'a> { /// If `require_alignment` is false, this function will automatically allocate a new aligned buffer /// and copy over the data if any array data in the input `buf` is not properly aligned. /// (Properly aligned array data will remain zero-copy.) -/// Under the hood it will use [`arrow_data::ArrayDataBuilder::build_aligned`] to construct [`arrow_data::ArrayData`]. +/// Under the hood it will use [`arrow_data::ArrayDataBuilder::align_buffers`] to construct [`arrow_data::ArrayData`]. pub fn read_record_batch( buf: &Buffer, batch: crate::RecordBatch, @@ -1046,7 +1046,7 @@ impl FileDecoder { /// If `require_alignment` is false (the default), this decoder will automatically allocate a /// new aligned buffer and copy over the data if any array data in the input `buf` is not /// properly aligned. (Properly aligned array data will remain zero-copy.) - /// Under the hood it will use [`arrow_data::ArrayDataBuilder::build_aligned`] to construct + /// Under the hood it will use [`arrow_data::ArrayDataBuilder::align_buffers`] to construct /// [`arrow_data::ArrayData`]. pub fn with_require_alignment(mut self, require_alignment: bool) -> Self { self.require_alignment = require_alignment; diff --git a/arrow-ipc/src/reader/stream.rs b/arrow-ipc/src/reader/stream.rs index bfecf7b6ffa5..e18d32cea2f6 100644 --- a/arrow-ipc/src/reader/stream.rs +++ b/arrow-ipc/src/reader/stream.rs @@ -101,7 +101,7 @@ impl StreamDecoder { /// If `require_alignment` is false (the default), this decoder will automatically allocate a /// new aligned buffer and copy over the data if any array data in the input `buf` is not /// properly aligned. (Properly aligned array data will remain zero-copy.) - /// Under the hood it will use [`arrow_data::ArrayDataBuilder::build_aligned`] to construct + /// Under the hood it will use [`arrow_data::ArrayDataBuilder::align_buffers`] to construct /// [`arrow_data::ArrayData`]. pub fn with_require_alignment(mut self, require_alignment: bool) -> Self { self.require_alignment = require_alignment; diff --git a/arrow-schema/src/fields.rs b/arrow-schema/src/fields.rs index 93638181d9ae..3d19cb963f96 100644 --- a/arrow-schema/src/fields.rs +++ b/arrow-schema/src/fields.rs @@ -534,55 +534,6 @@ impl UnionFields { Ok(Self(out.into())) } - /// Create a new [`UnionFields`] from a [`Fields`] and array of type_ids - /// - /// See - /// - /// # Deprecated - /// - /// Use [`UnionFields::try_new`] instead. This method panics on invalid input, - /// while `try_new` returns a `Result`. - /// - /// # Panics - /// - /// Panics if any type_id appears more than once (duplicate type ids). - /// - /// ``` - /// use arrow_schema::{DataType, Field, UnionFields}; - /// // Create a new UnionFields with type id mapping - /// // 1 -> DataType::UInt8 - /// // 3 -> DataType::Utf8 - /// UnionFields::try_new( - /// vec![1, 3], - /// vec![ - /// Field::new("field1", DataType::UInt8, false), - /// Field::new("field3", DataType::Utf8, false), - /// ], - /// ); - /// ``` - #[deprecated(since = "57.0.0", note = "Use `try_new` instead")] - pub fn new(type_ids: T, fields: F) -> Self - where - F: IntoIterator, - F::Item: Into, - T: IntoIterator, - { - let fields = fields.into_iter().map(Into::into); - let mut set = 0_u128; - type_ids - .into_iter() - .inspect(|&idx| { - let mask = 1_u128 << idx; - if (set & mask) != 0 { - panic!("duplicate type id: {idx}"); - } else { - set |= mask; - } - }) - .zip(fields) - .collect() - } - /// Return size of this instance in bytes. pub fn size(&self) -> usize { self.iter() @@ -614,13 +565,13 @@ impl UnionFields { /// ``` /// use arrow_schema::{DataType, Field, UnionFields}; /// - /// let fields = UnionFields::new( + /// let fields = UnionFields::try_new( /// vec![1, 3], /// vec![ /// Field::new("field1", DataType::UInt8, false), /// Field::new("field3", DataType::Utf8, false), /// ], - /// ); + /// ).unwrap(); /// /// assert!(fields.get(0).is_some()); /// assert!(fields.get(1).is_some()); From cfc2b88d1d4e95e2de0dc5eeba4c23694058a903 Mon Sep 17 00:00:00 2001 From: ClSlaid Date: Wed, 3 Jun 2026 05:09:37 +0800 Subject: [PATCH 39/51] Add coalesce inline-view filter benchmarks (#10050) This is a benchmark-only companion patch for https://github.com/apache/arrow-rs/pull/9755. It keeps the functional changes out of this PR and only adds benchmark coverage in `arrow/benches/coalesce_kernels.rs` so the coalesce inline-view filter work can be tested independently. Benchmark coverage included: - filter and take coalesce benchmarks - primitive schemas - single-column `Utf8View` and `BinaryView` - mixed primitive + `Utf8View` and primitive + `BinaryView` schemas - filter cases for short inline strings with `max_string_len=8` - filter/take cases for longer view strings, including `max_string_len=20`, `30`, and `128` depending on scenario Coverage note: - The filter benchmarks cover the main short-inline path targeted by #9755 for both `Utf8View` and `BinaryView`. - The take benchmarks cover `Utf8View`/`BinaryView` and mixed schemas, but do not add `max_string_len=8` take variants. This patch keeps the benchmark changes aligned with the benchmark patch currently carried by #9755. Validation: ```text cargo fmt --package arrow cargo bench --bench coalesce_kernels -- --list git diff --check ``` --- arrow/benches/coalesce_kernels.rs | 98 +++++++++++++++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/arrow/benches/coalesce_kernels.rs b/arrow/benches/coalesce_kernels.rs index 0816d1a2e8bb..4affcc346e89 100644 --- a/arrow/benches/coalesce_kernels.rs +++ b/arrow/benches/coalesce_kernels.rs @@ -51,6 +51,13 @@ fn add_all_filter_benchmarks(c: &mut Criterion) { true, )])); + // Single BinaryViewArray + let single_binaryview_schema = SchemaRef::new(Schema::new(vec![Field::new( + "value", + DataType::BinaryView, + true, + )])); + // Mixed primitive, StringViewArray let mixed_utf8view_schema = SchemaRef::new(Schema::new(vec![ Field::new("int32_val", DataType::Int32, true), @@ -58,6 +65,13 @@ fn add_all_filter_benchmarks(c: &mut Criterion) { Field::new("utf8view_val", DataType::Utf8View, true), ])); + // Mixed primitive, BinaryViewArray + let mixed_binaryview_schema = SchemaRef::new(Schema::new(vec![ + Field::new("int32_val", DataType::Int32, true), + Field::new("float_val", DataType::Float64, true), + Field::new("binaryview_val", DataType::BinaryView, true), + ])); + // Mixed primitive, StringArray let mixed_utf8_schema = SchemaRef::new(Schema::new(vec![ Field::new("int32_val", DataType::Int32, true), @@ -106,6 +120,42 @@ fn add_all_filter_benchmarks(c: &mut Criterion) { } .build(); + FilterBenchmarkBuilder { + c, + name: "single_utf8view (max_string_len=8)", + batch_size, + num_output_batches: 50, + null_density, + selectivity, + max_string_len: 8, + schema: &single_schema, + } + .build(); + + FilterBenchmarkBuilder { + c, + name: "single_binaryview", + batch_size, + num_output_batches: 50, + null_density, + selectivity, + max_string_len: 30, + schema: &single_binaryview_schema, + } + .build(); + + FilterBenchmarkBuilder { + c, + name: "single_binaryview (max_string_len=8)", + batch_size, + num_output_batches: 50, + null_density, + selectivity, + max_string_len: 8, + schema: &single_binaryview_schema, + } + .build(); + // Model mostly short strings, but some longer ones FilterBenchmarkBuilder { c, @@ -119,6 +169,18 @@ fn add_all_filter_benchmarks(c: &mut Criterion) { } .build(); + FilterBenchmarkBuilder { + c, + name: "mixed_utf8view (max_string_len=8)", + batch_size, + num_output_batches: 20, + null_density, + selectivity, + max_string_len: 8, + schema: &mixed_utf8view_schema, + } + .build(); + // Model mostly longer strings FilterBenchmarkBuilder { c, @@ -132,6 +194,42 @@ fn add_all_filter_benchmarks(c: &mut Criterion) { } .build(); + FilterBenchmarkBuilder { + c, + name: "mixed_binaryview (max_string_len=20)", + batch_size, + num_output_batches: 20, + null_density, + selectivity, + max_string_len: 20, + schema: &mixed_binaryview_schema, + } + .build(); + + FilterBenchmarkBuilder { + c, + name: "mixed_binaryview (max_string_len=8)", + batch_size, + num_output_batches: 20, + null_density, + selectivity, + max_string_len: 8, + schema: &mixed_binaryview_schema, + } + .build(); + + FilterBenchmarkBuilder { + c, + name: "mixed_binaryview (max_string_len=128)", + batch_size, + num_output_batches: 20, + null_density, + selectivity, + max_string_len: 128, + schema: &mixed_binaryview_schema, + } + .build(); + FilterBenchmarkBuilder { c, name: "mixed_utf8", From 259cff297c5a6b1015ad9ee02ebeb61b53f39a70 Mon Sep 17 00:00:00 2001 From: Adam Reeve Date: Wed, 3 Jun 2026 11:49:20 +1200 Subject: [PATCH 40/51] feat(parquet): add uses_key_retriever method to FileDecryptionProperties (#9895) # Which issue does this PR close? - Closes #9721. # Are these changes tested? Yes, includes new unit tests. # Are there any user-facing changes? Yes, there is a new public API method. --- parquet/src/encryption/decrypt.rs | 8 +++++++ parquet/tests/encryption/encryption.rs | 33 ++++++++++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/parquet/src/encryption/decrypt.rs b/parquet/src/encryption/decrypt.rs index 0066523419de..bfa587277908 100644 --- a/parquet/src/encryption/decrypt.rs +++ b/parquet/src/encryption/decrypt.rs @@ -432,6 +432,14 @@ impl FileDecryptionProperties { } (column_names, column_keys) } + + /// Whether these decryption properties use a key retriever. + /// When false, explicit keys were provided up front and can + /// be retrieved without providing key metadata, rather than + /// resolved on demand. + pub fn uses_key_retriever(&self) -> bool { + matches!(self.keys, DecryptionKeys::ViaRetriever(_)) + } } impl std::fmt::Debug for FileDecryptionProperties { diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index edd26f29619c..2ac94b8c3786 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -1474,3 +1474,36 @@ fn test_decrypt_page_index( Ok(()) } + +#[test] +fn test_decryption_properties_uses_key_retriever() { + let key_retriever = TestKeyRetriever::new() + .with_key( + AES_128_FOOTER_KEY_NAME.to_owned(), + AES_128_FOOTER_KEY.to_vec(), + ) + .with_key( + AES_128_KEY_NAMES[0].to_owned(), + AES_128_COLUMN_KEYS[0].to_vec(), + ); + + let properties_with_retriever = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + assert!(properties_with_retriever.uses_key_retriever()); + + let properties_with_keys = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.to_vec()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].to_vec()) + .build() + .unwrap(); + + assert!(!properties_with_keys.uses_key_retriever()); + + let uniform_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.to_vec()) + .build() + .unwrap(); + + assert!(!uniform_properties.uses_key_retriever()); +} From f03e1bc028630561af7d4adfc000916337ec6ce5 Mon Sep 17 00:00:00 2001 From: pchintar <89355405+pchintar@users.noreply.github.com> Date: Wed, 3 Jun 2026 06:06:16 +0530 Subject: [PATCH 41/51] fix(ipc): handle duplicate projection indices in IPC reader (#9952) # Which issue does this PR close? - Closes #9950 . # Rationale for this change The current IPC reader does not correctly handle duplicate projection indices. `Schema::project`(in `arrow-schema/src/schema.rs`) and `RecordBatch::project`(in `arrow-array/src/record_batch.rs`) both map each requested index directly, preserve the projection order and allow duplicate indices such as: ```rust id="n4pq0f" vec![1, 1] ``` However, the IPC reader currently uses: ```rust id="gjklyo" projection.iter().position(|p| p == &idx) ``` which only returns the first matching entry. As a result, only one column is decoded even though the projected schema contains multiple fields, leading to schema/column count mismatches when constructing the `RecordBatch`. This also affects reordered duplicate projections such as: ```rust id="jlwmku" vec![2, 0, 2] ``` # What changes are included in this PR? * Updated IPC projection handling in `arrow-ipc/src/reader.rs` to preserve all matching projection entries * Reused the decoded array for duplicate projection indices instead of decoding the same field multiple times * Preserved projection order for reordered duplicate projections # Are these changes tested? Yes. Added `test_projection_duplicate_indices`, which verifies: * duplicate projections (`vec![1, 1]`) * reordered duplicate projections (`vec![2, 0, 2]`) The test compares IPC projection results against `RecordBatch::project`. The test fails before the fix and passes after it. All existing `arrow-ipc` tests also pass `cargo test -p arrow-ipc --lib` # Are there any user-facing changes? No. --- arrow-ipc/src/reader.rs | 43 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 38 insertions(+), 5 deletions(-) diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs index 6d1466febf0f..6d1e799d43c9 100644 --- a/arrow-ipc/src/reader.rs +++ b/arrow-ipc/src/reader.rs @@ -560,11 +560,20 @@ impl<'a> RecordBatchDecoder<'a> { let mut arrays = vec![]; // project fields for (idx, field) in schema.fields().iter().enumerate() { - // Create array for projected field - if let Some(proj_idx) = projection.iter().position(|p| p == &idx) { - let child = self.create_array(field, &mut variadic_counts)?; - arrays.push((proj_idx, child)); - } else { + // A projected field can appear more than once, so collect all matching positions. + let mut child = None; + for (proj_idx, projected_idx) in projection.iter().enumerate() { + if *projected_idx == idx { + if child.is_none() { + child = Some(self.create_array(field, &mut variadic_counts)?); + } + + // Reuse the decoded array for duplicate projection entries. + arrays.push((proj_idx, child.as_ref().unwrap().clone())); + } + } + + if child.is_none() { self.skip_field(field, &mut variadic_counts)?; } } @@ -2297,6 +2306,30 @@ mod tests { } } + #[test] + fn test_projection_duplicate_indices() { + let schema = create_test_projection_schema(); + let batch = create_test_projection_batch_data(&schema); + + // Write the batch to IPC + let mut buf = Vec::new(); + { + let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + + // Verify duplicate([1, 1]) and reordered([2, 0, 2]) projection indices + for projection in [vec![1, 1], vec![2, 0, 2]] { + let reader = + FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection.clone())); + let read_batch = reader.unwrap().next().unwrap().unwrap(); + + let expected_batch = batch.project(&projection).unwrap(); + assert_eq!(read_batch, expected_batch); + } + } + #[test] fn test_arrow_single_float_row() { let schema = Schema::new(vec![ From f001b222483ad2a18e6bb4563c3505e9f4280f66 Mon Sep 17 00:00:00 2001 From: RIchard Baah <137434454+Rich-T-kid@users.noreply.github.com> Date: Wed, 3 Jun 2026 07:34:00 -0400 Subject: [PATCH 42/51] [#10029][benchmarks] arrow-flight roundtrip as well as encode/decode (#10031) # Which issue does this PR close? - Contributes towards closing #10029. # Rationale for this change Provides benchmarks for arrow-flight crate. benchmarks for round trip as well as encode/decode individually. # What changes are included in this PR? Adds three criterion benches under arrow-flight/benchmarks/ (roundtrip.rs, flight_encode.rs, flight_decode.rs), each sweeping a tunable matrix of rows, cols, and column types (fixed Int64, variable StringArray, nested List, dict DictionaryArray) built via a shared common::build_batch helper. # Are these changes tested? n/a # Are there any user-facing changes? no --- Cargo.lock | 2 + arrow-flight/Cargo.toml | 6 ++ arrow-flight/benches/common/mod.rs | 154 +++++++++++++++++++++++++++++ arrow-flight/benches/flight.rs | 87 ++++++++++++++++ 4 files changed, 249 insertions(+) create mode 100644 arrow-flight/benches/common/mod.rs create mode 100644 arrow-flight/benches/flight.rs diff --git a/Cargo.lock b/Cargo.lock index 28fe6a43a5c5..9e5963a70a0b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -348,6 +348,7 @@ dependencies = [ "base64", "bytes", "clap", + "criterion", "futures", "http", "http-body", @@ -1063,6 +1064,7 @@ dependencies = [ "serde", "serde_json", "tinytemplate", + "tokio", "walkdir", ] diff --git a/arrow-flight/Cargo.toml b/arrow-flight/Cargo.toml index 8f95e1995a67..8e399fbc5a52 100644 --- a/arrow-flight/Cargo.toml +++ b/arrow-flight/Cargo.toml @@ -76,6 +76,7 @@ cli = ["arrow-array/chrono-tz", "arrow-cast/prettyprint", "tonic/tls-webpki-root [dev-dependencies] arrow-cast = { workspace = true, features = ["prettyprint"] } assert_cmd = "2.0.8" +criterion = { workspace = true, default-features = false, features = ["async_tokio"] } http = "1.1.0" http-body = "1.0.0" hyper-util = "0.1" @@ -105,3 +106,8 @@ required-features = ["flight-sql", "tls-ring"] name = "flight_sql_client_cli" path = "tests/flight_sql_client_cli.rs" required-features = ["cli", "flight-sql", "tls-ring"] + +[[bench]] +name = "flight" +path = "benches/flight.rs" +harness = false \ No newline at end of file diff --git a/arrow-flight/benches/common/mod.rs b/arrow-flight/benches/common/mod.rs new file mode 100644 index 000000000000..eb8bea8dc591 --- /dev/null +++ b/arrow-flight/benches/common/mod.rs @@ -0,0 +1,154 @@ +// 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 std::sync::{Arc, RwLock}; + +use arrow_array::{ + Array, ArrayRef, DictionaryArray, Int32Array, Int64Array, ListArray, RecordBatch, StringArray, + types::Int32Type, +}; +use arrow_buffer::OffsetBuffer; +use arrow_flight::{ + Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo, + HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket, + flight_service_server::{FlightService, FlightServiceServer}, +}; +use arrow_schema::{DataType, Field, Schema}; +use bytes::Bytes; +use futures::{StreamExt, TryStreamExt, stream::BoxStream}; +use hyper_util::rt::TokioIo; +use tonic::{ + Request, Response, Status, Streaming, + transport::{Channel, Endpoint, Server}, +}; + +pub type Builder = fn(usize) -> ArrayRef; + +pub const TYPES: &[(&str, Builder)] = &[ + ("fixed", fixed), + ("nested", nested), + ("variable", variable), + ("dict", dict), +]; + +fn fixed(n: usize) -> ArrayRef { + Arc::new(Int64Array::from_iter_values(0..n as i64)) +} + +fn variable(n: usize) -> ArrayRef { + Arc::new(StringArray::from_iter_values( + (0..n).map(|i| format!("variable_string_{i}{}", "_".repeat(i % 16))), + )) +} + +fn nested(n: usize) -> ArrayRef { + let values = Int32Array::from_iter_values(0..(n * 4) as i32); + let offsets = OffsetBuffer::::from_lengths(std::iter::repeat_n(4usize, n)); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + Arc::new(ListArray::new(field, offsets, Arc::new(values), None)) +} + +fn dict(n: usize) -> ArrayRef { + let keys = Int32Array::from_iter_values((0..n).map(|i| (i % 32) as i32)); + let values = StringArray::from_iter_values((0..32).map(|i| format!("dictionary_value_{i:03}"))); + Arc::new(DictionaryArray::::try_new(keys, Arc::new(values)).unwrap()) +} + +pub fn build_batch(name: &str, rows: usize, cols: usize, build: Builder) -> RecordBatch { + let arrays: Vec = (0..cols).map(|_| build(rows)).collect(); + let fields: Vec = arrays + .iter() + .enumerate() + .map(|(i, a)| Field::new(format!("column_{i}_{name}"), a.data_type().clone(), false)) + .collect(); + RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap() +} + +#[derive(Clone, Default)] +pub struct BenchServer { + frames: Arc>>, +} + +impl BenchServer { + #[allow(dead_code)] + pub fn set_frames(&self, frames: Vec) { + *self.frames.write().unwrap() = frames; + } +} + +fn unimpl() -> Result { + Err(Status::unimplemented("")) +} + +#[rustfmt::skip] +#[tonic::async_trait] +impl FlightService for BenchServer { + type HandshakeStream = BoxStream<'static, Result>; + type ListFlightsStream = BoxStream<'static, Result>; + type DoGetStream = BoxStream<'static, Result>; + type DoPutStream = BoxStream<'static, Result>; + type DoActionStream = BoxStream<'static, Result>; + type ListActionsStream = BoxStream<'static, Result>; + type DoExchangeStream = BoxStream<'static, Result>; + + async fn do_get(&self, _: Request) -> Result, Status> { + let frames = self.frames.read().unwrap().clone(); + Ok(Response::new(futures::stream::iter(frames.into_iter().map(Ok)).boxed())) + } + + async fn do_put(&self, req: Request>) -> Result, Status> { + let _: Vec = req.into_inner().try_collect().await?; + let ack = PutResult { app_metadata: Bytes::new() }; + Ok(Response::new(futures::stream::iter([Ok(ack)]).boxed())) + } + + async fn do_exchange(&self, req: Request>) -> Result, Status> { + Ok(Response::new(req.into_inner().boxed())) + } + + async fn handshake(&self, _: Request>) -> Result, Status> { unimpl() } + async fn list_flights(&self, _: Request) -> Result, Status> { unimpl() } + async fn get_flight_info(&self, _: Request) -> Result, Status> { unimpl() } + async fn poll_flight_info(&self, _: Request) -> Result, Status> { unimpl() } + async fn get_schema(&self, _: Request) -> Result, Status> { unimpl() } + async fn do_action(&self, _: Request) -> Result, Status> { unimpl() } + async fn list_actions(&self, _: Request) -> Result, Status> { unimpl() } +} +#[allow(dead_code)] +pub async fn start_server() -> (Channel, BenchServer) { + const DUMMY_URL: &str = "http://localhost:50051"; + + let bench_server = BenchServer::default(); + + let (client, server) = tokio::io::duplex(1024 * 1024); + + let mut client = Some(client); + let channel = Endpoint::try_from(DUMMY_URL) + .expect("Invalid dummy URL for building an endpoint. This should never happen") + .connect_with_connector_lazy(tower::service_fn(move |_| { + let client = client + .take() + .expect("Client taken twice. This should never happen"); + async move { Ok::<_, std::io::Error>(TokioIo::new(client)) } + })); + tokio::spawn( + Server::builder() + .add_service(FlightServiceServer::new(bench_server.clone())) + .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server))), + ); + (channel, bench_server) +} diff --git a/arrow-flight/benches/flight.rs b/arrow-flight/benches/flight.rs new file mode 100644 index 000000000000..4841e9dd9822 --- /dev/null +++ b/arrow-flight/benches/flight.rs @@ -0,0 +1,87 @@ +// 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 arrow_array::RecordBatch; +use arrow_flight::{FlightClient, FlightData, encode::FlightDataEncoderBuilder}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use futures::TryStreamExt; +use tonic::transport::Channel; + +mod common; +use common::{TYPES, build_batch, start_server}; + +const ROWS: [usize; 2] = [8 * 1024, 64 * 1024]; +const COLS: [usize; 2] = [1, 8]; + +fn bench_encode(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut g = c.benchmark_group("encode"); + + for &(name, build) in TYPES { + for &rows in &ROWS { + for &cols in &COLS { + let batch = build_batch(name, rows, cols, build); + let id = BenchmarkId::new(name, format!("{rows}x{cols}")); + g.throughput(Throughput::Bytes(batch.get_array_memory_size() as u64)); + g.bench_with_input(id, &batch, |b, batch| { + b.to_async(&rt).iter(|| async { + let _: Vec = FlightDataEncoderBuilder::new() + .build(futures::stream::iter([Ok(batch.clone())])) + .try_collect() + .await + .unwrap(); + }); + }); + } + } + } +} + +async fn roundtrip(channel: Channel, batch: RecordBatch) { + let mut client = FlightClient::new(channel); + let frames = FlightDataEncoderBuilder::new().build(futures::stream::iter([Ok(batch)])); + let _: Vec = client + .do_exchange(frames) + .await + .unwrap() + .try_collect() + .await + .unwrap(); +} + +fn bench_roundtrip(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let (channel, _) = rt.block_on(start_server()); + let mut g = c.benchmark_group("roundtrip"); + + for &(name, build) in TYPES { + for &rows in &ROWS { + for &cols in &COLS { + let batch = build_batch(name, rows, cols, build); + let id = BenchmarkId::new(name, format!("{rows}x{cols}")); + g.throughput(Throughput::Bytes(batch.get_array_memory_size() as u64)); + g.bench_with_input(id, &batch, |b, batch| { + b.to_async(&rt) + .iter(|| roundtrip(channel.clone(), batch.clone())); + }); + } + } + } +} + +criterion_group!(benches, bench_encode, bench_roundtrip); +criterion_main!(benches); From 4e7a2fa553e2e5e7385f6c4e77984a354d40c813 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 3 Jun 2026 14:15:20 +0100 Subject: [PATCH 43/51] Improve `take_bytes` perf in the null cases between 10-25% (#9625) # Which issue does this PR close? - Closes #NNN. # Rationale for this change Just improves performance, I was profiling some things downstream and got curious about how it works. # What changes are included in this PR? The main idea is to use a two-pass approach: 1. Compute byte offsets and collects (start, end) byte ranges 2. Copy byte data via raw pointer writes (`copy_byte_ranges`) This PR also reduces the branching from 4 (one for each nullability combination) to only two. # Are these changes tested? Existing tests # Are there any user-facing changes? None --------- Signed-off-by: Adam Gutglick Co-authored-by: Andrew Lamb Co-authored-by: Claude Opus 4.8 (1M context) --- arrow-select/src/take.rs | 241 ++++++++++++++++++++++++++------------- 1 file changed, 164 insertions(+), 77 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index cbb65ac915dd..c245247417c6 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -495,6 +495,7 @@ fn take_bytes( array: &GenericByteArray, indices: &PrimitiveArray, ) -> Result, ArrowError> { + let mut values: Vec = Vec::new(); let mut offsets = Vec::with_capacity(indices.len() + 1); offsets.push(T::Offset::default()); @@ -502,92 +503,116 @@ fn take_bytes( let mut capacity = 0; let nulls = take_nulls(array.nulls(), indices); - let (offsets, values) = if array.null_count() == 0 && indices.null_count() == 0 { - offsets.reserve(indices.len()); - for index in indices.values() { - let index = index.as_usize(); - capacity += input_offsets[index + 1].as_usize() - input_offsets[index].as_usize(); - offsets.push( - T::Offset::from_usize(capacity) - .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?, - ); - } - let mut values = Vec::with_capacity(capacity); - - for index in indices.values() { - values.extend_from_slice(array.value(index.as_usize()).as_ref()); - } - (offsets, values) - } else if indices.null_count() == 0 { - offsets.reserve(indices.len()); - for index in indices.values() { - let index = index.as_usize(); - if array.is_valid(index) { - capacity += input_offsets[index + 1].as_usize() - input_offsets[index].as_usize(); + // Branch on output nulls — `None` means every output slot is valid. + match nulls.as_ref().filter(|n| n.null_count() > 0) { + // Fast path: no nulls in output, every index is valid. + None => { + for index in indices.values() { + let index = index.as_usize(); + let start = input_offsets[index].as_usize(); + let end = input_offsets[index + 1].as_usize(); + capacity += end - start; + offsets.push( + T::Offset::from_usize(capacity) + .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?, + ); } - offsets.push( - T::Offset::from_usize(capacity) - .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?, - ); - } - let mut values = Vec::with_capacity(capacity); - for index in indices.values() { - let index = index.as_usize(); - if array.is_valid(index) { - values.extend_from_slice(array.value(index).as_ref()); - } - } - (offsets, values) - } else if array.null_count() == 0 { - offsets.reserve(indices.len()); - for (i, index) in indices.values().iter().enumerate() { - let index = index.as_usize(); - if indices.is_valid(i) { - capacity += input_offsets[index + 1].as_usize() - input_offsets[index].as_usize(); + values.reserve(capacity); + + let dst = values.spare_capacity_mut(); + debug_assert!(dst.len() >= capacity); + let mut offset = 0; + + for index in indices.values() { + // SAFETY: in-bounds proven by the first loop's bounds-checked offset access. + // dst asserted above to include the required capacity. + unsafe { + let data: &[u8] = array.value_unchecked(index.as_usize()).as_ref(); + std::ptr::copy_nonoverlapping( + data.as_ptr(), + dst.get_unchecked_mut(offset..).as_mut_ptr().cast::(), + data.len(), + ); + offset += data.len(); + } } - offsets.push( - T::Offset::from_usize(capacity) - .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?, - ); - } - let mut values = Vec::with_capacity(capacity); - for (i, index) in indices.values().iter().enumerate() { - if indices.is_valid(i) { - values.extend_from_slice(array.value(index.as_usize()).as_ref()); + // SAFETY: wrote exactly `capacity` bytes above; reserved on line above. + unsafe { + values.set_len(capacity); } } - (offsets, values) - } else { - let nulls = nulls.as_ref().unwrap(); - offsets.reserve(indices.len()); - for (i, index) in indices.values().iter().enumerate() { - let index = index.as_usize(); - if nulls.is_valid(i) { - capacity += input_offsets[index + 1].as_usize() - input_offsets[index].as_usize(); + // Nullable path: only process valid (non-null) output positions. + Some(output_nulls) => { + let mut source_ranges = Vec::with_capacity(indices.len() - output_nulls.null_count()); + let mut last_filled = 0; + + // Pre-fill offsets; we overwrite valid positions below. + offsets.resize(indices.len() + 1, T::Offset::default()); + + // Pass 1: find all valid ranges that need to be copied. + for i in output_nulls.valid_indices() { + let current_offset = T::Offset::from_usize(capacity) + .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?; + // Fill offsets for skipped null slots so they get zero-length ranges. + if last_filled < i { + offsets[last_filled + 1..=i].fill(current_offset); + } + + // SAFETY: `i` comes from a validity bitmap over `indices`, so it is in-bounds. + let index = unsafe { indices.value_unchecked(i) }.as_usize(); + let start = input_offsets[index].as_usize(); + let end = input_offsets[index + 1].as_usize(); + capacity += end - start; + offsets[i + 1] = T::Offset::from_usize(capacity) + .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?; + + source_ranges.push((start, end)); + last_filled = i + 1; } - offsets.push( - T::Offset::from_usize(capacity) - .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?, + + // Fill trailing null offsets after the last valid position. + let final_offset = T::Offset::from_usize(capacity) + .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?; + offsets[last_filled + 1..].fill(final_offset); + // Pass 2: copy byte data for all collected ranges. + values.reserve(capacity); + debug_assert_eq!( + source_ranges.iter().map(|(s, e)| e - s).sum::(), + capacity, + "capacity must equal total bytes across all ranges" ); - } - let mut values = Vec::with_capacity(capacity); - - for (i, index) in indices.values().iter().enumerate() { - // check index is valid before using index. The value in - // NULL index slots may not be within bounds of array - let index = index.as_usize(); - if nulls.is_valid(i) { - values.extend_from_slice(array.value(index).as_ref()); + + let src = array.value_data(); + let src = src.as_ptr(); + let dst = values.spare_capacity_mut(); + debug_assert!(dst.len() >= capacity); + + let mut offset = 0; + + for (start, end) in source_ranges.into_iter() { + let value_len = end - start; + // SAFETY: caller guarantees each (start, end) is in-bounds of `src`. + // `dst` asserted above to include the required capacity. + // The regions don't overlap (src is input, dst is a fresh allocation). + unsafe { + std::ptr::copy_nonoverlapping( + src.add(start), + dst.get_unchecked_mut(offset..).as_mut_ptr().cast::(), + value_len, + ); + offset += value_len; + } } + // SAFETY: caller guarantees `capacity` == total bytes across all ranges, + // so the loop above wrote exactly `capacity` bytes. + unsafe { values.set_len(capacity) }; } - (offsets, values) }; - T::Offset::from_usize(values.len()) - .ok_or_else(|| ArrowError::OffsetOverflowError(values.len()))?; - + // SAFETY: offsets are monotonically increasing and in-bounds of `values`, + // and `nulls` (if present) has length == `indices.len()`. let array = unsafe { let offsets = OffsetBuffer::new_unchecked(offsets.into()); GenericByteArray::::new_unchecked(offsets, values.into(), nulls) @@ -1728,6 +1753,41 @@ mod tests { assert_eq!(result.as_ref(), &expected); } + /// Take from a *sliced* byte array, i.e. one whose value offsets do not + /// start at zero. This exercises copying byte data out of an array with a + /// non-zero base offset for both the no-null fast path and the nullable + /// path (null indices and selected null values). + #[test] + fn test_take_bytes_sliced_values() { + let values = StringArray::from(vec![ + Some("aaa"), + Some("bbb"), + None, + Some("ccccc"), + Some("dd"), + None, + Some("eeee"), + ]); + // Slice so the underlying value offsets no longer start at 0: + // sliced == [None, "ccccc", "dd", None, "eeee"] + let sliced = values.slice(2, 5); + + // Fast path: every output slot is valid (no null indices, no null + // values selected). + let indices = Int32Array::from(vec![1, 2, 4, 1]); + let result = take(&sliced, &indices, None).unwrap(); + let expected = + StringArray::from(vec![Some("ccccc"), Some("dd"), Some("eeee"), Some("ccccc")]); + assert_eq!(result.as_string::(), &expected); + + // Nullable path: a null index (position 1) and selected null values + // (sliced indices 0 and 3 are null). + let indices = Int32Array::from(vec![Some(1), None, Some(0), Some(4), Some(3)]); + let result = take(&sliced, &indices, None).unwrap(); + let expected = StringArray::from(vec![Some("ccccc"), None, None, Some("eeee"), None]); + assert_eq!(result.as_string::(), &expected); + } + fn _test_byte_view() where T: ByteViewType, @@ -2808,11 +2868,38 @@ mod tests { assert_eq!(array.len(), 3); } + /// Fixture for the offset-overflow tests: a single large value plus the + /// number of times it must be selected so the cumulative offset exceeds + /// `i32::MAX`. Using a large value keeps the index count (and the test + /// runtime) small. + fn offset_overflow_fixture() -> (StringArray, usize) { + let value_len = 1_000_000; + let values = StringArray::from(vec![Some("a".repeat(value_len))]); + let n = i32::MAX as usize / value_len + 1; + (values, n) + } + #[test] fn test_take_bytes_offset_overflow() { - let indices = Int32Array::from(vec![0; (i32::MAX >> 4) as usize]); - let text = ('a'..='z').collect::(); - let values = StringArray::from(vec![Some(text.clone())]); + let (values, n) = offset_overflow_fixture(); + let indices = Int32Array::from(vec![0; n]); + assert!(matches!( + take(&values, &indices, None), + Err(ArrowError::OffsetOverflowError(_)) + )); + } + + /// The offset-overflow error must also be produced on the nullable code + /// path (when the output contains nulls), not only on the no-null fast path. + #[test] + fn test_take_bytes_offset_overflow_nullable() { + let (values, n) = offset_overflow_fixture(); + // A null index forces the output to contain nulls, exercising the + // nullable code path. + let validity = + NullBuffer::from_iter(std::iter::once(false).chain(std::iter::repeat_n(true, n))); + let indices = Int32Array::new(vec![0i32; n + 1].into(), Some(validity)); + assert!(matches!( take(&values, &indices, None), Err(ArrowError::OffsetOverflowError(_)) From 58bdc7d557df323b5c15b5779f8ea23b05bf0ccf Mon Sep 17 00:00:00 2001 From: theirix Date: Wed, 3 Jun 2026 14:18:47 +0100 Subject: [PATCH 44/51] arrow-buffer: i256: Implement num_traits wrapping shift (#9418) # Which issue does this PR close? - Closes #9417 # Rationale for this change A follow-up to #8976 Implement some missing traits - [WrappingShl](https://docs.rs/num-traits/latest/num_traits/ops/wrapping/trait.WrappingShl.html) and [WrappingShr](https://docs.rs/num-traits/latest/num_traits/ops/wrapping/trait.WrappingShl.html) # What changes are included in this PR? - num_traits' WrappingShl implementation for `usize` - `Shl` and `Shr` trait implementation for all scalar numeric types, not only for `u8` # Are these changes tested? - Unit tests # Are there any user-facing changes? --- arrow-buffer/src/bigint/mod.rs | 71 +++++++++++++++++++++++++++++++++- 1 file changed, 69 insertions(+), 2 deletions(-) diff --git a/arrow-buffer/src/bigint/mod.rs b/arrow-buffer/src/bigint/mod.rs index 15faed43a130..396597d7c61d 100644 --- a/arrow-buffer/src/bigint/mod.rs +++ b/arrow-buffer/src/bigint/mod.rs @@ -20,8 +20,8 @@ use crate::bigint::div::div_rem; use num_bigint::BigInt; use num_traits::{ Bounded, CheckedAdd, CheckedDiv, CheckedMul, CheckedNeg, CheckedRem, CheckedSub, FromPrimitive, - Num, One, Signed, ToPrimitive, WrappingAdd, WrappingMul, WrappingNeg, WrappingSub, Zero, - cast::AsPrimitive, + Num, One, Signed, ToPrimitive, WrappingAdd, WrappingMul, WrappingNeg, WrappingShl, WrappingShr, + WrappingSub, Zero, cast::AsPrimitive, }; use std::cmp::Ordering; use std::num::ParseIntError; @@ -807,6 +807,55 @@ impl Shr for i256 { } } +impl WrappingShl for i256 { + #[inline] + fn wrapping_shl(&self, rhs: u32) -> i256 { + // Limit shift to 256 (max valid shift for i256) + (*self).shl(rhs as u8) + } +} + +impl WrappingShr for i256 { + #[inline] + fn wrapping_shr(&self, rhs: u32) -> i256 { + // Limit shift to 256 (max valid shift for i256) + (*self).shr(rhs as u8) + } +} + +// Define Shl and Shr for specified integer types using +// an existing Shl and Shr implementation +macro_rules! define_standard_shift { + // Handle multiple types + ($trait_name:ident, $method:ident, [$($t:ty),+]) => { + $(define_standard_shift!($trait_name, $method, $t);)+ + }; + // Handle single type + ($trait_name:ident, $method:ident, $t:ty) => { + impl $trait_name<$t> for i256 { + type Output = i256; + + #[inline] + fn $method(self, rhs: $t) -> Self::Output { + let rhs = u8::try_from(rhs).expect("rhs overflow for shift"); + // Other possible overflows are handled by Shl implementation + self.$method(rhs) + } + } + }; +} + +define_standard_shift!( + Shl, + shl, + [u16, u32, u64, u128, usize, i16, i32, i64, i128, isize] +); +define_standard_shift!( + Shr, + shr, + [u16, u32, u64, u128, usize, i16, i32, i64, i128, isize] +); + macro_rules! define_as_primitive { ($native_ty:ty) => { impl AsPrimitive for $native_ty { @@ -1190,9 +1239,19 @@ mod tests { let (expected, _) = i256::from_bigint_with_overflow(bl.clone() << shift); assert_eq!(actual.to_string(), expected.to_string()); + let wrapping_actual = ::wrapping_shl(&il, shift as u32); + assert_eq!(wrapping_actual.to_string(), expected.to_string()); + let actual = il >> shift; let (expected, _) = i256::from_bigint_with_overflow(bl.clone() >> shift); assert_eq!(actual.to_string(), expected.to_string()); + + let wrapping_actual = ::wrapping_shr(&il, shift as u32); + assert_eq!(wrapping_actual.to_string(), expected.to_string()); + + // Check wrapping of the shift argument + let wrapping_actual = ::wrapping_shr(&il, 512 + shift as u32); + assert_eq!(wrapping_actual.to_string(), expected.to_string()); } } @@ -1535,6 +1594,14 @@ mod tests { assert_eq!(::max_value(), i256::MAX); } + #[should_panic] + #[test] + fn test_shl_panic_on_arg_overflow() { + let value = i256::from(123); + let rhs = std::hint::black_box(500); + let _ = value << rhs; + } + #[test] fn test_numtraits_from_str_radix() { assert_eq!( From 9949226f0ab644c701e6ed283db32989bbf6b006 Mon Sep 17 00:00:00 2001 From: mwish Date: Wed, 3 Jun 2026 21:33:22 +0800 Subject: [PATCH 45/51] perf(parquet): LevelInfoBuilder batch write when no repetition childs (#10037) # Which issue does this PR close? - Closes #10023 . # Rationale for this change Parquet writer writes lists element one by one, this is extremly slow. This patch batches writes. # What changes are included in this PR? Batches write when writing list with maximum rep level. # Are these changes tested? Covered by existing # Are there any user-facing changes? No --- parquet/src/arrow/arrow_writer/levels.rs | 191 +++++++++++++++++++++-- 1 file changed, 180 insertions(+), 11 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/levels.rs b/parquet/src/arrow/arrow_writer/levels.rs index 10f90f707c08..e577cf73a807 100644 --- a/parquet/src/arrow/arrow_writer/levels.rs +++ b/parquet/src/arrow/arrow_writer/levels.rs @@ -118,6 +118,7 @@ enum LevelInfoBuilder { LevelContext, // Context OffsetBuffer, // Offsets Option, // Nulls + bool, // is_last_level (child has no nested rep) ), /// A large list array LargeList( @@ -125,6 +126,7 @@ enum LevelInfoBuilder { LevelContext, // Context OffsetBuffer, // Offsets Option, // Nulls + bool, // is_last_level (child has no nested rep) ), /// A fixed size list array FixedSizeList( @@ -223,22 +225,31 @@ impl LevelInfoBuilder { DataType::List(_) => { let list = array.as_list(); let child = Self::try_new(child.as_ref(), ctx, list.values())?; + let is_last = child.child_has_no_nested_rep(); let offsets = list.offsets().clone(); - Self::List(Box::new(child), ctx, offsets, list.nulls().cloned()) + Self::List( + Box::new(child), + ctx, + offsets, + list.nulls().cloned(), + is_last, + ) } DataType::LargeList(_) => { let list = array.as_list(); let child = Self::try_new(child.as_ref(), ctx, list.values())?; + let is_last = child.child_has_no_nested_rep(); let offsets = list.offsets().clone(); let nulls = list.nulls().cloned(); - Self::LargeList(Box::new(child), ctx, offsets, nulls) + Self::LargeList(Box::new(child), ctx, offsets, nulls, is_last) } DataType::Map(_, _) => { let map = array.as_map(); let entries = Arc::new(map.entries().clone()) as ArrayRef; let child = Self::try_new(child.as_ref(), ctx, &entries)?; + let is_last = child.child_has_no_nested_rep(); let offsets = map.offsets().clone(); - Self::List(Box::new(child), ctx, offsets, map.nulls().cloned()) + Self::List(Box::new(child), ctx, offsets, map.nulls().cloned(), is_last) } DataType::FixedSizeList(_, size) => { let list = array.as_fixed_size_list(); @@ -274,8 +285,8 @@ impl LevelInfoBuilder { fn finish(self) -> Vec { match self { LevelInfoBuilder::Primitive(v) => vec![v], - LevelInfoBuilder::List(v, _, _, _) - | LevelInfoBuilder::LargeList(v, _, _, _) + LevelInfoBuilder::List(v, _, _, _, _) + | LevelInfoBuilder::LargeList(v, _, _, _, _) | LevelInfoBuilder::FixedSizeList(v, _, _, _) | LevelInfoBuilder::ListView(v, _, _, _, _) | LevelInfoBuilder::LargeListView(v, _, _, _, _) => v.finish(), @@ -287,11 +298,11 @@ impl LevelInfoBuilder { fn write(&mut self, range: Range) { match self { LevelInfoBuilder::Primitive(info) => Self::write_leaf(info, range), - LevelInfoBuilder::List(child, ctx, offsets, nulls) => { - Self::write_list(child, ctx, offsets, nulls.as_ref(), range) + LevelInfoBuilder::List(child, ctx, offsets, nulls, is_last) => { + Self::write_list(child, ctx, offsets, nulls.as_ref(), range, *is_last) } - LevelInfoBuilder::LargeList(child, ctx, offsets, nulls) => { - Self::write_list(child, ctx, offsets, nulls.as_ref(), range) + LevelInfoBuilder::LargeList(child, ctx, offsets, nulls, is_last) => { + Self::write_list(child, ctx, offsets, nulls.as_ref(), range, *is_last) } LevelInfoBuilder::FixedSizeList(child, ctx, size, nulls) => { Self::write_fixed_size_list(child, ctx, *size, nulls.as_ref(), range) @@ -308,6 +319,19 @@ impl LevelInfoBuilder { } } + /// Returns `true` if the child contains no nested repetition levels, meaning + /// each child element produces exactly one rep_level entry in the leaf. + /// This is true for `Primitive` children and `Struct` trees with no list descendants. + fn child_has_no_nested_rep(&self) -> bool { + match self { + LevelInfoBuilder::Primitive(_) => true, + LevelInfoBuilder::Struct(children, _, _) => { + children.iter().all(|c| c.child_has_no_nested_rep()) + } + _ => false, + } + } + /// Write `range` elements from ListArray `array` /// /// Note: MapArrays are `ListArray` under the hood and so are dispatched to this method @@ -317,6 +341,7 @@ impl LevelInfoBuilder { offsets: &[O], nulls: Option<&NullBuffer>, range: Range, + is_last_level: bool, ) { // Fast path: entire list array is null; emit bulk null rep/def levels if nulls.is_some_and(|nulls| nulls.null_count() == nulls.len()) { @@ -327,6 +352,18 @@ impl LevelInfoBuilder { return; } + // Fast path for "last-level list": when the child has no nested rep_levels, + // each child element produces exactly one rep_level entry. We can batch + // contiguous non-empty list slots into a single child.write() call, then + // fix up the rep_levels at list-slot boundaries using offsets directly. + // + // Kept as a separate function so the compiler can optimize write_list's + // hot loop independently (function body size affects codegen quality). + if is_last_level { + Self::write_list_last_level(child, ctx, offsets, nulls, range); + return; + } + let offsets = &offsets[range.start..range.end + 1]; let write_non_null_slice = @@ -427,6 +464,138 @@ impl LevelInfoBuilder { } } + /// Optimized write path for lists whose child has no nested repetition levels. + /// + /// When the child is a leaf (or a struct of leaves), each child element maps to + /// exactly one rep_level entry. This lets us batch contiguous non-empty list + /// slots into a single `child.write()` call, then stamp the list-start markers + /// at positions computed directly from offsets — avoiding per-slot `write` + + /// reverse-scan overhead. + fn write_list_last_level( + child: &mut LevelInfoBuilder, + ctx: &LevelContext, + offsets: &[O], + nulls: Option<&NullBuffer>, + range: Range, + ) { + let null_offset = range.start; + let offsets = &offsets[range.start..range.end + 1]; + let list_start_rep = ctx.rep_level - 1; + + let emit_nulls = |child: &mut LevelInfoBuilder, count: usize| { + child.visit_leaves(|leaf| { + leaf.append_rep_level_run(list_start_rep, count); + leaf.append_def_level_run(ctx.def_level - 2, count); + }); + }; + + let emit_empties = |child: &mut LevelInfoBuilder, count: usize| { + child.visit_leaves(|leaf| { + leaf.append_rep_level_run(list_start_rep, count); + leaf.append_def_level_run(ctx.def_level - 1, count); + }); + }; + + let emit_non_empty_run = |child: &mut LevelInfoBuilder, run_offsets: &[O]| { + debug_assert!(run_offsets.len() >= 2); + let values_start = run_offsets[0].as_usize(); + let values_end = run_offsets[run_offsets.len() - 1].as_usize(); + debug_assert!(values_end > values_start); + + // Write all leaf values in one batch. Since the child has no nested + // rep, this emits (values_end - values_start) rep_levels all equal + // to ctx.rep_level (= "continuation within list"). + child.write(values_start..values_end); + + // The first element of each list slot needs rep_level = + // list_start_rep to mark a new list boundary. Because there's a 1:1 + // mapping between child elements and rep_level entries, the position + // of each slot's first element is directly computable from offsets. + child.visit_leaves(|leaf| { + let rep_levels = leaf.rep_levels.materialize_mut().unwrap(); + let batch_len = values_end - values_start; + let batch_base = rep_levels.len() - batch_len; + + for slot_offset in run_offsets.iter().take(run_offsets.len() - 1) { + let list_start_pos = batch_base + (slot_offset.as_usize() - values_start); + rep_levels[list_start_pos] = list_start_rep; + } + }); + }; + + // Classify each slot, detect run boundaries, flush on transition. + #[derive(Clone, Copy, PartialEq)] + enum SlotKind { + Null, + Empty, + NonEmpty, + } + + let num_slots = offsets.len() - 1; + if num_slots == 0 { + return; + } + + macro_rules! classify { + ($i:expr, $nulls:expr) => { + if !$nulls.is_valid($i + null_offset) { + SlotKind::Null + } else if offsets[$i] == offsets[$i + 1] { + SlotKind::Empty + } else { + SlotKind::NonEmpty + } + }; + } + + macro_rules! flush_run { + ($kind:expr, $start:expr, $end:expr) => { + match $kind { + SlotKind::Null => emit_nulls(child, $end - $start), + SlotKind::Empty => emit_empties(child, $end - $start), + SlotKind::NonEmpty => emit_non_empty_run(child, &offsets[$start..$end + 1]), + } + }; + } + + match nulls { + Some(nulls) => { + let mut run_kind = classify!(0, nulls); + let mut run_start: usize = 0; + for i in 1..num_slots { + let kind = classify!(i, nulls); + if kind != run_kind { + flush_run!(run_kind, run_start, i); + run_kind = kind; + run_start = i; + } + } + flush_run!(run_kind, run_start, num_slots); + } + None => { + let mut run_kind = if offsets[0] == offsets[1] { + SlotKind::Empty + } else { + SlotKind::NonEmpty + }; + let mut run_start: usize = 0; + for i in 1..num_slots { + let kind = if offsets[i] == offsets[i + 1] { + SlotKind::Empty + } else { + SlotKind::NonEmpty + }; + if kind != run_kind { + flush_run!(run_kind, run_start, i); + run_kind = kind; + run_start = i; + } + } + flush_run!(run_kind, run_start, num_slots); + } + } + } + /// Write `range` elements from ListViewArray `array` fn write_list_view( child: &mut LevelInfoBuilder, @@ -734,8 +903,8 @@ impl LevelInfoBuilder { fn visit_leaves(&mut self, visit: impl Fn(&mut ArrayLevels) + Copy) { match self { LevelInfoBuilder::Primitive(info) => visit(info), - LevelInfoBuilder::List(c, _, _, _) - | LevelInfoBuilder::LargeList(c, _, _, _) + LevelInfoBuilder::List(c, _, _, _, _) + | LevelInfoBuilder::LargeList(c, _, _, _, _) | LevelInfoBuilder::FixedSizeList(c, _, _, _) | LevelInfoBuilder::ListView(c, _, _, _, _) | LevelInfoBuilder::LargeListView(c, _, _, _, _) => c.visit_leaves(visit), From 6c397977687380f5e3d13d3d061b0369fe31cb3a Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 3 Jun 2026 09:38:18 -0400 Subject: [PATCH 46/51] Improve email created by create_tarball.sh script (#9944) # Which issue does this PR close? # Rationale for this change There are several issues I found with the email template created by the current template: 1. It uses links to tags (rather than git shas) which can potentially be changed 2. It does not include the actual SHA values (only a link to a place to download the sha values) which means in theory it is not clear what exact artifact is being voted on 3. It does not include a link to the issue used to do release coordiation # What changes are included in this PR? Fix the above issues Example new output: ``` --------------------------------------------------------- To: dev@arrow.apache.org Subject: [VOTE][RUST] Release Apache Arrow Rust 57.3.1 RC1 Hi, I would like to propose a release of Apache Arrow Rust Implementation, version 57.3.1. This release candidate is based on commit: da8975cfacdf8623892a7937dc5c5e6515a05483 [1]. The SHA256 of the release candidate is: 067a4c47c515d57b283f431d426c46c0f48601a2017202a490d2a234e0cd2fb4 The proposed release tarball and signatures are hosted at [2]. The changelog is located at [3]. The release tracking issue is: [4] Please download, verify checksums and signatures, run the unit tests, and vote on the release. There is a script [4] that automates some of the verification. The vote will be open for at least 72 hours. [ ] +1 Release this as Apache Arrow Rust 57.3.1 [ ] +0 [ ] -1 Do not release this as Apache Arrow Rust 57.3.1 because... [1]: https://github.com/apache/arrow-rs/tree/da8975cfacdf8623892a7937dc5c5e6515a05483 [2]: https://dist.apache.org/repos/dist/dev/arrow/apache-arrow-rs-57.3.1-rc1 [3]: https://github.com/apache/arrow-rs/blob/da8975cfacdf8623892a7937dc5c5e6515a05483/CHANGELOG.md [4]: https://github.com/apache/arrow-rs/blob/master/dev/release/verify-release-candidate.sh [5]: RELEASE_ISSUE ``` # Are these changes tested? I tested this script while creating release candidates for 57.3.1 and 56.2.1 and it worked well - https://github.com/apache/arrow-rs/issues/9858 - https://github.com/apache/arrow-rs/issues/9857 # Are there any user-facing changes? --- dev/release/create-tarball.sh | 58 +++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 27 deletions(-) diff --git a/dev/release/create-tarball.sh b/dev/release/create-tarball.sh index b75313b6f0d6..332de47715bf 100755 --- a/dev/release/create-tarball.sh +++ b/dev/release/create-tarball.sh @@ -65,8 +65,9 @@ else tar=tar fi -if ! git -C "${SOURCE_TOP_DIR}" rev-list --max-count=1 ${tag}; then +if ! commit_sha=$(git -C "${SOURCE_TOP_DIR}" rev-list --max-count=1 "${tag}"); then echo "Cannot continue: unknown git tag: $tag" + exit 1 fi @@ -79,6 +80,26 @@ url="https://dist.apache.org/repos/dist/dev/arrow/${release}-rc${rc}" echo "Attempting to create ${tarball} from tag ${tag}" +# create containing the files in git at $tag +# the files in the tarball are prefixed with {release} +# (e.g. apache-arrow-rs-4.0.1) +mkdir -p ${distdir} +(cd "${SOURCE_TOP_DIR}" && \ + git archive ${tag} --prefix ${release}/ \ + | gzip > ${tarball}) + +echo "Running rat license checker on ${tarball}" +${SOURCE_DIR}/run-rat.sh ${tarball} + +echo "Signing tarball and creating checksums" +gpg --armor --output ${tarball}.asc --detach-sig ${tarball} +# create signing with relative path of tarball +# so that they can be verified with a command such as +# shasum --check apache-arrow-rs-4.1.0-rc2.tar.gz.sha512 +(cd ${distdir} && shasum -a 256 ${tarname}) > ${tarball}.sha256 +(cd ${distdir} && shasum -a 512 ${tarname}) > ${tarball}.sha512 +sha256=$(cut -d ' ' -f 1 "${tarball}.sha256") + echo "Draft email for dev@arrow.apache.org mailing list" echo "" echo "---------------------------------------------------------" @@ -88,14 +109,17 @@ Subject: [VOTE][RUST] Release Apache Arrow Rust ${version} RC${rc} Hi, -I would like to propose a release of Apache Arrow Rust Implementation, version ${tag}. +I would like to propose a release of Apache Arrow Rust Implementation, version ${version}. -This release candidate is based on commit: ${tag} [1] +This release candidate is based on commit: ${commit_sha} [1]. +The SHA256 of the release candidate is: ${sha256} The proposed release tarball and signatures are hosted at [2]. The changelog is located at [3]. +The release tracking issue is: [4] + Please download, verify checksums and signatures, run the unit tests, and vote on the release. There is a script [4] that automates some of the verification. @@ -106,35 +130,15 @@ The vote will be open for at least 72 hours. [ ] +0 [ ] -1 Do not release this as Apache Arrow Rust ${version} because... -[1]: https://github.com/apache/arrow-rs/tree/${tag} +[1]: https://github.com/apache/arrow-rs/tree/${commit_sha} [2]: ${url} -[3]: https://github.com/apache/arrow-rs/blob/${tag}/CHANGELOG.md +[3]: https://github.com/apache/arrow-rs/blob/${commit_sha}/CHANGELOG.md [4]: https://github.com/apache/arrow-rs/blob/master/dev/release/verify-release-candidate.sh +[5]: RELEASE_ISSUE MAIL echo "---------------------------------------------------------" - - -# create containing the files in git at $tag -# the files in the tarball are prefixed with {release} -# (e.g. apache-arrow-rs-4.0.1) -mkdir -p ${distdir} -(cd "${SOURCE_TOP_DIR}" && \ - git archive ${tag} --prefix ${release}/ \ - | gzip > ${tarball}) - -echo "Running rat license checker on ${tarball}" -${SOURCE_DIR}/run-rat.sh ${tarball} - -echo "Signing tarball and creating checksums" -gpg --armor --output ${tarball}.asc --detach-sig ${tarball} -# create signing with relative path of tarball -# so that they can be verified with a command such as -# shasum --check apache-arrow-rs-4.1.0-rc2.tar.gz.sha512 -(cd ${distdir} && shasum -a 256 ${tarname}) > ${tarball}.sha256 -(cd ${distdir} && shasum -a 512 ${tarname}) > ${tarball}.sha512 - echo "Uploading to apache dist/dev to ${url}" svn co --depth=empty https://dist.apache.org/repos/dist/dev/arrow ${SOURCE_TOP_DIR}/dev/dist svn add ${distdir} -svn ci -m "Apache Arrow Rust ${version} ${rc}" ${distdir} \ No newline at end of file +svn ci -m "Apache Arrow Rust ${version} ${rc}" ${distdir} From 97f4b1460cd32b5b8c3b91ed554238e12e3e68cf Mon Sep 17 00:00:00 2001 From: theirix Date: Wed, 3 Jun 2026 14:45:53 +0100 Subject: [PATCH 47/51] arrow-buffer: i256: implement ilog (#9453) # Which issue does this PR close? - Closes #9452. # Rationale for this change Implementation of integer logarithm. There is no matching `num_traits` trait, but this implementation provides a good motivation for such a trait. # What changes are included in this PR? - No external dependencies - Checked methods (log, log2, log10) - Unchecked methods (panic by design) # Are these changes tested? - Unit tests # Are there any user-facing changes? No --------- Co-authored-by: Andrew Lamb Co-authored-by: Claude Opus 4.8 (1M context) --- arrow-buffer/src/bigint/mod.rs | 223 +++++++++++++++++++++++++++++++++ 1 file changed, 223 insertions(+) diff --git a/arrow-buffer/src/bigint/mod.rs b/arrow-buffer/src/bigint/mod.rs index 396597d7c61d..e1db0c6141e0 100644 --- a/arrow-buffer/src/bigint/mod.rs +++ b/arrow-buffer/src/bigint/mod.rs @@ -614,6 +614,101 @@ impl i256 { let n = (n.high >> 64) as i64; // throw away the lower 192 bits (n as f64) * f64::powi(2.0, 192 - (k as i32)) // convert to f64 and scale it, as we left-shift k bit previous, so we need to scale it by 2^(192-k) } + + /// Computes the `base` logarithm of the number `self` + /// Returns `None` if `self` is less than or equal to zero, or if `base` is less than 2. + #[inline] + pub fn checked_ilog(self, base: i256) -> Option { + if base == Self::from(10) { + // Faster implementation for base 10 + return self.checked_ilog10(); + } + + if self <= Self::ZERO { + return None; + } + if base <= Self::ONE { + return None; + } + if self < base { + return Some(0); + } + + let mut val = 1; + let mut base_exp = base; + + let boundary = self.checked_div(base)?; + while base_exp <= boundary { + val += 1; + base_exp = base_exp.checked_mul(base)?; + } + Some(val) + } + + /// Computes the `base` logarithm of the number `self` + /// Panic if `self` is less than or equal to zero, or if `base` is less than 2. + #[inline] + pub fn ilog(self, base: i256) -> u32 { + self.checked_ilog(base) + .unwrap_or_else(|| panic!("ilog overflow with {self} and base {base}")) + } + + /// Computes the decimal logarithm of the number `self` + /// Returns `None` if `self` is less than or equal to zero. + #[inline] + pub fn checked_ilog10(self) -> Option { + if self <= Self::ZERO { + return None; + } + if self < Self::from(10) { + return Some(0); + } + + // Layered approach to calculate logarithm using i128 log operations only + // Consult int_log10.rs stdlib implementiation for u128 + let pow_64: i256 = i256::from(10).checked_pow(64).unwrap(); + let pow_32: i256 = i256::from(10).checked_pow(32).unwrap(); + if self >= pow_64 { + let value = self.checked_div(pow_64)?; + // self is between 10^64 and 10^77 (~i256::MAX). + // `value` is 14 digits max (10^77 / 10^64 = 10^13), + // so it fits to `low` u128 + debug_assert!(value.high == 0); + Some(64 + value.low.checked_ilog10()?) + } else if self >= pow_32 { + let value = self.checked_div(pow_32)?; + // self is between 10^32 and 10^64. + // `value` is 33 digits max (10^64/10^32=10^32) + // so it fits to `low` 128-bit value + debug_assert!(value.high == 0); + Some(32 + value.low.checked_ilog10()?) + } else { + // self fits within u128 (high == 0 and self > 0). + self.low.checked_ilog10() + } + } + + /// Computes the decimal logarithm of the number `self` + /// Panics if `self` is less than or equal to zero. + #[inline] + pub fn ilog10(self) -> u32 { + self.checked_ilog10() + .unwrap_or_else(|| panic!("ilog10 overflow with {self}")) + } + + /// Computes the binary logarithm of the number `self` + /// Returns `None` if `self` is less than or equal to zero. + #[inline] + pub fn checked_ilog2(self) -> Option { + self.checked_ilog(i256::from(2)) + } + + /// Computes the base 2 logarithm of the number, rounded down. + #[inline] + pub fn ilog2(self) -> u32 { + self.checked_ilog2() + .unwrap_or_else(|| panic!("ilog2 overflow with {self}")) + } } /// Temporary workaround due to lack of stable const array slicing @@ -1647,4 +1742,132 @@ mod tests { assert_eq!(i256::MAX.trailing_zeros(), 0); assert_eq!(i256::from(-1).trailing_zeros(), 0); } + + #[test] + fn test_ilog() { + let value = i256::from(128); + + // log2 + assert_eq!(value.ilog(i256::from(2)), 7); + assert_eq!(value.ilog2(), 7); + + // log10 + assert_eq!(value.ilog(i256::from(10)), 2); + assert_eq!(value.ilog10(), 2); + + // negative base + assert_eq!(value.checked_ilog(i256::from(-2)), None); + assert_eq!(value.checked_ilog(i256::from(-10)), None); + assert_eq!(value.checked_ilog(i256::from(0)), None); + assert_eq!(value.checked_ilog(i256::from(1)), None); + + // negative self + let neg_value = i256::from(-128); + assert_eq!(neg_value.checked_ilog(i256::from(2)), None); + assert_eq!(neg_value.checked_ilog(i256::from(10)), None); + assert_eq!(neg_value.checked_ilog10(), None); + assert_eq!(neg_value.checked_ilog2(), None); + + // zero self + assert_eq!(i256::ZERO.checked_ilog(i256::from(2)), None); + assert_eq!(i256::ZERO.checked_ilog(i256::from(10)), None); + assert_eq!(i256::ZERO.checked_ilog10(), None); + assert_eq!(i256::ZERO.checked_ilog2(), None); + + // self == base, matches std: `n.ilog(n) == 1` + assert_eq!(i256::from(2).checked_ilog(i256::from(2)), Some(1)); + assert_eq!(i256::from(3).checked_ilog(i256::from(3)), Some(1)); + assert_eq!(i256::from(5).checked_ilog(i256::from(5)), Some(1)); + assert_eq!(i256::from(1000).checked_ilog(i256::from(1000)), Some(1)); + assert_eq!(i256::from(2).checked_ilog2(), Some(1)); + assert_eq!(i256::from(2).ilog2(), 1); + // base 10 goes through the checked_ilog10 fast path + assert_eq!(i256::from(10).checked_ilog(i256::from(10)), Some(1)); + + // self < base is 0 + assert_eq!(i256::from(3).checked_ilog(i256::from(5)), Some(0)); + + // cross-check small results (0 and 1) against u128::ilog + for base in [2i64, 3, 5, 7, 1000] { + for v in 1i64..64 { + let want = (v as u128).ilog(base as u128); + assert_eq!( + i256::from(v).checked_ilog(i256::from(base)), + Some(want), + "checked_ilog({v}, {base})" + ); + } + } + + let value = i256::from_parts(100000000, 1234); + assert_eq!(value.checked_ilog(i256::from(10)), Some(41)); + assert_eq!(value.checked_ilog10(), Some(41)); + + // Large i256 values + let large = i256::from_parts(100000000, i128::MAX); + // log2 of 2 powered to approximately 255 should be 254 + assert_eq!(large.checked_ilog(i256::from(2)), Some(254)); + + // log10(large)=76 + assert_eq!(large.checked_ilog(i256::from(10)), Some(76)); + assert_eq!(large.checked_ilog10(), Some(76)); + + // log5(large) + assert_eq!(large.checked_ilog(i256::from(5)), Some(109)); + + // Maximum representable value is 2^254 + assert!(i256::from(2).checked_pow(255).is_none()); + let value = i256::from(2).checked_pow(254).expect("construct"); + assert_eq!(value.checked_ilog(i256::from(2)), Some(254)); + + // Logarithm of a maximum representable value is 254 + assert_eq!(i256::MAX.checked_ilog(i256::from(2)), Some(254)); + } + + #[test] + fn test_ilog10() { + // Edge cases + assert_eq!(i256::ZERO.checked_ilog10(), None); + assert_eq!(i256::MINUS_ONE.checked_ilog10(), None); + assert_eq!(i256::MAX.checked_ilog10(), Some(76)); + assert_eq!(i256::from(10).checked_ilog10(), Some(1)); + + // small values + assert_eq!(i256::from(1).checked_ilog10(), Some(0)); + assert_eq!(i256::from(9).checked_ilog10(), Some(0)); + + // case with high == 0 + assert_eq!(i256::from(100).checked_ilog10(), Some(2)); + // case with high == 0 and full low + assert_eq!(i256::from_parts(u128::MAX, 0).checked_ilog10(), Some(38)); + + // case with high > 0 + assert_eq!(i256::from_parts(0, 1).checked_ilog10(), Some(38)); + + // case with non-null high and low, slow branch + let pow50 = i256::from(10).checked_pow(50).unwrap(); + assert_eq!(pow50.checked_ilog10(), Some(50)); + + // case with non-null high and low, fast branch + let pow64 = i256::from(10).checked_pow(64).unwrap(); + assert_eq!(pow64.checked_ilog10(), Some(64)); + } + + #[test] + #[should_panic(expected = "ilog10 overflow")] + fn test_ilog10_zero_panics() { + let _ = i256::ZERO.ilog10(); + } + + #[test] + #[should_panic(expected = "ilog overflow")] + fn test_ilog_zero_panics() { + let _ = i256::ZERO.ilog(i256::from(5)); + } + + #[test] + #[should_panic(expected = "ilog2 overflow")] + fn test_ilog2_zero_panics() { + let _ = i256::ZERO.ilog2(); + } } From 2a1d40d35c3dd5cc274dd97baa413a5384006fb2 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 3 Jun 2026 21:17:25 +0100 Subject: [PATCH 48/51] Reduce Miri runtime even more (#9650) # Rationale for this change Follow up to #9629, as noted there miri runtime is dominated by the following 3 tests: 1. `test_from_bitwise_binary_op` 2. `sort::tests::fuzz_partition_validity` 3. `sort::tests::test_fuzz_random_strings` # What changes are included in this PR? Under Miri, this PR reduces the amount of variations tested in 1, and ignores the latter two. # Are these changes tested? They are tests! # Are there any user-facing changes? No --------- Co-authored-by: Jeffrey Vo --- arrow-buffer/src/buffer/boolean.rs | 7 ++++++- arrow-ord/src/sort.rs | 2 ++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index 52df67080a3c..52f1e3fbb510 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -1073,7 +1073,12 @@ mod tests { let input_buffer_left = BooleanBuffer::from(&input_bools_left[..]); let input_buffer_right = BooleanBuffer::from(&input_bools_right[..]); - for left_offset in 0..200 { + #[cfg(miri)] // Takes too long otherwise + let left_offsets = [0, 1, 7, 8, 63, 64, 65]; + #[cfg(not(miri))] + let left_offsets = 0..200; + + for left_offset in left_offsets { for right_offset in [0, 4, 5, 17, 33, 24, 45, 64, 65, 100, 200] { for len_offset in [0, 1, 44, 100, 256, 300, 512] { let len = 1024 - len_offset - left_offset.max(right_offset); // ensure we don't go out of bounds diff --git a/arrow-ord/src/sort.rs b/arrow-ord/src/sort.rs index 773d9cf12e13..114ca05bb7eb 100644 --- a/arrow-ord/src/sort.rs +++ b/arrow-ord/src/sort.rs @@ -4892,6 +4892,7 @@ mod tests { } #[test] + #[cfg_attr(miri, ignore)] // Takes too long fn fuzz_partition_validity() { let mut rng = StdRng::seed_from_u64(0xF00D_CAFE); for _ in 0..1_000 { @@ -5100,6 +5101,7 @@ mod tests { // Fuzz testing: generate random UTF-8 strings and verify sort correctness #[test] + #[cfg_attr(miri, ignore)] // Takes too long fn test_fuzz_random_strings() { let mut rng = StdRng::seed_from_u64(42); // Fixed seed for reproducibility From fa4b19a928656a15351ac8b629945f28e2402c73 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 3 Jun 2026 17:09:58 -0500 Subject: [PATCH 49/51] fix(parquet): bound page byte size for large variable-width values - closes https://github.com/apache/arrow-rs/issues/10061 The column writer only checks the data/dictionary page byte limit *after* each `write_batch_size` mini-batch, so a batch of large variable-width values piles into a single oversized page before the check fires (we've observed multi-GiB data pages and large dictionary-page overshoot at default settings). Make the mini-batch size byte-budget aware in the generic column writer: - `ColumnValueEncoder::count_values_within_byte_budget{,_gather}` (default `None` = "no estimate, stay batched"), with a concrete impl on `ColumnValueEncoderImpl` driven by `plain_encoded_byte_size`. Fixed-width physical types answer in one division; only variable-width BYTE_ARRAY/FLBA walk values, stopping at the first that overruns. - `LevelDataRef::value_count` converts a chunk's level span into a leaf value count (O(1) for flat columns, def-level scan when nullable/nested). - `ByteBudgetChunker` picks the largest sub-batch that fits one page budget. The common case (small or fixed-width values) returns the whole chunk with no value inspection, so the hot path is unchanged. During dictionary encoding it sizes against the dictionary page's remaining budget instead, since the data page then holds only small RLE indices. - `write_batch_internal` consults the chunker per chunk and, only when a chunk would overflow, routes through `write_granular_chunk`, which sub-batches so the post-write page check fires in time. Repeated/nested columns step on record (rep == 0) boundaries so a record never spans pages. Includes the `ColumnWriterImpl`-level regression tests (data page, list, nullable, FLBA, dictionary spill, dictionary page bound). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/column/writer/byte_budget_chunker.rs | 204 +++++++ parquet/src/column/writer/encoder.rs | 141 +++++ parquet/src/column/writer/mod.rs | 531 +++++++++++++++++- 3 files changed, 870 insertions(+), 6 deletions(-) create mode 100644 parquet/src/column/writer/byte_budget_chunker.rs diff --git a/parquet/src/column/writer/byte_budget_chunker.rs b/parquet/src/column/writer/byte_budget_chunker.rs new file mode 100644 index 000000000000..56c7b4c6f86d --- /dev/null +++ b/parquet/src/column/writer/byte_budget_chunker.rs @@ -0,0 +1,204 @@ +// 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. + +//! See [`ByteBudgetChunker`] for byte-budget-aware mini-batch sizing. + +use crate::basic::Type; +use crate::column::writer::LevelDataRef; +use crate::column::writer::encoder::ColumnValueEncoder; +use crate::file::properties::WriterProperties; +use crate::schema::types::ColumnDescriptor; + +/// Picks byte-budget-aware mini-batch sizes for one column. +/// +/// The parquet column writer checks the data page byte limit only *after* +/// each mini-batch finishes writing. Mini-batches are sized in rows +/// (`write_batch_size`, default 1024), so for BYTE_ARRAY columns whose +/// values are large (e.g. multi-MiB blobs) a single mini-batch can buffer +/// GiB into one page before the limit is consulted. +/// +/// This isolates the per-chunk decision that prevents that: given a chunk's +/// level data and the input values, pick the largest `sub_batch_size` such +/// that one mini-batch fits in one page byte budget. For the overwhelmingly +/// common case (small or fixed-width values) the answer is just `chunk_size` +/// and the decision is O(1) on the column type — only when the input might +/// overflow does the chunker consult the encoder's byte estimate. +pub(crate) struct ByteBudgetChunker { + /// Configured data page byte limit for the column. + page_byte_limit: usize, + /// Max definition level of the column; a level equal to this marks a + /// present (non-null) leaf value. Used to count values per chunk. + max_def_level: i16, + /// `true` when no chunk of `base_batch_size` values can ever overflow + /// `page_byte_limit` regardless of input. Set once at column open from + /// the physical type's known per-value byte size; lets the per-chunk + /// decision short-circuit with no work for every numeric, bool, or + /// narrow `FIXED_LEN_BYTE_ARRAY` column. + static_always_fits: bool, + /// Configured dictionary page byte limit for the column. + dict_page_byte_limit: usize, + /// As [`Self::static_always_fits`] but for the dictionary page: `true` + /// when one `base_batch_size` mini-batch of this fixed-width type cannot + /// overshoot `dict_page_byte_limit` by more than one mini-batch's worth. + static_dict_always_fits: bool, +} + +impl ByteBudgetChunker { + #[inline] + pub(crate) fn new( + descr: &ColumnDescriptor, + props: &WriterProperties, + base_batch_size: usize, + ) -> Self { + let page_byte_limit = props.column_data_page_size_limit(descr.path()); + let dict_page_byte_limit = props.column_dictionary_page_size_limit(descr.path()); + let static_bytes_per_value = match descr.physical_type() { + Type::BOOLEAN => Some(1), + Type::INT32 | Type::FLOAT => Some(std::mem::size_of::()), + Type::INT64 | Type::DOUBLE => Some(std::mem::size_of::()), + Type::INT96 => Some(12), + Type::FIXED_LEN_BYTE_ARRAY => Some(descr.type_length().max(0) as usize), + Type::BYTE_ARRAY => None, + }; + let static_fits = |limit: usize| { + static_bytes_per_value + .map(|b| b.saturating_mul(base_batch_size) <= limit) + .unwrap_or(false) + }; + Self { + page_byte_limit, + max_def_level: descr.max_def_level(), + static_always_fits: static_fits(page_byte_limit), + dict_page_byte_limit, + static_dict_always_fits: static_fits(dict_page_byte_limit), + } + } + + /// Decide how many levels at the start of a chunk belong in one + /// mini-batch, so the mini-batch cannot overflow whichever page is + /// currently accumulating value bytes: the data page when plain-encoding, + /// or the *dictionary* page while dictionary-encoding. A returned value + /// smaller than `chunk_size` triggers granular sub-batching in + /// `write_batch_internal`. + /// + /// While dictionary-encoding, the data page holds only small RLE indices, + /// but the dictionary page accumulates the distinct values themselves — + /// so it is the dictionary page's remaining budget that must bound the + /// mini-batch. The per-mini-batch dictionary spill check would otherwise + /// let one mini-batch of large values balloon the dictionary page. + /// + /// Returns `chunk_size` immediately (no value inspection) when the chunk + /// is empty, or when the column is a fixed-width type whose mini-batches + /// statically cannot overshoot the relevant page. + /// + /// `#[inline]`: this is a tiny per-chunk dispatcher; the actual byte + /// inspection lives in the out-of-line `byte_budget_sub_batch_size`. + #[inline] + pub(crate) fn pick_sub_batch_size( + &self, + encoder: &E, + values: &E::Values, + value_indices: Option<&[usize]>, + chunk_def: LevelDataRef<'_>, + values_offset: usize, + chunk_size: usize, + ) -> usize { + if chunk_size == 0 { + return chunk_size; + } + let budget = if encoder.has_dictionary() { + if self.static_dict_always_fits { + return chunk_size; + } + // Bound the mini-batch by the dictionary page's *remaining* + // budget (it accumulates across mini-batches until it spills). + match encoder.estimated_dict_page_size() { + Some(used) => self.dict_page_byte_limit.saturating_sub(used), + None => return chunk_size, + } + } else { + if self.static_always_fits { + return chunk_size; + } + self.page_byte_limit + }; + self.byte_budget_sub_batch_size::( + values, + value_indices, + chunk_def, + values_offset, + chunk_size, + budget, + ) + } + + /// Inspect value sizes to decide how many of the chunk's values fit in + /// `budget` bytes (the data page or dictionary page remaining budget). + /// + /// `#[inline(never)]` keeps this slow path out of the hot + /// `write_batch_internal` loop; numeric and bool columns never reach it. + #[inline(never)] + fn byte_budget_sub_batch_size( + &self, + values: &E::Values, + value_indices: Option<&[usize]>, + chunk_def: LevelDataRef<'_>, + values_offset: usize, + chunk_size: usize, + budget: usize, + ) -> usize { + // How many of this chunk's levels carry an actual value. For a + // non-nullable, unrepeated column every level is a value, so + // `value_count` is O(1) (`Absent`/`Uniform` def levels); only + // nullable or nested columns pay the O(chunk_size) def-level scan. + let vals_in_chunk = chunk_def.value_count(chunk_size, self.max_def_level); + if vals_in_chunk == 0 { + return chunk_size; + } + // Ask the encoder how many of the next values fit in one page byte + // budget. Dispatch on whether the caller supplied gather indices; + // this mirrors how `write_mini_batch` picks `write_gather` vs + // `write`. + let fit = match value_indices { + Some(idx) => { + let end = (values_offset + vals_in_chunk).min(idx.len()); + let start = values_offset.min(end); + E::count_values_within_byte_budget_gather(values, &idx[start..end], budget) + } + None => { + E::count_values_within_byte_budget(values, values_offset, vals_in_chunk, budget) + } + }; + match fit { + None => chunk_size, + Some(values_per_subbatch) => { + // Convert the value count back into a level count. For a + // non-nullable column this is a no-op; for nullable/nested + // columns scale by the chunk's observed value-to-level + // ratio. + let levels_per_subbatch = if vals_in_chunk == chunk_size { + values_per_subbatch + } else { + (values_per_subbatch * chunk_size) + .div_ceil(vals_in_chunk) + .max(1) + }; + chunk_size.min(levels_per_subbatch.max(1)) + } + } + } +} diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 2ea3376ae708..d9adacff4101 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -90,6 +90,42 @@ pub trait ColumnValueEncoder { /// Write the values at the indexes in `indices` to this [`ColumnValueEncoder`] fn write_gather(&mut self, values: &Self::Values, indices: &[usize]) -> Result<()>; + /// Returns the largest `k` such that the first `k` values in + /// `values[offset..offset + len]` encode to at most `byte_budget` + /// bytes — i.e. how many values fit in a single page byte budget. + /// + /// Returns `len` if every value fits. Returns at least 1 if a single + /// value alone exceeds the budget, matching parquet's "at least one + /// value per data page" rule. + /// + /// `None` means "no cheap estimate available"; the caller stays on + /// the batched fast path and lets the post-write + /// `should_add_data_page` check handle bounding. + /// + /// Implementations should short-circuit aggressively: the typical + /// case is "everything fits, return `len`", and the next-most-common + /// case is "one wide value, return 1." The variable-width walk only + /// needs to be precise when the chunk is genuinely near the budget. + fn count_values_within_byte_budget( + _values: &Self::Values, + _offset: usize, + _len: usize, + _byte_budget: usize, + ) -> Option { + None + } + + /// As [`Self::count_values_within_byte_budget`] but using gather + /// `indices` rather than a contiguous range. Returns the number of + /// `indices` that fit, not the maximum index value. + fn count_values_within_byte_budget_gather( + _values: &Self::Values, + _indices: &[usize], + _byte_budget: usize, + ) -> Option { + None + } + /// Returns the number of buffered values fn num_values(&self) -> usize; @@ -247,6 +283,39 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { self.write_slice(&slice) } + fn count_values_within_byte_budget( + values: &[T::T], + offset: usize, + len: usize, + byte_budget: usize, + ) -> Option { + // Clamp so that a caller-supplied `len` that overruns the input + // (e.g. a level/value mismatch the encoder will reject later) + // returns an estimate instead of panicking here. + let end = (offset + len).min(values.len()); + let start = offset.min(end); + count_within_budget::( + end - start, + byte_budget, + values[start..end].iter().map(Some), + ) + } + + fn count_values_within_byte_budget_gather( + values: &[T::T], + indices: &[usize], + byte_budget: usize, + ) -> Option { + // `values.get` yields `None` for an out-of-range index (defensive + // against a level/value mismatch the encoder rejects later); such a + // position is counted but contributes no bytes. + count_within_budget::( + indices.len(), + byte_budget, + indices.iter().map(|&i| values.get(i)), + ) + } + fn num_values(&self) -> usize { self.num_values } @@ -411,3 +480,75 @@ where } } } + +/// Plain-encoded byte cost of a single value of type `T::T`. +/// +/// Derived from [`ParquetValueType::dict_encoding_size`] (which returns +/// `(per-value overhead, value-bytes)`) so we don't add a parallel +/// per-value-size hook to the trait. Mirrors the dispatch in +/// `KeyStorage::push` (`encodings/encoding/dict_encoder.rs`). +/// +/// Placed at the end of the module deliberately. Inserting it above the +/// `ColumnValueEncoder` trait shifts the trait and `ColumnValueEncoderImpl` +/// within the compiled module enough to perturb downstream code placement, +/// which measurably regresses unrelated arrow-writer string benchmarks +/// (~5-9% on `string` / `string_and_binary_view`). Defining it last keeps +/// the hot encoder code at the offsets it has on `main`. +#[inline] +fn plain_encoded_byte_size(value: &T::T) -> usize { + let (overhead, bytes) = value.dict_encoding_size(); + match ::PHYSICAL_TYPE { + // Plain BYTE_ARRAY = 4-byte length prefix + payload. + Type::BYTE_ARRAY => overhead + bytes, + // Plain FLBA = raw bytes only; `dict_encoding_size`'s length prefix + // is irrelevant here, so the encoder passes `type_length` directly. + Type::FIXED_LEN_BYTE_ARRAY => bytes, + // Numeric/bool are short-circuited by the caller via + // `mem::size_of`, so this is unreachable in practice; fall back to + // `overhead` defensively. + _ => overhead, + } +} + +/// How many leading values fit in `byte_budget` bytes, shared by the two +/// `ColumnValueEncoder::count_values_within_byte_budget*` methods (one walks a +/// contiguous slice, the other gathers by index). +/// +/// `n` is the answer when everything fits; `vals` yields each candidate value, +/// or `None` for a position that should still be counted but contributes no +/// bytes (an out-of-range gather index). The boundary value that crosses the +/// budget is included in the count so the caller's page-flush check trips on +/// this mini-batch rather than leaving a sliver for the next page; this also +/// catches a lone outlier wherever it lands among small values. +/// +/// Defined at the end of the module alongside `plain_encoded_byte_size` for +/// the same reason — see that function's note on code placement and the +/// `string` / `string_and_binary_view` benchmarks. +#[inline] +fn count_within_budget<'a, T: DataType>( + n: usize, + byte_budget: usize, + vals: impl Iterator>, +) -> Option +where + T::T: 'a, +{ + // Fixed-size physical types have a constant per-value byte cost, so the + // answer is one division — no walk needed. + let phys = ::PHYSICAL_TYPE; + if phys != Type::BYTE_ARRAY && phys != Type::FIXED_LEN_BYTE_ARRAY { + let per = std::mem::size_of::().max(1); + return Some((byte_budget / per).max(1).min(n)); + } + // Variable-width: accumulate, exit at the first value past the budget. + let mut cum: usize = 0; + for (i, v) in vals.enumerate() { + if let Some(v) = v { + cum = cum.saturating_add(plain_encoded_byte_size::(v)); + } + if cum > byte_budget { + return Some(i + 1); + } + } + Some(n) +} diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 4e53230bbf89..090350f53aae 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -49,8 +49,11 @@ use crate::file::properties::{ use crate::file::statistics::{Statistics, ValueStatistics}; use crate::schema::types::{ColumnDescPtr, ColumnDescriptor}; +mod byte_budget_chunker; pub(crate) mod encoder; +use byte_budget_chunker::ByteBudgetChunker; + macro_rules! downcast_writer { ($e:expr, $i:ident, $b:expr) => { match $e { @@ -374,6 +377,24 @@ impl<'a> LevelDataRef<'a> { Self::Uniform { value, .. } => Self::Uniform { value, count: len }, } } + + /// Count of positions in this slice that represent an actual value + /// (definition level equal to `max_def`). `Absent` means the column has + /// `max_def == 0` and every position is a value, so the implicit count + /// is the caller-supplied `total`. + pub(crate) fn value_count(self, total: usize, max_def: i16) -> usize { + match self { + Self::Absent => total, + Self::Materialized(values) => values.iter().filter(|&&d| d == max_def).count(), + Self::Uniform { value, count } => { + if value == max_def { + count + } else { + 0 + } + } + } + } } /// Typed column writer for a primitive column. @@ -545,6 +566,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { } else { self.props.write_batch_size() }; + let chunker = ByteBudgetChunker::new(&self.descr, &self.props, base_batch_size); while levels_offset < num_levels { let mut end_offset = num_levels.min(levels_offset + base_batch_size); @@ -555,14 +577,45 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { } } - values_offset += self.write_mini_batch( + let chunk_size = end_offset - levels_offset; + let chunk_def = def_levels.slice(levels_offset, chunk_size); + let chunk_rep = rep_levels.slice(levels_offset, chunk_size); + + // Key decision point: can we write this whole chunk as one + // mini-batch (the common case — small or fixed-width values, no + // further page-size accounting needed), or must we fall back to + // byte-budget-aware sub-batching to keep a page from overshooting + // `data_page_size_limit`? `pick_sub_batch_size` returns + // `chunk_size` for the former. + let sub_batch_size = chunker.pick_sub_batch_size( + &self.encoder, values, - values_offset, value_indices, - end_offset - levels_offset, - def_levels.slice(levels_offset, end_offset - levels_offset), - rep_levels.slice(levels_offset, end_offset - levels_offset), - )?; + chunk_def, + values_offset, + chunk_size, + ); + + if sub_batch_size >= chunk_size { + values_offset += self.write_mini_batch( + values, + values_offset, + value_indices, + chunk_size, + chunk_def, + chunk_rep, + )?; + } else { + values_offset += self.write_granular_chunk( + values, + values_offset, + value_indices, + chunk_size, + chunk_def, + chunk_rep, + sub_batch_size, + )?; + } levels_offset = end_offset; } @@ -713,6 +766,69 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { }) } + /// Writes a chunk in `sub_batch_size`-level sub-batches, checking the + /// data page byte limit after each. This keeps the page size close to + /// `data_page_size_limit` instead of overshooting it by a whole chunk. + /// + /// For repeated/nested columns sub-batches step from one `rep == 0` + /// boundary to the next so a record never spans data pages, matching + /// the parquet format rule. + /// + /// Returns the total number of values consumed across all sub-batches. + /// + /// `#[inline(never)]` keeps this slow path — only reached for + /// variable-width columns whose values need page splitting — out of + /// the hot `write_batch_internal` loop. + #[allow(clippy::too_many_arguments)] + #[inline(never)] + fn write_granular_chunk( + &mut self, + values: &E::Values, + values_offset: usize, + value_indices: Option<&[usize]>, + chunk_size: usize, + chunk_def: LevelDataRef<'_>, + chunk_rep: LevelDataRef<'_>, + sub_batch_size: usize, + ) -> Result { + // The chunker always sizes a sub-batch to at least one level, so each + // iteration below makes progress (`sub_end > sub_start`). + debug_assert!(sub_batch_size >= 1, "chunker must size at least one level"); + let mut values_consumed = 0; + let mut sub_start = 0; + while sub_start < chunk_size { + let sub_end = match chunk_rep { + LevelDataRef::Materialized(levels) => { + // Pack up to `sub_batch_size` levels per mini-batch, then + // extend to the next record boundary (rep == 0) so a + // record never spans data pages. Packing whole records + // rather than stepping one record at a time avoids + // calling `write_mini_batch` per record: records average + // only a handful of levels, so a record-at-a-time step + // would issue many more mini-batches than necessary. + let mut e = (sub_start + sub_batch_size).min(chunk_size); + while e < chunk_size && levels[e] != 0 { + e += 1; + } + e + } + _ => (sub_start + sub_batch_size).min(chunk_size), + }; + let sub_len = sub_end - sub_start; + let written = self.write_mini_batch( + values, + values_offset + values_consumed, + value_indices, + sub_len, + chunk_def.slice(sub_start, sub_len), + chunk_rep.slice(sub_start, sub_len), + )?; + values_consumed += written; + sub_start = sub_end; + } + Ok(values_consumed) + } + /// Creates a new streaming level encoder appropriate for the writer version. fn create_level_encoder(max_level: i16, props: &WriterProperties) -> LevelEncoder { match props.writer_version() { @@ -2676,6 +2792,310 @@ mod tests { assert_eq!(other_values, vec![10]); } + #[test] + fn test_column_writer_caps_page_size_for_large_byte_array_values() { + // Regression: the post-write data page byte limit check only fires + // at mini-batch boundaries, so a 1024-row mini-batch of multi-MiB + // BYTE_ARRAY values used to buffer multiple GiB into a single page + // before the limit was even consulted. With the threshold-based + // granular mode this batch should split into ~one page per value. + let value_size = 64 * 1024; // 64 KiB per value + let page_byte_limit = 16 * 1024; // 16 KiB page limit + let num_rows = 64; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + // Default write_batch_size (1024) — without the fix this + // buffers the entire input into a single ~4 MiB page. + .build(); + + let data: Vec<_> = (0..num_rows) + .map(|i| ByteArray::from(vec![i as u8; value_size])) + .collect(); + let pages = write_and_collect_pages::(props, 0, 0, &data, None, None); + + // Every value must end up somewhere. + let total_values: u32 = pages.data_pages.iter().map(|(_, n)| n).sum(); + assert_eq!(total_values as usize, num_rows); + // Without the fix this assertion fired with one ~4 MiB page; the + // threshold splits the input so that no page holds more than a + // single oversized value's worth of bytes. + assert!( + pages.data_pages.len() >= num_rows / 2, + "expected pages to be cut close to one per value, got {:?}", + pages.data_pages, + ); + // Each page must be bounded by roughly one value's worth of bytes; + // parquet allows a single oversized value to occupy a page by + // itself but never lets us pile many of them together. + for (size, _) in &pages.data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound ({}B) — pages {:?}", + value_size + 64, + pages.data_pages, + ); + } + } + + #[test] + fn test_column_writer_caps_page_size_for_large_values_in_list() { + // Coverage for the Materialized-rep branch of + // `write_granular_chunk`. The flat-column regression test + // exercises the per-level step; this exercises the + // record-by-record step used when rep levels are present. + // + // Column is `list` (max_def = 1, max_rep = 1) + // with 3 records of 3 large blobs each. The page byte limit is + // smaller than a single blob, so granular mode kicks in, and the + // Materialized-rep arm of `write_granular_chunk` steps from one + // `rep == 0` boundary to the next so a record never spans pages. + let value_size = 32 * 1024; + let page_byte_limit = 16 * 1024; + let values_per_record = 3; + let num_records = 3; + let num_values = values_per_record * num_records; + + // rep levels: 0, 1, 1, 0, 1, 1, 0, 1, 1 + let mut rep_levels = Vec::with_capacity(num_values); + for _ in 0..num_records { + rep_levels.push(0i16); + rep_levels.extend(std::iter::repeat_n(1i16, values_per_record - 1)); + } + let def_levels = vec![1i16; num_values]; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(); + + let data: Vec<_> = (0..num_values) + .map(|i| ByteArray::from(vec![i as u8; value_size])) + .collect(); + let pages = write_and_collect_pages::( + props, + 1, + 1, + &data, + Some(&def_levels), + Some(&rep_levels), + ); + let data_pages = pages.data_pages; + + // The Materialized-rep arm groups levels by record, and each + // record's bytes blow the page byte limit on its own, so we get + // exactly one page per record. + assert_eq!( + data_pages.len(), + num_records, + "expected one data page per record, got {data_pages:?}" + ); + for (bytes, n_values) in &data_pages { + assert_eq!( + *n_values as usize, values_per_record, + "each page must hold a whole record's leaves, got {data_pages:?}" + ); + // Each page is one full record (its leaves cannot be split), + // so allow up to `values_per_record` blobs of payload plus a + // small fudge for level encoding overhead. + let upper_bound = values_per_record * (value_size + 16); + assert!( + *bytes <= upper_bound, + "page size {bytes} exceeds whole-record bound ({upper_bound}); pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_caps_page_size_with_nullable_large_values() { + // Coverage for `LevelDataRef::value_count` on Materialized def + // levels: a nullable column with mixed nulls and large values. + // `value_count` must return the actual non-null count so the + // byte estimate reflects bytes that will actually be written, + // not the level count. + let value_size = 32 * 1024; + let page_byte_limit = 16 * 1024; + let num_levels = 32; + + // Alternating null / non-null: 16 nulls and 16 values. + let def_levels: Vec = (0..num_levels as i16).map(|i| i % 2).collect(); + let num_values = def_levels.iter().filter(|&&d| d == 1).count(); + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(); + + let data: Vec<_> = (0..num_values) + .map(|i| ByteArray::from(vec![i as u8; value_size])) + .collect(); + let pages = + write_and_collect_pages::(props, 1, 0, &data, Some(&def_levels), None); + let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect(); + + // With 16 actual values of 32 KiB each and a 16 KiB page limit, + // every non-null value should get its own page (plus possibly + // adjacent nulls). At minimum, the number of pages must be + // roughly the value count, not 1 (which is what `main` produced). + assert!( + data_pages.len() >= num_values / 2, + "expected at least {} pages for {num_values} large values, got {} pages: {data_pages:?}", + num_values / 2, + data_pages.len(), + ); + // No page contains more than ~one value's worth of payload bytes. + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_dict_enabled_large_values_post_spill() { + // While dictionary encoding is active, `has_dictionary()` short- + // circuits `estimated_value_bytes` — the byte estimate is plain- + // encoded size but dict-encoded pages only store small RLE + // indices, so we'd otherwise shrink pages spuriously. Once the + // dictionary spills (each value is large + unique), plain + // encoding takes over and the byte-budget sub-batch kicks in. + // + // This test makes sure the writer survives that transition and + // produces bounded pages thereafter. + let value_size = 64 * 1024; + let page_byte_limit = 16 * 1024; + let num_rows = 32; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(true) + // Force a small dict so it spills quickly even though + // each value here is unique. + .set_dictionary_page_size_limit(1024) + .set_data_page_size_limit(page_byte_limit) + // Small mini-batches so dict fallback happens part-way + // through the input, leaving subsequent mini-batches to + // exercise the post-spill plain-encoding path that the + // page-size fix actually targets. + .set_write_batch_size(4) + .build(); + + let data: Vec<_> = (0..num_rows) + .map(|i| ByteArray::from(vec![i as u8; value_size])) + .collect(); + let pages = write_and_collect_pages::(props, 0, 0, &data, None, None); + let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect(); + + // After spill, plain encoding writes one ~64 KiB value per page. + // Without the fix, post-spill writes still buffered all 32 + // values into a single ~2 MiB page. + assert!( + data_pages.len() >= num_rows / 2, + "expected >= {} data pages after dict spill, got {} ({data_pages:?})", + num_rows / 2, + data_pages.len(), + ); + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_caps_dictionary_page_size() { + // A column of large *distinct* values with dictionary encoding on: + // the dictionary page accumulates the values themselves, and its + // spill check runs only once per mini-batch. Without bounding the + // dictionary-encoding mini-batch, one `write_batch_size` mini-batch + // would intern `write_batch_size * value_size` bytes into the + // dictionary page before the check fires (~16 MiB here). The chunker + // must sub-batch the dictionary-encoding phase too. + let value_size = 8 * 1024; + let dict_page_limit = 64 * 1024; + let num_rows = 2048; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(dict_page_limit) + .build(); + + let data: Vec<_> = (0..num_rows) + .map(|i| { + // each value distinct, so the dictionary cannot dedup them + let mut v = vec![0u8; value_size]; + v[..8].copy_from_slice(&(i as u64).to_le_bytes()); + ByteArray::from(v) + }) + .collect(); + let pages = write_and_collect_pages::(props, 0, 0, &data, None, None); + let dict_page_size = pages.dict_page_size; + + assert!( + dict_page_size > 0, + "expected the column to dictionary-encode" + ); + // Bounded near the limit (~2x from the post-mini-batch check). Before + // the fix the dictionary page reached num_rows * value_size (~16 MiB, + // 256x the limit). + assert!( + dict_page_size <= 3 * dict_page_limit, + "dictionary page {dict_page_size} exceeds 3x the {dict_page_limit} limit", + ); + } + + #[test] + fn test_column_writer_caps_page_size_for_fixed_len_byte_array() { + // Coverage for `ParquetValueType::byte_size` override on + // `FixedLenByteArray`. With `type_length = 1`, each plain-encoded + // value is one byte, so a 4-byte page byte limit forces the + // sub-batch sizer to write ~4 values per page rather than one + // page for the whole batch. + let page_byte_limit = 4; + let num_values = 128; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(); + + let data: Vec<_> = (0..num_values) + .map(|i| { + let mut fla = FixedLenByteArray::default(); + fla.set_data(Bytes::from(vec![i as u8])); + fla + }) + .collect(); + let pages = + write_and_collect_pages::(props, 0, 0, &data, None, None); + let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect(); + + // Without the fix this is a single 128-byte page; with the fix + // the byte budget caps each page at ~`page_byte_limit` bytes. + assert!( + data_pages.len() >= num_values / 8, + "expected pages capped by byte budget, got {data_pages:?}" + ); + for size in &data_pages { + assert!( + *size <= page_byte_limit * 4, + "page size {size} larger than expected; pages {data_pages:?}" + ); + } + } + #[test] fn test_bool_statistics() { let stats = statistics_roundtrip::(&[true, false, false, true]); @@ -4309,6 +4729,69 @@ mod tests { get_typed_column_writer::(column_writer) } + /// Pages collected by [`write_and_collect_pages`]. + struct CollectedPages { + /// `(compressed byte size, value count)` for every data page, in order. + data_pages: Vec<(usize, u32)>, + /// Largest dictionary page seen, or 0 if the column wasn't dict-encoded. + dict_page_size: usize, + } + + /// Writes `data` (with optional def/rep levels) through a raw + /// `ColumnWriterImpl` configured by `props`, then re-reads the file and + /// returns its page layout. Shared by the page-size regression tests so + /// each only has to express its props, input, and assertions. + fn write_and_collect_pages( + props: WriterProperties, + max_def_level: i16, + max_rep_level: i16, + data: &[T::T], + def_levels: Option<&[i16]>, + rep_levels: Option<&[i16]>, + ) -> CollectedPages { + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let mut writer = + get_test_column_writer::(page_writer, max_def_level, max_rep_level, Arc::new(props)); + writer.write_batch(data, def_levels, rep_levels).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let read_props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(read_props), + ) + .unwrap(), + ); + + let mut collected = CollectedPages { + data_pages: Vec::new(), + dict_page_size: 0, + }; + while let Some(page) = page_reader.get_next_page().unwrap() { + match page.page_type() { + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => { + collected + .data_pages + .push((page.buffer().len(), page.num_values())); + } + PageType::DICTIONARY_PAGE => { + collected.dict_page_size = collected.dict_page_size.max(page.buffer().len()); + } + _ => {} + } + } + collected + } + /// Returns column reader. fn get_test_column_reader( page_reader: Box, @@ -4717,6 +5200,42 @@ mod tests { } } + #[test] + fn test_level_data_ref_value_count() { + // `value_count` is what the byte-budget chunker uses to convert a + // chunk's level span into a leaf-value count. It must work for any + // column shape — flat, nullable, or nested — because the leaf + // values array is decoupled from the rep/def level stream. + let max_def = 2; + // Non-nullable / unrepeated: no def levels materialized — every + // level is a value. + assert_eq!(LevelDataRef::Absent.value_count(64, max_def), 64); + // Uniform run of present values, and of nulls. + assert_eq!( + LevelDataRef::Uniform { + value: max_def, + count: 40 + } + .value_count(40, max_def), + 40 + ); + assert_eq!( + LevelDataRef::Uniform { + value: max_def - 1, + count: 40 + } + .value_count(40, max_def), + 0 + ); + // Materialized def levels (nullable / nested): only levels equal to + // `max_def` are values; empty-list / null levels are not. + let levels = [2i16, 0, 2, 1, 2, 2, 0]; + assert_eq!( + LevelDataRef::Materialized(&levels).value_count(levels.len(), max_def), + 4 + ); + } + #[test] fn test_uniform_def_levels_all_null() { // All-null column: def_level=0 (null) for every slot, no values written. From 0eee8782aa92aab9f9bb7603e1fe5288663b9590 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 3 Jun 2026 17:10:13 -0500 Subject: [PATCH 50/51] fix(parquet): bound page size for arrow byte-array column writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement `ColumnValueEncoder::count_values_within_byte_budget_gather` for `ByteArrayEncoder`, the encoder real `ArrowWriter` users hit, so the page-size bound from the previous commit also fires for arrow string/binary columns (the generic path only covered `ColumnValueEncoderImpl`). The impl stays off the hot path for small values via cheap O(1) upper bounds before any per-value walk: - Offset-backed arrays (`Utf8`/`LargeUtf8`/`Binary`/`LargeBinary`): the span `offsets[last+1] - offsets[first]` bounds the chunk payload in O(1); exact even for nullable columns (skipped positions add zero), so sparse `indices` skip the per-value walk too. - View arrays (`Utf8View`/`BinaryView`): lengths live in the low 32 bits of each view word, so an O(1) `n * (max_value_len + 4)` bound skips the scan in the common case; otherwise scan lengths with no data-buffer deref. - Dictionary input: treated as always-fitting — dict-encoded arrow input implies values small enough to dedup, the opposite of the blob case this targets, and a per-key walk measurably regressed the bench. Includes the arrow-writer unit tests for granular-mode round-trip and the all-null string column. Co-Authored-By: Claude Opus 4.8 (1M context) --- parquet/src/arrow/arrow_writer/byte_array.rs | 181 ++++++++++++++++++- parquet/src/arrow/arrow_writer/mod.rs | 118 ++++++++++++ 2 files changed, 298 insertions(+), 1 deletion(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 9cb0718b4d84..145431c26465 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -30,10 +30,12 @@ use crate::geospatial::statistics::GeospatialStatistics; use crate::schema::types::ColumnDescPtr; use crate::util::bit_util::num_required_bits; use crate::util::interner::{Interner, Storage}; +use arrow_array::types::ByteArrayType; use arrow_array::{ Array, ArrayAccessor, BinaryArray, BinaryViewArray, DictionaryArray, FixedSizeBinaryArray, - LargeBinaryArray, LargeStringArray, StringArray, StringViewArray, + GenericByteArray, LargeBinaryArray, LargeStringArray, StringArray, StringViewArray, }; +use arrow_buffer::{ArrowNativeType, Buffer}; use arrow_schema::DataType; macro_rules! downcast_dict_impl { @@ -481,6 +483,84 @@ impl ColumnValueEncoder for ByteArrayEncoder { Ok(()) } + fn count_values_within_byte_budget_gather( + values: &Self::Values, + indices: &[usize], + byte_budget: usize, + ) -> Option { + // `ByteArrayEncoder` only ever writes via `write_gather`, so this + // is the relevant method. + // + // Two-stage walk for the simple offset-buffer byte array types: + // 1. If indices are contiguous, compute the total payload in + // O(1) via a single subtraction on the offsets buffer. + // When the total fits the budget — the overwhelmingly + // common "small values" case — return immediately. + // 2. Otherwise, walk per-value byte sizes from the offsets + // buffer (still cheap, no slice/UTF-8 construction) and + // exit at the first value that pushes the cumulative sum + // past the budget. This bounds skewed distributions: an + // outlier value is caught wherever it lands in the chunk. + let count = match values.data_type() { + DataType::Utf8 => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::LargeUtf8 => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::Binary => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::LargeBinary => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + // View arrays carry each value's length in the low 32 bits of + // its u128 view word, so lengths are scannable without touching + // any data buffer — and the common small-value case skips even + // that scan via an O(1) conservative bound. + DataType::Utf8View => { + let array = values.as_any().downcast_ref::().unwrap(); + count_within_budget_views( + array.views(), + indices, + byte_budget, + max_view_value_len(array.data_buffers()), + ) + } + DataType::BinaryView => { + let array = values.as_any().downcast_ref::().unwrap(); + count_within_budget_views( + array.views(), + indices, + byte_budget, + max_view_value_len(array.data_buffers()), + ) + } + // The values in an arrow dictionary are already small and + // deduplicated, so there is nothing to bound — treat every + // chunk as fitting and stay on the batched path. (A per-value + // walk through dict keys on every chunk also measured ~+30-80% + // slower than `main`.) + DataType::Dictionary(_, _) => indices.len(), + // Every byte-array type `ByteArrayEncoder` is constructed for + // has an explicit arm above. A `Dictionary(value = FixedSizeBinary)` + // column hits the `Dictionary(_, _)` arm (its `values.data_type()` + // is `Dictionary`), and a bare `FixedSizeBinary` column is routed + // to the generic column writer, never this encoder — so no other + // type can reach here. + data_type => unreachable!("ByteArrayEncoder cannot be constructed for {data_type:?}"), + }; + Some(count) + } + fn num_values(&self) -> usize { match &self.dict_encoder { Some(encoder) => encoder.indices.len(), @@ -593,6 +673,105 @@ where } } +/// Upper bound on any single value's byte length in a view array. +fn max_view_value_len(buffers: &[Buffer]) -> usize { + /// Bytes that fit inline in a u128 view word (the rest is len + prefix). + const MAX_INLINE_VIEW_LEN: usize = 12; + // An out-of-line view's data is a contiguous slice of exactly one data + // buffer, so it cannot exceed the largest buffer; inline views hold at + // most `MAX_INLINE_VIEW_LEN`. Loose (a value is usually far smaller than + // a whole buffer) but O(number of buffers) and always sound. + buffers + .iter() + .map(|b| b.len()) + .max() + .unwrap_or(0) + .max(MAX_INLINE_VIEW_LEN) +} + +/// Number of leading `indices` whose cumulative plain-encoded size fits +/// `byte_budget` (boundary value included), for view arrays (`Utf8View`, +/// `BinaryView`). +fn count_within_budget_views( + views: &[u128], + indices: &[usize], + byte_budget: usize, + max_value_len: usize, +) -> usize { + // Each plain-encoded BYTE_ARRAY value carries a 4-byte length prefix, so + // the budget is compared against `value_len + size_of::()` — the + // bytes actually written to the page, not just the payload. + // + // Stage 1: O(1) conservative bound. View arrays have no prefix-sum + // offsets buffer, so the exact span subtraction used by + // `count_within_budget_offsets` is unavailable; instead bound every + // value by `max_value_len`. Skips the walk for the common small-value + // case (what view arrays are built for, and where there is nothing to + // bound). + let per_value = max_value_len + std::mem::size_of::(); + if indices.len().saturating_mul(per_value) <= byte_budget { + return indices.len(); + } + // Stage 2: exact per-value scan, reading each length from the low 32 + // bits of its u128 view word (no data-buffer dereference). + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let len = (views[*idx] as u32) as usize; + cum = cum.saturating_add(len + std::mem::size_of::()); + if cum > byte_budget { + return i + 1; + } + } + indices.len() +} + +/// Number of leading `indices` whose cumulative plain-encoded size fits +/// `byte_budget` (boundary value included), for offset-buffer byte arrays +/// (`Utf8`/`LargeUtf8`/`Binary`/`LargeBinary`). +/// +/// `indices` are assumed sorted ascending — they always are here, since +/// they come from `non_null_indices`, which is built in array order. +fn count_within_budget_offsets( + values: &GenericByteArray, + indices: &[usize], + byte_budget: usize, +) -> usize { + if indices.is_empty() { + return 0; + } + let n = indices.len(); + let first = indices[0]; + let last = indices[n - 1]; + let offsets = values.value_offsets(); + // Each plain-encoded value carries a 4-byte length prefix on the page. + let prefix_overhead = std::mem::size_of::(); + + // Stage 1: O(1) span upper bound. The span `offsets[last+1] - + // offsets[first]` covers every array position in `[first, last]`, a + // superset of `indices` — and the skipped positions in a nullable + // column are nulls with zero offset delta, so the span still equals the + // exact payload. If it fits the budget, every value fits. Covers the + // common small-value case for both non-null and (sparse) nullable + // columns. + if last >= first { + let payload = (offsets[last + 1] - offsets[first]).as_usize(); + if payload + n * prefix_overhead <= byte_budget { + return n; + } + } + + // Stage 2: scan per-index lengths from the offsets buffer. + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let len = (offsets[idx + 1] - offsets[*idx]).as_usize() + prefix_overhead; + cum = cum.saturating_add(len); + if cum > byte_budget { + return i + 1; + } + } + n +} + /// Computes the min and max for the provided array and indices /// /// This is a free function so it can be used with `downcast_op!` diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 79542caed9b7..86e6cf081fed 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -4904,6 +4904,124 @@ mod tests { assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4); } + #[test] + fn test_arrow_writer_granular_mode_roundtrip() { + // Granular mode subdivides chunks and writes more pages than + // `main`. Make sure the data we write back is bit-identical to + // what went in — page-count assertions elsewhere only prove + // pages were cut, not that the encoded data is correct. + // + // Mix value sizes so that the cumulative-byte-budget cutoff + // lands mid-chunk, exercising both batched and granular paths + // within the same `write_batch_internal` call. + let small = "tiny".to_string(); + let big = "x".repeat(64 * 1024); + let strings: Vec = (0..256) + .map(|i| { + if i % 16 == 0 { + big.clone() + } else { + small.clone() + } + }) + .collect(); + + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(strings.clone())) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap(); + let read = reader.next().unwrap().unwrap(); + assert!(reader.next().is_none(), "expected one batch"); + let col = read + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col.len(), strings.len()); + for (i, expected) in strings.iter().enumerate() { + assert_eq!( + col.value(i), + expected.as_str(), + "value mismatch at index {i}" + ); + } + } + + #[test] + fn test_arrow_writer_all_null_string_column() { + // The `LevelDataRef::value_count` Uniform branch with + // `value != max_def` (entirely-null chunk) must return 0 so the + // sub-batch sizer short-circuits to batch mode without trying + // to estimate byte budgets for non-existent values. + let num_rows = 1024; + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8, + true, + )])); + let nulls: Vec> = vec![None; num_rows]; + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(nulls)) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + // Re-parse the file: row group has one column, every row is + // null, all data pages report `num_rows / page_count` rows. + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let row_group = metadata.row_group(0); + let col_meta = row_group.column(0); + assert_eq!(row_group.num_rows() as usize, num_rows); + // Statistics record `null_count = num_rows` — proves every value + // was written as null. + if let Some(stats) = col_meta.statistics() { + assert_eq!( + stats.null_count_opt().unwrap_or(0) as usize, + num_rows, + "expected all-null column to report null_count = num_rows" + ); + } + + let mut reader = + SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap(); + let mut total_values = 0u32; + while let Some(page) = reader.get_next_page().unwrap() { + if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) { + total_values += page.num_values(); + } + } + assert_eq!( + total_values as usize, num_rows, + "expected every level position to be represented in some page" + ); + } + struct WriteBatchesShape { num_batches: usize, rows_per_batch: usize, From c65423731c50ed5b7ff46f6a63f55fa2c8b661cd Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 3 Jun 2026 17:10:13 -0500 Subject: [PATCH 51/51] test(parquet): page-layout coverage for byte-budget page sizing Add declarative `LayoutTest` cases covering the arrow write path's page layout under the new byte budget, replacing hand-rolled page-reading loops with exact page counts/sizes: - large `Utf8` strings and `Utf8View` strings (one page per value) - large values inside a list column (record-by-record stepping) - nullable large values (def-level value counting) - dictionary spill then plain-encode transition - FixedSizeBinary byte budget Also updates the existing `test_string` dict-spill expectations: the dictionary page is now bounded at its limit and spills one mini-batch earlier instead of overshooting. Co-Authored-By: Claude Opus 4.8 (1M context) --- parquet/tests/arrow_writer_layout.rs | 365 ++++++++++++++++++++++++++- 1 file changed, 357 insertions(+), 8 deletions(-) diff --git a/parquet/tests/arrow_writer_layout.rs b/parquet/tests/arrow_writer_layout.rs index b9d997beb289..1c63a3144391 100644 --- a/parquet/tests/arrow_writer_layout.rs +++ b/parquet/tests/arrow_writer_layout.rs @@ -19,14 +19,16 @@ use arrow::array::{Int32Array, StringArray}; use arrow::record_batch::RecordBatch; -use arrow_array::builder::{Int32Builder, ListBuilder}; +use arrow_array::builder::{BinaryBuilder, Int32Builder, ListBuilder}; +use arrow_array::types::Int32Type; +use arrow_array::{DictionaryArray, FixedSizeBinaryArray, StringViewArray}; use bytes::Bytes; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; use parquet::basic::{Encoding, PageType}; use parquet::file::metadata::PageIndexPolicy; use parquet::file::metadata::ParquetMetaData; -use parquet::file::properties::{ReaderProperties, WriterProperties}; +use parquet::file::properties::{EnabledStatistics, ReaderProperties, WriterProperties}; use parquet::file::reader::SerializedPageReader; use parquet::schema::types::ColumnPath; use std::sync::Arc; @@ -408,16 +410,16 @@ fn test_string() { columns: vec![ColumnChunk { pages: vec![ Page { - rows: 130, + rows: 126, page_header_size: 38, - compressed_size: 138, + compressed_size: 114, encoding: Encoding::RLE_DICTIONARY, page_type: PageType::DATA_PAGE, }, Page { - rows: 1250, + rows: 1254, page_header_size: 40, - compressed_size: 10000, + compressed_size: 10032, encoding: Encoding::PLAIN, page_type: PageType::DATA_PAGE, }, @@ -429,10 +431,16 @@ fn test_string() { page_type: PageType::DATA_PAGE, }, ], + // The byte-budget chunker sub-batches the dictionary + // phase. The mini-batch deliberately includes the value + // that crosses the 1000-byte limit so the spill triggers + // on this chunk rather than carrying a sliver into the + // next page, giving a 126-row dictionary page at 1008 + // bytes. dictionary_page: Some(Page { - rows: 130, + rows: 126, page_header_size: 38, - compressed_size: 1040, + compressed_size: 1008, encoding: Encoding::PLAIN, page_type: PageType::DICTIONARY_PAGE, }), @@ -599,3 +607,344 @@ fn test_per_column_data_page_size_limit() { assert_eq!(col_a_page_count, 16); assert_eq!(col_b_page_count, 1); } + +#[test] +fn test_fixed_size_binary() { + // FixedSizeBinary values larger than the data page byte limit. + let value_size = 1024usize; + let num_rows = 64usize; + let values: Vec = (0..num_rows) + .flat_map(|i| vec![i as u8; value_size]) + .collect(); + let array = + Arc::new(FixedSizeBinaryArray::try_new(value_size as i32, values.into(), None).unwrap()) + as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(4096) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + // 12 pages of 5 values (5 * 1024 = 5120 B, the boundary + // value pushes each page just past the 4096 B limit) plus + // a final page with the remaining 4 values. + pages: (0..12) + .map(|_| Page { + rows: 5, + page_header_size: 157, + compressed_size: 5120, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .chain(std::iter::once(Page { + rows: 4, + page_header_size: 157, + compressed_size: 4096, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + })) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_dictionary() { + // Arrow `DictionaryArray` input. + let num_rows = 2000; + let dict_values = StringArray::from_iter_values(["alpha", "beta", "gamma", "delta"]); + let keys = Int32Array::from_iter_values((0..num_rows).map(|i| i % 4)); + let array = + Arc::new(DictionaryArray::::try_new(keys, Arc::new(dict_values)).unwrap()) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(1000) + .set_data_page_size_limit(1000) + .set_write_batch_size(10) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: vec![Page { + rows: 2000, + page_header_size: 40, + compressed_size: 505, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }], + dictionary_page: Some(Page { + rows: 4, + page_header_size: 38, + compressed_size: 35, + encoding: Encoding::PLAIN, + page_type: PageType::DICTIONARY_PAGE, + }), + }], + }], + }, + }); +} + +#[test] +fn test_large_string() { + // Large `Utf8` values (64 KiB each) with a 16 KiB data page limit. + // + // Each value already exceeds the page byte budget, so the byte-budget + // chunker in `ByteArrayEncoder` (the offsets-buffer scan in + // `count_within_budget_offsets`) cuts one value per page instead of + // buffering the whole ~2 MiB column into a single page. This drives the + // real `ArrowWriter` user path; the lower-level column writer is covered + // by `test_column_writer_caps_page_size_for_large_byte_array_values`. + let value_size = 64 * 1024; + let strings: Vec = (0..32).map(|_| "x".repeat(value_size)).collect(); + let array = Arc::new(StringArray::from(strings)) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + // Disable statistics so page headers stay small and the layout is + // determined purely by the page-splitting logic under test. + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + // One 64 KiB value per page (4-byte length prefix + value). + pages: (0..32) + .map(|_| Page { + rows: 1, + page_header_size: 21, + compressed_size: 65540, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_large_string_view() { + // Same bytes and expected layout as `test_large_string`, but the input + // is a `Utf8View` array. View arrays expose no contiguous offsets + // buffer, so the arrow writer bounds pages via the view-specific scan + // (`count_within_budget_views`, reading each value's length from the + // low 32 bits of its view word) rather than the offsets scan. This + // confirms that path caps pages identically. + let value_size = 64 * 1024; + let strings: Vec = (0..32).map(|_| "x".repeat(value_size)).collect(); + let array = Arc::new(StringViewArray::from_iter_values( + strings.iter().map(|s| s.as_str()), + )) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: (0..32) + .map(|_| Page { + rows: 1, + page_header_size: 21, + compressed_size: 65540, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_large_values_in_list() { + // `list` with large leaf values, driving the record-by-record + // (Materialized-rep) arm of the byte-budget chunker through the arrow + // path. Because repetition levels are present, a list element's leaves + // can never span pages, so the chunker steps from one `rep == 0` + // boundary to the next. Three records of three 32 KiB blobs each, with + // a 16 KiB page limit, yield exactly one page per record (a whole + // record's ~96 KiB of leaves stays together even though it blows the + // budget — it cannot be split). The raw-writer analogue is + // `test_column_writer_caps_page_size_for_large_values_in_list`. + let value_size = 32 * 1024; + let mut builder = ListBuilder::new(BinaryBuilder::new()); + let mut byte = 0u8; + for _ in 0..3 { + for _ in 0..3 { + builder.values().append_value(vec![byte; value_size]); + byte = byte.wrapping_add(1); + } + builder.append(true); + } + let array = Arc::new(builder.finish()) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + // One record (3 leaves + rep/def levels) per page. + pages: (0..3) + .map(|_| Page { + rows: 1, + page_header_size: 21, + compressed_size: 98328, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_nullable_large_values() { + // Nullable `Utf8` column alternating null / 32 KiB value. The byte + // budget must count only the non-null values (the def-level value + // count), not the level count — otherwise the estimate would be wrong + // for sparse columns. With a 16 KiB page limit each 32 KiB value still + // gets its own page (carrying its adjacent leading null), giving 16 + // two-row pages. Mirrors the raw-writer + // `test_column_writer_caps_page_size_with_nullable_large_values`. + let value_size = 32 * 1024; + let big = "x".repeat(value_size); + let values: Vec> = (0..32) + .map(|i| if i % 2 == 1 { Some(big.clone()) } else { None }) + .collect(); + let array = Arc::new(StringArray::from(values)) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + // Each page holds one null + one 32 KiB value (2 rows). + pages: (0..16) + .map(|_| Page { + rows: 2, + page_header_size: 21, + compressed_size: 32778, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_dictionary_spill_large_values() { + // Dictionary encoding is enabled, but each value is large (64 KiB) and + // unique, so the dictionary spills almost immediately (1 KiB dict page + // limit). After the spill, plain encoding takes over and the byte-budget + // sub-batch bounds each page to a single value. The first value is + // interned into the dictionary page (one RLE_DICTIONARY data page + // referencing it); the remaining 31 fall back to PLAIN, one per page. + // Mirrors `test_column_writer_dict_enabled_large_values_post_spill`, + // exercising the same dict→plain transition via the arrow path. + let value_size = 64 * 1024; + let strings: Vec = (0..32) + .map(|i| format!("{i:05}") + &"x".repeat(value_size - 5)) + .collect(); + let array = Arc::new(StringArray::from(strings)) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + // Tiny dict limit so it spills after the first (already oversized) + // value, leaving the rest to exercise the post-spill plain path. + .set_dictionary_page_size_limit(1024) + .set_data_page_size_limit(16 * 1024) + .set_write_batch_size(4) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: std::iter::once(Page { + // The single value interned before the dict spilled. + rows: 1, + page_header_size: 17, + compressed_size: 2, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }) + .chain((0..31).map(|_| Page { + // Post-spill plain-encoded values, one per page. + rows: 1, + page_header_size: 21, + compressed_size: 65540, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + })) + .collect(), + dictionary_page: Some(Page { + rows: 1, + page_header_size: 21, + compressed_size: 65540, + encoding: Encoding::PLAIN, + page_type: PageType::DICTIONARY_PAGE, + }), + }], + }], + }, + }); +}