re_hdf5 0.36.0

//! HDF5 → Arrow conversion: dtype mapping, dataset reads, attribute mapping.
//!
//! Per-row representation: every dataset row is exactly one instance. The row's
//! inner Arrow value type encodes the per-row payload, and a component column is
//! always an outer `ListArray` with one element per row (offsets step by 1):
//!
//! - 0-D / 1-D dataset → each row is one scalar;
//! - 2-D `[N, K]` → each row is one `FixedSizeList<K>`;
//! - 3-D+ `[N, d1, …, dk]` → each row is one `List` blob of `d1*…*dk` values
//!   in row-major order (no shape metadata is emitted; consumers recover the
//!   shape via `list_datasets`).
//!
//! This keeps a dataset's standalone-component shape identical to its
//! struct-field shape ([`crate::Hdf5Config::use_structs`]): both come from
//! [`read_row_values`].

use std::sync::Arc;

use arrow::array::{
    Array as _, ArrayRef, FixedSizeListArray, Float32Array, Float64Array, Int8Array, Int16Array,
    Int32Array, Int64Array, ListArray, StringArray, StructArray, UInt8Array, UInt16Array,
    UInt32Array, UInt64Array,
};
use arrow::buffer::{OffsetBuffer, ScalarBuffer};
use arrow::datatypes::{Field, Fields};
use hdf5_pure::DType;
use re_sdk_types::{ComponentDescriptor, ComponentIdentifier};

use crate::config::IndexType;
use crate::error::Hdf5Error;
use crate::walk::{DatasetDesc, H5Path};

/// The element types this reader maps to Arrow: the 10 numeric widths + strings.
pub(crate) fn supported_dtype(dtype: &DType) -> bool {
    is_numeric_dtype(dtype) | matches!(dtype, DType::String | DType::VariableLengthString)
}

pub(crate) fn is_numeric_dtype(dtype: &DType) -> bool {
    matches!(
        dtype,
        DType::I8
            | DType::I16
            | DType::I32
            | DType::I64
            | DType::U8
            | DType::U16
            | DType::U32
            | DType::U64
            | DType::F32
            | DType::F64
    )
}

/// Element type of an HDF5 dataset, as exposed by [`crate::DatasetInfo::dtype`].
///
/// An owned mirror of the subset of `hdf5-pure`'s `DType` this reader maps, so the
/// young `hdf5-pure` type never leaks into the public API. `Display` yields
/// numpy-style names (`"uint8"`, `"float64"`, `"string"`) — deliberately not
/// `DType`'s `Display`, which prints short names (`"u8"`, `"f64"`).
///
/// `#[non_exhaustive]`: the mapped set is expected to grow (compound, enum, array,
/// …), so adding variants stays non-breaking.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum DatasetDtype {
    Int8,
    Int16,
    Int32,
    Int64,
    UInt8,
    UInt16,
    UInt32,
    UInt64,
    Float32,
    Float64,
    String,

    /// An element type this reader does not map (compound, enum, array, reference,
    /// opaque, …).
    Unsupported,
}

impl DatasetDtype {
    fn as_numpy_str(self) -> &'static str {
        match self {
            Self::Int8 => "int8",
            Self::Int16 => "int16",
            Self::Int32 => "int32",
            Self::Int64 => "int64",
            Self::UInt8 => "uint8",
            Self::UInt16 => "uint16",
            Self::UInt32 => "uint32",
            Self::UInt64 => "uint64",
            Self::Float32 => "float32",
            Self::Float64 => "float64",
            Self::String => "string",
            Self::Unsupported => "unsupported",
        }
    }
}

impl std::fmt::Display for DatasetDtype {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.as_numpy_str())
    }
}

impl From<&DType> for DatasetDtype {
    fn from(dtype: &DType) -> Self {
        // The listed types and the `_` arm both map to `Unsupported`, but keep them
        // apart: only the listed ones are types we know about.
        #[expect(clippy::match_same_arms)]
        match dtype {
            DType::I8 => Self::Int8,
            DType::I16 => Self::Int16,
            DType::I32 => Self::Int32,
            DType::I64 => Self::Int64,
            DType::U8 => Self::UInt8,
            DType::U16 => Self::UInt16,
            DType::U32 => Self::UInt32,
            DType::U64 => Self::UInt64,
            DType::F32 => Self::Float32,
            DType::F64 => Self::Float64,
            DType::String | DType::VariableLengthString => Self::String,

            //TODO(ab): support compound, enum, array and object-reference types?
            DType::Compound(_)
            | DType::Enum(_)
            | DType::Array(..)
            | DType::ObjectReference
            | DType::Other(_) => Self::Unsupported,

            // `DType` is `#[non_exhaustive]`. A variant added upstream lands here
            // instead of breaking the build.
            _ => Self::Unsupported,
        }
    }
}

