Support for substrait plans

Author: LiaCastaneda

Cargo.lock

@@ -44,6 +44,7 @@ datafusion = { git = "https://github.com/apache/datafusion", branch = "main", fe
     "avro",
 ] }
 datafusion-proto = { git = "https://github.com/apache/datafusion", branch = "main" }
+datafusion-substrait = { git = "https://github.com/apache/datafusion", branch = "main" }
 env_logger = "0.11"
 futures = "0.3"
 itertools = "0.14"
@@ -53,6 +54,9 @@ log = "0.4"
 rand = "0.8"
 uuid = { version = "1.6", features = ["v4"] }
 
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+
 object_store = { version = "0.12.0", features = [
   "aws",
   "gcp",

Cargo.toml

@@ -160,9 +160,10 @@ impl ExecutionPlan for DistributedExplainExec {
             self.distributed_stages.as_str(),
         ]);
 
-        let batch =
-            RecordBatch::try_new(schema.clone(), vec![Arc::new(plan_types), Arc::new(plans)])
-                .map_err(|e| datafusion::error::DataFusionError::ArrowError(Box::new(e), None))?;
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(plan_types), Arc::new(plans)],
+        ).map_err(|e| datafusion::error::DataFusionError::ArrowError(Box::new(e), None))?;
 
         // Use MemoryStream which is designed for DataFusion execution plans
         let stream = MemoryStream::try_new(vec![batch], schema, None)?;

src/explain.rs

@@ -134,6 +134,12 @@ pub trait FlightSqlHandler: Send + Sync {
         request: Request<FlightDescriptor>,
     ) -> Result<Response<FlightInfo>, Status>;
 
+    async fn get_flight_info_substrait(
+        &self,
+        query: arrow_flight::sql::CommandStatementSubstraitPlan,
+        request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status>;
+
     async fn do_get_statement(
         &self,
         ticket: arrow_flight::sql::TicketStatementQuery,
@@ -163,6 +169,20 @@ impl FlightSqlService for FlightSqlServ {
                 error!("Error in do_flight_info_statement: {:?}", e);
             })
     }
