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

//! h5py-style indexing (`ds[1, 2:10:3, ...]`) mapped onto hyperslab
//! selections, so only the selected elements are read.
//!
//! The rules and error messages follow h5py's `selections.py`: integers
//! (negative from the end) drop their axis, slices must have a positive
//! step, one `Ellipsis` fills the unmentioned axes, a single increasing list
//! of integers may index one axis, and strings name compound fields.
//! Everything else (`None`/`np.newaxis`, boolean masks, several index lists)
//! is refused with the error h5py gives.

use clawhdf5_format::selection::Selection;
use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyEllipsis, PySlice, PyString, PyTuple};

/// The selection along one axis.
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum Axis {
    /// A single index: the axis is dropped from the result.
    Index(u64),
    /// `start, start + step, ...`, `count` of them.
    Slice { start: u64, step: u64, count: u64 },
    /// Increasing, distinct indices.
    List(Vec<u64>),
}

impl Axis {
    fn len(&self) -> u64 {
        match self {
            Axis::Index(_) => 1,
            Axis::Slice { count, .. } => *count,
            Axis::List(v) => v.len() as u64,
        }
    }
}

/// A parsed index expression.
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct Plan {
    /// One entry per dataset axis.
    pub axes: Vec<Axis>,
    /// Compound field names to keep (empty: all).
    pub fields: Vec<String>,
    /// For a scalar dataset: `ds[()]` gives a scalar, `ds[...]` a 0-d array.
    /// For other datasets: every axis was an integer, so h5py gives a scalar.
    pub scalar: bool,
}

impl Plan {
    /// The shape of the result.
    pub fn out_shape(&self) -> Vec<usize> {
        self.axes
            .iter()
            .filter(|a| !matches!(a, Axis::Index(_)))
            .map(|a| a.len() as usize)
            .collect()
    }

    /// Whether the selection is empty.
    pub fn is_empty(&self) -> bool {
        self.axes.iter().any(|a| a.len() == 0)
    }

    /// The hyperslab reads that make up this selection, each with the shape
    /// of its block (index axes kept at length 1). More than one only when an
    /// axis is indexed by a list: one read per run of consecutive indices,
    /// concatenated along `list_axis` afterwards.
    pub fn reads(&self, dims: &[u64]) -> (Vec<(Selection, Vec<usize>)>, Option<usize>) {
        let list_axis = self.axes.iter().position(|a| matches!(a, Axis::List(_)));
        let runs: Vec<(u64, u64)> = match list_axis.map(|i| &self.axes[i]) {
            Some(Axis::List(idx)) => consecutive_runs(idx),
            _ => vec![(0, 0)],
        };
        let mut out = Vec::with_capacity(runs.len());
        for (run_start, run_len) in runs {
            let mut start = Vec::with_capacity(dims.len());
            let mut stride = Vec::with_capacity(dims.len());
            let mut count = Vec::with_capacity(dims.len());
            for axis in &self.axes {
                let (s, st, c) = match axis {
                    Axis::Index(i) => (*i, 1, 1),
                    Axis::Slice { start, step, count } => (*start, *step, *count),
                    Axis::List(_) => (run_start, 1, run_len),
                };
                start.push(s);
                // A stride only matters between blocks; keep it >= 1.
                stride.push(if c <= 1 { 1 } else { st });
                count.push(c);
            }
            let block_shape: Vec<usize> = count.iter().map(|&c| c as usize).collect();
            let whole = start.iter().all(|&s| s == 0)
                && stride.iter().all(|&s| s == 1)
                && count.as_slice() == dims;
            let sel = if whole {
                Selection::All
            } else {
                let block = vec![1; dims.len()];
                Selection::Hyperslab {
                    start,
                    stride,
                    count,
                    block,
                }
            };
            out.push((sel, block_shape));
        }
        (out, list_axis)
    }
}

