Revert "[ENH][PERF] Optimize block decoding (#5396)"

This reverts commit f3def7b.

Author: jairad26

Cargo.lock

@@ -37,7 +37,6 @@ sea-query = "0.32"
 sea-query-binder = "0.7"
 serde = { version = "1.0.215", features = ["derive", "rc"] }
 serde_json = "1.0.133"
-serde_bytes = "0.11.17"
 setsum = "0.7"
 sprs = "0.11"
 tantivy = "0.22.0"

Cargo.toml

@@ -8,7 +8,6 @@ path = "src/lib.rs"
 
 [dependencies]
 serde = { workspace = true }
-serde_bytes = { workspace = true }
 arrow = { workspace = true }
 thiserror = { workspace = true }
 uuid = { workspace = true }

rust/blockstore/Cargo.toml

@@ -1,15 +1,13 @@
 use std::cmp::Ordering;
 use std::collections::HashMap;
-use std::io::Read;
+use std::io::SeekFrom;
 use std::ops::{Bound, RangeBounds};
-use std::sync::Arc;
 
 use crate::arrow::types::{ArrowReadableKey, ArrowReadableValue};
 use arrow::array::ArrayData;
 use arrow::buffer::Buffer;
-use arrow::ipc::convert::fb_to_schema;
-use arrow::ipc::reader::{read_footer_length, FileDecoder};
-use arrow::ipc::{root_as_footer, root_as_message, Footer, MessageHeader, MetadataVersion};
+use arrow::ipc::reader::read_footer_length;
+use arrow::ipc::{root_as_footer, root_as_message, MessageHeader, MetadataVersion};
 use arrow::util::bit_util;
 use arrow::{
     array::{Array, StringArray},
@@ -68,8 +66,9 @@ impl<'de> Deserialize<'de> for RecordBatchWrapper {
     where
         D: serde::Deserializer<'de>,
     {
-        let data: &'de [u8] = serde_bytes::deserialize(deserializer)?;
-        let rb = Block::load_record_batch(data, false).map_err(D::Error::custom)?;
+        let data = Vec::<u8>::deserialize(deserializer)?;
+        let reader = std::io::Cursor::new(data);
+        let rb = Block::load_record_batch(reader, false).map_err(D::Error::custom)?;
         Ok(RecordBatchWrapper(rb))
     }
 }
@@ -522,15 +521,20 @@ impl Block {
 
     /// Load a block from bytes in Arrow IPC format with the given id
     pub fn from_bytes(bytes: &[u8], id: Uuid) -> Result<Self, BlockLoadError> {
-        Self::load_from_bytes(bytes, id, false)
+        Self::from_bytes_internal(bytes, id, false)
     }
 
     /// Load a block from bytes in Arrow IPC format with the given id and validate the layout
     /// ### Notes
     /// - This method should be used in tests to ensure that the layout of the IPC file is as expected
     /// - The validation is not performant and should not be used in production code
     pub fn from_bytes_with_validation(bytes: &[u8], id: Uuid) -> Result<Self, BlockLoadError> {
-        Self::load_from_bytes(bytes, id, true)
+        Self::from_bytes_internal(bytes, id, true)
+    }
+
+    fn from_bytes_internal(bytes: &[u8], id: Uuid, validate: bool) -> Result<Self, BlockLoadError> {
+        let cursor = std::io::Cursor::new(bytes);
+        Self::load_with_reader(cursor, id, validate)
     }
 
     /// Load a block from the given path with the given id and validate the layout
@@ -554,47 +558,41 @@ impl Block {
                 return Err(BlockLoadError::IOError(e));
             }
         };
-        let mut reader = std::io::BufReader::new(file);
-        let mut target_buffer = Vec::with_capacity(reader.get_ref().metadata()?.len() as usize);
-        reader.read_to_end(&mut target_buffer)?;
-        Self::load_from_bytes(&target_buffer, id, validate)
+        let reader = std::io::BufReader::new(file);
+        Self::load_with_reader(reader, id, validate)
     }
 
