datafusion_orc/
file_source.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::physical_exec::OrcOpener;
19use datafusion::common::tree_node::TreeNodeRecursion;
20use datafusion::common::DataFusionError;
21use datafusion::datasource::physical_plan::{FileOpener, FileScanConfig, FileSource};
22use datafusion::datasource::table_schema::TableSchema;
23use datafusion::object_store::ObjectStore;
24use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
25use datafusion::physical_plan::projection::ProjectionExprs;
26use datafusion::physical_plan::PhysicalExpr;
27use std::sync::Arc;
28
29#[derive(Debug, Clone)]
30pub struct OrcSource {
31    metrics: ExecutionPlanMetricsSet,
32    batch_size: usize,
33    table_schema: TableSchema,
34    projection: ProjectionExprs,
35}
36
37impl OrcSource {
38    pub fn new(table_schema: TableSchema) -> Self {
39        let table_schema_ref = table_schema.table_schema();
40        let projection = ProjectionExprs::from_indices(
41            &(0..table_schema_ref.fields().len()).collect::<Vec<_>>(),
42            table_schema_ref,
43        );
44        Self {
45            metrics: ExecutionPlanMetricsSet::default(),
46            batch_size: 1024,
47            table_schema,
48            projection,
49        }
50    }
51}
52
53impl FileSource for OrcSource {
54    fn create_file_opener(
55        &self,
56        object_store: Arc<dyn ObjectStore>,
57        config: &FileScanConfig,
58        _partition: usize,
59    ) -> Result<Arc<dyn FileOpener>, DataFusionError> {
60        OrcOpener::try_new(
61            object_store,
62            self.table_schema.table_schema().clone(),
63            config.batch_size.unwrap_or(self.batch_size),
64            self.projection.clone(),
65        )
66        .map(|f| Arc::new(f) as Arc<dyn FileOpener>)
67    }
68
69    fn table_schema(&self) -> &TableSchema {
70        &self.table_schema
71    }
72
73    fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
74        Arc::new(Self {
75            batch_size,
76            ..self.clone()
77        })
78    }
79
80    fn projection(&self) -> Option<&ProjectionExprs> {
81        Some(&self.projection)
82    }
83
84    fn metrics(&self) -> &ExecutionPlanMetricsSet {
85        &self.metrics
86    }
87
88    fn file_type(&self) -> &str {
89        "orc"
90    }
91
92    fn try_pushdown_projection(
93        &self,
94        projection: &ProjectionExprs,
95    ) -> Result<Option<Arc<dyn FileSource>>, DataFusionError> {
96        let mut source = self.clone();
97        source.projection = self.projection.try_merge(projection)?;
98        Ok(Some(Arc::new(source)))
99    }
100
101    fn apply_expressions(
102        &self,
103        f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion, DataFusionError>,
104    ) -> Result<TreeNodeRecursion, DataFusionError> {
105        datafusion::physical_plan::apply_expression_roots(self.projection.iter(), f)
106    }
107}
datafusion_orc/file_source.rs

datafusion_orc/
file_source.rs
