Skip to content

Commit 43d912c

Browse files
askoaaskoaalamb
authored
Better construction of RecordBatchOptions (#2729)
* 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]>
1 parent 0ebd71e commit 43d912c

File tree

2 files changed

+31
-16
lines changed

2 files changed

+31
-16
lines changed

arrow/src/ipc/reader.rs

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -578,10 +578,7 @@ pub fn read_record_batch(
578578
let mut node_index = 0;
579579
let mut arrays = vec![];
580580

581-
let options = RecordBatchOptions {
582-
row_count: Some(batch.length() as usize),
583-
..Default::default()
584-
};
581+
let options = RecordBatchOptions::new().with_row_count(Some(batch.length() as usize));
585582

586583
if let Some(projection) = projection {
587584
// project fields
@@ -1692,10 +1689,9 @@ mod tests {
16921689
#[test]
16931690
fn test_no_columns_batch() {
16941691
let schema = Arc::new(Schema::new(vec![]));
1695-
let options = RecordBatchOptions {
1696-
match_field_names: true,
1697-
row_count: Some(10),
1698-
};
1692+
let options = RecordBatchOptions::new()
1693+
.with_match_field_names(true)
1694+
.with_row_count(Some(10));
16991695
let input_batch =
17001696
RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
17011697
let output_batch = roundtrip_ipc_stream(&input_batch);

arrow/src/record_batch.rs

Lines changed: 27 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ impl RecordBatch {
8080
/// # }
8181
/// ```
8282
pub fn try_new(schema: SchemaRef, columns: Vec<ArrayRef>) -> Result<Self> {
83-
let options = RecordBatchOptions::default();
83+
let options = RecordBatchOptions::new();
8484
Self::try_new_impl(schema, columns, &options)
8585
}
8686

@@ -413,15 +413,29 @@ pub struct RecordBatchOptions {
413413
pub row_count: Option<usize>,
414414
}
415415

416-
impl Default for RecordBatchOptions {
417-
fn default() -> Self {
416+
impl RecordBatchOptions {
417+
pub fn new() -> Self {
418418
Self {
419419
match_field_names: true,
420420
row_count: None,
421421
}
422422
}
423+
/// Sets the row_count of RecordBatchOptions and returns self
424+
pub fn with_row_count(mut self, row_count: Option<usize>) -> Self {
425+
self.row_count = row_count;
426+
self
427+
}
428+
/// Sets the match_field_names of RecordBatchOptions and returns self
429+
pub fn with_match_field_names(mut self, match_field_names: bool) -> Self {
430+
self.match_field_names = match_field_names;
431+
self
432+
}
433+
}
434+
impl Default for RecordBatchOptions {
435+
fn default() -> Self {
436+
Self::new()
437+
}
423438
}
424-
425439
impl From<&StructArray> for RecordBatch {
426440
/// Create a record batch from struct array, where each field of
427441
/// the `StructArray` becomes a `Field` in the schema.
@@ -901,10 +915,7 @@ mod tests {
901915
.to_string()
902916
.contains("must either specify a row count or at least one column"));
903917

904-
let options = RecordBatchOptions {
905-
row_count: Some(10),
906-
..Default::default()
907-
};
918+
let options = RecordBatchOptions::new().with_row_count(Some(10));
908919

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

0 commit comments

Comments
 (0)