@@ -35,16 +35,18 @@ use crate::{Error, ErrorKind, Result};
35
35
#[ derive( Clone ) ]
36
36
pub struct EqualityDeleteFileWriterBuilder < B : FileWriterBuilder > {
37
37
inner : B ,
38
+ config : EqualityDeleteWriterConfig ,
38
39
}
39
40
40
41
impl < B : FileWriterBuilder > EqualityDeleteFileWriterBuilder < B > {
41
42
/// Create a new `EqualityDeleteFileWriterBuilder` using a `FileWriterBuilder`.
42
- pub fn new ( inner : B ) -> Self {
43
- Self { inner }
43
+ pub fn new ( inner : B , config : EqualityDeleteWriterConfig ) -> Self {
44
+ Self { inner, config }
44
45
}
45
46
}
46
47
47
48
/// Config for `EqualityDeleteWriter`.
49
+ #[ derive( Clone ) ]
48
50
pub struct EqualityDeleteWriterConfig {
49
51
// Field ids used to determine row equality in equality delete files.
50
52
equality_ids : Vec < i32 > ,
@@ -108,14 +110,13 @@ impl EqualityDeleteWriterConfig {
108
110
#[ async_trait:: async_trait]
109
111
impl < B : FileWriterBuilder > IcebergWriterBuilder for EqualityDeleteFileWriterBuilder < B > {
110
112
type R = EqualityDeleteFileWriter < B > ;
111
- type C = EqualityDeleteWriterConfig ;
112
113
113
- async fn build ( self , config : Self :: C ) -> Result < Self :: R > {
114
+ async fn build ( self ) -> Result < Self :: R > {
114
115
Ok ( EqualityDeleteFileWriter {
115
116
inner_writer : Some ( self . inner . clone ( ) . build ( ) . await ?) ,
116
- projector : config. projector ,
117
- equality_ids : config. equality_ids ,
118
- partition_value : config. partition_value ,
117
+ projector : self . config . projector ,
118
+ equality_ids : self . config . equality_ids ,
119
+ partition_value : self . config . partition_value ,
119
120
} )
120
121
}
121
122
}
@@ -386,20 +387,19 @@ mod test {
386
387
EqualityDeleteWriterConfig :: new ( equality_ids, Arc :: new ( schema) , None ) . unwrap ( ) ;
387
388
let delete_schema =
388
389
arrow_schema_to_schema ( equality_config. projected_arrow_schema_ref ( ) ) . unwrap ( ) ;
389
- let projector = equality_config. projector . clone ( ) ;
390
390
391
391
// prepare writer
392
392
let pb = ParquetWriterBuilder :: new (
393
393
WriterProperties :: builder ( ) . build ( ) ,
394
- Arc :: new ( delete_schema) ,
395
394
file_io. clone ( ) ,
396
395
location_gen,
397
396
file_name_gen,
397
+ Arc :: new ( delete_schema) ,
398
398
) ;
399
-
400
- let mut equality_delete_writer = EqualityDeleteFileWriterBuilder :: new ( pb)
401
- . build ( equality_config)
399
+ let mut equality_delete_writer = EqualityDeleteFileWriterBuilder :: new ( pb, equality_config)
400
+ . build ( )
402
401
. await ?;
402
+ let projector = equality_delete_writer. projector . clone ( ) ;
403
403
404
404
// write
405
405
equality_delete_writer. write ( to_write. clone ( ) ) . await ?;
@@ -556,13 +556,13 @@ mod test {
556
556
557
557
let pb = ParquetWriterBuilder :: new (
558
558
WriterProperties :: builder ( ) . build ( ) ,
559
- Arc :: new ( delete_schema) ,
560
559
file_io. clone ( ) ,
561
560
location_gen,
562
561
file_name_gen,
562
+ Arc :: new ( delete_schema) ,
563
563
) ;
564
- let mut equality_delete_writer = EqualityDeleteFileWriterBuilder :: new ( pb)
565
- . build ( config )
564
+ let mut equality_delete_writer = EqualityDeleteFileWriterBuilder :: new ( pb, config )
565
+ . build ( )
566
566
. await ?;
567
567
568
568
// prepare data
@@ -640,10 +640,12 @@ mod test {
640
640
Some ( b"" ) ,
641
641
Some ( b"zzzz" ) ,
642
642
] ) ) as ArrayRef ;
643
- let to_write = RecordBatch :: try_new ( delete_arrow_schema. clone ( ) , vec ! [
644
- col0, col1, col2, col3, col4, col5, col6, col7, col8, col9, col10, col11, col12, col13,
645
- ] )
646
- . unwrap ( ) ;
643
+ let to_write =
644
+ RecordBatch :: try_new ( Arc :: new ( schema_to_arrow_schema ( & schema) . unwrap ( ) ) , vec ! [
645
+ col0, col1, col2, col3, col4, col5, col6, col7, col8, col9, col10, col11, col12,
646
+ col13,
647
+ ] )
648
+ . unwrap ( ) ;
647
649
equality_delete_writer. write ( to_write. clone ( ) ) . await ?;
648
650
let res = equality_delete_writer. close ( ) . await ?;
649
651
assert_eq ! ( res. len( ) , 1 ) ;
0 commit comments