spatialrust_records/
stream.rs

1//! Chunked spatial-record sources and sinks.
2
3use spatialrust_core::{PointBuffer, PointBufferSet, PointCloud, SpatialMetadata, SpatialTensor};
4
5use crate::{RecordProvenance, RecordsError, RecordsResult, SchemaDescriptor, SpatialRecord};
6
7/// Pull-based source of versioned spatial records.
8pub trait SpatialRecordSource {
9    /// Returns the schema contract for every emitted record.
10    fn schema(&self) -> &SchemaDescriptor;
11
12    /// Returns the next record, or `None` when exhausted.
13    fn next_record(&mut self) -> Option<RecordsResult<SpatialRecord>>;
14}
15
16/// Push-based sink for versioned spatial records.
17pub trait SpatialRecordSink {
18    /// Accepts one record.
19    fn write_record(&mut self, record: &SpatialRecord) -> RecordsResult<()>;
20
21    /// Finalizes the sink. Default is a no-op.
22    fn finish(&mut self) -> RecordsResult<()> {
23        Ok(())
24    }
25}
26
27/// Splits one in-memory cloud into fixed-size record chunks.
28pub struct MemoryChunkSource {
29    schema: SchemaDescriptor,
30    metadata: SpatialMetadata,
31    provenance: RecordProvenance,
32    cloud: PointCloud,
33    chunk_size: usize,
34    offset: usize,
35}
36
37impl MemoryChunkSource {
38    /// Creates a chunked source over `cloud` using `chunk_size` points per record.
39    pub fn try_new(
40        schema: SchemaDescriptor,
41        cloud: PointCloud,
42        chunk_size: usize,
43    ) -> RecordsResult<Self> {
44        Self::try_new_with_provenance(schema, cloud, chunk_size, RecordProvenance::default())
45    }
46
47    /// Creates a chunked source with explicit lineage copied to every chunk.
48    pub fn try_new_with_provenance(
49        schema: SchemaDescriptor,
50        cloud: PointCloud,
51        chunk_size: usize,
52        provenance: RecordProvenance,
53    ) -> RecordsResult<Self> {
54        if chunk_size == 0 {
55            return Err(RecordsError::InvalidConfiguration("chunk_size must be positive".into()));
56        }
57        if cloud.schema() != schema.point_schema() {
58            return Err(RecordsError::SchemaMismatch(
59                "chunk source cloud schema must match descriptor".into(),
60            ));
61        }
62        cloud.validate()?;
63        Ok(Self {
64            metadata: cloud.metadata().clone(),
65            provenance,
66            schema,
67            cloud,
68            chunk_size,
69            offset: 0,
70        })
71    }
72
73    /// Creates a source using [`SpatialTensor`]’s default chunk size as a hint.
74    pub fn try_with_default_chunk(
75        schema: SchemaDescriptor,
76        cloud: PointCloud,
77    ) -> RecordsResult<Self> {
78        let _ = SpatialTensor::new(&cloud, spatialrust_core::DEFAULT_SPATIAL_TENSOR_CHUNK_SIZE)?;
79        Self::try_new(schema, cloud, spatialrust_core::DEFAULT_SPATIAL_TENSOR_CHUNK_SIZE)
80    }
81}
82
83impl SpatialRecordSource for MemoryChunkSource {
84    fn schema(&self) -> &SchemaDescriptor {
85        &self.schema
86    }
87
88    fn next_record(&mut self) -> Option<RecordsResult<SpatialRecord>> {
89        if self.offset >= self.cloud.len() {
90            return None;
91        }
92        let end = (self.offset + self.chunk_size).min(self.cloud.len());
93        let range = self.offset..end;
94        self.offset = end;
95        Some(slice_cloud(&self.schema, &self.cloud, &self.metadata, &self.provenance, range))
96    }
97}
98
99/// Collects records into an owned point cloud.
100#[derive(Clone, Debug, Default)]
101pub struct MemoryChunkSink {
102    schema: Option<SchemaDescriptor>,
103    buffers: PointBufferSet,
104    metadata: SpatialMetadata,
105    provenance: RecordProvenance,
106    len: usize,
107}
108
109impl MemoryChunkSink {
110    /// Creates an empty sink.
111    #[must_use]
112    pub fn new() -> Self {
113        Self::default()
114    }
115
116    /// Consumes the sink into one assembled record.
117    pub fn into_record(self) -> RecordsResult<Option<SpatialRecord>> {
118        let Some(schema) = self.schema else {
119            return Ok(None);
120        };
121        if self.len == 0 {
122            let cloud = PointCloud::try_from_parts(
123                schema.point_schema().clone(),
124                PointBufferSet::new(),
125                self.metadata,
126            )?;
127            return Ok(Some(SpatialRecord::try_new_with_provenance(
128                schema,
129                cloud,
130                self.provenance,
131            )?));
132        }
133        let cloud =
134            PointCloud::try_from_parts(schema.point_schema().clone(), self.buffers, self.metadata)?;
135        Ok(Some(SpatialRecord::try_new_with_provenance(schema, cloud, self.provenance)?))
136    }
137}
138
139impl SpatialRecordSink for MemoryChunkSink {
140    fn write_record(&mut self, record: &SpatialRecord) -> RecordsResult<()> {
141        match &self.schema {
142            None => {
143                self.schema = Some(record.schema().clone());
144                self.metadata = record.metadata().clone();
145                self.provenance = record.provenance().clone().without_sequence();
146            }
147            Some(schema) if schema != record.schema() => {
148                return Err(RecordsError::SchemaMismatch(
149                    "chunk sink requires a homogeneous schema across records".into(),
150                ));
151            }
152            Some(_) => {}
153        }
154        for field in record.schema().point_schema().fields() {
155            let source = record.cloud().field(&field.name)?;
156            match self.buffers.get_mut(&field.name) {
157                Some(dst) => append_buffer(dst, source)?,
158                None => {
159                    self.buffers.insert(field.name.clone(), clone_buffer(source)?);
160                }
161            }
162        }
163        self.len += record.cloud().len();
164        Ok(())
165    }
166}
167
168fn slice_cloud(
169    schema: &SchemaDescriptor,
170    cloud: &PointCloud,
171    metadata: &SpatialMetadata,
172    provenance: &RecordProvenance,
173    range: std::ops::Range<usize>,
174) -> RecordsResult<SpatialRecord> {
175    let mut buffers = PointBufferSet::new();
176    for field in schema.point_schema().fields() {
177        let source = cloud.field(&field.name)?;
178        buffers.insert(field.name.clone(), slice_buffer(source, &range)?);
179    }
180    let chunk =
181        PointCloud::try_from_parts(schema.point_schema().clone(), buffers, metadata.clone())?;
182    SpatialRecord::try_new_with_provenance(schema.clone(), chunk, provenance.clone())
183}
184
185fn slice_buffer(
186    buffer: &PointBuffer,
187    range: &std::ops::Range<usize>,
188) -> RecordsResult<PointBuffer> {
189    Ok(match buffer {
190        PointBuffer::F32(values) => PointBuffer::F32(values[range.clone()].to_vec()),
191        PointBuffer::F64(values) => PointBuffer::F64(values[range.clone()].to_vec()),
192        PointBuffer::U8(values) => PointBuffer::U8(values[range.clone()].to_vec()),
193        PointBuffer::U16(values) => PointBuffer::U16(values[range.clone()].to_vec()),
194        PointBuffer::U32(values) => PointBuffer::U32(values[range.clone()].to_vec()),
195        PointBuffer::I32(values) => PointBuffer::I32(values[range.clone()].to_vec()),
196    })
197}
198
199fn clone_buffer(buffer: &PointBuffer) -> RecordsResult<PointBuffer> {
200    Ok(match buffer {
201        PointBuffer::F32(values) => PointBuffer::F32(values.clone()),
202        PointBuffer::F64(values) => PointBuffer::F64(values.clone()),
203        PointBuffer::U8(values) => PointBuffer::U8(values.clone()),
204        PointBuffer::U16(values) => PointBuffer::U16(values.clone()),
205        PointBuffer::U32(values) => PointBuffer::U32(values.clone()),
206        PointBuffer::I32(values) => PointBuffer::I32(values.clone()),
207    })
208}
209
210fn append_buffer(dst: &mut PointBuffer, src: &PointBuffer) -> RecordsResult<()> {
211    match (dst, src) {
212        (PointBuffer::F32(dst), PointBuffer::F32(src)) => dst.extend_from_slice(src),
213        (PointBuffer::F64(dst), PointBuffer::F64(src)) => dst.extend_from_slice(src),
214        (PointBuffer::U8(dst), PointBuffer::U8(src)) => dst.extend_from_slice(src),
215        (PointBuffer::U16(dst), PointBuffer::U16(src)) => dst.extend_from_slice(src),
216        (PointBuffer::U32(dst), PointBuffer::U32(src)) => dst.extend_from_slice(src),
217        (PointBuffer::I32(dst), PointBuffer::I32(src)) => dst.extend_from_slice(src),
218        (dst, src) => {
219            return Err(RecordsError::SchemaMismatch(format!(
220                "cannot append {:?} into {:?}",
221                src.dtype(),
222                dst.dtype()
223            )));
224        }
225    }
226    Ok(())
227}
228
229#[cfg(test)]
230mod tests {
231    use super::{MemoryChunkSink, MemoryChunkSource, SpatialRecordSink, SpatialRecordSource};
232    use crate::{SchemaDescriptor, SchemaVersion, SpatialRecord};
233    use spatialrust_core::{
234        PointBuffer, PointBufferSet, PointCloud, SpatialMetadata, StandardSchemas,
235    };
236
237    #[test]
238    fn memory_chunk_roundtrip() {
239        let mut buffers = PointBufferSet::new();
240        buffers.insert("x", PointBuffer::from_f32(vec![0.0, 1.0, 2.0, 3.0, 4.0]));
241        buffers.insert("y", PointBuffer::from_f32(vec![0.0; 5]));
242        buffers.insert("z", PointBuffer::from_f32(vec![1.0; 5]));
243        let cloud = PointCloud::try_from_parts(
244            StandardSchemas::point_xyz(),
245            buffers,
246            SpatialMetadata::default(),
247        )
248        .unwrap();
249        let schema =
250            SchemaDescriptor::try_new("point", SchemaVersion::new(1, 0), cloud.schema().clone())
251                .unwrap();
252        let provenance = crate::RecordProvenance::try_new("memory-source")
253            .unwrap()
254            .with_stream_id("points")
255            .with_sequence(Some(7));
256        let mut source = MemoryChunkSource::try_new_with_provenance(
257            schema.clone(),
258            cloud,
259            2,
260            provenance.clone(),
261        )
262        .unwrap();
263        let mut sink = MemoryChunkSink::new();
264        let mut chunks = 0;
265        while let Some(record) = source.next_record() {
266            sink.write_record(&record.unwrap()).unwrap();
267            chunks += 1;
268        }
269        assert_eq!(chunks, 3);
270        let assembled = sink.into_record().unwrap().unwrap();
271        assert_eq!(assembled.cloud().len(), 5);
272        assert_eq!(assembled.schema(), &schema);
273        assert_eq!(assembled.provenance().source_id, provenance.source_id);
274        assert_eq!(assembled.provenance().stream_id, provenance.stream_id);
275        assert_eq!(assembled.provenance().sequence, None);
276        let _ = SpatialRecord::try_new(schema, assembled.into_cloud());
277    }
278}
spatialrust_records/stream.rs

spatialrust_records/
stream.rs
