Merge pull request #9080 from RinChanNOWWW/read-parquet

feat(query): new table function `read_parquet` to read parquet files as a table.

Author: mergify[bot]

Cargo.lock

@@ -58,6 +58,7 @@ members = [
     "src/query/storages/stage",
     "src/query/storages/system",
     "src/query/storages/view",
+    "src/query/storages/parquet",
     "src/query/users",
     # databend-query
     "src/query/service",

Cargo.toml

@@ -10,6 +10,7 @@ edition = { workspace = true }
 storage-hdfs = ["opendal/services-hdfs"]
 
 [dependencies]
+common-arrow = { path = "../arrow" }
 common-auth = { path = "../auth" }
 common-base = { path = "../base" }
 common-exception = { path = "../exception" }

src/common/storage/Cargo.toml

@@ -0,0 +1,100 @@
+// 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.
+
+//! This module provides data structures for build column indexes.
+//! It's used by Fuse Engine and Parquet Engine.
+
+use common_arrow::arrow::datatypes::DataType as ArrowType;
+use common_arrow::arrow::datatypes::Field as ArrowField;
+use common_arrow::arrow::datatypes::Schema as ArrowSchema;
+use common_exception::ErrorCode;
+use common_exception::Result;
+
+#[derive(Debug, Clone)]
+pub struct ColumnLeaves {
+    pub column_leaves: Vec<ColumnLeaf>,
+}
+
+impl ColumnLeaves {
+    pub fn new_from_schema(schema: &ArrowSchema) -> Self {
+        let mut leaf_id = 0;
+        let mut column_leaves = Vec::with_capacity(schema.fields.len());
+
+        for field in &schema.fields {
+            let column_leaf = Self::traverse_fields_dfs(field, &mut leaf_id);
+            column_leaves.push(column_leaf);
+        }
+
+        Self { column_leaves }
+    }
+
+    fn traverse_fields_dfs(field: &ArrowField, leaf_id: &mut usize) -> ColumnLeaf {
+        match &field.data_type {
+            ArrowType::Struct(inner_fields) => {
+                let mut child_column_leaves = Vec::with_capacity(inner_fields.len());
+                let mut child_leaf_ids = Vec::with_capacity(inner_fields.len());
+                for inner_field in inner_fields {
+                    let child_column_leaf = Self::traverse_fields_dfs(inner_field, leaf_id);
+                    child_leaf_ids.extend(child_column_leaf.leaf_ids.clone());
+                    child_column_leaves.push(child_column_leaf);
+                }
+                ColumnLeaf::new(field.clone(), child_leaf_ids, Some(child_column_leaves))
+            }
+            _ => {
+                let column_leaf = ColumnLeaf::new(field.clone(), vec![*leaf_id], None);
+                *leaf_id += 1;
+                column_leaf
+            }
+        }
+    }
+
+    pub fn traverse_path<'a>(
+        column_leaves: &'a [ColumnLeaf],
+        path: &'a [usize],
+    ) -> Result<&'a ColumnLeaf> {
+        let column_leaf = &column_leaves[path[0]];
+        if path.len() > 1 {
+            return match &column_leaf.children {
+                Some(ref children) => Self::traverse_path(children, &path[1..]),
+                None => Err(ErrorCode::Internal(format!(
+                    "Cannot get column_leaf by path: {:?}",
+                    path
+                ))),
+            };
+        }
+        Ok(column_leaf)
+    }
+}
+
+/// `ColumnLeaf` contains all the leaf column ids of the column.
+/// For the nested types, it may contain more than one leaf column.
+#[derive(Debug, Clone)]
+pub struct ColumnLeaf {
+    pub field: ArrowField,
+    // `leaf_ids` is the indices of all the leaf columns in DFS order,
+    // through which we can find the meta information of the leaf columns.
+    pub leaf_ids: Vec<usize>,
+    // Optional children column for nested types.
+    pub children: Option<Vec<ColumnLeaf>>,
+}
+
+impl ColumnLeaf {
+    pub fn new(field: ArrowField, leaf_ids: Vec<usize>, children: Option<Vec<ColumnLeaf>>) -> Self {
+        Self {
+            field,
+            leaf_ids,
+            children,
+        }
+    }
+}

