Skip to content

feat: seekable reader #530

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
splix wants to merge 4 commits into
base: main
Choose a base branch
from
+357 −13
Draft

feat: seekable reader #530

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
35f8fdd
feat: seekable reader
splix Apr 5, 2026
10f82da
after code review
splix Apr 6, 2026
95a5d88
after code review
splix Apr 6, 2026
6d736c7
after code review
splix May 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions avro/src/error.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),

Expand Down
2 changes: 1 addition & 1 deletion avro/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down
180 changes: 170 additions & 10 deletions avro/src/reader/block.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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<R> {
Direct(R),
Tracking { inner: R, pos: u64 },
}

impl<R> PositionTracker<R> {
/// Byte offset consumed so far, or `None` when tracking is disabled.
pub(super) fn position(&self) -> Option<u64> {
match self {
Self::Direct(_) => None,
Self::Tracking { pos, .. } => Some(*pos),
}
}
}

impl<R: Read> Read for PositionTracker<R> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
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<usize> {
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<u8>) -> std::io::Result<usize> {
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(())
}
}
}
Comment thread
splix marked this conversation as resolved.
}

impl<R: Seek> Seek for PositionTracker<R> {
fn seek(&mut self, from: SeekFrom) -> std::io::Result<u64> {
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<R>,
/// Internal buffering to reduce allocation.
buf: Vec<u8>,
buf_idx: usize,
Expand All @@ -51,14 +159,27 @@ pub(super) struct Block<'r, R> {
pub(super) user_metadata: HashMap<String, Vec<u8>>,
names_refs: Names,
human_readable: bool,
/// Byte offset where data blocks begin (right after header and sync marker).
pub(super) data_start: Option<u64>,
/// Position and record count of the currently loaded block.
pub(super) current_block_info: Option<BlockPosition>,
}

impl<'r, R: Read> Block<'r, R> {
pub(super) fn new(
reader: R,
schemata: Vec<&'r Schema>,
human_readable: bool,
track_positions: bool,
) -> AvroResult<Block<'r, R>> {
let reader = if track_positions {
PositionTracker::Tracking {
inner: reader,
pos: 0,

@martin-g martin-g Jul 21, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the underlying reader R is already partially read then pos: 0 won't be correct.
I am not sure how feasible this is though.
We can add one more optional field to ReaderBuilder - initial_position, with a default value of 0.

}
} else {
PositionTracker::Direct(reader)
};
let mut block = Block {
reader,
codec: Codec::Null,
Expand All @@ -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)
}

Expand Down Expand Up @@ -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;
Expand All @@ -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 {

@martin-g martin-g Jul 21, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this be done earlier ?!
If there is a problem with the marker then any call to Reader::current_block() will return wrong data.

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() {
Expand Down Expand Up @@ -295,6 +431,26 @@ impl<'r, R: Read> Block<'r, R> {
}
}

impl<R: Read + Seek> 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).

@martin-g martin-g Jul 21, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://github.com/apache/avro-rs/pull/530/changes#diff-6d2c5562c5bd4d218417ad6c2eedd5e2a4f1cad8980fba3198a896aed8828adeR303-R306 returns Ok on EOF, not an Err.

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<String, Value>) -> AvroResult<Codec> {
let result = metadata
.get("avro.codec")
Expand Down Expand Up @@ -356,15 +512,15 @@ fn read_codec(metadata: &HashMap<String, Value>) -> AvroResult<Codec> {

#[cfg(test)]
mod tests {
use super::Block;
use super::{Block, PositionTracker};
use crate::error::Details;
use crate::{Codec, Schema};

#[test]
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,
Expand All @@ -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();
Expand All @@ -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,
Expand All @@ -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();
Expand Down
Loading
Loading