feat: replace StreamReader/StreamWriter with FileReader/FileWriter

Author: Devdutt Shenoi

src/parseable/staging/reader.rs

@@ -19,16 +19,14 @@
 
 use std::{
     fs::{remove_file, File},
-    io::{self, BufReader, Read, Seek, SeekFrom},
+    io::BufReader,
     path::PathBuf,
     sync::Arc,
-    vec::IntoIter,
 };
 
 use arrow_array::{RecordBatch, TimestampMillisecondArray};
-use arrow_ipc::{reader::StreamReader, root_as_message_unchecked, MessageHeader};
+use arrow_ipc::reader::FileReader;
 use arrow_schema::Schema;
-use byteorder::{LittleEndian, ReadBytesExt};
 use itertools::kmerge_by;
 use tracing::error;
 
@@ -39,7 +37,7 @@ use crate::{
 
 #[derive(Debug)]
 pub struct MergedRecordReader {
-    pub readers: Vec<StreamReader<BufReader<File>>>,
+    pub readers: Vec<FileReader<BufReader<File>>>,
 }
 
 impl MergedRecordReader {
@@ -53,7 +51,7 @@ impl MergedRecordReader {
                 remove_file(file).unwrap();
             } else {
                 let Ok(reader) =
-                    StreamReader::try_new(BufReader::new(File::open(file).unwrap()), None)
+                    FileReader::try_new(BufReader::new(File::open(file).unwrap()), None)
                 else {
                     error!("Invalid file detected, ignoring it: {:?}", file);
                     continue;
@@ -74,27 +72,6 @@ impl MergedRecordReader {
         )
         .unwrap()
     }
-}
-
-#[derive(Debug)]
-pub struct MergedReverseRecordReader {
-    pub readers: Vec<StreamReader<BufReader<OffsetReader<File>>>>,
-}
-
-impl MergedReverseRecordReader {
-    pub fn try_new(files: &[PathBuf]) -> Self {
-        let mut readers = Vec::with_capacity(files.len());
-        for file in files {
-            let Ok(reader) = get_reverse_reader(File::open(file).unwrap()) else {
-                error!("Invalid file detected, ignoring it: {:?}", file);
-                continue;
-            };
-
-            readers.push(reader);
-        }
-
-        Self { readers }
-    }
 
     pub fn merged_iter(
         self,
@@ -111,15 +88,6 @@ impl MergedReverseRecordReader {
         .map(|batch| reverse(&batch))
         .map(move |batch| adapt_batch(&schema, &batch))
     }
-
-    pub fn merged_schema(&self) -> Schema {
-        Schema::try_merge(
-            self.readers
-                .iter()
-                .map(|reader| reader.schema().as_ref().clone()),
-        )
-        .unwrap()
-    }
 }
 
 fn get_timestamp_millis(batch: &RecordBatch, time_partition: Option<String>) -> i64 {
@@ -138,6 +106,7 @@ fn get_timestamp_millis(batch: &RecordBatch, time_partition: Option<String>) ->
         None => get_default_timestamp_millis(batch),
     }
 }
+
 fn get_default_timestamp_millis(batch: &RecordBatch) -> i64 {
     match batch
         .column(0)
@@ -157,172 +126,18 @@ fn get_default_timestamp_millis(batch: &RecordBatch) -> i64 {
     }
 }
 
-/// OffsetReader takes in a reader and list of offset and sizes and
-/// provides a reader over the file by reading only the offsets
-/// from start of the list to end.
-///
-/// Safety Invariant: Reader is already validated and all offset and limit are valid to read.
-///
-/// On empty list the reader returns no bytes read.
-pub struct OffsetReader<R: Read + Seek> {
-    reader: R,
-    offset_list: IntoIter<(u64, usize)>,
-    current_offset: u64,
-    current_size: usize,
-    buffer: Vec<u8>,
-    buffer_position: usize,
-    finished: bool,
-}
-
-impl<R: Read + Seek> OffsetReader<R> {
-    fn new(reader: R, offset_list: Vec<(u64, usize)>) -> Self {
-        let mut offset_list = offset_list.into_iter();
-        let mut finished = false;
-
-        let (current_offset, current_size) = offset_list.next().unwrap_or_default();
-        if current_offset == 0 && current_size == 0 {
-            finished = true
-        }
-
-        OffsetReader {
-            reader,
-            offset_list,
-            current_offset,
-            current_size,
-            buffer: vec![0; 4096],
-            buffer_position: 0,
-            finished,
-        }
-    }
-}
-
-impl<R: Read + Seek> Read for OffsetReader<R> {
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        let offset = self.current_offset;
-        let size = self.current_size;
-
-        if self.finished {
-            return Ok(0);
-        }
-        // on empty buffer load current data represented by
-        // current_offset and current_size into self buffer
-        if self.buffer_position == 0 {
-            self.reader.seek(SeekFrom::Start(offset))?;
-            // resize for current message
-            if self.buffer.len() < size {
-                self.buffer.resize(size, 0)
-            }
-            self.reader.read_exact(&mut self.buffer[0..size])?;
-        }
-
-        let remaining_bytes = size - self.buffer_position;
-        let max_read = usize::min(remaining_bytes, buf.len());
-
-        // Copy data from the buffer to the provided buffer
-        let read_data = &self.buffer[self.buffer_position..self.buffer_position + max_read];
-        buf[..max_read].copy_from_slice(read_data);
-
-        self.buffer_position += max_read;
-
-        if self.buffer_position >= size {
-            // If we've read the entire section, move to the next offset
-            match self.offset_list.next() {
-                Some((offset, size)) => {
-                    self.current_offset = offset;
-                    self.current_size = size;
-                    self.buffer_position = 0;
-                }
-                None => {
-                    // iter is exhausted, no more read can be done
-                    self.finished = true
-                }
-            }
-        }
-
-        Ok(max_read)
-    }
-}
-
-pub fn get_reverse_reader<T: Read + Seek>(
-    mut reader: T,
-) -> Result<StreamReader<BufReader<OffsetReader<T>>>, io::Error> {
-    let mut offset = 0;
-    let mut messages = Vec::new();
-
-    while let Some(res) = find_limit_and_type(&mut reader).transpose() {
-        match res {
-            Ok((header, size)) => {
-                messages.push((header, offset, size));
-                offset += size;
-            }
-            Err(err) if err.kind() == io::ErrorKind::UnexpectedEof => break,
-            Err(err) => return Err(err),
-        }
-    }
-
-    // reverse everything leaving the first because it has schema message.
-    messages[1..].reverse();
-    let messages = messages
-        .into_iter()
-        .map(|(_, offset, size)| (offset as u64, size))
-        .collect();
-
-    // reset reader
-    reader.rewind()?;
-
-    Ok(StreamReader::try_new(BufReader::new(OffsetReader::new(reader, messages)), None).unwrap())
-}
-
-// return limit for
-fn find_limit_and_type(
-    reader: &mut (impl Read + Seek),
-) -> Result<Option<(MessageHeader, usize)>, io::Error> {
-    let mut size = 0;
-    let marker = reader.read_u32::<LittleEndian>()?;
-    size += 4;
-
-    if marker != 0xFFFFFFFF {
-        return Err(io::Error::new(
-            io::ErrorKind::InvalidData,
-            "Invalid Continuation Marker",
-        ));
-    }
-
-    let metadata_size = reader.read_u32::<LittleEndian>()? as usize;
-    size += 4;
-
-    if metadata_size == 0x00000000 {
-        return Ok(None);
-    }
-
-    let mut message = vec![0u8; metadata_size];
-    reader.read_exact(&mut message)?;
-    size += metadata_size;
-
-    let message = unsafe { root_as_message_unchecked(&message) };
-    let header = message.header_type();
-    let message_size = message.bodyLength();
-    size += message_size as usize;
-
-    let padding = (8 - (size % 8)) % 8;
-    reader.seek(SeekFrom::Current(padding as i64 + message_size))?;
-    size += padding;
-
-    Ok(Some((header, size)))
-}
-
 #[cfg(test)]
 mod tests {
-    use std::{io::Cursor, sync::Arc};
+    use std::{fs::File, path::Path, sync::Arc};
 
     use arrow_array::{
         cast::AsArray, types::Int64Type, Array, Float64Array, Int64Array, RecordBatch, StringArray,
     };
-    use arrow_ipc::writer::{
-        write_message, DictionaryTracker, IpcDataGenerator, IpcWriteOptions, StreamWriter,
+    use arrow_ipc::{
+        reader::FileReader,
+        writer::{write_message, DictionaryTracker, FileWriter, IpcDataGenerator, IpcWriteOptions},
     };
-
-    use super::get_reverse_reader;
+    use temp_dir::TempDir;
 
     fn rb(rows: usize) -> RecordBatch {
         let array1: Arc<dyn Array> = Arc::new(Int64Array::from_iter(0..(rows as i64)));
@@ -339,42 +154,48 @@ mod tests {
         .unwrap()
     }
 
-    fn write_mem(rbs: &[RecordBatch]) -> Vec<u8> {
-        let buf = Vec::new();
-        let mut writer = StreamWriter::try_new(buf, &rbs[0].schema()).unwrap();
+    fn write_file(rbs: &[RecordBatch], path: &Path) {
+        let file = File::create(path).unwrap();
+        let mut writer = FileWriter::try_new_buffered(file, &rbs[0].schema()).unwrap();
 
         for rb in rbs {
             writer.write(rb).unwrap()
         }
 
-        writer.into_inner().unwrap()
+        writer.finish().unwrap();
     }
 
     #[test]
     fn test_empty_row() {
+        let temp_dir = TempDir::new().unwrap();
+        let path = temp_dir.path().join("test.arrows");
         let rb = rb(0);
-        let buf = write_mem(&[rb]);
-        let reader = Cursor::new(buf);
-        let mut reader = get_reverse_reader(reader).unwrap();
+        write_file(&[rb], &path);
+        let reader = File::open(path).unwrap();
+        let mut reader = FileReader::try_new_buffered(reader, None).unwrap();
         let rb = reader.next().unwrap().unwrap();
         assert_eq!(rb.num_rows(), 0);
     }
 
     #[test]
     fn test_one_row() {
+        let temp_dir = TempDir::new().unwrap();
+        let path = temp_dir.path().join("test.arrows");
         let rb = rb(1);
-        let buf = write_mem(&[rb]);
-        let reader = Cursor::new(buf);
-        let mut reader = get_reverse_reader(reader).unwrap();
+        write_file(&[rb], &path);
+        let reader = File::open(path).unwrap();
+        let mut reader = FileReader::try_new_buffered(reader, None).unwrap();
         let rb = reader.next().unwrap().unwrap();
         assert_eq!(rb.num_rows(), 1);
     }
 
     #[test]
     fn test_multiple_row_multiple_rbs() {
-        let buf = write_mem(&[rb(1), rb(2), rb(3)]);
-        let reader = Cursor::new(buf);
-        let mut reader = get_reverse_reader(reader).unwrap();
+        let temp_dir = TempDir::new().unwrap();
+        let path = temp_dir.path().join("test.arrows");
+        write_file(&[rb(1), rb(2), rb(3)], &path);
+        let reader = File::open(path).unwrap();
+        let mut reader = FileReader::try_new_buffered(reader, None).unwrap();
         let rb = reader.next().unwrap().unwrap();
         assert_eq!(rb.num_rows(), 3);
         let col1_val: Vec<i64> = rb
@@ -394,40 +215,42 @@ mod tests {
 
     #[test]
     fn manual_write() {
+        let temp_dir = TempDir::new().unwrap();
+        let path = temp_dir.path().join("test.arrows");
         let error_on_replacement = true;
         let options = IpcWriteOptions::default();
         let mut dictionary_tracker = DictionaryTracker::new(error_on_replacement);
         let data_gen = IpcDataGenerator {};
-
-        let mut buf = Vec::new();
+        let mut file = File::create(&path).unwrap();
         let rb1 = rb(1);
 
         let schema = data_gen.schema_to_bytes_with_dictionary_tracker(
             &rb1.schema(),
             &mut dictionary_tracker,
             &options,
         );
-        write_message(&mut buf, schema, &options).unwrap();
+        write_message(&mut file, schema, &options).unwrap();
 
         for i in (1..=3).cycle().skip(1).take(10000) {
             let (_, encoded_message) = data_gen
                 .encoded_batch(&rb(i), &mut dictionary_tracker, &options)
                 .unwrap();
-            write_message(&mut buf, encoded_message, &options).unwrap();
+            write_message(&mut file, encoded_message, &options).unwrap();
         }
 
         let schema = data_gen.schema_to_bytes_with_dictionary_tracker(
             &rb1.schema(),
             &mut dictionary_tracker,
             &options,
         );
-        write_message(&mut buf, schema, &options).unwrap();
+        write_message(&mut file, schema, &options).unwrap();
 
-        let buf = Cursor::new(buf);
-        let reader = get_reverse_reader(buf).unwrap().flatten();
+        let reader = File::open(path).unwrap();
+        let reader = FileReader::try_new_buffered(reader, None).unwrap();
 
         let mut sum = 0;
         for rb in reader {
+            let rb = rb.unwrap();
             sum += 1;
             assert!(rb.num_rows() > 0);
         }