Continuing implementation with fixes and improvements

Lint fixes

Author: matthewgapp

datafusion-cli/Cargo.lock

@@ -753,6 +753,11 @@ impl DFField {
         self.field = f.into();
         self
     }
+
+    pub fn with_qualifier(mut self, qualifier: impl Into<OwnedTableReference>) -> Self {
+        self.qualifier = Some(qualifier.into());
+        self
+    }
 }
 
 impl From<FieldRef> for DFField {

datafusion/common/src/dfschema.rs

@@ -50,6 +50,7 @@ use crate::physical_expr::create_physical_expr;
 use crate::physical_optimizer::optimizer::PhysicalOptimizerRule;
 use crate::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
 use crate::physical_plan::analyze::AnalyzeExec;
+use crate::physical_plan::continuance::ContinuanceExec;
 use crate::physical_plan::explain::ExplainExec;
 use crate::physical_plan::expressions::{Column, PhysicalSortExpr};
 use crate::physical_plan::filter::FilterExec;
@@ -58,6 +59,7 @@ use crate::physical_plan::joins::SortMergeJoinExec;
 use crate::physical_plan::joins::{CrossJoinExec, NestedLoopJoinExec};
 use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use crate::physical_plan::projection::ProjectionExec;
+use crate::physical_plan::recursive_query::RecursiveQueryExec;
 use crate::physical_plan::repartition::RepartitionExec;
 use crate::physical_plan::sorts::sort::SortExec;
 use crate::physical_plan::unnest::UnnestExec;
@@ -87,7 +89,9 @@ use datafusion_expr::expr::{
 };
 use datafusion_expr::expr_rewriter::{unalias, unnormalize_cols};
 use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary;
-use datafusion_expr::{DescribeTable, DmlStatement, StringifiedPlan, WriteOp};
+use datafusion_expr::{
+    DescribeTable, DmlStatement, NamedRelation, RecursiveQuery, StringifiedPlan, WriteOp,
+};
 use datafusion_expr::{WindowFrame, WindowFrameBound};
 use datafusion_physical_expr::expressions::Literal;
 use datafusion_sql::utils::window_expr_common_partition_keys;
@@ -1316,6 +1320,29 @@ impl DefaultPhysicalPlanner {
                         Ok(plan)
                     }
                 }
+                LogicalPlan::RecursiveQuery(RecursiveQuery { name, static_term, recursive_term, is_distinct }) => {
+                    let static_term = self.create_initial_plan(static_term, session_state).await?;
+                    let recursive_term = self.create_initial_plan(recursive_term, session_state).await?;
+
+                    Ok(Arc::new(RecursiveQueryExec::new(name.clone(), static_term, recursive_term, *is_distinct)))
+                }
+                LogicalPlan::NamedRelation(NamedRelation {name, schema}) => {
+                    // Named relations is how we represent access to any sort of dynamic data provider. They
+                    // differ from tables in the sense that they can start existing dynamically during the
+                    // execution of a query and then disappear before it even finishes.
+                    //
+                    // This system allows us to replicate the tricky behavior of classical databases where a
+                    // temporary "working table" (as it is called in Postgres) can be used when dealing with
+                    // complex operations (such as recursive CTEs) and then can be dropped. Since DataFusion
+                    // at its core is heavily stream-based and vectorized, we try to avoid using 'real' tables
+                    // and let the streams take care of the data flow in this as well.
+
+                    // Since the actual "input"'s will be only available to us at runtime (through task context)
+                    // we can't really do any sort of meaningful validation here.
+                    let schema = SchemaRef::new(schema.as_ref().to_owned().into());
+                    Ok(Arc::new(ContinuanceExec::new(name.clone(), schema)))
+                }
+
             };
             exec_plan
         }.boxed()

datafusion/core/src/physical_planner.rs

@@ -45,4 +45,5 @@ object_store = "0.7.0"
 parking_lot = "0.12"
 rand = "0.8"
 tempfile = "3"
