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..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}, + 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}, @@ -35,10 +35,118 @@ use crate::{ util, }; +/// 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. +#[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 fn offset(&self) -> u64 { + self.offset + } + + /// Total number of records in this block. + pub fn message_count(&self) -> usize { + self.message_count + } +} + +/// 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. +/// +/// `Tracking` mirrors those same delegations and accumulates bytes consumed, +/// exposing the offset via [`Self::position`]. Tracking is opt-in. +#[derive(Debug, Clone)] +pub(super) enum PositionTracker { + Direct(R), + Tracking { inner: R, pos: u64 }, +} + +impl PositionTracker { + /// 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 { + 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 { + match self { + Self::Direct(r) => r.seek(from), + Self::Tracking { inner, pos } => { + let new_pos = inner.seek(from)?; + *pos = new_pos; + Ok(new_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 +159,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: Option, + /// Position and record count of the currently loaded block. + pub(super) current_block_info: Option, } impl<'r, R: Read> Block<'r, R> { @@ -58,7 +170,16 @@ 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, codec: Codec::Null, @@ -71,9 +192,12 @@ impl<'r, R: Read> Block<'r, R> { user_metadata: Default::default(), names_refs: Default::default(), human_readable, + data_start: None, + current_block_info: None, }; block.read_header()?; + block.data_start = block.reader.position(); Ok(block) } @@ -142,6 +266,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; @@ -162,7 +287,18 @@ impl<'r, R: Read> Block<'r, R> { // 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() { @@ -295,6 +431,26 @@ 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.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(()) + } +} + fn read_codec(metadata: &HashMap) -> AvroResult { let result = metadata .get("avro.codec") @@ -356,7 +512,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 +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: &[0x02, 0x01], + reader: PositionTracker::Direct(&[0x02, 0x01]), buf: vec![], buf_idx: 0, message_count: 0, @@ -375,6 +531,8 @@ mod tests { user_metadata: Default::default(), names_refs: Default::default(), human_readable: false, + data_start: None, + current_block_info: None, }; let err = block.read_block_next().unwrap_err().into_details(); @@ -385,7 +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: &[0x01, 0x00], + reader: PositionTracker::Direct(&[0x01, 0x00]), buf: vec![], buf_idx: 0, message_count: 0, @@ -396,6 +554,8 @@ mod tests { user_metadata: Default::default(), names_refs: Default::default(), human_readable: false, + data_start: None, + 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..56263c6a 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; @@ -76,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, @@ -159,6 +168,42 @@ impl Iterator for Reader<'_, R> { } } +impl Reader<'_, R> { + /// The currently loaded block's position and record count. + /// + /// 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). + /// + /// 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 + } +} + +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 + /// 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<()> { + let seek_status = self.block.seek_to_block(offset); + self.errored = seek_status.is_err(); + seek_status + } +} + /// Wrapper around [`Reader`] where the iterator deserializes `T`. pub struct ReaderDeser<'a, R, T> { inner: Reader<'a, R>, @@ -366,4 +411,140 @@ 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::builder(Cursor::new(&data)).seekable(true).build()?; + 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(), Some(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::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(()) + } }