/// Open `desc`'s dataset for value reads.
///
/// The returned handle is owned (no borrow of `file`) and carries its own chunk
/// cache, so callers streaming a dataset window-by-window should keep it open
/// across windows.
pub(crate) fn open_dataset(
    file: &hdf5_pure::File,
    desc: &DatasetDesc,
) -> Result<hdf5_pure::Dataset, Hdf5Error> {
    file.dataset(&desc.path.as_hdf5())
        .map_err(|source| Hdf5Error::read_dataset(&desc.path, source))
}

/// Estimated bytes of one row (inner elements × element size), for sizing
/// row windows. Strings use a fixed per-element guess.
//TODO(RR-5285): derive string estimates from the heap references' exact lengths
// instead of the guess — long strings currently oversize windows proportionally.
pub(crate) fn row_byte_estimate(desc: &DatasetDesc) -> usize {
    let element_bytes = match desc.dtype {
        DType::I8 | DType::U8 => 1,
        DType::I16 | DType::U16 => 2,
        DType::I32 | DType::U32 | DType::F32 => 4,
        DType::I64 | DType::U64 | DType::F64 => 8,
        // Strings and anything else: a guess is fine, this only sizes windows.
        _ => 16,
    };
    #[expect(clippy::cast_possible_truncation)]
    let elements_per_row = desc.shape[1..].iter().product::<u64>() as usize;
    elements_per_row.max(1).saturating_mul(element_bytes)
}

/// Read the row window `[start_row, start_row + num_rows)` of a dataset into
/// its per-row value array (see module docs) and the matching `Field` (named
/// after the dataset leaf).
///
/// The single core builder used by both the standalone and the struct-packed
/// emit paths. A 0-D dataset counts as one row (pass `(0, 1)`).
pub(crate) fn read_row_values(
    dataset: &hdf5_pure::Dataset,
    desc: &DatasetDesc,
    start_row: usize,
    num_rows: usize,
) -> Result<(Field, ArrayRef), Hdf5Error> {
    re_tracing::profile_function!();

    let flat = read_flat_values(dataset, desc, start_row, num_rows)?;

    let values: ArrayRef = match desc.shape.len() {
        // Each row is one scalar (a 0-D dataset yields a single row).
        0 | 1 => flat,

        // Each row is one fixed-size list of `K` (`K` lives in the Arrow type).
        2 => {
            let k = i32::try_from(desc.shape[1]).map_err(|_err| Hdf5Error::ListTooLong {
                length: desc.shape[1],
            })?;
            let item_field = Arc::new(Field::new("item", flat.data_type().clone(), true));
            Arc::new(FixedSizeListArray::try_new(item_field, k, flat, None)?)
        }

        // Each row is one blob of the row's `d1*…*dk` raw row-major values
        // (one instance per frame for image-like datasets, not one per pixel).
        _ => {
            #[expect(clippy::cast_possible_truncation)]
            let per_row = desc.shape[1..].iter().product::<u64>() as usize;
            let item_field = Arc::new(Field::new("item", flat.data_type().clone(), true));
            let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(per_row, num_rows));
            Arc::new(ListArray::try_new(item_field, offsets, flat, None)?)
        }
    };

    let field = Field::new(desc.name(), values.data_type().clone(), true);
    Ok((field, values))
}

/// Read a dataset's raw values for a row window into a flat (row-major) Arrow array.
///
/// The windowed `read_*_rows` calls delegate to a whole read when the window
/// covers every row, so a full-range window costs no more than one.
fn read_flat_values(
    dataset: &hdf5_pure::Dataset,
    desc: &DatasetDesc,
    start_row: usize,
    num_rows: usize,
) -> Result<ArrayRef, Hdf5Error> {
    let read_err = |source| Hdf5Error::read_dataset(&desc.path, source);
    let (start, count) = (start_row as u64, num_rows as u64);

    macro_rules! read {
        ($rows_fn:ident, $array:ident) => {
            Arc::new($array::from(
                dataset.$rows_fn(start, count).map_err(read_err)?,
            )) as ArrayRef
        };
    }

    Ok(match &desc.dtype {
        DType::I8 => read!(read_i8_rows, Int8Array),
        DType::I16 => read!(read_i16_rows, Int16Array),
        DType::I32 => read!(read_i32_rows, Int32Array),
        DType::I64 => read!(read_i64_rows, Int64Array),
        DType::U8 => read!(read_u8_rows, UInt8Array),
        DType::U16 => read!(read_u16_rows, UInt16Array),
        DType::U32 => read!(read_u32_rows, UInt32Array),
        DType::U64 => read!(read_u64_rows, UInt64Array),
        DType::F32 => read!(read_f32_rows, Float32Array),
        DType::F64 => read!(read_f64_rows, Float64Array),
        DType::String | DType::VariableLengthString => {
            read!(read_string_rows, StringArray)
        }
        unsupported => {
            // Planning filters unsupported dtypes out before any read.
            return Err(Hdf5Error::UnsupportedElementType {
                dtype: unsupported.to_string(),
            });
        }
    })
}

