[Parquet] Return error from RleDecoder::reload rather than panic

parquet/src/arrow/array_reader/byte_array.rs

@@ -274,7 +274,7 @@ impl ByteArrayDecoder {
                 validate_utf8,
             )),
             Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => ByteArrayDecoder::Dictionary(
-                ByteArrayDecoderDictionary::new(data, num_levels, num_values),
+                ByteArrayDecoderDictionary::new(data, num_levels, num_values)?,
             ),
             Encoding::DELTA_LENGTH_BYTE_ARRAY => ByteArrayDecoder::DeltaLength(
                 ByteArrayDecoderDeltaLength::new(data, validate_utf8)?,
@@ -563,10 +563,10 @@ pub struct ByteArrayDecoderDictionary {
 }
 
 impl ByteArrayDecoderDictionary {
-    fn new(data: Bytes, num_levels: usize, num_values: Option<usize>) -> Self {
-        Self {
-            decoder: DictIndexDecoder::new(data, num_levels, num_values),
-        }
+    fn new(data: Bytes, num_levels: usize, num_values: Option<usize>) -> Result<Self> {
+        Ok(Self {
+            decoder: DictIndexDecoder::new(data, num_levels, num_values)?,
+        })
     }
 
     fn read<I: OffsetSizeTrait>(

parquet/src/arrow/array_reader/byte_array_dictionary.rs

@@ -293,7 +293,7 @@ where
             Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => {
                 let bit_width = data[0];
                 let mut decoder = RleDecoder::new(bit_width);
-                decoder.set_data(data.slice(1..));
+                decoder.set_data(data.slice(1..))?;
                 MaybeDictionaryDecoder::Dict {
                     decoder,
                     max_remaining_values: num_values.unwrap_or(num_levels),

parquet/src/arrow/array_reader/byte_view_array.rs

@@ -236,7 +236,7 @@ impl ByteViewArrayDecoder {
             Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => {
                 ByteViewArrayDecoder::Dictionary(ByteViewArrayDecoderDictionary::new(
                     data, num_levels, num_values,
-                ))
+                )?)
             }
             Encoding::DELTA_LENGTH_BYTE_ARRAY => ByteViewArrayDecoder::DeltaLength(
                 ByteViewArrayDecoderDeltaLength::new(data, validate_utf8)?,
@@ -426,10 +426,10 @@ pub struct ByteViewArrayDecoderDictionary {
 }
 
 impl ByteViewArrayDecoderDictionary {
-    fn new(data: Bytes, num_levels: usize, num_values: Option<usize>) -> Self {
-        Self {
-            decoder: DictIndexDecoder::new(data, num_levels, num_values),
-        }
+    fn new(data: Bytes, num_levels: usize, num_values: Option<usize>) -> Result<Self> {
+        Ok(Self {
+            decoder: DictIndexDecoder::new(data, num_levels, num_values)?,
+        })
     }
 
     /// Reads the next indexes from self.decoder

parquet/src/arrow/array_reader/fixed_len_byte_array.rs

@@ -381,7 +381,7 @@ impl ColumnValueDecoder for ValueDecoder {
                 offset: 0,
             },
             Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => Decoder::Dict {
-                decoder: DictIndexDecoder::new(data, num_levels, num_values),
+                decoder: DictIndexDecoder::new(data, num_levels, num_values)?,
             },
             Encoding::DELTA_BYTE_ARRAY => Decoder::Delta {
                 decoder: DeltaByteArrayDecoder::new(data)?,

parquet/src/arrow/decoder/dictionary_index.rs

@@ -42,18 +42,18 @@ pub struct DictIndexDecoder {
 impl DictIndexDecoder {
     /// Create a new [`DictIndexDecoder`] with the provided data page, the number of levels
     /// associated with this data page, and the number of non-null values (if known)
-    pub fn new(data: Bytes, num_levels: usize, num_values: Option<usize>) -> Self {
+    pub fn new(data: Bytes, num_levels: usize, num_values: Option<usize>) -> Result<Self> {
         let bit_width = data[0];
         let mut decoder = RleDecoder::new(bit_width);
-        decoder.set_data(data.slice(1..));
+        decoder.set_data(data.slice(1..))?;
 
-        Self {
+        Ok(Self {
             decoder,
             index_buf: Box::new([0; 1024]),
             index_buf_len: 0,
             index_offset: 0,
             max_remaining_values: num_values.unwrap_or(num_levels),
-        }
+        })
     }
 
     /// Read up to `len` values, returning the number of values read

parquet/src/arrow/record_reader/definition_levels.rs

@@ -131,11 +131,12 @@ impl DefinitionLevelBufferDecoder {
 impl ColumnLevelDecoder for DefinitionLevelBufferDecoder {
     type Buffer = DefinitionLevelBuffer;
 
-    fn set_data(&mut self, encoding: Encoding, data: Bytes) {
+    fn set_data(&mut self, encoding: Encoding, data: Bytes) -> Result<()> {
         match &mut self.decoder {
             MaybePacked::Packed(d) => d.set_data(encoding, data),
-            MaybePacked::Fallback(d) => d.set_data(encoding, data),
-        }
+            MaybePacked::Fallback(d) => d.set_data(encoding, data)?,
+        };
+        Ok(())
     }
 }
 

parquet/src/column/reader.rs

@@ -451,7 +451,7 @@ where
                                 self.rep_level_decoder
                                     .as_mut()
                                     .unwrap()
-                                    .set_data(rep_level_encoding, level_data);
+                                    .set_data(rep_level_encoding, level_data)?;
                             }
 
                             if max_def_level > 0 {
@@ -466,7 +466,7 @@ where
                                 self.def_level_decoder
                                     .as_mut()
                                     .unwrap()
-                                    .set_data(def_level_encoding, level_data);
+                                    .set_data(def_level_encoding, level_data)?;
                             }
 
                             self.values_decoder.set_data(
@@ -512,7 +512,7 @@ where
                                 self.rep_level_decoder.as_mut().unwrap().set_data(
                                     Encoding::RLE,
                                     buf.slice(..rep_levels_byte_len as usize),
-                                );
+                                )?;
                             }
 
                             // DataPage v2 only supports RLE encoding for definition
@@ -524,7 +524,7 @@ where
                                         rep_levels_byte_len as usize
                                             ..(rep_levels_byte_len + def_levels_byte_len) as usize,
                                     ),
-                                );
+                                )?;
                             }
 
                             self.values_decoder.set_data(

parquet/src/column/reader/decoder.rs

@@ -32,7 +32,7 @@ pub trait ColumnLevelDecoder {
     type Buffer;
 
     /// Set data for this [`ColumnLevelDecoder`]
-    fn set_data(&mut self, encoding: Encoding, data: Bytes);
+    fn set_data(&mut self, encoding: Encoding, data: Bytes) -> Result<()>;
 }
 
 pub trait RepetitionLevelDecoder: ColumnLevelDecoder {
@@ -266,15 +266,15 @@ enum LevelDecoder {
 }
 
 impl LevelDecoder {
-    fn new(encoding: Encoding, data: Bytes, bit_width: u8) -> Self {
+    fn new(encoding: Encoding, data: Bytes, bit_width: u8) -> Result<Self> {
         match encoding {
             Encoding::RLE => {
                 let mut decoder = RleDecoder::new(bit_width);
-                decoder.set_data(data);
-                Self::Rle(decoder)
+                decoder.set_data(data)?;
+                Ok(Self::Rle(decoder))
             }
             #[allow(deprecated)]
-            Encoding::BIT_PACKED => Self::Packed(BitReader::new(data), bit_width),
+            Encoding::BIT_PACKED => Ok(Self::Packed(BitReader::new(data), bit_width)),
             _ => unreachable!("invalid level encoding: {}", encoding),
         }
     }
@@ -310,8 +310,9 @@ impl DefinitionLevelDecoderImpl {
 impl ColumnLevelDecoder for DefinitionLevelDecoderImpl {
     type Buffer = Vec<i16>;
 
-    fn set_data(&mut self, encoding: Encoding, data: Bytes) {
-        self.decoder = Some(LevelDecoder::new(encoding, data, self.bit_width))
+    fn set_data(&mut self, encoding: Encoding, data: Bytes) -> Result<()> {
+        self.decoder = Some(LevelDecoder::new(encoding, data, self.bit_width)?);
+        Ok(())
     }
 }
 
@@ -413,10 +414,11 @@ impl RepetitionLevelDecoderImpl {
 impl ColumnLevelDecoder for RepetitionLevelDecoderImpl {
     type Buffer = Vec<i16>;
 
-    fn set_data(&mut self, encoding: Encoding, data: Bytes) {
-        self.decoder = Some(LevelDecoder::new(encoding, data, self.bit_width));
+    fn set_data(&mut self, encoding: Encoding, data: Bytes) -> Result<()> {
+        self.decoder = Some(LevelDecoder::new(encoding, data, self.bit_width)?);
         self.buffer_len = 0;
         self.buffer_offset = 0;
+        Ok(())
     }
 }
 
@@ -499,14 +501,14 @@ mod tests {
         let data = Bytes::from(encoder.consume());
 
         let mut decoder = RepetitionLevelDecoderImpl::new(1);
-        decoder.set_data(Encoding::RLE, data.clone());
+        decoder.set_data(Encoding::RLE, data.clone()).unwrap();
         let (_, levels) = decoder.skip_rep_levels(100, 4).unwrap();
         assert_eq!(levels, 4);
 
         // The length of the final bit packed run is ambiguous, so without the correct
         // levels limit, it will decode zero padding
         let mut decoder = RepetitionLevelDecoderImpl::new(1);
-        decoder.set_data(Encoding::RLE, data);
+        decoder.set_data(Encoding::RLE, data).unwrap();
         let (_, levels) = decoder.skip_rep_levels(100, 6).unwrap();
         assert_eq!(levels, 6);
     }
@@ -525,7 +527,7 @@ mod tests {
             let data = Bytes::from(encoder.consume());
 
             let mut decoder = RepetitionLevelDecoderImpl::new(5);
-            decoder.set_data(Encoding::RLE, data);
+            decoder.set_data(Encoding::RLE, data).unwrap();
 
             let total_records = encoded.iter().filter(|x| **x == 0).count();
             let mut remaining_records = total_records;

parquet/src/encodings/decoding.rs

@@ -393,7 +393,7 @@ impl<T: DataType> Decoder<T> for DictDecoder<T> {
             ));
         }
         let mut rle_decoder = RleDecoder::new(bit_width);
-        rle_decoder.set_data(data.slice(1..));
+        rle_decoder.set_data(data.slice(1..))?;
         self.num_values = num_values;
         self.rle_decoder = Some(rle_decoder);
         Ok(())
@@ -473,7 +473,7 @@ impl<T: DataType> Decoder<T> for RleValueDecoder<T> {
 
         self.decoder = RleDecoder::new(1);
         self.decoder
-            .set_data(data.slice(I32_SIZE..I32_SIZE + data_size));
+            .set_data(data.slice(I32_SIZE..I32_SIZE + data_size))?;
         self.values_left = num_values;
         Ok(())
     }