CoCalc -- tests.rs

GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-parquet/src/arrow/write/nested/dremel/tests.rs
⁶⁹⁴⁰ views
1
use super::*;
2

3
mod def {
4
    use super::*;
5
    use crate::write::pages::{ListNested, PrimitiveNested, StructNested};
6

7
    fn test(nested: Vec<Nested>, expected: Vec<u16>) {
8
        let mut iter = BufferedDremelIter::new(&nested).map(|d| d.def);
9
        // assert_eq!(iter.size_hint().0, expected.len());
10
        let result = iter.by_ref().collect::<Vec<_>>();
11
        assert_eq!(result, expected);
12
        // assert_eq!(iter.size_hint().0, 0);
13
    }
14

15
    #[test]
16
    fn struct_dbl_optional() {
17
        let a = [true, true, true, false, true, true];
18
        let b = [true, false, true, false, false, true];
19
        let nested = vec![
20
            Nested::Struct(StructNested {
21
                is_optional: true,
22
                validity: Some(a.into()),
23
                length: 6,
24
            }),
25
            Nested::Primitive(PrimitiveNested {
26
                validity: Some(b.into()),
27
                is_optional: true,
28
                length: 6,
29
            }),
30
        ];
31
        let expected = vec![2, 1, 2, 0, 1, 2];
32

33
        test(nested, expected)
34
    }
35

36
    #[test]
37
    fn struct_optional() {
38
        let b = [
39
            true, false, true, true, false, true, false, false, true, true,
40
        ];
41
        let nested = vec![
42
            Nested::Struct(StructNested {
43
                is_optional: true,
44
                validity: None,
45
                length: 10,
46
            }),
47
            Nested::Primitive(PrimitiveNested {
48
                validity: Some(b.into()),
49
                is_optional: true,
50
                length: 10,
51
            }),
52
        ];
53
        let expected = vec![2, 1, 2, 2, 1, 2, 1, 1, 2, 2];
54

55
        test(nested, expected)
56
    }
57

58
    #[test]
59
    fn nested_edge_simple() {
60
        let nested = vec![
61
            Nested::List(ListNested {
62
                is_optional: true,
63
                offsets: vec![0, 2].try_into().unwrap(),
64
                validity: None,
65
            }),
66
            Nested::Primitive(PrimitiveNested {
67
                validity: None,
68
                is_optional: true,
69
                length: 2,
70
            }),
71
        ];
72
        let expected = vec![3, 3];
73

74
        test(nested, expected)
75
    }
76

77
    #[test]
78
    fn struct_optional_1() {
79
        let b = [
80
            true, false, true, true, false, true, false, false, true, true,
81
        ];
82
        let nested = vec![
83
            Nested::Struct(StructNested {
84
                validity: None,
85
                is_optional: true,
86
                length: 10,
87
            }),
88
            Nested::Primitive(PrimitiveNested {
89
                validity: Some(b.into()),
90
                is_optional: true,
91
                length: 10,
92
            }),
93
        ];
94
        let expected = vec![2, 1, 2, 2, 1, 2, 1, 1, 2, 2];
95

96
        test(nested, expected)
97
    }
98

99
    #[test]
100
    fn struct_optional_optional() {
101
        let nested = vec![
102
            Nested::Struct(StructNested {
103
                is_optional: true,
104
                validity: None,
105
                length: 10,
106
            }),
107
            Nested::Primitive(PrimitiveNested {
108
                validity: None,
109
                is_optional: true,
110
                length: 10,
111
            }),
112
        ];
113
        let expected = vec![2, 2, 2, 2, 2, 2, 2, 2, 2, 2];
114

115
        test(nested, expected)
116
    }
117

118
    #[test]
119
    fn l1_required_required() {
120
        let nested = vec![
121
            // [[0, 1], [], [2, 0, 3], [4, 5, 6], [], [7, 8, 9], [], [10]]
122
            Nested::List(ListNested {
123
                is_optional: false,
124
                offsets: vec![0, 2, 2, 5, 8, 8, 11, 11, 12].try_into().unwrap(),
125
                validity: None,
126
            }),
127
            Nested::Primitive(PrimitiveNested {
128
                validity: None,
129
                is_optional: false,
130
                length: 12,
131
            }),
132
        ];
133
        let expected = vec![1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1];
134

135
        test(nested, expected)
136
    }
137

138
    #[test]
139
    fn l1_optional_optional() {
140
        // [[0, 1], None, [2, None, 3], [4, 5, 6], [], [7, 8, 9], None, [10]]
141

142
        let v0 = [true, false, true, true, true, true, false, true];
143
        let v1 = [
144
            true, true, //[0, 1]
145
            true, false, true, //[2, None, 3]
146
            true, true, true, //[4, 5, 6]
147
            true, true, true, //[7, 8, 9]
148
            true, //[10]
149
        ];
150
        let nested = vec![
151
            Nested::List(ListNested {
152
                is_optional: true,
153
                offsets: vec![0, 2, 2, 5, 8, 8, 11, 11, 12].try_into().unwrap(),
154
                validity: Some(v0.into()),
155
            }),
156
            Nested::Primitive(PrimitiveNested {
157
                validity: Some(v1.into()),
158
                is_optional: true,
159
                length: 12,
160
            }),
161
        ];
162
        let expected = vec![3, 3, 0, 3, 2, 3, 3, 3, 3, 1, 3, 3, 3, 0, 3];
163

164
        test(nested, expected)
165
    }
166

167
    #[test]
168
    fn l2_required_required_required() {
169
        /*
170
        [
171
            [
172
                [1,2,3],
173
                [4,5,6,7],
174
            ],
175
            [
176
                [8],
177
                [9, 10]
178
            ]
179
        ]
180
        */
181
        let nested = vec![
182
            Nested::List(ListNested {
183
                is_optional: false,
184
                offsets: vec![0, 2, 4].try_into().unwrap(),
185
                validity: None,
186
            }),
187
            Nested::List(ListNested {
188
                is_optional: false,
189
                offsets: vec![0, 3, 7, 8, 10].try_into().unwrap(),
190
                validity: None,
191
            }),
192
            Nested::Primitive(PrimitiveNested {
193
                validity: None,
194
                is_optional: false,
195
                length: 10,
196
            }),
197
        ];
198
        let expected = vec![2, 2, 2, 2, 2, 2, 2, 2, 2, 2];
199

200
        test(nested, expected)
201
    }
202

203
    #[test]
204
    fn l2_optional_required_required() {
205
        let a = [true, false, true, true];
206
        /*
207
        [
208
            [
209
                [1,2,3],
210
                [4,5,6,7],
211
            ],
212
            None,
213
            [
214
                [8],
215
                [],
216
                [9, 10]
217
            ]
218
        ]
219
        */
220
        let nested = vec![
221
            Nested::List(ListNested {
222
                is_optional: true,
223
                offsets: vec![0, 2, 2, 2, 5].try_into().unwrap(),
224
                validity: Some(a.into()),
225
            }),
226
            Nested::List(ListNested {
227
                is_optional: false,
228
                offsets: vec![0, 3, 7, 8, 8, 10].try_into().unwrap(),
229
                validity: None,
230
            }),
231
            Nested::Primitive(PrimitiveNested {
232
                validity: None,
233
                is_optional: false,
234
                length: 10,
235
            }),
236
        ];
237
        let expected = vec![3, 3, 3, 3, 3, 3, 3, 0, 1, 3, 2, 3, 3];
238

239
        test(nested, expected)
240
    }
241

242
    mod fixedlist {
243
        use super::*;
244

245
        #[test]
246
        fn fsl() {
247
            /* [ [ 1, 2 ], None, [ None, 3 ] ] */
248
            let a = [true, false, true];
249
            let b = [true, true, false, false, false, true];
250
            let nested = vec![
251
                Nested::fixed_size_list(Some(a.into()), true, 2, 3),
252
                Nested::primitive(Some(b.into()), true, 6),
253
            ];
254
            let expected = vec![3, 3, 0, 2, 3];
255

256
            test(nested, expected)
257
        }
258

259
        #[test]
260
        fn fsl_fsl() {
261
            // [
262
            //    [ [ 1, 2, 3 ], [ 4, 5, 6 ] ],
263
            //    None,
264
            //    [ None, [ 7, None, 9 ] ],
265
            // ]
266
            let a = [true, false, true];
267
            let b = [true, true, true, true, false, true];
268
            let c = [
269
                true, true, true, true, true, true, false, false, false, false, false, false,
270
                false, false, false, true, false, true,
271
            ];
272
            let nested = vec![
273
                Nested::fixed_size_list(Some(a.into()), true, 2, 3),
274
                Nested::fixed_size_list(Some(b.into()), true, 3, 6),
275
                Nested::primitive(Some(c.into()), true, 18),
276
            ];
277
            let expected = vec![5, 5, 5, 5, 5, 5, 0, 2, 5, 4, 5];
278

279
            test(nested, expected)
280
        }
281

282
        #[test]
283
        fn fsl_fsl_1() {
284
            // [
285
            //     [ [1, 5, 2], [42, 13, 37] ],
286
            //     None,
287
            //     [ None, [3, 1, 3] ]
288
            // ]
289
            let a = [true, false, true];
290
            let b = [true, true, false, false, false, true];
291
            let c = [
292
                true, true, true, true, true, true, false, false, false, false, false, false,
293
                false, false, false, true, true, true,
294
            ];
295
            let nested = vec![
296
                Nested::fixed_size_list(Some(a.into()), true, 2, 3),
297
                Nested::fixed_size_list(Some(b.into()), true, 3, 6),
298
                Nested::primitive(Some(c.into()), true, 18),
299
            ];
300
            let expected = vec![5, 5, 5, 5, 5, 5, 0, 2, 5, 5, 5];
301

302
            test(nested, expected)
303
        }
304
    }
305

306
    mod simple {
307
        use super::*;
308

309
        #[test]
310
        fn none() {
311
            /* [ None ] */
312
            let a = [false];
313
            let b = [];
314
            let nested = vec![
315
                Nested::List(ListNested {
316
                    is_optional: true,
317
                    offsets: vec![0, 0].try_into().unwrap(),
318
                    validity: Some(a.into()),
319
                }),
320
                Nested::List(ListNested {
321
                    is_optional: true,
322
                    offsets: vec![0].try_into().unwrap(),
323
                    validity: Some(b.into()),
324
                }),
325
                Nested::Primitive(PrimitiveNested {
326
                    validity: None,
327
                    is_optional: false,
328
                    length: 0,
329
                }),
330
            ];
331
            let expected = vec![0];
332

333
            test(nested, expected)
334
        }
335

336
        #[test]
337
        fn empty() {
338
            /* [ [ ] ] */
339
            let a = [true];
340
            let b = [];
341
            let nested = vec![
342
                Nested::List(ListNested {
343
                    is_optional: true,
344
                    offsets: vec![0, 0].try_into().unwrap(),
345
                    validity: Some(a.into()),
346
                }),
347
                Nested::List(ListNested {
348
                    is_optional: true,
349
                    offsets: vec![0].try_into().unwrap(),
350
                    validity: Some(b.into()),
351
                }),
352
                Nested::Primitive(PrimitiveNested {
353
                    validity: None,
354
                    is_optional: false,
355
                    length: 0,
356
                }),
357
            ];
358
            let expected = vec![1];
359

360
            test(nested, expected)
361
        }
362

363
        #[test]
364
        fn list_none() {
365
            /* [ [ None ] ] */
366
            let a = [true];
367
            let b = [false];
368
            let nested = vec![
369
                Nested::List(ListNested {
370
                    is_optional: true,
371
                    offsets: vec![0, 1].try_into().unwrap(),
372
                    validity: Some(a.into()),
373
                }),
374
                Nested::List(ListNested {
375
                    is_optional: true,
376
                    offsets: vec![0, 0].try_into().unwrap(),
377
                    validity: Some(b.into()),
378
                }),
379
                Nested::Primitive(PrimitiveNested {
380
                    validity: None,
381
                    is_optional: false,
382
                    length: 0,
383
                }),
384
            ];
385
            let expected = vec![2];
386

387
            test(nested, expected)
388
        }
389

390
        #[test]
391
        fn list_empty() {
392
            /* [ [ [] ] ] */
393
            let a = [true];
394
            let b = [true];
395
            let nested = vec![
396
                Nested::List(ListNested {
397
                    is_optional: true,
398
                    offsets: vec![0, 1].try_into().unwrap(),
399
                    validity: Some(a.into()),
400
                }),
401
                Nested::List(ListNested {
402
                    is_optional: true,
403
                    offsets: vec![0, 0].try_into().unwrap(),
404
                    validity: Some(b.into()),
405
                }),
406
                Nested::Primitive(PrimitiveNested {
407
                    validity: None,
408
                    is_optional: false,
409
                    length: 0,
410
                }),
411
            ];
412
            let expected = vec![3];
413

414
            test(nested, expected)
415
        }
416

417
        #[test]
418
        fn list_list_one() {
419
            /* [ [ [ 1 ] ] ] */
420
            let a = [true];
421
            let b = [true];
422
            let nested = vec![
423
                Nested::List(ListNested {
424
                    is_optional: true,
425
                    offsets: vec![0, 1].try_into().unwrap(),
426
                    validity: Some(a.into()),
427
                }),
428
                Nested::List(ListNested {
429
                    is_optional: true,
430
                    offsets: vec![0, 1].try_into().unwrap(),
431
                    validity: Some(b.into()),
432
                }),
433
                Nested::Primitive(PrimitiveNested {
434
                    validity: None,
435
                    is_optional: false,
436
                    length: 1,
437
                }),
438
            ];
439
            let expected = vec![4];
440

441
            test(nested, expected)
442
        }
443
    }
444

445
    #[test]
446
    fn l2_optional_optional_required() {
447
        let a = [true, false, true];
448
        let b = [true, true, true, true, false];
449
        /*
450
        [
451
            [
452
                [1,2,3],
453
                [4,5,6,7],
454
            ],
455
            None,
456
            [
457
                [8],
458
                [],
459
                None,
460
            ],
461
        ]
462
        */
463
        let nested = vec![
464
            Nested::List(ListNested {
465
                is_optional: true,
466
                offsets: vec![0, 2, 2, 5].try_into().unwrap(),
467
                validity: Some(a.into()),
468
            }),
469
            Nested::List(ListNested {
470
                is_optional: true,
471
                offsets: vec![0, 3, 7, 8, 8, 8].try_into().unwrap(),
472
                validity: Some(b.into()),
473
            }),
474
            Nested::Primitive(PrimitiveNested {
475
                validity: None,
476
                is_optional: false,
477
                length: 8,
478
            }),
479
        ];
480
        let expected = vec![4, 4, 4, 4, 4, 4, 4, 0, 4, 3, 2];
481

482
        test(nested, expected)
483
    }
484

485
    #[test]
486
    fn l2_optional_optional_optional() {
487
        let a = [true, false, true];
488
        let b = [true, true, true, false];
489
        let c = [true, true, true, true, false, true, true, true];
490
        /*
491
        [
492
            [
493
                [1,2,3],
494
                [4,None,6,7],
495
            ],
496
            None,
497
            [
498
                [8],
499
                None,
500
            ],
501
        ]
502
        */
503
        let nested = vec![
504
            Nested::List(ListNested {
505
                is_optional: true,
506
                offsets: vec![0, 2, 2, 4].try_into().unwrap(),
507
                validity: Some(a.into()),
508
            }),
509
            Nested::List(ListNested {
510
                is_optional: true,
511
                offsets: vec![0, 3, 7, 8, 8].try_into().unwrap(),
512
                validity: Some(b.into()),
513
            }),
514
            Nested::Primitive(PrimitiveNested {
515
                validity: Some(c.into()),
516
                is_optional: true,
517
                length: 8,
518
            }),
519
        ];
520
        let expected = vec![5, 5, 5, 5, 4, 5, 5, 0, 5, 2];
521

522
        test(nested, expected)
523
    }
524

525
    /*
526
        [{"a": "a"}, {"a": "b"}],
527
        None,
528
        [{"a": "b"}, None, {"a": "b"}],
529
        [{"a": None}, {"a": None}, {"a": None}],
530
        [],
531
        [{"a": "d"}, {"a": "d"}, {"a": "d"}],
532
        None,
533
        [{"a": "e"}],
534
    */
535
    #[test]
536
    fn nested_list_struct_nullable() {
537
        let a = [
538
            true, true, true, false, true, false, false, false, true, true, true, true,
539
        ];
540
        let b = [
541
            true, true, true, false, true, true, true, true, true, true, true, true,
542
        ];
543
        let c = [true, false, true, true, true, true, false, true];
544
        let nested = vec![
545
            Nested::List(ListNested {
546
                is_optional: true,
547
                offsets: vec![0, 2, 2, 5, 8, 8, 11, 11, 12].try_into().unwrap(),
548
                validity: Some(c.into()),
549
            }),
550
            Nested::Struct(StructNested {
551
                validity: Some(b.into()),
552
                is_optional: true,
553
                length: 12,
554
            }),
555
            Nested::Primitive(PrimitiveNested {
556
                validity: Some(a.into()),
557
                is_optional: true,
558
                length: 12,
559
            }),
560
        ];
561
        let expected = vec![4, 4, 0, 4, 2, 4, 3, 3, 3, 1, 4, 4, 4, 0, 4];
562

563
        test(nested, expected)
564
    }
565

566
    #[test]
567
    fn nested_list_struct_nullable1() {
568
        let c = [true, false];
569
        let nested = vec![
570
            Nested::List(ListNested {
571
                is_optional: true,
572
                offsets: vec![0, 1, 1].try_into().unwrap(),
573
                validity: Some(c.into()),
574
            }),
575
            Nested::Struct(StructNested {
576
                validity: None,
577
                is_optional: true,
578
                length: 1,
579
            }),
580
            Nested::Primitive(PrimitiveNested {
581
                validity: None,
582
                is_optional: true,
583
                length: 1,
584
            }),
585
        ];
586
        let expected = vec![4, 0];
587

588
        test(nested, expected)
589
    }
590

591
    #[test]
592
    fn nested_struct_list_nullable() {
593
        // [
594
        //     { "a": [] },
595
        //     { "a", [] },
596
        // ]
597
        let a = [true, false, true, true, true, true, false, true];
598
        let b = [
599
            true, true, true, false, true, true, true, true, true, true, true, true,
600
        ];
601
        let nested = vec![
602
            Nested::Struct(StructNested {
603
                validity: None,
604
                is_optional: true,
605
                length: 8,
606
            }),
607
            Nested::List(ListNested {
608
                is_optional: true,
609
                offsets: vec![0, 2, 2, 5, 8, 8, 11, 11, 12].try_into().unwrap(),
610
                validity: Some(a.into()),
611
            }),
612
            Nested::Primitive(PrimitiveNested {
613
                validity: Some(b.into()),
614
                is_optional: true,
615
                length: 12,
616
            }),
617
        ];
618
        let expected = vec![4, 4, 1, 4, 3, 4, 4, 4, 4, 2, 4, 4, 4, 1, 4];
619

620
        test(nested, expected)
621
    }
622

623
    #[test]
624
    fn nested_struct_list_nullable1() {
625
        let a = [true, true, false];
626
        let nested = vec![
627
            Nested::Struct(StructNested {
628
                validity: None,
629
                is_optional: true,
630
                length: 3,
631
            }),
632
            Nested::List(ListNested {
633
                is_optional: true,
634
                offsets: vec![0, 1, 1, 1].try_into().unwrap(),
635
                validity: Some(a.into()),
636
            }),
637
            Nested::Primitive(PrimitiveNested {
638
                validity: None,
639
                is_optional: true,
640
                length: 1,
641
            }),
642
        ];
643
        let expected = vec![4, 2, 1];
644

645
        test(nested, expected)
646
    }
647

648
    #[test]
649
    fn nested_list_struct_list_nullable1() {
650
        /*
651
        [
652
            [{"a": ["b"]}, None],
653
        ]
654
        */
655

656
        let a = [true];
657
        let b = [true, false];
658
        let c = [true, false];
659
        let d = [true];
660
        let nested = vec![
661
            Nested::List(ListNested {
662
                is_optional: true,
663
                offsets: vec![0, 2].try_into().unwrap(),
664
                validity: Some(a.into()),
665
            }),
666
            Nested::Struct(StructNested {
667
                validity: Some(b.into()),
668
                is_optional: true,
669
                length: 2,
670
            }),
671
            Nested::List(ListNested {
672
                is_optional: true,
673
                offsets: vec![0, 1, 1].try_into().unwrap(),
674
                validity: Some(c.into()),
675
            }),
676
            Nested::Primitive(PrimitiveNested {
677
                validity: Some(d.into()),
678
                is_optional: true,
679
                length: 1,
680
            }),
681
        ];
682
        /*
683
                0 6
684
                1 6
685
                0 0
686
                0 6
687
                1 2
688
        */
689
        let expected = vec![6, 2];
690

691
        test(nested, expected)
692
    }
693

694
    #[test]
695
    fn nested_list_struct_list_nullable() {
696
        /*
697
            [
698
            [{"a": ["a"]}, {"a": ["b"]}],
699
            None,
700
            [{"a": ["b"]}, None, {"a": ["b"]}],
701
            [{"a": None}, {"a": None}, {"a": None}],
702
            [],
703
            [{"a": ["d"]}, {"a": [None]}, {"a": ["c", "d"]}],
704
            None,
705
            [{"a": []}],
706
        ]
707
            */
708
        let a = [true, false, true, true, true, true, false, true];
709
        let b = [
710
            true, true, true, false, true, true, true, true, true, true, true, true,
711
        ];
712
        let c = [
713
            true, true, true, false, true, false, false, false, true, true, true, true,
714
        ];
715
        let d = [true, true, true, true, true, false, true, true];
716
        let nested = vec![
717
            Nested::List(ListNested {
718
                is_optional: true,
719
                offsets: vec![0, 2, 2, 5, 8, 8, 11, 11, 12].try_into().unwrap(),
720
                validity: Some(a.into()),
721
            }),
722
            Nested::Struct(StructNested {
723
                validity: Some(b.into()),
724
                is_optional: true,
725
                length: 12,
726
            }),
727
            Nested::List(ListNested {
728
                is_optional: true,
729
                offsets: vec![0, 1, 2, 3, 3, 4, 4, 4, 4, 5, 6, 8, 8]
730
                    .try_into()
731
                    .unwrap(),
732
                validity: Some(c.into()),
733
            }),
734
            Nested::Primitive(PrimitiveNested {
735
                validity: Some(d.into()),
736
                is_optional: true,
737
                length: 8,
738
            }),
739
        ];
740
        let expected = vec![6, 6, 0, 6, 2, 6, 3, 3, 3, 1, 6, 5, 6, 6, 0, 4];
741

742
        test(nested, expected)
743
    }
744
}
745

