feat: support query ORC file as rows of variants. (#18415)

Author: youngsofun

Cargo.lock

@@ -21,8 +21,7 @@ use crate::plan::StageTableInfo;
 #[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, Eq, Debug)]
 pub struct OrcTableInfo {
     pub stage_table_info: StageTableInfo,
-    pub arrow_schema: arrow_schema::SchemaRef,
-    pub schema_from: String,
+    pub schema: Option<(arrow_schema::SchemaRef, String)>,
 }
 
 impl OrcTableInfo {

src/query/catalog/src/plan/datasource/datasource_info/orc.rs

@@ -1774,6 +1774,14 @@ impl TableContext for QueryContext {
                 }
             }
             FileFormatParams::Orc(..) => {
+                let is_variant =
+                    match max_column_position {
+                        0 => false,
+                        1 => true,
+                        _ => return Err(ErrorCode::SemanticError(
+                            "[QUERY-CTX] Query from ORC file only support $1 as column position",
+                        )),
+                    };
                 let schema = Arc::new(TableSchema::empty());
                 let info = StageTableInfo {
                     schema,
@@ -1785,7 +1793,7 @@ impl TableContext for QueryContext {
                     default_exprs: None,
                     copy_into_table_options: Default::default(),
                     stage_root,
-                    is_variant: false,
+                    is_variant,
                 };
                 OrcTable::try_create(info).await
             }

src/query/service/src/sessions/query_ctx.rs

@@ -7,12 +7,14 @@ publish = { workspace = true }
 edition = { workspace = true }
 
 [dependencies]
+arrow-array = { workspace = true }
 databend-common-ast = { workspace = true }
 databend-common-catalog = { workspace = true }
 databend-common-exception = { workspace = true }
 databend-common-expression = { workspace = true }
 databend-common-functions = { workspace = true }
 databend-common-meta-app = { workspace = true }
+jiff = { workspace = true }
 serde = { workspace = true }
 typetag = { workspace = true }
 

src/query/storages/common/stage/Cargo.toml

@@ -21,7 +21,7 @@ use databend_common_expression::DataSchema;
 use databend_common_expression::TableDataType;
 use jiff::tz::TimeZone;
 
-pub fn read_record_batch(
+pub fn read_record_batch_to_variant_column(
     record_batch: RecordBatch,
     builder: &mut BinaryColumnBuilder,
     tz: &TimeZone,
@@ -40,7 +40,7 @@ pub fn read_record_batch(
     Ok(())
 }
 
-pub fn record_batch_to_block(
+pub fn record_batch_to_variant_block(
     record_batch: RecordBatch,
     tz: &TimeZone,
     typ: &TableDataType,
@@ -50,7 +50,7 @@ pub fn record_batch_to_block(
         record_batch.num_rows(),
         record_batch.get_array_memory_size(),
     );
-    read_record_batch(record_batch, &mut builder, tz, typ, schema)?;
+    read_record_batch_to_variant_column(record_batch, &mut builder, tz, typ, schema)?;
     let column = builder.build();
     let num_rows = column.len();
     Ok(DataBlock::new(

src/query/storages/parquet/src/parquet_variant_table/recordbatch_to_variant.rs renamed to src/query/storages/common/stage/src/read/columnar/arrow_to_variant.rs

@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub mod arrow_to_variant;
 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 projection::project_columnar;

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

@@ -320,7 +320,7 @@ impl IcebergTable {
             let data_schema = Arc::new(data_schema);
             pipeline.add_source(
                 |output| {
-                    ORCSource::try_create(
+                    ORCSource::try_create_with_schema(
                         output,
                         ctx.clone(),
                         Arc::new(op.clone()),

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

@@ -31,6 +31,7 @@ bytes = { workspace = true }
 chrono = { workspace = true }
 dashmap = { workspace = true }
 futures-util = { workspace = true }
+jiff = { workspace = true }
 log = { workspace = true }
 opendal = { workspace = true }
 orc-rust = { workspace = true }

src/query/storages/orc/Cargo.toml

@@ -14,3 +14,4 @@
 
 pub(crate) mod decoder;
 pub(crate) mod source;
+pub(crate) mod variant_decoder;

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

@@ -35,22 +35,52 @@ use orc_rust::projection::ProjectionMask;
 use orc_rust::ArrowReaderBuilder;
 
 use crate::chunk_reader_impl::OrcChunkReader;
+use crate::hashable_schema::HashableSchema;
 use crate::strip::StripeInMemory;
 use crate::utils::map_orc_error;
 
+pub struct InferredSchema {
+    arrow_schema: arrow_schema::SchemaRef,
+    schema_from: Option<String>,
+    projection: Projection,
+}
+
+impl InferredSchema {
+    fn check_file_schema(&self, arrow_schema: arrow_schema::SchemaRef, path: &str) -> Result<()> {
+        if self.arrow_schema.fields != arrow_schema.fields {
+            return Err(ErrorCode::TableSchemaMismatch(format!(
+                "{} get diff schema in file '{}'. Expected schema: {:?}, actual: {:?}",
+                self.schema_from
+                    .as_ref()
+                    .map_or(String::new(), |schema_from| {
+                        format!("infer schema from '{}', but ", schema_from)
+                    }),
+                path,
+                self.arrow_schema,
+                arrow_schema
+            )));
+        }
+        Ok(())
+    }
+}
+
+pub struct ReadingFile {
+    path: String,
+    stripe_factory: Box<StripeFactory<OrcChunkReader>>,
+    size: usize,
+    schema: Option<HashableSchema>,
+}
+
 pub struct ORCSource {
     table_ctx: Arc<dyn TableContext>,
     op_registry: Arc<dyn OperatorRegistry>,
-    pub(crate) reader: Option<(String, Box<StripeFactory<OrcChunkReader>>, usize)>,
+    pub reader: Option<ReadingFile>,
     scan_progress: Arc<Progress>,
-
-    arrow_schema: arrow_schema::SchemaRef,
-    schema_from: Option<String>,
-    projection: Projection,
+    inferred_schema: Option<InferredSchema>,
 }
 
 impl ORCSource {
-    pub fn try_create(
+    pub fn try_create_with_schema(
         output: Arc<OutputPort>,
         table_ctx: Arc<dyn TableContext>,
         op_registry: Arc<dyn OperatorRegistry>,
@@ -65,27 +95,29 @@ impl ORCSource {
             op_registry,
             scan_progress,
             reader: None,
-            arrow_schema,
-            schema_from,
-            projection,
+            inferred_schema: Some(InferredSchema {
+                arrow_schema,
+                schema_from,
+                projection,
+            }),
         })
     }
 
-    fn check_file_schema(&self, arrow_schema: arrow_schema::SchemaRef, path: &str) -> Result<()> {
-        if self.arrow_schema.fields != arrow_schema.fields {
-            return Err(ErrorCode::TableSchemaMismatch(format!(
-                "{}get diff schema in file '{}'. Expected schema: {:?}, actual: {:?}",
-                self.schema_from
-                    .as_ref()
-                    .map_or(String::new(), |schema_from| {
-                        format!("infer schema from '{}', but ", schema_from)
-                    }),
-                path,
-                self.arrow_schema,
-                arrow_schema
-            )));
-        }
-        Ok(())
+    pub fn try_create(
+        output: Arc<OutputPort>,
+        table_ctx: Arc<dyn TableContext>,
+        op_registry: Arc<dyn OperatorRegistry>,
+        inferred_schema: Option<InferredSchema>,
+    ) -> Result<ProcessorPtr> {
+        let scan_progress = table_ctx.get_scan_progress();
+
+        AsyncSourcer::create(table_ctx.clone(), output, ORCSource {
+            table_ctx,
+            op_registry,
+            scan_progress,
+            reader: None,
+            inferred_schema,
+        })
     }
 
     async fn next_part(&mut self) -> Result<bool> {
@@ -105,20 +137,32 @@ impl ORCSource {
         let builder = ArrowReaderBuilder::try_new_async(file)
             .await
             .map_err(|e| map_orc_error(e, path))?;
-        let projection = if let Projection::Columns(projection) = &self.projection {
-            ProjectionMask::roots(
-                builder.file_metadata().root_data_type(),
-                projection.iter().map(|index| index + 1),
-            )
-        } else {
-            ProjectionMask::all()
-        };
+        let mut projection = ProjectionMask::all();
+        if let Some(schema) = &self.inferred_schema {
+            if let Projection::Columns(p) = &schema.projection {
+                projection = ProjectionMask::roots(
+                    builder.file_metadata().root_data_type(),
+                    p.iter().map(|index| index + 1),
+                );
+            }
+        }
+
         let reader = builder.with_projection(projection).build_async();
         let (factory, schema) = reader.into_parts();
-        let factory = factory.unwrap();
-        self.check_file_schema(schema, path)?;
+        let stripe_factory = factory.unwrap();
+        let schema = if let Some(inferred_schema) = &self.inferred_schema {
+            inferred_schema.check_file_schema(schema, path)?;
+            None
+        } else {
+            Some(HashableSchema::try_create(schema)?)
+        };
 
-        self.reader = Some((path.to_string(), factory, size));
+        self.reader = Some(ReadingFile {
+            path: path.to_string(),
+            stripe_factory,
+            size,
+            schema,
+        });
         Ok(true)
     }
 }
@@ -134,8 +178,9 @@ impl AsyncSource for ORCSource {
             if self.reader.is_none() && !self.next_part().await? {
                 return Ok(None);
             }
-            if let Some((path, factory, 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()))?;
@@ -144,10 +189,10 @@ impl AsyncSource for ORCSource {
                         self.reader = None;
                         let progress_values = ProgressValues {
                             rows: 0,
-                            bytes: size,
+                            bytes: file.size,
                         };
                         self.scan_progress.incr(&progress_values);
-                        Profile::record_usize_profile(ProfileStatisticsName::ScanBytes, size);
+                        Profile::record_usize_profile(ProfileStatisticsName::ScanBytes, file.size);
                         Profile::record_usize_profile(ProfileStatisticsName::ScanPartitions, 1);
                         continue;
                     }
@@ -157,11 +202,17 @@ impl AsyncSource for ORCSource {
                             bytes: 0,
                         };
                         self.scan_progress.incr(&progress_values);
-                        self.reader = Some((path.clone(), Box::new(factory), size));
+
                         let meta = Box::new(StripeInMemory {
-                            path,
+                            path: file.path.clone(),
                             stripe,
-                            schema: None,
+                            schema: file.schema.clone(),
+                        });
+                        self.reader = Some(ReadingFile {
+                            path: file.path.clone(),
+                            stripe_factory: Box::new(factory),
+                            size: file.size,
+                            schema: file.schema.clone(),
                         });
                         return Ok(Some(DataBlock::empty_with_meta(meta)));
                     }