Prune columns of row groups.

And fix clippy.

Author: RinChanNOWWW

src/query/storages/parquet/src/parquet_part.rs

@@ -14,6 +14,7 @@
 
 use std::any::Any;
 use std::collections::hash_map::DefaultHasher;
+use std::collections::HashMap;
 use std::hash::Hash;
 use std::hash::Hasher;
 use std::sync::Arc;
@@ -91,9 +92,9 @@ impl From<Compression> for ParquetCompression {
     }
 }
 
-impl Into<Compression> for ParquetCompression {
-    fn into(self) -> Compression {
-        match self {
+impl From<ParquetCompression> for Compression {
+    fn from(value: ParquetCompression) -> Self {
+        match value {
             ParquetCompression::Uncompressed => Compression::Uncompressed,
             ParquetCompression::Snappy => Compression::Snappy,
             ParquetCompression::Gzip => Compression::Gzip,
@@ -117,7 +118,7 @@ pub struct ColumnMeta {
 pub struct ParquetRowGroupPart {
     pub location: String,
     pub num_rows: usize,
-    pub column_metas: Vec<ColumnMeta>,
+    pub column_metas: HashMap<usize, ColumnMeta>,
 }
 
 #[typetag::serde(name = "parquet_row_group")]
@@ -144,7 +145,7 @@ impl ParquetRowGroupPart {
     pub fn create(
         location: String,
         num_rows: usize,
-        column_metas: Vec<ColumnMeta>,
+        column_metas: HashMap<usize, ColumnMeta>,
     ) -> Arc<Box<dyn PartInfo>> {
         Arc::new(Box::new(ParquetRowGroupPart {
             location,

src/query/storages/parquet/src/parquet_reader/deserialize.rs

@@ -50,7 +50,7 @@ impl ParquetReader {
             let mut metas = Vec::with_capacity(indices.len());
             let mut chunks = Vec::with_capacity(indices.len());
             for index in indices {
-                let column_meta = &part.column_metas[*index];
+                let column_meta = &part.column_metas[index];
                 let cnt = cnt_map.get_mut(index).unwrap();
                 *cnt -= 1;
                 let column_chunk = if cnt > &mut 0 {

src/query/storages/parquet/src/parquet_reader/meta.rs

@@ -16,7 +16,6 @@ use std::fs::File;
 
 use common_arrow::arrow::io::parquet::read as pread;
 use common_arrow::parquet::metadata::FileMetaData;
-use common_datavalues::DataField;
 use common_datavalues::DataSchema;
 use common_exception::ErrorCode;
 use common_exception::Result;
@@ -38,18 +37,10 @@ impl ParquetReader {
 
     #[inline]
     pub fn infer_schema(meta: &FileMetaData) -> Result<DataSchema> {
-        // Do not use `pread::infer_schema(meta)` becuase it will use metadata `ARROW:schema`.
-        // There maybe dictionary types in the schema, which is not supported by Databend.
-        // So we need to convert the primitive schema directly.
-        let field = pread::schema::parquet_to_arrow_schema(meta.schema().fields())
-            .into_iter()
-            .map(|mut f| {
-                // Need to change all the field name to lowercase.
-                f.name = f.name.to_lowercase();
-                DataField::from(&f)
-            })
-            .collect::<Vec<_>>();
-
-        Ok(DataSchema::new(field))
+        let mut arrow_schema = pread::infer_schema(meta)?;
+        arrow_schema.fields.iter_mut().for_each(|f| {
+            f.name = f.name.to_lowercase();
+        });
+        Ok(DataSchema::from(arrow_schema))
     }
 }

src/query/storages/parquet/src/parquet_reader/mod.rs

@@ -28,6 +28,10 @@ use common_storage::ColumnLeaves;
 use opendal::Operator;
 pub use read::IndexedChunk;
 
+/// The reader to parquet files with a projected schema.
+///
+/// **ALERT**: dictionary type is not supported yet.
+/// If there are dictionary pages in the parquet file, the reading process may fail.
 #[derive(Clone)]
 pub struct ParquetReader {
     operator: Operator,
@@ -62,12 +66,7 @@ impl ParquetReader {
         let schema_descriptors = to_parquet_schema(&arrow_schema)?;
 
         // Project schema
-        let projected_schema = match projection {
-            Projection::Columns(ref indices) => DataSchemaRef::new(schema.project(indices)),
-            Projection::InnerColumns(ref path_indices) => {
-                DataSchemaRef::new(schema.inner_project(path_indices))
-            }
-        };
+        let projected_schema = DataSchemaRef::new(projection.project_schema(&schema));
         // Project column leaves
         let projected_column_leaves = ColumnLeaves {
             column_leaves: projection
@@ -101,4 +100,8 @@ impl ParquetReader {
     pub fn schema(&self) -> DataSchemaRef {
         self.projected_schema.clone()
     }
+
+    pub fn columns_to_read(&self) -> &HashSet<usize> {
+        &self.columns_to_read
+    }
 }

src/query/storages/parquet/src/parquet_reader/read.rs

@@ -26,7 +26,7 @@ impl ParquetReader {
         let mut chunks = Vec::with_capacity(self.columns_to_read.len());
 
         for index in &self.columns_to_read {
-            let meta = &part.column_metas[*index];
+            let meta = &part.column_metas[index];
             let op = self.operator.clone();
             let chunk =
                 Self::sync_read_column(op.object(&part.location), meta.offset, meta.length)?;

src/query/storages/parquet/src/parquet_source.rs

@@ -248,7 +248,7 @@ impl Processor for ParquetSource {
         match std::mem::replace(&mut self.state, State::Finish) {
             State::ReadDataPrewhere(Some(part)) => {
                 let rg_part = ParquetRowGroupPart::from_part(&part)?;
-                let chunks = self.prewhere_reader.sync_read_columns_data(&rg_part)?;
+                let chunks = self.prewhere_reader.sync_read_columns_data(rg_part)?;
                 if self.prewhere_filter.is_some() {
                     self.state = State::PrewhereFilter(part, chunks);
                 } else {

src/query/storages/parquet/src/table_function/read.rs

@@ -12,6 +12,7 @@
 //  See the License for the specific language governing permissions and
 //  limitations under the License.
 
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use common_catalog::plan::DataSourcePlan;
@@ -126,24 +127,33 @@ impl ParquetTable {
                 part.location.clone()
             })
             .collect::<Vec<_>>();
+
+        let columns_to_read =
+            PushDownInfo::projection_of_push_downs(&plan.source_info.schema(), &plan.push_downs);
+        let max_io_requests = self.adjust_io_request(&ctx, columns_to_read.len())?;
         let ctx_ref = ctx.clone();
+        // `dummy_reader` is only used for prune columns in row groups.
+        let dummy_reader = ParquetReader::create(
+            self.operator.clone(),
+            plan.source_info.schema(),
+            columns_to_read,
+        )?;
         pipeline.set_on_init(move || {
             let mut partitions = Vec::with_capacity(locations.len());
             for location in &locations {
-                let file_meta = ParquetReader::read_meta(&location)?;
+                let file_meta = ParquetReader::read_meta(location)?;
                 for rg in &file_meta.row_groups {
-                    let column_metas = rg
-                        .columns()
-                        .iter()
-                        .map(|c| {
-                            let (offset, length) = c.byte_range();
-                            ColumnMeta {
-                                offset,
-                                length,
-                                compression: c.compression().into(),
-                            }
-                        })
-                        .collect();
+                    let mut column_metas =
+                        HashMap::with_capacity(dummy_reader.columns_to_read().len());
+                    for index in dummy_reader.columns_to_read() {
+                        let c = &rg.columns()[*index];
+                        let (offset, length) = c.byte_range();
+                        column_metas.insert(*index, ColumnMeta {
+                            offset,
+                            length,
+                            compression: c.compression().into(),
+                        });
+                    }
 
                     partitions.push(ParquetRowGroupPart::create(
                         location.clone(),
@@ -157,17 +167,13 @@ impl ParquetTable {
             Ok(())
         });
 
-        let columns_to_read =
-            PushDownInfo::projection_of_push_downs(&plan.schema(), &plan.push_downs);
-        let max_io_requests = self.adjust_io_request(&ctx, columns_to_read.len())?;
-
         // If there is a `PrewhereInfo`, the final output should be `PrehwereInfo.output_columns`.
         // `PrewhereInfo.output_columns` should be a subset of `PushDownInfo.projection`.
         let output_projection = match PushDownInfo::prewhere_of_push_downs(&plan.push_downs) {
             None => {
                 PushDownInfo::projection_of_push_downs(&self.table_info.schema(), &plan.push_downs)
             }
-            Some(v) => v.output_columns.clone(),
+            Some(v) => v.output_columns,
         };
         let output_schema = Arc::new(output_projection.project_schema(&plan.source_info.schema()));