impl skip PrimitiveValueDecoder

Author: suxiaogang223

src/array_decoder/decimal.rs

@@ -112,6 +112,12 @@ struct DecimalScaleRepairDecoder {
 }
 
 impl PrimitiveValueDecoder<i128> for DecimalScaleRepairDecoder {
+    fn skip(&mut self, n: usize) -> Result<()> {
+        self.varint_iter.skip(n)?;
+        self.scale_iter.skip(n)?;
+        Ok(())
+    }
+
     fn decode(&mut self, out: &mut [i128]) -> Result<()> {
         // TODO: can probably optimize, reuse buffers?
         let mut varint = vec![0; out.len()];

src/array_decoder/mod.rs

@@ -74,6 +74,12 @@ pub trait ArrayBatchDecoder: Send {
         batch_size: usize,
         parent_present: Option<&NullBuffer>,
     ) -> Result<ArrayRef>;
+
+    /// Skip the next `n` values without decoding them, failing if it cannot skip the enough values.
+    fn skip_values(&mut self, n: usize) -> Result<()> {
+        // TODO: implement
+        Ok(())
+    }
 }
 
 struct PrimitiveArrayDecoder<T: ArrowPrimitiveType> {
@@ -122,6 +128,24 @@ impl<T: ArrowPrimitiveType> ArrayBatchDecoder for PrimitiveArrayDecoder<T> {
         let array = Arc::new(array) as ArrayRef;
         Ok(array)
     }
+
+    fn skip_rows(&mut self, row_count: usize) -> Result<()> {
+        // If we have a present stream, we need to decode it to know how many
+        // non-null values to skip in the data stream
+        let non_null_count = if let Some(ref mut present) = self.present {
+            let mut present_buffer = vec![false; row_count];
+            present.inner.decode(&mut present_buffer)?;
+            // Count non-null values (where present is true)
+            present_buffer.iter().filter(|&&v| v).count()
+        } else {
+            // No nulls, so all rows have values
+            row_count
+        };
+
+        // Skip the data stream for non-null values only
+        self.iter.skip(non_null_count)?;
+        Ok(())
+    }
 }
 
 type Int64ArrayDecoder = PrimitiveArrayDecoder<Int64Type>;
@@ -167,6 +191,24 @@ impl ArrayBatchDecoder for BooleanArrayDecoder {
         };
         Ok(Arc::new(array))
     }
+
+    fn skip_rows(&mut self, row_count: usize) -> Result<()> {
+        // If we have a present stream, we need to decode it to know how many
+        // non-null values to skip in the data stream
+        let non_null_count = if let Some(ref mut present) = self.present {
+            let mut present_buffer = vec![false; row_count];
+            present.inner.decode(&mut present_buffer)?;
+            // Count non-null values (where present is true)
+            present_buffer.iter().filter(|&&v| v).count()
+        } else {
+            // No nulls, so all rows have values
+            row_count
+        };
+
+        // Skip the data stream for non-null values only
+        self.iter.skip(non_null_count)?;
+        Ok(())
+    }
 }
 
 struct PresentDecoder {
@@ -514,14 +556,11 @@ impl NaiveStripeDecoder {
         })
     }
 
