From 30b345b1e664b3341ed79bdbc586c209370e9091 Mon Sep 17 00:00:00 2001 From: 0lai0 Date: Tue, 9 Jun 2026 16:29:07 +0800 Subject: [PATCH 1/2] ParquetReader generic over f32/f64 with Arrow cast --- qdp/qdp-core/src/lib.rs | 2 +- qdp/qdp-core/src/reader.rs | 32 ++- qdp/qdp-core/src/readers/parquet.rs | 413 +++++++++++++++++----------- qdp/qdp-core/src/remote.rs | 2 +- qdp/qdp-core/tests/parquet_f32.rs | 161 +++++++++++ 5 files changed, 441 insertions(+), 169 deletions(-) create mode 100644 qdp/qdp-core/tests/parquet_f32.rs diff --git a/qdp/qdp-core/src/lib.rs b/qdp/qdp-core/src/lib.rs index ac5dd5fe94..822fba96e1 100644 --- a/qdp/qdp-core/src/lib.rs +++ b/qdp/qdp-core/src/lib.rs @@ -37,7 +37,7 @@ mod profiling; pub use error::{MahoutError, Result, cuda_error_to_string}; pub use gpu::memory::Precision; -pub use reader::{FloatElem, NullHandling, handle_float64_nulls}; +pub use reader::{FloatElem, NullHandling, handle_float32_nulls, handle_float64_nulls}; pub use types::{Dtype, Encoding}; // Throughput/latency pipeline runner: single path using QdpEngine and encode_batch in Rust. diff --git a/qdp/qdp-core/src/reader.rs b/qdp/qdp-core/src/reader.rs index a51fd334a1..11c430341f 100644 --- a/qdp/qdp-core/src/reader.rs +++ b/qdp/qdp-core/src/reader.rs @@ -45,7 +45,7 @@ //! } //! ``` -use arrow::array::{Array, Float64Array}; +use arrow::array::{Array, Float32Array, Float64Array}; use crate::error::Result; @@ -53,7 +53,7 @@ use crate::error::Result; /// /// Keeps f32 file data as `Vec` end-to-end once readers implement /// `DataReader`; today most readers use the default `T = f64`. -pub trait FloatElem: Copy + Send + Sync + 'static {} +pub trait FloatElem: Copy + Default + Send + Sync + 'static {} impl FloatElem for f32 {} impl FloatElem for f64 {} @@ -96,6 +96,34 @@ pub fn handle_float64_nulls( Ok(()) } +/// Append values from a `Float32Array` into `output`, applying the given null policy. +/// +/// When there are no nulls the fast path copies the underlying buffer directly. +pub fn handle_float32_nulls( + output: &mut Vec, + float_array: &Float32Array, + null_handling: NullHandling, +) -> crate::error::Result<()> { + if float_array.null_count() == 0 { + output.extend_from_slice(float_array.values()); + } else { + match null_handling { + NullHandling::FillZero => { + output.extend(float_array.iter().map(|opt| opt.unwrap_or(0.0))); + } + NullHandling::Reject => { + return Err(crate::error::MahoutError::InvalidInput( + "Null value encountered in Float32Array. \ + Use NullHandling::FillZero to replace nulls with 0.0, \ + or clean the data at the source." + .to_string(), + )); + } + } + } + Ok(()) +} + /// Generic data reader interface for batch quantum data. /// /// Implementations should read data in the format: diff --git a/qdp/qdp-core/src/readers/parquet.rs b/qdp/qdp-core/src/readers/parquet.rs index 0c9a3e162e..475b88d1e7 100644 --- a/qdp/qdp-core/src/readers/parquet.rs +++ b/qdp/qdp-core/src/readers/parquet.rs @@ -17,24 +17,180 @@ //! Parquet format reader implementation. use std::fs::File; +use std::marker::PhantomData; use std::path::Path; -use arrow::array::{Array, FixedSizeListArray, Float64Array, ListArray}; +use arrow::array::{Array, FixedSizeListArray, Float32Array, Float64Array, ListArray}; +use arrow::compute; use arrow::datatypes::DataType; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use crate::error::{MahoutError, Result}; -use crate::reader::{DataReader, NullHandling, StreamingDataReader, handle_float64_nulls}; +use crate::reader::{ + DataReader, FloatElem, NullHandling, StreamingDataReader, handle_float32_nulls, + handle_float64_nulls, +}; -/// Reader for Parquet files containing List or FixedSizeList columns. -pub struct ParquetReader { +// --------------------------------------------------------------------------- +// Internal sealed helper trait +// --------------------------------------------------------------------------- + +/// Zero-copy or Arrow-cast extraction from an Arrow array into an output `Vec`. +/// +/// Implemented for `f32` and `f64` only: +/// - same dtype → `extend_from_slice` directly from the Arrow buffer (zero alloc) +/// - cross dtype → `arrow::compute::cast` first, then `extend_from_slice` +/// - f32 → f64: exact, NaN preserved +/// - f64 → f32: values outside f32 range become ±Inf, NaN preserved +pub(crate) trait ArrowPrimitive: FloatElem { + fn extend_from_arrow_array( + output: &mut Vec, + array: &dyn Array, + null_handling: NullHandling, + ) -> Result<()>; + + fn collect_from_arrow_array( + array: &dyn Array, + null_handling: NullHandling, + ) -> Result> { + let mut out = Vec::new(); + Self::extend_from_arrow_array(&mut out, array, null_handling)?; + Ok(out) + } +} + +impl ArrowPrimitive for f64 { + fn extend_from_arrow_array( + output: &mut Vec, + array: &dyn Array, + null_handling: NullHandling, + ) -> Result<()> { + match array.data_type() { + DataType::Float64 => { + let arr = array + .as_any() + .downcast_ref::() + .ok_or_else(|| MahoutError::Io("Float64 downcast failed".to_string()))?; + handle_float64_nulls(output, arr, null_handling)?; + } + DataType::Float32 => { + let casted = compute::cast(array, &DataType::Float64) + .map_err(|e| MahoutError::Io(format!("Arrow cast f32→f64: {e}")))?; + let arr = casted + .as_any() + .downcast_ref::() + .ok_or_else(|| MahoutError::Io("Cast to Float64 failed".to_string()))?; + handle_float64_nulls(output, arr, null_handling)?; + } + other => { + return Err(MahoutError::InvalidInput(format!( + "Expected Float32 or Float64 values, got {other:?}" + ))); + } + } + Ok(()) + } +} + +impl ArrowPrimitive for f32 { + fn extend_from_arrow_array( + output: &mut Vec, + array: &dyn Array, + null_handling: NullHandling, + ) -> Result<()> { + match array.data_type() { + DataType::Float32 => { + let arr = array + .as_any() + .downcast_ref::() + .ok_or_else(|| MahoutError::Io("Float32 downcast failed".to_string()))?; + handle_float32_nulls(output, arr, null_handling)?; + } + DataType::Float64 => { + // f64 → f32: values outside f32 range become ±Inf; NaN is preserved + let casted = compute::cast(array, &DataType::Float32) + .map_err(|e| MahoutError::Io(format!("Arrow cast f64→f32: {e}")))?; + let arr = casted + .as_any() + .downcast_ref::() + .ok_or_else(|| MahoutError::Io("Cast to Float32 failed".to_string()))?; + handle_float32_nulls(output, arr, null_handling)?; + } + other => { + return Err(MahoutError::InvalidInput(format!( + "Expected Float32 or Float64 values, got {other:?}" + ))); + } + } + Ok(()) + } +} + +// --------------------------------------------------------------------------- +// Shared schema validator +// --------------------------------------------------------------------------- + +fn validate_float_list_schema(field: &arrow::datatypes::Field) -> Result<()> { + match field.data_type() { + DataType::List(child_field) => { + if !matches!( + child_field.data_type(), + DataType::Float32 | DataType::Float64 + ) { + return Err(MahoutError::InvalidInput(format!( + "Expected List or List column, got List<{:?}>", + child_field.data_type() + ))); + } + } + DataType::FixedSizeList(child_field, _) => { + if !matches!( + child_field.data_type(), + DataType::Float32 | DataType::Float64 + ) { + return Err(MahoutError::InvalidInput(format!( + "Expected FixedSizeList or FixedSizeList column, got FixedSizeList<{:?}>", + child_field.data_type() + ))); + } + } + _ => { + return Err(MahoutError::InvalidInput(format!( + "Expected List or FixedSizeList column, got {:?}", + field.data_type() + ))); + } + } + Ok(()) +} + +fn validate_float_list_or_scalar_schema(field: &arrow::datatypes::Field) -> Result<()> { + match field.data_type() { + DataType::Float32 | DataType::Float64 => Ok(()), + _ => validate_float_list_schema(field), + } +} + +// --------------------------------------------------------------------------- +// ParquetReader +// --------------------------------------------------------------------------- + +/// Reader for Parquet files containing `List` or +/// `FixedSizeList` columns. +/// +/// Generic over `T` (`f32` or `f64`): +/// - same dtype as the file → zero-copy path via `extend_from_slice` +/// - different dtype → `arrow::compute::cast` (f64→f32: overflow → ±Inf; NaN preserved) +pub struct ParquetReader { reader: Option, sample_size: Option, total_rows: usize, null_handling: NullHandling, + _phantom: PhantomData, } -impl ParquetReader { +#[allow(private_bounds)] +impl ParquetReader { /// Create a new Parquet reader. /// /// # Arguments @@ -48,7 +204,6 @@ impl ParquetReader { ) -> Result { let path = path.as_ref(); - // Verify file exists match path.try_exists() { Ok(false) => { return Err(MahoutError::Io(format!( @@ -85,31 +240,7 @@ impl ParquetReader { ))); } - let field = &schema.fields()[0]; - match field.data_type() { - DataType::List(child_field) => { - if !matches!(child_field.data_type(), DataType::Float64) { - return Err(MahoutError::InvalidInput(format!( - "Expected List column, got List<{:?}>", - child_field.data_type() - ))); - } - } - DataType::FixedSizeList(child_field, _) => { - if !matches!(child_field.data_type(), DataType::Float64) { - return Err(MahoutError::InvalidInput(format!( - "Expected FixedSizeList column, got FixedSizeList<{:?}>", - child_field.data_type() - ))); - } - } - _ => { - return Err(MahoutError::InvalidInput(format!( - "Expected List or FixedSizeList column, got {:?}", - field.data_type() - ))); - } - } + validate_float_list_schema(&schema.fields()[0])?; let total_rows = builder.metadata().file_metadata().num_rows() as usize; @@ -125,20 +256,21 @@ impl ParquetReader { sample_size: None, total_rows, null_handling, + _phantom: PhantomData, }) } } -impl DataReader for ParquetReader { - fn read_batch(&mut self) -> Result<(Vec, usize, usize)> { +impl DataReader for ParquetReader { + fn read_batch(&mut self) -> Result<(Vec, usize, usize)> { let reader = self .reader .take() .ok_or_else(|| MahoutError::InvalidInput("Reader already consumed".to_string()))?; - let mut all_data = Vec::new(); + let mut all_data: Vec = Vec::new(); let mut num_samples = 0; - let mut sample_size = None; + let mut sample_size: Option = None; for batch_result in reader { let batch = batch_result @@ -159,14 +291,7 @@ impl DataReader for ParquetReader { for i in 0..list_array.len() { let value_array = list_array.value(i); - let float_array = value_array - .as_any() - .downcast_ref::() - .ok_or_else(|| { - MahoutError::Io("List values must be Float64".to_string()) - })?; - - let current_size = float_array.len(); + let current_size = value_array.len(); if let Some(expected_size) = sample_size { if current_size != expected_size { @@ -180,8 +305,11 @@ impl DataReader for ParquetReader { all_data.reserve(current_size * self.total_rows); } - handle_float64_nulls(&mut all_data, float_array, self.null_handling)?; - + T::extend_from_arrow_array( + &mut all_data, + &*value_array, + self.null_handling, + )?; num_samples += 1; } } @@ -201,18 +329,12 @@ impl DataReader for ParquetReader { } let values = list_array.values(); - let float_array = values - .as_any() - .downcast_ref::() - .ok_or_else(|| MahoutError::Io("Values must be Float64".to_string()))?; - - handle_float64_nulls(&mut all_data, float_array, self.null_handling)?; - + T::extend_from_arrow_array(&mut all_data, values, self.null_handling)?; num_samples += list_array.len(); } _ => { return Err(MahoutError::Io(format!( - "Expected List or FixedSizeList, got {:?}", + "Expected List or FixedSizeList, got {:?}", column.data_type() ))); } @@ -236,20 +358,27 @@ impl DataReader for ParquetReader { } } -/// Streaming Parquet reader for List and FixedSizeList columns. +// --------------------------------------------------------------------------- +// ParquetStreamingReader +// --------------------------------------------------------------------------- + +/// Streaming Parquet reader for `List` and +/// `FixedSizeList` columns. /// -/// Reads Parquet files in chunks without loading entire file into memory. -/// Supports efficient streaming for large files via Producer-Consumer pattern. -pub struct ParquetStreamingReader { +/// Reads Parquet files in chunks without loading the entire file into memory. +/// Supports efficient streaming for large files via the Producer-Consumer pattern. +pub struct ParquetStreamingReader { reader: parquet::arrow::arrow_reader::ParquetRecordBatchReader, sample_size: Option, - leftover_data: Vec, + leftover_data: Vec, leftover_cursor: usize, pub total_rows: usize, null_handling: NullHandling, + _phantom: PhantomData, } -impl ParquetStreamingReader { +#[allow(private_bounds)] +impl ParquetStreamingReader { /// Create a new streaming Parquet reader. /// /// # Arguments @@ -263,7 +392,6 @@ impl ParquetStreamingReader { ) -> Result { let path = path.as_ref(); - // Verify file exists match path.try_exists() { Ok(false) => { return Err(MahoutError::Io(format!( @@ -300,34 +428,7 @@ impl ParquetStreamingReader { ))); } - let field = &schema.fields()[0]; - match field.data_type() { - DataType::List(child_field) => { - if !matches!(child_field.data_type(), DataType::Float64) { - return Err(MahoutError::InvalidInput(format!( - "Expected List column, got List<{:?}>", - child_field.data_type() - ))); - } - } - DataType::FixedSizeList(child_field, _) => { - if !matches!(child_field.data_type(), DataType::Float64) { - return Err(MahoutError::InvalidInput(format!( - "Expected FixedSizeList column, got FixedSizeList<{:?}>", - child_field.data_type() - ))); - } - } - DataType::Float64 => { - // Scalar Float64 for basis encoding (one index per sample) - } - _ => { - return Err(MahoutError::InvalidInput(format!( - "Expected Float64, List, or FixedSizeList column, got {:?}", - field.data_type() - ))); - } - } + validate_float_list_or_scalar_schema(&schema.fields()[0])?; let total_rows = builder.metadata().file_metadata().num_rows() as usize; @@ -344,6 +445,7 @@ impl ParquetStreamingReader { leftover_cursor: 0, total_rows, null_handling, + _phantom: PhantomData, }) } @@ -353,13 +455,13 @@ impl ParquetStreamingReader { } } -impl DataReader for ParquetStreamingReader { - fn read_batch(&mut self) -> Result<(Vec, usize, usize)> { +impl DataReader for ParquetStreamingReader { + fn read_batch(&mut self) -> Result<(Vec, usize, usize)> { let mut all_data = Vec::new(); let mut num_samples = 0; loop { - let mut buffer = vec![0.0; 1024 * 1024]; // 1M elements buffer + let mut buffer = vec![T::default(); 1024 * 1024]; let written = self.read_chunk(&mut buffer)?; if written == 0 { break; @@ -384,8 +486,8 @@ impl DataReader for ParquetStreamingReader { } } -impl StreamingDataReader for ParquetStreamingReader { - fn read_chunk(&mut self, buffer: &mut [f64]) -> Result { +impl StreamingDataReader for ParquetStreamingReader { + fn read_chunk(&mut self, buffer: &mut [T]) -> Result { let mut written = 0; let buf_cap = buffer.len(); let calc_limit = |ss: usize| -> usize { @@ -439,24 +541,16 @@ impl StreamingDataReader for ParquetStreamingReader { continue; } - let mut batch_values = Vec::new(); + let mut batch_values: Vec = Vec::new(); let mut current_sample_size = None; for i in 0..list_array.len() { let value_array = list_array.value(i); - let float_array = value_array - .as_any() - .downcast_ref::() - .ok_or_else(|| { - MahoutError::Io("List values must be Float64".to_string()) - })?; - if i == 0 { - current_sample_size = Some(float_array.len()); + current_sample_size = Some(value_array.len()); } - - handle_float64_nulls( + T::extend_from_arrow_array( &mut batch_values, - float_array, + &*value_array, self.null_handling, )?; } @@ -482,55 +576,25 @@ impl StreamingDataReader for ParquetStreamingReader { } let current_sample_size = *size as usize; - let values = list_array.values(); - let float_array = values - .as_any() - .downcast_ref::() - .ok_or_else(|| { - MahoutError::Io( - "FixedSizeList values must be Float64".to_string(), - ) - })?; - - let mut batch_values = Vec::new(); - handle_float64_nulls( - &mut batch_values, - float_array, - self.null_handling, - )?; + let batch_values = + T::collect_from_arrow_array(values, self.null_handling)?; (current_sample_size, batch_values) } - DataType::Float64 => { - // Scalar Float64 for basis encoding (one index per sample) - let float_array = column - .as_any() - .downcast_ref::() - .ok_or_else(|| { - MahoutError::Io( - "Failed to downcast to Float64Array".to_string(), - ) - })?; - - if float_array.is_empty() { + DataType::Float32 | DataType::Float64 => { + // Scalar float for basis encoding (one index per sample) + if column.is_empty() { continue; } - let current_sample_size = 1; - - let mut batch_values = Vec::new(); - handle_float64_nulls( - &mut batch_values, - float_array, - self.null_handling, - )?; - + let batch_values = + T::collect_from_arrow_array(&**column, self.null_handling)?; (current_sample_size, batch_values) } _ => { return Err(MahoutError::Io(format!( - "Expected Float64, List, or FixedSizeList, got {:?}", + "Expected Float32/Float64, List, or FixedSizeList, got {:?}", column.data_type() ))); } @@ -581,6 +645,11 @@ impl StreamingDataReader for ParquetStreamingReader { self.total_rows } } + +// --------------------------------------------------------------------------- +// Unit tests +// --------------------------------------------------------------------------- + #[cfg(test)] mod tests { use super::*; @@ -634,7 +703,7 @@ mod tests { #[test] fn test_parquet_reader_missing_file() { let path = std::env::temp_dir().join(format!("missing_{}.parquet", std::process::id())); - let result = ParquetReader::new(&path, None, NullHandling::FillZero); + let result = ParquetReader::::new(&path, None, NullHandling::FillZero); assert!(matches!(result, Err(MahoutError::Io(_)))); } @@ -644,13 +713,17 @@ mod tests { let array = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; let file = write_test_parquet(schema, vec![array]); - let result = ParquetReader::new(file.path(), None, NullHandling::FillZero); + let result = ParquetReader::::new(file.path(), None, NullHandling::FillZero); assert!(result.is_err()); let err_msg = match result { Err(e) => e.to_string(), Ok(_) => panic!(), }; - assert!(err_msg.contains("Expected List or FixedSizeList column")); + assert!( + err_msg.contains("Expected List") + || err_msg.contains("Expected FixedSizeList") + || err_msg.contains("Float32/Float64") + ); } #[test] @@ -663,7 +736,7 @@ mod tests { let arr2 = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; let file = write_test_parquet(schema, vec![arr1, arr2]); - let result = ParquetReader::new(file.path(), None, NullHandling::FillZero); + let result = ParquetReader::::new(file.path(), None, NullHandling::FillZero); assert!(result.is_err()); let err_msg = match result { Err(e) => e.to_string(), @@ -685,13 +758,13 @@ mod tests { let file = write_test_parquet(schema, vec![array]); - let result = ParquetReader::new(file.path(), None, NullHandling::FillZero); + let result = ParquetReader::::new(file.path(), None, NullHandling::FillZero); assert!(result.is_err()); let err_msg = match result { Err(e) => e.to_string(), Ok(_) => panic!(), }; - assert!(err_msg.contains("Expected List column")); + assert!(err_msg.contains("Expected List or List")); } #[test] @@ -709,7 +782,8 @@ mod tests { let file = write_test_parquet(schema, vec![array]); - let mut reader = ParquetReader::new(file.path(), None, NullHandling::FillZero).unwrap(); + let mut reader = + ParquetReader::::new(file.path(), None, NullHandling::FillZero).unwrap(); let (data, num_samples, sample_size) = reader.read_batch().unwrap(); assert_eq!(data, vec![1.0, 2.0, 3.0, 4.0]); assert_eq!(num_samples, 2); @@ -733,7 +807,8 @@ mod tests { let file = write_test_parquet(schema, vec![array]); - let mut reader = ParquetReader::new(file.path(), None, NullHandling::FillZero).unwrap(); + let mut reader = + ParquetReader::::new(file.path(), None, NullHandling::FillZero).unwrap(); let (data, num_samples, sample_size) = reader.read_batch().unwrap(); assert_eq!(data, vec![5.0, 6.0, 7.0, 8.0]); assert_eq!(num_samples, 2); @@ -755,7 +830,8 @@ mod tests { let file = write_test_parquet(schema, vec![array]); - let mut reader = ParquetReader::new(file.path(), None, NullHandling::FillZero).unwrap(); + let mut reader = + ParquetReader::::new(file.path(), None, NullHandling::FillZero).unwrap(); let result = reader.read_batch(); assert!(result.is_err()); let err_msg = match result { @@ -767,6 +843,7 @@ mod tests { #[test] fn test_parquet_streaming_reader_scalar_f64() { + use arrow::array::Float64Builder; let schema = Arc::new(Schema::new(vec![Field::new( "data", DataType::Float64, @@ -778,7 +855,7 @@ mod tests { let file = write_test_parquet(schema, vec![array]); let mut streaming_reader = - ParquetStreamingReader::new(file.path(), None, NullHandling::FillZero).unwrap(); + ParquetStreamingReader::::new(file.path(), None, NullHandling::FillZero).unwrap(); assert_eq!(streaming_reader.total_rows(), 3); let mut buffer = vec![0.0; 2]; let written1 = streaming_reader.read_chunk(&mut buffer).unwrap(); @@ -812,7 +889,8 @@ mod tests { let file = write_test_parquet(schema, vec![array]); let mut reader = - ParquetStreamingReader::new(file.path(), Some(1), NullHandling::FillZero).unwrap(); + ParquetStreamingReader::::new(file.path(), Some(1), NullHandling::FillZero) + .unwrap(); let (data, num_samples, sample_size) = reader.read_batch().unwrap(); assert_eq!(data, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); assert_eq!(num_samples, 3); @@ -829,7 +907,8 @@ mod tests { let array = Arc::new(builder.finish()) as ArrayRef; let file = write_test_parquet(schema, vec![array]); - let mut reader = ParquetReader::new(file.path(), None, NullHandling::FillZero).unwrap(); + let mut reader = + ParquetReader::::new(file.path(), None, NullHandling::FillZero).unwrap(); let result = reader.read_batch(); assert!(result.is_err()); let err_msg = match result { @@ -850,7 +929,7 @@ mod tests { let file = write_test_parquet(schema, vec![array]); let mut reader = - ParquetStreamingReader::new(file.path(), None, NullHandling::FillZero).unwrap(); + ParquetStreamingReader::::new(file.path(), None, NullHandling::FillZero).unwrap(); let result = reader.read_batch(); assert!(result.is_err()); let err_msg = match result { @@ -899,7 +978,8 @@ mod tests { #[test] fn test_parquet_reader_list_f64_null_fill_zero() { let file = write_list_parquet_with_nulls(); - let mut reader = ParquetReader::new(file.path(), None, NullHandling::FillZero).unwrap(); + let mut reader = + ParquetReader::::new(file.path(), None, NullHandling::FillZero).unwrap(); let (data, num_samples, sample_size) = reader.read_batch().unwrap(); assert_eq!(data, vec![1.0, 0.0, 3.0, 4.0]); assert_eq!(num_samples, 2); @@ -909,7 +989,8 @@ mod tests { #[test] fn test_parquet_reader_list_f64_null_reject() { let file = write_list_parquet_with_nulls(); - let mut reader = ParquetReader::new(file.path(), None, NullHandling::Reject).unwrap(); + let mut reader = + ParquetReader::::new(file.path(), None, NullHandling::Reject).unwrap(); let result = reader.read_batch(); assert!(result.is_err()); let err_msg = match result { @@ -922,7 +1003,8 @@ mod tests { #[test] fn test_parquet_reader_fixed_size_list_null_fill_zero() { let file = write_fixed_size_list_parquet_with_nulls(); - let mut reader = ParquetReader::new(file.path(), None, NullHandling::FillZero).unwrap(); + let mut reader = + ParquetReader::::new(file.path(), None, NullHandling::FillZero).unwrap(); let (data, num_samples, sample_size) = reader.read_batch().unwrap(); assert_eq!(data, vec![1.0, 0.0, 3.0, 4.0]); assert_eq!(num_samples, 2); @@ -932,7 +1014,8 @@ mod tests { #[test] fn test_parquet_reader_fixed_size_list_null_reject() { let file = write_fixed_size_list_parquet_with_nulls(); - let mut reader = ParquetReader::new(file.path(), None, NullHandling::Reject).unwrap(); + let mut reader = + ParquetReader::::new(file.path(), None, NullHandling::Reject).unwrap(); let result = reader.read_batch(); assert!(result.is_err()); let err_msg = match result { @@ -949,7 +1032,7 @@ mod tests { let file = TempTestFile::new(); let path = file.path().to_path_buf(); drop(file); - let result = ParquetStreamingReader::new(&path, None, NullHandling::FillZero); + let result = ParquetStreamingReader::::new(&path, None, NullHandling::FillZero); assert!(matches!(result, Err(MahoutError::Io(_)))); } @@ -959,7 +1042,7 @@ mod tests { let array = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; let file = write_test_parquet(schema, vec![array]); - let result = ParquetStreamingReader::new(file.path(), None, NullHandling::FillZero); + let result = ParquetStreamingReader::::new(file.path(), None, NullHandling::FillZero); assert!(result.is_err()); let err_msg = match result { Err(e) => e.to_string(), @@ -978,7 +1061,7 @@ mod tests { let arr2 = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; let file = write_test_parquet(schema, vec![arr1, arr2]); - let result = ParquetStreamingReader::new(file.path(), None, NullHandling::FillZero); + let result = ParquetStreamingReader::::new(file.path(), None, NullHandling::FillZero); assert!(result.is_err()); let err_msg = match result { Err(e) => e.to_string(), @@ -1000,12 +1083,12 @@ mod tests { let file = write_test_parquet(schema, vec![array]); - let result = ParquetStreamingReader::new(file.path(), None, NullHandling::FillZero); + let result = ParquetStreamingReader::::new(file.path(), None, NullHandling::FillZero); assert!(result.is_err()); let err_msg = match result { Err(e) => e.to_string(), Ok(_) => panic!(), }; - assert!(err_msg.contains("Expected List")); + assert!(err_msg.contains("Expected List or List")); } } diff --git a/qdp/qdp-core/src/remote.rs b/qdp/qdp-core/src/remote.rs index a2577eb6c0..c800fe52f8 100644 --- a/qdp/qdp-core/src/remote.rs +++ b/qdp/qdp-core/src/remote.rs @@ -256,7 +256,7 @@ mod tests { // Verify it's a valid parquet that our reader can parse. use crate::reader::DataReader; - let mut reader = crate::readers::ParquetReader::new( + let mut reader = crate::readers::ParquetReader::::new( &resolved.path, None, crate::reader::NullHandling::FillZero, diff --git a/qdp/qdp-core/tests/parquet_f32.rs b/qdp/qdp-core/tests/parquet_f32.rs new file mode 100644 index 0000000000..9925bc45b9 --- /dev/null +++ b/qdp/qdp-core/tests/parquet_f32.rs @@ -0,0 +1,161 @@ +// +// 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. + +//! Acceptance tests for ParquetReader — issue #1340. + +use std::fs; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use arrow::array::{ArrayRef, Float32Builder, Float64Builder, ListBuilder, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use parquet::arrow::ArrowWriter; +use qdp_core::reader::{DataReader, NullHandling}; +use qdp_core::readers::parquet::ParquetReader; + +static FILE_COUNTER: AtomicUsize = AtomicUsize::new(0); + +struct TempFile(std::path::PathBuf); + +impl TempFile { + fn path(&self) -> &std::path::Path { + &self.0 + } +} + +impl Drop for TempFile { + fn drop(&mut self) { + let _ = fs::remove_file(&self.0); + } +} + +fn write_list_parquet(schema: Arc, arrays: Vec) -> TempFile { + let n = FILE_COUNTER.fetch_add(1, Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "mahout_parquet_f32_{}_{}.parquet", + std::process::id(), + n, + )); + let batch = RecordBatch::try_new(schema.clone(), arrays).unwrap(); + let file = fs::File::create(&path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + TempFile(path) +} + +// --------------------------------------------------------------------------- +// Acceptance test 1: f32 column read as f32 (zero-copy path) +// --------------------------------------------------------------------------- + +/// ParquetReader:: on a List file → values come back as Vec, +/// no precision loss, correct count. +#[test] +fn test_f32_column_read_as_f32() { + let item_field = Arc::new(Field::new("item", DataType::Float32, true)); + let list_field = Field::new("data", DataType::List(item_field), true); + let schema = Arc::new(Schema::new(vec![list_field])); + + let mut builder = ListBuilder::new(Float32Builder::new()); + builder.values().append_slice(&[1.0_f32, 2.5_f32, 3.75_f32]); + builder.append(true); + builder.values().append_slice(&[4.0_f32, 5.5_f32, 6.25_f32]); + builder.append(true); + let array = Arc::new(builder.finish()) as ArrayRef; + + let tmp = write_list_parquet(schema, vec![array]); + + let mut reader = ParquetReader::::new(tmp.path(), None, NullHandling::FillZero).unwrap(); + let (data, num_samples, sample_size) = reader.read_batch().unwrap(); + + assert_eq!(num_samples, 2); + assert_eq!(sample_size, 3); + assert_eq!(data, vec![1.0_f32, 2.5, 3.75, 4.0, 5.5, 6.25]); +} + +// --------------------------------------------------------------------------- +// Acceptance test 2: f64 column cast to f32 +// --------------------------------------------------------------------------- + +/// ParquetReader:: on a List file → Arrow cast applied. +/// Normal values: cast is precise within f32 range. +/// Overflow (f64 > f32::MAX): → +Inf (Arrow safe cast behaviour). +/// NaN: preserved. +#[test] +fn test_f64_column_cast_to_f32() { + let item_field = Arc::new(Field::new("item", DataType::Float64, true)); + let list_field = Field::new("data", DataType::List(item_field), true); + let schema = Arc::new(Schema::new(vec![list_field])); + + let overflow = f64::from(f32::MAX) * 2.0; // overflows f32 → +Inf after cast + let nan = f64::NAN; + + let mut builder = ListBuilder::new(Float64Builder::new()); + builder.values().append_slice(&[1.0_f64, -2.0_f64]); + builder.append(true); + builder.values().append_value(overflow); + builder.values().append_value(nan); + builder.append(true); + let array = Arc::new(builder.finish()) as ArrayRef; + + let tmp = write_list_parquet(schema, vec![array]); + + let mut reader = ParquetReader::::new(tmp.path(), None, NullHandling::FillZero).unwrap(); + let (data, num_samples, sample_size) = reader.read_batch().unwrap(); + + assert_eq!(num_samples, 2); + assert_eq!(sample_size, 2); + assert_eq!(data[0], 1.0_f32); + assert_eq!(data[1], -2.0_f32); + assert!( + data[2].is_infinite() && data[2] > 0.0, + "expected +Inf, got {}", + data[2] + ); + assert!(data[3].is_nan(), "expected NaN, got {}", data[3]); +} + +// --------------------------------------------------------------------------- +// Acceptance test 3: unsupported column type → InvalidInput with dtype in message +// --------------------------------------------------------------------------- + +/// ParquetReader on a List file must fail at construction with an +/// InvalidInput error whose message mentions the actual column dtype. +#[test] +fn test_unsupported_column_type_returns_error_with_dtype() { + let item_field = Arc::new(Field::new("item", DataType::Int32, true)); + let list_field = Field::new("data", DataType::List(item_field), true); + let schema = Arc::new(Schema::new(vec![list_field])); + + let mut builder = arrow::array::ListBuilder::new(arrow::array::Int32Builder::new()); + builder.values().append_value(42); + builder.append(true); + let array = Arc::new(builder.finish()) as ArrayRef; + + let tmp = write_list_parquet(schema, vec![array]); + + let result_f32 = ParquetReader::::new(tmp.path(), None, NullHandling::FillZero); + let result_f64 = ParquetReader::::new(tmp.path(), None, NullHandling::FillZero); + + for result in [result_f32.map(|_| ()), result_f64.map(|_| ())] { + assert!(result.is_err(), "expected error for Int32 column"); + let msg = result.unwrap_err().to_string(); + assert!( + msg.contains("Int32") || msg.contains("int32"), + "error message should contain the dtype, got: {msg}" + ); + } +} From 2a581605956e3092b629585d1a7a946490b540d9 Mon Sep 17 00:00:00 2001 From: 0lai0 Date: Thu, 11 Jun 2026 00:53:57 +0800 Subject: [PATCH 2/2] address comment --- qdp/qdp-core/src/reader.rs | 92 +++++--- qdp/qdp-core/src/readers/parquet.rs | 354 +++++++++++++++------------- qdp/qdp-core/tests/parquet_f32.rs | 140 ++++++++++- 3 files changed, 383 insertions(+), 203 deletions(-) diff --git a/qdp/qdp-core/src/reader.rs b/qdp/qdp-core/src/reader.rs index 11c430341f..46db0a6c1b 100644 --- a/qdp/qdp-core/src/reader.rs +++ b/qdp/qdp-core/src/reader.rs @@ -45,20 +45,39 @@ //! } //! ``` -use arrow::array::{Array, Float32Array, Float64Array}; +use arrow::array::{Array, Float32Array, Float64Array, PrimitiveArray}; +use arrow::datatypes::{ArrowPrimitiveType, Float32Type, Float64Type}; -use crate::error::Result; +use crate::error::{MahoutError, Result}; + +/// Maps a Rust float primitive to its Arrow array type. +/// +/// `pub(crate)` seals `FloatElem`: external callers cannot implement `ArrowPrimitive` +/// and therefore cannot implement `FloatElem` for new types. +pub(crate) trait ArrowPrimitive { + type ArrowType: ArrowPrimitiveType; +} + +impl ArrowPrimitive for f32 { + type ArrowType = Float32Type; +} + +impl ArrowPrimitive for f64 { + type ArrowType = Float64Type; +} /// Scalar element type for [`DataReader`] output (`f32` or `f64` only). /// -/// Keeps f32 file data as `Vec` end-to-end once readers implement -/// `DataReader`; today most readers use the default `T = f64`. -pub trait FloatElem: Copy + Default + Send + Sync + 'static {} +/// Sealed by the `pub(crate) ArrowPrimitive` supertrait — no external implementations +/// are possible. Keeps f32 file data as `Vec` end-to-end; today most readers +/// use the default `T = f64`. +#[allow(private_bounds)] +pub trait FloatElem: ArrowPrimitive + Copy + Default + Send + Sync + 'static {} impl FloatElem for f32 {} impl FloatElem for f64 {} -/// Policy for handling null values in Float64 arrays. +/// Policy for handling null values in float arrays. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub enum NullHandling { /// Replace nulls with 0.0 (backward-compatible default). @@ -68,34 +87,48 @@ pub enum NullHandling { Reject, } -/// Append values from a `Float64Array` into `output`, applying the given null policy. +/// Append values from a primitive array into `output`, applying the given null policy. /// /// When there are no nulls the fast path copies the underlying buffer directly. -pub fn handle_float64_nulls( - output: &mut Vec, - float_array: &Float64Array, +pub(crate) fn handle_primitive_nulls( + output: &mut Vec, + array: &PrimitiveArray

