laddu_data/io/

mod.rs

use serde::{Deserialize, Serialize};
use std::fmt::Display;
use std::sync::Arc;

use crate::{LadduDataError, LadduDataResult, data::EventBatch, schema::Schema};

/// In-memory event sources and sinks.
pub mod memory;
mod output;
/// Parquet event sources and sinks.
pub mod parquet;
/// ROOT event sources.
pub mod root;

mod source;

pub use output::{OutputMode, OutputPath};

pub(crate) use source::{SourceBuild, SourceBuildOptions, build_source};

/// Adds stable operation/resource context to a source failure.
pub(crate) fn source_error(
    operation: impl AsRef<str>,
    resource: impl Display,
    cause: impl Display,
) -> LadduDataError {
    LadduDataError::Source(format!("{} `{resource}`: {cause}", operation.as_ref()))
}

/// Adds stable operation/resource context to a sink failure.
pub(crate) fn sink_error(
    operation: impl AsRef<str>,
    resource: impl Display,
    cause: impl Display,
) -> LadduDataError {
    LadduDataError::Sink(format!("{} `{resource}`: {cause}", operation.as_ref()))
}

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

    #[test]
    fn source_and_sink_context_have_stable_operation_resource_format() {
        assert!(matches!(
            source_error("read ROOT tree", "events.root::events", "branch failed"),
            LadduDataError::Source(message)
                if message == "read ROOT tree `events.root::events`: branch failed"
        ));
        assert!(matches!(
            sink_error("write Parquet file", "events.parquet", "disk full"),
            LadduDataError::Sink(message)
                if message == "write Parquet file `events.parquet`: disk full"
        ));
    }
}

#[cfg(test)]
mod contract_tests;

#[cfg(feature = "mpi")]
/// Distribution of event I/O across MPI ranks.
#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
pub enum Distribution {
    /// Single-process I/O.
    #[default]
    Serial,
    /// MPI-distributed I/O with explicit rank metadata.
    Mpi {
        /// Zero-based rank.
        rank: usize,
        /// Number of ranks.
        nranks: usize,
        /// Work-partitioning strategy.
        partitioning: Partitioning,
    },
}

#[cfg(feature = "mpi")]
impl Distribution {
    /// Creates serial distribution.
    pub fn serial() -> Self {
        Self::Serial
    }

    /// Creates MPI distribution from a communicator.
    pub fn from_world<C>(world: &C) -> Self
    where
        C: mpi::topology::Communicator,
    {
        Self::Mpi {
            rank: world.rank() as usize,
            nranks: world.size() as usize,
            partitioning: Partitioning::default(),
        }
    }

    /// Returns the current rank.
    pub fn rank(self) -> usize {
        match self {
            Self::Serial => 0,
            Self::Mpi { rank, .. } => rank,
        }
    }

    /// Returns the number of ranks.
    pub fn nranks(self) -> usize {
        match self {
            Self::Serial => 1,
            Self::Mpi { nranks, .. } => nranks,
        }
    }

    /// Returns the partitioning strategy.
    pub fn partitioning(self) -> Partitioning {
        match self {
            Self::Serial => Partitioning::Contiguous,
            Self::Mpi { partitioning, .. } => partitioning,
        }
    }

    /// Returns this distribution with a new partitioning strategy.
    pub fn with_partitioning(self, partitioning: Partitioning) -> Self {
        match self {
            Self::Serial => Self::Serial,
            Self::Mpi { rank, nranks, .. } => Self::Mpi {
                rank,
                nranks,
                partitioning,
            },
        }
    }
}

/// Strategy for partitioning input rows across ranks.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum Partitioning {
    /// Each rank reads a contiguous global row range.
    #[default]
    Contiguous,

    /// Each rank reads whole source fragments, such as files or row groups, round-robin.
    FileGroups,

    /// Rank r keeps rows where global_row % nranks == r.
    /// Deterministic, but usually slower.
    Rows,
}

/// Options controlling event-source reads.
#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
pub struct ReadPlan {
    /// Optional maximum output batch size.
    pub chunk_size: Option<usize>,

    #[cfg(feature = "mpi")]
    /// MPI distribution.
    pub distribution: Distribution,
}

impl ReadPlan {
    /// Creates a serial read plan.
    pub fn serial() -> Self {
        Self::default()
    }

    /// Returns the current rank.
    pub fn rank(&self) -> usize {
        #[cfg(feature = "mpi")]
        {
            self.distribution.rank()
        }

        #[cfg(not(feature = "mpi"))]
        {
            0
        }
    }

    /// Returns the number of ranks.
    pub fn nranks(&self) -> usize {
        #[cfg(feature = "mpi")]
        {
            self.distribution.nranks()
        }

        #[cfg(not(feature = "mpi"))]
        {
            1
        }
    }

