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

laddu_data/io/
memory.rs
