feat(sql): migrate to DataFusion-based streaming SQL planner

Author: luoluoyuyu

Cargo.lock

@@ -52,8 +52,36 @@ arrow-array = "52"
 arrow-ipc = "52"
 arrow-schema = "52"
 proctitle = "0.1"
+unicase = "2.7"
+petgraph = "0.7"
+itertools = "0.14"
+strum = { version = "0.26", features = ["derive"] }
+datafusion-functions-aggregate = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'}
+
+typify = { git = 'https://github.com/ArroyoSystems/typify.git', branch = 'arroyo' }
+parquet = {git = 'https://github.com/ArroyoSystems/arrow-rs', branch = '55.2.0/parquet'}
+arrow-json = {git = 'https://github.com/ArroyoSystems/arrow-rs', branch = '55.2.0/json'}
+datafusion = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'}
+datafusion-common = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'}
+datafusion-execution = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'}
+datafusion-expr = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'}
+datafusion-physical-expr = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'}
+datafusion-physical-plan = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'}
+datafusion-proto = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'}
+datafusion-functions = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'}
+datafusion-functions-window = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'}
+
+sqlparser = { git = "https://github.com/FunctionStream/sqlparser-rs", branch = "0.6.0/function-sql-parser" }
+
+cornucopia_async = { git = "https://github.com/ArroyoSystems/cornucopia", branch = "sqlite" }
+cornucopia = { git = "https://github.com/ArroyoSystems/cornucopia", branch = "sqlite" }
+jiter = {git = "https://github.com/ArroyoSystems/jiter", branch = "disable_python" }
+
 
 [features]
 default = ["incremental-cache", "python"]
 incremental-cache = ["wasmtime/incremental-cache"]
 python = []
+
+[patch."https://github.com/ArroyoSystems/sqlparser-rs"]
+sqlparser = { git = "https://github.com/FunctionStream/sqlparser-rs", branch = "0.6.0/function-sql-parser" }

Cargo.toml

@@ -14,7 +14,7 @@ use super::Analysis;
 use crate::coordinator::execution_context::ExecutionContext;
 use crate::coordinator::statement::{
     CreateFunction, CreatePythonFunction, DropFunction, ShowFunctions, StartFunction, Statement,
-    StatementVisitor, StatementVisitorContext, StatementVisitorResult, StopFunction,
+    StatementVisitor, StatementVisitorContext, StatementVisitorResult, StopFunction, StreamingSql,
 };
 use std::fmt;
 
@@ -115,4 +115,13 @@ impl StatementVisitor for Analyzer<'_> {
     ) -> StatementVisitorResult {
         StatementVisitorResult::Analyze(Box::new(stmt.clone()))
     }
+
+    fn visit_streaming_sql(
+        &self,
+        stmt: &StreamingSql,
+        _context: &StatementVisitorContext,
+    ) -> StatementVisitorResult {
+        // TODO: add semantic analysis for streaming SQL (schema validation, etc.)
+        StatementVisitorResult::Analyze(Box::new(StreamingSql::new(stmt.statement.clone())))
+    }
 }

src/coordinator/analyze/analyzer.rs

@@ -20,6 +20,7 @@ use crate::coordinator::execution::Executor;
 use crate::coordinator::plan::{LogicalPlanVisitor, LogicalPlanner, PlanNode};
 use crate::coordinator::statement::Statement;
 use crate::runtime::taskexecutor::TaskManager;
+use crate::sql::planner::StreamSchemaProvider;
 
 use super::execution_context::ExecutionContext;
 
@@ -90,7 +91,8 @@ impl Coordinator {
     }
 
     fn step_build_logical_plan(&self, analysis: &Analysis) -> Result<Box<dyn PlanNode>> {
-        let visitor = LogicalPlanVisitor::new();
+        let schema_provider = StreamSchemaProvider::new();
+        let visitor = LogicalPlanVisitor::new(schema_provider);
         let plan = visitor.visit(analysis);
         Ok(plan)
     }

src/coordinator/coordinator.rs

@@ -14,6 +14,7 @@ use crate::coordinator::dataset::{ExecuteResult, ShowFunctionsResult, empty_reco
 use crate::coordinator::plan::{
     CreateFunctionPlan, CreatePythonFunctionPlan, DropFunctionPlan, PlanNode, PlanVisitor,
     PlanVisitorContext, PlanVisitorResult, ShowFunctionsPlan, StartFunctionPlan, StopFunctionPlan,
+    StreamingSqlPlan,
 };
 use crate::coordinator::statement::{ConfigSource, FunctionSource};
 use crate::runtime::taskexecutor::TaskManager;
@@ -200,4 +201,17 @@ impl PlanVisitor for Executor {
 
         PlanVisitorResult::Execute(result)
     }
+
+    fn visit_streaming_sql_plan(
+        &self,
+        plan: &StreamingSqlPlan,
+        _context: &PlanVisitorContext,
+    ) -> PlanVisitorResult {
+        // TODO: apply rewrite_plan for streaming transformations, then execute
+        let result = Err(ExecuteError::Internal(format!(
+            "Streaming SQL execution not yet implemented. LogicalPlan:\n{}",
+            plan.logical_plan.display_indent()
+        )));
+        PlanVisitorResult::Execute(result)
+    }
 }

src/coordinator/execution/executor.rs

@@ -23,5 +23,5 @@ pub use coordinator::Coordinator;
 pub use dataset::{DataSet, ShowFunctionsResult};
 pub use statement::{
     CreateFunction, CreatePythonFunction, DropFunction, PythonModule, ShowFunctions, StartFunction,
-    Statement, StopFunction,
+    Statement, StopFunction, StreamingSql,
 };

src/coordinator/mod.rs

