Skip to main content

mito2/sst/parquet/
read_columns.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::{HashMap, HashSet};
16
17use parquet::arrow::ProjectionMask;
18use parquet::basic::{ConvertedType, Type as PhysicalType};
19use parquet::schema::types::{ColumnDescriptor, SchemaDescriptor};
20
21/// A nested field access path inside one parquet root column.
22pub type ParquetNestedPath = Vec<String>;
23
24/// The parquet columns to read.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct ParquetReadColumns {
27    /// Root parquet column indices in the same order as `cols`.
28    ///
29    /// Most readers need these indices as a borrowed slice for Arrow schema
30    /// projection or parquet root-column projection. Keeping them here avoids
31    /// repeatedly collecting `cols.iter().map(|col| col.root_index)`.
32    root_indices: Vec<usize>,
33    cols: Vec<ParquetReadColumn>,
34    has_nested: bool,
35}
36
37impl ParquetReadColumns {
38    /// Builds parquet read columns from deduplicated, normalized input.
39    ///
40    /// `cols` must not contain duplicate root indices, and nested paths must
41    /// already be merged. Empty `nested_paths` means reading the whole root column.
42    ///
43    /// This constructor does not validate or merge input.
44    pub fn from_deduped(cols: Vec<ParquetReadColumn>) -> Self {
45        let has_nested = cols.iter().any(|col| !col.nested_paths.is_empty());
46        let root_indices = cols.iter().map(|col| col.root_index).collect();
47        Self {
48            root_indices,
49            cols,
50            has_nested,
51        }
52    }
53
54    /// Builds root-column projections from root indices that are already
55    /// deduplicated.
56    ///
57    /// Note: this constructor does not check for duplicates.
58    pub fn from_deduped_root_indices(root_indices: impl IntoIterator<Item = usize>) -> Self {
59        let root_indices = root_indices.into_iter().collect::<Vec<_>>();
60        let cols = root_indices
61            .iter()
62            .copied()
63            .map(ParquetReadColumn::new)
64            .collect();
65        Self {
66            root_indices,
67            cols,
68            has_nested: false,
69        }
70    }
71
72    pub fn columns(&self) -> &[ParquetReadColumn] {
73        &self.cols
74    }
75
76    pub fn has_nested(&self) -> bool {
77        self.has_nested
78    }
79
80    pub fn root_indices_iter(&self) -> impl Iterator<Item = usize> + '_ {
81        self.root_indices.iter().copied()
82    }
83
84    /// Returns root parquet column indices.
85    pub fn root_indices(&self) -> &[usize] {
86        &self.root_indices
87    }
88}
89
90/// Read requirement for a single parquet root column.
91///
92/// `root_index` identifies the root column in the parquet schema.
93///
94/// If `nested_paths` is empty, the whole root column is read. Otherwise, only
95/// leaves under the specified nested paths are read.
96///
97/// To construct a [`ParquetReadColumn`]:
98/// - `ParquetReadColumn::new(0)` reads the whole root column at index `0`.
99/// - `ParquetReadColumn::new(0).with_nested_paths(vec![vec!["j".into(), "b".into()]])`
100///   reads only leaves under `j.b`.
101#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct ParquetReadColumn {
103    /// Root field index in the parquet schema.
104    root_index: usize,
105    /// Nested paths to read under this root column.
106    ///
107    /// Each path includes the root column itself. For example, for a root
108    /// column `j`, path `["j", "a", "b"]` refers to `j.a.b`.
109    ///
110    /// If empty, the whole root column is read.
111    nested_paths: Vec<ParquetNestedPath>,
112}
113
114impl ParquetReadColumn {
115    pub fn new(root_index: usize) -> Self {
116        Self {
117            root_index,
118            nested_paths: vec![],
119        }
120    }
121
122    pub fn with_nested_paths(self, nested_paths: Vec<ParquetNestedPath>) -> Self {
123        Self {
124            nested_paths,
125            ..self
126        }
127    }
128
129    /// Merges additional nested paths into this root column.
130    pub fn merge_nested_paths(&mut self, nested_paths: Vec<ParquetNestedPath>) {
131        let reads_whole_root = self.nested_paths.is_empty() || nested_paths.is_empty();
132        if reads_whole_root {
133            // Empty nested paths means reading the whole root column.
134            self.nested_paths = vec![];
135        } else {
136            self.nested_paths.extend(nested_paths);
137        }
138    }
139
140    pub fn root_index(&self) -> usize {
141        self.root_index
142    }
143
144    pub fn nested_paths(&self) -> &[ParquetNestedPath] {
145        &self.nested_paths
146    }
147}
148
149/// Projection plan built for a parquet file.
150#[derive(Clone)]
151pub struct ProjectionMaskPlan {
152    /// `mask` is the projection mask applied to the parquet reader.
153    pub mask: ProjectionMask,
154    /// A boolean mask in output schema order indicating whether each projected root
155    /// has data read from the current parquet file.
156    ///
157    /// - `true`: parquet reads either requested nested data or a fallback parent
158    ///   for this root.
159    /// - `false`: this root is absent and must be synthesized later.
160    ///
161    /// The length of `projected_root_presence` is always equal to the
162    /// number of fields in the output schema.
163    pub projected_root_presence: Vec<bool>,
164}
165
166/// Builds a projection mask plan for reading a parquet file.
167///
168/// `parquet_read_cols` defines the requested root columns and optional
169/// nested paths to read.
170///
171/// `parquet_schema_desc` is the schema descriptor of the current parquet
172/// file. It is used to resolve requested nested paths to actual leaf
173/// column indices.
174///
175/// See [`ProjectionMaskPlan`] for the returned value.
176///
177/// For example, if the query requests `j.a` and `k`, but the current
178/// parquet file only contains leaves under `j.b` and `k`, then the
179/// returned plan keeps `k` in the projection mask and marks `j` as
180/// not present in the output, so it can be synthesized during
181/// post-processing.
182pub fn build_projection_plan(
183    parquet_read_cols: &ParquetReadColumns,
184    parquet_schema_desc: &SchemaDescriptor,
185) -> ProjectionMaskPlan {
186    if !parquet_read_cols.has_nested() {
187        let mask =
188            ProjectionMask::roots(parquet_schema_desc, parquet_read_cols.root_indices_iter());
189        return ProjectionMaskPlan {
190            mask,
191            projected_root_presence: vec![true; parquet_read_cols.columns().len()],
192        };
193    }
194
195    let (matched_leaves, matched_roots) =
196        build_parquet_leaves_indices(parquet_schema_desc, parquet_read_cols);
197
198    let projected_root_presence = parquet_read_cols
199        .columns()
200        .iter()
201        .map(|col| matched_roots.contains(&col.root_index()))
202        .collect();
203
204    let mask = ProjectionMask::leaves(parquet_schema_desc, matched_leaves);
205    ProjectionMaskPlan {
206        mask,
207        projected_root_presence,
208    }
209}
210
211/// Builds parquet leaf-column indices for reading a parquet file.
212///
213/// Returns `(matched_leaves, matched_roots)`:
214/// - `matched_leaves`: matched leaf-column indices in the current parquet file schema.
215/// - `matched_roots`: root-field indices read from the current parquet file schema.
216fn build_parquet_leaves_indices(
217    parquet_schema_desc: &SchemaDescriptor,
218    projection: &ParquetReadColumns,
219) -> (Vec<usize>, HashSet<usize>) {
220    let mut map = HashMap::with_capacity(projection.cols.len());
221    for col in &projection.cols {
222        map.insert(col.root_index, col);
223    }
224
225    let mut matched_leaves = HashSet::new();
226    let mut matched_roots = HashSet::with_capacity(projection.cols.len());
227
228    // Tracks whether each requested nested path matched by prefix without fallback.
229    let mut prefix_matched = HashMap::<usize, Vec<bool>>::new();
230    for col in &projection.cols {
231        prefix_matched.insert(col.root_index, vec![false; col.nested_paths.len()]);
232    }
233
234    // First select root/nested prefix matches.
235    for (leaf_idx, leaf_col) in parquet_schema_desc.columns().iter().enumerate() {
236        let root_idx = parquet_schema_desc.get_column_root_idx(leaf_idx);
237        let Some(col) = map.get(&root_idx) else {
238            continue;
239        };
240        if col.nested_paths.is_empty() {
241            matched_leaves.insert(leaf_idx);
242            matched_roots.insert(root_idx);
243            continue;
244        }
245
246        let leaf_path = leaf_col.path().parts();
247        let mut matched = false;
248        for (path_idx, _) in col
249            .nested_paths
250            .iter()
251            .enumerate()
252            .filter(|(_, nested_path)| leaf_path.starts_with(nested_path))
253        {
254            prefix_matched.get_mut(&root_idx).unwrap()[path_idx] = true;
255            matched = true;
256        }
257
258        if matched {
259            matched_leaves.insert(leaf_idx);
260            matched_roots.insert(root_idx);
261        }
262    }
263
264    // Then fallback prefix misses to their nearest variant parent.
265    // TODO(fys): Gate fallback planning on the root being JSON2. A raw Binary
266    // leaf is a JSONB variant only under a JSON2 root; plain struct Binary
267    // children should not enter this fallback path.
268    for col in &projection.cols {
269        for (path_idx, nested_path) in col.nested_paths.iter().enumerate() {
270            if prefix_matched[&col.root_index][path_idx] {
271                continue;
272            }
273
274            let Some(leaf_idx) =
275                find_nearest_variant_parent(parquet_schema_desc, col.root_index, nested_path)
276            else {
277                continue;
278            };
279
280            matched_leaves.insert(leaf_idx);
281            matched_roots.insert(col.root_index);
282        }
283    }
284
285    let mut matched_leaves = matched_leaves.into_iter().collect::<Vec<_>>();
286    matched_leaves.sort_unstable();
287    (matched_leaves, matched_roots)
288}
289
290fn find_nearest_variant_parent(
291    parquet_schema_desc: &SchemaDescriptor,
292    root_idx: usize,
293    nested_path: &[String],
294) -> Option<usize> {
295    // TODO(fys): Build a variant path index if fallback lookup becomes hot.
296    if nested_path.len() <= 1 {
297        return None;
298    }
299
300    // JSON2 root columns are always structured fields. Fallback only applies to
301    // variant leaves below the root.
302    for parent_len in (2..nested_path.len()).rev() {
303        let parent_path = &nested_path[..parent_len];
304        for (leaf_idx, leaf_col) in parquet_schema_desc.columns().iter().enumerate() {
305            if parquet_schema_desc.get_column_root_idx(leaf_idx) != root_idx {
306                continue;
307            }
308            if leaf_col.path().parts() == parent_path && is_variant_leaf(leaf_col) {
309                return Some(leaf_idx);
310            }
311        }
312    }
313
314    None
315}
316
317fn is_variant_leaf(leaf_col: &ColumnDescriptor) -> bool {
318    // TODO(fys): Recognize JSON2 variant type from Arrow extension metadata
319    // if the parquet Arrow schema preserves it for nested fields.
320    matches!(
321        leaf_col.physical_type(),
322        PhysicalType::BYTE_ARRAY | PhysicalType::FIXED_LEN_BYTE_ARRAY
323    ) && leaf_col.logical_type_ref().is_none()
324        && leaf_col.converted_type() == ConvertedType::NONE
325}
326
327#[cfg(test)]
328mod tests {
329    use std::sync::Arc;
330
331    use parquet::basic::{ConvertedType, LogicalType, Repetition};
332    use parquet::schema::types::Type;
333
334    use super::*;
335
336    #[test]
337    fn test_build_projection_mask_without_nested_paths() {
338        let parquet_schema_desc = build_test_nested_parquet_schema();
339        let projection = ParquetReadColumns::from_deduped_root_indices([0, 1]);
340
341        let plan = build_projection_plan(&projection, &parquet_schema_desc);
342
343        assert_eq!(vec![true, true], plan.projected_root_presence);
344        assert_eq!(
345            ProjectionMask::roots(&parquet_schema_desc, [0, 1]),
346            plan.mask
347        );
348    }
349
350    #[test]
351    fn test_reads_whole_root() {
352        let parquet_schema_desc = build_test_nested_parquet_schema();
353
354        let projection = ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0)]);
355
356        let (matched_leaves, matched_roots) =
357            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
358        assert_eq!(vec![0, 1, 2], matched_leaves);
359        assert_eq!(HashSet::from([0]), matched_roots);
360    }
361
362    #[test]
363    fn test_filters_nested_paths() {
364        let parquet_schema_desc = build_test_nested_parquet_schema();
365
366        let projection = ParquetReadColumns::from_deduped(vec![
367            ParquetReadColumn::new(0)
368                .with_nested_paths(vec![vec!["j".to_string(), "b".to_string()]]),
369            ParquetReadColumn::new(1),
370        ]);
371
372        let (matched_leaves, matched_roots) =
373            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
374        assert_eq!(vec![1, 2, 3], matched_leaves);
375        assert_eq!(HashSet::from([0, 1]), matched_roots);
376    }
377
378    #[test]
379    fn test_reads_middle_level_path() {
380        let parquet_schema_desc = build_test_nested_parquet_schema();
381
382        let projection = ParquetReadColumns::from_deduped(vec![
383            ParquetReadColumn::new(0)
384                .with_nested_paths(vec![vec!["j".to_string(), "b".to_string()]]),
385        ]);
386
387        let (matched_leaves, matched_roots) =
388            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
389        assert_eq!(vec![1, 2], matched_leaves);
390        assert_eq!(HashSet::from([0]), matched_roots);
391    }
392
393    #[test]
394    fn test_parent_path_covers_redundant_child_path() {
395        let parquet_schema_desc = build_test_nested_parquet_schema();
396        let nested_paths = vec![
397            vec!["j".to_string(), "b".to_string()],
398            vec!["j".to_string(), "b".to_string(), "c".to_string()],
399        ];
400
401        let read_column = ParquetReadColumn::new(0).with_nested_paths(nested_paths);
402        let projection = ParquetReadColumns::from_deduped(vec![read_column]);
403
404        let (matched_leaves, matched_roots) =
405            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
406        assert_eq!(vec![1, 2], matched_leaves);
407        assert_eq!(HashSet::from([0]), matched_roots);
408    }
409
410    #[test]
411    fn test_reads_leaf_level_path() {
412        let parquet_schema_desc = build_test_nested_parquet_schema();
413
414        let projection =
415            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
416                vec![vec!["j".to_string(), "b".to_string(), "c".to_string()]],
417            )]);
418
419        let (matched_leaves, matched_roots) =
420            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
421        assert_eq!(vec![1], matched_leaves);
422        assert_eq!(HashSet::from([0]), matched_roots);
423    }
424
425    #[test]
426    fn test_build_projection_mask_with_unmatched_roots() {
427        let parquet_schema_desc = build_test_nested_parquet_schema();
428
429        let projection = ParquetReadColumns::from_deduped(vec![
430            ParquetReadColumn::new(0)
431                .with_nested_paths(vec![vec!["j".to_string(), "missing".to_string()]]),
432            ParquetReadColumn::new(1),
433        ]);
434
435        let plan = build_projection_plan(&projection, &parquet_schema_desc);
436
437        assert_eq!(vec![false, true], plan.projected_root_presence);
438        assert_eq!(
439            ProjectionMask::leaves(&parquet_schema_desc, vec![3]),
440            plan.mask
441        );
442    }
443
444    #[test]
445    fn test_merges_mixed_paths() {
446        let parquet_schema_desc = build_test_nested_parquet_schema();
447
448        let projection =
449            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
450                vec![
451                    vec!["j".to_string(), "a".to_string()],
452                    vec!["j".to_string(), "b".to_string(), "d".to_string()],
453                ],
454            )]);
455
456        let (matched_leaves, matched_roots) =
457            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
458        assert_eq!(vec![0, 2], matched_leaves);
459        assert_eq!(HashSet::from([0]), matched_roots);
460    }
461
462    #[test]
463    fn test_merge_nested_paths_extends_paths() {
464        let mut col = ParquetReadColumn::new(0)
465            .with_nested_paths(vec![vec!["j".to_string(), "a".to_string()]]);
466
467        col.merge_nested_paths(vec![vec!["j".to_string(), "b".to_string()]]);
468
469        assert_eq!(
470            &[
471                vec!["j".to_string(), "a".to_string()],
472                vec!["j".to_string(), "b".to_string()],
473            ],
474            col.nested_paths()
475        );
476    }
477
478    #[test]
479    fn test_merge_nested_paths_with_whole_root() {
480        let mut col = ParquetReadColumn::new(0)
481            .with_nested_paths(vec![vec!["j".to_string(), "a".to_string()]]);
482
483        col.merge_nested_paths(vec![]);
484
485        assert!(col.nested_paths().is_empty());
486    }
487
488    #[test]
489    fn test_fallback_to_nearest_variant_parent() {
490        let parquet_schema_desc = build_test_variant_parent_schema();
491        let projection =
492            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
493                vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]],
494            )]);
495
496        let plan = build_projection_plan(&projection, &parquet_schema_desc);
497
498        assert_eq!(vec![true], plan.projected_root_presence);
499        assert_eq!(
500            ProjectionMask::leaves(&parquet_schema_desc, vec![0]),
501            plan.mask
502        );
503    }
504
505    #[test]
506    fn test_prefix_match_prevents_variant_parent_fallback() {
507        let parquet_schema_desc = build_test_variant_parent_schema();
508        let projection =
509            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
510                vec![vec!["j".to_string(), "b".to_string(), "x".to_string()]],
511            )]);
512
513        let plan = build_projection_plan(&projection, &parquet_schema_desc);
514
515        assert_eq!(vec![true], plan.projected_root_presence);
516        assert_eq!(
517            ProjectionMask::leaves(&parquet_schema_desc, vec![1]),
518            plan.mask
519        );
520    }
521
522    #[test]
523    fn test_mixed_prefix_and_fallback_paths() {
524        let parquet_schema_desc = build_test_variant_parent_schema();
525        let nested_paths = vec![
526            vec!["j".to_string(), "a".to_string(), "x".to_string()],
527            vec!["j".to_string(), "b".to_string(), "x".to_string()],
528        ];
529        let projection = ParquetReadColumns::from_deduped(vec![
530            ParquetReadColumn::new(0).with_nested_paths(nested_paths),
531        ]);
532
533        let plan = build_projection_plan(&projection, &parquet_schema_desc);
534
535        assert_eq!(vec![true], plan.projected_root_presence);
536        assert_eq!(
537            ProjectionMask::leaves(&parquet_schema_desc, vec![0, 1]),
538            plan.mask
539        );
540    }
541
542    #[test]
543    fn test_fallback_selects_multiple_variant_parents() {
544        let parquet_schema_desc = build_test_two_variant_parents_schema();
545        let nested_paths = vec![
546            vec!["j".to_string(), "a".to_string(), "y".to_string()],
547            vec!["j".to_string(), "b".to_string(), "d".to_string()],
548            vec!["j".to_string(), "a".to_string(), "x".to_string()],
549        ];
550        let projection = ParquetReadColumns::from_deduped(vec![
551            ParquetReadColumn::new(0).with_nested_paths(nested_paths),
552        ]);
553
554        let plan = build_projection_plan(&projection, &parquet_schema_desc);
555
556        assert_eq!(vec![true], plan.projected_root_presence);
557        assert_eq!(
558            ProjectionMask::leaves(&parquet_schema_desc, vec![0, 1]),
559            plan.mask
560        );
561    }
562
563    #[test]
564    fn test_nested_paths_fallback_to_variant_parent_by_default() {
565        let parquet_schema_desc = build_test_variant_parent_schema();
566        let projection =
567            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
568                vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]],
569            )]);
570
571        let plan = build_projection_plan(&projection, &parquet_schema_desc);
572
573        assert_eq!(vec![true], plan.projected_root_presence);
574        assert_eq!(
575            ProjectionMask::leaves(&parquet_schema_desc, vec![0]),
576            plan.mask
577        );
578    }
579
580    #[test]
581    fn test_non_variant_parent_does_not_fallback() {
582        let parquet_schema_desc = build_test_nested_parquet_schema();
583        let projection =
584            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
585                vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]],
586            )]);
587
588        let plan = build_projection_plan(&projection, &parquet_schema_desc);
589
590        assert_eq!(vec![false], plan.projected_root_presence);
591    }
592
593    #[test]
594    fn test_utf8_parent_does_not_fallback() {
595        let parquet_schema_desc = build_test_utf8_parent_schema();
596        let projection =
597            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
598                vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]],
599            )]);
600
601        let plan = build_projection_plan(&projection, &parquet_schema_desc);
602
603        assert_eq!(vec![false], plan.projected_root_presence);
604    }
605
606    #[test]
607    fn test_root_variant_does_not_fallback() {
608        let parquet_schema_desc = build_test_root_variant_schema();
609        let projection = ParquetReadColumns::from_deduped(vec![
610            ParquetReadColumn::new(0)
611                .with_nested_paths(vec![vec!["j".to_string(), "a".to_string()]]),
612        ]);
613
614        let plan = build_projection_plan(&projection, &parquet_schema_desc);
615
616        assert_eq!(vec![false], plan.projected_root_presence);
617    }
618
619    // Test schema:
620    // schema
621    // |- j
622    // |  |- a: INT64
623    // |  `- b
624    // |     |- c: INT64
625    // |     `- d: INT64
626    // `- k: INT64
627    fn build_test_nested_parquet_schema() -> SchemaDescriptor {
628        let leaf_a = Arc::new(
629            Type::primitive_type_builder("a", parquet::basic::Type::INT64)
630                .with_repetition(Repetition::REQUIRED)
631                .build()
632                .unwrap(),
633        );
634        let leaf_c = Arc::new(
635            Type::primitive_type_builder("c", parquet::basic::Type::INT64)
636                .with_repetition(Repetition::REQUIRED)
637                .build()
638                .unwrap(),
639        );
640        let leaf_d = Arc::new(
641            Type::primitive_type_builder("d", parquet::basic::Type::INT64)
642                .with_repetition(Repetition::REQUIRED)
643                .build()
644                .unwrap(),
645        );
646        let group_b = Arc::new(
647            Type::group_type_builder("b")
648                .with_repetition(Repetition::REQUIRED)
649                .with_fields(vec![leaf_c, leaf_d])
650                .build()
651                .unwrap(),
652        );
653        let root_j = Arc::new(
654            Type::group_type_builder("j")
655                .with_repetition(Repetition::REQUIRED)
656                .with_fields(vec![leaf_a, group_b])
657                .build()
658                .unwrap(),
659        );
660        let root_k = Arc::new(
661            Type::primitive_type_builder("k", parquet::basic::Type::INT64)
662                .with_repetition(Repetition::REQUIRED)
663                .build()
664                .unwrap(),
665        );
666        let schema = Arc::new(
667            Type::group_type_builder("schema")
668                .with_fields(vec![root_j, root_k])
669                .build()
670                .unwrap(),
671        );
672
673        SchemaDescriptor::new(schema)
674    }
675
676    // Test schema:
677    // schema
678    // `- j
679    //    |- a: BYTE_ARRAY
680    //    `- b
681    //       `- x: INT64
682    fn build_test_variant_parent_schema() -> SchemaDescriptor {
683        let leaf_a = Arc::new(
684            Type::primitive_type_builder("a", parquet::basic::Type::BYTE_ARRAY)
685                .with_repetition(Repetition::REQUIRED)
686                .build()
687                .unwrap(),
688        );
689        let leaf_x = Arc::new(
690            Type::primitive_type_builder("x", parquet::basic::Type::INT64)
691                .with_repetition(Repetition::REQUIRED)
692                .build()
693                .unwrap(),
694        );
695        let group_b = Arc::new(
696            Type::group_type_builder("b")
697                .with_repetition(Repetition::REQUIRED)
698                .with_fields(vec![leaf_x])
699                .build()
700                .unwrap(),
701        );
702        let root_j = Arc::new(
703            Type::group_type_builder("j")
704                .with_repetition(Repetition::REQUIRED)
705                .with_fields(vec![leaf_a, group_b])
706                .build()
707                .unwrap(),
708        );
709        let schema = Arc::new(
710            Type::group_type_builder("schema")
711                .with_fields(vec![root_j])
712                .build()
713                .unwrap(),
714        );
715
716        SchemaDescriptor::new(schema)
717    }
718
719    // Test schema:
720    // schema
721    // `- j
722    //    |- a: BYTE_ARRAY
723    //    `- b: BYTE_ARRAY
724    fn build_test_two_variant_parents_schema() -> SchemaDescriptor {
725        let leaf_a = Arc::new(
726            Type::primitive_type_builder("a", parquet::basic::Type::BYTE_ARRAY)
727                .with_repetition(Repetition::REQUIRED)
728                .build()
729                .unwrap(),
730        );
731        let leaf_b = Arc::new(
732            Type::primitive_type_builder("b", parquet::basic::Type::BYTE_ARRAY)
733                .with_repetition(Repetition::REQUIRED)
734                .build()
735                .unwrap(),
736        );
737        let root_j = Arc::new(
738            Type::group_type_builder("j")
739                .with_repetition(Repetition::REQUIRED)
740                .with_fields(vec![leaf_a, leaf_b])
741                .build()
742                .unwrap(),
743        );
744        let schema = Arc::new(
745            Type::group_type_builder("schema")
746                .with_fields(vec![root_j])
747                .build()
748                .unwrap(),
749        );
750
751        SchemaDescriptor::new(schema)
752    }
753
754    fn build_test_utf8_parent_schema() -> SchemaDescriptor {
755        let leaf_a = Arc::new(
756            Type::primitive_type_builder("a", parquet::basic::Type::BYTE_ARRAY)
757                .with_repetition(Repetition::REQUIRED)
758                .with_logical_type(Some(LogicalType::String))
759                .with_converted_type(ConvertedType::UTF8)
760                .build()
761                .unwrap(),
762        );
763        let root_j = Arc::new(
764            Type::group_type_builder("j")
765                .with_repetition(Repetition::REQUIRED)
766                .with_fields(vec![leaf_a])
767                .build()
768                .unwrap(),
769        );
770        let schema = Arc::new(
771            Type::group_type_builder("schema")
772                .with_fields(vec![root_j])
773                .build()
774                .unwrap(),
775        );
776
777        SchemaDescriptor::new(schema)
778    }
779
780    // Test schema:
781    // schema
782    // `- j: BYTE_ARRAY
783    fn build_test_root_variant_schema() -> SchemaDescriptor {
784        let root_j = Arc::new(
785            Type::primitive_type_builder("j", parquet::basic::Type::BYTE_ARRAY)
786                .with_repetition(Repetition::REQUIRED)
787                .build()
788                .unwrap(),
789        );
790        let schema = Arc::new(
791            Type::group_type_builder("schema")
792                .with_fields(vec![root_j])
793                .build()
794                .unwrap(),
795        );
796
797        SchemaDescriptor::new(schema)
798    }
799}