clawhdf5/crates/clawhdf5-py/src/dataset.rs

//! PyDataset — h5py-style read access to HDF5 datasets.
//!
//! `ds[key]` parses the key into hyperslab selections (see `select`) and
//! reads them through the facade's `read_selection`, which decodes only the
//! chunks a small selection touches (see its docs for when it decodes the
//! whole dataset instead); the
//! bytes it returns become the numpy array's buffer without a copy (see
//! `convert`). All file access and decoding runs with the GIL released, so
//! Python threads reading the same or different datasets run in parallel.

use std::sync::Arc;

use clawhdf5_format::datatype::Datatype;
use clawhdf5_format::object_header::ObjectHeader;
use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyList, PyTuple};

use crate::attrs::PyAttrs;
use crate::convert::{Converter, Elements, resolve_vl};
use crate::select::{self, Plan};
use crate::{PyEmpty, node, to_py_err};

/// A dataset in a file opened for reading.
///
/// ```python
/// ds = f['group/dataset']
/// ds.shape, ds.dtype, ds.attrs['units']
/// block = ds[10:20, ::2]      # a small selection reads only its chunks
/// ```
#[pyclass(name = "Dataset")]
pub struct PyDataset {
    file: Arc<clawhdf5_rs::File>,
    path: String,
    /// Where the dataset's object header is: reads open it from here rather
    /// than resolve `path` again.
    addr: u64,
    /// `None` for a dataset with a null dataspace (h5py's `Empty`).
    shape: Option<Vec<u64>>,
    /// The chunk shape, for a chunked dataset.
    chunks: Option<Vec<u64>>,
    datatype: Datatype,
    /// Why the datatype cannot be read into numpy, if it cannot.
    conv: Result<Converter, String>,
}

impl PyDataset {
    pub(crate) fn open(
        py: Python<'_>,
        file: Arc<clawhdf5_rs::File>,
        path: String,
        addr: u64,
        hdr: &ObjectHeader,
    ) -> PyResult<Self> {
        crate::no_panic(|| {
            let null = node::is_null(&node::dataspace(&file, hdr)?);
            let (shape, datatype) = {
                let ds = file.dataset_at(addr).map_err(to_py_err)?;
                let shape = if null {
                    None
                } else {
                    Some(ds.shape().map_err(to_py_err)?)
                };
                (shape, ds.raw_datatype().map_err(to_py_err)?)
            };
            let conv = Converter::new(py, &datatype, file.superblock().offset_size)
                .map_err(|e| e.value(py).to_string());
            let chunks = shape
                .as_ref()
                .and_then(|s| node::chunk_shape(&file, hdr, s.len()));
            Ok(Self {
                file,
                path,
                addr,
                shape,
                chunks,
                datatype,
                conv,
            })
        })
    }

    fn converter(&self) -> PyResult<&Converter> {
        self.conv
            .as_ref()
            .map_err(|msg| PyTypeError::new_err(format!("{}: {msg}", node::name(&self.path))))
    }