, null_handling: NullHandling, -) -> crate::error::Result<()> { - if float_array.null_count() == 0 { - output.extend_from_slice(float_array.values()); +) -> Result<()> +where + P::Native: Default, +{ + if array.null_count() == 0 { + output.extend_from_slice(array.values()); } else { match null_handling { NullHandling::FillZero => { - output.extend(float_array.iter().map(|opt| opt.unwrap_or(0.0))); + output.extend(array.iter().map(|opt| opt.unwrap_or_default())); } NullHandling::Reject => { - return Err(crate::error::MahoutError::InvalidInput( - "Null value encountered in Float64Array. \ + return Err(MahoutError::InvalidInput(format!( + "Null value encountered in {:?} array. \ Use NullHandling::FillZero to replace nulls with 0.0, \ - or clean the data at the source." - .to_string(), - )); + or clean the data at the source.", + P::DATA_TYPE, + ))); } } } Ok(()) } +/// Append values from a `Float64Array` into `output`, applying the given null policy. +/// +/// When there are no nulls the fast path copies the underlying buffer directly. +pub fn handle_float64_nulls( + output: &mut Vec, + float_array: &Float64Array, + null_handling: NullHandling, +) -> Result<()> { + handle_primitive_nulls::(output, float_array, null_handling) +} + /// Append values from a `Float32Array` into `output`, applying the given null policy. /// /// When there are no nulls the fast path copies the underlying buffer directly. @@ -103,25 +136,8 @@ pub fn handle_float32_nulls( output: &mut Vec, float_array: &Float32Array, null_handling: NullHandling, -) -> crate::error::Result<()> { - if float_array.null_count() == 0 { - output.extend_from_slice(float_array.values()); - } else { - match null_handling { - NullHandling::FillZero => { - output.extend(float_array.iter().map(|opt| opt.unwrap_or(0.0))); - } - NullHandling::Reject => { - return Err(crate::error::MahoutError::InvalidInput( - "Null value encountered in Float32Array. \ - Use NullHandling::FillZero to replace nulls with 0.0, \ - or clean the data at the source." - .to_string(), - )); - } - } - } - Ok(()) +) -> Result<()> { + handle_primitive_nulls::(output, float_array, null_handling) } /// Generic data reader interface for batch quantum data. diff --git a/qdp/qdp-core/src/readers/parquet.rs b/qdp/qdp-core/src/readers/parquet.rs index 475b88d1e7..30f2deead7 100644 --- a/qdp/qdp-core/src/readers/parquet.rs +++ b/qdp/qdp-core/src/readers/parquet.rs @@ -20,123 +20,29 @@ use std::fs::File; use std::marker::PhantomData; use std::path::Path; -use arrow::array::{Array, FixedSizeListArray, Float32Array, Float64Array, ListArray}; +use arrow::array::{Array, ArrayRef, FixedSizeListArray, ListArray, PrimitiveArray}; use arrow::compute; -use arrow::datatypes::DataType; +use arrow::datatypes::{ArrowPrimitiveType, DataType}; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use crate::error::{MahoutError, Result}; use crate::reader::{ - DataReader, FloatElem, NullHandling, StreamingDataReader, handle_float32_nulls, - handle_float64_nulls, + ArrowPrimitive, DataReader, FloatElem, NullHandling, StreamingDataReader, + handle_primitive_nulls, }; // --------------------------------------------------------------------------- -// Internal sealed helper trait +// Module-level helpers // --------------------------------------------------------------------------- -/// Zero-copy or Arrow-cast extraction from an Arrow array into an output `Vec`. -/// -/// Implemented for `f32` and `f64` only: -/// - same dtype → `extend_from_slice` directly from the Arrow buffer (zero alloc) -/// - cross dtype → `arrow::compute::cast` first, then `extend_from_slice` -/// - f32 → f64: exact, NaN preserved -/// - f64 → f32: values outside f32 range become ±Inf, NaN preserved -pub(crate) trait ArrowPrimitive: FloatElem { - fn extend_from_arrow_array( - output: &mut Vec, - array: &dyn Array, - null_handling: NullHandling, - ) -> Result<()>; - - fn collect_from_arrow_array( - array: &dyn Array, - null_handling: NullHandling, - ) -> Result> { - let mut out = Vec::new(); - Self::extend_from_arrow_array(&mut out, array, null_handling)?; - Ok(out) - } -} - -impl ArrowPrimitive for f64 { - fn extend_from_arrow_array( - output: &mut Vec, - array: &dyn Array, - null_handling: NullHandling, - ) -> Result<()> { - match array.data_type() { - DataType::Float64 => { - let arr = array - .as_any() - .downcast_ref::() - .ok_or_else(|| MahoutError::Io("Float64 downcast failed".to_string()))?; - handle_float64_nulls(output, arr, null_handling)?; - } - DataType::Float32 => { - let casted = compute::cast(array, &DataType::Float64) - .map_err(|e| MahoutError::Io(format!("Arrow cast f32→f64: {e}")))?; - let arr = casted - .as_any() - .downcast_ref::() - .ok_or_else(|| MahoutError::Io("Cast to Float64 failed".to_string()))?; - handle_float64_nulls(output, arr, null_handling)?; - } - other => { - return Err(MahoutError::InvalidInput(format!( - "Expected Float32 or Float64 values, got {other:?}" - ))); - } - } - Ok(()) - } -} - -impl ArrowPrimitive for f32 { - fn extend_from_arrow_array( - output: &mut Vec, - array: &dyn Array, - null_handling: NullHandling, - ) -> Result<()> { - match array.data_type() { - DataType::Float32 => { - let arr = array - .as_any() - .downcast_ref::() - .ok_or_else(|| MahoutError::Io("Float32 downcast failed".to_string()))?; - handle_float32_nulls(output, arr, null_handling)?; - } - DataType::Float64 => { - // f64 → f32: values outside f32 range become ±Inf; NaN is preserved - let casted = compute::cast(array, &DataType::Float32) - .map_err(|e| MahoutError::Io(format!("Arrow cast f64→f32: {e}")))?; - let arr = casted - .as_any() - .downcast_ref::() - .ok_or_else(|| MahoutError::Io("Cast to Float32 failed".to_string()))?; - handle_float32_nulls(output, arr, null_handling)?; - } - other => { - return Err(MahoutError::InvalidInput(format!( - "Expected Float32 or Float64 values, got {other:?}" - ))); - } - } - Ok(()) - } +fn is_supported_float(dt: &DataType) -> bool { + matches!(dt, DataType::Float32 | DataType::Float64) } -// --------------------------------------------------------------------------- -// Shared schema validator -// --------------------------------------------------------------------------- - fn validate_float_list_schema(field: &arrow::datatypes::Field) -> Result<()> { match field.data_type() { DataType::List(child_field) => { - if !matches!( - child_field.data_type(), - DataType::Float32 | DataType::Float64 - ) { + if !is_supported_float(child_field.data_type()) { return Err(MahoutError::InvalidInput(format!( "Expected List or List column, got List<{:?}>", child_field.data_type() @@ -144,19 +50,18 @@ fn validate_float_list_schema(field: &arrow::datatypes::Field) -> Result<()> { } } DataType::FixedSizeList(child_field, _) => { - if !matches!( - child_field.data_type(), - DataType::Float32 | DataType::Float64 - ) { + if !is_supported_float(child_field.data_type()) { return Err(MahoutError::InvalidInput(format!( - "Expected FixedSizeList or FixedSizeList column, got FixedSizeList<{:?}>", + "Expected FixedSizeList or FixedSizeList column, \ + got FixedSizeList<{:?}>", child_field.data_type() ))); } } _ => { return Err(MahoutError::InvalidInput(format!( - "Expected List or FixedSizeList column, got {:?}", + "Expected List or FixedSizeList column, \ + got {:?}", field.data_type() ))); } @@ -171,6 +76,92 @@ fn validate_float_list_or_scalar_schema(field: &arrow::datatypes::Field) -> Resu } } +/// Returns the element DataType from a List, FixedSizeList, or scalar float field. +fn element_dtype(field: &arrow::datatypes::Field) -> Option { + match field.data_type() { + DataType::List(child) | DataType::FixedSizeList(child, _) => { + Some(child.data_type().clone()) + } + dt if is_supported_float(dt) => Some(dt.clone()), + _ => None, + } +} + +/// Extracts the offset-adjusted flat values slice from a `ListArray`. +/// +/// `ListArray::values()` returns the full backing child array; for a sliced +/// `ListArray` the first valid element starts at `offsets[0]`, not index 0. +/// Omitting this adjustment would read stale data outside the array's range. +fn list_flat_values(arr: &ListArray) -> ArrayRef { + let offsets = arr.offsets(); + let start = offsets[0] as usize; + let end = offsets[arr.len()] as usize; + arr.values().slice(start, end - start) +} + +/// Extracts the offset-adjusted flat values slice from a `FixedSizeListArray`. +/// +/// `FixedSizeListArray::values()` returns the full backing child array; for a sliced +/// array the valid range starts at `offset * value_size`, not at index 0. +fn fixed_size_list_flat_values(arr: &FixedSizeListArray) -> ArrayRef { + let size = arr.value_length() as usize; + let start = arr.offset() * size; + let end = (arr.offset() + arr.len()) * size; + arr.values().slice(start, end - start) +} + +/// Cast `array` to `P::DATA_TYPE` if needed, then append all values to `output`. +/// +/// Same dtype → zero-copy extend from the Arrow buffer. +/// Cross dtype (f64→f32) → `arrow::compute::cast` once, then extend. +/// - f64→f32: values outside f32 range become ±Inf; NaN preserved. +fn extend_floats( + output: &mut Vec, + array: &dyn Array, + null_handling: NullHandling, +) -> Result<()> +where + P::Native: Default, +{ + let target_dt = P::DATA_TYPE; + let casted; + let effective: &dyn Array = if array.data_type() == &target_dt { + array + } else if is_supported_float(array.data_type()) { + casted = compute::cast(array, &target_dt).map_err(|e| { + MahoutError::InvalidInput(format!( + "Arrow cast {:?}→{:?}: {e}", + array.data_type(), + target_dt + )) + })?; + &*casted + } else { + return Err(MahoutError::InvalidInput(format!( + "Expected Float32 or Float64 values, got {:?}", + array.data_type() + ))); + }; + + let arr = effective + .as_any() + .downcast_ref::>() + .ok_or_else(|| MahoutError::InvalidInput(format!("{:?} downcast failed", target_dt)))?; + handle_primitive_nulls::

