ADD: Add absolute seek support for DBN

Author: jparkdatabento

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## 0.60.0 - TBD
+
+### Enhancements
+- Added `RecordDecoder::seek_to` for absolute byte-offset seeking
+- Added `is_compressed` helpers for dynamic DBN readers
+
 ## 0.59.0 - 2026-06-02
 
 ### Enhancements

rust/dbn/src/decode/dbn/async.rs

@@ -10,8 +10,8 @@ use crate::{
     decode::{
         dbn::fsm::{DbnFsm, ProcessResult},
         zstd::zstd_decoder,
-        AsyncDecodeRecord, AsyncDecodeRecordRef, AsyncSkipBytes, DbnMetadata, VersionUpgradePolicy,
-        ZSTD_FILE_BUFFER_CAPACITY,
+        AsyncDecodeRecord, AsyncDecodeRecordRef, AsyncDynReader, AsyncSkipBytes, DbnMetadata,
+        VersionUpgradePolicy, ZSTD_FILE_BUFFER_CAPACITY,
     },
     HasRType, Metadata, RecordRef, Result, DBN_VERSION,
 };
@@ -458,6 +458,73 @@ where
     }
 }
 
+impl<R> RecordDecoder<R>
+where
+    R: io::AsyncReadExt + io::AsyncSeekExt + Unpin,
+{
+    /// Seeks to absolute byte position `pos` and clears buffered decoder state.
+    ///
+    /// # Warning
+    /// Callers are responsible for ensuring `pos` is a valid record boundary.
+    ///
+    /// # Cancel safety
+    /// This method may not be cancellation safe, depending on the cancellation safety
+    /// of `seek()` of the inner reader `R`.
+    ///
+    /// # Errors
+    /// This function returns an error if it fails to seek in the inner reader.
+    pub async fn seek_to(&mut self, pos: u64) -> crate::Result<()> {
+        self.reader
+            .seek(std::io::SeekFrom::Start(pos))
+            .await
+            .map(drop)
+            .map_err(|err| crate::Error::io(err, format!("seeking to byte offset {pos}")))?;
+        self.fsm.reset_for_seek();
+        Ok(())
+    }
+}
+
+impl<R> RecordDecoder<AsyncDynReader<R>>
+where
+    R: io::AsyncReadExt + io::AsyncBufReadExt + Unpin,
+{
+    /// Seeks to absolute byte position `pos` and clears buffered decoder state.
+    ///
+    /// # Warning
+    /// Callers are responsible for ensuring `pos` is a valid record boundary.
+    ///
+    /// # Cancel safety
+    /// This method is not cancel safe.
+    ///
+    /// # Errors
+    /// This function returns an error if it fails to seek in the inner reader, or if
+    /// the input is Zstandard-compressed.
+    pub async fn seek_to(&mut self, pos: u64) -> crate::Result<()>
+    where
+        R: io::AsyncSeekExt,
+    {
+        if self.reader.is_compressed() {
+            return Err(crate::Error::BadArgument {
+                param_name: "self".to_owned(),
+                desc: "absolute seek is unsupported for zstd-compressed input".to_owned(),
+            });
+        }
+        self.reader
+            .get_mut()
+            .seek(std::io::SeekFrom::Start(pos))
+            .await
+            .map(drop)
+            .map_err(|err| crate::Error::io(err, format!("seeking to byte offset {pos}")))?;
+        self.fsm.reset_for_seek();
+        Ok(())
+    }
+
+    /// Returns whether the input is Zstandard-compressed.
+    pub fn is_compressed(&self) -> bool {
+        self.reader.is_compressed()
+    }
+}
+
 impl<R> RecordDecoder<R>
 where
     R: AsyncSkipBytes + io::AsyncReadExt + Unpin,
@@ -763,6 +830,57 @@ mod tests {
         assert!(decoder.decode_record::<MboMsg>().await.unwrap().is_none());
     }
 