+
+    async fn get_flight_info_substrait_plan(
+        &self,
+        query: arrow_flight::sql::CommandStatementSubstraitPlan,
+        request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status> {
+        self.handler
+            .get_flight_info_substrait(query, request)
+            .await
+            .inspect_err(|e| {
+                error!("Error in get_flight_info_substrait: {:?}", e);
+            })
+    }
+
     async fn do_get_statement(
         &self,
         ticket: arrow_flight::sql::TicketStatementQuery,

src/flight.rs

@@ -136,6 +136,27 @@ impl FlightRequestHandler {
         Ok(Response::new(flight_info))
     }
 
+    pub async fn handle_substrait_info_request(&self, substrait_plan: datafusion_substrait::substrait::proto::Plan) -> Result<Response<FlightInfo>, Status> {
+        let query_plan = self
+            .planner
+            .prepare_substrait_query(substrait_plan)
+            .await
+            .map_err(|e| Status::internal(format!("Could not prepare query {e:?}")))?;
+
+        debug!("get flight info: query id {}", query_plan.query_id);
+
+        let flight_info = self.create_flight_info_response(
+            query_plan.query_id,
+            query_plan.worker_addresses,
+            query_plan.final_stage_id,
+            query_plan.schema,
+            None  // Regular queries don't have explain data
+        )?;
+
+        trace!("get_flight_info_statement done");
+        Ok(Response::new(flight_info))
+    }
+
     /// Handle execution of EXPLAIN statement queries.
     ///
     /// This function does not execute the plan but returns all plans we want to display to the user.

src/flight_handlers.rs

@@ -19,8 +19,9 @@ use std::sync::Arc;
 
 use anyhow::{anyhow, Context};
 use arrow_flight::{
-    flight_service_server::FlightServiceServer, FlightDescriptor, FlightInfo, Ticket,
+    flight_service_server::FlightServiceServer, sql::TicketStatementQuery, FlightDescriptor, FlightInfo, Ticket
 };
+use datafusion_substrait::substrait;
 use parking_lot::Mutex;
 use prost::Message;
 use tokio::{
@@ -74,7 +75,7 @@ impl FlightSqlHandler for DfRayProxyHandler {
 
     async fn do_get_statement(
         &self,
-        ticket: arrow_flight::sql::TicketStatementQuery,
+        ticket: TicketStatementQuery,
         request: Request<Ticket>,
     ) -> Result<Response<crate::flight::DoGetStream>, Status> {
         trace!("do_get_statement");
@@ -98,6 +99,22 @@ impl FlightSqlHandler for DfRayProxyHandler {
                 .await
         }
     }
+
+    async fn get_flight_info_substrait(
+        &self,
+        substrait: arrow_flight::sql::CommandStatementSubstraitPlan,
+        _request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status> {
+        let plan = match &substrait.plan {
+            Some(substrait_plan) => {
+                substrait::proto::Plan::decode(substrait_plan.plan.as_ref())
+                    .map_err(|e| Status::invalid_argument(format!("Invalid Substrait plan: {e}")))?
+            }
+            None => return Err(Status::invalid_argument("Missing Substrait plan")),
+        };
+
+        self.flight_handler.handle_substrait_info_request(plan).await
+    }
 }
 
 /// DFRayProcessorService is a Arrow Flight service that serves streams of
@@ -213,8 +230,7 @@ mod tests {
     use super::*;
     // Test-specific imports
     use arrow_flight::{
-        sql::{CommandStatementQuery, TicketStatementQuery},
-        FlightDescriptor, Ticket,
+        sql::{CommandStatementQuery, TicketStatementQuery}, FlightDescriptor, Ticket
     };
     use prost::Message;
     use tonic::Request;
@@ -345,4 +361,5 @@ mod tests {
     // TODO: Add tests for regular (non-explain) queries
     // We might need to create integration or end-to-end test infrastructure for this because
     // they need workers
+
 }

src/proxy_service.rs

@@ -5,6 +5,7 @@ use arrow::datatypes::SchemaRef;
 use datafusion::{
     logical_expr::LogicalPlan, physical_plan::ExecutionPlan, prelude::SessionContext,
 };
+use datafusion_substrait::logical_plan::{consumer::from_substrait_plan};
 
 use crate::{
     explain::{is_explain_query, DistributedExplainExec},
@@ -65,6 +66,38 @@ impl QueryPlanner {
         Self
     }
 
+    /// Dispatch a distributed query plan to the workers.
+    async fn dispatch_query_plan(&self, base_result: QueryPlanBase) -> Result<QueryPlan> {
+        if base_result.distributed_stages.is_empty() {
+            return Err(anyhow!("No stages generated for query").into());
+        }
+
+        let worker_addrs = get_worker_addresses()?;
+
+        // The last stage produces the data returned to the client.
+        let final_stage = &base_result.distributed_stages
+            [base_result.distributed_stages.len() - 1];
+        let schema = Arc::clone(&final_stage.plan.schema());
+        let final_stage_id = final_stage.stage_id;
+
+        // Physically dispatch each stage to the worker pool, further dividing
+        // them into partition groups.
+        let final_workers = distribute_stages(
+            &base_result.query_id,
+            base_result.distributed_stages,
+            worker_addrs,
+        )
+        .await?;
+
+        Ok(QueryPlan {
+            query_id: base_result.query_id,
+            worker_addresses: final_workers,
+            final_stage_id,
+            schema,
+            explain_data: None,
+        })
+    }
+
     /// Common planning steps shared by both query and its EXPLAIN
     ///
     /// Prepare a query by parsing the SQL, planning it, and distributing the
@@ -78,9 +111,8 @@ impl QueryPlanner {
         let logical_plan = logical_planning(sql, &ctx).await?;
         let physical_plan = physical_planning(&logical_plan, &ctx).await?;
 
-        // divide the physical plan into chunks (stages) that we can distribute to workers
-        let (distributed_plan, distributed_stages) =
-            execution_planning(physical_plan.clone(), 8192, Some(2)).await?;
+        // divide the physical plan into chunks (stages) that we can distribute to workers later in dispatch_query_plan
+        let (distributed_plan, distributed_stages) = execution_planning(physical_plan.clone(), 8192, Some(2)).await?;
 
         Ok(QueryPlanBase {
             query_id,
@@ -92,37 +124,45 @@ impl QueryPlanner {
         })
     }
 
-    /// Prepare a distributed query
+    /// Prepare a distributed query (SQL entry point)
     pub async fn prepare_query(&self, sql: &str) -> Result<QueryPlan> {
         let base_result = self.prepare_query_base(sql, "REGULAR").await?;
+        self.dispatch_query_plan(base_result).await
+    }
 
-        if base_result.distributed_stages.is_empty() {
-            return Err(anyhow!("No stages generated for query").into());
-        }
+    /// Prepare a distributed query (Substrait entry point)
+    pub async fn prepare_substrait_query(&self, substrait_plan: datafusion_substrait::substrait::proto::Plan) -> Result<QueryPlan> {
+        let base_result = self.prepare_substrait_query_base(substrait_plan, "SUBSTRAIT").await?;
+        self.dispatch_query_plan(base_result).await
+    }
 
-        let worker_addrs = get_worker_addresses()?;
+    pub async fn prepare_substrait_query_base(
+        &self,
+        substrait_plan: datafusion_substrait::substrait::proto::Plan,
+        query_type: &str,
+    ) -> Result<QueryPlanBase> {
+        debug!("prepare_substrait_query_base: {} Substrait = {:#?}", query_type, substrait_plan);
+        
+        let query_id = uuid::Uuid::new_v4().to_string();
+        let ctx = get_ctx().map_err(|e| anyhow!("Could not create context: {e}"))?;
 
-        // gather some information we need to send back such that
-        // we can send a ticket to the client
-        let final_stage = &base_result.distributed_stages[base_result.distributed_stages.len() - 1];
-        let schema = Arc::clone(&final_stage.plan.schema());
-        let final_stage_id = final_stage.stage_id;
+        let logical_plan = from_substrait_plan(&ctx.state(), &substrait_plan)
+            .await
+            .map_err(|e| anyhow!("Failed to convert DataFusion Logical Plan: {e}"))?;
 
-        // distribute the stages to workers, further dividing them up
-        // into chunks of partitions (partition_groups)
-        let final_workers = distribute_stages(
-            &base_result.query_id,
-            base_result.distributed_stages,
-            worker_addrs,
-        )
-        .await?;
 
-        Ok(QueryPlan {
-            query_id: base_result.query_id,
-            worker_addresses: final_workers,
-            final_stage_id,
-            schema,
-            explain_data: None,
+        let physical_plan = physical_planning(&logical_plan, &ctx).await.map_err(|e| anyhow!("Failed to convert DataFusion Physical Plan: {e}"))?;
+
+        // divide the physical plan into chunks (stages) that we can distribute to workers later in dispatch_query_plan
+        let (distributed_plan, distributed_stages) = execution_planning(physical_plan.clone(), 8192, Some(2)).await?;
+
+        Ok(QueryPlanBase {
+            query_id,
+            session_context: ctx,
+            logical_plan,
+            physical_plan,
+            distributed_plan,
+            distributed_stages,
         })
     }
 
@@ -206,17 +246,18 @@ impl QueryPlanner {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::{fs::{File}, path::Path};
+    use std::io::BufReader;
+
 
     // //////////////////////////////////////////////////////////////
     // Test helper functions
     // //////////////////////////////////////////////////////////////
 
     /// Set up mock worker environment for testing
     fn setup_mock_worker_env() {
-        let mock_addrs = vec![
-            ("mock_worker_1".to_string(), "localhost:9001".to_string()),
-            ("mock_worker_2".to_string(), "localhost:9002".to_string()),
-        ];
+        let mock_addrs = [("mock_worker_1".to_string(), "localhost:9001".to_string()),
+            ("mock_worker_2".to_string(), "localhost:9002".to_string())];
         let mock_env_value = mock_addrs
             .iter()
             .map(|(name, addr)| format!("{}/{}", name, addr))
@@ -257,6 +298,34 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_prepare_substrait_query_base() {
+        let planner = QueryPlanner::new();
+
+        // Read the JSON plan and convert to binary Substrait protobuf bytes
+        let plan = serde_json::from_reader::<_, datafusion_substrait::substrait::proto::Plan>(BufReader::new(
+            File::open(Path::new("testdata/substrait/select_one.substrait.json")).expect("file not found"),
+        )).expect("failed to parse json");
+        
+        let result = planner.prepare_substrait_query_base(plan, "TEST").await;
+        
+        if result.is_ok() {
+            let query_plan_base = result.unwrap();
+            // verify all fields have values
+            assert!(!query_plan_base.query_id.is_empty());
+            assert!(!query_plan_base.distributed_stages.is_empty());
+            assert!(!query_plan_base.physical_plan.schema().fields().is_empty());
+            // logical plan of select 1 on empty relation
+            assert_eq!(query_plan_base.logical_plan.to_string(), "Projection: Int64(1) AS test_col\n  Values: (Int64(0))");
+            // physical plan of select 1 on empty releation is ProjectionExec
+            assert_eq!(query_plan_base.physical_plan.name(), "ProjectionExec");
+        } else {
+            // If worker discovery fails, we expect a specific error
+            let error_msg = format!("{:?}", result.unwrap_err());
+            assert!(error_msg.contains("worker") || error_msg.contains("address"));
+        }
+    }
+
     #[tokio::test]
     async fn test_prepare_explain() {
         let planner = QueryPlanner::new();