scovich commented on code in PR #7834:
URL: https://github.com/apache/arrow-rs/pull/7834#discussion_r2193577119
##########
arrow-avro/src/reader/mod.rs:
##########
@@ -329,64 +334,41 @@ impl<R> Reader<R> {
pub fn avro_header(&self) -> &Header {
&self.header
}
-}
-impl<R: BufRead> Reader<R> {
/// Reads the next [`RecordBatch`] from the Avro file or `Ok(None)` on EOF
fn read(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
- if self.finished {
- return Ok(None);
- }
- loop {
- if !self.block_data.is_empty() {
- let consumed = self.decoder.decode(&self.block_data)?;
- if consumed > 0 {
- self.block_data.drain(..consumed);
- }
- match self.decoder.flush()? {
- None => {
- if !self.block_data.is_empty() {
- break;
- }
- }
- Some(batch) => {
- return Ok(Some(batch));
- }
- }
- }
- let maybe_block = {
+ 'outer: while !self.finished && !self.decoder.batch_is_full() {
+ while self.block_cursor == self.block_data.len() {
let buf = self.reader.fill_buf()?;
if buf.is_empty() {
- None
- } else {
- let read_len = buf.len();
- let consumed_len = self.block_decoder.decode(buf)?;
- self.reader.consume(consumed_len);
- if consumed_len == 0 && read_len != 0 {
- return Err(ArrowError::ParseError(
- "Could not decode next Avro block from partial
data".to_string(),
- ));
- }
- self.block_decoder.flush()
+ self.finished = true;
+ break 'outer;
}
- };
- match maybe_block {
- Some(block) => {
- let block_data = if let Some(ref codec) = self.compression
{
+ // Try to decode another block from the buffered reader.
+ let consumed = self.block_decoder.decode(buf)?;
+ self.reader.consume(consumed);
+ if let Some(block) = self.block_decoder.flush() {
+ // Successfully decoded a block.
+ let block_data = if let Some(ref codec) =
self.header.compression()? {
codec.decompress(&block.data)?
} else {
block.data
};
self.block_data = block_data;
+ self.block_cursor = 0;
+ } else if consumed == 0 {
+ // The block decoder made no progress on a non-empty
buffer.
+ return Err(ArrowError::ParseError(
+ "Could not decode next Avro block from partial
data".to_string(),
+ ));
}
- None => {
- self.finished = true;
- if !self.block_data.is_empty() {
- let consumed = self.decoder.decode(&self.block_data)?;
- self.block_data.drain(..consumed);
- }
- return self.decoder.flush();
- }
+ }
+ // Try to decode more rows from the current block.
+ let consumed =
self.decoder.decode(&self.block_data[self.block_cursor..])?;
+ if consumed == 0 && self.block_cursor < self.block_data.len() {
+ self.block_cursor = self.block_data.len();
Review Comment:
Thanks! Based on all that, it sounds like the current PR is mostly correct
in its handling of empty records... with the future improvement to return the
correct number of empty records before dropping the block? :shipit:
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