+    #[tokio::test]
+    async fn test_seek_to_resets_buffered_records() {
+        let mut first = MboMsg {
+            hd: RecordHeader::new::<MboMsg>(rtype::MBO, 1, 100, 1),
+            ..Default::default()
+        };
+        first.order_id = 1;
+        let mut second = MboMsg {
+            hd: RecordHeader::new::<MboMsg>(rtype::MBO, 1, 101, 2),
+            ..Default::default()
+        };
+        second.order_id = 2;
+        let mut buffer = Vec::new();
+        buffer.extend_from_slice(first.as_ref());
+        buffer.extend_from_slice(second.as_ref());
+
+        let mut decoder = RecordDecoder::with_version(
+            std::io::Cursor::new(buffer),
+            DBN_VERSION,
+            VersionUpgradePolicy::AsIs,
+            false,
+        )
+        .unwrap();
+
+        assert_eq!(*decoder.decode::<MboMsg>().await.unwrap().unwrap(), first);
+        decoder
+            .seek_to(std::mem::size_of::<MboMsg>() as u64)
+            .await
+            .unwrap();
+        assert_eq!(*decoder.decode::<MboMsg>().await.unwrap().unwrap(), second);
+        decoder.seek_to(0).await.unwrap();
+        assert_eq!(*decoder.decode::<MboMsg>().await.unwrap().unwrap(), first);
+    }
+
+    #[tokio::test]
+    async fn test_seek_to_compressed_returns_bad_argument() {
+        let reader =
+            AsyncDynReader::from_file(format!("{TEST_DATA_PATH}/test_data.mbo.v3.dbn.zst"))
+                .await
+                .unwrap();
+        let mut decoder =
+            RecordDecoder::with_version(reader, DBN_VERSION, VersionUpgradePolicy::AsIs, false)
+                .unwrap();
+
+        assert!(matches!(
+            decoder.seek_to(0).await.unwrap_err(),
+            Error::BadArgument { param_name, desc }
+                if param_name == "self" && desc.contains("zstd-compressed")
+        ));
+    }
+
     #[tokio::test]
     async fn test_dbn_identity_with_ts_out() {
         let rec1 = WithTsOut {

rust/dbn/src/decode/dbn/fsm.rs

@@ -288,6 +288,14 @@ impl DbnFsm {
         self.needs_upgrade = Self::compute_needs_upgrade(self.upgrade_policy, None);
     }
 
+    /// Resets buffered record state after seeking while preserving decoder
+    /// configuration.
+    pub(crate) fn reset_for_seek(&mut self) {
+        self.state = State::Record;
+        self.buffer.reset();
+        self.compat_buffer.reset();
+    }
+
     /// Skips ahead `nbytes`. Returns the actual number of bytes skipped.
     ///
     /// Writable space is not reclaimed until the next call to

rust/dbn/src/decode/dbn/sync.rs

@@ -8,8 +8,8 @@ use crate::{
     decode::{
         dbn::fsm::{DbnFsm, ProcessResult},
         private::LastRecord,
-        DbnMetadata, DecodeRecord, DecodeRecordRef, DecodeStream, SkipBytes, StreamIterDecoder,
-        VersionUpgradePolicy,
+        DbnMetadata, DecodeRecord, DecodeRecordRef, DecodeStream, DynReader, SkipBytes,
+        StreamIterDecoder, VersionUpgradePolicy,
     },
     HasRType, Metadata, RecordRef, DBN_VERSION,
 };
@@ -368,6 +368,64 @@ where
     }
 }
 
+impl<R> RecordDecoder<R>
+where
+    R: io::Read + io::Seek,
+{
+    /// Seeks to absolute byte position `pos` and clears buffered decoder state.
+    ///
+    /// # Warning
+    /// Callers are responsible for ensuring `pos` is a valid record boundary.
+    ///
+    /// # Errors
+    /// This function returns an error if it fails to seek in the inner reader.
+    pub fn seek_to(&mut self, pos: u64) -> crate::Result<()> {
+        self.reader
+            .seek(io::SeekFrom::Start(pos))
+            .map(drop)
+            .map_err(|err| crate::Error::io(err, format!("seeking to byte offset {pos}")))?;
+        self.fsm.reset_for_seek();
+        Ok(())
+    }
+}
+
+impl<'a, R> RecordDecoder<DynReader<'a, R>>
+where
+    R: io::BufRead,
+{
+    /// Seeks to absolute byte position `pos` and clears buffered decoder state.
+    ///
+    /// # Warning
+    /// Callers are responsible for ensuring `pos` is a valid record boundary.
+    ///
+    /// # Errors
+    /// This function returns an error if it fails to seek in the inner reader, or if
+    /// the input is Zstandard-compressed.
+    pub fn seek_to(&mut self, pos: u64) -> crate::Result<()>
+    where
+        R: io::Seek,
+    {
+        if self.reader.is_compressed() {
+            return Err(crate::Error::BadArgument {
+                param_name: "self".to_owned(),
+                desc: "absolute seek is unsupported for zstd-compressed input".to_owned(),
+            });
+        }
+        self.reader
+            .get_mut()
+            .seek(io::SeekFrom::Start(pos))
+            .map(drop)
+            .map_err(|err| crate::Error::io(err, format!("seeking to byte offset {pos}")))?;
+        self.fsm.reset_for_seek();
+        Ok(())
+    }
+
+    /// Returns whether the input is Zstandard-compressed.
+    pub fn is_compressed(&self) -> bool {
+        self.reader.is_compressed()
+    }
+}
+
 impl<R> LastRecord for RecordDecoder<R>
 where
     R: io::Read,
@@ -618,6 +676,54 @@ mod tests {
         assert!(decoder.decode_record::<MboMsg>().unwrap().is_none());
     }
 
