datafusion_catalog/
view.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
18//! View data source which uses a LogicalPlan as it's input.
19
20use std::{borrow::Cow, sync::Arc};
21
22use crate::Session;
23use crate::TableProvider;
24
25use arrow::datatypes::SchemaRef;
26use async_trait::async_trait;
27use datafusion_common::Column;
28use datafusion_common::error::Result;
29use datafusion_expr::TableType;
30use datafusion_expr::{Expr, LogicalPlan};
31use datafusion_expr::{LogicalPlanBuilder, TableProviderFilterPushDown};
32use datafusion_physical_plan::ExecutionPlan;
33
34/// An implementation of `TableProvider` that uses another logical plan.
35#[derive(Debug)]
36pub struct ViewTable {
37    /// LogicalPlan of the view
38    logical_plan: LogicalPlan,
39    /// File fields + partition columns
40    table_schema: SchemaRef,
41    /// SQL used to create the view, if available
42    definition: Option<String>,
43}
44
45impl ViewTable {
46    /// Create new view that is executed at query runtime.
47    ///
48    /// Takes a `LogicalPlan` and optionally the SQL text of the `CREATE`
49    /// statement.
50    ///
51    /// Notes: the `LogicalPlan` is not validated or type coerced. If this is
52    /// needed it should be done after calling this function.
53    pub fn new(logical_plan: LogicalPlan, definition: Option<String>) -> Self {
54        let table_schema = Arc::clone(logical_plan.schema().inner());
55        Self {
56            logical_plan,
57            table_schema,
58            definition,
59        }
60    }
61
62    /// Get definition ref
63    pub fn definition(&self) -> Option<&String> {
64        self.definition.as_ref()
65    }
66
67    /// Get logical_plan ref
68    pub fn logical_plan(&self) -> &LogicalPlan {
69        &self.logical_plan
70    }
71}
72
73#[async_trait]
74impl TableProvider for ViewTable {
75    fn get_logical_plan(&'_ self) -> Option<Cow<'_, LogicalPlan>> {
76        Some(Cow::Borrowed(&self.logical_plan))
77    }
78
79    fn schema(&self) -> SchemaRef {
80        Arc::clone(&self.table_schema)
81    }
82
83    fn table_type(&self) -> TableType {
84        TableType::View
85    }
86
87    fn get_table_definition(&self) -> Option<&str> {
88        self.definition.as_deref()
89    }
90    fn supports_filters_pushdown(
91        &self,
92        filters: &[&Expr],
93    ) -> Result<Vec<TableProviderFilterPushDown>> {
94        // A filter is added on the View when given
95        Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
96    }
97
98    async fn scan(
99        &self,
100        state: &dyn Session,
101        projection: Option<&Vec<usize>>,
102        filters: &[Expr],
103        limit: Option<usize>,
104    ) -> Result<Arc<dyn ExecutionPlan>> {
105        let filter = filters.iter().cloned().reduce(|acc, new| acc.and(new));
106        let plan = self.logical_plan().clone();
107        let mut plan = LogicalPlanBuilder::from(plan);
108
109        if let Some(filter) = filter {
110            plan = plan.filter(filter)?;
111        }
112
113        let mut plan = if let Some(projection) = projection {
114            // avoiding adding a redundant projection (e.g. SELECT * FROM view)
115            let current_projection =
116                (0..plan.schema().fields().len()).collect::<Vec<usize>>();
117            if projection == &current_projection {
118                plan
119            } else {
120                let fields: Vec<Expr> = projection
121                    .iter()
122                    .map(|i| {
123                        Expr::Column(Column::from(
124                            self.logical_plan.schema().qualified_field(*i),
125                        ))
126                    })
127                    .collect();
128                plan.project(fields)?
129            }
130        } else {
131            plan
132        };
133
134        if let Some(limit) = limit {
135            plan = plan.limit(0, Some(limit))?;
136        }
137
138        state.create_physical_plan(&plan.build()?).await
139    }
140}
datafusion_catalog/view.rs

datafusion_catalog/
view.rs