    /// Read the selection described by `plan` into a numpy array.
    fn read_plan<'py>(&self, py: Python<'py>, plan: &Plan) -> PyResult<Bound<'py, PyAny>> {
        let conv = self.converter()?;
        let dims = self.shape.as_deref().unwrap_or(&[]);
        let out_shape = plan.out_shape();

        let arr = if plan.is_empty() {
            conv.empty(py, &out_shape)?
        } else {
            let (vl, elem_size, unit) = (conv.is_vl(), conv.elem_size, conv.vl_unit);
            let list_axis = plan.list_axis();
            let chunk_len = match (&self.chunks, list_axis) {
                (Some(c), Some(a)) => c.get(a).copied(),
                _ => None,
            };
            let (reads, list_axis) = plan.reads(dims, chunk_len, elem_size);
            let read_shape = plan.read_shape();
            let file = &*self.file;
            let addr = self.addr;
            // Everything below touches only Rust data: release the GIL.
            let read = || -> Result<Elements, ReadError> {
                let ds = file.dataset_at(addr)?;
                let mut blocks = Vec::with_capacity(reads.len());
                for read in reads {
                    let raw = ds.read_selection(&read.sel)?;
                    let mut shape = read.shape;
                    let want = shape.iter().product::<usize>() * elem_size;
                    if raw.len() != want {
                        return Err(ReadError::Other(format!(
                            "read {} bytes, expected {want}",
                            raw.len()
                        )));
                    }
                    let raw = match (&read.pick, list_axis) {
                        (Some(pick), Some(axis)) => {
                            let kept = select::gather_along(&raw, &shape, axis, pick, elem_size);
                            shape[axis] = pick.len();
                            kept
                        }
                        _ => raw,
                    };
                    blocks.push((raw, shape));
                }
                // Several blocks only for a list index: join their bytes
                // (every byte of every element, padding included) along
                // that axis before anything becomes numpy.
                let raw = match (blocks.len(), list_axis) {
                    (1, _) => blocks.pop().expect("one block").0,
                    (_, Some(axis)) => select::join_along(&blocks, axis, elem_size),
                    _ => {
                        return Err(ReadError::Other(
                            "several reads without an index list".into(),
                        ));
                    }
                };
                if !vl {
                    return Ok(Elements::Bytes(raw));
                }
                let sb = file.superblock();
                let n = read_shape.iter().product();
                resolve_vl(
                    file.as_bytes(),
                    &raw,
                    n,
                    sb.offset_size,
                    sb.length_size,
                    unit,
                )
                .map(Elements::Vl)
                .map_err(ReadError::Other)
            };
            let data = py
                .detach(|| {
                    std::panic::catch_unwind(std::panic::AssertUnwindSafe(read))
                        .unwrap_or_else(|p| Err(ReadError::Panic(crate::panic_text(&*p))))
                })
                .map_err(|e| e.into_py(&self.path))?;
            let joined = conv.to_array(py, data, &read_shape, false)?;
            // Drop the axes indexed by an integer (length 1 in the blocks).
            let mut shape = out_shape.clone();
            if let crate::convert::Layout::Subarray(sub) = &conv.layout {
                shape.extend_from_slice(sub);
            }
            joined.call_method1("reshape", (PyTuple::new(py, shape)?,))?
        };

        let arr = select_fields(py, arr, &plan.fields)?;
        if plan.scalar {
            return arr.get_item(PyTuple::empty(py));
        }
        Ok(arr)
    }
}

/// An error from the read closure, turned into a Python error with the GIL.
enum ReadError {
    Lib(clawhdf5_rs::Error),
    Other(String),
    Panic(String),
}

impl From<clawhdf5_rs::Error> for ReadError {
    fn from(e: clawhdf5_rs::Error) -> Self {
        ReadError::Lib(e)
    }
}

impl ReadError {
    fn into_py(self, path: &str) -> PyErr {
        match self {
            ReadError::Lib(e) => to_py_err(e),
            ReadError::Other(msg) => PyValueError::new_err(format!("{}: {msg}", node::name(path))),
            ReadError::Panic(msg) => crate::InternalError::new_err(format!(
                "{}: clawhdf5 internal error (please report it): {msg}",
                node::name(path)
            )),
        }
    }
}

/// Keep only the named compound fields, as h5py's `ds['x']` / `ds['x', 'y']`.
fn select_fields<'py>(
    py: Python<'py>,
    arr: Bound<'py, PyAny>,
    fields: &[String],
) -> PyResult<Bound<'py, PyAny>> {
    if fields.is_empty() {
        return Ok(arr);
    }
    let names = arr.getattr("dtype")?.getattr("names")?;
    if names.is_none() {
        return Err(PyValueError::new_err(
            "Field names only allowed for compound types",
        ));
    }
    let names: Vec<String> = names.extract()?;
    for f in fields {
        if !names.contains(f) {
            return Err(PyValueError::new_err(format!(
                "Field {f} does not appear in this type."
            )));
        }
    }
    let np = py.import("numpy")?;
    if let [one] = fields {
        return np.call_method1("ascontiguousarray", (arr.get_item(one)?,));
    }
    let picked = arr.get_item(PyList::new(py, fields)?)?;
    py.import("numpy.lib.recfunctions")?
        .call_method1("repack_fields", (picked,))
}

