Make every row group as a part.

Author: RinChanNOWWW

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

@@ -20,7 +20,7 @@ mod parquet_reader;
 mod parquet_source;
 mod table_function;
 
-pub use parquet_part::ParquetPart;
+pub use parquet_part::ParquetLocationPart;
 pub use parquet_reader::ParquetReader;
 pub use parquet_source::ParquetSource;
 pub use table_function::ParquetTable;

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

@@ -18,24 +18,25 @@ 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;
 
 #[derive(serde::Serialize, serde::Deserialize, PartialEq, Eq)]
-pub struct ParquetPart {
+pub struct ParquetLocationPart {
     pub location: String,
 }
 
-#[typetag::serde(name = "parquet")]
-impl PartInfo for ParquetPart {
+#[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::<ParquetPart>() {
+        match info.as_any().downcast_ref::<ParquetLocationPart>() {
             None => false,
             Some(other) => self == other,
         }
@@ -48,16 +49,115 @@ impl PartInfo for ParquetPart {
     }
 }
 
-impl ParquetPart {
+impl ParquetLocationPart {
     pub fn create(location: String) -> Arc<Box<dyn PartInfo>> {
-        Arc::new(Box::new(ParquetPart { location }))
+        Arc::new(Box::new(ParquetLocationPart { location }))
     }
 
-    pub fn from_part(info: &PartInfoPtr) -> Result<&ParquetPart> {
-        match info.as_any().downcast_ref::<ParquetPart>() {
+    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 FusePartInfo.",
+                "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 Into<Compression> for ParquetCompression {
+    fn into(self) -> Compression {
+        match self {
+            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 ParquetRowGroupPart {
+    pub location: String,
+    pub num_rows: usize,
+    pub column_metas: Vec<ColumnMeta>,
+}
+
+#[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::<ParquetRowGroupPart>() {
+            None => false,
+            Some(other) => self == other,
+        }
+    }
+
+    fn hash(&self) -> u64 {
+        let mut s = DefaultHasher::new();
+        self.location.hash(&mut s);
+        s.finish()
+    }
+}
+
+impl ParquetRowGroupPart {
+    pub fn create(
+        location: String,
+        num_rows: usize,
+        column_metas: Vec<ColumnMeta>,
+    ) -> Arc<Box<dyn PartInfo>> {
+        Arc::new(Box::new(ParquetRowGroupPart {
+            location,
+            num_rows,
+            column_metas,
+        }))
+    }
+
+    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 ParquetRowGroupPart.",
             )),
         }
     }

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

@@ -20,8 +20,7 @@ use common_arrow::arrow::datatypes::Field;
 use common_arrow::arrow::io::parquet::read::column_iter_to_arrays;
 use common_arrow::arrow::io::parquet::read::ArrayIter;
 use common_arrow::arrow::io::parquet::read::RowGroupDeserializer;
-use common_arrow::parquet::metadata::ColumnChunkMetaData;
-use common_arrow::parquet::metadata::RowGroupMetaData;
+use common_arrow::parquet::metadata::ColumnDescriptor;
 use common_arrow::parquet::read::BasicDecompressor;
 use common_arrow::parquet::read::PageMetaData;
 use common_arrow::parquet::read::PageReader;
@@ -30,21 +29,19 @@ use common_exception::ErrorCode;
 use common_exception::Result;
 use common_storage::ColumnLeaf;
 
+use crate::parquet_part::ColumnMeta;
+use crate::parquet_part::ParquetRowGroupPart;
 use crate::ParquetReader;
 
 impl ParquetReader {
     pub fn deserialize(
         &self,
-        rg: &RowGroupMetaData,
+        part: &ParquetRowGroupPart,
         chunks: Vec<(usize, Vec<u8>)>,
     ) -> 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 column_metas = self
-            .get_column_metas(rg)
-            .into_iter()
-            .collect::<HashMap<_, _>>();
         let column_leaves = &self.projected_column_leaves.column_leaves;
         let mut cnt_map = Self::build_projection_count_map(column_leaves);
 
@@ -53,52 +50,58 @@ 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 = 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 {
                     chunk_map.get(index).unwrap().clone()
                 } else {
                     chunk_map.remove(index).unwrap()
                 };
-                metas.push(column_meta);
+                let descriptor = &self.projected_column_descriptors[index];
+                metas.push((column_meta, descriptor));
                 chunks.push(column_chunk);
             }
             columns_array_iter.push(Self::to_array_iter(
                 metas,
                 chunks,
-                rg.num_rows(),
+                part.num_rows,
                 column_leaf.field.clone(),
             )?);
         }
 