746
mod rep {
747
    use super::super::super::super::pages::ListNested;
748
    use super::*;
749

750
    fn test(nested: Vec<Nested>, expected: Vec<u16>) {
751
        let mut iter = BufferedDremelIter::new(&nested).map(|d| d.rep);
752
        // assert_eq!(iter.size_hint().0, expected.len());
753
        assert_eq!(iter.by_ref().collect::<Vec<_>>(), expected);
754
        // assert_eq!(iter.size_hint().0, 0);
755
    }
756

757
    #[test]
758
    fn struct_required() {
759
        let nested = vec![
760
            Nested::structure(None, false, 10),
761
            Nested::primitive(None, true, 10),
762
        ];
763
        let expected = vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
764

765
        test(nested, expected)
766
    }
767

768
    #[test]
769
    fn struct_optional() {
770
        let nested = vec![
771
            Nested::structure(None, true, 10),
772
            Nested::primitive(None, true, 10),
773
        ];
774
        let expected = vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
775

776
        test(nested, expected)
777
    }
778

779
    #[test]
780
    fn l1() {
781
        // [
782
        //    [ 1, 2 ],
783
        //    [],
784
        //    [ 3, 4, 5 ],
785
        //    [ 6, 7, 8 ],
786
        //    [],
787
        //    [ 9, 10, 11 ],
788
        //    [],
789
        //    [ 12 ],
790
        // ]
791
        let nested = vec![
792
            Nested::list(
793
                None,
794
                false,
795
                vec![0, 2, 2, 5, 8, 8, 11, 11, 12].try_into().unwrap(),
796
            ),
797
            Nested::primitive(None, false, 12),
798
        ];
799
        let expected = vec![0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0];
800

801
        test(nested, expected)
802
    }
803

804
    #[test]
805
    fn l2() {
806
        let nested = vec![
807
            Nested::List(ListNested {
808
                is_optional: false,
809
                offsets: vec![0, 2, 2, 4].try_into().unwrap(),
810
                validity: None,
811
            }),
812
            Nested::List(ListNested {
813
                is_optional: false,
814
                offsets: vec![0, 3, 7, 8, 10].try_into().unwrap(),
815
                validity: None,
816
            }),
817
            Nested::primitive(None, false, 10),
818
        ];
819
        let expected = vec![0, 2, 2, 1, 2, 2, 2, 0, 0, 1, 2];
820

821
        test(nested, expected)
822
    }
823

824
    #[test]
825
    fn list_of_struct() {
826
        /*
827
        [
828
            [{"a": "b"}],[{"a": "c"}]
829
        ]
830
        */
831
        let nested = vec![
832
            Nested::List(ListNested {
833
                is_optional: true,
834
                offsets: vec![0, 1, 2].try_into().unwrap(),
835
                validity: None,
836
            }),
837
            Nested::structure(None, true, 2),
838
            Nested::primitive(None, true, 2),
839
        ];
840
        let expected = vec![0, 0];
841

842
        test(nested, expected)
843
    }
844

845
    #[test]
846
    fn list_struct_list() {
847
        let nested = vec![
848
            Nested::List(ListNested {
849
                is_optional: true,
850
                offsets: vec![0, 2, 3].try_into().unwrap(),
851
                validity: None,
852
            }),
853
            Nested::structure(None, true, 3),
854
            Nested::List(ListNested {
855
                is_optional: true,
856
                offsets: vec![0, 3, 6, 7].try_into().unwrap(),
857
                validity: None,
858
            }),
859
            Nested::primitive(None, true, 7),
860
        ];
861
        let expected = vec![0, 2, 2, 1, 2, 2, 0];
862

863
        test(nested, expected)
864
    }
865

866
    #[test]
867
    fn struct_list_optional() {
868
        /*
869
        {"f1": ["a", "b", None, "c"]}
870
        */
871
        let nested = vec![
872
            Nested::structure(None, true, 1),
873
            Nested::List(ListNested {
874
                is_optional: true,
875
                offsets: vec![0, 4].try_into().unwrap(),
876
                validity: None,
877
            }),
878
            Nested::primitive(None, true, 4),
879
        ];
880
        let expected = vec![0, 1, 1, 1];
881

882
        test(nested, expected)
883
    }
884

885
    #[test]
886
    fn l2_other() {
887
        let nested = vec![
888
            Nested::List(ListNested {
889
                is_optional: false,
890
                offsets: vec![0, 1, 1, 3, 5, 5, 8, 8, 9].try_into().unwrap(),
891
                validity: None,
892
            }),
893
            Nested::List(ListNested {
894
                is_optional: false,
895
                offsets: vec![0, 2, 4, 5, 7, 8, 9, 10, 11, 12].try_into().unwrap(),
896
                validity: None,
897
            }),
898
            Nested::primitive(None, false, 12),
899
        ];
900
        let expected = vec![0, 2, 0, 0, 2, 1, 0, 2, 1, 0, 0, 1, 1, 0, 0];
901

902
        test(nested, expected)
903
    }
904

905
    #[test]
906
    fn list_struct_list_1() {
907
        /*
908
        [
909
            [{"a": ["a"]}, {"a": ["b"]}],
910
            [],
911
            [{"a": ["b"]}, None, {"a": ["b"]}],
912
            [{"a": []}, {"a": []}, {"a": []}],
913
            [],
914
            [{"a": ["d"]}, {"a": ["a"]}, {"a": ["c", "d"]}],
915
            [],
916
            [{"a": []}],
917
        ]
918
        // reps: [0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 0, 1, 1, 2, 0, 0]
919
        */
920
        let a = [
921
            true, true, true, false, true, true, true, true, true, true, true, true,
922
        ];
923
        let nested = vec![
924
            Nested::List(ListNested {
925
                is_optional: true,
926
                offsets: vec![0, 2, 2, 5, 8, 8, 11, 11, 12].try_into().unwrap(),
927
                validity: None,
928
            }),
929
            Nested::structure(Some(a.into()), true, 12),
930
            Nested::List(ListNested {
931
                is_optional: true,
932
                offsets: vec![0, 1, 2, 2, 3, 3, 4, 4, 4, 4, 5, 7, 8]
933
                    .try_into()
934
                    .unwrap(),
935
                validity: None,
936
            }),
937
            Nested::primitive(None, true, 8),
938
        ];
939
        let expected = vec![0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 0, 1, 1, 2, 0, 0];
940

941
        test(nested, expected)
942
    }
943

944
    #[test]
945
    fn list_struct_list_2() {
946
        /*
947
        [
948
            [{"a": []}],
949
        ]
950
        // reps: [0]
951
        */
952
        let nested = vec![
953
            Nested::List(ListNested {
954
                is_optional: true,
955
                offsets: vec![0, 1].try_into().unwrap(),
956
                validity: None,
957
            }),
958
            Nested::structure(None, true, 12),
959
            Nested::List(ListNested {
960
                is_optional: true,
961
                offsets: vec![0, 0].try_into().unwrap(),
962
                validity: None,
963
            }),
964
            Nested::primitive(None, true, 0),
965
        ];
966
        let expected = vec![0];
967

968
        test(nested, expected)
969
    }
970

971
    #[test]
972
    fn list_struct_list_3() {
973
        let nested = vec![
974
            Nested::List(ListNested {
975
                is_optional: true,
976
                offsets: vec![0, 1, 1].try_into().unwrap(),
977
                validity: None,
978
            }),
979
            Nested::structure(None, true, 12),
980
            Nested::List(ListNested {
981
                is_optional: true,
982
                offsets: vec![0, 0].try_into().unwrap(),
983
                validity: None,
984
            }),
985
            Nested::primitive(None, true, 0),
986
        ];
987
        let expected = vec![0, 0];
988
        // [1, 0], [0]
989
        // pick last
990

991
        test(nested, expected)
992
    }
993
}
994

995
Product

Resources

Company