    /// Returns whether reads are distributed.
    pub fn is_distributed(&self) -> bool {
        self.nranks() > 1
    }

    /// Returns the low-level fragment-partitioning strategy.
    pub fn fragment_partitioning(&self) -> FragmentPartitioning {
        #[cfg(feature = "mpi")]
        {
            match self.distribution.partitioning() {
                Partitioning::Contiguous => FragmentPartitioning::Contiguous,
                Partitioning::FileGroups => FragmentPartitioning::RoundRobinFragments,
                Partitioning::Rows => FragmentPartitioning::StridedRows,
            }
        }

        #[cfg(not(feature = "mpi"))]
        {
            FragmentPartitioning::Contiguous
        }
    }
}

/// Options controlling event-sink writes.
#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
pub struct WritePlan {
    #[cfg(feature = "mpi")]
    /// MPI distribution.
    pub distribution: Distribution,
}

impl From<ReadPlan> for WritePlan {
    #[cfg_attr(not(feature = "mpi"), allow(unused_variables))]
    fn from(plan: ReadPlan) -> Self {
        Self {
            #[cfg(feature = "mpi")]
            distribution: plan.distribution,
        }
    }
}

impl WritePlan {
    /// Returns the current rank.
    pub fn rank(&self) -> usize {
        #[cfg(feature = "mpi")]
        {
            self.distribution.rank()
        }

        #[cfg(not(feature = "mpi"))]
        {
            0
        }
    }

    /// Returns the number of ranks.
    pub fn nranks(&self) -> usize {
        #[cfg(feature = "mpi")]
        {
            self.distribution.nranks()
        }

        #[cfg(not(feature = "mpi"))]
        {
            1
        }
    }

    /// Returns whether writes are distributed.
    pub fn is_distributed(&self) -> bool {
        self.nranks() > 1
    }
}

/// Low-level assignment of source fragments or rows.
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub enum FragmentPartitioning {
    /// Contiguous global row ranges.
    Contiguous,
    /// Whole fragments assigned round-robin.
    RoundRobinFragments,
    /// Individual rows assigned by global index modulo rank count.
    StridedRows,
}

/// Optional performance and planning capabilities of an [`EventSource`].
#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
pub struct SourceCapabilities {
    /// Exact event count is available cheaply.
    pub exact_len: bool,
    /// Exact weighted total is available cheaply.
    pub exact_weighted_total: bool,
    /// Arbitrary row ranges can be read.
    pub random_access: bool,
    /// Distributed row assignment is deterministic.
    pub deterministic_partitioning: bool,
    /// Filters can be pushed into the source.
    pub predicate_pushdown: bool,
    /// Column projection can be pushed into the source.
    pub projection_pushdown: bool,
    /// Batches can be streamed without full materialization.
    pub streaming: bool,
}

/// Sendable iterator of fallible event batches.
pub type EventBatchIter = Box<dyn Iterator<Item = LadduDataResult<EventBatch>> + Send>;

/// Thread-safe producer of schema-compatible event batches.
///
/// Repeated reads with the same plan must replay identical ordered rows and
/// contents. Batch boundaries may vary. This guarantees that separately
/// collected row-data columns, expression results, and weights stay aligned.
pub trait EventSource: Send + Sync {
    /// Returns the source schema.
    ///
    /// # Errors
    ///
    /// Returns [`LadduDataError`] when source metadata cannot be read or
    /// interpreted as a logical schema.
    fn schema(&self) -> LadduDataResult<Arc<Schema>>;

    /// Returns optional source capabilities.
    fn capabilities(&self) -> SourceCapabilities {
        SourceCapabilities::default()
    }

    /// Returns the exact event count when cheaply available.
    ///
    /// # Errors
    ///
    /// Returns [`LadduDataError`] when the source cannot read the metadata
    /// needed to determine its event count.
    fn num_events(&self) -> LadduDataResult<Option<u64>> {
        Ok(None)
    }

    /// Returns the exact sum of event weights when cheaply available.
    ///
    /// # Errors
    ///
    /// Returns [`LadduDataError`] when source weights or their metadata cannot
    /// be read.
    fn weighted_total(&self) -> LadduDataResult<Option<f64>> {
        Ok(None)
    }

    /// Opens a batch iterator using `plan`.
    ///
    /// # Errors
    ///
    /// Returns [`LadduDataError`] when `plan` is invalid or the source cannot
    /// initialize the requested read.
    fn batches(&self, plan: ReadPlan) -> LadduDataResult<EventBatchIter>;
}

/// Consumer of schema-compatible event batches.
pub trait EventSink: Send {
    /// Returns whether written batches remain resident in memory.
    fn retains_batches(&self) -> bool {
        false
    }