@@ -10,22 +10,26 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use tracing::debug;
+
 use crate::coordinator::analyze::analysis::Analysis;
 use crate::coordinator::plan::{
     CreateFunctionPlan, CreatePythonFunctionPlan, DropFunctionPlan, PlanNode, ShowFunctionsPlan,
-    StartFunctionPlan, StopFunctionPlan,
+    StartFunctionPlan, StopFunctionPlan, StreamingSqlPlan,
 };
 use crate::coordinator::statement::{
     CreateFunction, CreatePythonFunction, DropFunction, ShowFunctions, StartFunction,
-    StatementVisitor, StatementVisitorContext, StatementVisitorResult, StopFunction,
+    StatementVisitor, StatementVisitorContext, StatementVisitorResult, StopFunction, StreamingSql,
 };
+use crate::sql::planner::StreamSchemaProvider;
 
-#[derive(Debug, Default)]
-pub struct LogicalPlanVisitor;
+pub struct LogicalPlanVisitor {
+    schema_provider: StreamSchemaProvider,
+}
 
 impl LogicalPlanVisitor {
-    pub fn new() -> Self {
-        Self
+    pub fn new(schema_provider: StreamSchemaProvider) -> Self {
+        Self { schema_provider }
     }
 
     pub fn visit(&self, analysis: &Analysis) -> Box<dyn PlanNode> {
@@ -51,7 +55,6 @@ impl StatementVisitor for LogicalPlanVisitor {
         let config_source = stmt.get_config_source().cloned();
         let extra_props = stmt.get_extra_properties().clone();
 
-        // Name will be read from config file during execution
         StatementVisitorResult::Plan(Box::new(CreateFunctionPlan::new(
             function_source,
             config_source,
@@ -106,4 +109,22 @@ impl StatementVisitor for LogicalPlanVisitor {
             config_content,
         )))
     }
+
+    fn visit_streaming_sql(
+        &self,
+        stmt: &StreamingSql,
+        _context: &StatementVisitorContext,
+    ) -> StatementVisitorResult {
+        let sql_to_rel = datafusion::sql::planner::SqlToRel::new(&self.schema_provider);
+
+        match sql_to_rel.sql_statement_to_plan(stmt.statement.clone()) {
+            Ok(plan) => {
+                debug!("Logical plan:\n{}", plan.display_graphviz());
+                StatementVisitorResult::Plan(Box::new(StreamingSqlPlan::new(plan)))
+            }
+            Err(e) => {
+                panic!("Failed to convert SQL statement to logical plan: {e}");
+            }
+        }
+    }
 }

src/coordinator/plan/logical_plan_visitor.rs

@@ -18,6 +18,7 @@ mod optimizer;
 mod show_functions_plan;
 mod start_function_plan;
 mod stop_function_plan;
+mod streaming_sql_plan;
 mod visitor;
 
 pub use create_function_plan::CreateFunctionPlan;
@@ -28,6 +29,7 @@ pub use optimizer::LogicalPlanner;
 pub use show_functions_plan::ShowFunctionsPlan;
 pub use start_function_plan::StartFunctionPlan;
 pub use stop_function_plan::StopFunctionPlan;
+pub use streaming_sql_plan::StreamingSqlPlan;
 pub use visitor::{PlanVisitor, PlanVisitorContext, PlanVisitorResult};
 
 use std::fmt;

src/coordinator/plan/mod.rs

@@ -10,33 +10,23 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-mod sql_parser;
+use datafusion::logical_expr::LogicalPlan;
 
-pub use sql_parser::SqlParser;
+use super::{PlanNode, PlanVisitor, PlanVisitorContext, PlanVisitorResult};
 
 #[derive(Debug)]
-pub struct ParseError {
-    pub message: String,
+pub struct StreamingSqlPlan {
+    pub logical_plan: LogicalPlan,
 }
 
-impl std::fmt::Display for ParseError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "Parse error: {}", self.message)
+impl StreamingSqlPlan {
+    pub fn new(logical_plan: LogicalPlan) -> Self {
+        Self { logical_plan }
     }
 }
 
-impl std::error::Error for ParseError {}
-
-impl From<String> for ParseError {
-    fn from(message: String) -> Self {
-        ParseError { message }
-    }
-}
-
-impl ParseError {
-    pub fn new(message: impl Into<String>) -> Self {
-        Self {
-            message: message.into(),
-        }
+impl PlanNode for StreamingSqlPlan {
+    fn accept(&self, visitor: &dyn PlanVisitor, context: &PlanVisitorContext) -> PlanVisitorResult {
+        visitor.visit_streaming_sql_plan(self, context)
     }
 }

src/sql/parser/mod.rs …c/coordinator/plan/streaming_sql_plan.rssrc/sql/parser/mod.rs renamed to src/coordinator/plan/streaming_sql_plan.rs src/sql/parser/mod.rs renamed to src/coordinator/plan/streaming_sql_plan.rs

@@ -12,7 +12,7 @@
 
 use super::{
     CreateFunctionPlan, CreatePythonFunctionPlan, DropFunctionPlan, ShowFunctionsPlan,
-    StartFunctionPlan, StopFunctionPlan,
+    StartFunctionPlan, StopFunctionPlan, StreamingSqlPlan,
 };
 
 /// Context passed to PlanVisitor methods
@@ -84,4 +84,10 @@ pub trait PlanVisitor {
         plan: &CreatePythonFunctionPlan,
         context: &PlanVisitorContext,
     ) -> PlanVisitorResult;
+
+    fn visit_streaming_sql_plan(
+        &self,
+        plan: &StreamingSqlPlan,
+        context: &PlanVisitorContext,
+    ) -> PlanVisitorResult;
 }