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-arith/src/temporal.rs b/arrow-arith/src/temporal.rs index 96a3b43a3811..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, @@ -63,6 +62,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 +103,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 +116,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 +159,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, @@ -249,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), ) } @@ -500,6 +500,8 @@ impl ExtractDatePartExt for PrimitiveArray { | DatePart::Day | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 + | DatePart::DayOfWeekSunday1 + | DatePart::DayOfWeekMonday1 | DatePart::DayOfYear | DatePart::Hour | DatePart::Minute @@ -536,6 +538,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 +582,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 +616,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 +650,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 +684,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 +718,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 +999,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 +1633,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 +1871,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 +2018,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 +2223,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 +2266,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 +2274,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 ", 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-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-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-array/src/ffi.rs b/arrow-array/src/ffi.rs index f50dd3420baa..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) @@ -281,7 +287,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 +312,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 +687,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-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/bigint/mod.rs b/arrow-buffer/src/bigint/mod.rs index 15faed43a130..e1db0c6141e0 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; @@ -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 @@ -807,6 +902,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 +1334,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 +1689,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!( @@ -1580,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(); + } } diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index 420bbf59f3be..52f1e3fbb510 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 { @@ -1013,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 @@ -1245,4 +1310,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}"); + } + } } 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-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); + } } diff --git a/arrow-cast/src/cast/decimal.rs b/arrow-cast/src/cast/decimal.rs index 3553f2b6a76f..9d1465567e0d 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 @@ -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:?}" @@ -777,7 +783,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 +791,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 +808,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, @@ -962,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/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/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(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> { @@ -172,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 @@ -183,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 ( @@ -633,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): @@ -2313,11 +2329,19 @@ 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| { - (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), @@ -2409,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, @@ -3349,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); @@ -3603,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, @@ -3830,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, @@ -3940,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![ @@ -3958,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!( @@ -5707,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::(); @@ -5734,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), @@ -6077,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()); @@ -9021,6 +9348,41 @@ mod tests { } } + #[test] + fn test_cast_fixed_size_list_to_list_preserves_field_metadata() { + use std::collections::HashMap; + + let metadata: HashMap = + 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 @@ -10112,6 +10474,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]); @@ -10449,8 +10901,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)); @@ -10489,8 +10941,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)); @@ -10554,8 +11006,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")), @@ -10703,6 +11155,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) { 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!( diff --git a/arrow-data/src/data.rs b/arrow-data/src/data.rs index 34d5cb6d3aba..0ce98aa090df 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())], @@ -2198,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, @@ -2262,6 +2268,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/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-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-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); diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs index 1d5e06c6871c..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)?; } } @@ -754,7 +763,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 +1055,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; @@ -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![ 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-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-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-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 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. 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-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(()) + } } 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()); diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index e95d01f2b592..fcbce82d5d9d 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::*; @@ -1281,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 @@ -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]); 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![ 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(_)) diff --git a/arrow-string/src/concat_elements.rs b/arrow-string/src/concat_elements.rs index cd4676d28752..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 @@ -505,30 +515,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 +686,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/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:?}" ))), 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/arrow/benches/coalesce_kernels.rs b/arrow/benches/coalesce_kernels.rs index b85c5cc532db..4affcc346e89 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 @@ -50,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), @@ -57,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), @@ -105,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, @@ -118,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, @@ -131,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", @@ -158,7 +257,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 +469,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 +617,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 +655,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 +741,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 +934,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 => diff --git a/arrow/src/util/data_gen.rs b/arrow/src/util/data_gen.rs index e54ab3499472..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(), ), @@ -178,6 +183,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 +238,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 +712,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/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} 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-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-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/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/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-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 e95c774f2926..38c577564dea 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -15,25 +15,28 @@ // 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::BorrowedShreddingState; 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, @@ -43,17 +46,75 @@ pub(crate) enum ShreddedPathStep<'a> { 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. -pub(crate) fn follow_shredded_path_element<'a>( - shredding_state: &BorrowedShreddingState<'a>, +/// 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<'_>, _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() { @@ -69,7 +130,7 @@ pub(crate) fn follow_shredded_path_element<'a>( 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()); }; @@ -90,15 +151,33 @@ 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 { .. } => { - // 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()), + } } } } @@ -154,7 +233,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 { @@ -356,15 +435,17 @@ 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, 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; @@ -447,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", ); } @@ -1740,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. @@ -1776,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: @@ -2493,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); @@ -4138,13 +4451,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 +4636,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 +4654,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, 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/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}" + ); } } diff --git a/parquet-variant/src/variant.rs b/parquet-variant/src/variant.rs index accff009045a..c9f175c3a610 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: @@ -292,7 +298,41 @@ pub enum Variant<'m, 'v> { } // We don't want this to grow because it could hurt performance of a frequently-created type. +#[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); + +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 +837,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 +878,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 +1037,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 +1078,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 +1119,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 +1160,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 +1181,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 +1215,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 +1250,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 +1266,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 +1301,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 +1313,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()), 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") diff --git a/parquet/benches/arrow_writer.rs b/parquet/benches/arrow_writer.rs index 9b22bb04b3e7..80d3e7144b19 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, @@ -136,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, @@ -392,6 +439,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)); @@ -401,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)); 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); 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/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index f56f9570adfb..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 { @@ -165,7 +167,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 +176,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 +185,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 +199,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 +345,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 +353,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,10 +473,94 @@ 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(()) } + 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(), @@ -554,15 +640,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 +662,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()); } } @@ -587,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/levels.rs b/parquet/src/arrow/arrow_writer/levels.rs index 73f922137620..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( @@ -153,6 +155,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 { @@ -216,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(); @@ -267,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(), @@ -280,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) @@ -301,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 @@ -310,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()) { @@ -320,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 = @@ -420,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, @@ -676,20 +852,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()); } } @@ -704,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), diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 04a798391429..86e6cf081fed 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 @@ -4901,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, diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index 6c7890b75c45..5a0083b7164d 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, @@ -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 52152988166f..b89db361eda1 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); } } @@ -495,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; @@ -507,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(); @@ -524,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(); @@ -563,7 +553,6 @@ mod test { } #[test] - #[allow(deprecated)] fn test_metadata_read_write_roundtrip_page_index() { let parquet_bytes = create_parquet_file(); @@ -572,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(); @@ -581,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(); @@ -1042,6 +1031,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/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/arrow/schema/extension.rs b/parquet/src/arrow/schema/extension.rs index a2e9c32dee8c..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 @@ -112,9 +115,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 +168,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/arrow/schema/primitive.rs b/parquet/src/arrow/schema/primitive.rs index b440753cc83b..c70885355d74 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}; @@ -168,18 +170,20 @@ 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), - ( - 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 +192,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 +224,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 +263,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 +287,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 {}", @@ -313,16 +303,18 @@ 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 { scale, precision }), _) => { + (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(*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) => { + check_decimal_length(type_length)?; if type_length <= 16 { decimal_128_type(scale, precision) } else { @@ -330,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 @@ -348,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) + ); + } +} diff --git a/parquet/src/basic.rs b/parquet/src/basic.rs index 17e5b9c321bb..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,353 +186,165 @@ 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<'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 { - scale: val.scale, - precision: val.precision, - } - } - 6 => { - prot.skip_empty_struct()?; - Self::Date - } - 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, - } - } - 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, - } - } - 10 => { - let val = IntType::read_thrift(&mut *prot)?; - Self::Integer { - is_signed: val.is_signed, - bit_width: val.bit_width, - } - } - 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 { - specification_version: val.specification_version, - } - } - 17 => { - let val = GeometryType::read_thrift(&mut *prot)?; - Self::Geometry { - crs: 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 { - crs: val.crs.map(|s| s.to_owned()), - algorithm: 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) +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(IntType { + bit_width, + is_signed, + }) } -} -impl WriteThrift for LogicalType { - const ELEMENT_TYPE: ElementType = ElementType::Struct; + /// Create a [`LogicalType::Decimal`] variant with the given `scale` and `precision` + pub fn decimal(scale: i32, precision: i32) -> Self { + Self::Decimal(DecimalType { scale, precision }) + } - 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() + /// 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(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(TimestampType { + is_adjusted_to_u_t_c, + unit, + }) } -} -write_thrift_field!(LogicalType, FieldType::Struct); + /// Create a [`LogicalType::Variant`] variant with the given `specification_version` + pub fn variant(specification_version: Option) -> Self { + Self::Variant(VariantType { + specification_version, + }) + } + + /// Create a [`LogicalType::Geometry`] variant with the given `crs` + pub fn geometry(crs: Option) -> Self { + 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(GeographyType { crs, algorithm }) + } +} // ---------------------------------------------------------------------- // Mirrors thrift enum `FieldRepetitionType` @@ -1224,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 @@ -1397,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, @@ -1427,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, }, @@ -1511,26 +1323,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 +1340,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 +1661,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 +1681,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 +1757,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 +1780,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 +2009,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 +2038,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/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 5d14ac6856f9..090350f53aae 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}; @@ -48,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 { @@ -373,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. @@ -544,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); @@ -554,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; } @@ -712,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() { @@ -1522,9 +1639,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 +1658,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() { @@ -2675,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]); @@ -4308,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, @@ -4449,10 +4933,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() @@ -4719,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. 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/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/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/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 { 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/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/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/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] 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); + } } 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..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})") @@ -457,7 +458,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 +544,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 +556,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 +568,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 +580,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 +640,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 +664,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 +676,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 +772,7 @@ mod tests { "field", None, PhysicalType::BYTE_ARRAY, - Some(LogicalType::Geometry { crs: None }), + Some(LogicalType::geometry(None)), ConvertedType::NONE, Repetition::REQUIRED, ) @@ -804,9 +784,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 +796,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 +808,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 +862,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 +1150,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 +1159,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..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" )); @@ -1449,10 +1448,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 +1460,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,16 +1475,13 @@ 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 { 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'" ); } @@ -1524,10 +1514,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/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(); } 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/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!( 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![ 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 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, + }), + }], + }], + }, + }); +} diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index c26a0edee682..2ac94b8c3786 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(); + non_uniform_encryption_plaintext_footer(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); - 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(); - - 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); + } + + non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_128_FOOTER_KEY, + AES_128_KEY_NAME_KEY, + ); - verify_encryption_test_file_read(file, decryption_properties); + 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); + verify_encryption_test_file_read(file, decryption_properties); + } + + non_uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY, AES_128_KEY_NAME_KEY); + + 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, + ); + } + + 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, + ); - read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); + 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"); + 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 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(); + decrypt_page_index_non_uniform(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); - 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(); - - 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( @@ -1262,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()); +} 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 +} 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()); 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 },