-    fn load_from_bytes(bytes: &[u8], id: Uuid, validate: bool) -> Result<Self, BlockLoadError> {
-        let batch = Self::load_record_batch(bytes, validate)?;
+    fn load_with_reader<R>(reader: R, id: Uuid, validate: bool) -> Result<Self, BlockLoadError>
+    where
+        R: std::io::Read + std::io::Seek,
+    {
+        let batch = Self::load_record_batch(reader, validate)?;
+        // TODO: how to store / hydrate id?
         Ok(Self::from_record_batch(id, batch))
     }
 
-    fn load_record_batch(bytes: &[u8], validate: bool) -> Result<RecordBatch, BlockLoadError> {
+    fn load_record_batch<R>(mut reader: R, validate: bool) -> Result<RecordBatch, BlockLoadError>
+    where
+        R: std::io::Read + std::io::Seek,
+    {
         if validate {
-            verify_buffers_layout(bytes).map_err(BlockLoadError::ArrowLayoutVerificationError)?;
+            verify_buffers_layout(&mut reader)
+                .map_err(BlockLoadError::ArrowLayoutVerificationError)?;
         }
 
-        let footer =
-            read_arrow_footer(bytes).map_err(BlockLoadError::ArrowLayoutVerificationError)?;
-        let schema = footer
-            .schema()
-            .ok_or(BlockLoadError::ArrowLayoutVerificationError(
-                ArrowLayoutVerificationError::RecordBatchDecodeError,
-            ))?;
-        let schema = fb_to_schema(schema);
-        // Requiring alignment should always work for blocks since we write them with alignment
-        // This is just being defensive
-        let decoder =
-            FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
-        let (block, record_batch_offset, _, record_batch_len) = read_record_batch_range(footer)?;
-
-        // This incurs a copy of the buffer we should be able to avoid this
-        // but as is foyer hands a reference to the [u8]. So the end to end
-        // path involves up to two copies, kernel to user space copy when reading from disk cache
-        // and then a copy into this buffer. We could avoid the second copy by changing foyer to
-        // hand over ownership of the buffer, but that would be a larger change.
-        // This is something we can optimize later if it becomes a bottleneck.
-        let buffer =
-            Buffer::from(&bytes[record_batch_offset..record_batch_offset + record_batch_len]);
-        decoder
-            .read_record_batch(block, &buffer)?
-            .ok_or(BlockLoadError::NoRecordBatches)
+        let mut arrow_reader = arrow::ipc::reader::FileReader::try_new(&mut reader, None)
+            .map_err(BlockLoadError::ArrowError)?;
+
+        let batch = match arrow_reader.next() {
+            Some(Ok(batch)) => batch,
+            Some(Err(e)) => {
+                return Err(BlockLoadError::ArrowError(e));
+            }
+            None => {
+                return Err(BlockLoadError::NoRecordBatches);
+            }
+        };
+        Ok(batch)
     }
 }
 
@@ -739,25 +737,40 @@ impl ChromaError for ArrowLayoutVerificationError {
     }
 }
 
