alopex_dataframe/io/
streaming_csv.rs

1//! Bounded, quote-aware CSV source for `DataFrameStream`.
2//!
3//! This reader deliberately does not call `read_csv_with_options`: it frames one complete CSV
4//! record group, reserves conservative decode ownership, and only then invokes Arrow's CSV
5//! decoder for that group.
6
7use std::fs::File;
8use std::io::{BufReader, Cursor, Read};
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11
12use arrow::datatypes::{Schema, SchemaRef};
13use arrow::record_batch::RecordBatch;
14use arrow_csv::reader::{Format, ReaderBuilder};
15use regex::Regex;
16
17use crate::io::options::CsvReadOptions;
18use crate::physical::budget::{ResourceReservation, ResourceScope};
19use crate::physical::{
20    BatchOpenContext, BatchSource, BatchSourceFactory, SourceLimit, StreamBatch,
21};
22use crate::{DataFrameError, Result};
23
24/// Re-consumable bounded CSV source factory.
25#[derive(Debug, Clone)]
26pub struct CsvBatchSourceFactory {
27    path: PathBuf,
28    options: CsvReadOptions,
29}
30
31impl CsvBatchSourceFactory {
32    /// Construct a factory for a CSV path and read options.
33    pub fn new(path: impl AsRef<Path>, options: CsvReadOptions) -> Self {
34        Self {
35            path: path.as_ref().to_path_buf(),
36            options,
37        }
38    }
39
40    /// Construct a factory from a physical CSV scan. Predicate execution is rejected by the
41    /// streaming compiler until its budgeted batch operator is installed.
42    pub(crate) fn from_scan(
43        path: PathBuf,
44        predicate: Option<crate::Expr>,
45        projection: Option<Vec<String>>,
46    ) -> Self {
47        let options = CsvReadOptions {
48            predicate,
49            projection,
50            ..CsvReadOptions::default()
51        };
52        Self { path, options }
53    }
54}
55
56impl BatchSourceFactory for CsvBatchSourceFactory {
57    fn source_name(&self) -> &'static str {
58        "csv"
59    }
60
61    fn schema(&self) -> Result<SchemaRef> {
62        Err(DataFrameError::streaming_unsupported(
63            "csv_scan",
64            "schema_is_available_after_bounded_open",
65        ))
66    }
67
68    fn source_limits(&self) -> Vec<SourceLimit> {
69        vec![SourceLimit {
70            source: crate::physical::PlanSubject::CsvScan,
71            code: "stable_input_order",
72            description: "CSV streaming preserves physical input record order",
73        }]
74    }
75
76    fn open(&self, context: BatchOpenContext) -> Result<Box<dyn BatchSource>> {
77        validate_options(&self.options)?;
78        if self.options.predicate.is_some() {
79            return Err(DataFrameError::streaming_unsupported(
80                "csv_scan",
81                "predicate_batch_operator_not_installed",
82            ));
83        }
84
85        let block_bytes = bounded_block_bytes(context.options.memory_limit_bytes)?;
86        let (full_schema, schema_reservation) =
87            infer_schema_bounded(&self.path, &self.options, &context, block_bytes)?;
88        let projection_indices =
89            projection_indices(&full_schema, self.options.projection.as_deref())?;
90        let output_schema = projected_schema(&full_schema, projection_indices.as_deref());
91
92        let file = File::open(&self.path)
93            .map_err(|source| DataFrameError::io_with_path(source, &self.path))?;
94        let mut framer = BoundedCsvFramer::new(
95            file,
96            self.path.clone(),
97            self.options.quote_char,
98            block_bytes,
99        );
100        if self.options.has_header {
101            framer.skip_record(&context)?;
102        }
103
104        Ok(Box::new(CsvBatchSource {
105            context,
106            framer,
107            full_schema,
108            output_schema,
109            projection_indices,
110            options: self.options.clone(),
111            schema_reservation: Some(schema_reservation),
112        }))
113    }
114}
115
116struct CsvBatchSource {
117    context: BatchOpenContext,
118    framer: BoundedCsvFramer<File>,
119    full_schema: SchemaRef,
120    output_schema: SchemaRef,
121    projection_indices: Option<Vec<usize>>,
122    options: CsvReadOptions,
123    // Holds the bounded schema-inference allocation for the source lifetime.
124    schema_reservation: Option<ResourceReservation>,
125}
126
127impl BatchSource for CsvBatchSource {
128    fn schema(&self) -> SchemaRef {
129        self.output_schema.clone()
130    }
131
132    fn next_batch(&mut self) -> Result<Option<StreamBatch>> {
133        let Some(frame) = self
134            .framer
135            .next_frame(self.context.options.batch_rows.get(), &self.context)?
136        else {
137            return Ok(None);
138        };
139
140        let decode_bytes = conservative_decode_upper_bound(
141            frame.bytes.len(),
142            frame.rows,
143            self.full_schema.fields().len(),
144        );
145        let reservation = self
146            .context
147            .budget
148            .reserve_batch(ResourceScope::Decode, decode_bytes)?;
149
150        let batch = decode_frame(
151            &frame.bytes,
152            frame.rows,
153            self.full_schema.clone(),
154            &self.options,
155        )?;
156        let batch = project_batch(batch, self.projection_indices.as_deref())?;
157        drop(frame);
158
159        Ok(Some(StreamBatch::new(batch, reservation)))
160    }
161
162    fn close(&mut self) -> Result<()> {
163        self.schema_reservation.take();
164        Ok(())
165    }
166}
167
168fn infer_schema_bounded(
169    path: &Path,
170    options: &CsvReadOptions,
171    context: &BatchOpenContext,
172    block_bytes: usize,
173) -> Result<(SchemaRef, ResourceReservation)> {
174    let file = File::open(path).map_err(|source| DataFrameError::io_with_path(source, path))?;
175    let mut framer =
176        BoundedCsvFramer::new(file, path.to_path_buf(), options.quote_char, block_bytes);
177    let record_count = options
178        .infer_schema_length
179        .saturating_add(usize::from(options.has_header))
180        .max(1);
181    let Some(frame) = framer.next_frame(record_count, context)? else {
182        return Err(DataFrameError::schema_mismatch("CSV input is empty"));
183    };
184    let schema_reservation = context.budget.reserve(
185        ResourceScope::Source,
186        schema_reservation_bytes(frame.bytes.len()),
187    )?;
188    let mut reader = BufReader::with_capacity(
189        frame.bytes.len().max(1),
190        Cursor::new(frame.bytes.as_slice()),
191    );
192    let (schema, _) = csv_format(options, options.has_header)?
193        .infer_schema(&mut reader, Some(options.infer_schema_length))
194        .map_err(|source| DataFrameError::Arrow { source })?;
195    drop(frame);
196    Ok((Arc::new(schema), schema_reservation))
197}
198
199fn decode_frame(
200    bytes: &[u8],
201    rows: usize,
202    schema: SchemaRef,
203    options: &CsvReadOptions,
204) -> Result<RecordBatch> {
205    let reader = ReaderBuilder::new(schema)
206        .with_format(csv_format(options, false)?)
207        .with_batch_size(rows.max(1))
208        .build(Cursor::new(bytes))
209        .map_err(|source| DataFrameError::Arrow { source })?;
210    let mut batches = reader;
211    let batch = batches
212        .next()
213        .transpose()
214        .map_err(|source| DataFrameError::Arrow { source })?
215        .ok_or_else(|| DataFrameError::schema_mismatch("CSV frame contained no records"))?;
216    if batches.next().is_some() {
217        return Err(DataFrameError::schema_mismatch(
218            "CSV frame decoded to more than one bounded batch",
219        ));
220    }
221    Ok(batch)
222}
223
224fn csv_format(options: &CsvReadOptions, has_header: bool) -> Result<Format> {
225    let mut format = Format::default()
226        .with_header(has_header)
227        .with_delimiter(options.delimiter);
228    if let Some(quote_char) = options.quote_char {
229        format = format.with_quote(quote_char);
230    }
231    if !options.null_values.is_empty() {
232        let pattern = options
233            .null_values
234            .iter()
235            .map(|value| regex::escape(value))
236            .collect::<Vec<_>>()
237            .join("|");
238        let regex = Regex::new(&format!("^(?:{pattern})$")).map_err(|error| {
239            DataFrameError::configuration("null_values", format!("invalid regex: {error}"))
240        })?;
241        format = format.with_null_regex(regex);
242    }
243    Ok(format)
244}
245
246fn validate_options(options: &CsvReadOptions) -> Result<()> {
247    if options.delimiter == b'\0' {
248        return Err(DataFrameError::configuration(
249            "delimiter",
250            "delimiter must not be NUL (0x00)",
251        ));
252    }
253    if options.quote_char == Some(b'\0') {
254        return Err(DataFrameError::configuration(
255            "quote_char",
256            "quote_char must not be NUL (0x00)",
257        ));
258    }
259    Ok(())
260}
261
262fn projection_indices(
263    schema: &SchemaRef,
264    projection: Option<&[String]>,
265) -> Result<Option<Vec<usize>>> {
266    let Some(projection) = projection else {
267        return Ok(None);
268    };
269    projection
270        .iter()
271        .map(|name| {
272            schema
273                .fields()
274                .iter()
275                .position(|field| field.name() == name)
276                .ok_or_else(|| DataFrameError::column_not_found(name.clone()))
277        })
278        .collect::<Result<Vec<_>>>()
279        .map(Some)
280}
281
282fn projected_schema(schema: &SchemaRef, projection: Option<&[usize]>) -> SchemaRef {
283    let Some(projection) = projection else {
284        return schema.clone();
285    };
286    Arc::new(Schema::new(
287        projection
288            .iter()
289            .map(|index| schema.field(*index).clone())
290            .collect::<Vec<_>>(),
291    ))
292}
293
294fn project_batch(batch: RecordBatch, projection: Option<&[usize]>) -> Result<RecordBatch> {
295    match projection {
296        Some(indices) => batch.project(indices).map_err(|error| {
297            DataFrameError::schema_mismatch(format!("failed to project CSV batch: {error}"))
298        }),
299        None => Ok(batch),
300    }
301}
302
303fn bounded_block_bytes(memory_limit_bytes: u64) -> Result<usize> {
304    let bytes = (memory_limit_bytes / 32).clamp(1, 64 * 1024);
305    usize::try_from(bytes).map_err(|_| {
306        DataFrameError::configuration("memory_limit_bytes", "does not fit this platform")
307    })
308}
309
310fn conservative_decode_upper_bound(raw_bytes: usize, rows: usize, columns: usize) -> u64 {
311    let raw = u64::try_from(raw_bytes).unwrap_or(u64::MAX);
312    let cells = u64::try_from(rows)
313        .unwrap_or(u64::MAX)
314        .saturating_mul(u64::try_from(columns).unwrap_or(u64::MAX));
315    raw.saturating_mul(16)
316        .saturating_add(cells.saturating_mul(16))
317        .saturating_add(1024)
318}
319
320fn schema_reservation_bytes(raw_bytes: usize) -> u64 {
321    u64::try_from(raw_bytes)
322        .unwrap_or(u64::MAX)
323        .saturating_mul(2)
324        .saturating_add(128)
325}
326
327struct FramedBytes {
328    bytes: Vec<u8>,
329    rows: usize,
330    _reservation: ResourceReservation,
331}
332
333/// Fixed-capacity framing reader that understands escaped quotes and multiline quoted records.
334struct BoundedCsvFramer<R> {
335    reader: R,
336    path: PathBuf,
337    quote_char: Option<u8>,
338    byte_cap: usize,
339    eof: bool,
340}
341
342impl<R: Read> BoundedCsvFramer<R> {
343    fn new(reader: R, path: PathBuf, quote_char: Option<u8>, byte_cap: usize) -> Self {
344        Self {
345            reader,
346            path,
347            quote_char,
348            byte_cap,
349            eof: false,
350        }
351    }
352
353    fn next_frame(
354        &mut self,
355        max_rows: usize,
356        context: &BatchOpenContext,
357    ) -> Result<Option<FramedBytes>> {
358        if self.eof {
359            return Ok(None);
360        }
361        let reservation = context
362            .budget
363            .reserve(ResourceScope::Source, self.byte_cap as u64)?;
364        let mut bytes = Vec::with_capacity(self.byte_cap);
365        let mut tracker = QuoteTracker::new(self.quote_char);
366        let mut rows: usize = 0;
367
368        loop {
369            let Some(byte) = self.read_byte()? else {
370                self.eof = true;
371                if bytes.is_empty() {
372                    return Ok(None);
373                }
374                tracker.finish()?;
375                if bytes.last() != Some(&b'\n') {
376                    rows = rows.saturating_add(1);
377                }
378                break;
379            };
380
381            if bytes.len() == self.byte_cap {
382                self.discard_record_tail(tracker, byte)?;
383                return Err(self.record_limit_error(context));
384            }
385            bytes.push(byte);
386            if tracker.observe(byte) {
387                rows = rows.saturating_add(1);
388                if rows == max_rows {
389                    break;
390                }
391            }
392        }
393
394        Ok(Some(FramedBytes {
395            bytes,
396            rows,
397            _reservation: reservation,
398        }))
399    }
400
401    fn skip_record(&mut self, context: &BatchOpenContext) -> Result<()> {
402        let mut tracker = QuoteTracker::new(self.quote_char);
403        let mut bytes = 0usize;
404        loop {
405            let Some(byte) = self.read_byte()? else {
406                self.eof = true;
407                tracker.finish()?;
408                return Ok(());
409            };
410            bytes = bytes.saturating_add(1);
411            if bytes > self.byte_cap {
412                self.discard_record_tail(tracker, byte)?;
413                return Err(self.record_limit_error(context));
414            }
415            if tracker.observe(byte) {
416                return Ok(());
417            }
418        }
419    }
420
421    fn read_byte(&mut self) -> Result<Option<u8>> {
422        let mut byte = [0_u8; 1];
423        match self.reader.read(&mut byte) {
424            Ok(0) => Ok(None),
425            Ok(_) => Ok(Some(byte[0])),
426            Err(source) => Err(DataFrameError::io_with_path(source, &self.path)),
427        }
428    }
429
430    fn discard_record_tail(&mut self, mut tracker: QuoteTracker, first: u8) -> Result<()> {
431        if tracker.observe(first) {
432            return Ok(());
433        }
434        while let Some(byte) = self.read_byte()? {
435            if tracker.observe(byte) {
436                return Ok(());
437            }
438        }
439        self.eof = true;
440        tracker.finish()
441    }
442
443    fn record_limit_error(&self, context: &BatchOpenContext) -> DataFrameError {
444        DataFrameError::resource_limit_exceeded(
445            context.budget.memory_limit_bytes(),
446            context
447                .budget
448                .usage()
449                .reserved_bytes
450                .saturating_add(self.byte_cap as u64)
451                .saturating_add(1),
452            context.budget.max_in_flight_batches(),
453            context.budget.usage().reserved_batches,
454            ResourceScope::Source,
455        )
456    }
457}
458
459#[derive(Clone, Copy)]
460struct QuoteTracker {
461    quote: Option<u8>,
462    in_quotes: bool,
463    quote_pending: bool,
464}
465
466impl QuoteTracker {
467    fn new(quote: Option<u8>) -> Self {
468        Self {
469            quote,
470            in_quotes: false,
471            quote_pending: false,
472        }
473    }
474
475    fn observe(&mut self, byte: u8) -> bool {
476        let Some(quote) = self.quote else {
477            return byte == b'\n';
478        };
479
480        if self.quote_pending {
481            if byte == quote {
482                self.quote_pending = false;
483                return false;
484            }
485            self.quote_pending = false;
486            self.in_quotes = false;
487        }
488        if self.in_quotes {
489            if byte == quote {
490                self.quote_pending = true;
491            }
492            return false;
493        }
494        if byte == quote {
495            self.in_quotes = true;
496            false
497        } else {
498            byte == b'\n'
499        }
500    }
501
502    fn finish(&mut self) -> Result<()> {
503        if self.quote_pending {
504            self.quote_pending = false;
505            self.in_quotes = false;
506        }
507        if self.in_quotes {
508            return Err(DataFrameError::schema_mismatch(
509                "CSV input ended inside a quoted record",
510            ));
511        }
512        Ok(())
513    }
514}
515
516#[cfg(test)]
517mod tests {
518    use std::num::NonZeroUsize;
519
520    use arrow::array::Int64Array;
521
522    use super::CsvBatchSourceFactory;
523    use crate::io::CsvReadOptions;
524    use crate::physical::budget::StreamOptions;
525    use crate::physical::DataFrameStream;
526    use crate::DataFrameError;
527
528    fn options(memory_limit_bytes: u64, batch_rows: usize) -> StreamOptions {
529        StreamOptions::new(
530            memory_limit_bytes,
531            NonZeroUsize::new(1).unwrap(),
532            NonZeroUsize::new(batch_rows).unwrap(),
533        )
534    }
535
536    #[test]
537    fn quoted_newline_is_one_record_and_input_order_is_stable() {
538        let dir = tempfile::tempdir().unwrap();
539        let path = dir.path().join("quoted.csv");
540        std::fs::write(&path, "a,b\n1,\"first\nline\"\n2,second\n").unwrap();
541
542        let factory = CsvBatchSourceFactory::new(
543            &path,
544            CsvReadOptions::default().with_infer_schema_length(1),
545        );
546        let mut stream = DataFrameStream::from_factory(&factory, options(16 * 1024, 1)).unwrap();
547        let first = stream.next_batch().unwrap().unwrap();
548        let second = stream.next_batch().unwrap().unwrap();
549
550        assert_eq!(first.height(), 1);
551        assert_eq!(second.height(), 1);
552        let first_values = first.column("a").unwrap().to_arrow();
553        let second_values = second.column("a").unwrap().to_arrow();
554        assert_eq!(
555            first_values[0]
556                .as_any()
557                .downcast_ref::<Int64Array>()
558                .unwrap()
559                .value(0),
560            1
561        );
562        assert_eq!(
563            second_values[0]
564                .as_any()
565                .downcast_ref::<Int64Array>()
566                .unwrap()
567                .value(0),
568            2
569        );
570    }
571
572    #[test]
573    fn oversized_record_fails_without_accumulating_the_record() {
574        let dir = tempfile::tempdir().unwrap();
575        let path = dir.path().join("oversized.csv");
576        std::fs::write(&path, format!("a\n{}\n", "x".repeat(256))).unwrap();
577
578        let factory = CsvBatchSourceFactory::new(&path, CsvReadOptions::default());
579        let err = match DataFrameStream::from_factory(&factory, options(512, 1)) {
580            Ok(_) => panic!("oversized CSV record must not open a stream"),
581            Err(error) => error,
582        };
583        assert!(matches!(err, DataFrameError::ResourceLimitExceeded { .. }));
584    }
585
586    #[test]
587    fn schema_error_is_terminal_and_does_not_yield_another_batch() {
588        let dir = tempfile::tempdir().unwrap();
589        let path = dir.path().join("invalid.csv");
590        std::fs::write(&path, "a,b\n1,2\n3,4,5\n").unwrap();
591
592        let factory = CsvBatchSourceFactory::new(
593            &path,
594            CsvReadOptions::default().with_infer_schema_length(1),
595        );
596        let mut stream = DataFrameStream::from_factory(&factory, options(16 * 1024, 1)).unwrap();
597        assert!(stream.next_batch().unwrap().is_some());
598        assert!(matches!(
599            stream.next_batch(),
600            Err(DataFrameError::StreamFailed { .. })
601        ));
602        assert!(matches!(
603            stream.next_batch(),
604            Err(DataFrameError::StreamFailed { .. })
605        ));
606    }
607
608    #[test]
609    fn invalid_utf8_is_a_repeatable_decode_terminal() {
610        let dir = tempfile::tempdir().unwrap();
611        let path = dir.path().join("invalid-utf8.csv");
612        std::fs::write(&path, b"a\nok\n\xff\n").unwrap();
613
614        let factory = CsvBatchSourceFactory::new(
615            &path,
616            CsvReadOptions::default().with_infer_schema_length(1),
617        );
618        let mut stream = DataFrameStream::from_factory(&factory, options(16 * 1024, 1)).unwrap();
619        assert!(stream.next_batch().unwrap().is_some());
620        assert!(matches!(
621            stream.next_batch(),
622            Err(DataFrameError::StreamFailed {
623                code: "decode_error",
624                ..
625            })
626        ));
627        assert!(matches!(
628            stream.next_batch(),
629            Err(DataFrameError::StreamFailed {
630                code: "decode_error",
631                ..
632            })
633        ));
634    }
635}
alopex_dataframe/io/streaming_csv.rs

alopex_dataframe/io/
streaming_csv.rs
