diff --git a/src/common/compressed_training_file_reader.rs b/src/common/compressed_training_file_reader.rs index 9c7111e..8e5d444 100644 --- a/src/common/compressed_training_file_reader.rs +++ b/src/common/compressed_training_file_reader.rs @@ -52,6 +52,14 @@ impl CompressedTrainingDataFileReader { Ok(data) } + pub fn read_next_chunk_into(&mut self, buffer: &mut Vec) -> Result<()> { + let header = self.read_chunk_header()?; + buffer.resize(header.chunk_size as usize, 0); + self.file.read_exact(buffer)?; + self.read_bytes += header.chunk_size as u64; + Ok(()) + } + fn read_chunk_header(&mut self) -> Result
{ let mut buf = [0u8; HEADER_SIZE]; diff --git a/src/reader/compressed_reader.rs b/src/reader/compressed_reader.rs index 0a263a1..e1b6812 100644 --- a/src/reader/compressed_reader.rs +++ b/src/reader/compressed_reader.rs @@ -97,10 +97,11 @@ impl CompressedTrainingDataEntryReader { reader.is_end = true; return Err(CompressedReaderError::EndOfFile); } else { - reader.chunk = match reader.input_file.as_mut().unwrap().read_next_chunk() { - Ok(chunk) => chunk, - Err(e) => return Err(CompressedReaderError::BinpackError(e)), - }; + reader + .input_file + .as_mut() + .unwrap() + .read_next_chunk_into(&mut reader.chunk)?; } Ok(reader)