diff --git a/crates/polars-arrow/src/io/ipc/read/file.rs b/crates/polars-arrow/src/io/ipc/read/file.rs index 0eb60d0a566c..c16a19560613 100644 --- a/crates/polars-arrow/src/io/ipc/read/file.rs +++ b/crates/polars-arrow/src/io/ipc/read/file.rs @@ -36,32 +36,6 @@ pub struct FileMetadata { pub size: u64, } -fn read_dictionary_message( - reader: &mut R, - offset: u64, - data: &mut Vec, -) -> PolarsResult<()> { - let mut message_size: [u8; 4] = [0; 4]; - reader.seek(SeekFrom::Start(offset))?; - reader.read_exact(&mut message_size)?; - if message_size == CONTINUATION_MARKER { - reader.read_exact(&mut message_size)?; - }; - let message_length = i32::from_le_bytes(message_size); - - let message_length: usize = message_length - .try_into() - .map_err(|_| polars_err!(oos = OutOfSpecKind::NegativeFooterLength))?; - - data.clear(); - data.try_reserve(message_length)?; - reader - .by_ref() - .take(message_length as u64) - .read_to_end(data)?; - - Ok(()) -} /// Read the row count by summing the length of the of the record batches pub fn get_row_count(reader: &mut R) -> PolarsResult { let mut message_scratch: Vec = Default::default(); @@ -72,7 +46,7 @@ pub fn get_row_count(reader: &mut R) -> PolarsResult { blocks .into_iter() .map(|block| { - let message = get_message_from_block(reader, block, &mut message_scratch)?; + let message = get_message_from_block(reader, &block, &mut message_scratch)?; let record_batch = get_record_batch(message)?; record_batch.length().map_err(|e| e.into()) }) @@ -100,20 +74,18 @@ fn read_dictionary_block( message_scratch: &mut Vec, dictionary_scratch: &mut Vec, ) -> 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 length: u64 = block .meta_data_length .try_into() .map_err(|_| polars_err!(oos = OutOfSpecKind::UnexpectedNegativeInteger))?; - read_dictionary_message(reader, offset, message_scratch)?; - - let message = arrow_format::ipc::MessageRef::read_as_root(message_scratch.as_ref()) - .map_err(|err| polars_err!(oos = OutOfSpecKind::InvalidFlatbufferMessage(err)))?; - - let batch = get_dictionary_batch(&message)?; read_dictionary( batch, @@ -299,7 +271,7 @@ fn get_message_from_block_offset<'a, R: Read + Seek>( fn get_message_from_block<'a, R: Read + Seek>( reader: &mut R, - block: arrow_format::ipc::Block, + block: &arrow_format::ipc::Block, message_scratch: &'a mut Vec, ) -> PolarsResult> { let offset: u64 = block