Merge pull request aldanor#198 from Rikorose/feat_io_read

Add ByteReader implementing io::Read and io::Seek

Author: mulimoen

CHANGELOG.md

@@ -6,6 +6,8 @@
 
 - Support for HDF5 version 1.13.0.
 - Support field renaming via `#[hdf5(rename = "new_name")]` helper attribute.
+- Add a `ByteReader` which implements `std::io::{Read, Seek}` for 1D `u8`
+  datasets. Usage via `Dataset::as_byte_reader()`.
 
 ### Changed
 

src/error.rs

@@ -1,6 +1,7 @@
 use std::convert::Infallible;
 use std::error::Error as StdError;
 use std::fmt;
+use std::io;
 use std::ops::Deref;
 use std::panic;
 use std::ptr;
@@ -268,6 +269,12 @@ impl From<ShapeError> for Error {
     }
 }
 
+impl From<Error> for io::Error {
+    fn from(err: Error) -> Self {
+        Self::new(io::ErrorKind::Other, err)
+    }
+}
+
 pub fn h5check<T: H5ErrorCode>(value: T) -> Result<T> {
     H5ErrorCode::h5check(value)
 }

src/hl.rs

@@ -17,7 +17,7 @@ pub use self::{
         Attribute, AttributeBuilder, AttributeBuilderData, AttributeBuilderEmpty,
         AttributeBuilderEmptyShape,
     },
-    container::{Container, Reader, Writer},
+    container::{ByteReader, Container, Reader, Writer},
     dataset::{
         Dataset, DatasetBuilder, DatasetBuilderData, DatasetBuilderEmpty, DatasetBuilderEmptyShape,
     },

src/hl/container.rs

@@ -1,5 +1,6 @@
 use std::convert::TryInto;
 use std::fmt::{self, Debug};
+use std::io;
 use std::mem;
 use std::ops::Deref;
 
@@ -346,6 +347,99 @@ impl<'a> Writer<'a> {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct ByteReader {
+    obj: Container,
+    pos: u64,
+    dt: Datatype,
+    obj_space: Dataspace,
+    xfer: PropertyList,
+}
+
+impl ByteReader {
+    pub fn new(obj: &Container) -> Result<Self> {
+        ensure!(!obj.is_attr(), "ByteReader cannot be used on attribute datasets");
+
+        let obj = obj.clone();
+        let file_dtype = obj.dtype()?;
+        let mem_dtype = Datatype::from_type::<u8>()?;
+        file_dtype.ensure_convertible(&mem_dtype, Conversion::NoOp)?;
+
+        let obj_space = obj.space()?;
+        ensure!(obj_space.shape().len() == 1, "Only rank 1 datasets can be read via ByteReader");
+        let xfer = PropertyList::from_id(h5call!(H5Pcreate(*crate::globals::H5P_DATASET_XFER))?)?;
+        if !hdf5_types::USING_H5_ALLOCATOR {
+            crate::hl::plist::set_vlen_manager_libc(xfer.id())?;
+        }
+        Ok(ByteReader { obj, pos: 0, obj_space, dt: mem_dtype, xfer })
+    }
+
+    fn dataset_len(&self) -> usize {
+        self.obj_space.shape()[0]
+    }
+
+    fn remaining_len(&self) -> usize {
+        self.dataset_len().saturating_sub(self.pos as usize)
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.pos >= self.dataset_len() as u64
+    }
+}
+
+impl io::Read for ByteReader {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let pos = self.pos as usize;
+        let amt = std::cmp::min(buf.len(), self.remaining_len());
+        let selection = Selection::new(pos..pos + amt);
+        let out_shape = selection.out_shape(&self.obj_space.shape())?;
+        let fspace = self.obj_space.select(selection)?;
+        let mspace = Dataspace::try_new(&out_shape)?;
+        h5call!(H5Dread(
+            self.obj.id(),
+            self.dt.id(),
+            mspace.id(),
+            fspace.id(),
+            self.xfer.id(),
+            buf.as_mut_ptr().cast()
+        ))?;
+        self.pos += amt as u64;
+        Ok(out_shape[0])
+    }
+}
+
+impl io::Seek for ByteReader {
+    fn seek(&mut self, style: io::SeekFrom) -> io::Result<u64> {
+        let (base_pos, offset) = match style {
+            io::SeekFrom::Start(n) => {
+                self.pos = n;
+                return Ok(n);
+            }
+            io::SeekFrom::End(n) => (self.dataset_len() as u64, n),
+            io::SeekFrom::Current(n) => (self.pos, n),
+        };
+        let new_pos = if offset.is_negative() {
+            base_pos.checked_sub(offset.wrapping_abs() as u64)
+        } else {
+            base_pos.checked_add(offset as u64)
+        };
+        match new_pos {
+            Some(n) => {
+                self.pos = n;
+                Ok(self.pos)
+            }
+            None => Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "invalid seek to a negative or overflowing position",
+            )),
+        }
+    }
+
+    fn stream_position(&mut self) -> io::Result<u64> {
+        Ok(self.pos)
+    }
+}
+
 #[repr(transparent)]
 #[derive(Clone)]
 /// An object which can be read or written to.
