Jefffrey commented on code in PR #11751:
URL: https://github.com/apache/arrow-rs/pull/11751#discussion_r4480386375
##########
arrow-avro/src/reader/async_reader/mod.rs:
##########
@@ -422,59 +416,43 @@ impl<R: AsyncFileReader + Unpin + 'static>
AsyncAvroFileReader<R> {
if matches!(self.block_decoder.state(),
BlockDecoderState::Finished) {
// We've already flushed, so if no batch was produced,
we are simply done.
- self.finishing_partial_block = false;
self.start_flushing();
continue;
}
- // If we've tried the following stage before, and still
can't decode,
- // this means the file is truncated or corrupted.
- if self.finishing_partial_block {
- return self.finish_with_error(AvroError::EOF(
- "Unexpected EOF while reading last Avro
block".into(),
- ));
- }
-
- // Avro splitting case: block is incomplete, we need to:
+ // Avro splitting case: the data ran out in a block this
range owns, we need to:
// 1. Parse the length so we know how much to read
// 2. Fetch more data from the reader
- // 3. Create a new block data from the remaining slice and
the newly fetched data
- // 4. Continue decoding until end of block
- self.finishing_partial_block = true;
-
- // Mid-block, but we don't know how many bytes are missing
yet
- if matches!(
- self.block_decoder.state(),
- BlockDecoderState::Count | BlockDecoderState::Size
- ) {
- // Max VLQ header is 20 bytes (10 bytes each for count
and size).
- // Fetch just enough to complete it.
- const MAX_VLQ_HEADER_SIZE: u64 = 20;
- let fetch_end = (self.range.end +
MAX_VLQ_HEADER_SIZE).min(self.file_size);
-
- // If there is nothing more to fetch, error out
- if fetch_end == self.range.end {
- return self.finish_with_error(AvroError::EOF(
- "Unexpected EOF while reading Avro block
header".into(),
- ));
+ // 3. Continue decoding until end of block
+ // Every fetch extends `range.end`, so this terminates at
the end of the file.
+ let range_to_fetch = match self.block_decoder.state() {
+ // Mid-block, but we don't know how many bytes are
missing yet
+ BlockDecoderState::Count | BlockDecoderState::Size => {
+ // Max VLQ header is 20 bytes (10 bytes each for
count and size).
+ // Fetch just enough to complete it.
+ const MAX_VLQ_HEADER_SIZE: u64 = 20;
+ let fetch_end =
+ (self.range.end +
MAX_VLQ_HEADER_SIZE).min(self.file_size);
+
+ if fetch_end == self.range.end {
+ if matches!(self.block_decoder.state(),
BlockDecoderState::Size) {
Review Comment:
here we narrow to only size state, instead of before where either count or
size state would cause EOF; is this intentional?
here's an example test mocked by Codex which fails on this PR:
```rust
#[tokio::test]
async fn test_truncated_block_count_errors() {
use crate::writer::AvroWriter;
let schema = Schema::new(vec![Field::new("id", DataType::Int32,
false)]);
let batch = RecordBatch::try_new(
Arc::new(schema.clone()),
vec![Arc::new(Int32Array::from(vec![1]))],
)
.unwrap();
let mut bytes = Vec::new();
let mut writer = AvroWriter::new(&mut bytes, schema).unwrap();
writer.write(&batch).unwrap();
writer.finish().unwrap();
// Start the next block's count varint without its terminating byte.
bytes.push(0x80);
let file_size = bytes.len() as u64;
let reader =
AsyncAvroFileReader::builder(std::io::Cursor::new(bytes), file_size, 1024)
.try_build()
.await
.unwrap();
let result: Result<Vec<RecordBatch>, _> = reader.try_collect().await;
assert!(
result.is_err(),
"accepted truncated block count: {result:?}"
);
}
```
it seems to suggest an EOF error should be occurring, but with changes here
this isnt happening
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