src/common/storage/src/column_leaf.rs

@@ -68,3 +68,7 @@ mod utils;
 
 mod cache;
 pub use cache::FuseCachePolicy;
+
+mod column_leaf;
+pub use column_leaf::ColumnLeaf;
+pub use column_leaf::ColumnLeaves;

src/common/storage/src/lib.rs

@@ -16,6 +16,9 @@ use std::collections::BTreeMap;
 use std::fmt::Formatter;
 
 use common_datavalues::DataSchema;
+use common_exception::Result;
+use common_storage::ColumnLeaf;
+use common_storage::ColumnLeaves;
 
 #[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, Eq)]
 pub enum Projection {
@@ -49,6 +52,28 @@ impl Projection {
             Projection::InnerColumns(path_indices) => schema.inner_project(path_indices),
         }
     }
+
+    pub fn project_column_leaves<'a>(
+        &'a self,
+        column_leaves: &'a ColumnLeaves,
+    ) -> Result<Vec<&ColumnLeaf>> {
+        let column_leaves = match self {
+            Projection::Columns(indices) => indices
+                .iter()
+                .map(|idx| &column_leaves.column_leaves[*idx])
+                .collect(),
+            Projection::InnerColumns(path_indices) => {
+                let paths: Vec<&Vec<usize>> = path_indices.values().collect();
+                paths
+                    .iter()
+                    .map(|path| {
+                        ColumnLeaves::traverse_path(&column_leaves.column_leaves, path).unwrap()
+                    })
+                    .collect()
+            }
+        };
+        Ok(column_leaves)
+    }
 }
 
 impl core::fmt::Debug for Projection {

src/query/catalog/src/plan/projection.rs

@@ -14,6 +14,7 @@
 
 use std::fmt::Debug;
 
+use common_datavalues::DataSchemaRef;
 use common_meta_types::UserStageInfo;
 
 use crate::plan::Expression;
@@ -61,3 +62,31 @@ pub struct PushDownInfo {
     /// Optional stage info, used for COPY into <table> from stage
     pub stage: Option<StagePushDownInfo>,
 }
+
+impl PushDownInfo {
+    pub fn prewhere_of_push_downs(push_downs: &Option<PushDownInfo>) -> Option<PrewhereInfo> {
+        if let Some(PushDownInfo { prewhere, .. }) = push_downs {
+            prewhere.clone()
+        } else {
+            None
+        }
+    }
+
+    pub fn projection_of_push_downs(
+        schema: &DataSchemaRef,
+        push_downs: &Option<PushDownInfo>,
+    ) -> Projection {
+        if let Some(PushDownInfo {
+            projection: Some(prj),
+            ..
+        }) = push_downs
+        {
+            prj.clone()
+        } else {
+            let indices = (0..schema.fields().len())
+                .into_iter()
+                .collect::<Vec<usize>>();
+            Projection::Columns(indices)
+        }
+    }
+}

src/query/catalog/src/plan/pushdown.rs

@@ -65,6 +65,7 @@ common-storages-index = { path = "../storages/index" }
 common-storages-information-schema = { path = "../storages/information-schema" }
 common-storages-memory = { path = "../storages/memory" }
 common-storages-null = { path = "../storages/null" }
+common-storages-parquet = { path = "../storages/parquet" }
 common-storages-random = { path = "../storages/random" }
 common-storages-share = { path = "../storages/share" }
 common-storages-stage = { path = "../storages/stage" }

src/query/service/Cargo.toml

@@ -13,16 +13,14 @@
 //  limitations under the License.
 
 mod async_crash_me;
-mod memory_block_part;
-mod numbers_part;
-mod numbers_table;
+mod numbers;
 mod sync_crash_me;
 mod table_function;
 mod table_function_factory;
 
-pub use memory_block_part::generate_numbers_parts;
-pub use numbers_part::NumbersPartInfo;
-pub use numbers_table::NumbersTable;
+pub use numbers::generate_numbers_parts;
+pub use numbers::NumbersPartInfo;
+pub use numbers::NumbersTable;
 pub use table_function::TableFunction;
 pub use table_function_factory::TableArgs;
 pub use table_function_factory::TableFunctionFactory;