Avoid zero-filling IPC reads with typed buffer handling

Author: pchintar

arrow-ipc/src/reader.rs

@@ -72,6 +72,57 @@ fn read_buffer(
         }
     }
 }
+
+/// Source for IPC body buffers.
+///
+/// Most decode paths use an already materialized [`Buffer`] and slice into it
+/// using the offsets from the IPC metadata. Keeping this behind a small helper
+/// lets typed buffer reads share the same bounds handling as regular buffer
+/// reads.
+enum IpcBufferSource<'a> {
+    Buffer(&'a Buffer),
+}
+
+impl<'a> IpcBufferSource<'a> {
+    fn read_buffer(
+        &self,
+        buf: &crate::Buffer,
+        compression: Option<CompressionCodec>,
+        decompression_context: &mut DecompressionContext,
+    ) -> Result<Buffer, ArrowError> {
+        match self {
+            Self::Buffer(data) => read_buffer(buf, data, compression, decompression_context),
+        }
+    }
+    /// Reads a physical IPC buffer that is expected to contain `len` values of
+    /// type `T`.
+    ///
+    /// The returned value is still a [`Buffer`], not a [`ScalarBuffer<T>`]. This
+    /// preserves the existing alignment behavior: properly aligned buffers remain
+    /// zero-copy, while unaligned buffers are still handled later by
+    /// `ArrayDataBuilder::align_buffers` according to `require_alignment`.
+    fn read_typed_buffer<T: ArrowNativeType>(
+        &self,
+        buf: &crate::Buffer,
+        len: usize,
+        compression: Option<CompressionCodec>,
+        decompression_context: &mut DecompressionContext,
+    ) -> Result<Buffer, ArrowError> {
+        let byte_len = len
+            .checked_mul(std::mem::size_of::<T>())
+            .ok_or_else(|| ArrowError::IpcError("Buffer length overflow".to_string()))?;
+
+        let buffer = self.read_buffer(buf, compression, decompression_context)?;
+        // Some invalid or legacy IPC inputs may contain shorter buffers than
+        // implied by the schema. Preserve the existing behavior and let array
+        // construction/validation report the error.
+        if buffer.len() < byte_len {
+            return Ok(buffer);
+        }
+        Ok(buffer.slice_with_length(0, byte_len))
+    }
+}
+
 impl RecordBatchDecoder<'_> {
     /// Coordinates reading arrays based on data types.
     ///
