alopex_dataframe/io/
columnar_segment.rs

1//! Bounded adapter for provisioned V08 columnar segment cursors.
2//!
3//! The concrete V08/KVS implementation lives in `alopex-embedded`, which
4//! depends on this crate.  Keeping the cursor trait in `physical::streaming`
5//! avoids a reverse dependency while retaining the same ownership and terminal
6//! rules as CSV and Parquet sources.
7
8use std::fmt;
9use std::sync::Arc;
10
11use arrow::datatypes::SchemaRef;
12
13use crate::physical::budget::{ResourceReservation, ResourceScope};
14use crate::physical::{
15    record_batch_memory_size, BatchOpenContext, BatchSource, BatchSourceFactory,
16    ColumnarSegmentCursor, ColumnarSegmentProvider, SourceLimit, StreamBatch,
17};
18use crate::{DataFrameError, Result};
19
20/// Re-consumable bounded source factory over an adapter-provided V08 cursor.
21#[derive(Clone)]
22pub struct ColumnarSegmentBatchSourceFactory {
23    provider: Arc<dyn ColumnarSegmentProvider>,
24}
25
26impl fmt::Debug for ColumnarSegmentBatchSourceFactory {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        f.debug_struct("ColumnarSegmentBatchSourceFactory")
29            .field("source_name", &self.provider.source_name())
30            .finish_non_exhaustive()
31    }
32}
33
34impl ColumnarSegmentBatchSourceFactory {
35    /// Creates a factory from a read-only V08 segment provider.
36    pub fn new(provider: Arc<dyn ColumnarSegmentProvider>) -> Self {
37        Self { provider }
38    }
39}
40
41impl BatchSourceFactory for ColumnarSegmentBatchSourceFactory {
42    fn source_name(&self) -> &'static str {
43        self.provider.source_name()
44    }
45
46    fn schema(&self) -> Result<SchemaRef> {
47        Err(DataFrameError::streaming_unsupported(
48            "columnar_segment_scan",
49            "schema_is_available_after_bounded_open",
50        ))
51    }
52
53    fn source_limits(&self) -> Vec<SourceLimit> {
54        self.provider.source_limits()
55    }
56
57    fn open(&self, context: BatchOpenContext) -> Result<Box<dyn BatchSource>> {
58        // Opening a V08 cursor deserializes bounded header/schema/directory
59        // objects. Reserve the complete configured limit before that opaque
60        // allocation, then retain only the cursor's conservative footprint.
61        let mut metadata_reservation = context
62            .budget
63            .reserve(ResourceScope::Source, context.options.memory_limit_bytes)?;
64        let cursor = self.provider.open(context.options.memory_limit_bytes)?;
65        let retained_metadata = cursor.metadata_footprint_upper_bound();
66        if retained_metadata > context.options.memory_limit_bytes {
67            return Err(DataFrameError::resource_limit_exceeded(
68                context.options.memory_limit_bytes,
69                retained_metadata,
70                context.budget.max_in_flight_batches(),
71                context.budget.usage().reserved_batches,
72                ResourceScope::Source,
73            ));
74        }
75        metadata_reservation.shrink_to(retained_metadata);
76        let schema = cursor.schema();
77
78        Ok(Box::new(ColumnarSegmentBatchSource {
79            context,
80            cursor,
81            schema,
82            metadata_reservation: Some(metadata_reservation),
83        }))
84    }
85}
86
87struct ColumnarSegmentBatchSource {
88    context: BatchOpenContext,
89    cursor: Box<dyn ColumnarSegmentCursor>,
90    schema: SchemaRef,
91    metadata_reservation: Option<ResourceReservation>,
92}
93
94impl BatchSource for ColumnarSegmentBatchSource {
95    fn schema(&self) -> SchemaRef {
96        self.schema.clone()
97    }
98
99    fn next_batch(&mut self) -> Result<Option<StreamBatch>> {
100        let Some(plan) = self.cursor.next_batch_plan()? else {
101            return Ok(None);
102        };
103        let reservation_upper_bound = plan.reservation_upper_bound();
104        let reservation = self
105            .context
106            .budget
107            .reserve_batch(ResourceScope::Decode, reservation_upper_bound)?;
108        let Some(batch) = self.cursor.next_batch()? else {
109            drop(reservation);
110            return Err(DataFrameError::schema_mismatch(
111                "V08 cursor returned no batch after declaring a batch plan",
112            ));
113        };
114
115        let actual_arrow_bytes = record_batch_memory_size(&batch);
116        if actual_arrow_bytes > plan.arrow_allocation_upper_bound {
117            drop(reservation);
118            return Err(DataFrameError::resource_limit_exceeded(
119                self.context.budget.memory_limit_bytes(),
120                self.context
121                    .budget
122                    .usage()
123                    .reserved_bytes
124                    .saturating_add(actual_arrow_bytes),
125                self.context.budget.max_in_flight_batches(),
126                self.context.budget.usage().reserved_batches,
127                ResourceScope::Decode,
128            ));
129        }
130
131        Ok(Some(StreamBatch::new(batch, reservation)))
132    }
133
134    fn close(&mut self) -> Result<()> {
135        let result = self.cursor.close();
136        self.metadata_reservation.take();
137        result
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use std::collections::VecDeque;
144    use std::num::NonZeroUsize;
145    use std::sync::atomic::{AtomicUsize, Ordering};
146
147    use arrow::array::{ArrayRef, Int64Array};
148    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
149    use arrow::record_batch::RecordBatch;
150
151    use super::*;
152    use crate::physical::{
153        ColumnarSegmentBatchPlan, ColumnarSegmentCursor, ColumnarSegmentProvider, DataFrameStream,
154        PlanSubject, SourceLimit,
155    };
156    use crate::{DataFrameError, Result};
157
158    #[derive(Clone)]
159    struct TestProvider {
160        batches: Vec<RecordBatch>,
161        plan: ColumnarSegmentBatchPlan,
162        metadata_bytes: u64,
163        next_calls: Arc<AtomicUsize>,
164        close_calls: Arc<AtomicUsize>,
165    }
166
167    impl ColumnarSegmentProvider for TestProvider {
168        fn source_limits(&self) -> Vec<SourceLimit> {
169            vec![SourceLimit {
170                source: PlanSubject::ColumnarSegmentScan,
171                code: "requires_v08_chunked_layout",
172                description: "legacy V2 segment blobs are not streaming sources",
173            }]
174        }
175
176        fn open(&self, _metadata_limit_bytes: u64) -> Result<Box<dyn ColumnarSegmentCursor>> {
177            Ok(Box::new(TestCursor {
178                batches: self.batches.clone().into(),
179                plan: self.plan,
180                metadata_bytes: self.metadata_bytes,
181                next_calls: self.next_calls.clone(),
182                close_calls: self.close_calls.clone(),
183            }))
184        }
185    }
186
187    struct TestCursor {
188        batches: VecDeque<RecordBatch>,
189        plan: ColumnarSegmentBatchPlan,
190        metadata_bytes: u64,
191        next_calls: Arc<AtomicUsize>,
192        close_calls: Arc<AtomicUsize>,
193    }
194
195    impl ColumnarSegmentCursor for TestCursor {
196        fn schema(&self) -> SchemaRef {
197            schema()
198        }
199
200        fn metadata_footprint_upper_bound(&self) -> u64 {
201            self.metadata_bytes
202        }
203
204        fn next_batch_plan(&self) -> Result<Option<ColumnarSegmentBatchPlan>> {
205            Ok((!self.batches.is_empty()).then_some(self.plan))
206        }
207
208        fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
209            self.next_calls.fetch_add(1, Ordering::SeqCst);
210            Ok(self.batches.pop_front())
211        }
212
213        fn close(&mut self) -> Result<()> {
214            self.close_calls.fetch_add(1, Ordering::SeqCst);
215            Ok(())
216        }
217    }
218
219    fn schema() -> SchemaRef {
220        Arc::new(Schema::new(vec![Field::new(
221            "value",
222            DataType::Int64,
223            false,
224        )]))
225    }
226
227    fn batch(value: i64) -> RecordBatch {
228        RecordBatch::try_new(
229            schema(),
230            vec![Arc::new(Int64Array::from(vec![value])) as ArrayRef],
231        )
232        .unwrap()
233    }
234
235    fn options(limit: u64) -> crate::physical::budget::StreamOptions {
236        crate::physical::budget::StreamOptions::new(
237            limit,
238            NonZeroUsize::new(1).unwrap(),
239            NonZeroUsize::new(1).unwrap(),
240        )
241    }
242
243    #[test]
244    fn source_reserves_before_cursor_read_and_releases_on_exhaustion() {
245        let next_calls = Arc::new(AtomicUsize::new(0));
246        let close_calls = Arc::new(AtomicUsize::new(0));
247        let provider = TestProvider {
248            batches: vec![batch(1), batch(2)],
249            plan: ColumnarSegmentBatchPlan {
250                decoded_bytes: 128,
251                arrow_allocation_upper_bound: 1024,
252            },
253            metadata_bytes: 256,
254            next_calls: next_calls.clone(),
255            close_calls: close_calls.clone(),
256        };
257        let factory = ColumnarSegmentBatchSourceFactory::new(Arc::new(provider));
258        let mut stream = DataFrameStream::from_factory(&factory, options(4096)).unwrap();
259
260        assert_eq!(next_calls.load(Ordering::SeqCst), 0);
261        assert_eq!(stream.next_batch().unwrap().unwrap().height(), 1);
262        assert_eq!(stream.next_batch().unwrap().unwrap().height(), 1);
263        assert!(stream.next_batch().unwrap().is_none());
264        assert_eq!(close_calls.load(Ordering::SeqCst), 1);
265        assert_eq!(stream.budget().usage().reserved_bytes, 0);
266    }
267
268    #[test]
269    fn oversized_row_group_is_rejected_before_cursor_read() {
270        let next_calls = Arc::new(AtomicUsize::new(0));
271        let close_calls = Arc::new(AtomicUsize::new(0));
272        let provider = TestProvider {
273            batches: vec![batch(1)],
274            plan: ColumnarSegmentBatchPlan {
275                decoded_bytes: 2048,
276                arrow_allocation_upper_bound: 2048,
277            },
278            metadata_bytes: 128,
279            next_calls: next_calls.clone(),
280            close_calls: close_calls.clone(),
281        };
282        let factory = ColumnarSegmentBatchSourceFactory::new(Arc::new(provider));
283        let mut stream = DataFrameStream::from_factory(&factory, options(1024)).unwrap();
284
285        assert!(matches!(
286            stream.next_batch(),
287            Err(DataFrameError::StreamFailed {
288                code: "resource_limit_exceeded",
289                ..
290            })
291        ));
292        assert_eq!(next_calls.load(Ordering::SeqCst), 0);
293        assert_eq!(close_calls.load(Ordering::SeqCst), 1);
294    }
295}
alopex_dataframe/io/columnar_segment.rs

alopex_dataframe/io/
columnar_segment.rs
