You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "alamb (via GitHub)" <gi...@apache.org> on 2023/01/27 11:44:20 UTC
[GitHub] [arrow-rs] alamb commented on a diff in pull request #3604: Add Push-Based CSV Decoder
alamb commented on code in PR #3604:
URL: https://github.com/apache/arrow-rs/pull/3604#discussion_r1088858660
##########
arrow-csv/src/reader/mod.rs:
##########
@@ -444,38 +433,146 @@ impl<R: Read> Reader<R> {
}
}
+impl<R: BufRead> BufReader<R> {
+ fn read(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
+ loop {
+ let buf = self.reader.fill_buf()?;
+ let decoded = self.decoder.decode(buf)?;
+ if decoded == 0 {
+ break;
+ }
+ self.reader.consume(decoded);
+ }
+
+ self.decoder.flush()
+ }
+}
+
impl<R: BufRead> Iterator for BufReader<R> {
type Item = Result<RecordBatch, ArrowError>;
fn next(&mut self) -> Option<Self::Item> {
+ self.read().transpose()
+ }
+}
+
+/// A push-based interface for decoding CSV data from an arbitrary byte stream
+///
+/// See [`Reader`] for a higher-level interface for interface with [`Read`]
+///
+/// The push-based interface facilitates integration with sources that yield arbitrarily
Review Comment:
👍
##########
arrow-csv/src/reader/mod.rs:
##########
@@ -444,38 +433,146 @@ impl<R: Read> Reader<R> {
}
}
+impl<R: BufRead> BufReader<R> {
+ fn read(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
+ loop {
+ let buf = self.reader.fill_buf()?;
+ let decoded = self.decoder.decode(buf)?;
+ if decoded == 0 {
+ break;
+ }
+ self.reader.consume(decoded);
+ }
+
+ self.decoder.flush()
+ }
+}
+
impl<R: BufRead> Iterator for BufReader<R> {
type Item = Result<RecordBatch, ArrowError>;
fn next(&mut self) -> Option<Self::Item> {
+ self.read().transpose()
+ }
+}
+
+/// A push-based interface for decoding CSV data from an arbitrary byte stream
+///
+/// See [`Reader`] for a higher-level interface for interface with [`Read`]
+///
+/// The push-based interface facilitates integration with sources that yield arbitrarily
+/// delimited bytes ranges, such as [`BufRead`], or a chunked byte stream received from
+/// object storage
+///
+/// ```
+/// # use std::io::BufRead;
+/// # use arrow_array::RecordBatch;
+/// # use arrow_csv::ReaderBuilder;
+/// # use arrow_schema::{ArrowError, SchemaRef};
+/// #
+/// fn read_from_csv<R: BufRead>(
+/// mut reader: R,
+/// schema: SchemaRef,
+/// batch_size: usize,
+/// ) -> Result<impl Iterator<Item = Result<RecordBatch, ArrowError>>, ArrowError> {
+/// let mut decoder = ReaderBuilder::new()
+/// .with_schema(schema)
+/// .with_batch_size(batch_size)
+/// .build_decoder();
+///
+/// let mut next = move || {
+/// loop {
+/// let buf = reader.fill_buf()?;
+/// let decoded = decoder.decode(buf)?;
+/// if decoded == 0 {
+/// break;
+/// }
+///
+/// // Consume the number of bytes read
+/// reader.consume(decoded);
+/// }
+/// decoder.flush()
+/// };
+/// Ok(std::iter::from_fn(move || next().transpose()))
+/// }
+/// ```
+#[derive(Debug)]
+pub struct Decoder {
+ /// Explicit schema for the CSV file
+ schema: SchemaRef,
+
+ /// Optional projection for which columns to load (zero-based column indices)
+ projection: Option<Vec<usize>>,
+
+ /// Number of records per batch
+ batch_size: usize,
+
+ /// Rows to skip
+ to_skip: usize,
+
+ /// Current line number
+ line_number: usize,
+
+ /// End line number
+ end: usize,
+
+ /// A decoder for [`StringRecords`]
+ record_decoder: RecordDecoder,
+
+ /// datetime format used to parse datetime values, (format understood by chrono)
+ ///
+ /// For format refer to [chrono docs](https://docs.rs/chrono/0.4.19/chrono/format/strftime/index.html)
+ datetime_format: Option<String>,
+}
+
+impl Decoder {
+ /// Decode records from `buf` returning the number of bytes read
+ ///
+ /// This method returns once `batch_size` objects have been parsed since the
+ /// last call to [`Self::flush`], or `buf` is exhausted. Any remaining bytes
+ /// should be included in the next call to [`Self::decode`]
+ ///
+ /// There is no requirement that `buf` contains a whole number of records, facilitating
Review Comment:
👍
##########
arrow-csv/src/reader/records.rs:
##########
@@ -17,120 +17,172 @@
use arrow_schema::ArrowError;
use csv_core::{ReadRecordResult, Reader};
-use std::io::BufRead;
/// The estimated length of a field in bytes
const AVERAGE_FIELD_SIZE: usize = 8;
/// The minimum amount of data in a single read
const MIN_CAPACITY: usize = 1024;
-pub struct RecordReader<R> {
- reader: R,
+/// [`RecordDecoder`] provides a push-based interface to decoder [`StringRecords`]
+#[derive(Debug)]
+pub struct RecordDecoder {
delimiter: Reader,
+ /// The expected number of fields per row
num_columns: usize,
+ /// The current line number
line_number: usize,
+
+ /// Offsets delimiting field start positions
offsets: Vec<usize>,
+
+ /// The current offset into `self.offsets`
+ ///
+ /// We track this independently of Vec to avoid re-zeroing memory
+ offsets_len: usize,
+
+ /// The number of fields read for the current record
+ current_field: usize,
+
+ /// The number of rows buffered
+ num_rows: usize,
+
+ /// Decoded field data
data: Vec<u8>,
+
+ /// Offsets into data
+ ///
+ /// We track this independently of Vec to avoid re-zeroing memory
+ data_len: usize,
}
-impl<R: BufRead> RecordReader<R> {
- pub fn new(reader: R, delimiter: Reader, num_columns: usize) -> Self {
+impl RecordDecoder {
+ pub fn new(delimiter: Reader, num_columns: usize) -> Self {
Self {
- reader,
delimiter,
num_columns,
line_number: 1,
offsets: vec![],
+ offsets_len: 1, // The first offset is always 0
+ current_field: 0,
+ data_len: 0,
data: vec![],
+ num_rows: 0,
}
}
- /// Clears and then fills the buffers on this [`RecordReader`]
- /// returning the number of records read
- fn fill_buf(&mut self, to_read: usize) -> Result<usize, ArrowError> {
- // Reserve sufficient capacity in offsets
- self.offsets.resize(to_read * self.num_columns + 1, 0);
-
- let mut read = 0;
+ /// Decodes records from `input` returning the number of records and bytes read
Review Comment:
```suggestion
/// Decodes records from `input` returning the number of records and bytes read and last offset consumed
```
Is that what the second `usize` is?
##########
arrow-csv/src/reader/records.rs:
##########
@@ -17,120 +17,172 @@
use arrow_schema::ArrowError;
use csv_core::{ReadRecordResult, Reader};
-use std::io::BufRead;
/// The estimated length of a field in bytes
const AVERAGE_FIELD_SIZE: usize = 8;
/// The minimum amount of data in a single read
const MIN_CAPACITY: usize = 1024;
-pub struct RecordReader<R> {
- reader: R,
+/// [`RecordDecoder`] provides a push-based interface to decoder [`StringRecords`]
+#[derive(Debug)]
+pub struct RecordDecoder {
delimiter: Reader,
+ /// The expected number of fields per row
num_columns: usize,
+ /// The current line number
line_number: usize,
+
+ /// Offsets delimiting field start positions
offsets: Vec<usize>,
+
+ /// The current offset into `self.offsets`
+ ///
+ /// We track this independently of Vec to avoid re-zeroing memory
+ offsets_len: usize,
+
+ /// The number of fields read for the current record
+ current_field: usize,
+
+ /// The number of rows buffered
+ num_rows: usize,
+
+ /// Decoded field data
data: Vec<u8>,
+
+ /// Offsets into data
+ ///
+ /// We track this independently of Vec to avoid re-zeroing memory
+ data_len: usize,
}
-impl<R: BufRead> RecordReader<R> {
- pub fn new(reader: R, delimiter: Reader, num_columns: usize) -> Self {
+impl RecordDecoder {
+ pub fn new(delimiter: Reader, num_columns: usize) -> Self {
Self {
- reader,
delimiter,
num_columns,
line_number: 1,
offsets: vec![],
+ offsets_len: 1, // The first offset is always 0
+ current_field: 0,
+ data_len: 0,
data: vec![],
+ num_rows: 0,
}
}
- /// Clears and then fills the buffers on this [`RecordReader`]
- /// returning the number of records read
- fn fill_buf(&mut self, to_read: usize) -> Result<usize, ArrowError> {
- // Reserve sufficient capacity in offsets
- self.offsets.resize(to_read * self.num_columns + 1, 0);
-
- let mut read = 0;
+ /// Decodes records from `input` returning the number of records and bytes read
+ ///
+ /// Note: this expects to be called with an empty `input` to signal EOF
+ pub fn decode(
+ &mut self,
+ input: &[u8],
+ to_read: usize,
+ ) -> Result<(usize, usize), ArrowError> {
if to_read == 0 {
- return Ok(0);
+ return Ok((0, 0));
}
- // The current offset into `self.data`
- let mut output_offset = 0;
+ // Reserve sufficient capacity in offsets
+ self.offsets
+ .resize(self.offsets_len + to_read * self.num_columns, 0);
+
// The current offset into `input`
let mut input_offset = 0;
- // The current offset into `self.offsets`
- let mut field_offset = 1;
- // The number of fields read for the current row
- let mut field_count = 0;
-
- 'outer: loop {
- let input = self.reader.fill_buf()?;
-
- 'input: loop {
- // Reserve necessary space in output data based on best estimate
- let remaining_rows = to_read - read;
- let capacity = remaining_rows * self.num_columns * AVERAGE_FIELD_SIZE;
- let estimated_data = capacity.max(MIN_CAPACITY);
- self.data.resize(output_offset + estimated_data, 0);
-
- loop {
- let (result, bytes_read, bytes_written, end_positions) =
- self.delimiter.read_record(
- &input[input_offset..],
- &mut self.data[output_offset..],
- &mut self.offsets[field_offset..],
- );
-
- field_count += end_positions;
- field_offset += end_positions;
- input_offset += bytes_read;
- output_offset += bytes_written;
-
- match result {
- ReadRecordResult::End => break 'outer, // Reached end of file
- ReadRecordResult::InputEmpty => break 'input, // Input exhausted, need to read more
- ReadRecordResult::OutputFull => break, // Need to allocate more capacity
- ReadRecordResult::OutputEndsFull => {
- let line_number = self.line_number + read;
- return Err(ArrowError::CsvError(format!("incorrect number of fields for line {}, expected {} got more than {}", line_number, self.num_columns, field_count)));
+
+ // The number of rows decoded in this pass
+ let mut read = 0;
+
+ loop {
+ // Reserve necessary space in output data based on best estimate
+ let remaining_rows = to_read - read;
+ let capacity = remaining_rows * self.num_columns * AVERAGE_FIELD_SIZE;
+ let estimated_data = capacity.max(MIN_CAPACITY);
+ self.data.resize(self.data_len + estimated_data, 0);
+
+ // Try to read a record
+ loop {
+ let (result, bytes_read, bytes_written, end_positions) =
+ self.delimiter.read_record(
+ &input[input_offset..],
+ &mut self.data[self.data_len..],
+ &mut self.offsets[self.offsets_len..],
+ );
+
+ self.current_field += end_positions;
+ self.offsets_len += end_positions;
+ input_offset += bytes_read;
+ self.data_len += bytes_written;
+
+ match result {
+ ReadRecordResult::End | ReadRecordResult::InputEmpty => {
+ // Reached end of input
+ return Ok((read, input_offset));
+ }
+ // Need to allocate more capacity
+ ReadRecordResult::OutputFull => break,
+ ReadRecordResult::OutputEndsFull => {
+ return Err(ArrowError::CsvError(format!("incorrect number of fields for line {}, expected {} got more than {}", self.line_number, self.num_columns, self.current_field)));
+ }
+ ReadRecordResult::Record => {
+ if self.current_field != self.num_columns {
+ return Err(ArrowError::CsvError(format!("incorrect number of fields for line {}, expected {} got {}", self.line_number, self.num_columns, self.current_field)));
}
- ReadRecordResult::Record => {
- if field_count != self.num_columns {
- let line_number = self.line_number + read;
- return Err(ArrowError::CsvError(format!("incorrect number of fields for line {}, expected {} got {}", line_number, self.num_columns, field_count)));
- }
- read += 1;
- field_count = 0;
-
- if read == to_read {
- break 'outer; // Read sufficient rows
- }
-
- if input.len() == input_offset {
- // Input exhausted, need to read more
- // Without this read_record will interpret the empty input
- // byte array as indicating the end of the file
- break 'input;
- }
+ read += 1;
+ self.current_field = 0;
+ self.line_number += 1;
+ self.num_rows += 1;
+
+ if read == to_read {
+ // Read sufficient rows
+ return Ok((read, input_offset));
+ }
+
+ if input.len() == input_offset {
+ // Input exhausted, need to read more
+ // Without this read_record will interpret the empty input
+ // byte array as indicating the end of the file
+ return Ok((read, input_offset));
}
}
}
}
- self.reader.consume(input_offset);
- input_offset = 0;
}
- self.reader.consume(input_offset);
+ }
+
+ /// Returns the current number of buffered records
+ pub fn len(&self) -> usize {
+ self.num_rows
+ }
+
+ /// Returns true if the decoder is empty
+ pub fn is_empty(&self) -> bool {
+ self.num_rows == 0
+ }
+
+ /// Clears the current contents of the decoder
+ pub fn clear(&mut self) {
+ // This does not reset current_field to allow clearing part way through a record
+ self.offsets_len = 1;
Review Comment:
I don't understand what the usecase for `clear` is here -- how would it clear part way through a record and then pick back up
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org