laddu_data/io/
memory.rs

1use std::sync::Arc;
2
3use crate::{
4    LadduDataError, LadduDataResult,
5    data::{EventBatch, OwnedEvent},
6    io::{
7        DataFragment, EventSink, EventSource, FragmentedSource, ReadPlan, SinkState,
8        SourceCapabilities, WritePlan, fragmented_batches,
9    },
10    schema::Schema,
11};
12
13/// Replayable in-memory event source backed by one or more batches.
14#[derive(Clone, Debug)]
15pub struct MemorySource {
16    schema: Arc<Schema>,
17    batches: Arc<[EventBatch]>,
18}
19
20/// Key identifying one batch fragment in a [`MemorySource`].
21#[derive(Clone, Copy, Debug)]
22pub struct MemoryFragmentKey {
23    batch_index: usize,
24}
25
26impl MemorySource {
27    /// Creates an empty replayable source with the supplied schema.
28    ///
29    /// Empty derived datasets retain schema and source capabilities while
30    /// yielding no batches. This is useful for partitioning operations where
31    /// a valid empty result is distinct from an invalid source definition.
32    pub fn empty(schema: Arc<Schema>) -> Self {
33        Self {
34            schema,
35            batches: Arc::from([]),
36        }
37    }
38
39    /// Creates a source containing one batch.
40    pub fn new(batch: EventBatch) -> Self {
41        Self {
42            schema: Arc::clone(batch.schema()),
43            batches: Arc::from([batch]),
44        }
45    }
46
47    /// Validates and creates a source from nonempty schema-compatible batches.
48    ///
49    /// # Errors
50    ///
51    /// Returns [`LadduDataError`] when `batches` is empty or contains
52    /// incompatible schemas.
53    pub fn from_batches(batches: Vec<EventBatch>) -> LadduDataResult<Self> {
54        if batches.is_empty() {
55            return Err(LadduDataError::InvalidArgument(
56                "memory source requires at least one batch",
57            ));
58        }
59
60        let schema = Arc::clone(batches[0].schema());
61
62        for batch in &batches {
63            if batch.schema().as_ref() != schema.as_ref() {
64                return Err(LadduDataError::Schema(
65                    "memory source batches have different schemas".into(),
66                ));
67            }
68        }
69
70        Ok(Self {
71            schema,
72            batches: batches.into(),
73        })
74    }
75
76    /// Collects owned events into an in-memory source.
77    ///
78    /// # Errors
79    ///
80    /// Returns [`LadduDataError`] when an event does not match `schema` or
81    /// weighted and unweighted events are mixed.
82    pub fn from_events<I>(schema: Arc<Schema>, events: I) -> LadduDataResult<Self>
83    where
84        I: IntoIterator<Item = OwnedEvent>,
85    {
86        let batch = EventBatch::from_events(schema, events)?;
87        Ok(Self::new(batch))
88    }
89
90    /// Returns the shared schema.
91    pub fn schema_arc(&self) -> &Arc<Schema> {
92        &self.schema
93    }
94
95    /// Returns the backing batches.
96    pub fn batches_slice(&self) -> &[EventBatch] {
97        &self.batches
98    }
99
100    /// Consumes the source and returns its shared batches.
101    pub fn into_batches(self) -> Arc<[EventBatch]> {
102        self.batches
103    }
104
105    /// Consumes and concatenates all batches.
106    ///
107    /// # Errors
108    ///
109    /// Returns [`LadduDataError`] when no batches are present or their schemas
110    /// are incompatible.
111    pub fn into_batch(self) -> LadduDataResult<EventBatch> {
112        EventBatch::concat(&self.batches)
113    }
114}
115
116impl EventSource for MemorySource {
117    fn schema(&self) -> LadduDataResult<Arc<Schema>> {
118        Ok(Arc::clone(&self.schema))
119    }
120
121    fn capabilities(&self) -> SourceCapabilities {
122        SourceCapabilities {
123            exact_len: true,
124            exact_weighted_total: true,
125            random_access: true,
126            deterministic_partitioning: true,
127            predicate_pushdown: false,
128            projection_pushdown: false,
129            streaming: false,
130        }
131    }
132
133    fn num_events(&self) -> LadduDataResult<Option<u64>> {
134        Ok(Some(self.batches.iter().map(|b| b.len() as u64).sum()))
135    }
136
137    fn weighted_total(&self) -> LadduDataResult<Option<f64>> {
138        let total = self
139            .batches
140            .iter()
141            .map(|batch| (0..batch.len()).map(|i| batch.weights_at(i)).sum::<f64>())
142            .sum();
143
144        Ok(Some(total))
145    }
146
147    fn batches(
148        &self,
149        plan: ReadPlan,
150    ) -> LadduDataResult<Box<dyn Iterator<Item = LadduDataResult<EventBatch>> + Send>> {
151        fragmented_batches(Arc::new(self.clone()), plan)
152    }
153}
154
155impl FragmentedSource for MemorySource {
156    type Key = MemoryFragmentKey;
157
158    fn fragments(&self) -> LadduDataResult<Vec<DataFragment<Self::Key>>> {
159        let mut fragments = Vec::with_capacity(self.batches.len());
160        let mut global_start = 0_u64;
161
162        for (batch_index, batch) in self.batches.iter().enumerate() {
163            let rows = batch.len() as u64;
164
165            fragments.push(DataFragment {
166                key: MemoryFragmentKey { batch_index },
167                global_start,
168                rows,
169            });
170
171            global_start += rows;
172        }
173
174        Ok(fragments)
175    }
176
177    fn read_fragment_range(
178        &self,
179        key: &Self::Key,
180        local_start: usize,
181        local_len: usize,
182        chunk_size: Option<usize>,
183    ) -> LadduDataResult<Box<dyn Iterator<Item = LadduDataResult<EventBatch>> + Send>> {
184        if matches!(chunk_size, Some(0)) {
185            return Err(LadduDataError::InvalidArgument(
186                "chunk_size must be nonzero",
187            ));
188        }
189
190        let batch = self
191            .batches
192            .get(key.batch_index)
193            .ok_or_else(|| LadduDataError::Source("invalid memory batch index".into()))?
194            .clone();
195
196        let end = local_start
197            .checked_add(local_len)
198            .ok_or(LadduDataError::InvalidArgument(
199                "slice range overflows usize",
200            ))?;
201
202        if end > batch.len() {
203            return Err(LadduDataError::InvalidArgument(
204                "memory fragment range exceeds batch length",
205            ));
206        }
207
208        Ok(Box::new(MemoryRangeIter {
209            batch,
210            pos: local_start,
211            end,
212            chunk_size: chunk_size.unwrap_or(local_len.max(1)),
213        }))
214    }
215}
216
217struct MemoryRangeIter {
218    batch: EventBatch,
219    pos: usize,
220    end: usize,
221    chunk_size: usize,
222}
223
224impl Iterator for MemoryRangeIter {
225    type Item = LadduDataResult<EventBatch>;
226
227    fn next(&mut self) -> Option<Self::Item> {
228        if self.pos >= self.end {
229            return None;
230        }
231
232        let next = (self.pos + self.chunk_size).min(self.end);
233        let batch = self.batch.slice(self.pos, next);
234        self.pos = next;
235
236        Some(Ok(batch))
237    }
238}
239
240/// Event sink that collects written batches in memory.
241#[derive(Clone, Debug, Default)]
242pub struct MemorySink {
243    schema: Option<Arc<Schema>>,
244    batches: Vec<EventBatch>,
245    state: SinkState,
246}
247
248impl MemorySink {
249    /// Creates an empty sink.
250    pub fn new() -> Self {
251        Self::default()
252    }
253
254    /// Returns the schema supplied to the most recent write.
255    pub fn schema(&self) -> Option<&Arc<Schema>> {
256        self.schema.as_ref()
257    }
258
259    /// Returns collected batches.
260    pub fn batches(&self) -> &[EventBatch] {
261        &self.batches
262    }
263
264    /// Consumes the sink and returns collected batches.
265    pub fn into_batches(self) -> Vec<EventBatch> {
266        self.batches
267    }
268
269    /// Consumes the sink and builds an in-memory source.
270    ///
271    /// # Errors
272    ///
273    /// Returns [`LadduDataError`] when the sink contains no batches or
274    /// incompatible schemas.
275    pub fn into_source(self) -> LadduDataResult<MemorySource> {
276        MemorySource::from_batches(self.batches)
277    }
278
279    /// Consumes the sink and concatenates collected batches.
280    ///
281    /// # Errors
282    ///
283    /// Returns [`LadduDataError`] when no batches were collected or their
284    /// schemas are incompatible.
285    pub fn into_batch(self) -> LadduDataResult<EventBatch> {
286        EventBatch::concat(&self.batches)
287    }
288
289    /// Clears schema, batches, and completion state.
290    pub fn clear(&mut self) {
291        self.schema = None;
292        self.batches.clear();
293        self.state = SinkState::Idle;
294    }
295}
296
297impl EventSink for MemorySink {
298    fn retains_batches(&self) -> bool {
299        true
300    }
301
302    fn begin(&mut self, schema: Arc<Schema>, _plan: WritePlan) -> LadduDataResult<()> {
303        match self.state {
304            SinkState::Idle => {}
305            SinkState::Writing => {
306                return Err(LadduDataError::Sink(
307                    "memory sink already initialized".into(),
308                ));
309            }
310            SinkState::Failed => {
311                return Err(LadduDataError::Sink(
312                    "memory sink requires abort after failure".into(),
313                ));
314            }
315        }
316
317        self.schema = Some(schema);
318        self.batches.clear();
319        self.state = SinkState::Writing;
320        Ok(())
321    }
322
323    fn write_batch(&mut self, batch: &EventBatch) -> LadduDataResult<()> {
324        if !matches!(self.state, SinkState::Writing) {
325            return Err(LadduDataError::Sink(
326                match self.state {
327                    SinkState::Idle => "memory sink not initialized",
328                    SinkState::Failed => "memory sink requires abort after failure",
329                    SinkState::Writing => unreachable!(),
330                }
331                .into(),
332            ));
333        }
334
335        let schema = self
336            .schema
337            .as_ref()
338            .ok_or_else(|| LadduDataError::Sink("memory sink not initialized".into()))?;
339
340        if schema.as_ref() != batch.schema().as_ref() {
341            return Err(LadduDataError::Sink(
342                "batch schema does not match memory sink schema".into(),
343            ));
344        }
345
346        self.batches.push(batch.clone());
347        Ok(())
348    }
349
350    fn finish(&mut self) -> LadduDataResult<()> {
351        if matches!(self.state, SinkState::Failed) {
352            return Err(LadduDataError::Sink(
353                "memory sink requires abort after failure".into(),
354            ));
355        }
356        self.state = SinkState::Idle;
357        Ok(())
358    }
359
360    fn abort(&mut self) -> LadduDataResult<()> {
361        self.schema = None;
362        self.batches.clear();
363        self.state = SinkState::Idle;
364        Ok(())
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use laddu_physics::vectors::RealVec4;
371
372    use super::*;
373    use crate::data::EventBatchBuilder;
374
375    fn v(x: f64) -> RealVec4 {
376        RealVec4 {
377            e: x,
378            px: x,
379            py: x,
380            pz: x,
381        }
382    }
383
384    fn schema() -> Arc<Schema> {
385        Arc::new(Schema::new(["p"], ["id"], true).unwrap())
386    }
387
388    fn batch(start: usize, len: usize) -> EventBatch {
389        let schema = schema();
390        let mut builder = EventBatchBuilder::with_capacity(schema, len);
391
392        for i in start..start + len {
393            builder
394                .push_weighted([v(i as f64)], [i as f64], 1.0 + i as f64)
395                .unwrap();
396        }
397
398        builder.finish().unwrap()
399    }
400
401    fn concat_scalars(batches: Vec<EventBatch>) -> Vec<f64> {
402        EventBatch::concat(&batches)
403            .unwrap()
404            .scalar_column(0)
405            .to_vec()
406    }
407
408    #[test]
409    fn memory_source_reports_exact_capabilities_counts_and_weighted_total() {
410        let source = MemorySource::from_batches(vec![batch(0, 2), batch(2, 3)]).unwrap();
411
412        let caps = source.capabilities();
413
414        assert!(caps.exact_len);
415        assert!(caps.exact_weighted_total);
416        assert!(caps.random_access);
417        assert!(caps.deterministic_partitioning);
418        assert!(!caps.streaming);
419
420        assert_eq!(source.num_events().unwrap(), Some(5));
421        assert_eq!(
422            source.weighted_total().unwrap(),
423            Some(1.0 + 2.0 + 3.0 + 4.0 + 5.0)
424        );
425    }
426
427    #[test]
428    fn memory_source_coalesces_by_default_but_respects_chunked_read_plan() {
429        let source = MemorySource::from_batches(vec![batch(0, 2), batch(2, 3)]).unwrap();
430
431        let default_batches: Vec<EventBatch> = source
432            .batches(ReadPlan::default())
433            .unwrap()
434            .map(Result::unwrap)
435            .collect();
436
437        assert_eq!(default_batches.len(), 1);
438        assert_eq!(
439            default_batches[0].scalar_column(0),
440            &[0.0, 1.0, 2.0, 3.0, 4.0]
441        );
442
443        let chunked_batches: Vec<EventBatch> = source
444            .batches(ReadPlan {
445                chunk_size: Some(2),
446                #[cfg(feature = "mpi")]
447                distribution: Default::default(),
448            })
449            .unwrap()
450            .map(Result::unwrap)
451            .collect();
452
453        assert_eq!(
454            chunked_batches
455                .iter()
456                .map(EventBatch::len)
457                .collect::<Vec<_>>(),
458            vec![2, 2, 1]
459        );
460        assert_eq!(
461            concat_scalars(chunked_batches),
462            vec![0.0, 1.0, 2.0, 3.0, 4.0]
463        );
464    }
465
466    #[test]
467    fn memory_source_rejects_empty_or_schema_mismatched_batches() {
468        assert!(matches!(
469            MemorySource::from_batches(vec![]),
470            Err(LadduDataError::InvalidArgument(_))
471        ));
472
473        let first = batch(0, 1);
474
475        let other_schema = Arc::new(Schema::new(["q"], ["id"], true).unwrap());
476        let mut builder = EventBatchBuilder::new(other_schema);
477        builder.push_weighted([v(10.0)], [10.0], 1.0).unwrap();
478        let second = builder.finish().unwrap();
479
480        assert!(matches!(
481            MemorySource::from_batches(vec![first, second]),
482            Err(LadduDataError::Schema(_))
483        ));
484    }
485
486    #[test]
487    fn memory_sink_validates_lifecycle_schema_and_can_be_reused() {
488        let mut sink = MemorySink::new();
489        let first = batch(0, 2);
490
491        assert!(matches!(
492            sink.write_batch(&first),
493            Err(LadduDataError::Sink(_))
494        ));
495
496        sink.begin(Arc::clone(first.schema()), WritePlan::default())
497            .unwrap();
498        sink.write_batch(&first).unwrap();
499        sink.finish().unwrap();
500
501        assert_eq!(sink.batches().len(), 1);
502        assert_eq!(sink.batches()[0].scalar_column(0), &[0.0, 1.0]);
503
504        let mismatched_schema = Arc::new(Schema::new(["other"], ["id"], true).unwrap());
505        let mut builder = EventBatchBuilder::new(mismatched_schema);
506        builder.push_weighted([v(9.0)], [9.0], 9.0).unwrap();
507        let mismatched = builder.finish().unwrap();
508
509        assert!(matches!(
510            sink.write_batch(&mismatched),
511            Err(LadduDataError::Sink(_))
512        ));
513
514        sink.clear();
515        assert!(sink.schema().is_none());
516        assert!(sink.batches().is_empty());
517    }
518
519    #[test]
520    fn memory_sink_into_source_roundtrips_captured_batches() {
521        let mut sink = MemorySink::new();
522        let first = batch(0, 2);
523        let second = batch(2, 2);
524
525        sink.begin(Arc::clone(first.schema()), WritePlan::default())
526            .unwrap();
527        sink.write_batch(&first).unwrap();
528        sink.write_batch(&second).unwrap();
529        sink.finish().unwrap();
530
531        let source = sink.into_source().unwrap();
532        let merged = source.into_batch().unwrap();
533
534        assert_eq!(merged.scalar_column(0), &[0.0, 1.0, 2.0, 3.0]);
535        assert_eq!(merged.weights_column().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
536    }
537
538    #[test]
539    fn memory_sink_abort_discards_partial_batches_and_allows_reuse() {
540        let mut sink = MemorySink::new();
541        let first = batch(0, 2);
542
543        sink.begin(Arc::clone(first.schema()), WritePlan::default())
544            .unwrap();
545        sink.write_batch(&first).unwrap();
546        sink.abort().unwrap();
547        assert!(sink.schema().is_none());
548        assert!(sink.batches().is_empty());
549
550        sink.begin(Arc::clone(first.schema()), WritePlan::default())
551            .unwrap();
552        sink.finish().unwrap();
553    }
554}
laddu_data/io/memory.rs

laddu_data/io/
memory.rs
