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::DataFusionError;
20use datafusion::datasource::physical_plan::{FileOpener, FileScanConfig, FileSource};
21use datafusion::datasource::table_schema::TableSchema;
22use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
23use datafusion::physical_plan::projection::ProjectionExprs;
24use object_store::ObjectStore;
25use std::sync::Arc;
26
27#[derive(Debug, Clone)]
28pub struct OrcSource {
29    metrics: ExecutionPlanMetricsSet,
30    batch_size: usize,
31    table_schema: TableSchema,
32    projection: ProjectionExprs,
33}
34
35impl OrcSource {
36    pub fn new(table_schema: TableSchema) -> Self {
37        let table_schema_ref = table_schema.table_schema();
38        let projection = ProjectionExprs::from_indices(
39            &(0..table_schema_ref.fields().len()).collect::<Vec<_>>(),
40            table_schema_ref,
41        );
42        Self {
43            metrics: ExecutionPlanMetricsSet::default(),
44            batch_size: 1024,
45            table_schema,
46            projection,
47        }
48    }
49}
50
51impl FileSource for OrcSource {
52    fn create_file_opener(
53        &self,
54        object_store: Arc<dyn ObjectStore>,
55        config: &FileScanConfig,
56        _partition: usize,
57    ) -> Result<Arc<dyn FileOpener>, DataFusionError> {
58        OrcOpener::try_new(
59            object_store,
60            self.table_schema.table_schema().clone(),
61            config.batch_size.unwrap_or(self.batch_size),
62            self.projection.clone(),
63        )
64        .map(|f| Arc::new(f) as Arc<dyn FileOpener>)
65    }
66
67    fn table_schema(&self) -> &TableSchema {
68        &self.table_schema
69    }
70
71    fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
72        Arc::new(Self {
73            batch_size,
74            ..self.clone()
75        })
76    }
77
78    fn projection(&self) -> Option<&ProjectionExprs> {
79        Some(&self.projection)
80    }
81
82    fn metrics(&self) -> &ExecutionPlanMetricsSet {
83        &self.metrics
84    }
85
86    fn file_type(&self) -> &str {
87        "orc"
88    }
89
90    fn try_pushdown_projection(
91        &self,
92        projection: &ProjectionExprs,
93    ) -> Result<Option<Arc<dyn FileSource>>, DataFusionError> {
94        let mut source = self.clone();
95        source.projection = self.projection.try_merge(projection)?;
96        Ok(Some(Arc::new(source)))
97    }
98}
datafusion_orc/file_source.rs

datafusion_orc/
file_source.rs