#[pymethods]
impl PyDataset {
    /// The shape of the dataset (`None` for an empty/null dataspace).
    #[getter]
    fn shape<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        match &self.shape {
            Some(s) => Ok(PyTuple::new(py, s)?.into_any()),
            None => Ok(py.None().into_bound(py)),
        }
    }

    /// The maximum shape (`None` per unlimited dimension), like h5py.
    #[getter]
    fn maxshape<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        crate::no_panic(|| {
            let Some(shape) = &self.shape else {
                return Ok(py.None().into_bound(py));
            };
            let max = self
                .file
                .dataset_at(self.addr)
                .and_then(|ds| ds.max_dimensions())
                .map_err(to_py_err)?
                .unwrap_or_else(|| shape.clone());
            let items: Vec<Option<u64>> = max
                .into_iter()
                .map(|d| (d != u64::MAX).then_some(d))
                .collect();
            Ok(PyTuple::new(py, items)?.into_any())
        })
    }

    /// The dataset's numpy dtype, as h5py reports it.
    #[getter]
    fn dtype<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        Ok(self.converter()?.dtype.bind(py).clone())
    }

    #[getter]
    fn ndim(&self) -> usize {
        self.shape.as_ref().map_or(0, Vec::len)
    }

    /// Number of elements (`None` for an empty/null dataspace, as h5py).
    #[getter]
    fn size(&self) -> Option<u64> {
        self.shape.as_ref().map(|s| s.iter().product())
    }

    /// The dataset's full name, e.g. `/group/data`.
    #[getter]
    fn name(&self) -> String {
        node::name(&self.path)
    }

    /// The dataset's attributes (read-only, dict-like).
    #[getter]
    fn attrs(&self) -> PyResult<PyAttrs> {
        PyAttrs::read(Arc::clone(&self.file), self.addr, &self.path)
    }

    /// Read with h5py indexing: integers, slices with positive steps,
    /// `...`, one increasing list of integers, and compound field names.
    /// A selection whose bounding box covers at most half the dataset reads
    /// only the chunks (or contiguous rows) it overlaps.
    fn __getitem__<'py>(
        &self,
        py: Python<'py>,
        key: &Bound<'py, PyAny>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let Some(dims) = &self.shape else {
            let is_empty_tuple = key.cast::<PyTuple>().is_ok_and(|t| t.is_empty());
            let is_ellipsis = key.is_instance_of::<pyo3::types::PyEllipsis>();
            if is_empty_tuple || is_ellipsis {
                let empty = PyEmpty::new(self.converter()?.dtype.clone_ref(py));
                return Ok(empty.into_pyobject(py)?.into_any());
            }
            return Err(PyValueError::new_err("Empty datasets cannot be sliced"));
        };
        let plan = select::parse(key, dims)?;
        self.read_plan(py, &plan)
    }

    /// `numpy.asarray(ds)` reads the whole dataset.
    #[pyo3(signature = (dtype=None, copy=None))]
    fn __array__<'py>(
        &self,
        py: Python<'py>,
        dtype: Option<&Bound<'py, PyAny>>,
        copy: Option<bool>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let _ = copy; // every read is a fresh array
        let Some(dims) = &self.shape else {
            return Err(PyValueError::new_err("an empty dataset has no array value"));
        };
        let ellipsis = pyo3::types::PyEllipsis::get(py).to_owned().into_any();
        let plan = select::parse(&ellipsis, dims)?;
        let arr = self.read_plan(py, &plan)?;
        match dtype {
            Some(dt) => arr.call_method1("astype", (dt,)),
            None => Ok(arr),
        }
    }

    fn __len__(&self) -> PyResult<usize> {
        match self.shape.as_deref() {
            Some([first, ..]) => Ok(*first as usize),
            _ => Err(PyTypeError::new_err(
                "Attempt to take len() of scalar dataset",
            )),
        }
    }

    fn __repr__(&self, py: Python<'_>) -> String {
        let dtype = match &self.conv {
            Ok(c) => c
                .dtype
                .bind(py)
                .str()
                .map(|s| s.to_string())
                .unwrap_or_default(),
            Err(_) => format!("{:?}", self.datatype),
        };
        let shape = match &self.shape {
            Some(s) => format!("{s:?}"),
            None => "None".to_string(),
        };
        format!(
            "<HDF5 dataset \"{}\": shape {shape}, type \"{dtype}\">",
            node::name(&self.path)
        )
    }
}