Add unified and configurable null handling (#1101)

* Add unified and configurable null handling

* Remove unused import of `Array` in I/O utilities

Author: ryankert01

qdp/qdp-core/src/encoding/mod.rs

@@ -141,7 +141,8 @@ pub(crate) fn stream_encode<E: ChunkEncoder>(
     encoder: E,
 ) -> Result<*mut DLManagedTensor> {
     // Initialize reader
-    let mut reader_core = crate::io::ParquetBlockReader::new(path, None)?;
+    let mut reader_core =
+        crate::io::ParquetBlockReader::new(path, None, crate::reader::NullHandling::FillZero)?;
     let num_samples = reader_core.total_rows;
 
     // Allocate output state vector

qdp/qdp-core/src/io.rs

@@ -26,45 +26,43 @@ use std::fs::File;
 use std::path::Path;
 use std::sync::Arc;
 
-use arrow::array::{Array, ArrayRef, Float64Array, RecordBatch};
+use arrow::array::{ArrayRef, Float64Array, RecordBatch};
 use arrow::datatypes::{DataType, Field, Schema};
 use parquet::arrow::ArrowWriter;
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 use parquet::file::properties::WriterProperties;
 
 use crate::error::{MahoutError, Result};
+use crate::reader::{NullHandling, handle_float64_nulls};
 
 /// Converts an Arrow Float64Array to Vec<f64>.
-pub fn arrow_to_vec(array: &Float64Array) -> Vec<f64> {
-    if array.null_count() == 0 {
-        array.values().to_vec()
-    } else {
-        array.iter().map(|opt| opt.unwrap_or(0.0)).collect()
-    }
+pub fn arrow_to_vec(array: &Float64Array, null_handling: NullHandling) -> Result<Vec<f64>> {
+    let mut result = Vec::with_capacity(array.len());
+    handle_float64_nulls(&mut result, array, null_handling)?;
+    Ok(result)
 }
 
 /// Flattens multiple Arrow Float64Arrays into a single Vec<f64>.
-pub fn arrow_to_vec_chunked(arrays: &[Float64Array]) -> Vec<f64> {
+pub fn arrow_to_vec_chunked(
+    arrays: &[Float64Array],
+    null_handling: NullHandling,
+) -> Result<Vec<f64>> {
     let total_len: usize = arrays.iter().map(|a| a.len()).sum();
     let mut result = Vec::with_capacity(total_len);
 
     for array in arrays {
-        if array.null_count() == 0 {
-            result.extend_from_slice(array.values());
-        } else {
-            result.extend(array.iter().map(|opt| opt.unwrap_or(0.0)));
-        }
+        handle_float64_nulls(&mut result, array, null_handling)?;
     }
 
-    result
+    Ok(result)
 }
 
 /// Reads Float64 data from a Parquet file.
 ///
 /// Expects a single Float64 column. For zero-copy access, use [`read_parquet_to_arrow`].
 pub fn read_parquet<P: AsRef<Path>>(path: P) -> Result<Vec<f64>> {
     let chunks = read_parquet_to_arrow(path)?;
-    Ok(arrow_to_vec_chunked(&chunks))
+    arrow_to_vec_chunked(&chunks, NullHandling::FillZero)
 }
 
 /// Writes Float64 data to a Parquet file.
@@ -228,7 +226,7 @@ pub fn write_arrow_to_parquet<P: AsRef<Path>>(
 /// Add OOM protection for very large files
 pub fn read_parquet_batch<P: AsRef<Path>>(path: P) -> Result<(Vec<f64>, usize, usize)> {
     use crate::reader::DataReader;
-    let mut reader = crate::readers::ParquetReader::new(path, None)?;
+    let mut reader = crate::readers::ParquetReader::new(path, None, NullHandling::FillZero)?;
     reader.read_batch()
 }
 
@@ -244,7 +242,7 @@ pub fn read_parquet_batch<P: AsRef<Path>>(path: P) -> Result<(Vec<f64>, usize, u
 /// Add OOM protection for very large files
 pub fn read_arrow_ipc_batch<P: AsRef<Path>>(path: P) -> Result<(Vec<f64>, usize, usize)> {
     use crate::reader::DataReader;
-    let mut reader = crate::readers::ArrowIPCReader::new(path)?;
+    let mut reader = crate::readers::ArrowIPCReader::new(path, NullHandling::FillZero)?;
     reader.read_batch()
 }
 

qdp/qdp-core/src/lib.rs

@@ -34,6 +34,7 @@ mod profiling;
 
 pub use error::{MahoutError, Result, cuda_error_to_string};
 pub use gpu::memory::Precision;
+pub use reader::{NullHandling, handle_float64_nulls};
 
 // Throughput/latency pipeline runner: single path using QdpEngine and encode_batch in Rust.
 #[cfg(target_os = "linux")]

qdp/qdp-core/src/pipeline_runner.rs

@@ -26,7 +26,7 @@ use crate::QdpEngine;
 use crate::dlpack::DLManagedTensor;
 use crate::error::{MahoutError, Result};
 use crate::io;
-use crate::reader::StreamingDataReader;
+use crate::reader::{NullHandling, StreamingDataReader};
 use crate::readers::ParquetStreamingReader;
 
 /// Configuration for throughput/latency pipeline runs (Python run_throughput_pipeline_py).
@@ -39,6 +39,7 @@ pub struct PipelineConfig {
     pub encoding_method: String,
     pub seed: Option<u64>,
     pub warmup_batches: usize,
+    pub null_handling: NullHandling,
 }
 
 impl Default for PipelineConfig {
@@ -51,6 +52,7 @@ impl Default for PipelineConfig {
             encoding_method: "amplitude".to_string(),
             seed: None,
             warmup_batches: 0,
+            null_handling: NullHandling::FillZero,
         }
     }
 }
@@ -154,12 +156,23 @@ fn path_extension_lower(path: &Path) -> Option<String> {
 
 /// Dispatches by path extension to the appropriate io reader. Returns (data, num_samples, sample_size).
 /// Unsupported or missing extension returns Err with message listing supported formats.
-fn read_file_by_extension(path: &Path) -> Result<(Vec<f64>, usize, usize)> {
+fn read_file_by_extension(
+    path: &Path,
+    null_handling: NullHandling,
+) -> Result<(Vec<f64>, usize, usize)> {
     let ext_lower = path_extension_lower(path);
     let ext = ext_lower.as_deref();
     match ext {
-        Some("parquet") => io::read_parquet_batch(path),
-        Some("arrow") | Some("feather") | Some("ipc") => io::read_arrow_ipc_batch(path),
+        Some("parquet") => {
+            use crate::reader::DataReader;
+            let mut reader = crate::readers::ParquetReader::new(path, None, null_handling)?;
+            reader.read_batch()
+        }
+        Some("arrow") | Some("feather") | Some("ipc") => {
+            use crate::reader::DataReader;
+            let mut reader = crate::readers::ArrowIPCReader::new(path, null_handling)?;
+            reader.read_batch()
+        }
         Some("npy") => io::read_numpy_batch(path),
         Some("pt") | Some("pth") => io::read_torch_batch(path),
         Some("pb") => io::read_tensorflow_batch(path),
@@ -211,7 +224,7 @@ impl PipelineIterator {
         batch_limit: usize,
     ) -> Result<Self> {
         let path = path.as_ref();
-        let (data, num_samples, sample_size) = read_file_by_extension(path)?;
+        let (data, num_samples, sample_size) = read_file_by_extension(path, config.null_handling)?;
         let vector_len = vector_len(config.num_qubits, &config.encoding_method);
 
         // Dimension validation at construction.
@@ -263,7 +276,11 @@ impl PipelineIterator {
             )));
         }
 
-        let mut reader = ParquetStreamingReader::new(path, Some(DEFAULT_PARQUET_ROW_GROUP_SIZE))?;
+        let mut reader = ParquetStreamingReader::new(
+            path,
+            Some(DEFAULT_PARQUET_ROW_GROUP_SIZE),
+            config.null_handling,
+        )?;
         let vector_len = vector_len(config.num_qubits, &config.encoding_method);
 
         // Read first chunk to learn sample_size; reuse as initial buffer.

qdp/qdp-core/src/reader.rs

@@ -45,8 +45,48 @@
 //! }
 //! ```
 
+use arrow::array::{Array, Float64Array};
+
 use crate::error::Result;
 
+/// Policy for handling null values in Float64 arrays.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub enum NullHandling {
+    /// Replace nulls with 0.0 (backward-compatible default).
+    #[default]
+    FillZero,
+    /// Return an error when a null is encountered.
+    Reject,
+}
+
+/// 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<f64>,
+    float_array: &Float64Array,
+    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 Float64Array. \
+                     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:

qdp/qdp-core/src/readers/arrow_ipc.rs

@@ -24,20 +24,22 @@ use arrow::datatypes::DataType;
 use arrow::ipc::reader::FileReader as ArrowFileReader;
 
 use crate::error::{MahoutError, Result};
-use crate::reader::DataReader;
+use crate::reader::{DataReader, NullHandling, handle_float64_nulls};
 
 /// Reader for Arrow IPC files containing FixedSizeList<Float64> or List<Float64> columns.
 pub struct ArrowIPCReader {
     path: std::path::PathBuf,
     read: bool,
+    null_handling: NullHandling,
 }
 
 impl ArrowIPCReader {
     /// Create a new Arrow IPC reader.
     ///
     /// # Arguments
     /// * `path` - Path to the Arrow IPC file (.arrow or .feather)
-    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
+    /// * `null_handling` - Policy for null values (defaults to `FillZero`)
+    pub fn new<P: AsRef<Path>>(path: P, null_handling: NullHandling) -> Result<Self> {
         let path = path.as_ref();
 
         // Verify file exists
@@ -64,6 +66,7 @@ impl ArrowIPCReader {
         Ok(Self {
             path: path.to_path_buf(),
             read: false,
+            null_handling,
         })
     }
 }
@@ -136,11 +139,7 @@ impl DataReader for ArrowIPCReader {
                         .downcast_ref::<Float64Array>()
                         .ok_or_else(|| MahoutError::Io("Values must be Float64".to_string()))?;
 
-                    if float_array.null_count() == 0 {
-                        all_data.extend_from_slice(float_array.values());
-                    } else {
-                        all_data.extend(float_array.iter().map(|opt| opt.unwrap_or(0.0)));
-                    }
+                    handle_float64_nulls(&mut all_data, float_array, self.null_handling)?;
 
                     num_samples += list_array.len();
                 }
@@ -182,11 +181,7 @@ impl DataReader for ArrowIPCReader {
                             all_data.reserve(new_capacity);
                         }
 
-                        if float_array.null_count() == 0 {
-                            all_data.extend_from_slice(float_array.values());
-                        } else {
-                            all_data.extend(float_array.iter().map(|opt| opt.unwrap_or(0.0)));
-                        }
+                        handle_float64_nulls(&mut all_data, float_array, self.null_handling)?;
 
                         num_samples += 1;
                     }