-    /// Skip the specified number of rows by decoding and discarding them
+    /// Skip the specified number of rows by calling skip_rows on each decoder
     fn skip_rows(&mut self, count: usize) -> Result<()> {
-        // Decode in batches to avoid large memory allocations
-        let mut remaining = count;
-        while remaining > 0 {
-            let chunk = self.batch_size.min(remaining);
-            let _ = self.inner_decode_next_batch(chunk)?;
-            remaining -= chunk;
+        // Call skip_rows on each decoder to efficiently skip rows
+        for decoder in &mut self.decoders {
+            decoder.skip_rows(count)?;
         }
         Ok(())
     }

src/array_decoder/timestamp.rs

@@ -308,6 +308,11 @@ impl TimestampNanosecondAsDecimalWithTzDecoder {
 }
 
 impl PrimitiveValueDecoder<i128> for TimestampNanosecondAsDecimalWithTzDecoder {
+    fn skip(&mut self, n: usize) -> Result<()> {
+        self.0.skip(n)?;
+        Ok(())
+    }
+
     fn decode(&mut self, out: &mut [i128]) -> Result<()> {
         self.0.decode(out)?;
         for x in out.iter_mut() {

src/encoding/boolean.rs

@@ -55,6 +55,11 @@ impl<R: Read> BooleanDecoder<R> {
 }
 
 impl<R: Read> PrimitiveValueDecoder<bool> for BooleanDecoder<R> {
+    fn skip(&mut self, n: usize) -> Result<()> {
+        self.decoder.skip(n)?;
+        Ok(())
+    }
+
     // TODO: can probably implement this better
     fn decode(&mut self, out: &mut [bool]) -> Result<()> {
         for x in out.iter_mut() {

src/encoding/decimal.rs

@@ -36,6 +36,13 @@ impl<R: Read> UnboundedVarintStreamDecoder<R> {
 }
 
 impl<R: Read> PrimitiveValueDecoder<i128> for UnboundedVarintStreamDecoder<R> {
+    fn skip(&mut self, n: usize) -> Result<()> {
+        for _ in 0..n {
+            read_varint_zigzagged::<i128, _, SignedEncoding>(&mut self.reader)?;
+        }
+        Ok(())
+    }
+
     fn decode(&mut self, out: &mut [i128]) -> Result<()> {
         for x in out.iter_mut() {
             *x = read_varint_zigzagged::<i128, _, SignedEncoding>(&mut self.reader)?;

src/encoding/float.rs

@@ -22,7 +22,7 @@ use bytes::{Bytes, BytesMut};
 use snafu::ResultExt;
 
 use crate::{
-    error::{IoSnafu, Result},
+    error::{IoSnafu, OutOfSpecSnafu, Result},
     memory::EstimateMemory,
 };
 
@@ -36,12 +36,12 @@ pub trait Float:
 impl Float for f32 {}
 impl Float for f64 {}
 
-pub struct FloatDecoder<F: Float, R: std::io::Read> {
+pub struct FloatDecoder<F: Float, R: std::io::Read + std::io::Seek> {
     reader: R,
     phantom: std::marker::PhantomData<F>,
 }
 
-impl<F: Float, R: std::io::Read> FloatDecoder<F, R> {
+impl<F: Float, R: std::io::Read + std::io::Seek> FloatDecoder<F, R> {
     pub fn new(reader: R) -> Self {
         Self {
             reader,
@@ -50,7 +50,28 @@ impl<F: Float, R: std::io::Read> FloatDecoder<F, R> {
     }
 }
 
-impl<F: Float, R: std::io::Read> PrimitiveValueDecoder<F> for FloatDecoder<F, R> {
+impl<F: Float, R: std::io::Read + std::io::Seek> PrimitiveValueDecoder<F> for FloatDecoder<F, R> {
+    fn skip(&mut self, n: usize) -> Result<()> {
+        let bytes_to_skip = n * std::mem::size_of::<F>();
+        let cur_pos = self.reader.stream_position().context(IoSnafu)?;
+        let seek_pos = self
+            .reader
+            .seek(std::io::SeekFrom::Current(bytes_to_skip as i64))
+            .context(IoSnafu)?;
+        let actual_skip = seek_pos - cur_pos;
+        if actual_skip != bytes_to_skip as u64 {
+            return OutOfSpecSnafu {
+                msg: format!(
+                    "failed to skip float values, expected to skip {} bytes, but actually skipped {}",
+                    bytes_to_skip, actual_skip
+                ),
+            }
+            .fail();
+        }
+
+        Ok(())
+    }
+
     fn decode(&mut self, out: &mut [F]) -> Result<()> {
         let bytes = must_cast_slice_mut::<F, u8>(out);
         self.reader.read_exact(bytes).context(IoSnafu)?;

src/encoding/mod.rs

@@ -51,7 +51,7 @@ pub trait PrimitiveValueEncoder<V: Copy>: EstimateMemory {
 }
 
 pub trait PrimitiveValueDecoder<V> {
-    /// Skip the next `n` values without decoding them. Failing if it cannot skip the enough values.
+    /// Skip the next `n` values without decoding them, failing if it cannot skip the enough values.
     fn skip(&mut self, n: usize) -> Result<()>;
 
     /// Decode out.len() values into out at a time, failing if it cannot fill
@@ -98,14 +98,25 @@ mod tests {
     use super::*;
 
     /// Emits numbers increasing from 0.
-    struct DummyDecoder;
+    struct DummyDecoder {
+        value: i32,
+    }
+
+    impl DummyDecoder {
+        fn new() -> Self {
+            Self { value: 0 }
+        }
+    }
 
     impl PrimitiveValueDecoder<i32> for DummyDecoder {
-        fn skip(&mut self, n: usize) -> Result<()> {
+        fn skip(&mut self, _n: usize) -> Result<()> {
+            self.value += 1;
             Ok(())
         }
         fn decode(&mut self, out: &mut [i32]) -> Result<()> {
-            let values = (0..out.len()).map(|x| x as i32).collect::<Vec<_>>();
+            let values = (0..out.len())
+                .map(|x| self.value + x as i32)
+                .collect::<Vec<_>>();
             out.copy_from_slice(&values);
             Ok(())
         }
@@ -128,7 +139,7 @@ mod tests {
     proptest! {
         #[test]
         fn decode_spaced_proptest(present: Vec<bool>) {
-            let mut decoder = DummyDecoder;
+            let mut decoder = DummyDecoder::new();
             let mut out = vec![-1; present.len()];
             decoder.decode_spaced(&mut out, &NullBuffer::from(present.clone())).unwrap();
             let expected = gen_spaced_dummy_decoder_expected(&present);
@@ -138,7 +149,7 @@ mod tests {
 
     #[test]
     fn decode_spaced_edge_cases() {
-        let mut decoder = DummyDecoder;
+        let mut decoder = DummyDecoder::new();
         let len = 10;
 
         // all present

src/encoding/rle.rs

@@ -61,6 +61,7 @@ pub trait GenericRle<V: Copy> {
 
 impl<V: Copy, G: GenericRle<V> + sealed::Rle> PrimitiveValueDecoder<V> for G {
     fn skip(&mut self, n: usize) -> Result<()> {
+        // Delegate to the GenericRle implementation
         self.skip_values(n)
     }
 

src/encoding/timestamp.rs

@@ -50,6 +50,12 @@ impl<T: ArrowTimestampType> TimestampDecoder<T> {
 }
 
 impl<T: ArrowTimestampType> PrimitiveValueDecoder<T::Native> for TimestampDecoder<T> {
+    fn skip(&mut self, n: usize) -> Result<()> {
+        self.data.skip(n)?;
+        self.secondary.skip(n)?;
+        Ok(())
+    }
+
     fn decode(&mut self, out: &mut [T::Native]) -> Result<()> {
         // TODO: can probably optimize, reuse buffers?
         let mut data = vec![0; out.len()];
@@ -90,6 +96,12 @@ impl TimestampNanosecondAsDecimalDecoder {
 }
 
 impl PrimitiveValueDecoder<i128> for TimestampNanosecondAsDecimalDecoder {
+    fn skip(&mut self, n: usize) -> Result<()> {
+        self.data.skip(n)?;
+        self.secondary.skip(n)?;
+        Ok(())
+    }
+
     fn decode(&mut self, out: &mut [i128]) -> Result<()> {
         // TODO: can probably optimize, reuse buffers?
         let mut data = vec![0; out.len()];