-        let mut deserializer = RowGroupDeserializer::new(columns_array_iter, rg.num_rows(), None);
+        let mut deserializer = RowGroupDeserializer::new(columns_array_iter, part.num_rows, None);
 
         self.try_next_block(&mut deserializer)
     }
 
-    /// Combine multiple columns into one arrow array.
     /// 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<&ColumnChunkMetaData>,
+        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, chunk)| {
+            .map(|(&(meta, descriptor), chunk)| {
                 let pages = PageReader::new_with_page_meta(
                     std::io::Cursor::new(chunk),
-                    PageMetaData::from(meta),
+                    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![]),
-                    &meta.descriptor().descriptor.primitive_type,
+                    &descriptor.descriptor.primitive_type,
                 )
             })
             .unzip();

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

@@ -15,9 +15,7 @@
 use std::fs::File;
 
 use common_arrow::arrow::io::parquet::read as pread;
-use common_arrow::parquet::metadata::ColumnChunkMetaData;
 use common_arrow::parquet::metadata::FileMetaData;
-use common_arrow::parquet::metadata::RowGroupMetaData;
 use common_datavalues::DataField;
 use common_datavalues::DataSchema;
 use common_exception::ErrorCode;
@@ -54,22 +52,4 @@ impl ParquetReader {
 
         Ok(DataSchema::new(field))
     }
-
-    pub fn get_column_metas<'a>(
-        &self,
-        rg: &'a RowGroupMetaData,
-    ) -> Vec<(usize, &'a ColumnChunkMetaData)> {
-        let columns = rg.columns();
-        let column_leaves = &self.projected_column_leaves.column_leaves;
-        let mut column_metas =
-            Vec::with_capacity(column_leaves.iter().map(|col| col.leaf_ids.len()).sum());
-        for column in column_leaves {
-            let indices = &column.leaf_ids;
-            for index in indices {
-                let column_meta = &columns[*index];
-                column_metas.push((*index, column_meta));
-            }
-        }
-        column_metas
-    }
 }

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

@@ -15,8 +15,12 @@
 mod deserialize;
 mod meta;
 mod read;
+use std::collections::HashMap;
+use std::collections::HashSet;
 use std::sync::Arc;
 
+use common_arrow::arrow::io::parquet::write::to_parquet_schema;
+use common_arrow::parquet::metadata::ColumnDescriptor;
 use common_catalog::plan::Projection;
 use common_datavalues::DataSchemaRef;
 use common_exception::Result;
@@ -27,8 +31,23 @@ pub use read::IndexedChunk;
 #[derive(Clone)]
 pub struct ParquetReader {
     operator: Operator,
+    /// The indices of columns need to read by this reader.
+    ///
+    /// Use [`HashSet`] to avoid duplicate indices.
+    /// Duplicate indices will exist when there are nested types.
+    ///
+    /// For example:
+    ///
+    /// ```sql
+    /// select a, a.b, a.c from t;
+    /// ```
+    columns_to_read: HashSet<usize>,
+    /// The schema of the [`common_datablocks::DataBlock`] this reader produces.
     projected_schema: DataSchemaRef,
+    /// [`ColumnLeaves`] corresponding to the `projected_schema`.
     projected_column_leaves: ColumnLeaves,
+    /// [`ColumnDescriptor`]s corresponding to the `projected_schema`.
+    projected_column_descriptors: HashMap<usize, ColumnDescriptor>,
 }
 
 impl ParquetReader {
@@ -40,26 +59,42 @@ impl 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)?;
 
-        // Projected schema and column leaves.
+        // 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))
             }
         };
+        // Project column leaves
         let projected_column_leaves = ColumnLeaves {
             column_leaves: projection
                 .project_column_leaves(&column_leaves)?
                 .iter()
                 .map(|&leaf| leaf.clone())
                 .collect(),
         };
+        let column_leaves = &projected_column_leaves.column_leaves;
+        // Project column descriptors and collect columns to read
+        let mut projected_column_descriptors = HashMap::with_capacity(column_leaves.len());
+        let mut columns_to_read =
+            HashSet::with_capacity(column_leaves.iter().map(|leaf| leaf.leaf_ids.len()).sum());
+        for column_leaf in column_leaves {
+            for index in &column_leaf.leaf_ids {
+                columns_to_read.insert(*index);
+                projected_column_descriptors
+                    .insert(*index, schema_descriptors.columns()[*index].clone());
+            }
+        }
 
         Ok(Arc::new(ParquetReader {
             operator,
+            columns_to_read,
             projected_schema,
             projected_column_leaves,
+            projected_column_descriptors,
         }))
     }