/// Wrap a dataset's per-row values into its standalone component: one
/// component per dataset, one row per dataset row.
pub(crate) fn values_to_component(
    name: &str,
    field: Field,
    values: ArrayRef,
) -> Result<(ComponentDescriptor, ListArray), Hdf5Error> {
    let list = wrap_one_per_row(field.with_name("item"), values)?;
    Ok((partial_descriptor(name)?, list))
}

/// Read a whole (small) dataset as its standalone component; used for the 0-D
/// static scalars.
pub(crate) fn read_dataset_to_list(
    file: &hdf5_pure::File,
    desc: &DatasetDesc,
) -> Result<(ComponentDescriptor, ListArray), Hdf5Error> {
    let dataset = open_dataset(file, desc)?;
    #[expect(clippy::cast_possible_truncation)]
    let num_rows = desc.shape.first().copied().unwrap_or(1) as usize;
    let (field, values) = read_row_values(&dataset, desc, 0, num_rows)?;
    values_to_component(desc.name(), field, values)
}

/// Struct emit path: assemble the per-dataset row values into one
/// `List<Struct>` component named `data`, one struct per row.
pub(crate) fn build_struct_component(
    columns: Vec<(Field, ArrayRef)>,
) -> Result<(ComponentDescriptor, ListArray), Hdf5Error> {
    let (fields, arrays): (Vec<_>, Vec<_>) = columns
        .into_iter()
        .map(|(field, array)| (Arc::new(field), array))
        .unzip();

    let struct_array = StructArray::try_new(Fields::from(fields), arrays, None)?;
    let item_field = Field::new("item", struct_array.data_type().clone(), true);
    let list = wrap_one_per_row(item_field, Arc::new(struct_array))?;
    Ok((ComponentDescriptor::partial("data"), list))
}

/// Wrap a per-row value array into the outer one-element-per-row `ListArray`.
fn wrap_one_per_row(item_field: Field, values: ArrayRef) -> Result<ListArray, Hdf5Error> {
    let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(1_usize, values.len()));
    Ok(ListArray::try_new(
        Arc::new(item_field),
        offsets,
        values,
        None,
    )?)
}

/// Read the 1-D index dataset and scale its values to nanoseconds.
///
/// Integer dtypes are cast to `i64` and then multiplied. Float dtypes are
/// multiplied in `f64` first and rounded to `i64` after, preserving sub-second
/// precision (`1.5` s → `1_500_000_000` ns) for float-seconds indices.
pub(crate) fn read_index_to_ns(
    file: &hdf5_pure::File,
    path: &H5Path,
    index_type: IndexType,
) -> Result<ScalarBuffer<i64>, Hdf5Error> {
    re_tracing::profile_function!();

    let read_err = |source| Hdf5Error::read_dataset(path, source);
    let dataset = file.dataset(&path.as_hdf5()).map_err(read_err)?;
    let multiplier = index_type.ns_multiplier();

    fn scale_ints<T: Into<i64>>(values: Vec<T>, multiplier: i64) -> Vec<i64> {
        values
            .into_iter()
            .map(|value| value.into() * multiplier)
            .collect()
    }

    #[expect(clippy::cast_possible_truncation)]
    fn scale_floats<T: Into<f64>>(values: Vec<T>, multiplier: i64) -> Vec<i64> {
        #[expect(clippy::cast_precision_loss)]
        let multiplier = multiplier as f64;
        values
            .into_iter()
            .map(|value| (value.into() * multiplier).round() as i64)
            .collect()
    }

    let values: Vec<i64> = match dataset.dtype().map_err(read_err)? {
        DType::I8 => scale_ints(dataset.read_i8().map_err(read_err)?, multiplier),
        DType::I16 => scale_ints(dataset.read_i16().map_err(read_err)?, multiplier),
        DType::I32 => scale_ints(dataset.read_i32().map_err(read_err)?, multiplier),
        DType::I64 => scale_ints(dataset.read_i64().map_err(read_err)?, multiplier),
        DType::U8 => scale_ints(dataset.read_u8().map_err(read_err)?, multiplier),
        DType::U16 => scale_ints(dataset.read_u16().map_err(read_err)?, multiplier),
        DType::U32 => scale_ints(dataset.read_u32().map_err(read_err)?, multiplier),
        #[expect(clippy::cast_possible_wrap)]
        DType::U64 => dataset
            .read_u64()
            .map_err(read_err)?
            .into_iter()
            .map(|value| value as i64 * multiplier)
            .collect(),
        DType::F32 => scale_floats(dataset.read_f32().map_err(read_err)?, multiplier),
        DType::F64 => scale_floats(dataset.read_f64().map_err(read_err)?, multiplier),
        non_numeric => {
            // Planning validates this before any read; kept as a real error for safety.
            return Err(Hdf5Error::IndexNotNumeric {
                path: path.to_string(),
                dtype: non_numeric.to_string(),
            });
        }
    };

    Ok(ScalarBuffer::from(values))
}