@@ -93,25 +144,46 @@ impl RecordBatchDecoder<'_> {
         let data_type = field.data_type();
         match data_type {
             Utf8 | Binary | LargeBinary | LargeUtf8 => {
+                // Binary and string arrays use fixed-width offset buffers
+                // followed by the variable-width value bytes buffer.
+                // Read the offsets through the next_typed_buffer::<T>() helper.
                 let field_node = self.next_node(field)?;
+                let null_buffer = self.next_buffer()?;
+                let len = field_node.length() as usize + 1;
+                let offsets = match data_type {
+                    Utf8 | Binary => self.next_typed_buffer::<i32>(len)?,
+                    LargeBinary | LargeUtf8 => self.next_typed_buffer::<i64>(len)?,
+                    _ => unreachable!(),
+                };
+
                 let buffers = [
-                    self.next_buffer()?,
-                    self.next_buffer()?,
-                    self.next_buffer()?,
+                    null_buffer,
+                    offsets,
+                    self.next_buffer()?, // value bytes
                 ];
+
                 self.create_primitive_array(field_node, data_type, &buffers)
             }
+            // The first buffer after the null bitmap is the fixed-width view
+            // buffer. Any remaining buffers are variadic data buffers.
             BinaryView | Utf8View => {
                 let count = variadic_counts
                     .pop_front()
                     .ok_or(ArrowError::IpcError(format!(
                         "Missing variadic count for {data_type} column"
                     )))?;
-                let count = count + 2; // view and null buffer.
-                let buffers = (0..count)
-                    .map(|_| self.next_buffer())
-                    .collect::<Result<Vec<_>, _>>()?;
+
                 let field_node = self.next_node(field)?;
+                let len = field_node.length() as usize;
+
+                let mut buffers = Vec::with_capacity(count as usize + 2);
+                buffers.push(self.next_buffer()?); // null buffer
+                buffers.push(self.next_typed_buffer::<u128>(len)?); // views
+
+                for _ in 0..count {
+                    buffers.push(self.next_buffer()?); // variadic data buffers
+                }
+
                 self.create_primitive_array(field_node, data_type, &buffers)
             }
             FixedSizeBinary(_) => {
@@ -121,17 +193,37 @@ impl RecordBatchDecoder<'_> {
             }
             List(list_field) | LargeList(list_field) | Map(list_field, _) => {
                 let list_node = self.next_node(field)?;
-                let list_buffers = [self.next_buffer()?, self.next_buffer()?];
+                let null_buffer = self.next_buffer()?;
+
+                let offset_len = list_node.length() as usize + 1;
+                let offsets = match data_type {
+                    List(_) | Map(_, _) => self.next_typed_buffer::<i32>(offset_len)?,
+                    LargeList(_) => self.next_typed_buffer::<i64>(offset_len)?,
+                    _ => unreachable!(),
+                };
+
+                let list_buffers = [null_buffer, offsets];
                 let values = self.create_array(list_field, variadic_counts)?;
                 self.create_list_array(list_node, data_type, &list_buffers, values)
             }
             ListView(list_field) | LargeListView(list_field) => {
                 let list_node = self.next_node(field)?;
-                let list_buffers = [
-                    self.next_buffer()?, // null buffer
-                    self.next_buffer()?, // offsets
-                    self.next_buffer()?, // sizes
-                ];
+                let null_buffer = self.next_buffer()?;
+
+                let len = list_node.length() as usize;
+                let (offsets, sizes) = match data_type {
+                    ListView(_) => (
+                        self.next_typed_buffer::<i32>(len)?,
+                        self.next_typed_buffer::<i32>(len)?,
+                    ),
+                    LargeListView(_) => (
+                        self.next_typed_buffer::<i64>(len)?,
+                        self.next_typed_buffer::<i64>(len)?,
+                    ),
+                    _ => unreachable!(),
+                };
+
+                let list_buffers = [null_buffer, offsets, sizes];
                 let values = self.create_array(list_field, variadic_counts)?;
                 self.create_list_view_array(list_node, data_type, &list_buffers, values)
             }
@@ -170,10 +262,29 @@ impl RecordBatchDecoder<'_> {
 
                 self.create_array_from_builder(builder)
             }
