From 35f8fdd2793c5e149cca1ec113deaff01540427a Mon Sep 17 00:00:00 2001 From: Igor Artamonov Date: Mon, 6 Apr 2026 00:51:17 +0100 Subject: [PATCH 1/4] feat: seekable reader --- avro/src/error.rs | 3 + avro/src/lib.rs | 2 +- avro/src/reader/block.rs | 99 ++++++++++++++++++++++++-- avro/src/reader/mod.rs | 147 ++++++++++++++++++++++++++++++++++++++- 4 files changed, 243 insertions(+), 8 deletions(-) diff --git a/avro/src/error.rs b/avro/src/error.rs index 2acb217b..1ad79f39 100644 --- a/avro/src/error.rs +++ b/avro/src/error.rs @@ -554,6 +554,9 @@ pub enum Details { #[error("Failed to read block marker bytes: {0}")] ReadBlockMarker(#[source] std::io::Error), + #[error("Failed to seek to block: {0}")] + SeekToBlock(#[source] std::io::Error), + #[error("Read into buffer failed: {0}")] ReadIntoBuf(#[source] std::io::Error), diff --git a/avro/src/lib.rs b/avro/src/lib.rs index 3f509897..586c6687 100644 --- a/avro/src/lib.rs +++ b/avro/src/lib.rs @@ -99,7 +99,7 @@ pub use error::Error; reason = "Still need to export it until we remove it completely" )] pub use reader::{ - Reader, + BlockPosition, Reader, datum::{from_avro_datum, from_avro_datum_reader_schemata, from_avro_datum_schemata}, read_marker, single_object::{GenericSingleObjectReader, SpecificSingleObjectReader}, diff --git a/avro/src/reader/block.rs b/avro/src/reader/block.rs index 5e9ae1b1..12fabfa4 100644 --- a/avro/src/reader/block.rs +++ b/avro/src/reader/block.rs @@ -17,7 +17,7 @@ use std::{ collections::HashMap, - io::{ErrorKind, Read}, + io::{ErrorKind, Read, Seek, SeekFrom}, str::FromStr, }; @@ -35,10 +35,57 @@ use crate::{ util, }; +/// Position and size of a single Avro data block within the file stream. +/// +/// Captured automatically as blocks are read during forward iteration. +/// Use with [`super::Reader::seek_to_block`] to jump back to a previously-read block. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BlockPosition { + /// Byte offset in the stream where this block starts (before the object-count varint). + pub offset: u64, + /// Total number of records in this block. + pub message_count: usize, +} + +/// Wraps an inner reader and tracks the current byte position. +/// +/// Avoids requiring `Seek` just to know how many bytes have been consumed. +/// When the inner reader also implements `Seek`, seeking updates the tracked position. +#[derive(Debug, Clone)] +struct PositionTracker { + inner: R, + pos: u64, +} + +impl PositionTracker { + fn new(inner: R) -> Self { + Self { inner, pos: 0 } + } + + fn position(&self) -> u64 { + self.pos + } +} + +impl Read for PositionTracker { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let n = self.inner.read(buf)?; + self.pos += n as u64; + Ok(n) + } +} + +impl Seek for PositionTracker { + fn seek(&mut self, from: SeekFrom) -> std::io::Result { + self.pos = self.inner.seek(from)?; + Ok(self.pos) + } +} + /// Internal Block reader. #[derive(Debug, Clone)] pub(super) struct Block<'r, R> { - reader: R, + reader: PositionTracker, /// Internal buffering to reduce allocation. buf: Vec, buf_idx: usize, @@ -51,6 +98,10 @@ pub(super) struct Block<'r, R> { pub(super) user_metadata: HashMap>, names_refs: Names, human_readable: bool, + /// Byte offset where data blocks begin (right after header and sync marker). + pub(super) data_start: u64, + /// Position and record count of the currently loaded block. + pub(super) current_block_info: Option, } impl<'r, R: Read> Block<'r, R> { @@ -60,7 +111,7 @@ impl<'r, R: Read> Block<'r, R> { human_readable: bool, ) -> AvroResult> { let mut block = Block { - reader, + reader: PositionTracker::new(reader), codec: Codec::Null, writer_schema: Schema::Null, schemata, @@ -71,9 +122,12 @@ impl<'r, R: Read> Block<'r, R> { user_metadata: Default::default(), names_refs: Default::default(), human_readable, + data_start: 0, + current_block_info: None, }; block.read_header()?; + block.data_start = block.reader.position(); Ok(block) } @@ -142,6 +196,7 @@ impl<'r, R: Read> Block<'r, R> { /// the block. The objects are stored in an internal buffer to the `Reader`. fn read_block_next(&mut self) -> AvroResult<()> { assert!(self.is_empty(), "Expected self to be empty!"); + let block_start = self.reader.position(); match util::read_usize(&mut self.reader).map_err(Error::into_details) { Ok(block_len) => { self.message_count = block_len; @@ -156,6 +211,11 @@ impl<'r, R: Read> Block<'r, R> { return Err(Details::GetBlockMarker.into()); } + self.current_block_info = Some(BlockPosition { + offset: block_start, + message_count: block_len, + }); + // NOTE (JAB): This doesn't fit this Reader pattern very well. // `self.buf` is a growable buffer that is reused as the reader is iterated. // For non `Codec::Null` variants, `decompress` will allocate a new `Vec` @@ -295,6 +355,23 @@ impl<'r, R: Read> Block<'r, R> { } } +impl Block<'_, R> { + /// Seek the underlying stream to `offset` and read the block there. + /// Validates the sync marker to confirm it's a real block boundary. + pub(super) fn seek_to_block(&mut self, offset: u64) -> AvroResult<()> { + self.reader + .seek(SeekFrom::Start(offset)) + .map_err(Details::SeekToBlock)?; + + self.buf.clear(); + self.buf_idx = 0; + self.message_count = 0; + self.current_block_info = None; + + self.read_block_next() + } +} + fn read_codec(metadata: &HashMap) -> AvroResult { let result = metadata .get("avro.codec") @@ -356,7 +433,7 @@ fn read_codec(metadata: &HashMap) -> AvroResult { #[cfg(test)] mod tests { - use super::Block; + use super::{Block, PositionTracker}; use crate::error::Details; use crate::{Codec, Schema}; @@ -364,7 +441,10 @@ mod tests { fn avro_rs_586_negative_block_size() { let mut block = Block::<'_, &[u8]> { // Block header with an object count of 1 and a block size of -1 - reader: &[0x02, 0x01], + reader: PositionTracker { + inner: &[0x02, 0x01], + pos: 0, + }, buf: vec![], buf_idx: 0, message_count: 0, @@ -375,6 +455,8 @@ mod tests { user_metadata: Default::default(), names_refs: Default::default(), human_readable: false, + data_start: 0, + current_block_info: None, }; let err = block.read_block_next().unwrap_err().into_details(); @@ -385,7 +467,10 @@ mod tests { fn avro_rs_586_negative_block_len() { let mut block = Block::<'_, &[u8]> { // Block header with an object count of -1 and a block size of 0 - reader: &[0x01, 0x00], + reader: PositionTracker { + inner: &[0x01, 0x00], + pos: 0, + }, buf: vec![], buf_idx: 0, message_count: 0, @@ -396,6 +481,8 @@ mod tests { user_metadata: Default::default(), names_refs: Default::default(), human_readable: false, + data_start: 0, + current_block_info: None, }; let err = block.read_block_next().unwrap_err().into_details(); diff --git a/avro/src/reader/mod.rs b/avro/src/reader/mod.rs index 7002fb86..802f34c6 100644 --- a/avro/src/reader/mod.rs +++ b/avro/src/reader/mod.rs @@ -21,9 +21,14 @@ mod block; pub mod datum; pub mod single_object; -use std::{collections::HashMap, io::Read, marker::PhantomData}; +use std::{ + collections::HashMap, + io::{Read, Seek}, + marker::PhantomData, +}; use block::Block; +pub use block::BlockPosition; use bon::bon; use serde::de::DeserializeOwned; @@ -159,6 +164,39 @@ impl Iterator for Reader<'_, R> { } } +impl Reader<'_, R> { + /// The currently loaded block's position and record count. + /// + /// Returns `None` before the first record has been read. + /// Updated automatically each time a new block is loaded during iteration. + pub fn current_block(&self) -> Option { + self.block.current_block_info + } + + /// Byte offset where data blocks begin (right after the file header). + /// + /// This is the offset of the first data block — equivalent to the position + /// that would be returned by `current_block().offset` for block 0. + pub fn data_start(&self) -> u64 { + self.block.data_start + } + + /// Seek to the data block at the given byte offset and load it. + /// + /// The offset must point to the start of a valid data block (before its + /// object-count varint). The block is read, decompressed, and its sync + /// marker is validated against the file header. After this call, [`Iterator::next`] + /// yields the first record in that block. + /// + /// Typically the caller saves offsets from [`current_block`](Self::current_block) + /// during forward iteration and later passes them here to jump back. + pub fn seek_to_block(&mut self, offset: u64) -> AvroResult<()> { + self.block.seek_to_block(offset)?; + self.errored = false; + Ok(()) + } +} + /// Wrapper around [`Reader`] where the iterator deserializes `T`. pub struct ReaderDeser<'a, R, T> { inner: Reader<'a, R>, @@ -366,4 +404,111 @@ mod tests { panic!("Expected an error in the reading of the codec!"); } } + + /// Write an Avro file with multiple blocks and verify we can seek between them. + #[test] + fn test_seek_to_block() -> TestResult { + use crate::writer::Writer; + + let schema = Schema::parse_str(SCHEMA)?; + let mut writer = Writer::new(&schema, Vec::new())?; + + // Block 0: records with a=10, a=20 + let mut r = Record::new(&schema).unwrap(); + r.put("a", 10i64); + r.put("b", "b0_r0"); + writer.append_value(r)?; + let mut r = Record::new(&schema).unwrap(); + r.put("a", 20i64); + r.put("b", "b0_r1"); + writer.append_value(r)?; + writer.flush()?; + + // Block 1: records with a=30, a=40 + let mut r = Record::new(&schema).unwrap(); + r.put("a", 30i64); + r.put("b", "b1_r0"); + writer.append_value(r)?; + let mut r = Record::new(&schema).unwrap(); + r.put("a", 40i64); + r.put("b", "b1_r1"); + writer.append_value(r)?; + writer.flush()?; + + // Block 2: records with a=50 + let mut r = Record::new(&schema).unwrap(); + r.put("a", 50i64); + r.put("b", "b2_r0"); + writer.append_value(r)?; + writer.flush()?; + + let data = writer.into_inner()?; + + // Read forward and collect block positions + let mut reader = Reader::new(Cursor::new(&data))?; + let mut block_offsets: Vec = Vec::new(); + let mut all_values: Vec = Vec::new(); + + assert!(reader.current_block().is_none()); + + while let Some(value) = reader.next() { + all_values.push(value?); + let pos = reader.current_block().expect("block info after read"); + if block_offsets + .last() + .is_none_or(|last| last.offset != pos.offset) + { + block_offsets.push(pos); + } + } + + assert_eq!(all_values.len(), 5); + assert_eq!(block_offsets.len(), 3); + assert_eq!(block_offsets[0].message_count, 2); + assert_eq!(block_offsets[1].message_count, 2); + assert_eq!(block_offsets[2].message_count, 1); + assert_eq!(reader.data_start(), block_offsets[0].offset); + + // Seek back to block 1 and read its records + reader.seek_to_block(block_offsets[1].offset)?; + let v1 = reader.next().unwrap()?; + assert_eq!(v1, all_values[2]); + let v2 = reader.next().unwrap()?; + assert_eq!(v2, all_values[3]); + + // Seek back to block 0 + reader.seek_to_block(block_offsets[0].offset)?; + let v0 = reader.next().unwrap()?; + assert_eq!(v0, all_values[0]); + + // Seek to block 2 + reader.seek_to_block(block_offsets[2].offset)?; + let v4 = reader.next().unwrap()?; + assert_eq!(v4, all_values[4]); + + assert!(reader.next().is_none()); + + Ok(()) + } + + /// Seeking to an invalid offset should fail with a sync marker error. + #[test] + fn test_seek_to_invalid_offset() -> TestResult { + use crate::writer::Writer; + + let schema = Schema::parse_str(SCHEMA)?; + let mut writer = Writer::new(&schema, Vec::new())?; + let mut r = Record::new(&schema).unwrap(); + r.put("a", 1i64); + r.put("b", "x"); + writer.append_value(r)?; + writer.flush()?; + let data = writer.into_inner()?; + + let mut reader = Reader::new(Cursor::new(&data))?; + let result = reader.seek_to_block(7); + assert!(result.is_err()); + + Ok(()) + } } From 10f82da1e375813675574dc4e37f0cde6e399688 Mon Sep 17 00:00:00 2001 From: Igor Artamonov Date: Mon, 6 Apr 2026 01:14:01 +0100 Subject: [PATCH 2/4] after code review --- avro/src/reader/block.rs | 17 +++++++++++++++-- avro/src/reader/mod.rs | 8 +++++--- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/avro/src/reader/block.rs b/avro/src/reader/block.rs index 12fabfa4..ccbf64d7 100644 --- a/avro/src/reader/block.rs +++ b/avro/src/reader/block.rs @@ -35,7 +35,7 @@ use crate::{ util, }; -/// Position and size of a single Avro data block within the file stream. +/// Byte offset and record count of a single Avro data block. /// /// Captured automatically as blocks are read during forward iteration. /// Use with [`super::Reader::seek_to_block`] to jump back to a previously-read block. @@ -358,6 +358,8 @@ impl<'r, R: Read> Block<'r, R> { impl Block<'_, R> { /// Seek the underlying stream to `offset` and read the block there. /// Validates the sync marker to confirm it's a real block boundary. + /// Returns an error if no valid block can be read at the offset + /// (e.g., the offset is at or past EOF). pub(super) fn seek_to_block(&mut self, offset: u64) -> AvroResult<()> { self.reader .seek(SeekFrom::Start(offset)) @@ -368,7 +370,18 @@ impl Block<'_, R> { self.message_count = 0; self.current_block_info = None; - self.read_block_next() + // read_block_next treats UnexpectedEof as a clean end-of-stream + // (returns Ok with message_count=0). That's correct for forward + // iteration but wrong here — the caller asked for a specific block. + self.read_block_next()?; + if self.is_empty() { + return Err(Details::SeekToBlock(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + format!("no block at offset {offset}"), + )) + .into()); + } + Ok(()) } } diff --git a/avro/src/reader/mod.rs b/avro/src/reader/mod.rs index 802f34c6..5e5946b0 100644 --- a/avro/src/reader/mod.rs +++ b/avro/src/reader/mod.rs @@ -164,11 +164,11 @@ impl Iterator for Reader<'_, R> { } } -impl Reader<'_, R> { +impl Reader<'_, R> { /// The currently loaded block's position and record count. /// - /// Returns `None` before the first record has been read. - /// Updated automatically each time a new block is loaded during iteration. + /// Returns `None` only before the first block is loaded (via iteration or + /// [`seek_to_block`](Self::seek_to_block)). Always `Some` afterward. pub fn current_block(&self) -> Option { self.block.current_block_info } @@ -180,7 +180,9 @@ impl Reader<'_, R> { pub fn data_start(&self) -> u64 { self.block.data_start } +} +impl Reader<'_, R> { /// Seek to the data block at the given byte offset and load it. /// /// The offset must point to the start of a valid data block (before its From 95a5d887210d94f0e2e2230b4021897006c6f364 Mon Sep 17 00:00:00 2001 From: Igor Artamonov Date: Mon, 6 Apr 2026 14:36:41 +0100 Subject: [PATCH 3/4] after code review --- avro/src/reader/block.rs | 10 ---------- avro/src/reader/mod.rs | 6 +++--- 2 files changed, 3 insertions(+), 13 deletions(-) diff --git a/avro/src/reader/block.rs b/avro/src/reader/block.rs index ccbf64d7..724bc5ab 100644 --- a/avro/src/reader/block.rs +++ b/avro/src/reader/block.rs @@ -370,17 +370,7 @@ impl Block<'_, R> { self.message_count = 0; self.current_block_info = None; - // read_block_next treats UnexpectedEof as a clean end-of-stream - // (returns Ok with message_count=0). That's correct for forward - // iteration but wrong here — the caller asked for a specific block. self.read_block_next()?; - if self.is_empty() { - return Err(Details::SeekToBlock(std::io::Error::new( - std::io::ErrorKind::UnexpectedEof, - format!("no block at offset {offset}"), - )) - .into()); - } Ok(()) } } diff --git a/avro/src/reader/mod.rs b/avro/src/reader/mod.rs index 5e5946b0..583aac05 100644 --- a/avro/src/reader/mod.rs +++ b/avro/src/reader/mod.rs @@ -193,9 +193,9 @@ impl Reader<'_, R> { /// Typically the caller saves offsets from [`current_block`](Self::current_block) /// during forward iteration and later passes them here to jump back. pub fn seek_to_block(&mut self, offset: u64) -> AvroResult<()> { - self.block.seek_to_block(offset)?; - self.errored = false; - Ok(()) + let seek_status = self.block.seek_to_block(offset); + self.errored = seek_status.is_err(); + seek_status } } From 6d736c795746d7235b84fb97d19d7b210ba000f7 Mon Sep 17 00:00:00 2001 From: Igor Artamonov Date: Sun, 24 May 2026 01:46:47 +0100 Subject: [PATCH 4/4] after code review --- avro/src/reader/block.rs | 162 ++++++++++++++++++++++++++++----------- avro/src/reader/mod.rs | 70 ++++++++++++----- 2 files changed, 168 insertions(+), 64 deletions(-) diff --git a/avro/src/reader/block.rs b/avro/src/reader/block.rs index 724bc5ab..73fd63bc 100644 --- a/avro/src/reader/block.rs +++ b/avro/src/reader/block.rs @@ -15,16 +15,16 @@ // specific language governing permissions and limitations // under the License. +use log::warn; +use serde::de::DeserializeOwned; +use serde_json::from_slice; +use std::io::IoSliceMut; use std::{ collections::HashMap, io::{ErrorKind, Read, Seek, SeekFrom}, str::FromStr, }; -use log::warn; -use serde::de::DeserializeOwned; -use serde_json::from_slice; - use crate::{ AvroResult, Codec, Error, decode::{decode, decode_internal}, @@ -41,44 +41,105 @@ use crate::{ /// Use with [`super::Reader::seek_to_block`] to jump back to a previously-read block. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct BlockPosition { + // Byte offset in the stream where this block starts (before the object-count varint). + offset: u64, + // Total number of records in this block. + message_count: usize, +} + +impl BlockPosition { /// Byte offset in the stream where this block starts (before the object-count varint). - pub offset: u64, + pub fn offset(&self) -> u64 { + self.offset + } + /// Total number of records in this block. - pub message_count: usize, + pub fn message_count(&self) -> usize { + self.message_count + } } -/// Wraps an inner reader and tracks the current byte position. +/// Optional byte-offset tracker over an inner reader. +/// +/// `Direct` is a transparent pass-through so the inner type's optimized +/// `read_exact`/`read_to_end`/`read_vectored` are reached unchanged. /// -/// Avoids requiring `Seek` just to know how many bytes have been consumed. -/// When the inner reader also implements `Seek`, seeking updates the tracked position. +/// `Tracking` mirrors those same delegations and accumulates bytes consumed, +/// exposing the offset via [`Self::position`]. Tracking is opt-in. #[derive(Debug, Clone)] -struct PositionTracker { - inner: R, - pos: u64, +pub(super) enum PositionTracker { + Direct(R), + Tracking { inner: R, pos: u64 }, } impl PositionTracker { - fn new(inner: R) -> Self { - Self { inner, pos: 0 } - } - - fn position(&self) -> u64 { - self.pos + /// Byte offset consumed so far, or `None` when tracking is disabled. + pub(super) fn position(&self) -> Option { + match self { + Self::Direct(_) => None, + Self::Tracking { pos, .. } => Some(*pos), + } } } impl Read for PositionTracker { fn read(&mut self, buf: &mut [u8]) -> std::io::Result { - let n = self.inner.read(buf)?; - self.pos += n as u64; - Ok(n) + match self { + Self::Direct(r) => r.read(buf), + Self::Tracking { inner, pos } => { + let n = inner.read(buf)?; + *pos += n as u64; + Ok(n) + } + } + } + + fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> std::io::Result { + match self { + Self::Direct(r) => r.read_vectored(bufs), + Self::Tracking { inner, pos } => { + let n = inner.read_vectored(bufs)?; + *pos += n as u64; + Ok(n) + } + } + } + + fn read_to_end(&mut self, buf: &mut Vec) -> std::io::Result { + match self { + Self::Direct(r) => r.read_to_end(buf), + Self::Tracking { inner, pos } => { + let n = inner.read_to_end(buf)?; + *pos += n as u64; + Ok(n) + } + } + } + + fn read_exact(&mut self, buf: &mut [u8]) -> std::io::Result<()> { + match self { + Self::Direct(r) => r.read_exact(buf), + Self::Tracking { inner, pos } => { + inner.read_exact(buf)?; + // `read_exact` either fills `buf` completely or errors, so on + // success exactly `buf.len()` bytes were consumed. + *pos += buf.len() as u64; + Ok(()) + } + } } } impl Seek for PositionTracker { fn seek(&mut self, from: SeekFrom) -> std::io::Result { - self.pos = self.inner.seek(from)?; - Ok(self.pos) + match self { + Self::Direct(r) => r.seek(from), + Self::Tracking { inner, pos } => { + let new_pos = inner.seek(from)?; + *pos = new_pos; + Ok(new_pos) + } + } } } @@ -99,7 +160,7 @@ pub(super) struct Block<'r, R> { names_refs: Names, human_readable: bool, /// Byte offset where data blocks begin (right after header and sync marker). - pub(super) data_start: u64, + pub(super) data_start: Option, /// Position and record count of the currently loaded block. pub(super) current_block_info: Option, } @@ -109,9 +170,18 @@ impl<'r, R: Read> Block<'r, R> { reader: R, schemata: Vec<&'r Schema>, human_readable: bool, + track_positions: bool, ) -> AvroResult> { + let reader = if track_positions { + PositionTracker::Tracking { + inner: reader, + pos: 0, + } + } else { + PositionTracker::Direct(reader) + }; let mut block = Block { - reader: PositionTracker::new(reader), + reader, codec: Codec::Null, writer_schema: Schema::Null, schemata, @@ -122,7 +192,7 @@ impl<'r, R: Read> Block<'r, R> { user_metadata: Default::default(), names_refs: Default::default(), human_readable, - data_start: 0, + data_start: None, current_block_info: None, }; @@ -211,18 +281,24 @@ impl<'r, R: Read> Block<'r, R> { return Err(Details::GetBlockMarker.into()); } - self.current_block_info = Some(BlockPosition { - offset: block_start, - message_count: block_len, - }); - // NOTE (JAB): This doesn't fit this Reader pattern very well. // `self.buf` is a growable buffer that is reused as the reader is iterated. // For non `Codec::Null` variants, `decompress` will allocate a new `Vec` // and replace `buf` with the new one, instead of reusing the same buffer. // We can address this by using some "limited read" type to decode directly // into the buffer. But this is fine, for now. - self.codec.decompress(&mut self.buf) + let next = self.codec.decompress(&mut self.buf); + + // Make sure the position points only to a valid block + self.current_block_info = match next { + Ok(_) => block_start.map(|offset| BlockPosition { + offset, + message_count: block_len, + }), + Err(_) => None, + }; + + next } Err(Details::ReadVariableIntegerBytes(io_err)) => { if let ErrorKind::UnexpectedEof = io_err.kind() { @@ -361,15 +437,15 @@ impl Block<'_, R> { /// Returns an error if no valid block can be read at the offset /// (e.g., the offset is at or past EOF). pub(super) fn seek_to_block(&mut self, offset: u64) -> AvroResult<()> { - self.reader - .seek(SeekFrom::Start(offset)) - .map_err(Details::SeekToBlock)?; - self.buf.clear(); self.buf_idx = 0; self.message_count = 0; self.current_block_info = None; + self.reader + .seek(SeekFrom::Start(offset)) + .map_err(Details::SeekToBlock)?; + self.read_block_next()?; Ok(()) } @@ -444,10 +520,7 @@ mod tests { fn avro_rs_586_negative_block_size() { let mut block = Block::<'_, &[u8]> { // Block header with an object count of 1 and a block size of -1 - reader: PositionTracker { - inner: &[0x02, 0x01], - pos: 0, - }, + reader: PositionTracker::Direct(&[0x02, 0x01]), buf: vec![], buf_idx: 0, message_count: 0, @@ -458,7 +531,7 @@ mod tests { user_metadata: Default::default(), names_refs: Default::default(), human_readable: false, - data_start: 0, + data_start: None, current_block_info: None, }; @@ -470,10 +543,7 @@ mod tests { fn avro_rs_586_negative_block_len() { let mut block = Block::<'_, &[u8]> { // Block header with an object count of -1 and a block size of 0 - reader: PositionTracker { - inner: &[0x01, 0x00], - pos: 0, - }, + reader: PositionTracker::Direct(&[0x01, 0x00]), buf: vec![], buf_idx: 0, message_count: 0, @@ -484,7 +554,7 @@ mod tests { user_metadata: Default::default(), names_refs: Default::default(), human_readable: false, - data_start: 0, + data_start: None, current_block_info: None, }; diff --git a/avro/src/reader/mod.rs b/avro/src/reader/mod.rs index 583aac05..56263c6a 100644 --- a/avro/src/reader/mod.rs +++ b/avro/src/reader/mod.rs @@ -81,11 +81,15 @@ impl<'a, R: Read> Reader<'a, R> { reader_schema: Option<&'a Schema>, schemata: Option>, #[builder(default = is_human_readable())] human_readable: bool, + /// Track byte offsets so [`Reader::current_block`] and + /// [`Reader::data_start`] return useful values. Off by default. + #[builder(default = false)] + seekable: bool, ) -> AvroResult> { let schemata = schemata.unwrap_or_else(|| reader_schema.map(|rs| vec![rs]).unwrap_or_default()); - let block = Block::new(reader, schemata, human_readable)?; + let block = Block::new(reader, schemata, human_readable, seekable)?; let mut reader = Reader { block, reader_schema, @@ -167,17 +171,18 @@ impl Iterator for Reader<'_, R> { impl Reader<'_, R> { /// The currently loaded block's position and record count. /// - /// Returns `None` only before the first block is loaded (via iteration or - /// [`seek_to_block`](Self::seek_to_block)). Always `Some` afterward. + /// Returns `None` when position tracking was not enabled (see the + /// `seekable` builder option) or before the first block has been + /// loaded via iteration or [`seek_to_block`](Self::seek_to_block). pub fn current_block(&self) -> Option { self.block.current_block_info } /// Byte offset where data blocks begin (right after the file header). /// - /// This is the offset of the first data block — equivalent to the position - /// that would be returned by `current_block().offset` for block 0. - pub fn data_start(&self) -> u64 { + /// Equivalent to the offset of block 0. `None` when position tracking (`seekable`) was + /// not enabled at construction. + pub fn data_start(&self) -> Option { self.block.data_start } } @@ -447,7 +452,7 @@ mod tests { let data = writer.into_inner()?; // Read forward and collect block positions - let mut reader = Reader::new(Cursor::new(&data))?; + let mut reader = Reader::builder(Cursor::new(&data)).seekable(true).build()?; let mut block_offsets: Vec = Vec::new(); let mut all_values: Vec = Vec::new(); @@ -458,7 +463,7 @@ mod tests { let pos = reader.current_block().expect("block info after read"); if block_offsets .last() - .is_none_or(|last| last.offset != pos.offset) + .is_none_or(|last| last.offset() != pos.offset()) { block_offsets.push(pos); } @@ -466,25 +471,25 @@ mod tests { assert_eq!(all_values.len(), 5); assert_eq!(block_offsets.len(), 3); - assert_eq!(block_offsets[0].message_count, 2); - assert_eq!(block_offsets[1].message_count, 2); - assert_eq!(block_offsets[2].message_count, 1); - assert_eq!(reader.data_start(), block_offsets[0].offset); + assert_eq!(block_offsets[0].message_count(), 2); + assert_eq!(block_offsets[1].message_count(), 2); + assert_eq!(block_offsets[2].message_count(), 1); + assert_eq!(reader.data_start(), Some(block_offsets[0].offset())); // Seek back to block 1 and read its records - reader.seek_to_block(block_offsets[1].offset)?; + reader.seek_to_block(block_offsets[1].offset())?; let v1 = reader.next().unwrap()?; assert_eq!(v1, all_values[2]); let v2 = reader.next().unwrap()?; assert_eq!(v2, all_values[3]); // Seek back to block 0 - reader.seek_to_block(block_offsets[0].offset)?; + reader.seek_to_block(block_offsets[0].offset())?; let v0 = reader.next().unwrap()?; assert_eq!(v0, all_values[0]); // Seek to block 2 - reader.seek_to_block(block_offsets[2].offset)?; + reader.seek_to_block(block_offsets[2].offset())?; let v4 = reader.next().unwrap()?; assert_eq!(v4, all_values[4]); @@ -507,9 +512,38 @@ mod tests { writer.flush()?; let data = writer.into_inner()?; - let mut reader = Reader::new(Cursor::new(&data))?; - let result = reader.seek_to_block(7); - assert!(result.is_err()); + let mut reader = Reader::builder(Cursor::new(&data)).seekable(true).build()?; + assert!(reader.seek_to_block(7).is_err()); + + Ok(()) + } + + /// Seeking to outside the data + #[test] + fn test_seek_after_eof() -> TestResult { + use crate::writer::Writer; + + let schema = Schema::parse_str(SCHEMA)?; + let mut writer = Writer::new(&schema, Vec::new())?; + let mut r = Record::new(&schema).unwrap(); + r.put("a", 1i64); + r.put("b", "x"); + writer.append_value(r)?; + writer.flush()?; + let data = writer.into_inner()?; + + let mut reader = Reader::builder(Cursor::new(&data)).seekable(true).build()?; + + let eof = data.len() as u64; + let _ = reader.seek_to_block(eof); + assert_eq!(reader.current_block(), None); + let next = reader.next(); + assert!(next.is_none()); + + let _ = reader.seek_to_block(eof + 1); + assert_eq!(reader.current_block(), None); + let next = reader.next(); + assert!(next.is_none()); Ok(()) }