Fix Correlated Subquery With Depth Larger Than One

benchmarks/src/cancellation.rs

@@ -33,6 +33,7 @@ use datafusion::execution::TaskContext;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::*;
+use datafusion::sql::planner::PlannerContext;
 use datafusion_common::instant::Instant;
 use futures::TryStreamExt;
 use object_store::ObjectStore;
@@ -185,7 +186,10 @@ async fn datafusion(store: Arc<dyn ObjectStore>) -> Result<()> {
     .await?;
 
     println!("Creating logical plan...");
-    let logical_plan = ctx.state().create_logical_plan(query).await?;
+    let logical_plan = ctx
+        .state()
+        .create_logical_plan(query, &mut PlannerContext::new())
+        .await?;
 
     println!("Creating physical plan...");
     let physical_plan = Arc::new(CoalescePartitionsExec::new(

datafusion-cli/src/cli_context.rs

@@ -23,6 +23,7 @@ use datafusion::{
     execution::{context::SessionState, TaskContext},
     logical_expr::LogicalPlan,
     prelude::SessionContext,
+    sql::planner::PlannerContext,
 };
 use object_store::ObjectStore;
 
@@ -51,6 +52,7 @@ pub trait CliSessionContext {
     async fn execute_logical_plan(
         &self,
         plan: LogicalPlan,
+        planner_context: Option<PlannerContext>,
     ) -> Result<DataFrame, DataFusionError>;
 }
 
@@ -92,7 +94,8 @@ impl CliSessionContext for SessionContext {
     async fn execute_logical_plan(
         &self,
         plan: LogicalPlan,
+        planner_context: Option<PlannerContext>,
     ) -> Result<DataFrame, DataFusionError> {
-        self.execute_logical_plan(plan).await
+        self.execute_logical_plan(plan, planner_context).await
     }
 }

datafusion-cli/src/exec.rs

@@ -26,6 +26,7 @@ use crate::{
     object_storage::get_object_store,
     print_options::{MaxRows, PrintOptions},
 };
+use datafusion::sql::planner::PlannerContext;
 use futures::StreamExt;
 use std::collections::HashMap;
 use std::fs::File;
@@ -231,10 +232,13 @@ pub(super) async fn exec_and_print(
         let adjusted =
             AdjustedPrintOptions::new(print_options.clone()).with_statement(&statement);
 
-        let plan = create_plan(ctx, statement).await?;
+        let mut planner_context = PlannerContext::new();
+        let plan = create_plan(ctx, statement, &mut planner_context).await?;
         let adjusted = adjusted.with_plan(&plan);
 
-        let df = ctx.execute_logical_plan(plan).await?;
+        let df = ctx
+            .execute_logical_plan(plan, Some(planner_context))
+            .await?;
         let physical_plan = df.create_physical_plan().await?;
 
         // Track memory usage for the query result if it's bounded
@@ -348,8 +352,12 @@ fn config_file_type_from_str(ext: &str) -> Option<ConfigFileType> {
 async fn create_plan(
     ctx: &dyn CliSessionContext,
     statement: Statement,
+    planner_context: &mut PlannerContext,
 ) -> Result<LogicalPlan, DataFusionError> {
-    let mut plan = ctx.session_state().statement_to_plan(statement).await?;
+    let mut plan = ctx
+        .session_state()
+        .statement_to_plan(statement, planner_context)
+        .await?;
 
     // Note that cmd is a mutable reference so that create_external_table function can remove all
     // datafusion-cli specific options before passing through to datafusion. Otherwise, datafusion
@@ -453,7 +461,10 @@ mod tests {
 
     async fn create_external_table_test(location: &str, sql: &str) -> Result<()> {
         let ctx = SessionContext::new();
-        let plan = ctx.state().create_logical_plan(sql).await?;
+        let plan = ctx
+            .state()
+            .create_logical_plan(sql, &mut PlannerContext::new())
+            .await?;
 
         if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &plan {
             let format = config_file_type_from_str(&cmd.file_type);
@@ -479,7 +490,10 @@ mod tests {
         let ctx = SessionContext::new();
         // AWS CONFIG register.
 
-        let plan = ctx.state().create_logical_plan(sql).await?;
+        let plan = ctx
+            .state()
+            .create_logical_plan(sql, &mut PlannerContext::new())
+            .await?;
 
         if let LogicalPlan::Copy(cmd) = &plan {
             let format = config_file_type_from_str(&cmd.file_type.get_ext());

datafusion-cli/src/object_storage.rs

@@ -493,7 +493,10 @@ mod tests {
         );
 
         let ctx = SessionContext::new();
-        let mut plan = ctx.state().create_logical_plan(&sql).await?;
+        let mut plan = ctx
+            .state()
+            .create_logical_plan(&sql, &mut PlannerContext::new())
+            .await?;
 
         if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &mut plan {
             ctx.register_table_options_extension_from_scheme(scheme);
@@ -538,7 +541,10 @@ mod tests {
         );
 
         let ctx = SessionContext::new();
-        let mut plan = ctx.state().create_logical_plan(&sql).await?;
+        let mut plan = ctx
+            .state()
+            .create_logical_plan(&sql, &mut PlannerContext::new())
+            .await?;
 
         if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &mut plan {
             ctx.register_table_options_extension_from_scheme(scheme);
@@ -564,7 +570,10 @@ mod tests {
             ) LOCATION '{location}'"
         );
 
-        let mut plan = ctx.state().create_logical_plan(&sql).await?;
+        let mut plan = ctx
+            .state()
+            .create_logical_plan(&sql, &mut PlannerContext::new())
+            .await?;
 
         if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &mut plan {
             ctx.register_table_options_extension_from_scheme(scheme);
@@ -592,7 +601,10 @@ mod tests {
         let sql = format!("CREATE EXTERNAL TABLE test STORED AS PARQUET OPTIONS('aws.access_key_id' '{access_key_id}', 'aws.secret_access_key' '{secret_access_key}', 'aws.oss.endpoint' '{endpoint}') LOCATION '{location}'");
 
         let ctx = SessionContext::new();
-        let mut plan = ctx.state().create_logical_plan(&sql).await?;
+        let mut plan = ctx
+            .state()
+            .create_logical_plan(&sql, &mut PlannerContext::new())
+            .await?;
 
         if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &mut plan {
             ctx.register_table_options_extension_from_scheme(scheme);
@@ -629,7 +641,10 @@ mod tests {
         let sql = format!("CREATE EXTERNAL TABLE test STORED AS PARQUET OPTIONS('gcp.service_account_path' '{service_account_path}', 'gcp.service_account_key' '{service_account_key}', 'gcp.application_credentials_path' '{application_credentials_path}') LOCATION '{location}'");
 
         let ctx = SessionContext::new();
-        let mut plan = ctx.state().create_logical_plan(&sql).await?;
+        let mut plan = ctx
+            .state()
+            .create_logical_plan(&sql, &mut PlannerContext::new())
+            .await?;
 
         if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &mut plan {
             ctx.register_table_options_extension_from_scheme(scheme);

datafusion/core/src/execution/context/mod.rs

@@ -87,6 +87,7 @@ use datafusion_session::SessionStore;
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
+use datafusion_sql::planner::PlannerContext;
 use object_store::ObjectStore;
 use parking_lot::RwLock;
 use url::Url;
@@ -620,10 +621,14 @@ impl SessionContext {
         sql: &str,
         options: SQLOptions,
     ) -> Result<DataFrame> {
-        let plan = self.state().create_logical_plan(sql).await?;
+        let plan = self
+            .state()
+            .create_logical_plan(sql, &mut PlannerContext::new())
+            .await?;
         options.verify_plan(&plan)?;
 
-        self.execute_logical_plan(plan).await
+        // TODO: fetch planner_context
+        self.execute_logical_plan(plan, None).await
     }
 
     /// Creates logical expressions from SQL query text.
@@ -659,7 +664,11 @@ impl SessionContext {
     /// If you wish to limit the type of plan that can be run from
     /// SQL, see [`Self::sql_with_options`] and
     /// [`SQLOptions::verify_plan`].
-    pub async fn execute_logical_plan(&self, plan: LogicalPlan) -> Result<DataFrame> {
+    pub async fn execute_logical_plan(
+        &self,
+        plan: LogicalPlan,
+        planner_context: Option<PlannerContext>,
+    ) -> Result<DataFrame> {
         match plan {
             LogicalPlan::Ddl(ddl) => {
                 // Box::pin avoids allocating the stack space within this function's frame
@@ -734,7 +743,10 @@ impl SessionContext {
                     .remove_prepared(deallocate.name.as_str())?;
                 self.return_empty_dataframe()
             }
-            plan => Ok(DataFrame::new(self.state(), plan)),
+            plan => Ok(DataFrame::new(
+                self.with_planner_context_state(planner_context),
+                plan,
+            )),
         }
     }
 
@@ -1647,6 +1659,17 @@ impl SessionContext {
         state
     }
 
+    /// Return a new  [`SessionState`] with specific [`PlannerContext`].
+    pub fn with_planner_context_state(
+        &self,
+        planner_context: Option<PlannerContext>,
+    ) -> SessionState {
+        let mut state = self.state.read().clone();
+        state.set_planner_context(planner_context);
+        state.execution_props_mut().start_execution();
+        state
+    }
+
     /// Get reference to [`SessionState`]
     pub fn state_ref(&self) -> Arc<RwLock<SessionState>> {
         Arc::clone(&self.state)

datafusion/core/src/execution/session_state.rs

@@ -179,6 +179,7 @@ pub struct SessionState {
     /// Cache logical plans of prepared statements for later execution.
     /// Key is the prepared statement name.
     prepared_plans: HashMap<String, Arc<PreparedPlan>>,
+    planner_context: Option<PlannerContext>,
 }
 
 impl Debug for SessionState {
@@ -207,6 +208,7 @@ impl Debug for SessionState {
             .field("aggregate_functions", &self.aggregate_functions)
             .field("window_functions", &self.window_functions)
             .field("prepared_plans", &self.prepared_plans)
+            .field("planner_context", &self.planner_context)
             .finish()
     }
 }
@@ -353,6 +355,11 @@ impl SessionState {
         self.function_factory.as_ref()
     }
 
+    /// Set the planner_context.
+    pub fn set_planner_context(&mut self, planner_context: Option<PlannerContext>) {
+        self.planner_context = planner_context;
+    }
+
     /// Get the table factories
     pub fn table_factories(&self) -> &HashMap<String, Arc<dyn TableProviderFactory>> {
         &self.table_factories
@@ -461,6 +468,7 @@ impl SessionState {
     pub async fn statement_to_plan(
         &self,
         statement: Statement,
+        planner_context: &mut PlannerContext,
     ) -> datafusion_common::Result<LogicalPlan> {
         let references = self.resolve_table_references(&statement)?;
 
@@ -482,7 +490,7 @@ impl SessionState {
         }
 
         let query = SqlToRel::new_with_options(&provider, self.get_parser_options());
-        query.statement_to_plan(statement)
+        query.statement_to_plan(statement, planner_context)
     }
 
     fn get_parser_options(&self) -> ParserOptions {
@@ -514,10 +522,11 @@ impl SessionState {
     pub async fn create_logical_plan(
         &self,
         sql: &str,
+        planner_context: &mut PlannerContext,
     ) -> datafusion_common::Result<LogicalPlan> {
         let dialect = self.config.options().sql_parser.dialect.as_str();
         let statement = self.sql_to_statement(sql, dialect)?;
-        let plan = self.statement_to_plan(statement).await?;
+        let plan = self.statement_to_plan(statement, planner_context).await?;
         Ok(plan)
     }
 
@@ -1383,6 +1392,7 @@ impl SessionStateBuilder {
             runtime_env,
             function_factory,
             prepared_plans: HashMap::new(),
+            planner_context: None,
         };
 
         if let Some(file_formats) = file_formats {

datafusion/core/tests/sql/sql_api.rs

@@ -162,7 +162,9 @@ async fn empty_statement_returns_error() {
     let state = ctx.state();
 
     // Give it an empty string which contains no statements
-    let plan_res = state.create_logical_plan("").await;
+    let plan_res = state
+        .create_logical_plan("", &mut PlannerContext::new())
+        .await;
     assert_eq!(
         plan_res.unwrap_err().strip_backtrace(),
         "Error during planning: No SQL statements were provided in the query string"
@@ -180,6 +182,7 @@ async fn multiple_statements_returns_error() {
     let plan_res = state
         .create_logical_plan(
             "INSERT INTO test (x) VALUES (1); INSERT INTO test (x) VALUES (2)",
+            &mut PlannerContext::new(),
         )
         .await;
     assert_eq!(
@@ -199,7 +202,10 @@ async fn ddl_can_not_be_planned_by_session_state() {
 
     // can not create a logical plan for catalog DDL
     let sql = "DROP TABLE test";
-    let plan = state.create_logical_plan(sql).await.unwrap();
+    let plan = state
+        .create_logical_plan(sql, &mut PlannerContext::new())
+        .await
+        .unwrap();
     let physical_plan = state.create_physical_plan(&plan).await;
     assert_eq!(
         physical_plan.unwrap_err().strip_backtrace(),