From 81f800fd05856787f5bd373e1138385b8a441dcb Mon Sep 17 00:00:00 2001
From: Yang Xiufeng <yangxiufeng.c@gmail.com>
Date: Tue, 25 Jun 2024 21:16:01 +0800
Subject: [PATCH] fix: panic when read column of type array<float>

---
 src/array_decoder/mod.rs   |  4 ++--
 src/reader/decode/float.rs | 39 ++++++++++++++------------------------
 2 files changed, 16 insertions(+), 27 deletions(-)

diff --git a/src/array_decoder/mod.rs b/src/array_decoder/mod.rs
index 755e920c..ee17392d 100644
--- a/src/array_decoder/mod.rs
+++ b/src/array_decoder/mod.rs
@@ -404,7 +404,7 @@ pub fn array_decoder_factory(
                 }
             );
             let iter = stripe.stream_map().get(column, Kind::Data);
-            let iter = Box::new(FloatIter::new(iter, stripe.number_of_rows()));
+            let iter = Box::new(FloatIter::new(iter));
             let present = get_present_vec(column, stripe)?
                 .map(|iter| Box::new(iter.into_iter()) as Box<dyn Iterator<Item = bool> + Send>);
             Box::new(Float32ArrayDecoder::new(iter, present))
@@ -418,7 +418,7 @@ pub fn array_decoder_factory(
                 }
             );
             let iter = stripe.stream_map().get(column, Kind::Data);
-            let iter = Box::new(FloatIter::new(iter, stripe.number_of_rows()));
+            let iter = Box::new(FloatIter::new(iter));
             let present = get_present_vec(column, stripe)?
                 .map(|iter| Box::new(iter.into_iter()) as Box<dyn Iterator<Item = bool> + Send>);
             Box::new(Float64ArrayDecoder::new(iter, present))
diff --git a/src/reader/decode/float.rs b/src/reader/decode/float.rs
index 20d2496d..c681d88b 100644
--- a/src/reader/decode/float.rs
+++ b/src/reader/decode/float.rs
@@ -32,26 +32,18 @@ impl Float for f64 {
 /// An iterator
 pub struct FloatIter<T: Float, R: std::io::Read> {
     reader: R,
-    remaining: usize,
     phantom: std::marker::PhantomData<T>,
 }
 
 impl<T: Float, R: std::io::Read> FloatIter<T, R> {
     /// Returns a new [`FloatIter`]
     #[inline]
-    pub fn new(reader: R, length: usize) -> Self {
+    pub fn new(reader: R) -> Self {
         Self {
             reader,
-            remaining: length,
             phantom: Default::default(),
         }
     }
-
-    /// The number of items remaining
-    #[inline]
-    pub fn len(&self) -> usize {
-        self.remaining
-    }
 }
 
 impl<T: Float, R: std::io::Read> Iterator for FloatIter<T, R> {
@@ -59,26 +51,23 @@ impl<T: Float, R: std::io::Read> Iterator for FloatIter<T, R> {
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        if self.remaining == 0 {
-            return None;
-        }
         let mut chunk: T::OBytes = Default::default();
-        if let Err(err) = self
+        match self
             .reader
-            .read_exact(chunk.as_mut())
+            .read(chunk.as_mut())
             .context(error::DecodeFloatSnafu)
         {
-            return Some(Err(err));
+            Err(err) => {
+                return Some(Err(err));
+            }
+            Ok(n) => {
+                if n == 0 {
+                    return None;
+                }
+            }
         };
-        self.remaining -= 1;
         Some(Ok(T::from_le_bytes(chunk)))
     }
-
-    #[inline]
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        let remaining = self.len();
-        (remaining, Some(remaining))
-    }
 }
 
 #[cfg(test)]
@@ -100,7 +89,7 @@ mod tests {
         let bytes = float_to_bytes(&input);
         let bytes = Cursor::new(bytes);
 
-        let iter = FloatIter::<F, _>::new(bytes, input.len());
+        let iter = FloatIter::<F, _>::new(bytes);
         let actual = iter.collect::<Result<Vec<_>>>().unwrap();
 
         assert_eq!(input, actual);
@@ -131,7 +120,7 @@ mod tests {
         let bytes = float_to_bytes(&[f32::NAN]);
         let bytes = Cursor::new(bytes);
 
-        let iter = FloatIter::<f32, _>::new(bytes, 1);
+        let iter = FloatIter::<f32, _>::new(bytes);
         let actual = iter.collect::<Result<Vec<_>>>().unwrap();
         assert_eq!(actual.len(), 1);
         assert!(actual[0].is_nan());
@@ -142,7 +131,7 @@ mod tests {
         let bytes = float_to_bytes(&[f64::NAN]);
         let bytes = Cursor::new(bytes);
 
-        let iter = FloatIter::<f64, _>::new(bytes, 1);
+        let iter = FloatIter::<f64, _>::new(bytes);
         let actual = iter.collect::<Result<Vec<_>>>().unwrap();
         assert_eq!(actual.len(), 1);
         assert!(actual[0].is_nan());