    /// Begins a write operation.
    ///
    /// # Errors
    ///
    /// Returns [`LadduDataError`] when the plan or schema is unsupported or
    /// output initialization fails.
    fn begin(&mut self, schema: Arc<Schema>, plan: WritePlan) -> LadduDataResult<()>;

    /// Writes one batch.
    ///
    /// # Errors
    ///
    /// Returns [`LadduDataError`] when the batch schema is incompatible or the
    /// output cannot be written.
    fn write_batch(&mut self, batch: &EventBatch) -> LadduDataResult<()>;

    /// Finishes and flushes the write operation.
    ///
    /// # Errors
    ///
    /// Returns [`LadduDataError`] when buffered output cannot be finalized.
    fn finish(&mut self) -> LadduDataResult<()>;

    /// Aborts the current write operation and releases any backend resources.
    ///
    /// Aborting is idempotent: calling it while the sink is idle is a no-op.
    /// Backends leave any partially written output in place; callers should
    /// treat such files as incomplete.
    ///
    /// # Errors
    ///
    /// Returns [`LadduDataError`] when backend resources cannot be released.
    fn abort(&mut self) -> LadduDataResult<()> {
        Ok(())
    }
}

/// Internal lifecycle shared by concrete event sinks.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub(crate) enum SinkState {
    #[default]
    Idle,
    Writing,
    Failed,
}

/// Metadata describing one addressable source fragment.
#[derive(Clone, Debug)]
pub struct DataFragment<K> {
    /// Source-specific fragment key.
    pub key: K,
    /// Global row offset.
    pub global_start: u64,
    /// Number of rows.
    pub rows: u64,
}

/// Planned read of one source fragment.
#[derive(Clone, Debug)]
pub struct FragmentRead<K> {
    /// Source-specific fragment key.
    pub key: K,
    /// Rows selected from the fragment.
    pub selection: FragmentSelection,
}

/// Row selection within one source fragment.
#[derive(Clone, Copy, Debug)]
pub enum FragmentSelection {
    /// Contiguous local row range.
    Range {
        /// First local row.
        local_start: usize,
        /// Number of local rows.
        local_len: usize,
    },
    /// Rows assigned by global index modulo rank count.
    StridedRows {
        /// Global offset of the fragment.
        global_start: u64,
        /// Number of rows in the fragment.
        rows: usize,
        /// Current rank.
        rank: usize,
        /// Number of ranks.
        nranks: usize,
    },
}

/// Event source composed of independently addressable fragments.
pub trait FragmentedSource: Send + Sync {
    /// Source-specific fragment key.
    type Key: Clone + Send + Sync + 'static;

    /// Lists all fragments in global row order.
    ///
    /// # Errors
    ///
    /// Returns [`LadduDataError`] when fragment metadata cannot be read.
    fn fragments(&self) -> LadduDataResult<Vec<DataFragment<Self::Key>>>;

    /// Reads a contiguous range within one fragment.
    ///
    /// # Errors
    ///
    /// Returns [`LadduDataError`] when the key, range, or chunk size is invalid
    /// or fragment data cannot be read.
    fn read_fragment_range(
        &self,
        key: &Self::Key,
        local_start: usize,
        local_len: usize,
        chunk_size: Option<usize>,
    ) -> LadduDataResult<EventBatchIter>;
}

/// Creates a planned batch iterator for a fragmented source.
///
/// # Errors
///
/// Returns [`LadduDataError`] when the read plan is invalid, fragment metadata
/// cannot be loaded, or the iterator cannot be initialized.
pub fn fragmented_batches<S>(source: Arc<S>, plan: ReadPlan) -> LadduDataResult<EventBatchIter>
where
    S: FragmentedSource + 'static,
{
    let iter = FragmentBatchIter::new(source, plan)?;

    if plan.chunk_size.is_none() {
        Ok(Box::new(CoalescedBatchIter::new(iter)))
    } else {
        Ok(Box::new(iter))
    }
}

/// Assigns source fragments or rows according to a read plan.
///
/// # Errors
///
/// Returns [`LadduDataError`] when rank settings are invalid or fragment sizes
/// cannot be represented on this platform.
pub fn plan_fragments<K: Clone>(
    fragments: &[DataFragment<K>],
    plan: ReadPlan,
) -> LadduDataResult<Vec<FragmentRead<K>>> {
    let total_rows: u64 = fragments.iter().map(|f| f.rows).sum();
    let rank = plan.rank();
    let nranks = plan.nranks();

    if nranks == 1 {
        return fragments
            .iter()
            .map(|f| {
                Ok(FragmentRead {
                    key: f.key.clone(),
                    selection: FragmentSelection::Range {
                        local_start: 0,
                        local_len: usize_from_u64(f.rows)?,
                    },
                })
            })
            .collect();
    }

    match plan.fragment_partitioning() {
        FragmentPartitioning::Contiguous => contiguous_plan(fragments, total_rows, rank, nranks),
        FragmentPartitioning::RoundRobinFragments => {
            round_robin_fragment_plan(fragments, rank, nranks)
        }
        FragmentPartitioning::StridedRows => strided_row_plan(fragments, rank, nranks),
    }
}