-fn read_arrow_footer(bytes: &[u8]) -> Result<Footer<'_>, ArrowLayoutVerificationError> {
+/// Verifies that the buffers in the IPC file are 64 byte aligned
+/// and stored in Arrow in the way we expect.
+/// All non-benchmark test code should use this by loading the block
+/// with verification enabled.
+fn verify_buffers_layout<R>(mut reader: R) -> Result<(), ArrowLayoutVerificationError>
+where
+    R: std::io::Read + std::io::Seek,
+{
+    // Read the IPC file and verify that the buffers are 64 byte aligned
+    // by inspecting the offsets, this is required since our
+    // size calculation assumes that the buffers are 64 byte aligned
     // Space for ARROW_MAGIC (6 bytes) and length (4 bytes)
     let mut footer_buffer = [0; 10];
-    let trailer_start = bytes.len() - 10;
-    footer_buffer.copy_from_slice(&bytes[trailer_start..]);
-    let footer_len =
-        read_footer_length(footer_buffer).map_err(ArrowLayoutVerificationError::ArrowError)?;
+    reader
+        .seek(SeekFrom::End(-10))
+        .map_err(ArrowLayoutVerificationError::IOError)?;
+    reader
+        .read_exact(&mut footer_buffer)
+        .map_err(ArrowLayoutVerificationError::IOError)?;
+
+    let footer_len = read_footer_length(footer_buffer);
+    let footer_len = footer_len.map_err(ArrowLayoutVerificationError::ArrowError)?;
 
     // read footer
-    let footer_data = &bytes[trailer_start - footer_len..trailer_start];
+    let mut footer_data = vec![0; footer_len];
+    reader
+        .seek(SeekFrom::End(-10 - footer_len as i64))
+        .map_err(ArrowLayoutVerificationError::IOError)?;
+    reader
+        .read_exact(&mut footer_data)
+        .map_err(ArrowLayoutVerificationError::IOError)?;
     let footer =
-        root_as_footer(footer_data).map_err(ArrowLayoutVerificationError::InvalidFlatbuffer)?;
+        root_as_footer(&footer_data).map_err(ArrowLayoutVerificationError::InvalidFlatbuffer)?;
 
-    Ok(footer)
-}
-
-fn read_record_batch_range(
-    footer: Footer<'_>,
-) -> Result<(&arrow::ipc::Block, usize, usize, usize), ArrowLayoutVerificationError> {
     // Read the record batch
     let record_batch_definitions = match footer.recordBatches() {
         Some(record_batch_definitions) => record_batch_definitions,
@@ -772,32 +785,25 @@ fn read_record_batch_range(
     }
 
     let record_batch_definition = record_batch_definitions.get(0);
-    let record_batch_offset = record_batch_definition.offset() as usize;
+    let record_batch_len = record_batch_definition.bodyLength() as usize
+        + record_batch_definition.metaDataLength() as usize;
     let record_batch_body_len = record_batch_definition.bodyLength() as usize;
-    let record_batch_len =
-        record_batch_body_len + record_batch_definition.metaDataLength() as usize;
-
-    Ok((
-        record_batch_definition,
-        record_batch_offset,
-        record_batch_body_len,
-        record_batch_len,
-    ))
-}
 
-/// Verifies that the buffers in the IPC file are 64 byte aligned
-/// and stored in Arrow in the way we expect.
-/// All non-benchmark test code should use this by loading the block
-/// with verification enabled.
-fn verify_buffers_layout(bytes: &[u8]) -> Result<(), ArrowLayoutVerificationError> {
-    // Read the IPC file and verify that the buffers are 64 byte aligned
-    // by inspecting the offsets, this is required since our
-    // size calculation assumes that the buffers are 64 byte aligned
-    // Space for ARROW_MAGIC (6 bytes) and length (4 bytes)
-    let footer = read_arrow_footer(bytes)?;
-    let (_, record_batch_offset, record_batch_body_len, record_batch_len) =
-        read_record_batch_range(footer)?;
-    let buffer = Buffer::from(&bytes[record_batch_offset..record_batch_offset + record_batch_len]);
+    // Read the actual record batch
+    let mut file_buffer = vec![0; record_batch_len];
+    match reader.seek(SeekFrom::Start(record_batch_definition.offset() as u64)) {
+        Ok(_) => {}
+        Err(e) => {
+            return Err(ArrowLayoutVerificationError::IOError(e));
+        }
+    }
+    match reader.read_exact(&mut file_buffer) {
+        Ok(_) => {}
+        Err(e) => {
+            return Err(ArrowLayoutVerificationError::IOError(e));
+        }
+    }
+    let buffer = Buffer::from(file_buffer);
 
     // This is borrowed from arrow-ipc parse_message.rs
     // https://arrow.apache.org/docs/format/Columnar.html#encapsulated-message-format