@@ -33,7 +33,7 @@ use serde_json::Value;
3333use crate :: {
3434 metadata:: SchemaVersion ,
3535 storage:: StreamType ,
36- utils:: arrow:: { get_field , get_timestamp_array , replace_columns } ,
36+ utils:: arrow:: { add_parseable_fields , get_field } ,
3737} ;
3838
3939use super :: { Event , DEFAULT_TIMESTAMP_KEY } ;
@@ -145,9 +145,10 @@ pub trait EventFormat: Sized {
145145 static_schema_flag : bool ,
146146 time_partition : Option < & String > ,
147147 schema_version : SchemaVersion ,
148+ p_custom_fields : & HashMap < String , String > ,
148149 ) -> Result < ( RecordBatch , bool ) , AnyError > {
149150 let p_timestamp = self . get_p_timestamp ( ) ;
150- let ( data, mut schema, is_first) = self . to_data (
151+ let ( data, schema, is_first) = self . to_data (
151152 storage_schema,
152153 time_partition,
153154 schema_version,
@@ -161,16 +162,6 @@ pub trait EventFormat: Sized {
161162 ) ) ;
162163 } ;
163164
164- // add the p_timestamp field to the event schema to the 0th index
165- schema. insert (
166- 0 ,
167- Arc :: new ( Field :: new (
168- DEFAULT_TIMESTAMP_KEY ,
169- DataType :: Timestamp ( TimeUnit :: Millisecond , None ) ,
170- true ,
171- ) ) ,
172- ) ;
173-
174165 // prepare the record batch and new fields to be added
175166 let mut new_schema = Arc :: new ( Schema :: new ( schema) ) ;
176167 if !Self :: is_schema_matching ( new_schema. clone ( ) , storage_schema, static_schema_flag) {
@@ -179,12 +170,8 @@ pub trait EventFormat: Sized {
179170 new_schema =
180171 update_field_type_in_schema ( new_schema, None , time_partition, None , schema_version) ;
181172
182- let mut rb = Self :: decode ( data, new_schema. clone ( ) ) ?;
183- rb = replace_columns (
184- rb. schema ( ) ,
185- & rb,
186- & [ ( 0 , Arc :: new ( get_timestamp_array ( p_timestamp, rb. num_rows ( ) ) ) ) ] ,
187- ) ;
173+ let rb = Self :: decode ( data, new_schema. clone ( ) ) ?;
174+ let rb = add_parseable_fields ( rb, p_timestamp, p_custom_fields) ?;
188175
189176 Ok ( ( rb, is_first) )
190177 }
@@ -222,6 +209,7 @@ pub trait EventFormat: Sized {
222209 time_partition : Option < & String > ,
223210 schema_version : SchemaVersion ,
224211 stream_type : StreamType ,
212+ p_custom_fields : & HashMap < String , String > ,
225213 ) -> Result < Event , AnyError > ;
226214}
227215
0 commit comments