split block by default rows

Author: ariesdevil

src/query/datablocks/src/data_block.rs

@@ -292,6 +292,51 @@ impl DataBlock {
         Ok(DataBlock::create(schema.clone(), columns))
     }
 
+    pub fn from_chunk_with_row_limit<A: AsRef<dyn Array>>(
+        schema: &DataSchemaRef,
+        chuck: &Chunk<A>,
+        num_rows: usize,
+    ) -> Result<Vec<DataBlock>> {
+        let columns: Vec<ColumnRef> = chuck
+            .columns()
+            .iter()
+            .zip(schema.fields().iter())
+            .map(|(col, f)| match f.is_nullable() {
+                true => col.into_nullable_column(),
+                false => col.into_column(),
+            })
+            .collect();
+
+        let block = DataBlock::create(schema.clone(), columns);
+
+        let total_rows_in_block: usize = block.num_rows();
+
+        if total_rows_in_block >= num_rows {
+            let num_blocks = if total_rows_in_block % num_rows == 0 {
+                total_rows_in_block / num_rows
+            } else {
+                total_rows_in_block / num_rows + 1
+            };
+            let mut blocks: Vec<DataBlock> = Vec::with_capacity(num_blocks);
+
+            let mut offset = 0;
+            let mut remain_rows = total_rows_in_block;
+            while remain_rows >= num_rows {
+                let cut = block.slice(offset, num_rows);
+                blocks.push(cut);
+                offset += num_rows;
+                remain_rows -= num_rows;
+            }
+            if remain_rows > 0 {
+                blocks.push(block.slice(offset, remain_rows));
+            }
+
+            Ok(blocks)
+        } else {
+            Ok(vec![block])
+        }
+    }
+
     pub fn get_serializers(&self) -> Result<Vec<TypeSerializerImpl>> {
         let columns_size = self.num_columns();
 

src/query/pipeline/sources/src/processors/sources/input_formats/impls/input_format_parquet.rs

@@ -62,6 +62,8 @@ use crate::processors::sources::input_formats::input_split::FileInfo;
 use crate::processors::sources::input_formats::input_split::SplitInfo;
 use crate::processors::sources::input_formats::InputFormat;
 
+pub const DEFAULT_ROW_PER_BLOCK: usize = 1000 * 1000;
+
 pub struct InputFormatParquet;
 
 fn col_offset(meta: &ColumnChunkMetaData) -> i64 {
@@ -280,6 +282,7 @@ impl RowGroupInMemory {
             )?;
             column_chunks.push(array_iters);
         }
+
         match RowGroupDeserializer::new(column_chunks, self.meta.num_rows(), None).next() {
             None => Err(ErrorCode::Internal(
                 "deserialize from raw group: fail to get a chunk",
@@ -332,8 +335,13 @@ impl BlockBuilderTrait for ParquetBlockBuilder {
     fn deserialize(&mut self, mut batch: Option<RowGroupInMemory>) -> Result<Vec<DataBlock>> {
         if let Some(rg) = batch.as_mut() {
             let chunk = rg.get_arrow_chunk()?;
-            let block = DataBlock::from_chunk(&self.ctx.schema, &chunk)?;
-            Ok(vec![block])
+            let blocks = DataBlock::from_chunk_with_row_limit(
+                &self.ctx.schema,
+                &chunk,
+                DEFAULT_ROW_PER_BLOCK,
+            )?;
+
+            Ok(blocks)
         } else {
             Ok(vec![])
         }