datafusion_catalog/
cte_worktable.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//! CteWorkTable implementation used for recursive queries
19
20use std::sync::Arc;
21use std::{any::Any, borrow::Cow};
22
23use crate::Session;
24use arrow::datatypes::SchemaRef;
25use async_trait::async_trait;
26use datafusion_physical_plan::work_table::WorkTableExec;
27
28use datafusion_physical_plan::ExecutionPlan;
29
30use datafusion_common::error::Result;
31use datafusion_expr::{Expr, LogicalPlan, TableProviderFilterPushDown, TableType};
32
33use crate::TableProvider;
34
35/// The temporary working table where the previous iteration of a recursive query is stored
36/// Naming is based on PostgreSQL's implementation.
37/// See here for more details: www.postgresql.org/docs/11/queries-with.html#id-1.5.6.12.5.4
38#[derive(Debug)]
39pub struct CteWorkTable {
40    /// The name of the CTE work table
41    name: String,
42    /// This schema must be shared across both the static and recursive terms of a recursive query
43    table_schema: SchemaRef,
44}
45
46impl CteWorkTable {
47    /// construct a new CteWorkTable with the given name and schema
48    /// This schema must match the schema of the recursive term of the query
49    /// Since the scan method will contain an physical plan that assumes this schema
50    pub fn new(name: &str, table_schema: SchemaRef) -> Self {
51        Self {
52            name: name.to_owned(),
53            table_schema,
54        }
55    }
56
57    /// The user-provided name of the CTE
58    pub fn name(&self) -> &str {
59        &self.name
60    }
61
62    /// The schema of the recursive term of the query
63    pub fn schema(&self) -> SchemaRef {
64        Arc::clone(&self.table_schema)
65    }
66}
67
68#[async_trait]
69impl TableProvider for CteWorkTable {
70    fn as_any(&self) -> &dyn Any {
71        self
72    }
73
74    fn get_logical_plan(&'_ self) -> Option<Cow<'_, LogicalPlan>> {
75        None
76    }
77
78    fn schema(&self) -> SchemaRef {
79        Arc::clone(&self.table_schema)
80    }
81
82    fn table_type(&self) -> TableType {
83        TableType::Temporary
84    }
85
86    async fn scan(
87        &self,
88        _state: &dyn Session,
89        _projection: Option<&Vec<usize>>,
90        _filters: &[Expr],
91        _limit: Option<usize>,
92    ) -> Result<Arc<dyn ExecutionPlan>> {
93        // TODO: pushdown filters and limits
94        Ok(Arc::new(WorkTableExec::new(
95            self.name.clone(),
96            Arc::clone(&self.table_schema),
97        )))
98    }
99
100    fn supports_filters_pushdown(
101        &self,
102        filters: &[&Expr],
103    ) -> Result<Vec<TableProviderFilterPushDown>> {
104        // TODO: should we support filter pushdown?
105        Ok(vec![
106            TableProviderFilterPushDown::Unsupported;
107            filters.len()
108        ])
109    }
110}
datafusion_catalog/cte_worktable.rs

datafusion_catalog/
cte_worktable.rs
