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

datafusion_catalog/
cte_worktable.rs
