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()
    }

    /// The shape of the result before the integer-indexed axes are dropped
    /// (they have length 1 here): the shape of the joined reads.
    pub fn read_shape(&self) -> Vec<usize> {
        self.axes.iter().map(|a| a.len() as usize).collect()
    }

    /// The axis indexed by a list, if any.
    pub fn list_axis(&self) -> Option<usize> {
        self.axes.iter().position(|a| matches!(a, Axis::List(_)))
    }

    /// 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; those are joined along `list_axis`
    /// afterwards (`join_along`), after keeping each read's `pick` rows.
    ///
    /// A list is read in groups, each one hyperslab over a stretch of the
    /// axis, not once per index: every read decodes the chunks it touches
    /// (and lists the dataset's chunks), so a read per run of indices decoded
    /// the same chunk again and again. For a chunked dataset (`chunk_len` is
    /// the chunk's length along the list axis) a group ends only where a
    /// whole chunk holds no selected index, so no chunk is decoded twice or
    /// without need. Otherwise a group ends at a gap of more than
    /// [`MAX_GAP_BYTES`] of unselected data.
    pub fn reads(
        &self,
        dims: &[u64],
        chunk_len: Option<u64>,
        elem_size: usize,
    ) -> (Vec<Read>, Option<usize>) {
        let list_axis = self.list_axis();
        let groups: Vec<&[u64]> = match list_axis.map(|i| &self.axes[i]) {
            Some(Axis::List(idx)) => {
                let row_bytes = self.row_bytes(elem_size);
                group_indices(idx, |last, next| match chunk_len {
                    Some(c) if c > 0 => next / c <= last / c + 1,
                    _ => (next - last - 1).saturating_mul(row_bytes) <= MAX_GAP_BYTES,
                })
            }
            _ => vec![&[]],
        };
        let mut out = Vec::with_capacity(groups.len());
        for group in groups {
            let (first, span) = match (group.first(), group.last()) {
                (Some(&f), Some(&l)) => (f, l - f + 1),
                _ => (0, 0),
            };
            let pick = (span != group.len() as u64)
                .then(|| group.iter().map(|&i| (i - first) as usize).collect());
            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(_) => (first, 1, span),
                };
                start.push(s);
                // A stride only matters between blocks; keep it >= 1.
                stride.push(if c <= 1 { 1 } else { st });
                count.push(c);
            }
            let 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(Read { sel, shape, pick });
        }
        (out, list_axis)
    }

    /// Bytes of one step along the list axis within a read's bounding box.
    fn row_bytes(&self, elem_size: usize) -> u64 {
        self.axes
            .iter()
            .map(|a| match a {
                Axis::Slice { step, count, .. } if *count > 0 => (count - 1) * step + 1,
                _ => 1,
            })
            .fold(elem_size as u64, u64::saturating_mul)
    }
}

/// Unselected data a read of a non-chunked dataset copies through rather
/// than start another read.
pub(crate) const MAX_GAP_BYTES: u64 = 64 * 1024;

/// One hyperslab read of a selection.
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct Read {
    pub sel: Selection,
    /// The block's shape (index axes at length 1).
    pub shape: Vec<usize>,
    /// For a list: the positions along the list axis, within the block, to
    /// keep (`None`: all of them).
    pub pick: Option<Vec<usize>>,
}

/// Split increasing indices into groups; `joins(last, next)` says whether
/// `next` extends the group whose last index is `last`.
fn group_indices(idx: &[u64], joins: impl Fn(u64, u64) -> bool) -> Vec<&[u64]> {
    let mut groups = Vec::new();
    let mut from = 0;
    for k in 1..idx.len() {
        if !joins(idx[k - 1], idx[k]) {
            groups.push(&idx[from..k]);
            from = k;
        }
    }
    if from < idx.len() {
        groups.push(&idx[from..]);
    }
    groups
}

/// Keep the elements at positions `pick` along `axis` of a row-major block.
pub(crate) fn gather_along(
    bytes: &[u8],
    shape: &[usize],
    axis: usize,
    pick: &[usize],
    elem_size: usize,
) -> Vec<u8> {
    let outer: usize = shape[..axis].iter().product();
    let inner: usize = shape[axis + 1..].iter().product::<usize>() * elem_size;
    let len = shape[axis];
    let mut out = Vec::with_capacity(outer * pick.len() * inner);
    for o in 0..outer {
        for &p in pick {
            let at = (o * len + p) * inner;
            out.extend_from_slice(&bytes[at..at + inner]);
        }
    }
    out
}

