@@ -80,7 +80,7 @@ impl RecordBatch {
80
80
/// # }
81
81
/// ```
82
82
pub fn try_new ( schema : SchemaRef , columns : Vec < ArrayRef > ) -> Result < Self > {
83
- let options = RecordBatchOptions :: default ( ) ;
83
+ let options = RecordBatchOptions :: new ( ) ;
84
84
Self :: try_new_impl ( schema, columns, & options)
85
85
}
86
86
@@ -413,15 +413,29 @@ pub struct RecordBatchOptions {
413
413
pub row_count : Option < usize > ,
414
414
}
415
415
416
- impl Default for RecordBatchOptions {
417
- fn default ( ) -> Self {
416
+ impl RecordBatchOptions {
417
+ pub fn new ( ) -> Self {
418
418
Self {
419
419
match_field_names : true ,
420
420
row_count : None ,
421
421
}
422
422
}
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
+ }
423
438
}
424
-
425
439
impl From < & StructArray > for RecordBatch {
426
440
/// Create a record batch from struct array, where each field of
427
441
/// the `StructArray` becomes a `Field` in the schema.
@@ -901,10 +915,7 @@ mod tests {
901
915
. to_string( )
902
916
. contains( "must either specify a row count or at least one column" ) ) ;
903
917
904
- let options = RecordBatchOptions {
905
- row_count : Some ( 10 ) ,
906
- ..Default :: default ( )
907
- } ;
918
+ let options = RecordBatchOptions :: new ( ) . with_row_count ( Some ( 10 ) ) ;
908
919
909
920
let ok =
910
921
RecordBatch :: try_new_with_options ( schema. clone ( ) , vec ! [ ] , & options) . unwrap ( ) ;
@@ -929,4 +940,12 @@ mod tests {
929
940
) ;
930
941
assert_eq ! ( "Invalid argument error: Column 'a' is declared as non-nullable but contains null values" , format!( "{}" , maybe_batch. err( ) . unwrap( ) ) ) ;
931
942
}
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
+ }
932
951
}
0 commit comments