fn consecutive_runs(idx: &[u64]) -> Vec<(u64, u64)> {
    let mut runs: Vec<(u64, u64)> = Vec::new();
    for &i in idx {
        match runs.last_mut() {
            Some((s, n)) if *s + *n == i => *n += 1,
            _ => runs.push((i, 1)),
        }
    }
    runs
}

/// Parse `key` for a dataset of shape `dims`.
pub(crate) fn parse(key: &Bound<'_, PyAny>, dims: &[u64]) -> PyResult<Plan> {
    let items: Vec<Bound<'_, PyAny>> = match key.cast::<PyTuple>() {
        Ok(t) => t.iter().collect(),
        Err(_) => vec![key.clone()],
    };
    let mut fields = Vec::new();
    let mut args = Vec::new();
    for item in items {
        if let Ok(s) = item.cast::<PyString>() {
            fields.push(s.to_str()?.to_owned());
        } else {
            args.push(item);
        }
    }

    if args.iter().any(|a| a.is_none()) {
        return Err(PyTypeError::new_err(
            "Indexing with None (or np.newaxis) is not supported",
        ));
    }

    let rank = dims.len();
    if rank == 0 {
        return match args.as_slice() {
            [] => Ok(Plan {
                axes: vec![],
                fields,
                scalar: true,
            }),
            [a] if a.is_instance_of::<PyEllipsis>() => Ok(Plan {
                axes: vec![],
                fields,
                scalar: false,
            }),
            _ => Err(PyValueError::new_err(
                "Illegal slicing argument for scalar dataspace",
            )),
        };
    }

    // Expand the ellipsis (at most one) to full slices.
    let n_ellipsis = args
        .iter()
        .filter(|a| a.is_instance_of::<PyEllipsis>())
        .count();
    if n_ellipsis > 1 {
        return Err(PyValueError::new_err("Only one ellipsis may be used."));
    }
    let explicit = args.len() - n_ellipsis;
    if explicit > rank {
        return Err(PyValueError::new_err(format!(
            "{explicit} indexing arguments for {rank} dimensions"
        )));
    }
    let py = key.py();
    let mut expanded: Vec<Option<Bound<'_, PyAny>>> = Vec::with_capacity(rank);
    for a in args {
        if a.is_instance_of::<PyEllipsis>() {
            for _ in 0..(rank - explicit) {
                expanded.push(None);
            }
        } else {
            expanded.push(Some(a));
        }
    }
    while expanded.len() < rank {
        expanded.push(None);
    }

    let mut axes = Vec::with_capacity(rank);
    for (arg, &n) in expanded.iter().zip(dims) {
        axes.push(match arg {
            None => Axis::Slice {
                start: 0,
                step: 1,
                count: n,
            },
            Some(a) => parse_axis(py, a, n)?,
        });
    }
    if axes.iter().filter(|a| matches!(a, Axis::List(_))).count() > 1 {
        return Err(PyTypeError::new_err(
            "Only one indexing vector or array is currently allowed for fancy indexing",
        ));
    }
    let scalar = axes.iter().all(|a| matches!(a, Axis::Index(_)));
    Ok(Plan {
        axes,
        fields,
        scalar,
    })
}