+tokio = { version = "1.28", features = ["macros", "rt", "rt-multi-thread", "sync", "fs", "parking_lot"] }
 url = "2.2"

datafusion/execution/Cargo.toml

@@ -33,13 +33,14 @@ use crate::{
     runtime_env::{RuntimeConfig, RuntimeEnv},
 };
 
-use arrow::{error::Result as ArrowResult, record_batch::RecordBatch};
-use futures::channel::mpsc::Receiver as SingleChannelReceiver;
+use arrow::record_batch::RecordBatch;
+// use futures::channel::mpsc::Receiver as SingleChannelReceiver;
+use tokio::sync::mpsc::Receiver as SingleChannelReceiver;
 // use futures::lock::Mutex;
 use parking_lot::Mutex;
 // use futures::
 
-type RelationHandler = SingleChannelReceiver<ArrowResult<RecordBatch>>;
+type RelationHandler = SingleChannelReceiver<Result<RecordBatch>>;
 
 /// Task Execution Context
 ///

datafusion/execution/src/task.rs

@@ -54,6 +54,8 @@ use std::collections::{HashMap, HashSet};
 use std::convert::TryFrom;
 use std::sync::Arc;
 
+use super::plan::{NamedRelation, RecursiveQuery};
+
 /// Default table name for unnamed table
 pub const UNNAMED_TABLE: &str = "?table?";
 
@@ -120,6 +122,39 @@ impl LogicalPlanBuilder {
         }))
     }
 
+    /// A named temporary relation with a schema.
+    ///
+    /// This is used to represent a relation that does not exist at the
+    /// planning stage, but will be created at execution time with the
+    /// given schema.
+    pub fn named_relation(name: &str, schema: DFSchemaRef) -> Self {
+        Self::from(LogicalPlan::NamedRelation(NamedRelation {
+            name: name.to_string(),
+            schema,
+        }))
+    }
+
+    /// Convert a regular plan into a recursive query.
+    pub fn to_recursive_query(
+        &self,
+        name: String,
+        recursive_term: LogicalPlan,
+        is_distinct: bool,
+    ) -> Result<Self> {
+        // TODO: we need to do a bunch of validation here. Maybe more.
+        if is_distinct {
+            return Err(DataFusionError::NotImplemented(
+                "Recursive queries with distinct is not supported".to_string(),
+            ));
+        }
+        Ok(Self::from(LogicalPlan::RecursiveQuery(RecursiveQuery {
+            name,
+            static_term: Arc::new(self.plan.clone()),
+            recursive_term: Arc::new(recursive_term),
+            is_distinct,
+        })))
+    }
+
     /// Create a values list based relation, and the schema is inferred from data, consuming
     /// `value`. See the [Postgres VALUES](https://www.postgresql.org/docs/current/queries-values.html)
     /// documentation for more details.

datafusion/expr/src/logical_plan/builder.rs

