feat: support query metadata of ORC file. (#18428)

* chore: polish error message.

* feat: support query metadata of ORC file.

Author: youngsofun

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

@@ -136,7 +136,7 @@ async fn streaming_load_handler_inner(
         .headers()
         .get(HEADER_SQL)
         .ok_or(poem::Error::from_string(
-            "[HTTP-STREAMING-LOAD] Missing required 'sql' header in request",
+            format!("[HTTP-STREAMING-LOAD] Missing required header {HEADER_SQL} in request"),
             StatusCode::BAD_REQUEST,
         ))?;
     let sql = sql.to_str().map_err(|e| {

src/query/sql/src/planner/binder/copy_into_table.rs

@@ -691,7 +691,7 @@ pub async fn resolve_stage_location(
     // my_named_stage/abc/
     let names: Vec<&str> = location.splitn(2, '/').filter(|v| !v.is_empty()).collect();
     if names[0] == STAGE_PLACEHOLDER {
-        return Err(ErrorCode::BadArguments("placeholder @_databend_upload should be used in streaming_load handler or replaced in client."));
+        return Err(ErrorCode::BadArguments("placeholder @_databend_upload as location: should be used in streaming_load handler or replaced in client."));
     }
 
     let stage = if names[0] == "~" {

src/query/sql/src/planner/binder/insert.rs

@@ -194,7 +194,7 @@ impl Binder {
                 match location.as_str() {
                     STAGE_PLACEHOLDER => {
                         if self.ctx.get_session_type() != SessionType::HTTPStreamingLoad {
-                            return Err(ErrorCode::BadArguments("placeholder @_databend_upload should be used in streaming_load handler or replaced in client."));
+                            return Err(ErrorCode::BadArguments("placeholder @_databend_upload in query handler: should be used in streaming_load handler or replaced in client."));
                         }
                         let (required_source_schema, values_consts) = if let Some(value) = value {
                             self.prepared_values(value, &schema, settings).await?

src/query/storages/common/stage/src/read/columnar/internal_columns.rs

@@ -0,0 +1,47 @@
+// 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 databend_common_catalog::plan::InternalColumnType;
+use databend_common_expression::types::DataType;
+use databend_common_expression::types::NumberColumnBuilder;
+use databend_common_expression::Column;
+use databend_common_expression::DataBlock;
+use databend_common_expression::Scalar;
+
+pub fn add_internal_columns(
+    internal_columns: &[InternalColumnType],
+    path: String,
+    b: &mut DataBlock,
+    start_row: &mut u64,
+) {
+    for c in internal_columns {
+        match c {
+            InternalColumnType::FileName => {
+                b.add_const_column(Scalar::String(path.clone()), DataType::String);
+            }
+            InternalColumnType::FileRowNumber => {
+                let end_row = (*start_row) + b.num_rows() as u64;
+                b.add_column(Column::Number(
+                    NumberColumnBuilder::UInt64(((*start_row)..end_row).collect()).build(),
+                ));
+                *start_row = end_row;
+            }
+            _ => {
+                unreachable!(
+                    "except InternalColumnType::FileName or InternalColumnType::FileRowNumber"
+                );
+            }
+        }
+    }
+}

src/query/storages/common/stage/src/read/columnar/mod.rs

@@ -13,8 +13,10 @@
 // limitations under the License.
 
 pub mod arrow_to_variant;
+mod internal_columns;
 mod projection;
 
 pub use arrow_to_variant::read_record_batch_to_variant_column;
 pub use arrow_to_variant::record_batch_to_variant_block;
+pub use internal_columns::add_internal_columns;
 pub use projection::project_columnar;

src/query/storages/iceberg/src/table.rs

@@ -333,7 +333,12 @@ impl IcebergTable {
             )?;
             pipeline.try_resize(max_threads)?;
             pipeline.add_accumulating_transformer(|| {
-                StripeDecoder::new(ctx.clone(), data_schema.clone(), arrow_schema.clone())
+                StripeDecoder::new(
+                    ctx.clone(),
+                    data_schema.clone(),
+                    arrow_schema.clone(),
+                    vec![],
+                )
             });
         } else {
             let arrow_schema: Schema = table_schema.as_ref().into();

src/query/storages/orc/src/copy_into_table/processors/source.rs

@@ -30,24 +30,19 @@ use databend_common_pipeline_sources::AsyncSource;
 use databend_common_pipeline_sources::AsyncSourcer;
 use databend_storages_common_stage::SingleFilePartition;
 use opendal::Operator;
-use orc_rust::async_arrow_reader::StripeFactory;
 use orc_rust::ArrowReaderBuilder;
 
 use crate::chunk_reader_impl::OrcChunkReader;
 use crate::hashable_schema::HashableSchema;
+use crate::processors::source::ReadingFile;
 use crate::strip::StripeInMemory;
 use crate::utils::map_orc_error;
 
 pub struct ORCSourceForCopy {
     table_ctx: Arc<dyn TableContext>,
     scan_progress: Arc<Progress>,
     op: Operator,
-    reader: Option<(
-        String,
-        Box<StripeFactory<OrcChunkReader>>,
-        HashableSchema,
-        usize,
-    )>,
+    reader: Option<ReadingFile>,
 }
 
 impl ORCSourceForCopy {
@@ -85,10 +80,16 @@ impl ORCSourceForCopy {
             .map_err(|e| map_orc_error(e, &path))?;
         let reader = builder.build_async();
         let (factory, schema) = reader.into_parts();
-        let factory = factory.expect("factory must has been created");
+        let stripe_factory = factory.expect("factory must has been created");
         let schema = HashableSchema::try_create(schema)?;
 
-        self.reader = Some((path, factory, schema, size));
+        self.reader = Some(ReadingFile {
+            path: path.to_string(),
+            stripe_factory,
+            size,
+            schema: Some(schema),
+            rows: 0,
+        });
         Ok(true)
     }
 }
@@ -105,8 +106,9 @@ impl AsyncSource for ORCSourceForCopy {
                 return Ok(None);
             }
             let start = Instant::now();
-            if let Some((path, factory, schema, size)) = mem::take(&mut self.reader) {
-                let (factory, stripe) = factory
+            if let Some(file) = mem::take(&mut self.reader) {
+                let (factory, stripe) = file
+                    .stripe_factory
                     .read_next_stripe()
                     .await
                     .map_err(|e| ErrorCode::StorageOther(e.to_string()))?;
@@ -118,24 +120,31 @@ impl AsyncSource for ORCSourceForCopy {
                     }
                     Some(stripe) => {
                         let used = start.elapsed().as_secs_f32();
+                        let rows = stripe.number_of_rows();
+
                         let bytes = stripe.stream_map().inner.values().map(|b| b.len()).sum();
-                        let progress_values = ProgressValues {
-                            rows: stripe.number_of_rows(),
-                            bytes,
-                        };
+                        let progress_values = ProgressValues { rows, bytes };
                         Profile::record_usize_profile(ProfileStatisticsName::ScanBytes, bytes);
                         log::info!(
-                            "read new stripe of {} rows and {bytes} bytes from {path}, use {} secs",
+                            "read new stripe of {} rows and {bytes} bytes from {}, use {} secs",
                             stripe.number_of_rows(),
+                            file.path,
                             used
                         );
                         self.scan_progress.incr(&progress_values);
 
-                        self.reader = Some((path.clone(), Box::new(factory), schema.clone(), size));
+                        self.reader = Some(ReadingFile {
+                            path: file.path.clone(),
+                            stripe_factory: Box::new(factory),
+                            size: file.size,
+                            schema: file.schema.clone(),
+                            rows: (rows as u64) + file.rows,
+                        });
                         let meta = Box::new(StripeInMemory {
-                            path,
+                            path: file.path.clone(),
                             stripe,
-                            schema: Some(schema),
+                            schema: file.schema,
+                            start_row: file.rows,
                         });
                         return Ok(Some(DataBlock::empty_with_meta(meta)));
                     }

src/query/storages/orc/src/processors/decoder.rs

@@ -15,6 +15,7 @@
 use std::sync::Arc;
 
 use arrow_array::RecordBatch;
+use databend_common_catalog::plan::InternalColumnType;
 use databend_common_catalog::query_kind::QueryKind;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
@@ -24,6 +25,7 @@ use databend_common_expression::DataSchema;
 use databend_common_pipeline_transforms::processors::AccumulatingTransform;
 use databend_common_storage::CopyStatus;
 use databend_common_storage::FileStatus;
+use databend_storages_common_stage::add_internal_columns;
 use orc_rust::array_decoder::NaiveStripeDecoder;
 
 use crate::strip::StripeInMemory;
@@ -33,13 +35,15 @@ pub struct StripeDecoder {
     data_schema: Arc<DataSchema>,
     arrow_schema: arrow_schema::SchemaRef,
     copy_status: Option<Arc<CopyStatus>>,
+    internal_columns: Vec<InternalColumnType>,
 }
 
 impl StripeDecoder {
     pub fn new(
         table_ctx: Arc<dyn TableContext>,
         data_schema: Arc<DataSchema>,
         arrow_schema: arrow_schema::SchemaRef,
+        internal_columns: Vec<InternalColumnType>,
     ) -> Self {
         let copy_status = if matches!(table_ctx.get_query_kind(), QueryKind::CopyIntoTable) {
             Some(table_ctx.get_copy_status())
@@ -50,6 +54,7 @@ impl StripeDecoder {
             copy_status,
             arrow_schema,
             data_schema,
+            internal_columns,
         }
     }
 }
@@ -68,14 +73,22 @@ impl AccumulatingTransform for StripeDecoder {
         let batches: std::result::Result<Vec<RecordBatch>, _> = decoder.into_iter().collect();
         let batches = batches.map_err(|e| map_orc_error(e, &stripe.path))?;
         let mut blocks = vec![];
+        let mut start_row = stripe.start_row;
+
         for batch in batches {
-            let (block, _) = DataBlock::from_record_batch(self.data_schema.as_ref(), &batch)?;
+            let (mut block, _) = DataBlock::from_record_batch(self.data_schema.as_ref(), &batch)?;
             if let Some(copy_status) = &self.copy_status {
                 copy_status.add_chunk(&stripe.path, FileStatus {
                     num_rows_loaded: block.num_rows(),
                     error: None,
                 })
             }
+            add_internal_columns(
+                &self.internal_columns,
+                stripe.path.clone(),
+                &mut block,
+                &mut start_row,
+            );
             blocks.push(block);
         }
         Ok(blocks)

src/query/storages/orc/src/processors/source.rs

@@ -65,10 +65,11 @@ impl InferredSchema {
 }
 
 pub struct ReadingFile {
-    path: String,
-    stripe_factory: Box<StripeFactory<OrcChunkReader>>,
-    size: usize,
-    schema: Option<HashableSchema>,
+    pub path: String,
+    pub stripe_factory: Box<StripeFactory<OrcChunkReader>>,
+    pub size: usize,
+    pub schema: Option<HashableSchema>,
+    pub rows: u64,
 }
 
 pub struct ORCSource {
@@ -162,6 +163,7 @@ impl ORCSource {
             stripe_factory,
             size,
             schema,
+            rows: 0,
         });
         Ok(true)
     }
@@ -197,22 +199,22 @@ impl AsyncSource for ORCSource {
                         continue;
                     }
                     Some(stripe) => {
-                        let progress_values = ProgressValues {
-                            rows: stripe.number_of_rows(),
-                            bytes: 0,
-                        };
+                        let rows = stripe.number_of_rows();
+                        let progress_values = ProgressValues { rows, bytes: 0 };
                         self.scan_progress.incr(&progress_values);
 
                         let meta = Box::new(StripeInMemory {
                             path: file.path.clone(),
                             stripe,
                             schema: file.schema.clone(),
+                            start_row: file.rows,
                         });
                         self.reader = Some(ReadingFile {
                             path: file.path.clone(),
                             stripe_factory: Box::new(factory),
                             size: file.size,
                             schema: file.schema.clone(),
+                            rows: (rows as u64) + file.rows,
                         });
                         return Ok(Some(DataBlock::empty_with_meta(meta)));
                     }

src/query/storages/orc/src/processors/variant_decoder.rs

@@ -15,6 +15,7 @@
 use std::sync::Arc;
 
 use arrow_array::RecordBatch;
+use databend_common_catalog::plan::InternalColumnType;
 use databend_common_catalog::query_kind::QueryKind;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::ErrorCode;
@@ -26,6 +27,7 @@ use databend_common_expression::TableSchema;
 use databend_common_pipeline_transforms::processors::AccumulatingTransform;
 use databend_common_storage::CopyStatus;
 use databend_common_storage::FileStatus;
+use databend_storages_common_stage::add_internal_columns;
 use databend_storages_common_stage::record_batch_to_variant_block;
 use jiff::tz::TimeZone;
 use orc_rust::array_decoder::NaiveStripeDecoder;
@@ -36,16 +38,25 @@ use crate::utils::map_orc_error;
 pub struct StripeDecoderForVariantTable {
     copy_status: Option<Arc<CopyStatus>>,
     tz: TimeZone,
+    internal_columns: Vec<InternalColumnType>,
 }
 
 impl StripeDecoderForVariantTable {
-    pub fn new(table_ctx: Arc<dyn TableContext>, tz: TimeZone) -> Self {
+    pub fn new(
+        table_ctx: Arc<dyn TableContext>,
+        tz: TimeZone,
+        internal_columns: Vec<InternalColumnType>,
+    ) -> Self {
         let copy_status = if matches!(table_ctx.get_query_kind(), QueryKind::CopyIntoTable) {
             Some(table_ctx.get_copy_status())
         } else {
             None
         };
-        StripeDecoderForVariantTable { copy_status, tz }
+        StripeDecoderForVariantTable {
+            copy_status,
+            tz,
+            internal_columns,
+        }
     }
 }
 
@@ -75,14 +86,22 @@ impl AccumulatingTransform for StripeDecoderForVariantTable {
         let batches = batches.map_err(|e| map_orc_error(e, &stripe.path))?;
 
         let mut blocks = vec![];
+        let mut start_row = stripe.start_row;
         for batch in batches {
-            let block = record_batch_to_variant_block(batch, &self.tz, &typ, &schemas.data_schema)?;
+            let mut block =
+                record_batch_to_variant_block(batch, &self.tz, &typ, &schemas.data_schema)?;
             if let Some(copy_status) = &self.copy_status {
                 copy_status.add_chunk(&stripe.path, FileStatus {
                     num_rows_loaded: block.num_rows(),
                     error: None,
                 })
             }
+            add_internal_columns(
+                &self.internal_columns,
+                stripe.path.clone(),
+                &mut block,
+                &mut start_row,
+            );
             blocks.push(block);
         }
         Ok(blocks)