feat: Support streaming IPC scan from S3 object store (#25868)

Author: kdn36

crates/polars-arrow/src/io/ipc/read/file.rs

@@ -33,7 +33,7 @@ pub struct FileMetadata {
     pub blocks: Vec<arrow_format::ipc::Block>,
 
     /// Dictionaries associated to each dict_id
-    pub(crate) dictionaries: Option<Vec<arrow_format::ipc::Block>>,
+    pub dictionaries: Option<Vec<arrow_format::ipc::Block>>,
 
     /// The total size of the file in bytes
     pub size: u64,
@@ -78,27 +78,34 @@ pub(crate) fn get_dictionary_batch<'a>(
     }
 }
 
-fn read_dictionary_block<R: Read + Seek>(
+pub fn read_dictionary_block<R: Read + Seek>(
     reader: &mut R,
     metadata: &FileMetadata,
     block: &arrow_format::ipc::Block,
+    // When true, the underlying reader bytestream represents a standalone IPC Block
+    // rather than a complete IPC File.
+    force_zero_offset: bool,
     dictionaries: &mut Dictionaries,
     message_scratch: &mut Vec<u8>,
     dictionary_scratch: &mut Vec<u8>,
 ) -> PolarsResult<()> {
-    let message = get_message_from_block(reader, block, message_scratch)?;
-    let batch = get_dictionary_batch(&message)?;
-
-    let offset: u64 = block
-        .offset
-        .try_into()
-        .map_err(|_| polars_err!(oos = OutOfSpecKind::UnexpectedNegativeInteger))?;
+    let offset: u64 = if force_zero_offset {
+        0
+    } else {
+        block
+            .offset
+            .try_into()
+            .map_err(|_| polars_err!(oos = OutOfSpecKind::UnexpectedNegativeInteger))?
+    };
 
     let length: u64 = block
         .meta_data_length
         .try_into()
         .map_err(|_| polars_err!(oos = OutOfSpecKind::UnexpectedNegativeInteger))?;
 
+    let message = get_message_from_block_offset(reader, offset, message_scratch)?;
+    let batch = get_dictionary_batch(&message)?;
+
     read_dictionary(
         batch,
         &metadata.schema,
@@ -132,6 +139,7 @@ pub fn read_file_dictionaries<R: Read + Seek>(
             reader,
             metadata,
             block,
+            false,
             &mut dictionaries,
             &mut message_scratch,
             scratch,
@@ -281,19 +289,19 @@ pub(crate) fn get_record_batch(
     }
 }
 
-fn get_message_from_block_offset<'a, R: Read + Seek>(
+pub fn get_message_from_block_offset<'a, R: Read + Seek>(
     reader: &mut R,
     offset: u64,
     message_scratch: &'a mut Vec<u8>,
 ) -> PolarsResult<arrow_format::ipc::MessageRef<'a>> {
-    // read length
     reader.seek(SeekFrom::Start(offset))?;
     let mut meta_buf = [0; 4];
     reader.read_exact(&mut meta_buf)?;
     if meta_buf == CONTINUATION_MARKER {
         // continuation marker encountered, read message next
         reader.read_exact(&mut meta_buf)?;
     }
+
     let meta_len = i32::from_le_bytes(meta_buf)
         .try_into()
         .map_err(|_| polars_err!(oos = OutOfSpecKind::UnexpectedNegativeInteger))?;
@@ -337,15 +345,21 @@ pub fn read_batch<R: Read + Seek>(
     projection: Option<&[usize]>,
     limit: Option<usize>,
     index: usize,
+    // When true, the reader object is handled as an IPC Block.
+    block_mode: bool,
     message_scratch: &mut Vec<u8>,
     data_scratch: &mut Vec<u8>,
 ) -> PolarsResult<RecordBatchT<Box<dyn Array>>> {
     let block = metadata.blocks[index];
 
-    let offset: u64 = block
-        .offset
-        .try_into()
-        .map_err(|_| polars_err!(oos = OutOfSpecKind::NegativeFooterLength))?;
+    let offset: u64 = if block_mode {
+        0
+    } else {
+        block
+            .offset
+            .try_into()
+            .map_err(|_| polars_err!(oos = OutOfSpecKind::NegativeFooterLength))?
+    };
 
     let length: u64 = block
         .meta_data_length

crates/polars-arrow/src/io/ipc/read/mod.rs

@@ -1,13 +1,15 @@
 //! APIs to read Arrow's IPC format.
 //!
-//! The two important structs here are the [`FileReader`](reader::FileReader),
+//! The two important File-based structs here are the [`FileReader`](reader::FileReader),
 //! which provides arbitrary access to any of its messages, and the
 //! [`StreamReader`](stream::StreamReader), which only supports reading
 //! data in the order it was written in.
+//! In addition, there is a Block-based struct [`BlockReader`](reader::BlockReader), which
+//! enabled random access to a standalone IPC Block.
 use crate::array::Array;
 
 mod array;
-mod common;
+pub mod common;
 mod deserialize;
 mod error;
 pub(crate) mod file;
@@ -23,10 +25,10 @@ pub use common::{ProjectionInfo, prepare_projection};
 pub use error::OutOfSpecKind;
 pub use file::{
     FileMetadata, deserialize_footer, get_row_count, get_row_count_from_blocks, read_batch,
-    read_file_dictionaries, read_file_metadata,
+    read_dictionary_block, read_file_dictionaries, read_file_metadata,
 };
 use polars_utils::aliases::PlHashMap;
-pub use reader::FileReader;
+pub use reader::{BlockReader, FileReader};
 pub use schema::deserialize_schema;
 pub use stream::{StreamMetadata, StreamReader, StreamState, read_stream_metadata};
 

crates/polars-arrow/src/io/ipc/read/reader.rs

@@ -3,7 +3,7 @@ use std::io::{Read, Seek};
 use polars_error::PolarsResult;
 
 use super::common::*;
-use super::file::{get_message_from_block, get_record_batch};
+use super::file::{get_message_from_block, get_message_from_block_offset, get_record_batch};
 use super::{Dictionaries, FileMetadata, read_batch, read_file_dictionaries};
 use crate::array::Array;
 use crate::datatypes::ArrowSchema;
@@ -114,7 +114,7 @@ impl<R: Read + Seek> FileReader<R> {
         (self.data_scratch, self.message_scratch) = scratches;
     }
 
-    fn read_dictionaries(&mut self) -> PolarsResult<()> {
+    pub fn read_dictionaries(&mut self) -> PolarsResult<()> {
         if self.dictionaries.is_none() {
             self.dictionaries = Some(read_file_dictionaries(
                 &mut self.reader,
@@ -187,6 +187,7 @@ impl<R: Read + Seek> Iterator for FileReader<R> {
             self.projection.as_ref().map(|x| x.columns.as_ref()),
             Some(self.remaining),
             block,
+            false,
             &mut self.message_scratch,
             &mut self.data_scratch,
         );
@@ -201,3 +202,27 @@ impl<R: Read + Seek> Iterator for FileReader<R> {
         Some(chunk)
     }
 }
+
+/// A reader that has access to exactly one standalone IPC Block of an Arrow IPC file.
+/// The block contains either a `RecordBatch` or a `DictionaryBatch`.
+/// The `dictionaries` field must be initialized prior to decoding a `RecordBatch`.
+pub struct BlockReader<R: Read + Seek> {
+    pub reader: R,
+}
+
+impl<R: Read + Seek> BlockReader<R> {
+    pub fn new(reader: R) -> Self {
+        Self { reader }
+    }
+
+    /// Reads the record batch header and returns its length (i.e., number of rows).
+    pub fn record_batch_num_rows(&mut self, message_scratch: &mut Vec<u8>) -> PolarsResult<usize> {
+        let offset: u64 = 0;
+
+        let message = get_message_from_block_offset(&mut self.reader, offset, message_scratch)?;
+        let batch = get_record_batch(message)?;
+
+        let out = batch.length().map(|l| usize::try_from(l).unwrap())?;
+        Ok(out)
+    }
+}