Merge pull request #9154 from RinChanNOWWW/improve-read-parquet

feat: `read_parquet` read meta before read data.

Author: BohuTANG

Cargo.lock

@@ -27,7 +27,6 @@ common-storage = { path = "../../../common/storage" }
 
 async-trait = { version = "0.1.57", package = "async-trait-fn" }
 chrono = { workspace = true }
-futures = "0.3.24"
 glob = "0.3.0"
 opendal = "0.22"
 serde = { workspace = true }

src/query/storages/parquet/Cargo.toml

@@ -20,8 +20,7 @@ mod parquet_reader;
 mod parquet_source;
 mod table_function;
 
-pub use parquet_column::ParquetColumnMeta;
-pub use parquet_part::ParquetPartInfo;
+pub use parquet_part::ParquetLocationPart;
 pub use parquet_reader::ParquetReader;
-pub use parquet_source::ParquetTableSource;
+pub use parquet_source::ParquetSource;
 pub use table_function::ParquetTable;

src/query/storages/parquet/src/lib.rs

@@ -55,28 +55,3 @@ impl From<ParquetCompression> for Compression {
         }
     }
 }
-
-#[derive(serde::Serialize, serde::Deserialize, PartialEq, Eq)]
-pub struct ParquetColumnMeta {
-    pub offset: u64,
-    pub length: u64,
-    pub num_values: u64,
-
-    pub compression: Compression,
-}
-
-impl ParquetColumnMeta {
-    pub fn create(
-        offset: u64,
-        length: u64,
-        num_values: u64,
-        compression: Compression,
-    ) -> ParquetColumnMeta {
-        ParquetColumnMeta {
-            offset,
-            length,
-            num_values,
-            compression,
-        }
-    }
-}

src/query/storages/parquet/src/parquet_column.rs

@@ -19,31 +19,116 @@ use std::hash::Hash;
 use std::hash::Hasher;
 use std::sync::Arc;
 
+use common_arrow::parquet::compression::Compression as ParquetCompression;
 use common_catalog::plan::PartInfo;
 use common_catalog::plan::PartInfoPtr;
 use common_exception::ErrorCode;
 use common_exception::Result;
 
-use crate::ParquetColumnMeta;
+#[derive(serde::Serialize, serde::Deserialize, PartialEq, Eq)]
+pub struct ParquetLocationPart {
+    pub location: String,
+}
+
+#[typetag::serde(name = "parquet_location")]
+impl PartInfo for ParquetLocationPart {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn equals(&self, info: &Box<dyn PartInfo>) -> bool {
+        match info.as_any().downcast_ref::<ParquetLocationPart>() {
+            None => false,
+            Some(other) => self == other,
+        }
+    }
+
+    fn hash(&self) -> u64 {
+        let mut s = DefaultHasher::new();
+        self.location.hash(&mut s);
+        s.finish()
+    }
+}
+
+impl ParquetLocationPart {
+    pub fn create(location: String) -> Arc<Box<dyn PartInfo>> {
+        Arc::new(Box::new(ParquetLocationPart { location }))
+    }
+
+    pub fn from_part(info: &PartInfoPtr) -> Result<&ParquetLocationPart> {
+        match info.as_any().downcast_ref::<ParquetLocationPart>() {
+            Some(part_ref) => Ok(part_ref),
+            None => Err(ErrorCode::Internal(
+                "Cannot downcast from PartInfo to ParquetLocationPart.",
+            )),
+        }
+    }
+}
+
+#[derive(serde::Serialize, serde::Deserialize, Debug, Eq, PartialEq, Hash, Clone, Copy)]
+pub enum Compression {
+    Uncompressed,
+    Snappy,
+    Gzip,
+    Lzo,
+    Brotli,
+    Lz4,
+    Zstd,
+    Lz4Raw,
+}
+
+impl From<Compression> for ParquetCompression {
+    fn from(value: Compression) -> Self {
+        match value {
+            Compression::Uncompressed => ParquetCompression::Uncompressed,
+            Compression::Snappy => ParquetCompression::Snappy,
+            Compression::Gzip => ParquetCompression::Gzip,
+            Compression::Lzo => ParquetCompression::Lzo,
+            Compression::Brotli => ParquetCompression::Brotli,
+            Compression::Lz4 => ParquetCompression::Lz4,
+            Compression::Zstd => ParquetCompression::Zstd,
+            Compression::Lz4Raw => ParquetCompression::Lz4Raw,
+        }
+    }
+}
+
+impl From<ParquetCompression> for Compression {
+    fn from(value: ParquetCompression) -> Self {
+        match value {
+            ParquetCompression::Uncompressed => Compression::Uncompressed,
+            ParquetCompression::Snappy => Compression::Snappy,
+            ParquetCompression::Gzip => Compression::Gzip,
+            ParquetCompression::Lzo => Compression::Lzo,
+            ParquetCompression::Brotli => Compression::Brotli,
+            ParquetCompression::Lz4 => Compression::Lz4,
+            ParquetCompression::Zstd => Compression::Zstd,
+            ParquetCompression::Lz4Raw => Compression::Lz4Raw,
+        }
+    }
+}
+
+#[derive(serde::Serialize, serde::Deserialize, PartialEq, Eq)]
+pub struct ColumnMeta {
+    pub offset: u64,
+    pub length: u64,
+    pub compression: Compression,
+}
 
 #[derive(serde::Serialize, serde::Deserialize, PartialEq, Eq)]
