feat: streaming load support parquet. (#18017)

* chore: ParquetPart do not need to impl Eq.

* refactor: simplify ParquetCopySource.

* refactor: rename fn.

* refactor: read_blocks_from_binary return iter instead of vector.

* feat: streaming load support parquet.

Author: youngsofun

src/meta/app/src/principal/file_format.rs

@@ -78,6 +78,16 @@ impl FileFormatParams {
         }
     }
 
+    pub fn support_streaming_load(&self) -> bool {
+        matches!(
+            self,
+            FileFormatParams::Csv(_)
+                | FileFormatParams::Tsv(_)
+                | FileFormatParams::NdJson(_)
+                | FileFormatParams::Parquet(_)
+        )
+    }
+
     pub fn default_by_type(format_type: StageFileFormatType) -> Result<Self> {
         match format_type {
             StageFileFormatType::Parquet => {

src/query/service/src/pipelines/builders/builder_copy_into_table.rs

@@ -35,12 +35,12 @@ use databend_common_sql::executor::physical_plans::CopyIntoTable;
 use databend_common_sql::executor::physical_plans::CopyIntoTableSource;
 use databend_common_sql::plans::CopyIntoTableMode;
 use databend_common_storage::StageFileInfo;
+use databend_common_storages_stage::TransformNullIf;
 use log::debug;
 use log::info;
 
 use crate::pipelines::processors::transforms::TransformAddConstColumns;
 use crate::pipelines::processors::transforms::TransformCastSchema;
-use crate::pipelines::processors::transforms::TransformNullIf;
 use crate::pipelines::PipelineBuilder;
 use crate::sessions::QueryContext;
 
@@ -135,7 +135,7 @@ impl PipelineBuilder {
         let source_schema = if let Some(null_if) =
             Self::need_null_if_processor(plan, &source_schema, plan_required_source_schema)
         {
-            let func_ctx = ctx.get_function_context()?;
+            let func_ctx = Arc::new(ctx.get_function_context()?);
             main_pipeline.try_add_transformer(|| {
                 TransformNullIf::try_new(
                     source_schema.clone(),

src/query/service/src/pipelines/processors/transforms/mod.rs

@@ -34,7 +34,6 @@ mod transform_filter;
 mod transform_limit;
 mod transform_merge_block;
 mod transform_merge_sort;
-mod transform_null_if;
 mod transform_recursive_cte_scan;
 mod transform_recursive_cte_source;
 mod transform_resort_addon;
@@ -67,7 +66,6 @@ pub use transform_filter::TransformFilter;
 pub use transform_limit::TransformLimit;
 pub use transform_merge_block::TransformMergeBlock;
 pub use transform_merge_sort::*;
-pub use transform_null_if::TransformNullIf;
 pub use transform_recursive_cte_scan::TransformRecursiveCteScan;
 pub use transform_recursive_cte_source::TransformRecursiveCteSource;
 pub use transform_resort_addon::TransformResortAddOn;

src/query/service/src/servers/http/v1/streaming_load.rs

@@ -170,7 +170,7 @@ pub async fn streaming_load_handler_inner(
                     receiver,
                     ..
                 } => {
-                    if !matches!(&**format, FileFormatParams::Csv(_) | FileFormatParams::Tsv(_) | FileFormatParams::NdJson(_))  {
+                    if !format.support_streaming_load() {
                         return Err(poem::Error::from_string( format!( "[HTTP-STREAMING-LOAD] Unsupported file format: {}", format.get_type() ), StatusCode::BAD_REQUEST));
                     }
                     let (tx, rx) = tokio::sync::mpsc::channel(1);
@@ -185,17 +185,22 @@ pub async fn streaming_load_handler_inner(
                             id: http_context.query_id.clone(),
                             stats: query_context.get_write_progress().get_values(),
                         })),
-                        Ok(Err(cause)) => Err(poem::Error::from_string(
+                        Ok(Err(cause)) => {
+                            info!("[HTTP-STREAMING-LOAD] Query execution failed: {:?}", cause);
+                            Err(poem::Error::from_string(
                             format!(
                                 "[HTTP-STREAMING-LOAD] Query execution failed: {}",
                                 cause.display_with_sql(sql).message()
                             ),
                             StatusCode::BAD_REQUEST,
-                        )),
-                        Err(_) => Err(poem::Error::from_string(
-                            "[HTTP-STREAMING-LOAD] Internal server error: execution thread panicked",
-                            StatusCode::INTERNAL_SERVER_ERROR,
-                        )),
+                        ))},
+                        Err(err) => {
+                            info!("[HTTP-STREAMING-LOAD] Internal server error: {:?}", err);
+                            Err(poem::Error::from_string(
+                                "[HTTP-STREAMING-LOAD] Internal server error: execution thread panicked",
+                                StatusCode::INTERNAL_SERVER_ERROR,
+                            ))
+                        },
                     }
                 }
                 _non_supported_source => Err(poem::Error::from_string(

src/query/storages/delta/src/partition.rs

@@ -22,7 +22,7 @@ use databend_common_expression::Scalar;
 use databend_common_storages_parquet::ParquetPart;
 
 /// only support parquet for now: https://github.com/delta-io/delta/issues/87
-#[derive(serde::Serialize, serde::Deserialize, PartialEq, Eq, Debug, Clone)]
+#[derive(serde::Serialize, serde::Deserialize, PartialEq, Debug, Clone)]
 pub struct DeltaPartInfo {
     pub data: ParquetPart,
     pub partition_values: Vec<Scalar>,

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

@@ -12,8 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod projection;
 mod reader;
 mod source;
 mod table;
 
+pub(crate) use projection::CopyProjectionEvaluator;
 pub use table::ParquetTableForCopy;

src/query/storages/parquet/src/copy_into_table/projection.rs

@@ -0,0 +1,45 @@
+// Copyright 2021 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::sync::Arc;
+
+use databend_common_exception::Result;
+use databend_common_expression::BlockEntry;
+use databend_common_expression::DataBlock;
+use databend_common_expression::DataSchemaRef;
+use databend_common_expression::Evaluator;
+use databend_common_expression::Expr;
+use databend_common_expression::FunctionContext;
+use databend_common_functions::BUILTIN_FUNCTIONS;
+
+pub(crate) struct CopyProjectionEvaluator {
+    schema: DataSchemaRef,
+    func_ctx: Arc<FunctionContext>,
+}
+
+impl CopyProjectionEvaluator {
+    pub(crate) fn new(schema: DataSchemaRef, func_ctx: Arc<FunctionContext>) -> Self {
+        Self { schema, func_ctx }
+    }
+    pub(crate) fn project(&self, block: &DataBlock, projection: &[Expr]) -> Result<DataBlock> {
+        let evaluator = Evaluator::new(block, &self.func_ctx, &BUILTIN_FUNCTIONS);
+        let mut columns = Vec::with_capacity(projection.len());
+        for (field, expr) in self.schema.fields().iter().zip(projection.iter()) {
+            let value = evaluator.run(expr)?;
+            let column = BlockEntry::new(field.data_type().clone(), value);
+            columns.push(column);
+        }
+        Ok(DataBlock::new(columns, block.num_rows()))
+    }
+}

src/query/storages/parquet/src/copy_into_table/reader.rs

@@ -54,7 +54,7 @@ impl RowGroupReaderForCopy {
             InMemoryRowGroup::new(&part.location, op.clone(), &part.meta, None, *read_settings);
         let mut _sorter = None;
         self.row_group_reader_builder
-            .build(row_group, None, &mut _sorter, None, batch_size)
+            .fetch_and_build(row_group, None, &mut _sorter, None, batch_size)
             .await
     }
 

src/query/storages/parquet/src/copy_into_table/source.rs

@@ -22,28 +22,24 @@ use databend_common_base::runtime::profile::Profile;
 use databend_common_base::runtime::profile::ProfileStatisticsName;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
-use databend_common_expression::BlockEntry;
 use databend_common_expression::DataBlock;
 use databend_common_expression::DataSchemaRef;
-use databend_common_expression::Evaluator;
-use databend_common_expression::FunctionContext;
-use databend_common_functions::BUILTIN_FUNCTIONS;
+use databend_common_expression::Expr;
 use databend_common_pipeline_core::processors::Event;
 use databend_common_pipeline_core::processors::OutputPort;
 use databend_common_pipeline_core::processors::Processor;
 use databend_common_pipeline_core::processors::ProcessorPtr;
 use opendal::Operator;
 
+use crate::copy_into_table::projection::CopyProjectionEvaluator;
 use crate::copy_into_table::reader::RowGroupReaderForCopy;
 use crate::parquet_reader::policy::ReadPolicyImpl;
 use crate::read_settings::ReadSettings;
 use crate::ParquetPart;
 
-type SchemaIndex = usize;
-
 enum State {
     Init,
-    ReadRowGroup((SchemaIndex, ReadPolicyImpl)),
+    ReadRowGroup((Vec<Expr>, ReadPolicyImpl)),
 }
 
 pub struct ParquetCopySource {
@@ -59,9 +55,7 @@ pub struct ParquetCopySource {
     // Used to read parquet.
     row_group_readers: Arc<HashMap<usize, RowGroupReaderForCopy>>,
     operator: Operator,
-    schema: DataSchemaRef,
-    func_ctx: FunctionContext,
-
+    copy_projection_evaluator: CopyProjectionEvaluator,
     state: State,
     batch_size: usize,
 }
@@ -76,20 +70,20 @@ impl ParquetCopySource {
     ) -> Result<ProcessorPtr> {
         let scan_progress = ctx.get_scan_progress();
         let batch_size = ctx.get_settings().get_parquet_max_block_size()? as usize;
-        let func_ctx = ctx.get_function_context()?;
+        let func_ctx = Arc::new(ctx.get_function_context()?);
+        let copy_projection_evaluator = CopyProjectionEvaluator::new(schema, func_ctx);
 
         Ok(ProcessorPtr::create(Box::new(Self {
             output,
             scan_progress,
             ctx,
             operator,
             row_group_readers,
-            func_ctx,
             batch_size,
             generated_data: None,
             is_finished: false,
             state: State::Init,
-            schema,
+            copy_projection_evaluator,
         })))
     }
 }
@@ -141,23 +135,13 @@ impl Processor for ParquetCopySource {
 
     fn process(&mut self) -> Result<()> {
         match std::mem::replace(&mut self.state, State::Init) {
-            State::ReadRowGroup((schema_index, mut reader)) => {
+            State::ReadRowGroup((projection, mut reader)) => {
                 if let Some(block) = reader.as_mut().read_block()? {
-                    let projection = self
-                        .row_group_readers
-                        .get(&schema_index)
-                        .unwrap()
-                        .output_projection();
-                    let evaluator = Evaluator::new(&block, &self.func_ctx, &BUILTIN_FUNCTIONS);
-                    let mut columns = Vec::with_capacity(projection.len());
-                    for (field, expr) in self.schema.fields().iter().zip(projection.iter()) {
-                        let value = evaluator.run(expr)?;
-                        let column = BlockEntry::new(field.data_type().clone(), value);
-                        columns.push(column);
-                    }
-                    let block = DataBlock::new(columns, block.num_rows());
-                    self.generated_data = Some(block);
-                    self.state = State::ReadRowGroup((schema_index, reader));
+                    self.generated_data = Some(
+                        self.copy_projection_evaluator
+                            .project(&block, &projection)?,
+                    );
+                    self.state = State::ReadRowGroup((projection, reader));
                 }
                 // Else: The reader is finished. We should try to build another reader.
             }
@@ -178,6 +162,7 @@ impl Processor for ParquetCopySource {
                                 .row_group_readers
                                 .get(&schema_index)
                                 .expect("schema index must exist");
+                            let projection = builder.output_projection().to_vec();
                             let reader = builder
                                 .build_reader(
                                     part,
@@ -188,7 +173,7 @@ impl Processor for ParquetCopySource {
                                 .await?
                                 .expect("reader must exist");
                             {
-                                self.state = State::ReadRowGroup((schema_index, reader));
+                                self.state = State::ReadRowGroup((projection, reader));
                             }
                             // Else: keep in init state.
                         }
@@ -200,7 +185,6 @@ impl Processor for ParquetCopySource {
             }
             _ => unreachable!(),
         }
-
         Ok(())
     }
 }

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

@@ -23,6 +23,7 @@
 #![feature(core_intrinsics)]
 #![feature(int_roundings)]
 #![feature(box_patterns)]
+#![feature(result_flattening)]
 // FIXME: Remove this once the deprecated code is removed
 #![allow(deprecated)]
 
@@ -45,6 +46,7 @@ mod schema;
 pub use copy_into_table::ParquetTableForCopy;
 pub use parquet_part::ParquetFilePart;
 pub use parquet_part::ParquetPart;
+pub use parquet_reader::InmMemoryFile;
 pub use parquet_reader::ParquetFileReader;
 pub use parquet_reader::ParquetReaderBuilder;
 pub use parquet_reader::ParquetWholeFileReader;