-            // Create dictionary array from RecordBatch
             Dictionary(_, _) => {
                 let index_node = self.next_node(field)?;
-                let index_buffers = [self.next_buffer()?, self.next_buffer()?];
+                let null_buffer = self.next_buffer()?;
+
+                // Dictionary indices are fixed-width values. Read the index
+                // buffer through the next_typed_buffer::<T>() helper so length handling is
+                // based on the physical key type.
+                let len = index_node.length() as usize;
+                let indices = match data_type {
+                    Dictionary(key_type, _) => match key_type.as_ref() {
+                        Int8 => self.next_typed_buffer::<i8>(len)?,
+                        Int16 => self.next_typed_buffer::<i16>(len)?,
+                        Int32 => self.next_typed_buffer::<i32>(len)?,
+                        Int64 => self.next_typed_buffer::<i64>(len)?,
+                        UInt8 => self.next_typed_buffer::<u8>(len)?,
+                        UInt16 => self.next_typed_buffer::<u16>(len)?,
+                        UInt32 => self.next_typed_buffer::<u32>(len)?,
+                        UInt64 => self.next_typed_buffer::<u64>(len)?,
+                        t => unreachable!("Unsupported dictionary key type {t:?}"),
+                    },
+                    _ => unreachable!(),
+                };
+                let index_buffers = [null_buffer, indices];
 
                 #[allow(deprecated)]
                 let dict_id = field.dict_id().ok_or_else(|| {
@@ -206,13 +317,13 @@ impl RecordBatchDecoder<'_> {
                     self.next_buffer()?;
                 }
 
-                let type_ids: ScalarBuffer<i8> =
-                    self.next_buffer()?.slice_with_length(0, len).into();
-
+                // Union type ids and dense union offsets are fixed-width
+                // buffers. Read them through next_typed_buffer::<T>() before
+                // constructing the ScalarBuffers used by UnionArray.
+                let type_ids: ScalarBuffer<i8> = self.next_typed_buffer::<i8>(len)?.into();
                 let value_offsets = match mode {
                     UnionMode::Dense => {
-                        let offsets: ScalarBuffer<i32> =
-                            self.next_buffer()?.slice_with_length(0, len * 4).into();
+                        let offsets: ScalarBuffer<i32> = self.next_typed_buffer::<i32>(len)?.into();
                         Some(offsets)
                     }
                     UnionMode::Sparse => None,
@@ -253,7 +364,41 @@ impl RecordBatchDecoder<'_> {
             }
             _ => {
                 let field_node = self.next_node(field)?;
-                let buffers = [self.next_buffer()?, self.next_buffer()?];
+                let null_buffer = self.next_buffer()?;
+
+                // Primitive and primitive-like arrays use fixed-width physical
+                // buffers with widths that depend on the logical data type.
+                let len = field_node.length() as usize;
+                let values = match data_type {
+                    Int8 => self.next_typed_buffer::<i8>(len)?,
+                    Int16 => self.next_typed_buffer::<i16>(len)?,
+                    Int32 => self.next_typed_buffer::<i32>(len)?,
+                    Int64 => self.next_typed_buffer::<i64>(len)?,
+                    UInt8 => self.next_typed_buffer::<u8>(len)?,
+                    UInt16 => self.next_typed_buffer::<u16>(len)?,
+                    UInt32 => self.next_typed_buffer::<u32>(len)?,
+                    UInt64 => self.next_typed_buffer::<u64>(len)?,
+                    Float16 => self.next_typed_buffer::<u16>(len)?,
+                    Float32 => self.next_typed_buffer::<f32>(len)?,
+                    Float64 => self.next_typed_buffer::<f64>(len)?,
+                    Date32 | Time32(_) | Interval(IntervalUnit::YearMonth) | Decimal32(_, _) => {
+                        self.next_typed_buffer::<i32>(len)?
+                    }
+                    Date64 | Time64(_) | Timestamp(_, _) | Duration(_) | Decimal64(_, _) => {
+                        self.next_typed_buffer::<i64>(len)?
+                    }
+                    Decimal128(_, _) => self.next_typed_buffer::<i128>(len)?,
+                    Decimal256(_, _) => self.next_typed_buffer::<arrow_buffer::i256>(len)?,
+                    Interval(IntervalUnit::DayTime) => {
+                        self.next_typed_buffer::<arrow_buffer::IntervalDayTime>(len)?
+                    }
+                    Interval(IntervalUnit::MonthDayNano) => {
+                        self.next_typed_buffer::<arrow_buffer::IntervalMonthDayNano>(len)?
+                    }
+                    Boolean | FixedSizeBinary(_) => self.next_buffer()?,
+                    t => unreachable!("Unsupported primitive data type {t:?}"),
+                };
+                let buffers = [null_buffer, values];
                 self.create_primitive_array(field_node, data_type, &buffers)
             }
         }
@@ -454,8 +599,8 @@ pub struct RecordBatchDecoder<'a> {
     decompression_context: DecompressionContext,
     /// The format version
     version: MetadataVersion,
-    /// The raw data buffer
-    data: &'a Buffer,
+    /// Source of IPC body buffers
+    data: IpcBufferSource<'a>,
     /// The fields comprising this array
     nodes: VectorIter<'a, FieldNode>,
     /// The buffers comprising this array
@@ -500,7 +645,7 @@ impl<'a> RecordBatchDecoder<'a> {
             compression,
             decompression_context: DecompressionContext::new(),
             version: *metadata,
-            data: buf,
+            data: IpcBufferSource::Buffer(buf),
             nodes: field_nodes.iter(),
             buffers: buffers.iter(),
             projection: None,
@@ -615,9 +760,22 @@ impl<'a> RecordBatchDecoder<'a> {
         let buffer = self.buffers.next().ok_or_else(|| {
             ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
         })?;
-        read_buffer(
+        self.data
+            .read_buffer(buffer, self.compression, &mut self.decompression_context)
+    }
+    /// Advances to the next IPC buffer and trims it to the expected physical
+    /// length for `len` values of `T`.
+    ///
+    /// This keeps typed buffer length handling in one place while leaving final
+    /// array construction on the existing `ArrayDataBuilder` path.
+    fn next_typed_buffer<T: ArrowNativeType>(&mut self, len: usize) -> Result<Buffer, ArrowError> {
+        let buffer = self.buffers.next().ok_or_else(|| {
+            ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
+        })?;
+
+        self.data.read_typed_buffer::<T>(
             buffer,
-            self.data,
+            len,
             self.compression,
             &mut self.decompression_context,
         )
@@ -902,16 +1060,29 @@ fn get_dictionary_values(
     Ok(dictionary_values)
 }
 
-/// Read the data for a given block
+/// Read the data for a given IPC file block.
+///
+/// The returned buffer is fully initialized by the reader. This avoids first
+/// zero-filling the allocation and then immediately overwriting it with block
+/// data.
 fn read_block<R: Read + Seek>(mut reader: R, block: &Block) -> Result<Buffer, ArrowError> {
     reader.seek(SeekFrom::Start(block.offset() as u64))?;
     let body_len = block.bodyLength().to_usize().unwrap();
     let metadata_len = block.metaDataLength().to_usize().unwrap();
     let total_len = body_len.checked_add(metadata_len).unwrap();
 
-    let mut buf = MutableBuffer::from_len_zeroed(total_len);
-    reader.read_exact(&mut buf)?;
-    Ok(buf.into())
+    let mut buf = Vec::with_capacity(total_len);
+    reader
+        .by_ref()
+        .take(total_len as u64)
+        .read_to_end(&mut buf)?;
+    if buf.len() != total_len {
+        return Err(ArrowError::IpcError(format!(
+            "Expected IPC block of length {total_len}, got {}",
+            buf.len()
+        )));
+    }
+    Ok(Buffer::from_vec(buf))
 }
 
 /// Parse an encapsulated message
@@ -1809,16 +1980,14 @@ impl<R: Read> MessageReader<R> {
         }
     }
 
-    /// Reads the entire next message from the underlying reader which includes
-    /// the metadata length, the metadata, and the body.
+    /// Reads the next IPC message, including the metadata and body.
     ///
     /// # Returns
-    /// - `Ok(None)` if the the reader signals the end of stream with EOF on
-    ///   the first read
-    /// - `Err(_)` if the reader returns an error other than on the first
-    ///   read, or if the metadata length is invalid
-    /// - `Ok(Some(_))` with the Message and buffer containiner the
-    ///   body bytes otherwise.
+    /// - `Ok(None)` if the reader signals end-of-stream before reading a
+    ///   metadata length
+    /// - `Err(_)` if the reader returns an error or the IPC message is invalid
+    /// - `Ok(Some(_))` with the decoded message metadata and body bytes
+    ///   otherwise
     fn maybe_next(&mut self) -> Result<Option<(Message::Message<'_>, MutableBuffer)>, ArrowError> {
         let meta_len = self.read_meta_len()?;
         let Some(meta_len) = meta_len else {
@@ -1832,10 +2001,23 @@ impl<R: Read> MessageReader<R> {
             ArrowError::ParseError(format!("Unable to get root as message: {err:?}"))
         })?;
 
-        let mut buf = MutableBuffer::from_len_zeroed(message.bodyLength() as usize);
-        self.reader.read_exact(&mut buf)?;
-
-        Ok(Some((message, buf)))
+        // The message body is fully provided by the reader. Read it into a
+        // Vec to avoid zero-filling a MutableBuffer before immediately
+        // overwriting it with the body bytes.
+        let body_len = message.bodyLength() as usize;
+        let mut buf = Vec::with_capacity(body_len);
+        self.reader
+            .by_ref()
+            .take(body_len as u64)
+            .read_to_end(&mut buf)?;
+
+        if buf.len() != body_len {
+            return Err(ArrowError::IpcError(format!(
+                "Expected IPC message body of length {body_len}, got {}",
+                buf.len()
+            )));
+        }
+        Ok(Some((message, MutableBuffer::from(buf))))
     }
 
     /// Get a mutable reference to the underlying reader.