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

mito2/sst/parquet/
read_columns.rs
