1use std::collections::{HashMap, HashSet};
16
17use parquet::arrow::ProjectionMask;
18use parquet::basic::{ConvertedType, Type as PhysicalType};
19use parquet::schema::types::{ColumnDescriptor, SchemaDescriptor};
20
21pub type ParquetNestedPath = Vec<String>;
23
24#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct ParquetReadColumns {
27 root_indices: Vec<usize>,
33 cols: Vec<ParquetReadColumn>,
34 has_nested: bool,
35}
36
37impl ParquetReadColumns {
38 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 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 pub fn root_indices(&self) -> &[usize] {
86 &self.root_indices
87 }
88}
89
90#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct ParquetReadColumn {
103 root_index: usize,
105 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 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 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#[derive(Clone)]
151pub struct ProjectionMaskPlan {
152 pub mask: ProjectionMask,
154 pub projected_root_presence: Vec<bool>,
164}
165
166pub 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
211fn 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 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 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 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 if nested_path.len() <= 1 {
297 return None;
298 }
299
300 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 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 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 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 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 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}