fn contiguous_plan<K: Clone>(
    fragments: &[DataFragment<K>],
    total_rows: u64,
    rank: usize,
    nranks: usize,
) -> LadduDataResult<Vec<FragmentRead<K>>> {
    let rank_start = total_rows * rank as u64 / nranks as u64;
    let rank_end = total_rows * (rank as u64 + 1) / nranks as u64;

    let mut out = Vec::new();

    for f in fragments {
        let frag_start = f.global_start;
        let frag_end = f.global_start + f.rows;

        let start = rank_start.max(frag_start);
        let end = rank_end.min(frag_end);

        if start < end {
            out.push(FragmentRead {
                key: f.key.clone(),
                selection: FragmentSelection::Range {
                    local_start: usize_from_u64(start - frag_start)?,
                    local_len: usize_from_u64(end - start)?,
                },
            });
        }
    }

    Ok(out)
}

fn round_robin_fragment_plan<K: Clone>(
    fragments: &[DataFragment<K>],
    rank: usize,
    nranks: usize,
) -> LadduDataResult<Vec<FragmentRead<K>>> {
    let mut out = Vec::new();

    for (i, f) in fragments.iter().enumerate() {
        if i % nranks == rank {
            out.push(FragmentRead {
                key: f.key.clone(),
                selection: FragmentSelection::Range {
                    local_start: 0,
                    local_len: usize_from_u64(f.rows)?,
                },
            });
        }
    }

    Ok(out)
}

fn strided_row_plan<K: Clone>(
    fragments: &[DataFragment<K>],
    rank: usize,
    nranks: usize,
) -> LadduDataResult<Vec<FragmentRead<K>>> {
    fragments
        .iter()
        .map(|f| {
            Ok(FragmentRead {
                key: f.key.clone(),
                selection: FragmentSelection::StridedRows {
                    global_start: f.global_start,
                    rows: usize_from_u64(f.rows)?,
                    rank,
                    nranks,
                },
            })
        })
        .collect()
}

fn usize_from_u64(value: u64) -> LadduDataResult<usize> {
    usize::try_from(value).map_err(|_| LadduDataError::InvalidArgument("row count exceeds usize"))
}

pub(crate) struct FragmentBatchIter<S>
where
    S: FragmentedSource,
{
    source: Arc<S>,
    reads: Vec<FragmentRead<S::Key>>,
    state: FragmentBatchState,
    chunk_size: Option<usize>,
}

enum FragmentBatchState {
    NeedFragment {
        index: usize,
    },
    Reading {
        iter: EventBatchIter,
        next_index: usize,
    },
    Done,
}

impl<S> FragmentBatchIter<S>
where
    S: FragmentedSource,
{
    pub(crate) fn new(source: Arc<S>, plan: ReadPlan) -> LadduDataResult<Self> {
        let fragments = source.fragments()?;
        let reads = plan_fragments(&fragments, plan)?;

        Ok(Self {
            source,
            reads,
            state: FragmentBatchState::NeedFragment { index: 0 },
            chunk_size: plan.chunk_size,
        })
    }

    fn open_fragment(&self, read: FragmentRead<S::Key>) -> LadduDataResult<EventBatchIter> {
        match read.selection {
            FragmentSelection::Range {
                local_start,
                local_len,
            } => {
                self.source
                    .read_fragment_range(&read.key, local_start, local_len, self.chunk_size)
            }

            FragmentSelection::StridedRows {
                global_start,
                rows,
                rank,
                nranks,
            } => {
                let inner = self
                    .source
                    .read_fragment_range(&read.key, 0, rows, self.chunk_size);

                inner.and_then(|iter| {
                    let iter = StridedRowsBatchIter::new(iter, global_start, rank, nranks)?;
                    Ok(Box::new(iter) as EventBatchIter)
                })
            }
        }
    }
}