(output, arr, null_handling) +} + +fn collect_floats( + array: &dyn Array, + null_handling: NullHandling, +) -> Result> +where + P::Native: Default, +{ + let mut out = Vec::new(); + extend_floats::

(&mut out, array, null_handling)?; + Ok(out) +} + // --------------------------------------------------------------------------- // ParquetReader // --------------------------------------------------------------------------- @@ -189,8 +180,7 @@ pub struct ParquetReader { _phantom: PhantomData, } -#[allow(private_bounds)] -impl ParquetReader { +impl ParquetReader { /// Create a new Parquet reader. /// /// # Arguments @@ -242,6 +232,17 @@ impl ParquetReader { validate_float_list_schema(&schema.fields()[0])?; + // Warn on f64→f32 narrowing cast: overflow becomes ±Inf with no error. + if let Some(file_dt) = element_dtype(&schema.fields()[0]) { + let target_dt = <::ArrowType as ArrowPrimitiveType>::DATA_TYPE; + if file_dt == DataType::Float64 && target_dt == DataType::Float32 { + log::warn!( + "Parquet column is Float64 but reading as f32: values outside f32 range \ + become ±Inf. Use ParquetReader:: to preserve precision." + ); + } + } + let total_rows = builder.metadata().file_metadata().num_rows() as usize; let reader = if let Some(batch_size) = batch_size { @@ -261,7 +262,7 @@ impl ParquetReader { } } -impl DataReader for ParquetReader { +impl DataReader for ParquetReader { fn read_batch(&mut self) -> Result<(Vec, usize, usize)> { let reader = self .reader @@ -289,29 +290,31 @@ impl DataReader for ParquetReader { MahoutError::Io("Failed to downcast to ListArray".to_string()) })?; + // Validate all rows have a consistent sample size. for i in 0..list_array.len() { - let value_array = list_array.value(i); - let current_size = value_array.len(); - - if let Some(expected_size) = sample_size { - if current_size != expected_size { + let row_len = list_array.value_length(i) as usize; + if let Some(expected) = sample_size { + if row_len != expected { return Err(MahoutError::InvalidInput(format!( "Inconsistent sample sizes: expected {}, got {}", - expected_size, current_size + expected, row_len ))); } } else { - sample_size = Some(current_size); - all_data.reserve(current_size * self.total_rows); + sample_size = Some(row_len); + all_data.reserve(row_len * self.total_rows); } - - T::extend_from_arrow_array( - &mut all_data, - &*value_array, - self.null_handling, - )?; - num_samples += 1; } + + // Cast the entire flat buffer once (avoids N per-row allocations + // on cross-dtype reads) then extend all_data in one pass. + let flat = list_flat_values(list_array); + extend_floats::<::ArrowType>( + &mut all_data, + &*flat, + self.null_handling, + )?; + num_samples += list_array.len(); } DataType::FixedSizeList(_, size) => { let list_array = column @@ -328,13 +331,18 @@ impl DataReader for ParquetReader { all_data.reserve(current_size * batch.num_rows()); } - let values = list_array.values(); - T::extend_from_arrow_array(&mut all_data, values, self.null_handling)?; + let flat = fixed_size_list_flat_values(list_array); + extend_floats::<::ArrowType>( + &mut all_data, + &*flat, + self.null_handling, + )?; num_samples += list_array.len(); } _ => { - return Err(MahoutError::Io(format!( - "Expected List or FixedSizeList, got {:?}", + return Err(MahoutError::InvalidInput(format!( + "Expected List or FixedSizeList, \ + got {:?}", column.data_type() ))); } @@ -377,8 +385,7 @@ pub struct ParquetStreamingReader { _phantom: PhantomData, } -#[allow(private_bounds)] -impl ParquetStreamingReader { +impl ParquetStreamingReader { /// Create a new streaming Parquet reader. /// /// # Arguments @@ -430,6 +437,18 @@ impl ParquetStreamingReader { validate_float_list_or_scalar_schema(&schema.fields()[0])?; + // Warn on f64→f32 narrowing cast: overflow becomes ±Inf with no error. + if let Some(file_dt) = element_dtype(&schema.fields()[0]) { + let target_dt = <::ArrowType as ArrowPrimitiveType>::DATA_TYPE; + if file_dt == DataType::Float64 && target_dt == DataType::Float32 { + log::warn!( + "ParquetStreamingReader: Float64 column cast to f32 — values outside \ + f32 range become ±Inf. Use ParquetStreamingReader:: to preserve \ + precision." + ); + } + } + let total_rows = builder.metadata().file_metadata().num_rows() as usize; let batch_size = batch_size.unwrap_or(2048); @@ -455,13 +474,14 @@ impl ParquetStreamingReader { } } -impl DataReader for ParquetStreamingReader { +impl DataReader for ParquetStreamingReader { fn read_batch(&mut self) -> Result<(Vec, usize, usize)> { let mut all_data = Vec::new(); let mut num_samples = 0; + // Hoist buffer out of the loop to avoid re-allocating 1M elements per iteration. + let mut buffer = vec![T::default(); 1024 * 1024]; loop { - let mut buffer = vec![T::default(); 1024 * 1024]; let written = self.read_chunk(&mut buffer)?; if written == 0 { break; @@ -486,7 +506,7 @@ impl DataReader for ParquetStreamingReader } } -impl StreamingDataReader for ParquetStreamingReader { +impl StreamingDataReader for ParquetStreamingReader { fn read_chunk(&mut self, buffer: &mut [T]) -> Result { let mut written = 0; let buf_cap = buffer.len(); @@ -537,29 +557,32 @@ impl StreamingDataReader for ParquetStreamingR MahoutError::Io("Failed to downcast to ListArray".to_string()) })?; - if list_array.len() == 0 { + if list_array.is_empty() { continue; } - let mut batch_values: Vec = Vec::new(); - let mut current_sample_size = None; - for i in 0..list_array.len() { - let value_array = list_array.value(i); - if i == 0 { - current_sample_size = Some(value_array.len()); + let current_sample_size = list_array.value_length(0) as usize; + + // Validate all rows in this batch have a consistent sample size. + for i in 1..list_array.len() { + let row_len = list_array.value_length(i) as usize; + if row_len != current_sample_size { + return Err(MahoutError::InvalidInput(format!( + "Inconsistent sample sizes: expected {}, got {}", + current_sample_size, row_len + ))); } - T::extend_from_arrow_array( - &mut batch_values, - &*value_array, - self.null_handling, - )?; } - ( - current_sample_size - .expect("list_array.len() > 0 ensures at least one element"), - batch_values, - ) + // Cast the entire flat buffer once (avoids N per-row allocations + // on cross-dtype reads). + let flat = list_flat_values(list_array); + let batch_values = collect_floats::<::ArrowType>( + &*flat, + self.null_handling, + )?; + + (current_sample_size, batch_values) } DataType::FixedSizeList(_, size) => { let list_array = column @@ -571,14 +594,16 @@ impl StreamingDataReader for ParquetStreamingR ) })?; - if list_array.len() == 0 { + if list_array.is_empty() { continue; } let current_sample_size = *size as usize; - let values = list_array.values(); - let batch_values = - T::collect_from_arrow_array(values, self.null_handling)?; + let flat = fixed_size_list_flat_values(list_array); + let batch_values = collect_floats::<::ArrowType>( + &*flat, + self.null_handling, + )?; (current_sample_size, batch_values) } @@ -588,13 +613,16 @@ impl StreamingDataReader for ParquetStreamingR continue; } let current_sample_size = 1; - let batch_values = - T::collect_from_arrow_array(&**column, self.null_handling)?; + let batch_values = collect_floats::<::ArrowType>( + &**column, + self.null_handling, + )?; (current_sample_size, batch_values) } _ => { - return Err(MahoutError::Io(format!( - "Expected Float32/Float64, List, or FixedSizeList, got {:?}", + return Err(MahoutError::InvalidInput(format!( + "Expected Float32/Float64, List, or \ + FixedSizeList, got {:?}", column.data_type() ))); } @@ -719,10 +747,10 @@ mod tests { Err(e) => e.to_string(), Ok(_) => panic!(), }; + // The error must mention the actual dtype, not just a generic "Expected" substring. assert!( - err_msg.contains("Expected List") - || err_msg.contains("Expected FixedSizeList") - || err_msg.contains("Float32/Float64") + err_msg.contains("Int32"), + "error message should contain the column dtype, got: {err_msg}" ); } diff --git a/qdp/qdp-core/tests/parquet_f32.rs b/qdp/qdp-core/tests/parquet_f32.rs index 9925bc45b9..ecfd293dbc 100644 --- a/qdp/qdp-core/tests/parquet_f32.rs +++ b/qdp/qdp-core/tests/parquet_f32.rs @@ -20,11 +20,13 @@ use std::fs; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use arrow::array::{ArrayRef, Float32Builder, Float64Builder, ListBuilder, RecordBatch}; +use arrow::array::{ + ArrayRef, FixedSizeListBuilder, Float32Builder, Float64Builder, ListBuilder, RecordBatch, +}; use arrow::datatypes::{DataType, Field, Schema}; use parquet::arrow::ArrowWriter; use qdp_core::reader::{DataReader, NullHandling}; -use qdp_core::readers::parquet::ParquetReader; +use qdp_core::readers::parquet::{ParquetReader, ParquetStreamingReader}; static FILE_COUNTER: AtomicUsize = AtomicUsize::new(0); @@ -159,3 +161,137 @@ fn test_unsupported_column_type_returns_error_with_dtype() { ); } } + +// --------------------------------------------------------------------------- +// Acceptance test 4: FixedSizeList read as f32 (zero-copy path) +// --------------------------------------------------------------------------- + +/// ParquetReader:: on a FixedSizeList file → values come back as +/// Vec with no precision loss. +#[test] +fn test_fixed_size_list_f32_as_f32() { + let item_field = Arc::new(Field::new("item", DataType::Float32, true)); + let list_field = Field::new("data", DataType::FixedSizeList(item_field, 3), true); + let schema = Arc::new(Schema::new(vec![list_field])); + + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 3); + builder.values().append_slice(&[1.0_f32, 2.0_f32, 3.0_f32]); + builder.append(true); + builder.values().append_slice(&[4.0_f32, 5.0_f32, 6.0_f32]); + builder.append(true); + let array = Arc::new(builder.finish()) as ArrayRef; + + let tmp = write_list_parquet(schema, vec![array]); + + let mut reader = ParquetReader::::new(tmp.path(), None, NullHandling::FillZero).unwrap(); + let (data, num_samples, sample_size) = reader.read_batch().unwrap(); + + assert_eq!(num_samples, 2); + assert_eq!(sample_size, 3); + assert_eq!(data, vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]); +} + +// --------------------------------------------------------------------------- +// Acceptance test 5: FixedSizeList cast to f32 +// --------------------------------------------------------------------------- + +/// ParquetReader:: on a FixedSizeList file → Arrow cast applied. +#[test] +fn test_fixed_size_list_f64_cast_to_f32() { + let item_field = Arc::new(Field::new("item", DataType::Float64, true)); + let list_field = Field::new("data", DataType::FixedSizeList(item_field, 2), true); + let schema = Arc::new(Schema::new(vec![list_field])); + + let mut builder = FixedSizeListBuilder::new(Float64Builder::new(), 2); + builder.values().append_slice(&[1.5_f64, -2.5_f64]); + builder.append(true); + builder + .values() + .append_slice(&[f64::from(f32::MAX) * 2.0, f64::NAN]); + builder.append(true); + let array = Arc::new(builder.finish()) as ArrayRef; + + let tmp = write_list_parquet(schema, vec![array]); + + let mut reader = ParquetReader::::new(tmp.path(), None, NullHandling::FillZero).unwrap(); + let (data, num_samples, sample_size) = reader.read_batch().unwrap(); + + assert_eq!(num_samples, 2); + assert_eq!(sample_size, 2); + assert_eq!(data[0], 1.5_f32); + assert_eq!(data[1], -2.5_f32); + assert!( + data[2].is_infinite() && data[2] > 0.0, + "expected +Inf, got {}", + data[2] + ); + assert!(data[3].is_nan(), "expected NaN, got {}", data[3]); +} + +// --------------------------------------------------------------------------- +// Acceptance test 6: ParquetStreamingReader on f32 column +// --------------------------------------------------------------------------- + +/// ParquetStreamingReader:: on a List file → same values as +/// ParquetReader::. +#[test] +fn test_streaming_reader_list_f32() { + let item_field = Arc::new(Field::new("item", DataType::Float32, true)); + let list_field = Field::new("data", DataType::List(item_field), true); + let schema = Arc::new(Schema::new(vec![list_field])); + + let mut builder = ListBuilder::new(Float32Builder::new()); + builder.values().append_slice(&[1.0_f32, 2.5_f32, 3.75_f32]); + builder.append(true); + builder.values().append_slice(&[4.0_f32, 5.5_f32, 6.25_f32]); + builder.append(true); + let array = Arc::new(builder.finish()) as ArrayRef; + + let tmp = write_list_parquet(schema, vec![array]); + + let mut reader = + ParquetStreamingReader::::new(tmp.path(), None, NullHandling::FillZero).unwrap(); + let (data, num_samples, sample_size) = reader.read_batch().unwrap(); + + assert_eq!(num_samples, 2); + assert_eq!(sample_size, 3); + assert_eq!(data, vec![1.0_f32, 2.5, 3.75, 4.0, 5.5, 6.25]); +} + +// --------------------------------------------------------------------------- +// Acceptance test 7: ParquetStreamingReader on f64 column (cast path) +// --------------------------------------------------------------------------- + +/// ParquetStreamingReader:: on a List file → Arrow cast applied; +/// overflow → ±Inf, NaN preserved. +#[test] +fn test_streaming_reader_list_f64_cast_to_f32() { + let item_field = Arc::new(Field::new("item", DataType::Float64, true)); + let list_field = Field::new("data", DataType::List(item_field), true); + let schema = Arc::new(Schema::new(vec![list_field])); + + let mut builder = ListBuilder::new(Float64Builder::new()); + builder.values().append_slice(&[1.0_f64, -2.0_f64]); + builder.append(true); + builder.values().append_value(f64::from(f32::MAX) * 2.0); + builder.values().append_value(f64::NAN); + builder.append(true); + let array = Arc::new(builder.finish()) as ArrayRef; + + let tmp = write_list_parquet(schema, vec![array]); + + let mut reader = + ParquetStreamingReader::::new(tmp.path(), None, NullHandling::FillZero).unwrap(); + let (data, num_samples, sample_size) = reader.read_batch().unwrap(); + + assert_eq!(num_samples, 2); + assert_eq!(sample_size, 2); + assert_eq!(data[0], 1.0_f32); + assert_eq!(data[1], -2.0_f32); + assert!( + data[2].is_infinite() && data[2] > 0.0, + "expected +Inf, got {}", + data[2] + ); + assert!(data[3].is_nan(), "expected NaN, got {}", data[3]); +}