16
16
// under the License.
17
17
use crate :: decoder:: { VariantBasicType , VariantPrimitiveType } ;
18
18
use crate :: {
19
- ShortString , Variant , VariantDecimal16 , VariantDecimal4 , VariantDecimal8 , VariantMetadata ,
19
+ ShortString , Variant , VariantDecimal16 , VariantDecimal4 , VariantDecimal8 , VariantList ,
20
+ VariantMetadata , VariantObject ,
20
21
} ;
21
22
use arrow_schema:: ArrowError ;
22
23
use indexmap:: { IndexMap , IndexSet } ;
23
- use std:: collections:: { HashMap , HashSet } ;
24
+ use std:: collections:: HashSet ;
24
25
25
26
const BASIC_TYPE_BITS : u8 = 2 ;
26
27
const UNIX_EPOCH_DATE : chrono:: NaiveDate = chrono:: NaiveDate :: from_ymd_opt ( 1970 , 1 , 1 ) . unwrap ( ) ;
@@ -216,6 +217,57 @@ impl ValueBuffer {
216
217
self . append_slice ( value. as_bytes ( ) ) ;
217
218
}
218
219
220
+ fn append_object ( & mut self , metadata_builder : & mut MetadataBuilder , obj : VariantObject ) {
221
+ let mut object_builder = self . new_object ( metadata_builder) ;
222
+
223
+ for ( field_name, value) in obj. iter ( ) {
224
+ object_builder. insert ( field_name, value) ;
225
+ }
226
+
227
+ object_builder. finish ( ) . unwrap ( ) ;
228
+ }
229
+
230
+ fn try_append_object (
231
+ & mut self ,
232
+ metadata_builder : & mut MetadataBuilder ,
233
+ obj : VariantObject ,
234
+ ) -> Result < ( ) , ArrowError > {
235
+ let mut object_builder = self . new_object ( metadata_builder) ;
236
+
237
+ for res in obj. iter_try ( ) {
238
+ let ( field_name, value) = res?;
239
+ object_builder. try_insert ( field_name, value) ?;
240
+ }
241
+
242
+ object_builder. finish ( ) ?;
243
+
244
+ Ok ( ( ) )
245
+ }
246
+
247
+ fn append_list ( & mut self , metadata_builder : & mut MetadataBuilder , list : VariantList ) {
248
+ let mut list_builder = self . new_list ( metadata_builder) ;
249
+ for value in list. iter ( ) {
250
+ list_builder. append_value ( value) ;
251
+ }
252
+ list_builder. finish ( ) ;
253
+ }
254
+
255
+ fn try_append_list (
256
+ & mut self ,
257
+ metadata_builder : & mut MetadataBuilder ,
258
+ list : VariantList ,
259
+ ) -> Result < ( ) , ArrowError > {
260
+ let mut list_builder = self . new_list ( metadata_builder) ;
261
+ for res in list. iter_try ( ) {
262
+ let value = res?;
263
+ list_builder. try_append_value ( value) ?;
264
+ }
265
+
266
+ list_builder. finish ( ) ;
267
+
268
+ Ok ( ( ) )
269
+ }
270
+
219
271
fn offset ( & self ) -> usize {
220
272
self . 0 . len ( )
221
273
}
@@ -252,9 +304,31 @@ impl ValueBuffer {
252
304
variant : Variant < ' m , ' d > ,
253
305
metadata_builder : & mut MetadataBuilder ,
254
306
) {
255
- self . try_append_variant ( variant, metadata_builder) . unwrap ( ) ;
307
+ match variant {
308
+ Variant :: Null => self . append_null ( ) ,
309
+ Variant :: BooleanTrue => self . append_bool ( true ) ,
310
+ Variant :: BooleanFalse => self . append_bool ( false ) ,
311
+ Variant :: Int8 ( v) => self . append_int8 ( v) ,
312
+ Variant :: Int16 ( v) => self . append_int16 ( v) ,
313
+ Variant :: Int32 ( v) => self . append_int32 ( v) ,
314
+ Variant :: Int64 ( v) => self . append_int64 ( v) ,
315
+ Variant :: Date ( v) => self . append_date ( v) ,
316
+ Variant :: TimestampMicros ( v) => self . append_timestamp_micros ( v) ,
317
+ Variant :: TimestampNtzMicros ( v) => self . append_timestamp_ntz_micros ( v) ,
318
+ Variant :: Decimal4 ( decimal4) => self . append_decimal4 ( decimal4) ,
319
+ Variant :: Decimal8 ( decimal8) => self . append_decimal8 ( decimal8) ,
320
+ Variant :: Decimal16 ( decimal16) => self . append_decimal16 ( decimal16) ,
321
+ Variant :: Float ( v) => self . append_float ( v) ,
322
+ Variant :: Double ( v) => self . append_double ( v) ,
323
+ Variant :: Binary ( v) => self . append_binary ( v) ,
324
+ Variant :: String ( s) => self . append_string ( s) ,
325
+ Variant :: ShortString ( s) => self . append_short_string ( s) ,
326
+ Variant :: Object ( obj) => self . append_object ( metadata_builder, obj) ,
327
+ Variant :: List ( list) => self . append_list ( metadata_builder, list) ,
328
+ }
256
329
}
257
330
331
+ /// Appends a variant to the buffer
258
332
fn try_append_variant < ' m , ' d > (
259
333
& mut self ,
260
334
variant : Variant < ' m , ' d > ,
@@ -279,35 +353,8 @@ impl ValueBuffer {
279
353
Variant :: Binary ( v) => self . append_binary ( v) ,
280
354
Variant :: String ( s) => self . append_string ( s) ,
281
355
Variant :: ShortString ( s) => self . append_short_string ( s) ,
282
- Variant :: Object ( obj) => {
283
- let metadata_field_names = metadata_builder
284
- . field_names
285
- . iter ( )
286
- . enumerate ( )
287
- . map ( |( i, f) | ( f. clone ( ) , i) )
288
- . collect :: < HashMap < _ , _ > > ( ) ;
289
-
290
- let mut object_builder = self . new_object ( metadata_builder) ;
291
-
292
- // first add all object fields that exist in metadata builder
293
- let mut object_fields = obj. iter ( ) . collect :: < Vec < _ > > ( ) ;
294
-
295
- object_fields
296
- . sort_by_key ( |( field_name, _) | metadata_field_names. get ( field_name as & str ) ) ;
297
-
298
- for ( field_name, value) in object_fields {
299
- object_builder. insert ( field_name, value) ;
300
- }
301
-
302
- object_builder. finish ( ) ?;
303
- }
304
- Variant :: List ( list) => {
305
- let mut list_builder = self . new_list ( metadata_builder) ;
306
- for value in list. iter ( ) {
307
- list_builder. append_value ( value) ;
308
- }
309
- list_builder. finish ( ) ;
310
- }
356
+ Variant :: Object ( obj) => self . try_append_object ( metadata_builder, obj) ?,
357
+ Variant :: List ( list) => self . try_append_list ( metadata_builder, list) ?,
311
358
}
312
359
313
360
Ok ( ( ) )
0 commit comments