impl<S> Iterator for FragmentBatchIter<S>
where
    S: FragmentedSource,
{
    type Item = LadduDataResult<EventBatch>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let state = std::mem::replace(&mut self.state, FragmentBatchState::Done);
            match state {
                FragmentBatchState::Done => return None,
                FragmentBatchState::Reading {
                    mut iter,
                    next_index,
                } => match iter.next() {
                    Some(batch) => {
                        self.state = FragmentBatchState::Reading { iter, next_index };
                        return Some(batch);
                    }
                    None => {
                        self.state = FragmentBatchState::NeedFragment { index: next_index };
                    }
                },
                FragmentBatchState::NeedFragment { index } => {
                    let Some(read) = self.reads.get(index).cloned() else {
                        self.state = FragmentBatchState::Done;
                        return None;
                    };

                    let next_index = index.saturating_add(1);
                    match self.open_fragment(read) {
                        Ok(iter) => {
                            self.state = FragmentBatchState::Reading { iter, next_index };
                        }
                        Err(err) => {
                            // Opening one fragment is an item-level failure. Keep
                            // the next fragment available for a later call.
                            self.state = FragmentBatchState::NeedFragment { index: next_index };
                            return Some(Err(err));
                        }
                    }
                }
            }
        }
    }
}

pub(crate) struct SliceBatchIter<I> {
    inner: I,
    start: usize,
    end: usize,
    state: SliceBatchState,
}

#[derive(Clone, Copy)]
enum SliceBatchState {
    Reading { consumed: usize },
    Done,
}

impl<I> SliceBatchIter<I> {
    pub(crate) fn new(inner: I, start: usize, len: usize) -> LadduDataResult<Self> {
        let end = start
            .checked_add(len)
            .ok_or(LadduDataError::InvalidArgument(
                "slice range overflows usize",
            ))?;
        Ok(Self {
            inner,
            start,
            end,
            state: SliceBatchState::Reading { consumed: 0 },
        })
    }
}

impl<I> Iterator for SliceBatchIter<I>
where
    I: Iterator<Item = LadduDataResult<EventBatch>>,
{
    type Item = LadduDataResult<EventBatch>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let SliceBatchState::Reading { consumed } = self.state else {
                return None;
            };

            if consumed >= self.end {
                self.state = SliceBatchState::Done;
                return None;
            }

            let Some(item) = self.inner.next() else {
                self.state = SliceBatchState::Done;
                return None;
            };
            let batch = match item {
                Ok(batch) => batch,
                Err(err) => return Some(Err(err)),
            };

            let batch_start = consumed;
            let batch_end = batch_start + batch.len();
            self.state = SliceBatchState::Reading {
                consumed: batch_end,
            };

            let lo = self.start.max(batch_start);
            let hi = self.end.min(batch_end);

            if lo >= hi {
                continue;
            }

            let local_lo = lo - batch_start;
            let local_hi = hi - batch_start;

            return Some(Ok(batch.slice(local_lo, local_hi)));
        }
    }
}

pub(crate) struct StridedRowsBatchIter<I> {
    inner: I,
    global_start: u64,
    rank: usize,
    nranks: usize,
    state: StridedRowsBatchState,
}

#[derive(Clone, Copy)]
enum StridedRowsBatchState {
    Reading { consumed: u64 },
    Done,
}

impl<I> StridedRowsBatchIter<I> {
    pub(crate) fn new(
        inner: I,
        global_start: u64,
        rank: usize,
        nranks: usize,
    ) -> LadduDataResult<Self> {
        if nranks == 0 {
            return Err(LadduDataError::InvalidArgument("nranks must be nonzero"));
        }
        if rank >= nranks {
            return Err(LadduDataError::InvalidArgument(
                "rank must be less than nranks",
            ));
        }
        Ok(Self {
            inner,
            global_start,
            rank,
            nranks,
            state: StridedRowsBatchState::Reading { consumed: 0 },
        })
    }
}

impl<I> Iterator for StridedRowsBatchIter<I>
where
    I: Iterator<Item = LadduDataResult<EventBatch>>,
{
    type Item = LadduDataResult<EventBatch>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let StridedRowsBatchState::Reading { consumed } = self.state else {
                return None;
            };

            let Some(item) = self.inner.next() else {
                self.state = StridedRowsBatchState::Done;
                return None;
            };
            let batch = match item {
                Ok(batch) => batch,
                Err(err) => return Some(Err(err)),
            };

            let batch_global_start = self.global_start + consumed;
            self.state = StridedRowsBatchState::Reading {
                consumed: consumed.saturating_add(batch.len() as u64),
            };

            let rows: Vec<usize> = (0..batch.len())
                .filter(|&i| {
                    ((batch_global_start + i as u64) % self.nranks as u64) == self.rank as u64
                })
                .collect();

            if rows.is_empty() {
                continue;
            }

            return Some(Ok(batch.select(&rows)));
        }
    }
}

pub(crate) struct CoalescedBatchIter<I> {
    state: CoalescedBatchState<I>,
}

enum CoalescedBatchState<I> {
    Reading(I),
    Done,
}

impl<I> CoalescedBatchIter<I> {
    pub(crate) fn new(inner: I) -> Self {
        Self {
            state: CoalescedBatchState::Reading(inner),
        }
    }
}