+    #[test]
+    fn test_seek_to_resets_buffered_records() {
+        let mut first = MboMsg {
+            hd: RecordHeader::new::<MboMsg>(rtype::MBO, 1, 100, 1),
+            ..Default::default()
+        };
+        first.order_id = 1;
+        let mut second = MboMsg {
+            hd: RecordHeader::new::<MboMsg>(rtype::MBO, 1, 101, 2),
+            ..Default::default()
+        };
+        second.order_id = 2;
+        let mut buffer = Vec::new();
+        buffer.extend_from_slice(first.as_ref());
+        buffer.extend_from_slice(second.as_ref());
+
+        let mut decoder = RecordDecoder::with_version(
+            std::io::Cursor::new(buffer),
+            DBN_VERSION,
+            VersionUpgradePolicy::AsIs,
+            false,
+        )
+        .unwrap();
+
+        assert_eq!(*decoder.decode::<MboMsg>().unwrap().unwrap(), first);
+        decoder
+            .seek_to(std::mem::size_of::<MboMsg>() as u64)
+            .unwrap();
+        assert_eq!(*decoder.decode::<MboMsg>().unwrap().unwrap(), second);
+        decoder.seek_to(0).unwrap();
+        assert_eq!(*decoder.decode::<MboMsg>().unwrap().unwrap(), first);
+    }
+
+    #[test]
+    fn test_seek_to_compressed_returns_bad_argument() {
+        let reader =
+            DynReader::from_file(format!("{TEST_DATA_PATH}/test_data.mbo.v3.dbn.zst")).unwrap();
+        let mut decoder =
+            RecordDecoder::with_version(reader, DBN_VERSION, VersionUpgradePolicy::AsIs, false)
+                .unwrap();
+
+        assert!(matches!(
+            decoder.seek_to(0).unwrap_err(),
+            Error::BadArgument { param_name, desc }
+                if param_name == "self" && desc.contains("zstd-compressed")
+        ));
+    }
+
     #[test]
     fn test_dbn_identity_with_ts_out() -> Result<()> {
         let rec1 = WithTsOut {

rust/dbn/src/decode/dyn_reader.rs

@@ -101,6 +101,11 @@ where
             DynReaderImpl::Zstd(reader) => reader.get_ref(),
         }
     }
+
+    /// Returns whether the input is Zstandard-compressed.
+    pub fn is_compressed(&self) -> bool {
+        matches!(&self.0, DynReaderImpl::Zstd(_))
+    }
 }
 
 impl DynReader<'_, BufReader<File>> {
@@ -180,6 +185,8 @@ mod tests {
             DynReader::from_file(format!("{TEST_DATA_PATH}/test_data.mbo.v3.dbn")).unwrap();
         let mut compressed =
             DynReader::from_file(format!("{TEST_DATA_PATH}/test_data.mbo.v3.dbn.zst")).unwrap();
+        assert!(!uncompressed.is_compressed());
+        assert!(compressed.is_compressed());
         let mut uncompressed_res = Vec::new();
         uncompressed.read_to_end(&mut uncompressed_res).unwrap();
         let mut compressed_res = Vec::new();
@@ -300,6 +307,11 @@ mod r#async {
                 DynReaderImpl::Zstd(reader) => reader.get_ref(),
             }
         }
+
+        /// Returns whether the input is Zstandard-compressed.
+        pub fn is_compressed(&self) -> bool {
+            matches!(&self.0, DynReaderImpl::Zstd(_))
+        }
     }
 
     impl DynReader<BufReader<File>> {
@@ -385,6 +397,7 @@ mod r#async {
             }
         }
     }
+
     #[cfg(test)]
     mod tests {
         use crate::{
@@ -398,7 +411,7 @@ mod r#async {
         #[tokio::test]
         async fn test_decode_multiframe_zst() {
             let mut decoder = AsyncDbnRecordDecoder::with_version(
-                DynReader::from_file(&format!(
+                DynReader::from_file(format!(
                     "{TEST_DATA_PATH}/multi-frame.definition.v1.dbn.frag.zst"
                 ))
                 .await
@@ -414,5 +427,20 @@ mod r#async {
             }
             assert_eq!(count, 8);
         }
+
+        #[tokio::test]
+        async fn test_dyn_reader_is_compressed() {
+            let uncompressed =
+                DynReader::from_file(format!("{TEST_DATA_PATH}/test_data.mbo.v3.dbn"))
+                    .await
+                    .unwrap();
+            let compressed =
+                DynReader::from_file(format!("{TEST_DATA_PATH}/test_data.mbo.v3.dbn.zst"))
+                    .await
+                    .unwrap();
+
+            assert!(!uncompressed.is_compressed());
+            assert!(compressed.is_compressed());
+        }
     }
 }