Convert ArrowSchema to DataSchema only when convert ArrowChunk to DataBlock.

Author: RinChanNOWWW

src/common/arrow/src/lib.rs

@@ -16,6 +16,7 @@
 
 mod parquet_read;
 mod parquet_write;
+pub mod schema_projection;
 
 pub use arrow;
 pub use arrow_format;

src/common/arrow/src/schema_projection.rs

@@ -0,0 +1,50 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::BTreeMap;
+
+use arrow::datatypes::DataType;
+use arrow::datatypes::Field;
+use arrow::datatypes::Schema;
+
+/// Project a [`Schema`] by picking the fields at the given indices.
+pub fn project(schema: &Schema, indices: &[usize]) -> Schema {
+    let fields = indices
+        .iter()
+        .map(|idx| schema.fields[*idx].clone())
+        .collect::<Vec<_>>();
+    Schema::with_metadata(fields.into(), schema.metadata.clone())
+}
+
+/// Project a [`Schema`] with inner columns by path.
+pub fn inner_project(schema: &Schema, path_indices: &BTreeMap<usize, Vec<usize>>) -> Schema {
+    let paths: Vec<Vec<usize>> = path_indices.values().cloned().collect();
+    let fields = paths
+        .iter()
+        .map(|path| traverse_paths(&schema.fields, path))
+        .collect::<Vec<_>>();
+    Schema::with_metadata(fields.into(), schema.metadata.clone())
+}
+
+fn traverse_paths(fields: &[Field], path: &[usize]) -> Field {
+    assert!(!path.is_empty());
+    let field = &fields[path[0]];
+    if path.len() == 1 {
+        return field.clone();
+    }
+    if let DataType::Struct(inner_fields) = field.data_type() {
+        return traverse_paths(inner_fields, &path[1..]);
+    }
+    unreachable!("Unable to get field paths. Fields: {:?}", fields);
+}

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

@@ -45,7 +45,7 @@ impl ParquetReader {
         filter: Option<Bitmap>,
     ) -> Result<DataBlock> {
         let mut chunk_map: HashMap<usize, Vec<u8>> = chunks.into_iter().collect();
-        let mut columns_array_iter = Vec::with_capacity(self.projected_schema.num_fields());
+        let mut columns_array_iter = Vec::with_capacity(self.projected_arrow_schema.fields.len());
 
         let column_leaves = &self.projected_column_leaves.column_leaves;
         let mut cnt_map = Self::build_projection_count_map(column_leaves);
@@ -201,7 +201,7 @@ impl ParquetReader {
                 "deserializer from row group: fail to get a chunk",
             )),
             Some(Err(cause)) => Err(ErrorCode::from(cause)),
-            Some(Ok(chunk)) => DataBlock::from_chunk(&self.projected_schema, &chunk),
+            Some(Ok(chunk)) => DataBlock::from_chunk(&self.output_schema, &chunk),
         }
     }
 

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

@@ -14,9 +14,9 @@
 
 use std::fs::File;
 
+use common_arrow::arrow::datatypes::Schema as ArrowSchema;
 use common_arrow::arrow::io::parquet::read as pread;
 use common_arrow::parquet::metadata::FileMetaData;
-use common_datavalues::DataSchema;
 use common_exception::ErrorCode;
 use common_exception::Result;
 
@@ -36,11 +36,11 @@ impl ParquetReader {
     }
 
     #[inline]
-    pub fn infer_schema(meta: &FileMetaData) -> Result<DataSchema> {
+    pub fn infer_schema(meta: &FileMetaData) -> Result<ArrowSchema> {
         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))
+        Ok(arrow_schema)
     }
 }

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

@@ -21,9 +21,12 @@ use std::collections::HashMap;
 use std::collections::HashSet;
 use std::sync::Arc;
 
+use common_arrow::arrow::datatypes::Schema as ArrowSchema;
 use common_arrow::arrow::io::parquet::write::to_parquet_schema;
 use common_arrow::parquet::metadata::ColumnDescriptor;