impl<I> Iterator for CoalescedBatchIter<I>
where
    I: Iterator<Item = LadduDataResult<EventBatch>>,
{
    type Item = LadduDataResult<EventBatch>;

    fn next(&mut self) -> Option<Self::Item> {
        let CoalescedBatchState::Reading(mut inner) =
            std::mem::replace(&mut self.state, CoalescedBatchState::Done)
        else {
            return None;
        };
        let mut batches = Vec::new();

        for batch in &mut inner {
            match batch {
                Ok(batch) => batches.push(batch),
                Err(err) => return Some(Err(err)),
            }
        }

        if batches.is_empty() {
            None
        } else {
            Some(EventBatch::concat(&batches))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        data::{EventBatch, EventBatchBuilder},
        schema::Schema,
    };
    use std::path::PathBuf;

    fn v(x: f64) -> RealVec4 {
        RealVec4 {
            e: x,
            px: x,
            py: x,
            pz: x,
        }
    }

    fn schema() -> Arc<Schema> {
        Arc::new(Schema::new(["p"], ["id"], true).unwrap())
    }

    fn batch(start: usize, len: usize) -> EventBatch {
        let schema = schema();
        let mut builder = EventBatchBuilder::with_capacity(schema, len);

        for i in start..start + len {
            builder
                .push_weighted([v(i as f64)], [i as f64], 100.0 + i as f64)
                .unwrap();
        }

        builder.finish().unwrap()
    }

    fn concat_values(batches: Vec<EventBatch>) -> Vec<f64> {
        EventBatch::concat(&batches)
            .unwrap()
            .scalar_column(0)
            .to_vec()
    }

    struct FragmentOpenFailureSource;

    impl FragmentedSource for FragmentOpenFailureSource {
        type Key = usize;

        fn fragments(&self) -> LadduDataResult<Vec<DataFragment<Self::Key>>> {
            Ok(vec![
                DataFragment {
                    key: 0,
                    global_start: 0,
                    rows: 1,
                },
                DataFragment {
                    key: 1,
                    global_start: 1,
                    rows: 1,
                },
            ])
        }

        fn read_fragment_range(
            &self,
            key: &Self::Key,
            _local_start: usize,
            _local_len: usize,
            _chunk_size: Option<usize>,
        ) -> LadduDataResult<EventBatchIter> {
            if *key == 0 {
                return Err(LadduDataError::Source(
                    "first fragment failed to open".into(),
                ));
            }

            Ok(Box::new(vec![Ok(batch(10, 1))].into_iter()))
        }
    }

    #[test]
    fn fragment_open_error_is_an_item_and_later_fragments_remain_readable() {
        let mut iter = FragmentBatchIter::new(
            Arc::new(FragmentOpenFailureSource),
            ReadPlan {
                chunk_size: Some(1),
                #[cfg(feature = "mpi")]
                distribution: Default::default(),
            },
        )
        .unwrap();

        assert!(matches!(
            iter.next(),
            Some(Err(LadduDataError::Source(message))) if message == "first fragment failed to open"
        ));
        assert_eq!(iter.next().unwrap().unwrap().scalar_column(0), &[10.0]);
        assert!(iter.next().is_none());
        assert!(iter.next().is_none());
    }

    #[test]
    fn slice_batch_iter_slices_across_batch_boundaries_without_losing_alignment() {
        let inner = vec![Ok(batch(0, 3)), Ok(batch(3, 2)), Ok(batch(5, 4))].into_iter();

        let out: Vec<EventBatch> = SliceBatchIter::new(inner, 2, 5)
            .unwrap()
            .map(Result::unwrap)
            .collect();

        let values = concat_values(out);
        assert_eq!(values, vec![2.0, 3.0, 4.0, 5.0, 6.0]);
    }

    #[test]
    fn slice_batch_iter_skips_empty_batches_and_has_a_repeated_terminal_none() {
        let inner = vec![Ok(batch(0, 0)), Ok(batch(0, 2))].into_iter();
        let mut iter = SliceBatchIter::new(inner, 0, 2).unwrap();

        assert_eq!(iter.next().unwrap().unwrap().len(), 2);
        assert!(iter.next().is_none());
        assert!(iter.next().is_none());
    }

    #[test]
    fn strided_rows_batch_iter_uses_global_row_numbers_across_batches() {
        let inner = vec![Ok(batch(0, 4)), Ok(batch(4, 5))].into_iter();

        let out: Vec<EventBatch> = StridedRowsBatchIter::new(inner, 1, 1, 3)
            .unwrap()
            .map(Result::unwrap)
            .collect();

        // Global rows are 1..=9 because global_start = 1.
        // Rank 1 of 3 keeps global rows 1, 4, 7.
        // Those correspond to local scalar ids 0, 3, 6.
        assert_eq!(concat_values(out), vec![0.0, 3.0, 6.0]);
    }

    #[test]
    fn coalesced_batch_iter_concatenates_successes_and_propagates_first_error() {
        let success_inner = vec![Ok(batch(0, 2)), Ok(batch(2, 3))].into_iter();
        let mut success = CoalescedBatchIter::new(success_inner);

        let merged = success.next().unwrap().unwrap();
        assert_eq!(merged.scalar_column(0), &[0.0, 1.0, 2.0, 3.0, 4.0]);
        assert!(success.next().is_none());

        let error_inner = vec![
            Ok(batch(0, 1)),
            Err(LadduDataError::Source("boom".into())),
            Ok(batch(1, 1)),
        ]
        .into_iter();

        let mut error_iter = CoalescedBatchIter::new(error_inner);
        let err = error_iter.next().unwrap().unwrap_err();

        assert!(matches!(err, LadduDataError::Source(msg) if msg == "boom"));
        assert!(error_iter.next().is_none());
        assert!(error_iter.next().is_none());

        let mut empty =
            CoalescedBatchIter::new(Vec::<LadduDataResult<EventBatch>>::new().into_iter());
        assert!(empty.next().is_none());
        assert!(empty.next().is_none());
    }

    #[test]
    fn output_path_resolves_single_file_and_per_rank_names() {
        let plan = WritePlan::default();

        let single = OutputPath::new(PathBuf::from("events.parquet"))
            .resolve(plan, "parquet")
            .unwrap();

        assert_eq!(single, PathBuf::from("events.parquet"));

        let per_rank_with_extension = OutputPath::new(PathBuf::from("events.parquet"))
            .with_mode(OutputMode::PerRankFiles)
            .resolve(plan, "parquet")
            .unwrap();

        assert_eq!(
            per_rank_with_extension,
            PathBuf::from("events.rank00000-of00001.parquet")
        );

        let per_rank_without_extension = OutputPath::new(PathBuf::from("events"))
            .with_mode(OutputMode::PerRankFiles)
            .resolve(plan, "root")
            .unwrap();

        assert_eq!(
            per_rank_without_extension,
            PathBuf::from("events").join("part-rank00000-of00001.root")
        );
    }

    #[test]
    fn plan_fragments_serial_mode_keeps_all_fragments_in_order() {
        let fragments = vec![
            DataFragment {
                key: "a",
                global_start: 0,
                rows: 2,
            },
            DataFragment {
                key: "b",
                global_start: 2,
                rows: 3,
            },
        ];

        let reads = plan_fragments(&fragments, ReadPlan::default()).unwrap();

        assert_eq!(reads.len(), 2);

        match &reads[0].selection {
            FragmentSelection::Range {
                local_start,
                local_len,
            } => {
                assert_eq!((*local_start, *local_len), (0, 2));
            }
            _ => panic!("expected range read"),
        }

        match &reads[1].selection {
            FragmentSelection::Range {
                local_start,
                local_len,
            } => {
                assert_eq!((*local_start, *local_len), (0, 3));
            }
            _ => panic!("expected range read"),
        }
    }

    use laddu_physics::vectors::RealVec4;
    #[cfg(feature = "mpi")]
    use mpi::traits::*;
    #[cfg(feature = "mpi")]
    use mpi_test::mpi_test;

    #[cfg(feature = "mpi")]
    fn distributed_plan(
        partitioning: Partitioning,
        world: &impl mpi::topology::Communicator,
    ) -> ReadPlan {
        ReadPlan {
            chunk_size: None,
            distribution: Distribution::from_world(world).with_partitioning(partitioning),
        }
    }

    #[cfg(feature = "mpi")]
    fn expected_contiguous_global_range(total_rows: u64, rank: usize, nranks: usize) -> (u64, u64) {
        let start = total_rows * rank as u64 / nranks as u64;
        let end = total_rows * (rank as u64 + 1) / nranks as u64;
        (start, end)
    }

    #[cfg(feature = "mpi")]
    #[mpi_test(np = [2, 3, 4])]
    fn mpi_contiguous_plan_assigns_disjoint_ranges_covering_all_rows() {
        let universe = mpi::initialize().unwrap();
        let world = universe.world();

        let rank = world.rank() as usize;
        let nranks = world.size() as usize;

        let fragments = vec![
            DataFragment {
                key: "a",
                global_start: 0,
                rows: 4,
            },
            DataFragment {
                key: "b",
                global_start: 4,
                rows: 5,
            },
            DataFragment {
                key: "c",
                global_start: 9,
                rows: 3,
            },
        ];

        let total_rows = fragments.iter().map(|f| f.rows).sum::<u64>();
        let plan = distributed_plan(Partitioning::Contiguous, &world);
        let reads = plan_fragments(&fragments, plan).unwrap();

        let assigned_rows: u64 = reads
            .iter()
            .map(|read| match read.selection {
                FragmentSelection::Range { local_len, .. } => local_len as u64,
                FragmentSelection::StridedRows { .. } => panic!("expected range selection"),
            })
            .sum();

        let (expected_start, expected_end) =
            expected_contiguous_global_range(total_rows, rank, nranks);

        assert_eq!(assigned_rows, expected_end - expected_start);

        for read in reads {
            let fragment = fragments
                .iter()
                .find(|fragment| fragment.key == read.key)
                .unwrap();

            match read.selection {
                FragmentSelection::Range {
                    local_start,
                    local_len,
                } => {
                    let global_start = fragment.global_start + local_start as u64;
                    let global_end = global_start + local_len as u64;

                    assert!(expected_start <= global_start);
                    assert!(global_end <= expected_end);
                    assert!(fragment.global_start <= global_start);
                    assert!(global_end <= fragment.global_start + fragment.rows);
                }
                FragmentSelection::StridedRows { .. } => panic!("expected range selection"),
            }
        }
    }

    #[cfg(feature = "mpi")]
    #[mpi_test(np = [2, 3])]
    fn mpi_file_group_plan_assigns_fragment_by_rank_round_robin() {
        let universe = mpi::initialize().unwrap();
        let world = universe.world();

        let rank = world.rank() as usize;
        let nranks = world.size() as usize;

        let fragments = (0..8)
            .map(|i| DataFragment {
                key: i,
                global_start: 10 * i as u64,
                rows: 10,
            })
            .collect::<Vec<_>>();

        let plan = distributed_plan(Partitioning::FileGroups, &world);
        let reads = plan_fragments(&fragments, plan).unwrap();

        let keys = reads.iter().map(|read| read.key).collect::<Vec<_>>();
        let expected = (0..8).filter(|i| i % nranks == rank).collect::<Vec<_>>();

        assert_eq!(keys, expected);

        for read in reads {
            match read.selection {
                FragmentSelection::Range {
                    local_start,
                    local_len,
                } => {
                    assert_eq!(local_start, 0);
                    assert_eq!(local_len, 10);
                }
                FragmentSelection::StridedRows { .. } => panic!("expected range selection"),
            }
        }
    }

    #[cfg(feature = "mpi")]
    #[mpi_test(np = [2, 3, 4])]
    fn mpi_rows_plan_assigns_strided_row_selection_with_world_rank() {
        let universe = mpi::initialize().unwrap();
        let world = universe.world();

        let rank = world.rank() as usize;
        let nranks = world.size() as usize;

        let fragments = vec![
            DataFragment {
                key: "a",
                global_start: 0,
                rows: 4,
            },
            DataFragment {
                key: "b",
                global_start: 4,
                rows: 5,
            },
        ];

        let plan = distributed_plan(Partitioning::Rows, &world);
        let reads = plan_fragments(&fragments, plan).unwrap();

        assert_eq!(reads.len(), fragments.len());

        for (read, fragment) in reads.iter().zip(fragments.iter()) {
            assert_eq!(read.key, fragment.key);

            match read.selection {
                FragmentSelection::StridedRows {
                    global_start,
                    rows,
                    rank: selected_rank,
                    nranks: selected_nranks,
                } => {
                    assert_eq!(global_start, fragment.global_start);
                    assert_eq!(rows, fragment.rows as usize);
                    assert_eq!(selected_rank, rank);
                    assert_eq!(selected_nranks, nranks);
                }
                FragmentSelection::Range { .. } => panic!("expected strided selection"),
            }
        }
    }

    #[cfg(feature = "mpi")]
    #[mpi_test(np = [2, 3])]
    fn mpi_read_plan_and_write_plan_reflect_world_distribution() {
        let universe = mpi::initialize().unwrap();
        let world = universe.world();

        let read_plan = ReadPlan {
            chunk_size: Some(7),
            distribution: Distribution::from_world(&world).with_partitioning(Partitioning::Rows),
        };

        assert!(read_plan.is_distributed());
        assert_eq!(read_plan.rank(), world.rank() as usize);
        assert_eq!(read_plan.nranks(), world.size() as usize);

        match read_plan.fragment_partitioning() {
            FragmentPartitioning::StridedRows => {}
            _ => panic!("expected strided row partitioning"),
        }

        let write_plan = WritePlan::from(read_plan);

        assert!(write_plan.is_distributed());
        assert_eq!(write_plan.rank(), world.rank() as usize);
        assert_eq!(write_plan.nranks(), world.size() as usize);
    }
}