diff --git a/datafusion-federation/src/sql/mod.rs b/datafusion-federation/src/sql/mod.rs index 1d01a7f..4b099dc 100644 --- a/datafusion-federation/src/sql/mod.rs +++ b/datafusion-federation/src/sql/mod.rs @@ -25,6 +25,7 @@ use datafusion::{ SendableRecordBatchStream, }, sql::{ + sqlparser::ast::Statement, unparser::{plan_to_sql, Unparser}, TableReference, }, @@ -698,8 +699,8 @@ impl VirtualExecutionPlan { fn sql(&self) -> Result { // Find all table scans, recover the SQLTableSource, find the remote table name and replace the name of the TableScan table. let mut known_rewrites = HashMap::new(); - let mut ast = Unparser::new(self.executor.dialect().as_ref()) - .plan_to_sql(&rewrite_table_scans(&self.plan, &mut known_rewrites)?)?; + let plan = &rewrite_table_scans(&self.plan, &mut known_rewrites)?; + let mut ast = self.plan_to_sql(plan)?; if let Some(analyzer) = self.executor.ast_analyzer() { ast = analyzer(ast)?; @@ -707,6 +708,10 @@ impl VirtualExecutionPlan { Ok(format!("{ast}")) } + + fn plan_to_sql(&self, plan: &LogicalPlan) -> Result { + Unparser::new(self.executor.dialect().as_ref()).plan_to_sql(plan) + } } impl DisplayAs for VirtualExecutionPlan { @@ -758,9 +763,7 @@ impl ExecutionPlan for VirtualExecutionPlan { _partition: usize, _context: Arc, ) -> Result { - let ast = plan_to_sql(&self.plan)?; - let query = format!("{ast}"); - + let query = self.plan_to_sql(&self.plan)?.to_string(); self.executor.execute(query.as_str(), self.schema()) }