+use common_arrow::schema_projection as ap;
 use common_catalog::plan::Projection;
+use common_datavalues::DataSchema;
 use common_datavalues::DataSchemaRef;
 use common_exception::Result;
 use common_storage::ColumnLeaves;
@@ -51,7 +54,13 @@ pub struct ParquetReader {
     /// ```
     columns_to_read: HashSet<usize>,
     /// The schema of the [`common_datablocks::DataBlock`] this reader produces.
-    projected_schema: DataSchemaRef,
+    output_schema: DataSchemaRef,
+    /// The actual schema used to read parquet. It will be converted to [`common_datavalues::DataSchema`] when output [`common_datablocks::DataBlock`].
+    ///
+    /// The reason of using [`ArrowSchema`] to read parquet is that
+    /// There are some types that Databend not support such as Timestmap of nanoseconds.
+    /// Such types will be convert to supported types after deserialization.
+    projected_arrow_schema: ArrowSchema,
     /// [`ColumnLeaves`] corresponding to the `projected_schema`.
     projected_column_leaves: ColumnLeaves,
     /// [`ColumnDescriptor`]s corresponding to the `projected_schema`.
@@ -61,16 +70,53 @@ pub struct ParquetReader {
 impl ParquetReader {
     pub fn create(
         operator: Operator,
-        schema: DataSchemaRef,
+        schema: ArrowSchema,
         projection: Projection,
     ) -> Result<Arc<ParquetReader>> {
-        // Full schema and column leaves.
-        let arrow_schema = schema.to_arrow();
-        let column_leaves = ColumnLeaves::new_from_schema(&arrow_schema);
-        let schema_descriptors = to_parquet_schema(&arrow_schema)?;
+        let (
+            projected_arrow_schema,
+            projected_column_leaves,
+            projected_column_descriptors,
+            columns_to_read,
+        ) = Self::do_projection(&schema, &projection)?;
 
+        Ok(Arc::new(ParquetReader {
+            operator,
+            columns_to_read,
+            output_schema: Arc::new(DataSchema::from(&projected_arrow_schema)),
+            projected_arrow_schema,
+            projected_column_leaves,
+            projected_column_descriptors,
+        }))
+    }
+
+    pub fn output_schema(&self) -> DataSchemaRef {
+        self.output_schema.clone()
+    }
+
+    pub fn columns_to_read(&self) -> &HashSet<usize> {
+        &self.columns_to_read
+    }
+
+    /// Project the schema and get the needed column leaves.
+    #[allow(clippy::type_complexity)]
+    pub fn do_projection(
+        schema: &ArrowSchema,
+        projection: &Projection,
+    ) -> Result<(
+        ArrowSchema,
+        ColumnLeaves,
+        HashMap<usize, ColumnDescriptor>,
+        HashSet<usize>,
+    )> {
+        // Full schema and column leaves.
+        let column_leaves = ColumnLeaves::new_from_schema(schema);
+        let schema_descriptors = to_parquet_schema(schema)?;
         // Project schema
-        let projected_schema = DataSchemaRef::new(projection.project_schema(&schema));
+        let projected_arrow_schema = match projection {
+            Projection::Columns(indices) => ap::project(schema, indices),
+            Projection::InnerColumns(path_indices) => ap::inner_project(schema, path_indices),
+        };
         // Project column leaves
         let projected_column_leaves = ColumnLeaves {
             column_leaves: projection
@@ -91,21 +137,11 @@ impl ParquetReader {
                     .insert(*index, schema_descriptors.columns()[*index].clone());
             }
         }
-
-        Ok(Arc::new(ParquetReader {
-            operator,
-            columns_to_read,
-            projected_schema,
+        Ok((
+            projected_arrow_schema,
             projected_column_leaves,
             projected_column_descriptors,
-        }))
-    }
-
-    pub fn schema(&self) -> DataSchemaRef {
-        self.projected_schema.clone()
-    }
-
-    pub fn columns_to_read(&self) -> &HashSet<usize> {
-        &self.columns_to_read
+            columns_to_read,
+        ))
     }
 }

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

@@ -37,8 +37,7 @@ use crate::ParquetSource;
 
 impl ParquetTable {
     pub fn create_reader(&self, projection: Projection) -> Result<Arc<ParquetReader>> {
-        let table_schema = self.table_info.schema();
-        ParquetReader::create(self.operator.clone(), table_schema, projection)
+        ParquetReader::create(self.operator.clone(), self.arrow_schema.clone(), projection)
     }
 
     // Build the prewhere reader.
@@ -133,19 +132,15 @@ impl ParquetTable {
         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,
-        )?;
+        let (_, _, _, columns_to_read) =
+            ParquetReader::do_projection(&plan.source_info.schema().to_arrow(), &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)?;
                 for rg in &file_meta.row_groups {
-                    let mut column_metas =
-                        HashMap::with_capacity(dummy_reader.columns_to_read().len());
-                    for index in dummy_reader.columns_to_read() {
+                    let mut column_metas = HashMap::with_capacity(columns_to_read.len());
+                    for index in &columns_to_read {
                         let c = &rg.columns()[*index];
                         let (offset, length) = c.byte_range();
                         column_metas.insert(*index, ColumnMeta {
@@ -178,8 +173,11 @@ impl ParquetTable {
         let output_schema = Arc::new(output_projection.project_schema(&plan.source_info.schema()));
 
         let prewhere_reader = self.build_prewhere_reader(plan)?;
-        let prewhere_filter =
-            self.build_prewhere_filter_executor(ctx.clone(), plan, prewhere_reader.schema())?;
+        let prewhere_filter = self.build_prewhere_filter_executor(
+            ctx.clone(),
+            plan,
+            prewhere_reader.output_schema(),
+        )?;
         let remain_reader = self.build_remain_reader(plan)?;
 
         // Add source pipe.

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

@@ -18,6 +18,7 @@ use std::sync::Arc;
 use chrono::NaiveDateTime;
 use chrono::TimeZone;
 use chrono::Utc;
+use common_arrow::arrow::datatypes::Schema as ArrowSchema;
 use common_catalog::plan::DataSourcePlan;
 use common_catalog::plan::PartStatistics;
 use common_catalog::plan::Partitions;
@@ -27,6 +28,7 @@ use common_catalog::table::Table;
 use common_catalog::table_args::TableArgs;
 use common_catalog::table_function::TableFunction;
 use common_config::GlobalConfig;
+use common_datavalues::DataSchema;
 use common_datavalues::DataValue;
 use common_exception::ErrorCode;
 use common_exception::Result;
@@ -45,6 +47,7 @@ pub struct ParquetTable {
 
     file_locations: Vec<String>,
     pub(super) table_info: TableInfo,
+    pub(super) arrow_schema: ArrowSchema,
     pub(super) operator: Operator,
 }
 
@@ -105,14 +108,14 @@ impl ParquetTable {
         // Assume all parquet files have the same schema.
         // If not, throw error during reading.
         let first_meta = ParquetReader::read_meta(&file_locations[0])?;
-        let schema = ParquetReader::infer_schema(&first_meta)?;
+        let arrow_schema = ParquetReader::infer_schema(&first_meta)?;
 
         let table_info = TableInfo {
             ident: TableIdent::new(table_id, 0),
             desc: format!("'{}'.'{}'", database_name, table_func_name),
             name: table_func_name.to_string(),
             meta: TableMeta {
-                schema: Arc::new(schema),
+                schema: Arc::new(DataSchema::from(&arrow_schema)),
                 engine: "SystemReadParquet".to_string(),
                 // Assuming that created_on is unnecessary for function table,
                 // we could make created_on fixed to pass test_shuffle_action_try_into.
@@ -131,6 +134,7 @@ impl ParquetTable {
             table_args,
             file_locations,
             table_info,
+            arrow_schema,
             operator,
         }))
     }