laddu_data/io/
mod.rs

1use serde::{Deserialize, Serialize};
2use std::fmt::Display;
3use std::sync::Arc;
4
5use crate::{LadduDataError, LadduDataResult, data::EventBatch, schema::Schema};
6
7/// In-memory event sources and sinks.
8pub mod memory;
9mod output;
10/// Parquet event sources and sinks.
11pub mod parquet;
12/// ROOT event sources.
13pub mod root;
14
15mod source;
16
17pub use output::{OutputMode, OutputPath};
18
19pub(crate) use source::{SourceBuild, SourceBuildOptions, build_source};
20
21/// Adds stable operation/resource context to a source failure.
22pub(crate) fn source_error(
23    operation: impl AsRef<str>,
24    resource: impl Display,
25    cause: impl Display,
26) -> LadduDataError {
27    LadduDataError::Source(format!("{} `{resource}`: {cause}", operation.as_ref()))
28}
29
30/// Adds stable operation/resource context to a sink failure.
31pub(crate) fn sink_error(
32    operation: impl AsRef<str>,
33    resource: impl Display,
34    cause: impl Display,
35) -> LadduDataError {
36    LadduDataError::Sink(format!("{} `{resource}`: {cause}", operation.as_ref()))
37}
38
39#[cfg(test)]
40mod context_tests {
41    use super::*;
42
43    #[test]
44    fn source_and_sink_context_have_stable_operation_resource_format() {
45        assert!(matches!(
46            source_error("read ROOT tree", "events.root::events", "branch failed"),
47            LadduDataError::Source(message)
48                if message == "read ROOT tree `events.root::events`: branch failed"
49        ));
50        assert!(matches!(
51            sink_error("write Parquet file", "events.parquet", "disk full"),
52            LadduDataError::Sink(message)
53                if message == "write Parquet file `events.parquet`: disk full"
54        ));
55    }
56}
57
58#[cfg(test)]
59mod contract_tests;
60
61#[cfg(feature = "mpi")]
62/// Distribution of event I/O across MPI ranks.
63#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
64pub enum Distribution {
65    /// Single-process I/O.
66    #[default]
67    Serial,
68    /// MPI-distributed I/O with explicit rank metadata.
69    Mpi {
70        /// Zero-based rank.
71        rank: usize,
72        /// Number of ranks.
73        nranks: usize,
74        /// Work-partitioning strategy.
75        partitioning: Partitioning,
76    },
77}
78
79#[cfg(feature = "mpi")]
80impl Distribution {
81    /// Creates serial distribution.
82    pub fn serial() -> Self {
83        Self::Serial
84    }
85
86    /// Creates MPI distribution from a communicator.
87    pub fn from_world<C>(world: &C) -> Self
88    where
89        C: mpi::topology::Communicator,
90    {
91        Self::Mpi {
92            rank: world.rank() as usize,
93            nranks: world.size() as usize,
94            partitioning: Partitioning::default(),
95        }
96    }
97
98    /// Returns the current rank.
99    pub fn rank(self) -> usize {
100        match self {
101            Self::Serial => 0,
102            Self::Mpi { rank, .. } => rank,
103        }
104    }
105
106    /// Returns the number of ranks.
107    pub fn nranks(self) -> usize {
108        match self {
109            Self::Serial => 1,
110            Self::Mpi { nranks, .. } => nranks,
111        }
112    }
113
114    /// Returns the partitioning strategy.
115    pub fn partitioning(self) -> Partitioning {
116        match self {
117            Self::Serial => Partitioning::Contiguous,
118            Self::Mpi { partitioning, .. } => partitioning,
119        }
120    }
121
122    /// Returns this distribution with a new partitioning strategy.
123    pub fn with_partitioning(self, partitioning: Partitioning) -> Self {
124        match self {
125            Self::Serial => Self::Serial,
126            Self::Mpi { rank, nranks, .. } => Self::Mpi {
127                rank,
128                nranks,
129                partitioning,
130            },
131        }
132    }
133}
134
135/// Strategy for partitioning input rows across ranks.
136#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
137pub enum Partitioning {
138    /// Each rank reads a contiguous global row range.
139    #[default]
140    Contiguous,
141
142    /// Each rank reads whole source fragments, such as files or row groups, round-robin.
143    FileGroups,
144
145    /// Rank r keeps rows where global_row % nranks == r.
146    /// Deterministic, but usually slower.
147    Rows,
148}
149
150/// Options controlling event-source reads.
151#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
152pub struct ReadPlan {
153    /// Optional maximum output batch size.
154    pub chunk_size: Option<usize>,
155
156    #[cfg(feature = "mpi")]
157    /// MPI distribution.
158    pub distribution: Distribution,
159}
160
161impl ReadPlan {
162    /// Creates a serial read plan.
163    pub fn serial() -> Self {
164        Self::default()
165    }
166
167    /// Returns the current rank.
168    pub fn rank(&self) -> usize {
169        #[cfg(feature = "mpi")]
170        {
171            self.distribution.rank()
172        }
173
174        #[cfg(not(feature = "mpi"))]
175        {
176            0
177        }
178    }
179
180    /// Returns the number of ranks.
181    pub fn nranks(&self) -> usize {
182        #[cfg(feature = "mpi")]
183        {
184            self.distribution.nranks()
185        }
186
187        #[cfg(not(feature = "mpi"))]
188        {
189            1
190        }
191    }
192
193    /// Returns whether reads are distributed.
194    pub fn is_distributed(&self) -> bool {
195        self.nranks() > 1
196    }
197
198    /// Returns the low-level fragment-partitioning strategy.
199    pub fn fragment_partitioning(&self) -> FragmentPartitioning {
200        #[cfg(feature = "mpi")]
201        {
202            match self.distribution.partitioning() {
203                Partitioning::Contiguous => FragmentPartitioning::Contiguous,
204                Partitioning::FileGroups => FragmentPartitioning::RoundRobinFragments,
205                Partitioning::Rows => FragmentPartitioning::StridedRows,
206            }
207        }
208
209        #[cfg(not(feature = "mpi"))]
210        {
211            FragmentPartitioning::Contiguous
212        }
213    }
214}
215
216/// Options controlling event-sink writes.
217#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
218pub struct WritePlan {
219    #[cfg(feature = "mpi")]
220    /// MPI distribution.
221    pub distribution: Distribution,
222}
223
224impl From<ReadPlan> for WritePlan {
225    #[cfg_attr(not(feature = "mpi"), allow(unused_variables))]
226    fn from(plan: ReadPlan) -> Self {
227        Self {
228            #[cfg(feature = "mpi")]
229            distribution: plan.distribution,
230        }
231    }
232}
233
234impl WritePlan {
235    /// Returns the current rank.
236    pub fn rank(&self) -> usize {
237        #[cfg(feature = "mpi")]
238        {
239            self.distribution.rank()
240        }
241
242        #[cfg(not(feature = "mpi"))]
243        {
244            0
245        }
246    }
247
248    /// Returns the number of ranks.
249    pub fn nranks(&self) -> usize {
250        #[cfg(feature = "mpi")]
251        {
252            self.distribution.nranks()
253        }
254
255        #[cfg(not(feature = "mpi"))]
256        {
257            1
258        }
259    }
260
261    /// Returns whether writes are distributed.
262    pub fn is_distributed(&self) -> bool {
263        self.nranks() > 1
264    }
265}
266
267/// Low-level assignment of source fragments or rows.
268#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
269pub enum FragmentPartitioning {
270    /// Contiguous global row ranges.
271    Contiguous,
272    /// Whole fragments assigned round-robin.
273    RoundRobinFragments,
274    /// Individual rows assigned by global index modulo rank count.
275    StridedRows,
276}
277
278/// Optional performance and planning capabilities of an [`EventSource`].
279#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
280pub struct SourceCapabilities {
281    /// Exact event count is available cheaply.
282    pub exact_len: bool,
283    /// Exact weighted total is available cheaply.
284    pub exact_weighted_total: bool,
285    /// Arbitrary row ranges can be read.
286    pub random_access: bool,
287    /// Distributed row assignment is deterministic.
288    pub deterministic_partitioning: bool,
289    /// Filters can be pushed into the source.
290    pub predicate_pushdown: bool,
291    /// Column projection can be pushed into the source.
292    pub projection_pushdown: bool,
293    /// Batches can be streamed without full materialization.
294    pub streaming: bool,
295}
296
297/// Sendable iterator of fallible event batches.
298pub type EventBatchIter = Box<dyn Iterator<Item = LadduDataResult<EventBatch>> + Send>;
299
300/// Thread-safe producer of schema-compatible event batches.
301pub trait EventSource: Send + Sync {
302    /// Returns the source schema.
303    ///
304    /// # Errors
305    ///
306    /// Returns [`LadduDataError`] when source metadata cannot be read or
307    /// interpreted as a logical schema.
308    fn schema(&self) -> LadduDataResult<Arc<Schema>>;
309
310    /// Returns optional source capabilities.
311    fn capabilities(&self) -> SourceCapabilities {
312        SourceCapabilities::default()
313    }
314
315    /// Returns the exact event count when cheaply available.
316    ///
317    /// # Errors
318    ///
319    /// Returns [`LadduDataError`] when the source cannot read the metadata
320    /// needed to determine its event count.
321    fn num_events(&self) -> LadduDataResult<Option<u64>> {
322        Ok(None)
323    }
324
325    /// Returns the exact sum of event weights when cheaply available.
326    ///
327    /// # Errors
328    ///
329    /// Returns [`LadduDataError`] when source weights or their metadata cannot
330    /// be read.
331    fn weighted_total(&self) -> LadduDataResult<Option<f64>> {
332        Ok(None)
333    }
334
335    /// Opens a batch iterator using `plan`.
336    ///
337    /// # Errors
338    ///
339    /// Returns [`LadduDataError`] when `plan` is invalid or the source cannot
340    /// initialize the requested read.
341    fn batches(&self, plan: ReadPlan) -> LadduDataResult<EventBatchIter>;
342}
343
344/// Consumer of schema-compatible event batches.
345pub trait EventSink: Send {
346    /// Returns whether written batches remain resident in memory.
347    fn retains_batches(&self) -> bool {
348        false
349    }
350
351    /// Begins a write operation.
352    ///
353    /// # Errors
354    ///
355    /// Returns [`LadduDataError`] when the plan or schema is unsupported or
356    /// output initialization fails.
357    fn begin(&mut self, schema: Arc<Schema>, plan: WritePlan) -> LadduDataResult<()>;
358
359    /// Writes one batch.
360    ///
361    /// # Errors
362    ///
363    /// Returns [`LadduDataError`] when the batch schema is incompatible or the
364    /// output cannot be written.
365    fn write_batch(&mut self, batch: &EventBatch) -> LadduDataResult<()>;
366
367    /// Finishes and flushes the write operation.
368    ///
369    /// # Errors
370    ///
371    /// Returns [`LadduDataError`] when buffered output cannot be finalized.
372    fn finish(&mut self) -> LadduDataResult<()>;
373
374    /// Aborts the current write operation and releases any backend resources.
375    ///
376    /// Aborting is idempotent: calling it while the sink is idle is a no-op.
377    /// Backends leave any partially written output in place; callers should
378    /// treat such files as incomplete.
379    ///
380    /// # Errors
381    ///
382    /// Returns [`LadduDataError`] when backend resources cannot be released.
383    fn abort(&mut self) -> LadduDataResult<()> {
384        Ok(())
385    }
386}
387
388/// Internal lifecycle shared by concrete event sinks.
389#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
390pub(crate) enum SinkState {
391    #[default]
392    Idle,
393    Writing,
394    Failed,
395}
396
397/// Metadata describing one addressable source fragment.
398#[derive(Clone, Debug)]
399pub struct DataFragment<K> {
400    /// Source-specific fragment key.
401    pub key: K,
402    /// Global row offset.
403    pub global_start: u64,
404    /// Number of rows.
405    pub rows: u64,
406}
407
408/// Planned read of one source fragment.
409#[derive(Clone, Debug)]
410pub struct FragmentRead<K> {
411    /// Source-specific fragment key.
412    pub key: K,
413    /// Rows selected from the fragment.
414    pub selection: FragmentSelection,
415}
416
417/// Row selection within one source fragment.
418#[derive(Clone, Copy, Debug)]
419pub enum FragmentSelection {
420    /// Contiguous local row range.
421    Range {
422        /// First local row.
423        local_start: usize,
424        /// Number of local rows.
425        local_len: usize,
426    },
427    /// Rows assigned by global index modulo rank count.
428    StridedRows {
429        /// Global offset of the fragment.
430        global_start: u64,
431        /// Number of rows in the fragment.
432        rows: usize,
433        /// Current rank.
434        rank: usize,
435        /// Number of ranks.
436        nranks: usize,
437    },
438}
439
440/// Event source composed of independently addressable fragments.
441pub trait FragmentedSource: Send + Sync {
442    /// Source-specific fragment key.
443    type Key: Clone + Send + Sync + 'static;
444
445    /// Lists all fragments in global row order.
446    ///
447    /// # Errors
448    ///
449    /// Returns [`LadduDataError`] when fragment metadata cannot be read.
450    fn fragments(&self) -> LadduDataResult<Vec<DataFragment<Self::Key>>>;
451
452    /// Reads a contiguous range within one fragment.
453    ///
454    /// # Errors
455    ///
456    /// Returns [`LadduDataError`] when the key, range, or chunk size is invalid
457    /// or fragment data cannot be read.
458    fn read_fragment_range(
459        &self,
460        key: &Self::Key,
461        local_start: usize,
462        local_len: usize,
463        chunk_size: Option<usize>,
464    ) -> LadduDataResult<EventBatchIter>;
465}
466
467/// Creates a planned batch iterator for a fragmented source.
468///
469/// # Errors
470///
471/// Returns [`LadduDataError`] when the read plan is invalid, fragment metadata
472/// cannot be loaded, or the iterator cannot be initialized.
473pub fn fragmented_batches<S>(source: Arc<S>, plan: ReadPlan) -> LadduDataResult<EventBatchIter>
474where
475    S: FragmentedSource + 'static,
476{
477    let iter = FragmentBatchIter::new(source, plan)?;
478
479    if plan.chunk_size.is_none() {
480        Ok(Box::new(CoalescedBatchIter::new(iter)))
481    } else {
482        Ok(Box::new(iter))
483    }
484}
485
486/// Assigns source fragments or rows according to a read plan.
487///
488/// # Errors
489///
490/// Returns [`LadduDataError`] when rank settings are invalid or fragment sizes
491/// cannot be represented on this platform.
492pub fn plan_fragments<K: Clone>(
493    fragments: &[DataFragment<K>],
494    plan: ReadPlan,
495) -> LadduDataResult<Vec<FragmentRead<K>>> {
496    let total_rows: u64 = fragments.iter().map(|f| f.rows).sum();
497    let rank = plan.rank();
498    let nranks = plan.nranks();
499
500    if nranks == 1 {
501        return fragments
502            .iter()
503            .map(|f| {
504                Ok(FragmentRead {
505                    key: f.key.clone(),
506                    selection: FragmentSelection::Range {
507                        local_start: 0,
508                        local_len: usize_from_u64(f.rows)?,
509                    },
510                })
511            })
512            .collect();
513    }
514
515    match plan.fragment_partitioning() {
516        FragmentPartitioning::Contiguous => contiguous_plan(fragments, total_rows, rank, nranks),
517        FragmentPartitioning::RoundRobinFragments => {
518            round_robin_fragment_plan(fragments, rank, nranks)
519        }
520        FragmentPartitioning::StridedRows => strided_row_plan(fragments, rank, nranks),
521    }
522}
523
524fn contiguous_plan<K: Clone>(
525    fragments: &[DataFragment<K>],
526    total_rows: u64,
527    rank: usize,
528    nranks: usize,
529) -> LadduDataResult<Vec<FragmentRead<K>>> {
530    let rank_start = total_rows * rank as u64 / nranks as u64;
531    let rank_end = total_rows * (rank as u64 + 1) / nranks as u64;
532
533    let mut out = Vec::new();
534
535    for f in fragments {
536        let frag_start = f.global_start;
537        let frag_end = f.global_start + f.rows;
538
539        let start = rank_start.max(frag_start);
540        let end = rank_end.min(frag_end);
541
542        if start < end {
543            out.push(FragmentRead {
544                key: f.key.clone(),
545                selection: FragmentSelection::Range {
546                    local_start: usize_from_u64(start - frag_start)?,
547                    local_len: usize_from_u64(end - start)?,
548                },
549            });
550        }
551    }
552
553    Ok(out)
554}
555
556fn round_robin_fragment_plan<K: Clone>(
557    fragments: &[DataFragment<K>],
558    rank: usize,
559    nranks: usize,
560) -> LadduDataResult<Vec<FragmentRead<K>>> {
561    let mut out = Vec::new();
562
563    for (i, f) in fragments.iter().enumerate() {
564        if i % nranks == rank {
565            out.push(FragmentRead {
566                key: f.key.clone(),
567                selection: FragmentSelection::Range {
568                    local_start: 0,
569                    local_len: usize_from_u64(f.rows)?,
570                },
571            });
572        }
573    }
574
575    Ok(out)
576}
577
578fn strided_row_plan<K: Clone>(
579    fragments: &[DataFragment<K>],
580    rank: usize,
581    nranks: usize,
582) -> LadduDataResult<Vec<FragmentRead<K>>> {
583    fragments
584        .iter()
585        .map(|f| {
586            Ok(FragmentRead {
587                key: f.key.clone(),
588                selection: FragmentSelection::StridedRows {
589                    global_start: f.global_start,
590                    rows: usize_from_u64(f.rows)?,
591                    rank,
592                    nranks,
593                },
594            })
595        })
596        .collect()
597}
598
599fn usize_from_u64(value: u64) -> LadduDataResult<usize> {
600    usize::try_from(value).map_err(|_| LadduDataError::InvalidArgument("row count exceeds usize"))
601}
602
603pub(crate) struct FragmentBatchIter<S>
604where
605    S: FragmentedSource,
606{
607    source: Arc<S>,
608    reads: Vec<FragmentRead<S::Key>>,
609    state: FragmentBatchState,
610    chunk_size: Option<usize>,
611}
612
613enum FragmentBatchState {
614    NeedFragment {
615        index: usize,
616    },
617    Reading {
618        iter: EventBatchIter,
619        next_index: usize,
620    },
621    Done,
622}
623
624impl<S> FragmentBatchIter<S>
625where
626    S: FragmentedSource,
627{
628    pub(crate) fn new(source: Arc<S>, plan: ReadPlan) -> LadduDataResult<Self> {
629        let fragments = source.fragments()?;
630        let reads = plan_fragments(&fragments, plan)?;
631
632        Ok(Self {
633            source,
634            reads,
635            state: FragmentBatchState::NeedFragment { index: 0 },
636            chunk_size: plan.chunk_size,
637        })
638    }
639
640    fn open_fragment(&self, read: FragmentRead<S::Key>) -> LadduDataResult<EventBatchIter> {
641        match read.selection {
642            FragmentSelection::Range {
643                local_start,
644                local_len,
645            } => {
646                self.source
647                    .read_fragment_range(&read.key, local_start, local_len, self.chunk_size)
648            }
649
650            FragmentSelection::StridedRows {
651                global_start,
652                rows,
653                rank,
654                nranks,
655            } => {
656                let inner = self
657                    .source
658                    .read_fragment_range(&read.key, 0, rows, self.chunk_size);
659
660                inner.and_then(|iter| {
661                    let iter = StridedRowsBatchIter::new(iter, global_start, rank, nranks)?;
662                    Ok(Box::new(iter) as EventBatchIter)
663                })
664            }
665        }
666    }
667}
668
669impl<S> Iterator for FragmentBatchIter<S>
670where
671    S: FragmentedSource,
672{
673    type Item = LadduDataResult<EventBatch>;
674
675    fn next(&mut self) -> Option<Self::Item> {
676        loop {
677            let state = std::mem::replace(&mut self.state, FragmentBatchState::Done);
678            match state {
679                FragmentBatchState::Done => return None,
680                FragmentBatchState::Reading {
681                    mut iter,
682                    next_index,
683                } => match iter.next() {
684                    Some(batch) => {
685                        self.state = FragmentBatchState::Reading { iter, next_index };
686                        return Some(batch);
687                    }
688                    None => {
689                        self.state = FragmentBatchState::NeedFragment { index: next_index };
690                    }
691                },
692                FragmentBatchState::NeedFragment { index } => {
693                    let Some(read) = self.reads.get(index).cloned() else {
694                        self.state = FragmentBatchState::Done;
695                        return None;
696                    };
697
698                    let next_index = index.saturating_add(1);
699                    match self.open_fragment(read) {
700                        Ok(iter) => {
701                            self.state = FragmentBatchState::Reading { iter, next_index };
702                        }
703                        Err(err) => {
704                            // Opening one fragment is an item-level failure. Keep
705                            // the next fragment available for a later call.
706                            self.state = FragmentBatchState::NeedFragment { index: next_index };
707                            return Some(Err(err));
708                        }
709                    }
710                }
711            }
712        }
713    }
714}
715
716pub(crate) struct SliceBatchIter<I> {
717    inner: I,
718    start: usize,
719    end: usize,
720    state: SliceBatchState,
721}
722
723#[derive(Clone, Copy)]
724enum SliceBatchState {
725    Reading { consumed: usize },
726    Done,
727}
728
729impl<I> SliceBatchIter<I> {
730    pub(crate) fn new(inner: I, start: usize, len: usize) -> LadduDataResult<Self> {
731        let end = start
732            .checked_add(len)
733            .ok_or(LadduDataError::InvalidArgument(
734                "slice range overflows usize",
735            ))?;
736        Ok(Self {
737            inner,
738            start,
739            end,
740            state: SliceBatchState::Reading { consumed: 0 },
741        })
742    }
743}
744
745impl<I> Iterator for SliceBatchIter<I>
746where
747    I: Iterator<Item = LadduDataResult<EventBatch>>,
748{
749    type Item = LadduDataResult<EventBatch>;
750
751    fn next(&mut self) -> Option<Self::Item> {
752        loop {
753            let SliceBatchState::Reading { consumed } = self.state else {
754                return None;
755            };
756
757            if consumed >= self.end {
758                self.state = SliceBatchState::Done;
759                return None;
760            }
761
762            let Some(item) = self.inner.next() else {
763                self.state = SliceBatchState::Done;
764                return None;
765            };
766            let batch = match item {
767                Ok(batch) => batch,
768                Err(err) => return Some(Err(err)),
769            };
770
771            let batch_start = consumed;
772            let batch_end = batch_start + batch.len();
773            self.state = SliceBatchState::Reading {
774                consumed: batch_end,
775            };
776
777            let lo = self.start.max(batch_start);
778            let hi = self.end.min(batch_end);
779
780            if lo >= hi {
781                continue;
782            }
783
784            let local_lo = lo - batch_start;
785            let local_hi = hi - batch_start;
786
787            return Some(Ok(batch.slice(local_lo, local_hi)));
788        }
789    }
790}
791
792pub(crate) struct StridedRowsBatchIter<I> {
793    inner: I,
794    global_start: u64,
795    rank: usize,
796    nranks: usize,
797    state: StridedRowsBatchState,
798}
799
800#[derive(Clone, Copy)]
801enum StridedRowsBatchState {
802    Reading { consumed: u64 },
803    Done,
804}
805
806impl<I> StridedRowsBatchIter<I> {
807    pub(crate) fn new(
808        inner: I,
809        global_start: u64,
810        rank: usize,
811        nranks: usize,
812    ) -> LadduDataResult<Self> {
813        if nranks == 0 {
814            return Err(LadduDataError::InvalidArgument("nranks must be nonzero"));
815        }
816        if rank >= nranks {
817            return Err(LadduDataError::InvalidArgument(
818                "rank must be less than nranks",
819            ));
820        }
821        Ok(Self {
822            inner,
823            global_start,
824            rank,
825            nranks,
826            state: StridedRowsBatchState::Reading { consumed: 0 },
827        })
828    }
829}
830
831impl<I> Iterator for StridedRowsBatchIter<I>
832where
833    I: Iterator<Item = LadduDataResult<EventBatch>>,
834{
835    type Item = LadduDataResult<EventBatch>;
836
837    fn next(&mut self) -> Option<Self::Item> {
838        loop {
839            let StridedRowsBatchState::Reading { consumed } = self.state else {
840                return None;
841            };
842
843            let Some(item) = self.inner.next() else {
844                self.state = StridedRowsBatchState::Done;
845                return None;
846            };
847            let batch = match item {
848                Ok(batch) => batch,
849                Err(err) => return Some(Err(err)),
850            };
851
852            let batch_global_start = self.global_start + consumed;
853            self.state = StridedRowsBatchState::Reading {
854                consumed: consumed.saturating_add(batch.len() as u64),
855            };
856
857            let rows: Vec<usize> = (0..batch.len())
858                .filter(|&i| {
859                    ((batch_global_start + i as u64) % self.nranks as u64) == self.rank as u64
860                })
861                .collect();
862
863            if rows.is_empty() {
864                continue;
865            }
866
867            return Some(Ok(batch.select(&rows)));
868        }
869    }
870}
871
872pub(crate) struct CoalescedBatchIter<I> {
873    state: CoalescedBatchState<I>,
874}
875
876enum CoalescedBatchState<I> {
877    Reading(I),
878    Done,
879}
880
881impl<I> CoalescedBatchIter<I> {
882    pub(crate) fn new(inner: I) -> Self {
883        Self {
884            state: CoalescedBatchState::Reading(inner),
885        }
886    }
887}
888
889impl<I> Iterator for CoalescedBatchIter<I>
890where
891    I: Iterator<Item = LadduDataResult<EventBatch>>,
892{
893    type Item = LadduDataResult<EventBatch>;
894
895    fn next(&mut self) -> Option<Self::Item> {
896        let CoalescedBatchState::Reading(mut inner) =
897            std::mem::replace(&mut self.state, CoalescedBatchState::Done)
898        else {
899            return None;
900        };
901        let mut batches = Vec::new();
902
903        for batch in &mut inner {
904            match batch {
905                Ok(batch) => batches.push(batch),
906                Err(err) => return Some(Err(err)),
907            }
908        }
909
910        if batches.is_empty() {
911            None
912        } else {
913            Some(EventBatch::concat(&batches))
914        }
915    }
916}
917
918#[cfg(test)]
919mod tests {
920    use super::*;
921    use crate::{
922        data::{EventBatch, EventBatchBuilder},
923        schema::Schema,
924    };
925    use std::path::PathBuf;
926
927    fn v(x: f64) -> RealVec4 {
928        RealVec4 {
929            e: x,
930            px: x,
931            py: x,
932            pz: x,
933        }
934    }
935
936    fn schema() -> Arc<Schema> {
937        Arc::new(Schema::new(["p"], ["id"], true).unwrap())
938    }
939
940    fn batch(start: usize, len: usize) -> EventBatch {
941        let schema = schema();
942        let mut builder = EventBatchBuilder::with_capacity(schema, len);
943
944        for i in start..start + len {
945            builder
946                .push_weighted([v(i as f64)], [i as f64], 100.0 + i as f64)
947                .unwrap();
948        }
949
950        builder.finish().unwrap()
951    }
952
953    fn concat_values(batches: Vec<EventBatch>) -> Vec<f64> {
954        EventBatch::concat(&batches)
955            .unwrap()
956            .scalar_column(0)
957            .to_vec()
958    }
959
960    struct FragmentOpenFailureSource;
961
962    impl FragmentedSource for FragmentOpenFailureSource {
963        type Key = usize;
964
965        fn fragments(&self) -> LadduDataResult<Vec<DataFragment<Self::Key>>> {
966            Ok(vec![
967                DataFragment {
968                    key: 0,
969                    global_start: 0,
970                    rows: 1,
971                },
972                DataFragment {
973                    key: 1,
974                    global_start: 1,
975                    rows: 1,
976                },
977            ])
978        }
979
980        fn read_fragment_range(
981            &self,
982            key: &Self::Key,
983            _local_start: usize,
984            _local_len: usize,
985            _chunk_size: Option<usize>,
986        ) -> LadduDataResult<EventBatchIter> {
987            if *key == 0 {
988                return Err(LadduDataError::Source(
989                    "first fragment failed to open".into(),
990                ));
991            }
992
993            Ok(Box::new(vec![Ok(batch(10, 1))].into_iter()))
994        }
995    }
996
997    #[test]
998    fn fragment_open_error_is_an_item_and_later_fragments_remain_readable() {
999        let mut iter = FragmentBatchIter::new(
1000            Arc::new(FragmentOpenFailureSource),
1001            ReadPlan {
1002                chunk_size: Some(1),
1003                #[cfg(feature = "mpi")]
1004                distribution: Default::default(),
1005            },
1006        )
1007        .unwrap();
1008
1009        assert!(matches!(
1010            iter.next(),
1011            Some(Err(LadduDataError::Source(message))) if message == "first fragment failed to open"
1012        ));
1013        assert_eq!(iter.next().unwrap().unwrap().scalar_column(0), &[10.0]);
1014        assert!(iter.next().is_none());
1015        assert!(iter.next().is_none());
1016    }
1017
1018    #[test]
1019    fn slice_batch_iter_slices_across_batch_boundaries_without_losing_alignment() {
1020        let inner = vec![Ok(batch(0, 3)), Ok(batch(3, 2)), Ok(batch(5, 4))].into_iter();
1021
1022        let out: Vec<EventBatch> = SliceBatchIter::new(inner, 2, 5)
1023            .unwrap()
1024            .map(Result::unwrap)
1025            .collect();
1026
1027        let values = concat_values(out);
1028        assert_eq!(values, vec![2.0, 3.0, 4.0, 5.0, 6.0]);
1029    }
1030
1031    #[test]
1032    fn slice_batch_iter_skips_empty_batches_and_has_a_repeated_terminal_none() {
1033        let inner = vec![Ok(batch(0, 0)), Ok(batch(0, 2))].into_iter();
1034        let mut iter = SliceBatchIter::new(inner, 0, 2).unwrap();
1035
1036        assert_eq!(iter.next().unwrap().unwrap().len(), 2);
1037        assert!(iter.next().is_none());
1038        assert!(iter.next().is_none());
1039    }
1040
1041    #[test]
1042    fn strided_rows_batch_iter_uses_global_row_numbers_across_batches() {
1043        let inner = vec![Ok(batch(0, 4)), Ok(batch(4, 5))].into_iter();
1044
1045        let out: Vec<EventBatch> = StridedRowsBatchIter::new(inner, 1, 1, 3)
1046            .unwrap()
1047            .map(Result::unwrap)
1048            .collect();
1049
1050        // Global rows are 1..=9 because global_start = 1.
1051        // Rank 1 of 3 keeps global rows 1, 4, 7.
1052        // Those correspond to local scalar ids 0, 3, 6.
1053        assert_eq!(concat_values(out), vec![0.0, 3.0, 6.0]);
1054    }
1055
1056    #[test]
1057    fn coalesced_batch_iter_concatenates_successes_and_propagates_first_error() {
1058        let success_inner = vec![Ok(batch(0, 2)), Ok(batch(2, 3))].into_iter();
1059        let mut success = CoalescedBatchIter::new(success_inner);
1060
1061        let merged = success.next().unwrap().unwrap();
1062        assert_eq!(merged.scalar_column(0), &[0.0, 1.0, 2.0, 3.0, 4.0]);
1063        assert!(success.next().is_none());
1064
1065        let error_inner = vec![
1066            Ok(batch(0, 1)),
1067            Err(LadduDataError::Source("boom".into())),
1068            Ok(batch(1, 1)),
1069        ]
1070        .into_iter();
1071
1072        let mut error_iter = CoalescedBatchIter::new(error_inner);
1073        let err = error_iter.next().unwrap().unwrap_err();
1074
1075        assert!(matches!(err, LadduDataError::Source(msg) if msg == "boom"));
1076        assert!(error_iter.next().is_none());
1077        assert!(error_iter.next().is_none());
1078
1079        let mut empty =
1080            CoalescedBatchIter::new(Vec::<LadduDataResult<EventBatch>>::new().into_iter());
1081        assert!(empty.next().is_none());
1082        assert!(empty.next().is_none());
1083    }
1084
1085    #[test]
1086    fn output_path_resolves_single_file_and_per_rank_names() {
1087        let plan = WritePlan::default();
1088
1089        let single = OutputPath::new(PathBuf::from("events.parquet"))
1090            .resolve(plan, "parquet")
1091            .unwrap();
1092
1093        assert_eq!(single, PathBuf::from("events.parquet"));
1094
1095        let per_rank_with_extension = OutputPath::new(PathBuf::from("events.parquet"))
1096            .with_mode(OutputMode::PerRankFiles)
1097            .resolve(plan, "parquet")
1098            .unwrap();
1099
1100        assert_eq!(
1101            per_rank_with_extension,
1102            PathBuf::from("events.rank00000-of00001.parquet")
1103        );
1104
1105        let per_rank_without_extension = OutputPath::new(PathBuf::from("events"))
1106            .with_mode(OutputMode::PerRankFiles)
1107            .resolve(plan, "root")
1108            .unwrap();
1109
1110        assert_eq!(
1111            per_rank_without_extension,
1112            PathBuf::from("events").join("part-rank00000-of00001.root")
1113        );
1114    }
1115
1116    #[test]
1117    fn plan_fragments_serial_mode_keeps_all_fragments_in_order() {
1118        let fragments = vec![
1119            DataFragment {
1120                key: "a",
1121                global_start: 0,
1122                rows: 2,
1123            },
1124            DataFragment {
1125                key: "b",
1126                global_start: 2,
1127                rows: 3,
1128            },
1129        ];
1130
1131        let reads = plan_fragments(&fragments, ReadPlan::default()).unwrap();
1132
1133        assert_eq!(reads.len(), 2);
1134
1135        match &reads[0].selection {
1136            FragmentSelection::Range {
1137                local_start,
1138                local_len,
1139            } => {
1140                assert_eq!((*local_start, *local_len), (0, 2));
1141            }
1142            _ => panic!("expected range read"),
1143        }
1144
1145        match &reads[1].selection {
1146            FragmentSelection::Range {
1147                local_start,
1148                local_len,
1149            } => {
1150                assert_eq!((*local_start, *local_len), (0, 3));
1151            }
1152            _ => panic!("expected range read"),
1153        }
1154    }
1155
1156    use laddu_physics::vectors::RealVec4;
1157    #[cfg(feature = "mpi")]
1158    use mpi::traits::*;
1159    #[cfg(feature = "mpi")]
1160    use mpi_test::mpi_test;
1161
1162    #[cfg(feature = "mpi")]
1163    fn distributed_plan(
1164        partitioning: Partitioning,
1165        world: &impl mpi::topology::Communicator,
1166    ) -> ReadPlan {
1167        ReadPlan {
1168            chunk_size: None,
1169            distribution: Distribution::from_world(world).with_partitioning(partitioning),
1170        }
1171    }
1172
1173    #[cfg(feature = "mpi")]
1174    fn expected_contiguous_global_range(total_rows: u64, rank: usize, nranks: usize) -> (u64, u64) {
1175        let start = total_rows * rank as u64 / nranks as u64;
1176        let end = total_rows * (rank as u64 + 1) / nranks as u64;
1177        (start, end)
1178    }
1179
1180    #[cfg(feature = "mpi")]
1181    #[mpi_test(np = [2, 3, 4])]
1182    fn mpi_contiguous_plan_assigns_disjoint_ranges_covering_all_rows() {
1183        let universe = mpi::initialize().unwrap();
1184        let world = universe.world();
1185
1186        let rank = world.rank() as usize;
1187        let nranks = world.size() as usize;
1188
1189        let fragments = vec![
1190            DataFragment {
1191                key: "a",
1192                global_start: 0,
1193                rows: 4,
1194            },
1195            DataFragment {
1196                key: "b",
1197                global_start: 4,
1198                rows: 5,
1199            },
1200            DataFragment {
1201                key: "c",
1202                global_start: 9,
1203                rows: 3,
1204            },
1205        ];
1206
1207        let total_rows = fragments.iter().map(|f| f.rows).sum::<u64>();
1208        let plan = distributed_plan(Partitioning::Contiguous, &world);
1209        let reads = plan_fragments(&fragments, plan).unwrap();
1210
1211        let assigned_rows: u64 = reads
1212            .iter()
1213            .map(|read| match read.selection {
1214                FragmentSelection::Range { local_len, .. } => local_len as u64,
1215                FragmentSelection::StridedRows { .. } => panic!("expected range selection"),
1216            })
1217            .sum();
1218
1219        let (expected_start, expected_end) =
1220            expected_contiguous_global_range(total_rows, rank, nranks);
1221
1222        assert_eq!(assigned_rows, expected_end - expected_start);
1223
1224        for read in reads {
1225            let fragment = fragments
1226                .iter()
1227                .find(|fragment| fragment.key == read.key)
1228                .unwrap();
1229
1230            match read.selection {
1231                FragmentSelection::Range {
1232                    local_start,
1233                    local_len,
1234                } => {
1235                    let global_start = fragment.global_start + local_start as u64;
1236                    let global_end = global_start + local_len as u64;
1237
1238                    assert!(expected_start <= global_start);
1239                    assert!(global_end <= expected_end);
1240                    assert!(fragment.global_start <= global_start);
1241                    assert!(global_end <= fragment.global_start + fragment.rows);
1242                }
1243                FragmentSelection::StridedRows { .. } => panic!("expected range selection"),
1244            }
1245        }
1246    }
1247
1248    #[cfg(feature = "mpi")]
1249    #[mpi_test(np = [2, 3])]
1250    fn mpi_file_group_plan_assigns_fragment_by_rank_round_robin() {
1251        let universe = mpi::initialize().unwrap();
1252        let world = universe.world();
1253
1254        let rank = world.rank() as usize;
1255        let nranks = world.size() as usize;
1256
1257        let fragments = (0..8)
1258            .map(|i| DataFragment {
1259                key: i,
1260                global_start: 10 * i as u64,
1261                rows: 10,
1262            })
1263            .collect::<Vec<_>>();
1264
1265        let plan = distributed_plan(Partitioning::FileGroups, &world);
1266        let reads = plan_fragments(&fragments, plan).unwrap();
1267
1268        let keys = reads.iter().map(|read| read.key).collect::<Vec<_>>();
1269        let expected = (0..8).filter(|i| i % nranks == rank).collect::<Vec<_>>();
1270
1271        assert_eq!(keys, expected);
1272
1273        for read in reads {
1274            match read.selection {
1275                FragmentSelection::Range {
1276                    local_start,
1277                    local_len,
1278                } => {
1279                    assert_eq!(local_start, 0);
1280                    assert_eq!(local_len, 10);
1281                }
1282                FragmentSelection::StridedRows { .. } => panic!("expected range selection"),
1283            }
1284        }
1285    }
1286
1287    #[cfg(feature = "mpi")]
1288    #[mpi_test(np = [2, 3, 4])]
1289    fn mpi_rows_plan_assigns_strided_row_selection_with_world_rank() {
1290        let universe = mpi::initialize().unwrap();
1291        let world = universe.world();
1292
1293        let rank = world.rank() as usize;
1294        let nranks = world.size() as usize;
1295
1296        let fragments = vec![
1297            DataFragment {
1298                key: "a",
1299                global_start: 0,
1300                rows: 4,
1301            },
1302            DataFragment {
1303                key: "b",
1304                global_start: 4,
1305                rows: 5,
1306            },
1307        ];
1308
1309        let plan = distributed_plan(Partitioning::Rows, &world);
1310        let reads = plan_fragments(&fragments, plan).unwrap();
1311
1312        assert_eq!(reads.len(), fragments.len());
1313
1314        for (read, fragment) in reads.iter().zip(fragments.iter()) {
1315            assert_eq!(read.key, fragment.key);
1316
1317            match read.selection {
1318                FragmentSelection::StridedRows {
1319                    global_start,
1320                    rows,
1321                    rank: selected_rank,
1322                    nranks: selected_nranks,
1323                } => {
1324                    assert_eq!(global_start, fragment.global_start);
1325                    assert_eq!(rows, fragment.rows as usize);
1326                    assert_eq!(selected_rank, rank);
1327                    assert_eq!(selected_nranks, nranks);
1328                }
1329                FragmentSelection::Range { .. } => panic!("expected strided selection"),
1330            }
1331        }
1332    }
1333
1334    #[cfg(feature = "mpi")]
1335    #[mpi_test(np = [2, 3])]
1336    fn mpi_read_plan_and_write_plan_reflect_world_distribution() {
1337        let universe = mpi::initialize().unwrap();
1338        let world = universe.world();
1339
1340        let read_plan = ReadPlan {
1341            chunk_size: Some(7),
1342            distribution: Distribution::from_world(&world).with_partitioning(Partitioning::Rows),
1343        };
1344
1345        assert!(read_plan.is_distributed());
1346        assert_eq!(read_plan.rank(), world.rank() as usize);
1347        assert_eq!(read_plan.nranks(), world.size() as usize);
1348
1349        match read_plan.fragment_partitioning() {
1350            FragmentPartitioning::StridedRows => {}
1351            _ => panic!("expected strided row partitioning"),
1352        }
1353
1354        let write_plan = WritePlan::from(read_plan);
1355
1356        assert!(write_plan.is_distributed());
1357        assert_eq!(write_plan.rank(), world.rank() as usize);
1358        assert_eq!(write_plan.nranks(), world.size() as usize);
1359    }
1360}
laddu_data/io/mod.rs

laddu_data/io/
mod.rs
