alopex_dataframe/io/
streaming_parquet.rs

1//! Bounded Parquet source backed by one row-group reader at a time.
2//!
3//! Metadata is read once under a footer reservation. Each selected row group is then checked
4//! against the remaining resource budget before a Parquet reader is built, and every decoded
5//! Arrow batch receives its reservation before `next()` is called.
6
7use std::collections::VecDeque;
8use std::fs::File;
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11
12use arrow::datatypes::{Schema, SchemaRef};
13use parquet::arrow::arrow_reader::{
14    ArrowReaderMetadata, ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder,
15};
16use parquet::arrow::ProjectionMask;
17use parquet::file::metadata::ParquetMetaData;
18
19use crate::io::options::ParquetReadOptions;
20use crate::physical::budget::{ResourceReservation, ResourceScope};
21use crate::physical::{
22    BatchOpenContext, BatchSource, BatchSourceFactory, PlanSubject, SourceLimit, StreamBatch,
23};
24use crate::{DataFrameError, Result};
25
26/// Re-consumable bounded Parquet source factory.
27#[derive(Debug, Clone)]
28pub struct ParquetBatchSourceFactory {
29    path: PathBuf,
30    options: ParquetReadOptions,
31}
32
33impl ParquetBatchSourceFactory {
34    /// Construct a factory for a Parquet path and read options.
35    pub fn new(path: impl AsRef<Path>, options: ParquetReadOptions) -> Self {
36        Self {
37            path: path.as_ref().to_path_buf(),
38            options,
39        }
40    }
41
42    /// Construct a factory from a physical Parquet scan.
43    pub(crate) fn from_scan(
44        path: PathBuf,
45        predicate: Option<crate::Expr>,
46        columns: Option<Vec<String>>,
47    ) -> Self {
48        Self {
49            path,
50            options: ParquetReadOptions {
51                columns,
52                predicate,
53                ..ParquetReadOptions::default()
54            },
55        }
56    }
57}
58
59impl BatchSourceFactory for ParquetBatchSourceFactory {
60    fn source_name(&self) -> &'static str {
61        "parquet"
62    }
63
64    fn schema(&self) -> Result<SchemaRef> {
65        Err(DataFrameError::streaming_unsupported(
66            "parquet_scan",
67            "schema_is_available_after_bounded_open",
68        ))
69    }
70
71    fn source_limits(&self) -> Vec<SourceLimit> {
72        vec![
73            SourceLimit {
74                source: PlanSubject::ParquetScan,
75                code: "stable_row_group_order",
76                description: "Parquet streaming preserves selected row-group and row order",
77            },
78            SourceLimit {
79                source: PlanSubject::ParquetScan,
80                code: "row_group_must_fit_resource_bound",
81                description: "A selected row group whose declared upper bound exceeds the resource limit is rejected before page decode",
82            },
83        ]
84    }
85
86    fn open(&self, context: BatchOpenContext) -> Result<Box<dyn BatchSource>> {
87        if self.options.predicate.is_some() {
88            return Err(DataFrameError::streaming_unsupported(
89                "parquet_scan",
90                "predicate_batch_operator_not_installed",
91            ));
92        }
93
94        // `try_new` parses footer metadata. Reserve the complete configured bound first, then
95        // shrink to the metadata's actual retained size before any row-group reader is opened.
96        let mut footer_reservation = context
97            .budget
98            .reserve(ResourceScope::Source, context.options.memory_limit_bytes)?;
99        let file = File::open(&self.path)
100            .map_err(|source| DataFrameError::io_with_path(source, &self.path))?;
101        let builder = ParquetRecordBatchReaderBuilder::try_new(file)
102            .map_err(|source| DataFrameError::Parquet { source })?;
103        let metadata = builder.metadata().clone();
104        let reader_metadata = ArrowReaderMetadata::try_new(metadata.clone(), Default::default())
105            .map_err(|source| DataFrameError::Parquet { source })?;
106        let full_schema = reader_metadata.schema().clone();
107        let projection_indices = projection_indices(&full_schema, self.options.columns.as_deref())?;
108        let output_schema = projected_schema(&full_schema, projection_indices.as_deref());
109        drop(builder);
110
111        let row_groups =
112            selected_row_groups(&metadata, &self.options, projection_indices.as_deref())?;
113        // The initial reservation covers opaque footer parsing and all source-open allocations.
114        // Retain only their conservative footprint before decoded batches begin reserving bytes.
115        footer_reservation.shrink_to(
116            parquet_metadata_bytes(&metadata)
117                .saturating_add(schema_reservation_bytes(&full_schema))
118                .saturating_add(schema_reservation_bytes(&output_schema))
119                .saturating_add(
120                    u64::try_from(row_groups.len())
121                        .unwrap_or(u64::MAX)
122                        .saturating_mul(32),
123                )
124                .saturating_add(
125                    u64::try_from(projection_indices.as_ref().map_or(0, Vec::len))
126                        .unwrap_or(u64::MAX)
127                        .saturating_mul(8),
128                ),
129        );
130        Ok(Box::new(ParquetBatchSource {
131            context,
132            path: self.path.clone(),
133            reader_metadata,
134            options: self.options.clone(),
135            output_schema,
136            projection_indices,
137            row_groups,
138            current: None,
139            footer_reservation: Some(footer_reservation),
140        }))
141    }
142}
143
144struct ParquetBatchSource {
145    context: BatchOpenContext,
146    path: PathBuf,
147    reader_metadata: ArrowReaderMetadata,
148    options: ParquetReadOptions,
149    output_schema: SchemaRef,
150    projection_indices: Option<Vec<usize>>,
151    row_groups: VecDeque<RowGroupPlan>,
152    current: Option<CurrentRowGroup>,
153    footer_reservation: Option<ResourceReservation>,
154}
155
156struct RowGroupPlan {
157    index: usize,
158    decode_upper_bound: u64,
159}
160
161struct CurrentRowGroup {
162    reader: ParquetRecordBatchReader,
163    decode_upper_bound: u64,
164}
165
166impl BatchSource for ParquetBatchSource {
167    fn schema(&self) -> SchemaRef {
168        self.output_schema.clone()
169    }
170
171    fn next_batch(&mut self) -> Result<Option<StreamBatch>> {
172        loop {
173            if self.current.is_none() {
174                let Some(plan) = self.row_groups.pop_front() else {
175                    return Ok(None);
176                };
177                self.current = Some(self.open_row_group(plan)?);
178            }
179
180            let current = self.current.as_mut().expect("current row group is set");
181            let reservation = self
182                .context
183                .budget
184                .reserve_batch(ResourceScope::Decode, current.decode_upper_bound)?;
185            match current.reader.next() {
186                Some(Ok(batch)) => return Ok(Some(StreamBatch::new(batch, reservation))),
187                Some(Err(source)) => return Err(DataFrameError::Arrow { source }),
188                None => {
189                    drop(reservation);
190                    self.current = None;
191                }
192            }
193        }
194    }
195
196    fn close(&mut self) -> Result<()> {
197        self.current = None;
198        self.footer_reservation.take();
199        Ok(())
200    }
201}
202
203impl ParquetBatchSource {
204    fn open_row_group(&self, plan: RowGroupPlan) -> Result<CurrentRowGroup> {
205        let footer_bytes = self.context.budget.usage().reserved_bytes;
206        let observed = footer_bytes.saturating_add(plan.decode_upper_bound);
207        if observed > self.context.budget.memory_limit_bytes() {
208            return Err(DataFrameError::resource_limit_exceeded(
209                self.context.budget.memory_limit_bytes(),
210                observed,
211                self.context.budget.max_in_flight_batches(),
212                self.context
213                    .budget
214                    .usage()
215                    .reserved_batches
216                    .saturating_add(1),
217                ResourceScope::Decode,
218            ));
219        }
220
221        let file = File::open(&self.path)
222            .map_err(|source| DataFrameError::io_with_path(source, &self.path))?;
223        let mut builder =
224            ParquetRecordBatchReaderBuilder::new_with_metadata(file, self.reader_metadata.clone())
225                .with_batch_size(
226                    self.options
227                        .batch_size
228                        .min(self.context.options.batch_rows.get())
229                        .max(1),
230                )
231                .with_row_groups(vec![plan.index]);
232        if let Some(indices) = self.projection_indices.as_deref() {
233            let mask = ProjectionMask::roots(builder.parquet_schema(), indices.to_vec());
234            builder = builder.with_projection(mask);
235        }
236        let reader = builder
237            .build()
238            .map_err(|source| DataFrameError::Parquet { source })?;
239        Ok(CurrentRowGroup {
240            reader,
241            decode_upper_bound: plan.decode_upper_bound,
242        })
243    }
244}
245
246fn selected_row_groups(
247    metadata: &ParquetMetaData,
248    options: &ParquetReadOptions,
249    projection: Option<&[usize]>,
250) -> Result<VecDeque<RowGroupPlan>> {
251    let requested: Vec<usize> = match options.row_groups.as_deref() {
252        Some(indices) => indices.to_vec(),
253        None => (0..metadata.num_row_groups()).collect(),
254    };
255    requested
256        .into_iter()
257        .map(|index| {
258            let row_group = metadata.row_groups().get(index).ok_or_else(|| {
259                DataFrameError::configuration(
260                    "row_groups",
261                    format!("row group index {index} is outside the file metadata"),
262                )
263            })?;
264            let decode_upper_bound = selected_row_group_upper_bound(row_group, projection)?;
265            Ok(RowGroupPlan {
266                index,
267                decode_upper_bound,
268            })
269        })
270        .collect()
271}
272
273fn selected_row_group_upper_bound(
274    row_group: &parquet::file::metadata::RowGroupMetaData,
275    projection: Option<&[usize]>,
276) -> Result<u64> {
277    let selected = match projection {
278        Some(indices) => indices
279            .iter()
280            .map(|index| {
281                row_group.columns().get(*index).ok_or_else(|| {
282                    DataFrameError::schema_mismatch(
283                        "Parquet projection index exceeds row group schema",
284                    )
285                })
286            })
287            .collect::<Result<Vec<_>>>()?,
288        None => row_group.columns().iter().collect(),
289    };
290    Ok(selected.into_iter().fold(1024_u64, |bound, column| {
291        let compressed = u64::try_from(column.compressed_size()).unwrap_or(u64::MAX);
292        let uncompressed = u64::try_from(column.uncompressed_size()).unwrap_or(u64::MAX);
293        bound.saturating_add(compressed.max(uncompressed))
294    }))
295}
296
297fn parquet_metadata_bytes(metadata: &ParquetMetaData) -> u64 {
298    u64::try_from(metadata.memory_size()).unwrap_or(u64::MAX)
299}
300
301fn schema_reservation_bytes(schema: &SchemaRef) -> u64 {
302    schema.fields().iter().fold(256_u64, |bytes, field| {
303        bytes
304            .saturating_add(u64::try_from(field.name().len()).unwrap_or(u64::MAX))
305            .saturating_add(u64::try_from(field.data_type().to_string().len()).unwrap_or(u64::MAX))
306            .saturating_add(256)
307    })
308}
309
310fn projection_indices(
311    schema: &SchemaRef,
312    columns: Option<&[String]>,
313) -> Result<Option<Vec<usize>>> {
314    let Some(columns) = columns else {
315        return Ok(None);
316    };
317    columns
318        .iter()
319        .map(|name| {
320            schema
321                .fields()
322                .iter()
323                .position(|field| field.name() == name)
324                .ok_or_else(|| DataFrameError::column_not_found(name.clone()))
325        })
326        .collect::<Result<Vec<_>>>()
327        .map(Some)
328}
329
330fn projected_schema(schema: &SchemaRef, projection: Option<&[usize]>) -> SchemaRef {
331    let Some(projection) = projection else {
332        return schema.clone();
333    };
334    Arc::new(Schema::new(
335        projection
336            .iter()
337            .map(|index| schema.field(*index).clone())
338            .collect::<Vec<_>>(),
339    ))
340}
341
342#[cfg(test)]
343mod tests {
344    use std::num::NonZeroUsize;
345    use std::sync::Arc;
346
347    use arrow::array::{ArrayRef, Int64Array};
348    use arrow::datatypes::{DataType, Field, Schema};
349    use arrow::record_batch::RecordBatch;
350
351    use super::ParquetBatchSourceFactory;
352    use crate::io::{write_parquet, ParquetReadOptions};
353    use crate::physical::budget::StreamOptions;
354    use crate::physical::DataFrameStream;
355    use crate::{DataFrame, DataFrameError};
356
357    fn options(memory_limit_bytes: u64, batch_rows: usize) -> StreamOptions {
358        StreamOptions::new(
359            memory_limit_bytes,
360            NonZeroUsize::new(1).unwrap(),
361            NonZeroUsize::new(batch_rows).unwrap(),
362        )
363    }
364
365    fn parquet_fixture() -> (tempfile::TempDir, std::path::PathBuf) {
366        let dir = tempfile::tempdir().unwrap();
367        let path = dir.path().join("stream.parquet");
368        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)]));
369        let batch = RecordBatch::try_new(
370            schema,
371            vec![Arc::new(Int64Array::from(vec![1_i64, 2, 3, 4])) as ArrayRef],
372        )
373        .unwrap();
374        write_parquet(&path, &DataFrame::from_batches(vec![batch]).unwrap()).unwrap();
375        (dir, path)
376    }
377
378    #[test]
379    fn parquet_source_yields_incremental_stable_order_batches() {
380        let (_dir, path) = parquet_fixture();
381        let factory =
382            ParquetBatchSourceFactory::new(&path, ParquetReadOptions::default().with_batch_size(2));
383        let mut stream = DataFrameStream::from_factory(&factory, options(64 * 1024, 2)).unwrap();
384        let first = stream.next_batch().unwrap().unwrap();
385        let second = stream.next_batch().unwrap().unwrap();
386        assert_eq!(first.height(), 2);
387        assert_eq!(second.height(), 2);
388        let first_values = first.column("a").unwrap().to_arrow();
389        let second_values = second.column("a").unwrap().to_arrow();
390        assert_eq!(
391            first_values[0]
392                .as_any()
393                .downcast_ref::<Int64Array>()
394                .unwrap()
395                .value(0),
396            1
397        );
398        assert_eq!(
399            second_values[0]
400                .as_any()
401                .downcast_ref::<Int64Array>()
402                .unwrap()
403                .value(1),
404            4
405        );
406        assert!(stream.next_batch().unwrap().is_none());
407    }
408
409    #[test]
410    fn row_group_over_budget_fails_before_reader_build() {
411        let (_dir, path) = parquet_fixture();
412        let factory = ParquetBatchSourceFactory::new(&path, ParquetReadOptions::default());
413        let mut stream = DataFrameStream::from_factory(&factory, options(1024, 1)).unwrap();
414        assert!(matches!(
415            stream.next_batch(),
416            Err(DataFrameError::StreamFailed {
417                code: "resource_limit_exceeded",
418                ..
419            })
420        ));
421    }
422
423    #[test]
424    fn malformed_footer_fails_during_bounded_open() {
425        let dir = tempfile::tempdir().unwrap();
426        let path = dir.path().join("bad.parquet");
427        std::fs::write(&path, b"not a parquet file").unwrap();
428        let factory = ParquetBatchSourceFactory::new(&path, ParquetReadOptions::default());
429        assert!(matches!(
430            DataFrameStream::from_factory(&factory, options(64 * 1024, 1)),
431            Err(DataFrameError::Parquet { .. })
432        ));
433    }
434}
alopex_dataframe/io/streaming_parquet.rs

alopex_dataframe/io/
streaming_parquet.rs