/// The attribute types [`attr_to_component`] maps to Arrow, in the same order.
///
/// Planning drops the rest, so an unmapped attribute never reaches conversion.
/// `AttrValue` is `#[non_exhaustive]`, so a variant added upstream is unsupported
/// until it is listed here and in [`attr_to_component`].
pub(crate) fn supported_attr(value: &hdf5_pure::AttrValue) -> bool {
    use hdf5_pure::AttrValue;

    matches!(
        value,
        AttrValue::F64(_)
            | AttrValue::I32(_)
            | AttrValue::I64(_)
            | AttrValue::U32(_)
            | AttrValue::U64(_)
            | AttrValue::String(_)
            | AttrValue::AsciiString(_)
            | AttrValue::F64Array(_)
            | AttrValue::I64Array(_)
            | AttrValue::U64Array(_)
            | AttrValue::StringArray(_)
            | AttrValue::AsciiStringArray(_)
            | AttrValue::VarLenAsciiArray(_)
    )
}

/// Map one HDF5 attribute to a single-row static component.
///
/// Scalar variants become a one-scalar `List<primitive>` row; array variants a
/// single row whose value is a `FixedSizeList<L>` — the same one-per-row rule
/// as datasets.
///
/// Errors for an attribute type [`supported_attr`] rejects.
pub(crate) fn attr_to_component(
    name: &str,
    value: &hdf5_pure::AttrValue,
) -> Result<(ComponentDescriptor, ListArray), Hdf5Error> {
    use hdf5_pure::AttrValue;

    let values: ArrayRef = match value {
        AttrValue::F64(value) => Arc::new(Float64Array::from(vec![*value])),
        AttrValue::I32(value) => Arc::new(Int32Array::from(vec![*value])),
        AttrValue::I64(value) => Arc::new(Int64Array::from(vec![*value])),
        AttrValue::U32(value) => Arc::new(UInt32Array::from(vec![*value])),
        AttrValue::U64(value) => Arc::new(UInt64Array::from(vec![*value])),
        AttrValue::String(value) | AttrValue::AsciiString(value) => {
            Arc::new(StringArray::from(vec![value.as_str()]))
        }
        AttrValue::F64Array(values) => {
            one_row_fixed_size_list(Arc::new(Float64Array::from(values.clone())))?
        }
        AttrValue::I64Array(values) => {
            one_row_fixed_size_list(Arc::new(Int64Array::from(values.clone())))?
        }
        AttrValue::U64Array(values) => {
            one_row_fixed_size_list(Arc::new(UInt64Array::from(values.clone())))?
        }
        AttrValue::StringArray(values)
        | AttrValue::AsciiStringArray(values)
        | AttrValue::VarLenAsciiArray(values) => one_row_fixed_size_list(Arc::new(
            StringArray::from_iter_values(values.iter().map(String::as_str)),
        ))?,

        // Planning drops what `supported_attr` rejects, so this arm means the two
        // lists have drifted apart.
        _ => {
            return Err(Hdf5Error::UnsupportedAttributeType {
                name: name.to_owned(),
                type_name: value.type_name().to_owned(),
            });
        }
    };

    let item_field = Field::new("item", values.data_type().clone(), true);
    let list = wrap_one_per_row(item_field, values)?;
    Ok((partial_descriptor(name)?, list))
}

/// Wrap a length-`L` array into a length-1 `FixedSizeList<L>` array (one row).
fn one_row_fixed_size_list(inner: ArrayRef) -> Result<ArrayRef, Hdf5Error> {
    let len = i32::try_from(inner.len()).map_err(|_err| Hdf5Error::ListTooLong {
        length: inner.len() as u64,
    })?;
    let item_field = Arc::new(Field::new("item", inner.data_type().clone(), true));
    Ok(Arc::new(FixedSizeListArray::try_new(
        item_field, len, inner, None,
    )?))
}

fn partial_descriptor(name: &str) -> Result<ComponentDescriptor, Hdf5Error> {
    let component = ComponentIdentifier::try_new(name)
        .map_err(|source| Hdf5Error::invalid_component_name(name, source))?;
    Ok(ComponentDescriptor::partial(component))
}