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    /// Schema exposed by recursive self-references while planning the recursive term.
40    ///
41    /// This is a conservative work-table schema, not the final recursive query output
42    /// schema. For example, the SQL planner may mark fields nullable here so recursive
43    /// references do not inherit unsound anchor-term nullability assumptions.
44    table_schema: SchemaRef,
45}
46
47impl CteWorkTable {
48    /// Construct a new CteWorkTable with the given name and self-reference schema.
49    pub fn new(name: &str, table_schema: SchemaRef) -> Self {
50        Self {
51            name: name.to_owned(),
52            table_schema,
53        }
54    }
55
56    /// The user-provided name of the CTE
57    pub fn name(&self) -> &str {
58        &self.name
59    }
60
61    /// The schema exposed by scans of the recursive self-reference.
62    pub fn schema(&self) -> SchemaRef {
63        Arc::clone(&self.table_schema)
64    }
65}
66
67#[async_trait]
68impl TableProvider for CteWorkTable {
69    fn get_logical_plan(&'_ self) -> Option<Cow<'_, LogicalPlan>> {
70        None
71    }
72
73    fn schema(&self) -> SchemaRef {
74        Arc::clone(&self.table_schema)
75    }
76
77    fn table_type(&self) -> TableType {
78        TableType::Temporary
79    }
80
81    async fn scan(
82        &self,
83        state: &dyn Session,
84        projection: Option<&Vec<usize>>,
85        filters: &[Expr],
86        limit: Option<usize>,
87    ) -> Result<Arc<dyn ExecutionPlan>> {
88        let options = ScanArgs::default()
89            .with_projection(projection.map(|p| p.as_slice()))
90            .with_filters(Some(filters))
91            .with_limit(limit);
92        Ok(self.scan_with_args(state, options).await?.into_inner())
93    }
94
95    async fn scan_with_args<'a>(
96        &self,
97        _state: &dyn Session,
98        args: ScanArgs<'a>,
99    ) -> Result<ScanResult> {
100        Ok(ScanResult::new(Arc::new(WorkTableExec::new(
101            self.name.clone(),
102            Arc::clone(&self.table_schema),
103            args.projection().map(|p| p.to_vec()),
104        )?)))
105    }
106
107    fn supports_filters_pushdown(
108        &self,
109        filters: &[&Expr],
110    ) -> Result<Vec<TableProviderFilterPushDown>> {
111        // TODO: should we support filter pushdown?
112        Ok(vec![
113            TableProviderFilterPushDown::Unsupported;
114            filters.len()
115        ])
116    }
117}
datafusion_catalog/cte_worktable.rs

datafusion_catalog/
cte_worktable.rs