/// Join row-major blocks of `elem_size`-byte elements whose shapes differ
/// only along `axis` into one buffer, in order along that axis. Whole
/// elements are copied, so compound padding keeps the bytes that were read.
pub(crate) fn join_along(
    blocks: &[(Vec<u8>, Vec<usize>)],
    axis: usize,
    elem_size: usize,
) -> Vec<u8> {
    let Some((_, first)) = blocks.first() else {
        return Vec::new();
    };
    let outer: usize = first[..axis].iter().product();
    let inner: usize = first[axis + 1..].iter().product::<usize>() * elem_size;
    let total: usize = blocks.iter().map(|(_, s)| s[axis]).sum();
    let mut out = vec![0u8; outer * total * inner];
    let mut at = 0;
    for (bytes, shape) in blocks {
        let len = shape[axis] * inner;
        for o in 0..outer {
            let dst = (o * total) * inner + at;
            out[dst..dst + len].copy_from_slice(&bytes[o * len..(o + 1) * len]);
        }
        at += len;
    }
    out
}

/// 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 indices_group_by_chunk() {
        let chunked = |c: u64| move |last: u64, next: u64| next / c <= last / c + 1;
        // Chunks of 10: 3, 5 and 15 are in neighbouring chunks; 42 skips two.
        let idx = [3, 5, 15, 42, 43, 99];
        assert_eq!(
            group_indices(&idx, chunked(10)),
            vec![&[3, 5, 15][..], &[42, 43], &[99]]
        );
        assert_eq!(group_indices(&[], chunked(10)), Vec::<&[u64]>::new());
    }

    #[test]
    fn a_list_reads_once_per_group() {
        let plan = Plan {
            axes: vec![
                Axis::List(vec![0, 2, 3, 40]),
                Axis::Slice {
                    start: 0,
                    step: 1,
                    count: 5,
                },
            ],
            fields: vec![],
            scalar: false,
        };
        // Chunks of 8 rows: rows 0-3 are one read, row 40 another.
        let (reads, axis) = plan.reads(&[50, 5], Some(8), 4);
        assert_eq!(axis, Some(0));
        assert_eq!(reads.len(), 2);
        assert_eq!(reads[0].shape, vec![4, 5]);
        assert_eq!(reads[0].pick, Some(vec![0, 2, 3]));
        assert_eq!(reads[1].shape, vec![1, 5]);
        assert_eq!(reads[1].pick, None);
        // Not chunked: a gap under MAX_GAP_BYTES is read through.
        let (reads, _) = plan.reads(&[50, 5], None, 4);
        assert_eq!(reads.len(), 1);
        assert_eq!(reads[0].pick, Some(vec![0, 2, 3, 40]));
    }

    #[test]
    fn gather_keeps_picked_rows() {
        // A 2x3 block of 1-byte elements; keep columns 0 and 2.
        let block = [1, 2, 3, 4, 5, 6];
        assert_eq!(
            gather_along(&block, &[2, 3], 1, &[0, 2], 1),
            vec![1, 3, 4, 6]
        );
    }

    #[test]
    fn blocks_join_along_the_list_axis() {
        // Two 2x1 and 2x2 blocks of 1-byte elements, joined along axis 1.
        let a = (vec![1, 2], vec![2, 1]);
        let b = (vec![3, 4, 5, 6], vec![2, 2]);
        assert_eq!(join_along(&[a, b], 1, 1), vec![1, 3, 4, 2, 5, 6]);
        // Along axis 0 it is concatenation; 2-byte elements stay whole.
        let a = (vec![1, 2, 3, 4], vec![1, 2]);
        let b = (vec![5, 6, 7, 8], vec![1, 2]);
        assert_eq!(join_along(&[a, b], 0, 2), vec![1, 2, 3, 4, 5, 6, 7, 8]);
    }

    #[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], None, 8);
        assert_eq!(list, None);
        assert_eq!(
            reads,
            vec![Read {
                sel: Selection::All,
                shape: vec![4, 3],
                pick: None
            }]
        );
    }

    #[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], None, 8);
        assert_eq!(
            reads,
            vec![Read {
                sel: Selection::Hyperslab {
                    start: vec![2, 1],
                    stride: vec![1, 3],
                    count: vec![1, 2],
                    block: vec![1, 1],
                },
                shape: vec![1, 2],
                pick: None
            }]
        );
        assert_eq!(plan.out_shape(), vec![2]);
    }
}