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 datatypes::extension::json::JSON2_REMAINDER_FIELD_NAME;
18use parquet::arrow::ProjectionMask;
19use parquet::basic::{ConvertedType, Type as PhysicalType};
20use parquet::schema::types::{ColumnDescriptor, SchemaDescriptor};
21
22/// A nested field access path inside one parquet root column.
23pub type ParquetNestedPath = Vec<String>;
24
25/// The parquet columns to read.
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct ParquetReadColumns {
28    /// Root parquet column indices in the same order as `cols`.
29    ///
30    /// Most readers need these indices as a borrowed slice for Arrow schema
31    /// projection or parquet root-column projection. Keeping them here avoids
32    /// repeatedly collecting `cols.iter().map(|col| col.root_index)`.
33    root_indices: Vec<usize>,
34    cols: Vec<ParquetReadColumn>,
35    has_nested: bool,
36}
37
38impl ParquetReadColumns {
39    /// Builds parquet read columns from deduplicated, normalized input.
40    ///
41    /// `cols` must not contain duplicate root indices, and nested paths must
42    /// already be merged. Empty `nested_paths` means reading the whole root column.
43    ///
44    /// This constructor does not validate or merge input.
45    pub fn from_deduped(cols: Vec<ParquetReadColumn>) -> Self {
46        let has_nested = cols.iter().any(|col| !col.nested_paths.is_empty());
47        let root_indices = cols.iter().map(|col| col.root_index).collect();
48        Self {
49            root_indices,
50            cols,
51            has_nested,
52        }
53    }
54
55    /// Builds root-column projections from root indices that are already
56    /// deduplicated.
57    ///
58    /// Note: this constructor does not check for duplicates.
59    pub fn from_deduped_root_indices(root_indices: impl IntoIterator<Item = usize>) -> Self {
60        let root_indices = root_indices.into_iter().collect::<Vec<_>>();
61        let cols = root_indices
62            .iter()
63            .copied()
64            .map(ParquetReadColumn::new)
65            .collect();
66        Self {
67            root_indices,
68            cols,
69            has_nested: false,
70        }
71    }
72
73    pub fn columns(&self) -> &[ParquetReadColumn] {
74        &self.cols
75    }
76
77    pub fn has_nested(&self) -> bool {
78        self.has_nested
79    }
80
81    pub fn root_indices_iter(&self) -> impl Iterator<Item = usize> + '_ {
82        self.root_indices.iter().copied()
83    }
84
85    /// Returns root parquet column indices.
86    pub fn root_indices(&self) -> &[usize] {
87        &self.root_indices
88    }
89}
90
91/// Read requirement for a single parquet root column.
92///
93/// `root_index` identifies the root column in the parquet schema.
94///
95/// If `nested_paths` is empty, the whole root column is read. Otherwise, only
96/// leaves under the specified nested paths are read.
97///
98/// To construct a [`ParquetReadColumn`]:
99/// - `ParquetReadColumn::new(0)` reads the whole root column at index `0`.
100/// - `ParquetReadColumn::new(0).with_nested_paths(vec![vec!["j".into(), "b".into()]])`
101///   reads only leaves under `j.b`.
102#[derive(Debug, Clone, PartialEq, Eq)]
103pub struct ParquetReadColumn {
104    /// Root field index in the parquet schema.
105    root_index: usize,
106    /// Nested paths to read under this root column.
107    ///
108    /// Each path includes the root column itself. For example, for a root
109    /// column `j`, path `["j", "a", "b"]` refers to `j.a.b`.
110    ///
111    /// If empty, the whole root column is read.
112    nested_paths: Vec<ParquetNestedPath>,
113}
114
115impl ParquetReadColumn {
116    pub fn new(root_index: usize) -> Self {
117        Self {
118            root_index,
119            nested_paths: vec![],
120        }
121    }
122
123    pub fn with_nested_paths(self, nested_paths: Vec<ParquetNestedPath>) -> Self {
124        Self {
125            nested_paths,
126            ..self
127        }
128    }
129
130    /// Merges additional nested paths into this root column.
131    pub fn merge_nested_paths(&mut self, nested_paths: Vec<ParquetNestedPath>) {
132        let reads_whole_root = self.nested_paths.is_empty() || nested_paths.is_empty();
133        if reads_whole_root {
134            // Empty nested paths means reading the whole root column.
135            self.nested_paths = vec![];
136        } else {
137            self.nested_paths.extend(nested_paths);
138        }
139    }
140
141    pub fn root_index(&self) -> usize {
142        self.root_index
143    }
144
145    pub fn nested_paths(&self) -> &[ParquetNestedPath] {
146        &self.nested_paths
147    }
148}
149
150/// Projection plan built for a parquet file.
151#[derive(Clone)]
152pub struct ProjectionMaskPlan {
153    /// `mask` is the projection mask applied to the parquet reader.
154    pub mask: ProjectionMask,
155    /// A boolean mask in output schema order indicating whether each projected root
156    /// has data read from the current parquet file.
157    ///
158    /// - `true`: parquet reads either requested nested data or a fallback parent
159    ///   for this root.
160    /// - `false`: this root is absent and must be synthesized later.
161    ///
162    /// The length of `projected_root_presence` is always equal to the
163    /// number of fields in the output schema.
164    pub projected_root_presence: Vec<bool>,
165}
166
167/// Builds a projection mask plan for reading a parquet file.
168///
169/// `parquet_read_cols` defines the requested root columns and optional
170/// nested paths to read.
171///
172/// `parquet_schema_desc` is the schema descriptor of the current parquet
173/// file. It is used to resolve requested nested paths to actual leaf
174/// column indices.
175///
176/// See [`ProjectionMaskPlan`] for the returned value.
177///
178/// For example, if the query requests `j.a` and `k`, but the current
179/// parquet file only contains leaves under `j.b` and `k`, then the
180/// returned plan keeps `k` in the projection mask and marks `j` as
181/// not present in the output, so it can be synthesized during
182/// post-processing.
183pub(crate) fn build_projection_plan(
184    parquet_read_cols: &ParquetReadColumns,
185    parquet_schema_desc: &SchemaDescriptor,
186) -> ProjectionMaskPlan {
187    if !parquet_read_cols.has_nested() {
188        let mask =
189            ProjectionMask::roots(parquet_schema_desc, parquet_read_cols.root_indices_iter());
190        return ProjectionMaskPlan {
191            mask,
192            projected_root_presence: vec![true; parquet_read_cols.columns().len()],
193        };
194    }
195
196    let (matched_leaves, matched_roots) =
197        build_parquet_leaves_indices(parquet_schema_desc, parquet_read_cols);
198
199    let projected_root_presence = parquet_read_cols
200        .columns()
201        .iter()
202        .map(|col| matched_roots.contains(&col.root_index()))
203        .collect();
204
205    let mask = ProjectionMask::leaves(parquet_schema_desc, matched_leaves);
206    ProjectionMaskPlan {
207        mask,
208        projected_root_presence,
209    }
210}
211
212/// Builds parquet leaf-column indices for reading a parquet file.
213///
214/// Returns `(matched_leaves, matched_roots)`:
215/// - `matched_leaves`: matched leaf-column indices in the current parquet file schema.
216/// - `matched_roots`: root-field indices read from the current parquet file schema.
217fn build_parquet_leaves_indices(
218    parquet_schema_desc: &SchemaDescriptor,
219    projection: &ParquetReadColumns,
220) -> (Vec<usize>, HashSet<usize>) {
221    let mut map = HashMap::with_capacity(projection.cols.len());
222    for col in &projection.cols {
223        map.insert(col.root_index, col);
224    }
225
226    let mut matched_leaves = HashSet::new();
227    let mut matched_roots = HashSet::with_capacity(projection.cols.len());
228
229    // Tracks whether each requested nested path matched by prefix without fallback.
230    let mut prefix_matched = HashMap::<usize, Vec<bool>>::new();
231    for col in &projection.cols {
232        prefix_matched.insert(col.root_index, vec![false; col.nested_paths.len()]);
233    }
234
235    // First select root/nested prefix matches.
236    for (leaf_idx, leaf_col) in parquet_schema_desc.columns().iter().enumerate() {
237        let root_idx = parquet_schema_desc.get_column_root_idx(leaf_idx);
238        let Some(col) = map.get(&root_idx) else {
239            continue;
240        };
241        if col.nested_paths.is_empty() {
242            matched_leaves.insert(leaf_idx);
243            matched_roots.insert(root_idx);
244            continue;
245        }
246
247        let leaf_path = leaf_col.path().parts();
248        let mut matched = false;
249        for (path_idx, _) in col
250            .nested_paths
251            .iter()
252            .enumerate()
253            .filter(|(_, nested_path)| leaf_path.starts_with(nested_path))
254        {
255            prefix_matched.get_mut(&root_idx).unwrap()[path_idx] = true;
256            matched = true;
257        }
258
259        if matched {
260            matched_leaves.insert(leaf_idx);
261            matched_roots.insert(root_idx);
262        }
263    }
264
265    // Then include v2 remainder leaves or fallback prefix misses to their nearest variant parent.
266    // TODO(fys): Gate fallback planning on the root being JSON2. A raw Binary
267    // leaf is a JSONB variant only under a JSON2 root; plain struct Binary
268    // children should not enter this fallback path.
269    for col in &projection.cols {
270        let path_matches = &prefix_matched[&col.root_index];
271        let mut needs_remainder = false;
272        for (matched, nested_path) in path_matches.iter().zip(&col.nested_paths) {
273            if *matched {
274                if !needs_remainder {
275                    needs_remainder =
276                        path_points_to_struct(parquet_schema_desc, col.root_index, nested_path);
277                }
278                continue;
279            }
280
281            if let Some(leaf_idx) =
282                find_nearest_variant_parent(parquet_schema_desc, col.root_index, nested_path)
283            {
284                matched_leaves.insert(leaf_idx);
285                matched_roots.insert(col.root_index);
286            } else {
287                needs_remainder = true;
288            }
289        }
290
291        if needs_remainder {
292            let remainder_leaves = find_remainder_leaves(parquet_schema_desc, col.root_index);
293            if !remainder_leaves.is_empty() {
294                matched_leaves.extend(remainder_leaves);
295                matched_roots.insert(col.root_index);
296            }
297        }
298    }
299
300    let mut matched_leaves = matched_leaves.into_iter().collect::<Vec<_>>();
301    matched_leaves.sort_unstable();
302    (matched_leaves, matched_roots)
303}
304
305/// Returns whether a nested path points to an explicitly materialized object.
306///
307/// JSON2 v2 can split an object's children between its Struct field and the remainder,
308/// so reading the Struct leaves alone may produce an incomplete object.
309fn path_points_to_struct(
310    parquet_schema_desc: &SchemaDescriptor,
311    root_idx: usize,
312    path: &[String],
313) -> bool {
314    let Some(mut field) = parquet_schema_desc.root_schema().get_fields().get(root_idx) else {
315        return false;
316    };
317    for name in path.iter().skip(1) {
318        if !field.is_group() {
319            return false;
320        }
321        let Some(child) = field.get_fields().iter().find(|field| field.name() == name) else {
322            return false;
323        };
324        field = child;
325    }
326    field.is_group()
327}
328
329/// Finds the Parquet leaves backing a JSON2 v2 remainder field.
330///
331/// The remainder is a sibling of explicitly materialized fields, so prefix matching a
332/// requested path cannot find it. These leaves are needed when an explicit path is absent
333/// or an explicitly materialized object may have additional children in the remainder.
334fn find_remainder_leaves(parquet_schema_desc: &SchemaDescriptor, root_idx: usize) -> Vec<usize> {
335    parquet_schema_desc
336        .columns()
337        .iter()
338        .enumerate()
339        .filter_map(|(i, column)| {
340            let path = column.path().parts();
341            (parquet_schema_desc.get_column_root_idx(i) == root_idx
342                && path.get(1).is_some_and(|x| x == JSON2_REMAINDER_FIELD_NAME))
343            .then_some(i)
344        })
345        .collect::<Vec<_>>()
346}
347
348fn find_nearest_variant_parent(
349    parquet_schema_desc: &SchemaDescriptor,
350    root_idx: usize,
351    nested_path: &[String],
352) -> Option<usize> {
353    // TODO(fys): Build a variant path index if fallback lookup becomes hot.
354    if nested_path.len() <= 1 {
355        return None;
356    }
357
358    // JSON2 root columns are always structured fields. Fallback only applies to
359    // variant leaves below the root.
360    for parent_len in (2..nested_path.len()).rev() {
361        let parent_path = &nested_path[..parent_len];
362        for (leaf_idx, leaf_col) in parquet_schema_desc.columns().iter().enumerate() {
363            if parquet_schema_desc.get_column_root_idx(leaf_idx) != root_idx {
364                continue;
365            }
366            if leaf_col.path().parts() == parent_path && is_variant_leaf(leaf_col) {
367                return Some(leaf_idx);
368            }
369        }
370    }
371
372    None
373}
374
375fn is_variant_leaf(leaf_col: &ColumnDescriptor) -> bool {
376    // TODO(fys): Recognize JSON2 variant type from Arrow extension metadata
377    // if the parquet Arrow schema preserves it for nested fields.
378    matches!(
379        leaf_col.physical_type(),
380        PhysicalType::BYTE_ARRAY | PhysicalType::FIXED_LEN_BYTE_ARRAY
381    ) && leaf_col.logical_type_ref().is_none()
382        && leaf_col.converted_type() == ConvertedType::NONE
383}
384
385#[cfg(test)]
386mod tests {
387    use std::sync::Arc;
388
389    use parquet::basic::{ConvertedType, LogicalType, Repetition};
390    use parquet::errors::ParquetError;
391    use parquet::schema::types::Type;
392
393    use super::*;
394
395    #[test]
396    fn test_build_projection_mask_without_nested_paths() {
397        let parquet_schema_desc = build_test_nested_parquet_schema();
398        let projection = ParquetReadColumns::from_deduped_root_indices([0, 1]);
399
400        let plan = build_projection_plan(&projection, &parquet_schema_desc);
401
402        assert_eq!(vec![true, true], plan.projected_root_presence);
403        assert_eq!(
404            ProjectionMask::roots(&parquet_schema_desc, [0, 1]),
405            plan.mask
406        );
407    }
408
409    #[test]
410    fn test_reads_whole_root() {
411        let parquet_schema_desc = build_test_nested_parquet_schema();
412
413        let projection = ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0)]);
414
415        let (matched_leaves, matched_roots) =
416            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
417        assert_eq!(vec![0, 1, 2], matched_leaves);
418        assert_eq!(HashSet::from([0]), matched_roots);
419    }
420
421    #[test]
422    fn test_filters_nested_paths() {
423        let parquet_schema_desc = build_test_nested_parquet_schema();
424
425        let projection = ParquetReadColumns::from_deduped(vec![
426            ParquetReadColumn::new(0)
427                .with_nested_paths(vec![vec!["j".to_string(), "b".to_string()]]),
428            ParquetReadColumn::new(1),
429        ]);
430
431        let (matched_leaves, matched_roots) =
432            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
433        assert_eq!(vec![1, 2, 3], matched_leaves);
434        assert_eq!(HashSet::from([0, 1]), matched_roots);
435    }
436
437    #[test]
438    fn test_reads_middle_level_path() {
439        let parquet_schema_desc = build_test_nested_parquet_schema();
440
441        let projection = ParquetReadColumns::from_deduped(vec![
442            ParquetReadColumn::new(0)
443                .with_nested_paths(vec![vec!["j".to_string(), "b".to_string()]]),
444        ]);
445
446        let (matched_leaves, matched_roots) =
447            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
448        assert_eq!(vec![1, 2], matched_leaves);
449        assert_eq!(HashSet::from([0]), matched_roots);
450    }
451
452    #[test]
453    fn test_parent_path_covers_redundant_child_path() {
454        let parquet_schema_desc = build_test_nested_parquet_schema();
455        let nested_paths = vec![
456            vec!["j".to_string(), "b".to_string()],
457            vec!["j".to_string(), "b".to_string(), "c".to_string()],
458        ];
459
460        let read_column = ParquetReadColumn::new(0).with_nested_paths(nested_paths);
461        let projection = ParquetReadColumns::from_deduped(vec![read_column]);
462
463        let (matched_leaves, matched_roots) =
464            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
465        assert_eq!(vec![1, 2], matched_leaves);
466        assert_eq!(HashSet::from([0]), matched_roots);
467    }
468
469    #[test]
470    fn test_reads_leaf_level_path() {
471        let parquet_schema_desc = build_test_nested_parquet_schema();
472
473        let projection =
474            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
475                vec![vec!["j".to_string(), "b".to_string(), "c".to_string()]],
476            )]);
477
478        let (matched_leaves, matched_roots) =
479            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
480        assert_eq!(vec![1], matched_leaves);
481        assert_eq!(HashSet::from([0]), matched_roots);
482    }
483
484    #[test]
485    fn test_build_projection_mask_with_unmatched_roots() {
486        let parquet_schema_desc = build_test_nested_parquet_schema();
487
488        let projection = ParquetReadColumns::from_deduped(vec![
489            ParquetReadColumn::new(0)
490                .with_nested_paths(vec![vec!["j".to_string(), "missing".to_string()]]),
491            ParquetReadColumn::new(1),
492        ]);
493
494        let plan = build_projection_plan(&projection, &parquet_schema_desc);
495
496        assert_eq!(vec![false, true], plan.projected_root_presence);
497        assert_eq!(
498            ProjectionMask::leaves(&parquet_schema_desc, vec![3]),
499            plan.mask
500        );
501    }
502
503    #[test]
504    fn test_v2_routes_missing_path_to_remainder() -> Result<(), ParquetError> {
505        let parquet = build_test_v2_schema()?;
506        let projection =
507            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
508                vec![
509                    vec!["j".to_string(), "cold".to_string()],
510                    vec!["j".to_string(), "another".to_string()],
511                ],
512            )]);
513
514        let plan = build_projection_plan(&projection, &parquet);
515
516        assert_eq!(vec![true], plan.projected_root_presence);
517        assert_eq!(ProjectionMask::leaves(&parquet, [0, 1]), plan.mask);
518        Ok(())
519    }
520
521    #[test]
522    fn test_v2_explicit_path_does_not_read_remainder() -> Result<(), ParquetError> {
523        let parquet = build_test_v2_schema()?;
524        let projection = ParquetReadColumns::from_deduped(vec![
525            ParquetReadColumn::new(0)
526                .with_nested_paths(vec![vec!["j".to_string(), "hot".to_string()]]),
527        ]);
528
529        let plan = build_projection_plan(&projection, &parquet);
530
531        assert_eq!(vec![true], plan.projected_root_presence);
532        assert_eq!(ProjectionMask::leaves(&parquet, [3]), plan.mask);
533        Ok(())
534    }
535
536    #[test]
537    fn test_v2_container_path_reads_remainder() -> Result<(), ParquetError> {
538        let parquet = build_test_v2_schema()?;
539        let projection = ParquetReadColumns::from_deduped(vec![
540            ParquetReadColumn::new(0)
541                .with_nested_paths(vec![vec!["j".to_string(), "commit".to_string()]]),
542        ]);
543
544        let plan = build_projection_plan(&projection, &parquet);
545
546        assert_eq!(vec![true], plan.projected_root_presence);
547        assert_eq!(ProjectionMask::leaves(&parquet, [0, 1, 2]), plan.mask);
548        Ok(())
549    }
550
551    // A nested path under an explicit Variant is stored entirely in that Variant parent. The
552    // remainder may preserve `opaque: null`, but it cannot contain `opaque.leaf`, so reading the
553    // nearest Variant parent is sufficient.
554    #[test]
555    fn test_v2_variant_parent_path_reads_parent() -> Result<(), ParquetError> {
556        let parquet = build_test_v2_schema()?;
557        let projection =
558            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
559                vec![vec![
560                    "j".to_string(),
561                    "opaque".to_string(),
562                    "leaf".to_string(),
563                ]],
564            )]);
565
566        let plan = build_projection_plan(&projection, &parquet);
567
568        assert_eq!(vec![true], plan.projected_root_presence);
569        assert_eq!(ProjectionMask::leaves(&parquet, [4]), plan.mask);
570        Ok(())
571    }
572
573    #[test]
574    fn test_merges_mixed_paths() {
575        let parquet_schema_desc = build_test_nested_parquet_schema();
576
577        let projection =
578            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
579                vec![
580                    vec!["j".to_string(), "a".to_string()],
581                    vec!["j".to_string(), "b".to_string(), "d".to_string()],
582                ],
583            )]);
584
585        let (matched_leaves, matched_roots) =
586            build_parquet_leaves_indices(&parquet_schema_desc, &projection);
587        assert_eq!(vec![0, 2], matched_leaves);
588        assert_eq!(HashSet::from([0]), matched_roots);
589    }
590
591    #[test]
592    fn test_merge_nested_paths_extends_paths() {
593        let mut col = ParquetReadColumn::new(0)
594            .with_nested_paths(vec![vec!["j".to_string(), "a".to_string()]]);
595
596        col.merge_nested_paths(vec![vec!["j".to_string(), "b".to_string()]]);
597
598        assert_eq!(
599            &[
600                vec!["j".to_string(), "a".to_string()],
601                vec!["j".to_string(), "b".to_string()],
602            ],
603            col.nested_paths()
604        );
605    }
606
607    #[test]
608    fn test_merge_nested_paths_with_whole_root() {
609        let mut col = ParquetReadColumn::new(0)
610            .with_nested_paths(vec![vec!["j".to_string(), "a".to_string()]]);
611
612        col.merge_nested_paths(vec![]);
613
614        assert!(col.nested_paths().is_empty());
615    }
616
617    #[test]
618    fn test_fallback_to_nearest_variant_parent() {
619        let parquet_schema_desc = build_test_variant_parent_schema();
620        let projection =
621            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
622                vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]],
623            )]);
624
625        let plan = build_projection_plan(&projection, &parquet_schema_desc);
626
627        assert_eq!(vec![true], plan.projected_root_presence);
628        assert_eq!(
629            ProjectionMask::leaves(&parquet_schema_desc, vec![0]),
630            plan.mask
631        );
632    }
633
634    #[test]
635    fn test_prefix_match_prevents_variant_parent_fallback() {
636        let parquet_schema_desc = build_test_variant_parent_schema();
637        let projection =
638            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
639                vec![vec!["j".to_string(), "b".to_string(), "x".to_string()]],
640            )]);
641
642        let plan = build_projection_plan(&projection, &parquet_schema_desc);
643
644        assert_eq!(vec![true], plan.projected_root_presence);
645        assert_eq!(
646            ProjectionMask::leaves(&parquet_schema_desc, vec![1]),
647            plan.mask
648        );
649    }
650
651    #[test]
652    fn test_mixed_prefix_and_fallback_paths() {
653        let parquet_schema_desc = build_test_variant_parent_schema();
654        let nested_paths = vec![
655            vec!["j".to_string(), "a".to_string(), "x".to_string()],
656            vec!["j".to_string(), "b".to_string(), "x".to_string()],
657        ];
658        let projection = ParquetReadColumns::from_deduped(vec![
659            ParquetReadColumn::new(0).with_nested_paths(nested_paths),
660        ]);
661
662        let plan = build_projection_plan(&projection, &parquet_schema_desc);
663
664        assert_eq!(vec![true], plan.projected_root_presence);
665        assert_eq!(
666            ProjectionMask::leaves(&parquet_schema_desc, vec![0, 1]),
667            plan.mask
668        );
669    }
670
671    #[test]
672    fn test_fallback_selects_multiple_variant_parents() {
673        let parquet_schema_desc = build_test_two_variant_parents_schema();
674        let nested_paths = vec![
675            vec!["j".to_string(), "a".to_string(), "y".to_string()],
676            vec!["j".to_string(), "b".to_string(), "d".to_string()],
677            vec!["j".to_string(), "a".to_string(), "x".to_string()],
678        ];
679        let projection = ParquetReadColumns::from_deduped(vec![
680            ParquetReadColumn::new(0).with_nested_paths(nested_paths),
681        ]);
682
683        let plan = build_projection_plan(&projection, &parquet_schema_desc);
684
685        assert_eq!(vec![true], plan.projected_root_presence);
686        assert_eq!(
687            ProjectionMask::leaves(&parquet_schema_desc, vec![0, 1]),
688            plan.mask
689        );
690    }
691
692    #[test]
693    fn test_nested_paths_fallback_to_variant_parent_by_default() {
694        let parquet_schema_desc = build_test_variant_parent_schema();
695        let projection =
696            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
697                vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]],
698            )]);
699
700        let plan = build_projection_plan(&projection, &parquet_schema_desc);
701
702        assert_eq!(vec![true], plan.projected_root_presence);
703        assert_eq!(
704            ProjectionMask::leaves(&parquet_schema_desc, vec![0]),
705            plan.mask
706        );
707    }
708
709    #[test]
710    fn test_non_variant_parent_does_not_fallback() {
711        let parquet_schema_desc = build_test_nested_parquet_schema();
712        let projection =
713            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
714                vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]],
715            )]);
716
717        let plan = build_projection_plan(&projection, &parquet_schema_desc);
718
719        assert_eq!(vec![false], plan.projected_root_presence);
720    }
721
722    #[test]
723    fn test_utf8_parent_does_not_fallback() {
724        let parquet_schema_desc = build_test_utf8_parent_schema();
725        let projection =
726            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
727                vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]],
728            )]);
729
730        let plan = build_projection_plan(&projection, &parquet_schema_desc);
731
732        assert_eq!(vec![false], plan.projected_root_presence);
733    }
734
735    #[test]
736    fn test_root_variant_does_not_fallback() {
737        let parquet_schema_desc = build_test_root_variant_schema();
738        let projection = ParquetReadColumns::from_deduped(vec![
739            ParquetReadColumn::new(0)
740                .with_nested_paths(vec![vec!["j".to_string(), "a".to_string()]]),
741        ]);
742
743        let plan = build_projection_plan(&projection, &parquet_schema_desc);
744
745        assert_eq!(vec![false], plan.projected_root_presence);
746    }
747
748    // Test schema:
749    // schema
750    // |- j
751    // |  |- a: INT64
752    // |  `- b
753    // |     |- c: INT64
754    // |     `- d: INT64
755    // `- k: INT64
756    fn build_test_nested_parquet_schema() -> SchemaDescriptor {
757        let leaf_a = Arc::new(
758            Type::primitive_type_builder("a", parquet::basic::Type::INT64)
759                .with_repetition(Repetition::REQUIRED)
760                .build()
761                .unwrap(),
762        );
763        let leaf_c = Arc::new(
764            Type::primitive_type_builder("c", parquet::basic::Type::INT64)
765                .with_repetition(Repetition::REQUIRED)
766                .build()
767                .unwrap(),
768        );
769        let leaf_d = Arc::new(
770            Type::primitive_type_builder("d", parquet::basic::Type::INT64)
771                .with_repetition(Repetition::REQUIRED)
772                .build()
773                .unwrap(),
774        );
775        let group_b = Arc::new(
776            Type::group_type_builder("b")
777                .with_repetition(Repetition::REQUIRED)
778                .with_fields(vec![leaf_c, leaf_d])
779                .build()
780                .unwrap(),
781        );
782        let root_j = Arc::new(
783            Type::group_type_builder("j")
784                .with_repetition(Repetition::REQUIRED)
785                .with_fields(vec![leaf_a, group_b])
786                .build()
787                .unwrap(),
788        );
789        let root_k = Arc::new(
790            Type::primitive_type_builder("k", parquet::basic::Type::INT64)
791                .with_repetition(Repetition::REQUIRED)
792                .build()
793                .unwrap(),
794        );
795        let schema = Arc::new(
796            Type::group_type_builder("schema")
797                .with_fields(vec![root_j, root_k])
798                .build()
799                .unwrap(),
800        );
801
802        SchemaDescriptor::new(schema)
803    }
804
805    fn build_test_v2_schema() -> Result<SchemaDescriptor, ParquetError> {
806        let metadata = Arc::new(
807            Type::primitive_type_builder("metadata", parquet::basic::Type::BYTE_ARRAY)
808                .with_repetition(Repetition::REQUIRED)
809                .build()?,
810        );
811        let value = Arc::new(
812            Type::primitive_type_builder("value", parquet::basic::Type::BYTE_ARRAY)
813                .with_repetition(Repetition::REQUIRED)
814                .build()?,
815        );
816        let remainder = Arc::new(
817            Type::group_type_builder(JSON2_REMAINDER_FIELD_NAME)
818                .with_repetition(Repetition::OPTIONAL)
819                .with_logical_type(Some(LogicalType::Variant {
820                    specification_version: None,
821                }))
822                .with_fields(vec![metadata, value])
823                .build()?,
824        );
825        let operation = Arc::new(
826            Type::primitive_type_builder("operation", parquet::basic::Type::INT64)
827                .with_repetition(Repetition::OPTIONAL)
828                .build()?,
829        );
830        let commit = Arc::new(
831            Type::group_type_builder("commit")
832                .with_repetition(Repetition::OPTIONAL)
833                .with_fields(vec![operation])
834                .build()?,
835        );
836        let hot = Arc::new(
837            Type::primitive_type_builder("hot", parquet::basic::Type::INT64)
838                .with_repetition(Repetition::OPTIONAL)
839                .build()?,
840        );
841        // Normally there are no other explicit Variant fields exist if a remainder field is present.
842        // However, when structured values reach JSON2_MAX_STRUCTURED_DEPTH, there are. `opaque`
843        // models such a deep leaf without building a deeply nested test schema.
844        let opaque = Arc::new(
845            Type::primitive_type_builder("opaque", parquet::basic::Type::BYTE_ARRAY)
846                .with_repetition(Repetition::OPTIONAL)
847                .build()?,
848        );
849        let root = Arc::new(
850            Type::group_type_builder("j")
851                .with_repetition(Repetition::OPTIONAL)
852                .with_fields(vec![remainder, commit, hot, opaque])
853                .build()?,
854        );
855        Ok(SchemaDescriptor::new(Arc::new(
856            Type::group_type_builder("schema")
857                .with_fields(vec![root])
858                .build()?,
859        )))
860    }
861
862    // Test schema:
863    // schema
864    // `- j
865    //    |- a: BYTE_ARRAY
866    //    `- b
867    //       `- x: INT64
868    fn build_test_variant_parent_schema() -> SchemaDescriptor {
869        let leaf_a = Arc::new(
870            Type::primitive_type_builder("a", parquet::basic::Type::BYTE_ARRAY)
871                .with_repetition(Repetition::REQUIRED)
872                .build()
873                .unwrap(),
874        );
875        let leaf_x = Arc::new(
876            Type::primitive_type_builder("x", parquet::basic::Type::INT64)
877                .with_repetition(Repetition::REQUIRED)
878                .build()
879                .unwrap(),
880        );
881        let group_b = Arc::new(
882            Type::group_type_builder("b")
883                .with_repetition(Repetition::REQUIRED)
884                .with_fields(vec![leaf_x])
885                .build()
886                .unwrap(),
887        );
888        let root_j = Arc::new(
889            Type::group_type_builder("j")
890                .with_repetition(Repetition::REQUIRED)
891                .with_fields(vec![leaf_a, group_b])
892                .build()
893                .unwrap(),
894        );
895        let schema = Arc::new(
896            Type::group_type_builder("schema")
897                .with_fields(vec![root_j])
898                .build()
899                .unwrap(),
900        );
901
902        SchemaDescriptor::new(schema)
903    }
904
905    // Test schema:
906    // schema
907    // `- j
908    //    |- a: BYTE_ARRAY
909    //    `- b: BYTE_ARRAY
910    fn build_test_two_variant_parents_schema() -> SchemaDescriptor {
911        let leaf_a = Arc::new(
912            Type::primitive_type_builder("a", parquet::basic::Type::BYTE_ARRAY)
913                .with_repetition(Repetition::REQUIRED)
914                .build()
915                .unwrap(),
916        );
917        let leaf_b = Arc::new(
918            Type::primitive_type_builder("b", parquet::basic::Type::BYTE_ARRAY)
919                .with_repetition(Repetition::REQUIRED)
920                .build()
921                .unwrap(),
922        );
923        let root_j = Arc::new(
924            Type::group_type_builder("j")
925                .with_repetition(Repetition::REQUIRED)
926                .with_fields(vec![leaf_a, leaf_b])
927                .build()
928                .unwrap(),
929        );
930        let schema = Arc::new(
931            Type::group_type_builder("schema")
932                .with_fields(vec![root_j])
933                .build()
934                .unwrap(),
935        );
936
937        SchemaDescriptor::new(schema)
938    }
939
940    fn build_test_utf8_parent_schema() -> SchemaDescriptor {
941        let leaf_a = Arc::new(
942            Type::primitive_type_builder("a", parquet::basic::Type::BYTE_ARRAY)
943                .with_repetition(Repetition::REQUIRED)
944                .with_logical_type(Some(LogicalType::String))
945                .with_converted_type(ConvertedType::UTF8)
946                .build()
947                .unwrap(),
948        );
949        let root_j = Arc::new(
950            Type::group_type_builder("j")
951                .with_repetition(Repetition::REQUIRED)
952                .with_fields(vec![leaf_a])
953                .build()
954                .unwrap(),
955        );
956        let schema = Arc::new(
957            Type::group_type_builder("schema")
958                .with_fields(vec![root_j])
959                .build()
960                .unwrap(),
961        );
962
963        SchemaDescriptor::new(schema)
964    }
965
966    // Test schema:
967    // schema
968    // `- j: BYTE_ARRAY
969    fn build_test_root_variant_schema() -> SchemaDescriptor {
970        let root_j = Arc::new(
971            Type::primitive_type_builder("j", parquet::basic::Type::BYTE_ARRAY)
972                .with_repetition(Repetition::REQUIRED)
973                .build()
974                .unwrap(),
975        );
976        let schema = Arc::new(
977            Type::group_type_builder("schema")
978                .with_fields(vec![root_j])
979                .build()
980                .unwrap(),
981        );
982
983        SchemaDescriptor::new(schema)
984    }
985}
mito2/sst/parquet/read_columns.rs

mito2/sst/parquet/
read_columns.rs