@@ -34,9 +34,10 @@ pub use ddl::{
 pub use dml::{DmlStatement, WriteOp};
 pub use plan::{
     Aggregate, Analyze, CrossJoin, DescribeTable, Distinct, EmptyRelation, Explain,
-    Extension, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning,
-    PlanType, Prepare, Projection, Repartition, Sort, StringifiedPlan, Subquery,
-    SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window,
+    Extension, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, NamedRelation,
+    Partitioning, PlanType, Prepare, Projection, RecursiveQuery, Repartition, Sort,
+    StringifiedPlan, Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union,
+    Unnest, Values, Window,
 };
 pub use statement::{
     SetVariable, Statement, TransactionAccessMode, TransactionConclusion, TransactionEnd,

datafusion/expr/src/logical_plan/mod.rs

@@ -110,6 +110,8 @@ pub enum LogicalPlan {
     /// produces 0 or 1 row. This is used to implement SQL `SELECT`
     /// that has no values in the `FROM` clause.
     EmptyRelation(EmptyRelation),
+    /// A named temporary relation with a schema.
+    NamedRelation(NamedRelation),
     /// Produces the output of running another query.  This is used to
     /// implement SQL subqueries
     Subquery(Subquery),
@@ -152,6 +154,8 @@ pub enum LogicalPlan {
     /// Unnest a column that contains a nested list type such as an
     /// ARRAY. This is used to implement SQL `UNNEST`
     Unnest(Unnest),
+    /// A variadic query (e.g. "Recursive CTEs")
+    RecursiveQuery(RecursiveQuery),
 }
 
 impl LogicalPlan {
@@ -188,6 +192,11 @@ impl LogicalPlan {
             LogicalPlan::Copy(CopyTo { input, .. }) => input.schema(),
             LogicalPlan::Ddl(ddl) => ddl.schema(),
             LogicalPlan::Unnest(Unnest { schema, .. }) => schema,
+            LogicalPlan::NamedRelation(NamedRelation { schema, .. }) => schema,
+            LogicalPlan::RecursiveQuery(RecursiveQuery { static_term, .. }) => {
+                // we take the schema of the static term as the schema of the entire recursive query
+                static_term.schema()
+            }
         }
     }
 
@@ -230,6 +239,7 @@ impl LogicalPlan {
             LogicalPlan::Explain(_)
             | LogicalPlan::Analyze(_)
             | LogicalPlan::EmptyRelation(_)
+            | LogicalPlan::NamedRelation(_)
             | LogicalPlan::Ddl(_)
             | LogicalPlan::Dml(_)
             | LogicalPlan::Copy(_)
@@ -240,6 +250,10 @@ impl LogicalPlan {
             | LogicalPlan::TableScan(_) => {
                 vec![self.schema()]
             }
+            LogicalPlan::RecursiveQuery(RecursiveQuery { static_term, .. }) => {
+                // return only the schema of the static term
+                static_term.all_schemas()
+            }
             // return children schemas
             LogicalPlan::Limit(_)
             | LogicalPlan::Subquery(_)
@@ -371,6 +385,9 @@ impl LogicalPlan {
             }
             // plans without expressions
             LogicalPlan::EmptyRelation(_)
+            | LogicalPlan::NamedRelation(_)
+            // TODO: not sure if this should go here
+            | LogicalPlan::RecursiveQuery(_)
             | LogicalPlan::Subquery(_)
             | LogicalPlan::SubqueryAlias(_)
             | LogicalPlan::Limit(_)
@@ -415,8 +432,14 @@ impl LogicalPlan {
             LogicalPlan::Ddl(ddl) => ddl.inputs(),
             LogicalPlan::Unnest(Unnest { input, .. }) => vec![input],
             LogicalPlan::Prepare(Prepare { input, .. }) => vec![input],
+            LogicalPlan::RecursiveQuery(RecursiveQuery {
+                static_term,
+                recursive_term,
+                ..
+            }) => vec![static_term, recursive_term],
             // plans without inputs
             LogicalPlan::TableScan { .. }
+            | LogicalPlan::NamedRelation(_)
             | LogicalPlan::Statement { .. }
             | LogicalPlan::EmptyRelation { .. }
             | LogicalPlan::Values { .. }
@@ -492,6 +515,9 @@ impl LogicalPlan {
                     cross.left.head_output_expr()
                 }
             }
+            LogicalPlan::RecursiveQuery(RecursiveQuery { static_term, .. }) => {
+                static_term.head_output_expr()
+            }
             LogicalPlan::Union(union) => Ok(Some(Expr::Column(
                 union.schema.fields()[0].qualified_column(),
             ))),
@@ -511,6 +537,7 @@ impl LogicalPlan {
             }
             LogicalPlan::Subquery(_) => Ok(None),
             LogicalPlan::EmptyRelation(_)
+            | LogicalPlan::NamedRelation(_)
             | LogicalPlan::Prepare(_)
             | LogicalPlan::Statement(_)
             | LogicalPlan::Values(_)
@@ -839,6 +866,14 @@ impl LogicalPlan {
                     input: Arc::new(inputs[0].clone()),
                 }))
             }
+            LogicalPlan::RecursiveQuery(RecursiveQuery {
+                name, is_distinct, ..
+            }) => Ok(LogicalPlan::RecursiveQuery(RecursiveQuery {
+                name: name.clone(),
+                static_term: Arc::new(inputs[0].clone()),
+                recursive_term: Arc::new(inputs[1].clone()),
+                is_distinct: *is_distinct,
+            })),
             LogicalPlan::Analyze(a) => {
                 assert!(expr.is_empty());
                 assert_eq!(inputs.len(), 1);
@@ -877,6 +912,7 @@ impl LogicalPlan {
                 }))
             }
             LogicalPlan::EmptyRelation(_)