@@ -395,6 +489,14 @@ impl Container {
         Writer::new(self)
     }
 
+    /// Creates `ByteReader` which implements [`Read`](std::io::Read)
+    /// and [`Seek`](std::io::Seek).
+    ///
+    /// ``ByteReader`` only supports 1-D `u8` datasets.
+    pub fn as_byte_reader(&self) -> Result<ByteReader> {
+        ByteReader::new(&self)
+    }
+
     /// Returns the datatype of the dataset/attribute.
     pub fn dtype(&self) -> Result<Datatype> {
         if self.is_attr() {

src/lib.rs

@@ -58,7 +58,7 @@ mod export {
         hl::selection::{Hyperslab, Selection, SliceOrIndex},
         hl::{
             Attribute, AttributeBuilder, AttributeBuilderData, AttributeBuilderEmpty,
-            AttributeBuilderEmptyShape, Container, Conversion, Dataset, DatasetBuilder,
+            AttributeBuilderEmptyShape, ByteReader, Container, Conversion, Dataset, DatasetBuilder,
             DatasetBuilderData, DatasetBuilderEmpty, DatasetBuilderEmptyShape, Dataspace, Datatype,
             File, FileBuilder, Group, LinkInfo, LinkType, Location, LocationInfo, LocationToken,
             LocationType, Object, PropertyList, Reader, Writer,

tests/test_dataset.rs

@@ -1,10 +1,11 @@
 use std::convert::TryFrom;
 use std::fmt;
+use std::io::{Read, Seek, SeekFrom};
 
 use ndarray::{s, Array1, Array2, ArrayD, IxDyn, SliceInfo};
 use rand::prelude::{Rng, SeedableRng, SmallRng};
 
-use hdf5_types::TypeDescriptor;
+use hdf5_types::{H5Type, TypeDescriptor};
 
 mod common;
 
@@ -171,6 +172,85 @@ where
     Ok(())
 }
 
+fn test_byte_read_seek_impl(ds: &hdf5::Dataset, arr: &ArrayD<u8>, ndim: usize) -> hdf5::Result<()> {
+    let mut rng = SmallRng::seed_from_u64(42);
+    ds.write(arr)?;
+
+    // Read whole
+    let reader = ds.as_byte_reader();
+    let mut reader = if ndim != 1 {
+        assert!(reader.is_err());
+        return Ok(());
+    } else {
+        reader.unwrap()
+    };
+    let mut out_bytes = vec![0u8; arr.len()];
+    reader.read(&mut out_bytes.as_mut_slice()).expect("io::Read failed");
+    assert_eq!(out_bytes.as_slice(), arr.as_slice().unwrap());
+
+    // Read in chunks
+    let mut reader = reader.clone();
+    reader.seek(std::io::SeekFrom::Start(0)).expect("io::Seek failed");
+    let mut pos = 0;
+    while pos < arr.len() {
+        let chunk_len: usize = rng.gen_range(1..arr.len() + 1);
+        let mut chunk = vec![0u8; chunk_len];
+        let n_read = reader.read(&mut chunk).expect("io::Read failed");
+        if pos + chunk_len < arr.len() {
+            // We did not read until end. Thus, the chunk should be fully filled.
+            assert_eq!(chunk_len, n_read);
+        }
+        assert_eq!(&chunk[..n_read], arr.slice(s![pos..pos + n_read]).as_slice().unwrap());
+        pos += chunk_len;
+    }
+
+    // Seek to the begining and read again
+    reader.seek(SeekFrom::Start(0)).expect("io::Seek failed");
+    let mut out_bytes = vec![0u8; arr.len()];
+    reader.read(&mut out_bytes.as_mut_slice()).expect("io::Read failed");
+    assert_eq!(out_bytes.as_slice(), arr.as_slice().unwrap());
+
+    // Seek to a random position from start
+    let pos = rng.gen_range(0..arr.len() + 1) as u64;
+    let seeked_pos = reader.seek(SeekFrom::Start(pos)).expect("io::Seek failed") as usize;
+    let mut out_bytes = vec![0u8; arr.len() - seeked_pos];
+    reader.read(&mut out_bytes.as_mut_slice()).expect("io::Read failed");
+    assert_eq!(out_bytes.as_slice(), arr.slice(s![seeked_pos..]).as_slice().unwrap());
+
+    // Seek from current position
+    let orig_pos = reader.seek(SeekFrom::Start(pos)).expect("io::Seek failed") as i64;
+    let rel_pos = rng.gen_range(-(arr.len() as i64)..arr.len() as i64 + 1);
+    let pos_res = reader.seek(SeekFrom::Current(rel_pos));
+    if (rel_pos + orig_pos) < 0 {
+        assert!(pos_res.is_err()) // We cannot seek before start
+    } else {
+        let seeked_pos = pos_res.unwrap() as usize;
+        assert_eq!(rel_pos + orig_pos, seeked_pos as i64);
+        let mut out_bytes = vec![0u8; arr.len() - seeked_pos];
+        reader.read(&mut out_bytes.as_mut_slice()).expect("io::Read failed");
+        assert_eq!(out_bytes.as_slice(), arr.slice(s![seeked_pos..]).as_slice().unwrap());
+    }
+
+    // Seek to a random position from end
+    let pos = -(rng.gen_range(0..arr.len() + 1) as i64);
+    let seeked_pos = reader.seek(SeekFrom::End(pos)).expect("io::Seek failed") as usize;
+    assert_eq!(pos, seeked_pos as i64 - arr.len() as i64);
+    let mut out_bytes = vec![0u8; arr.len() - seeked_pos];
+    reader.read(&mut out_bytes.as_mut_slice()).expect("io::Read failed");
+    assert_eq!(out_bytes.as_slice(), arr.slice(s![seeked_pos..]).as_slice().unwrap());
+
+    // Seek before start
+    assert!(reader.seek(SeekFrom::End(-(arr.len() as i64) - 1)).is_err());
+
+    // Test stream position start
+    // Requires Rust 1.55.0: reader.rewind().expect("io::Seek::rewind failed");
+    assert_eq!(0, reader.seek(SeekFrom::Start(0)).unwrap());
+    assert_eq!(0, reader.stream_position().unwrap());
+    assert_eq!(0, reader.seek(SeekFrom::End(-(arr.len() as i64))).unwrap());
+    assert_eq!(0, reader.stream_position().unwrap());
+    Ok(())
+}
+
 fn test_read_write<T>() -> hdf5::Result<()>
 where
     T: hdf5::H5Type + fmt::Debug + PartialEq + Gen + Clone,
@@ -278,3 +358,24 @@ fn test_read_write_rename_fields() -> hdf5::Result<()> {
     test_read_write::<RenameEnum>()?;
     Ok(())
 }
+
+#[test]
+fn test_byte_read_seek() -> hdf5::Result<()> {
+    let mut rng = SmallRng::seed_from_u64(42);
+    let file = new_in_memory_file()?;
+
+    for ndim in 0..=2 {
+        for _ in 0..=20 {
+            let arr: ArrayD<u8> = gen_arr(&mut rng, ndim);
+
+            let ds: hdf5::Dataset = file.new_dataset::<u8>().shape(arr.shape()).create("x")?;
+            let ds = scopeguard::guard(ds, |ds| {
+                drop(ds);
+                drop(file.unlink("x"));
+            });
+
+            test_byte_read_seek_impl(&ds, &arr, ndim)?;
+        }
+    }
+    Ok(())
+}