fn parse_axis(py: Python<'_>, a: &Bound<'_, PyAny>, n: u64) -> PyResult<Axis> {
    if a.is_none() {
        return Err(PyTypeError::new_err(
            "Indexing with None (or np.newaxis) is not supported",
        ));
    }
    if let Ok(s) = a.cast::<PySlice>() {
        let n_isize = isize::try_from(n)
            .map_err(|_| PyValueError::new_err("dimension too large to slice"))?;
        let ind = s.indices(n_isize)?;
        if ind.step < 1 {
            return Err(PyValueError::new_err(format!(
                "Step must be >= 1 (got {})",
                ind.step
            )));
        }
        // `slicelength` is the number of elements selected, >= 0.
        let count = ind.slicelength as u64;
        let start = if count == 0 { 0 } else { ind.start as u64 };
        return Ok(Axis::Slice {
            start,
            step: ind.step as u64,
            count,
        });
    }
    let np = py.import("numpy")?;
    let is_bool =
        a.is_instance_of::<pyo3::types::PyBool>() || a.is_instance(&np.getattr("bool_")?)?;
    let is_array_like = a.is_instance(&np.getattr("ndarray")?)?
        || a.is_instance_of::<pyo3::types::PyList>()
        || a.is_instance_of::<PyTuple>();
    if !is_bool && !is_array_like && a.hasattr("__index__")? {
        let i: i128 = a.call_method0("__index__")?.extract()?;
        return Ok(Axis::Index(normalize(i, n)?));
    }
    if is_array_like {
        let arr = np.call_method1("asarray", (a,))?;
        let kind: String = arr.getattr("dtype")?.getattr("kind")?.extract()?;
        if kind == "b" {
            return Err(PyTypeError::new_err(
                "Boolean mask indexing is not supported by clawhdf5",
            ));
        }
        let ndim: usize = arr.getattr("ndim")?.extract()?;
        let size: usize = arr.getattr("size")?.extract()?;
        if size > 0 && kind != "i" && kind != "u" {
            return Err(PyTypeError::new_err(
                "Indexing arrays must have integer dtypes",
            ));
        }
        if ndim > 1 {
            return Err(PyTypeError::new_err(
                "Only 1-D integer lists or arrays can be used for fancy indexing",
            ));
        }
        let vals: Vec<i128> = arr.call_method0("tolist")?.extract()?;
        let mut idx = Vec::with_capacity(vals.len());
        for v in vals {
            idx.push(normalize(v, n)?);
        }
        if idx.windows(2).any(|w| w[0] >= w[1]) {
            return Err(PyTypeError::new_err(
                "Indexing elements must be in increasing order",
            ));
        }
        return Ok(Axis::List(idx));
    }
    Err(PyTypeError::new_err(format!(
        "Illegal index type for clawhdf5 datasets: {}",
        a.get_type().name()?
    )))
}

fn normalize(i: i128, n: u64) -> PyResult<u64> {
    let n_i = i128::from(n);
    let j = if i < 0 { i + n_i } else { i };
    if j < 0 || j >= n_i {
        let hi = n_i - 1;
        return Err(PyIndexError::new_err(format!(
            "Index ({i}) out of range for (0-{hi})"
        )));
    }
    Ok(j as u64)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn runs_group_consecutive_indices() {
        assert_eq!(
            consecutive_runs(&[1, 2, 3, 7, 9, 10]),
            vec![(1, 3), (7, 1), (9, 2)]
        );
        assert_eq!(consecutive_runs(&[]), vec![]);
    }

    #[test]
    fn full_selection_reads_everything() {
        let plan = Plan {
            axes: vec![
                Axis::Slice {
                    start: 0,
                    step: 1,
                    count: 4,
                },
                Axis::Slice {
                    start: 0,
                    step: 1,
                    count: 3,
                },
            ],
            fields: vec![],
            scalar: false,
        };
        let (reads, list) = plan.reads(&[4, 3]);
        assert_eq!(list, None);
        assert_eq!(reads, vec![(Selection::All, vec![4, 3])]);
    }

    #[test]
    fn index_and_step_map_to_a_hyperslab() {
        let plan = Plan {
            axes: vec![
                Axis::Index(2),
                Axis::Slice {
                    start: 1,
                    step: 3,
                    count: 2,
                },
            ],
            fields: vec![],
            scalar: false,
        };
        let (reads, _) = plan.reads(&[4, 8]);
        assert_eq!(
            reads,
            vec![(
                Selection::Hyperslab {
                    start: vec![2, 1],
                    stride: vec![1, 3],
                    count: vec![1, 2],
                    block: vec![1, 1],
                },
                vec![1, 2]
            )]
        );
        assert_eq!(plan.out_shape(), vec![2]);
    }
}