-pub struct ParquetPartInfo {
+pub struct ParquetRowGroupPart {
     pub location: String,
-    /// FusePartInfo itself is not versioned
-    /// the `format_version` is the version of the block which the `location` points to
-    pub format_version: u64,
-    pub nums_rows: usize,
-    pub columns_meta: HashMap<usize, ParquetColumnMeta>,
+    pub num_rows: usize,
+    pub column_metas: HashMap<usize, ColumnMeta>,
 }
 
-#[typetag::serde(name = "parquet")]
-impl PartInfo for ParquetPartInfo {
+#[typetag::serde(name = "parquet_row_group")]
+impl PartInfo for ParquetRowGroupPart {
     fn as_any(&self) -> &dyn Any {
         self
     }
 
     fn equals(&self, info: &Box<dyn PartInfo>) -> bool {
-        match info.as_any().downcast_ref::<ParquetPartInfo>() {
+        match info.as_any().downcast_ref::<ParquetRowGroupPart>() {
             None => false,
             Some(other) => self == other,
         }
@@ -56,26 +141,24 @@ impl PartInfo for ParquetPartInfo {
     }
 }
 
-impl ParquetPartInfo {
+impl ParquetRowGroupPart {
     pub fn create(
         location: String,
-        format_version: u64,
-        rows_count: u64,
-        columns_meta: HashMap<usize, ParquetColumnMeta>,
+        num_rows: usize,
+        column_metas: HashMap<usize, ColumnMeta>,
     ) -> Arc<Box<dyn PartInfo>> {
-        Arc::new(Box::new(ParquetPartInfo {
+        Arc::new(Box::new(ParquetRowGroupPart {
             location,
-            format_version,
-            columns_meta,
-            nums_rows: rows_count as usize,
+            num_rows,
+            column_metas,
         }))
     }
 
-    pub fn from_part(info: &PartInfoPtr) -> Result<&ParquetPartInfo> {
-        match info.as_any().downcast_ref::<ParquetPartInfo>() {
+    pub fn from_part(info: &PartInfoPtr) -> Result<&ParquetRowGroupPart> {
+        match info.as_any().downcast_ref::<ParquetRowGroupPart>() {
             Some(part_ref) => Ok(part_ref),
             None => Err(ErrorCode::Internal(
-                "Cannot downcast from PartInfo to FusePartInfo.",
+                "Cannot downcast from PartInfo to ParquetRowGroupPart.",
             )),
         }
     }

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

@@ -24,105 +24,97 @@ use common_arrow::parquet::metadata::ColumnDescriptor;
 use common_arrow::parquet::read::BasicDecompressor;
 use common_arrow::parquet::read::PageMetaData;
 use common_arrow::parquet::read::PageReader;
-use common_catalog::plan::PartInfoPtr;
 use common_datablocks::DataBlock;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_storage::ColumnLeaf;
 
-use crate::ParquetColumnMeta;
-use crate::ParquetPartInfo;
+use crate::parquet_part::ColumnMeta;
+use crate::parquet_part::ParquetRowGroupPart;
 use crate::ParquetReader;
 
 impl ParquetReader {
-    fn to_array_iter(
-        metas: Vec<&ParquetColumnMeta>,
-        chunks: Vec<Vec<u8>>,
-        rows: usize,
-        column_descriptors: Vec<&ColumnDescriptor>,
-        field: Field,
-    ) -> Result<ArrayIter<'static>> {
-        let columns = metas
-            .iter()
-            .zip(chunks.into_iter().zip(column_descriptors.iter()))
-            .map(|(meta, (chunk, column_descriptor))| {
-                let page_meta_data = PageMetaData {
-                    column_start: meta.offset,
-                    num_values: meta.num_values as i64,
-                    compression: meta.compression.into(),
-                    descriptor: column_descriptor.descriptor.clone(),
-                };
-                let pages = PageReader::new_with_page_meta(
-                    std::io::Cursor::new(chunk),
-                    page_meta_data,
-                    Arc::new(|_, _| true),
-                    vec![],
-                    usize::MAX,
-                );
-                Ok(BasicDecompressor::new(pages, vec![]))
-            })
-            .collect::<Result<Vec<_>>>()?;
-
-        let types = column_descriptors
-            .iter()
-            .map(|column_descriptor| &column_descriptor.descriptor.primitive_type)
-            .collect::<Vec<_>>();
-
-        Ok(column_iter_to_arrays(
-            columns,
-            types,
-            field,
-            Some(rows),
-            rows,
-        )?)
-    }
-
     pub fn deserialize(
         &self,
-        part: PartInfoPtr,
+        part: &ParquetRowGroupPart,
         chunks: Vec<(usize, Vec<u8>)>,
     ) -> Result<DataBlock> {
-        let part = ParquetPartInfo::from_part(&part)?;
         let mut chunk_map: HashMap<usize, Vec<u8>> = chunks.into_iter().collect();
-        let mut columns_array_iter = Vec::with_capacity(self.projection.len());
+        let mut columns_array_iter = Vec::with_capacity(self.projected_schema.num_fields());
+
+        let column_leaves = &self.projected_column_leaves.column_leaves;
+        let mut cnt_map = Self::build_projection_count_map(column_leaves);
 
-        let num_rows = part.nums_rows;
-        let columns = self.projection.project_column_leaves(&self.column_leaves)?;
-        let mut cnt_map = Self::build_projection_count_map(&columns);
-        for column in &columns {
-            let field = column.field.clone();
-            let indices = &column.leaf_ids;
-            let mut column_metas = Vec::with_capacity(indices.len());
-            let mut column_chunks = Vec::with_capacity(indices.len());
-            let mut column_descriptors = Vec::with_capacity(indices.len());
+        for column_leaf in column_leaves {
+            let indices = &column_leaf.leaf_ids;
+            let mut metas = Vec::with_capacity(indices.len());
+            let mut chunks = Vec::with_capacity(indices.len());
             for index in indices {
-                let column_meta = &part.columns_meta[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 {
                     chunk_map.get(index).unwrap().clone()
                 } else {
                     chunk_map.remove(index).unwrap()
                 };
-                let column_descriptor = &self.parquet_schema_descriptor.columns()[*index];
-                column_metas.push(column_meta);
-                column_chunks.push(column_chunk);
-                column_descriptors.push(column_descriptor);
+                let descriptor = &self.projected_column_descriptors[index];
+                metas.push((column_meta, descriptor));
+                chunks.push(column_chunk);
             }
             columns_array_iter.push(Self::to_array_iter(
-                column_metas,
-                column_chunks,
-                num_rows,
-                column_descriptors,
-                field,
+                metas,
+                chunks,
+                part.num_rows,
+                column_leaf.field.clone(),
             )?);
         }
 
-        let mut deserializer = RowGroupDeserializer::new(columns_array_iter, num_rows, None);
+        let mut deserializer = RowGroupDeserializer::new(columns_array_iter, part.num_rows, None);
 
         self.try_next_block(&mut deserializer)
     }
 
+    /// The number of columns can be greater than 1 because the it may be a nested type.
+    /// Combine multiple columns into one arrow array.
+    fn to_array_iter(
+        metas: Vec<(&ColumnMeta, &ColumnDescriptor)>,
+        chunks: Vec<Vec<u8>>,
+        rows: usize,
+        field: Field,
+    ) -> Result<ArrayIter<'static>> {
+        let (columns, types) = metas
+            .iter()
+            .zip(chunks.into_iter())
+            .map(|(&(meta, descriptor), chunk)| {
+                let pages = PageReader::new_with_page_meta(
+                    std::io::Cursor::new(chunk),
+                    PageMetaData {
+                        column_start: meta.offset,
+                        num_values: meta.length as i64,
+                        compression: meta.compression.into(),
+                        descriptor: descriptor.descriptor.clone(),
+                    },
+                    Arc::new(|_, _| true),
+                    vec![],
+                    usize::MAX,
+                );
+                (
+                    BasicDecompressor::new(pages, vec![]),
+                    &descriptor.descriptor.primitive_type,
+                )
+            })
+            .unzip();
+
+        Ok(column_iter_to_arrays(
+            columns,
+            types,
+            field,
+            Some(rows),
+            rows,
+        )?)
+    }
+
     fn try_next_block(&self, deserializer: &mut RowGroupDeserializer) -> Result<DataBlock> {
         match deserializer.next() {
             None => Err(ErrorCode::Internal(
@@ -134,7 +126,7 @@ impl ParquetReader {
     }
 
     // Build a map to record the count number of each leaf_id
-    fn build_projection_count_map(columns: &Vec<&ColumnLeaf>) -> HashMap<usize, usize> {
+    fn build_projection_count_map(columns: &[ColumnLeaf]) -> HashMap<usize, usize> {
         let mut cnt_map = HashMap::with_capacity(columns.len());
         for column in columns {
             for index in &column.leaf_ids {