fast-path parquet root projection without nested fields

Author: fengys1996

src/mito2/src/sst/parquet/read_columns.rs

@@ -23,6 +23,7 @@ pub type ParquetNestedPath = Vec<String>;
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ParquetReadColumns {
     cols: Vec<ParquetReadColumn>,
+    has_nested: bool,
 }
 
 impl ParquetReadColumns {
@@ -35,12 +36,23 @@ impl ParquetReadColumns {
             .into_iter()
             .map(ParquetReadColumn::new)
             .collect();
-        Self { cols }
+        Self {
+            cols,
+            has_nested: false,
+        }
     }
 
     pub fn columns(&self) -> &[ParquetReadColumn] {
         &self.cols
     }
+
+    pub fn has_nested(&self) -> bool {
+        self.has_nested
+    }
+
+    pub fn root_indices_iter(&self) -> impl Iterator<Item = usize> + '_ {
+        self.cols.iter().map(|col| col.root_index)
+    }
 }
 
 /// Read requirement for a single parquet root column.
@@ -141,6 +153,7 @@ mod tests {
                 root_index: 0,
                 nested_paths: vec![],
             }],
+            has_nested: false,
         };
 
         assert_eq!(
@@ -164,6 +177,7 @@ mod tests {
                     nested_paths: vec![],
                 },
             ],
+            has_nested: true,
         };
 
         assert_eq!(
@@ -181,6 +195,7 @@ mod tests {
                 root_index: 0,
                 nested_paths: vec![vec!["j".to_string(), "b".to_string()]],
             }],
+            has_nested: true,
         };
 
         assert_eq!(
@@ -198,6 +213,7 @@ mod tests {
                 root_index: 0,
                 nested_paths: vec![vec!["j".to_string(), "b".to_string(), "c".to_string()]],
             }],
+            has_nested: true,
         };
 
         assert_eq!(
@@ -218,6 +234,7 @@ mod tests {
                     vec!["j".to_string(), "b".to_string(), "d".to_string()],
                 ],
             }],
+            has_nested: true,
         };
 
         assert_eq!(

src/mito2/src/sst/parquet/reader.rs

@@ -424,9 +424,14 @@ impl ParquetReaderBuilder {
         let parquet_projection = ParquetReadColumns::from_deduped_root_indices(
             read_format.projection_indices().iter().copied(),
         );
-        let leaf_indices = build_parquet_leaves_indices(parquet_schema_desc, &parquet_projection);
-        let projection_mask =
-            ProjectionMask::leaves(parquet_schema_desc, leaf_indices.iter().copied());
+
+        let projection_mask = if parquet_projection.has_nested() {
+            let leaf_indices =
+                build_parquet_leaves_indices(parquet_schema_desc, &parquet_projection);
+            ProjectionMask::leaves(parquet_schema_desc, leaf_indices.iter().copied())
+        } else {
+            ProjectionMask::roots(parquet_schema_desc, parquet_projection.root_indices_iter())
+        };
         let selection = self
             .row_groups_to_read(&read_format, &parquet_meta, &mut metrics.filter_metrics)
             .await;