+            | LogicalPlan::NamedRelation(_)
             | LogicalPlan::Ddl(_)
             | LogicalPlan::Statement(_) => {
                 // All of these plan types have no inputs / exprs so should not be called
@@ -1040,6 +1076,9 @@ impl LogicalPlan {
                 }),
             LogicalPlan::TableScan(TableScan { fetch, .. }) => *fetch,
             LogicalPlan::EmptyRelation(_) => Some(0),
+            // TODO: not sure if this is correct
+            LogicalPlan::NamedRelation(_) => None,
+            LogicalPlan::RecursiveQuery(_) => None,
             LogicalPlan::Subquery(_) => None,
             LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => input.max_rows(),
             LogicalPlan::Limit(Limit { fetch, .. }) => *fetch,
@@ -1387,6 +1426,14 @@ impl LogicalPlan {
             fn fmt(&self, f: &mut Formatter) -> fmt::Result {
                 match self.0 {
                     LogicalPlan::EmptyRelation(_) => write!(f, "EmptyRelation"),
+                    LogicalPlan::NamedRelation(NamedRelation { name, .. }) => {
+                        write!(f, "NamedRelation: {}", name)
+                    }
+                    LogicalPlan::RecursiveQuery(RecursiveQuery {
+                        is_distinct, ..
+                    }) => {
+                        write!(f, "RecursiveQuery: is_distinct={}", is_distinct)
+                    }
                     LogicalPlan::Values(Values { ref values, .. }) => {
                         let str_values: Vec<_> = values
                             .iter()
@@ -1685,6 +1732,28 @@ pub struct EmptyRelation {
     pub schema: DFSchemaRef,
 }
 
+/// A named temporary relation with a known schema.
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub struct NamedRelation {
+    /// The relation name
+    pub name: String,
+    /// The schema description
+    pub schema: DFSchemaRef,
+}
+
+/// A variadic query operation
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub struct RecursiveQuery {
+    /// Name of the query
+    pub name: String,
+    /// The static term
+    pub static_term: Arc<LogicalPlan>,
+    /// The recursive term
+    pub recursive_term: Arc<LogicalPlan>,
+    /// Distinction
+    pub is_distinct: bool,
+}
+
 /// Values expression. See
 /// [Postgres VALUES](https://www.postgresql.org/docs/current/queries-values.html)
 /// documentation for more details.

datafusion/expr/src/logical_plan/plan.rs

@@ -366,6 +366,8 @@ impl OptimizerRule for CommonSubexprEliminate {
             | LogicalPlan::Dml(_)
             | LogicalPlan::Copy(_)
             | LogicalPlan::Unnest(_)
+            | LogicalPlan::NamedRelation(_)
+            | LogicalPlan::RecursiveQuery(_)
             | LogicalPlan::Prepare(_) => {
                 // apply the optimization to all inputs of the plan
                 utils::optimize_children(self, plan, config)?

datafusion/optimizer/src/common_subexpr_eliminate.rs

@@ -60,4 +60,5 @@ tempfile = "3"
 #[dev-dependencies]
 termtree = "0.4.1"
 tokio = { version = "1.28", features = ["macros", "rt", "rt-multi-thread", "sync", "fs", "parking_lot"] }
+tokio-stream = { version = "0.1.14" }
 uuid = { version = "^1.2", features = ["v4"] }