Skip to content

Commit

Permalink
Better construction of RecordBatchOptions (#2729)
Browse files Browse the repository at this point in the history
* include builder for RecordBatchOptions

* fix clippy warnings

* fix clippy warnings

* remove builder struct

* removed a wrong comment

* Update comment  in arrow/src/record_batch.rs

Co-authored-by: Andrew Lamb <[email protected]>

* Update comment in arrow/src/record_batch.rs

Co-authored-by: Andrew Lamb <[email protected]>

Co-authored-by: askoa <askoa@local>
Co-authored-by: Andrew Lamb <[email protected]>
  • Loading branch information
3 people authored Sep 16, 2022
1 parent 0ebd71e commit 43d912c
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 16 deletions.
12 changes: 4 additions & 8 deletions arrow/src/ipc/reader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -578,10 +578,7 @@ pub fn read_record_batch(
let mut node_index = 0;
let mut arrays = vec![];

let options = RecordBatchOptions {
row_count: Some(batch.length() as usize),
..Default::default()
};
let options = RecordBatchOptions::new().with_row_count(Some(batch.length() as usize));

if let Some(projection) = projection {
// project fields
Expand Down Expand Up @@ -1692,10 +1689,9 @@ mod tests {
#[test]
fn test_no_columns_batch() {
let schema = Arc::new(Schema::new(vec![]));
let options = RecordBatchOptions {
match_field_names: true,
row_count: Some(10),
};
let options = RecordBatchOptions::new()
.with_match_field_names(true)
.with_row_count(Some(10));
let input_batch =
RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
let output_batch = roundtrip_ipc_stream(&input_batch);
Expand Down
35 changes: 27 additions & 8 deletions arrow/src/record_batch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ impl RecordBatch {
/// # }
/// ```
pub fn try_new(schema: SchemaRef, columns: Vec<ArrayRef>) -> Result<Self> {
let options = RecordBatchOptions::default();
let options = RecordBatchOptions::new();
Self::try_new_impl(schema, columns, &options)
}

Expand Down Expand Up @@ -413,15 +413,29 @@ pub struct RecordBatchOptions {
pub row_count: Option<usize>,
}

impl Default for RecordBatchOptions {
fn default() -> Self {
impl RecordBatchOptions {
pub fn new() -> Self {
Self {
match_field_names: true,
row_count: None,
}
}
/// Sets the row_count of RecordBatchOptions and returns self
pub fn with_row_count(mut self, row_count: Option<usize>) -> Self {
self.row_count = row_count;
self
}
/// Sets the match_field_names of RecordBatchOptions and returns self
pub fn with_match_field_names(mut self, match_field_names: bool) -> Self {
self.match_field_names = match_field_names;
self
}
}
impl Default for RecordBatchOptions {
fn default() -> Self {
Self::new()
}
}

impl From<&StructArray> for RecordBatch {
/// Create a record batch from struct array, where each field of
/// the `StructArray` becomes a `Field` in the schema.
Expand Down Expand Up @@ -901,10 +915,7 @@ mod tests {
.to_string()
.contains("must either specify a row count or at least one column"));

let options = RecordBatchOptions {
row_count: Some(10),
..Default::default()
};
let options = RecordBatchOptions::new().with_row_count(Some(10));

let ok =
RecordBatch::try_new_with_options(schema.clone(), vec![], &options).unwrap();
Expand All @@ -929,4 +940,12 @@ mod tests {
);
assert_eq!("Invalid argument error: Column 'a' is declared as non-nullable but contains null values", format!("{}", maybe_batch.err().unwrap()));
}
#[test]
fn test_record_batch_options() {
let options = RecordBatchOptions::new()
.with_match_field_names(false)
.with_row_count(Some(20));
assert!(!options.match_field_names);
assert_eq!(options.row_count.unwrap(), 20)
}
}

0 comments on commit 43d912c

Please sign in to comment.