Add CSV Decoder::capacity (#3674) (#3677)

* Add CSV Decoder::capacity (#3674)

* Add test

* Remove unnecessary extern

* Add docs

arrow-csv/src/reader/mod.rs

@@ -438,10 +438,15 @@ impl<R: BufRead> BufReader<R> {
  loop {
  let buf = self.reader.fill_buf()?;
  let decoded = self.decoder.decode(buf)?;
- if decoded == 0 {
+ self.reader.consume(decoded);
+ // Yield if decoded no bytes or the decoder is full
+ //
+ // The capacity check avoids looping around and potentially
+ // blocking reading data in fill_buf that isn't needed
+ // to flush the next batch
+ if decoded == 0 || self.decoder.capacity() == 0 {
  break;
  }
- self.reader.consume(decoded);
  }
 
  self.decoder.flush()
@@ -574,6 +579,11 @@ impl Decoder {
  self.line_number += rows.len();
  Ok(Some(batch))
  }
+
+ /// Returns the number of records that can be read before requiring a call to [`Self::flush`]
+ pub fn capacity(&self) -> usize {
+ self.batch_size - self.record_decoder.len()
+ }
 }
 
 /// Parses a slice of [`StringRecords`] into a [RecordBatch]
@@ -2269,4 +2279,73 @@ mod tests {
  "Csv error: Encountered invalid UTF-8 data for line 1 and field 1",
  );
  }
+
+ struct InstrumentedRead<R> {
+ r: R,
+ fill_count: usize,
+ fill_sizes: Vec<usize>,
+ }
+
+ impl<R> InstrumentedRead<R> {
+ fn new(r: R) -> Self {
+ Self {
+ r,
+ fill_count: 0,
+ fill_sizes: vec![],
+ }
+ }
+ }
+
+ impl<R: Seek> Seek for InstrumentedRead<R> {
+ fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
+ self.r.seek(pos)
+ }
+ }
+
+ impl<R: BufRead> Read for InstrumentedRead<R> {
+ fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+ self.r.read(buf)
+ }
+ }
+
+ impl<R: BufRead> BufRead for InstrumentedRead<R> {
+ fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
+ self.fill_count += 1;
+ let buf = self.r.fill_buf()?;
+ self.fill_sizes.push(buf.len());
+ Ok(buf)
+ }
+
+ fn consume(&mut self, amt: usize) {
+ self.r.consume(amt)
+ }
+ }
+
+ #[test]
+ fn test_io() {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Utf8, false),
+ Field::new("b", DataType::Utf8, false),
+ ]));
+ let csv = "foo,bar\nbaz,foo\na,b\nc,d";
+ let mut read = InstrumentedRead::new(Cursor::new(csv.as_bytes()));
+ let reader = ReaderBuilder::new()
+ .with_schema(schema)
+ .with_batch_size(3)
+ .build_buffered(&mut read)
+ .unwrap();
+
+ let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
+ assert_eq!(batches.len(), 2);
+ assert_eq!(batches[0].num_rows(), 3);
+ assert_eq!(batches[1].num_rows(), 1);
+
+ // Expect 4 calls to fill_buf
+ // 1. Read first 3 rows
+ // 2. Read final row
+ // 3. Delimit and flush final row
+ // 4. Iterator finished
+ assert_eq!(&read.fill_sizes, &[23, 3, 0, 0]);
+ assert_eq!(read.fill_count, 4);
+ }
 }