fix: restore Stage plan during deserialization round-trip (#404)

Closes #402

  ### Problem

After #345, `parse_stage_proto` unconditionally sets `stage.plan = None`
during deserialization, discarding children that DataFusion's framework
correctly deserialized. This breaks any architecture
where the distributed plan goes through a serialize/deserialize
round-trip before `DistributedExec::prepare_plan()` runs — such as Arrow
Flight SQL's two-phase `GetFlightInfo`/`DoGet` protocol.

  ### Fix

`children()` on network boundary nodes already encodes plan presence: it
returns `[plan]` when `Some`, `[]` when `None`. DataFusion's
serialization framework preserves this through the round-trip and
passes the deserialized children as `inputs`. The fix infers plan
presence from `inputs` rather than ignoring them:

  ```rust
  plan: inputs.first().cloned(),
  ```

  ### Tests

- Updated two existing tests (`test_roundtrip_single_flight`,
`test_roundtrip_single_flight_coalesce`) to pass correct inputs for
stages with `plan: None`
- Added two new tests (`test_roundtrip_single_flight_with_plan`,
`test_roundtrip_single_flight_coalesce_with_plan`) that validate the
round-trip when a stage has `plan: Some(...)`
  
 ### Visibility changes
In a Flight SQL architecture, the coordinator server operates in two
distinct phases separated by a network boundary:

1. GetFlightInfo — optimizes the query and serializes the full physical
plan (including DFD's network boundary nodes) into a ticket
  2. DoGet — deserializes the plan from the ticket and executes it

This means the server must be able to serialize and deserialize DFD's
plan nodes independently of DFD's internal machinery. To do this, it
needs DistributedCodec in its PhysicalExtensionCodec chain — otherwise
the server has no way to round-trip NetworkShuffleExec,
NetworkCoalesceExec, etc. through the ticket.

Similarly, a Flight SQL server that also acts as a DFD worker needs to
implement the WorkerService gRPC trait. This requires
WorkerServiceServer to register the service with tonic,
FlightAppMetadata to read passthrough headers from requests, and
Worker::impl_execute_task to delegate task execution. The error
conversion functions (datafusion_error_to_tonic_status, etc.) are needed
to translate between DFD's error types and gRPC status codes in the
server's request handlers.

Currently these types are all pub(crate), which forces downstream Flight
SQL integrations to fork the crate just to access them. Making them
public allows DFD to be used as a library in these architectures without
modification.

Author: gabrielkerr

src/lib.rs

@@ -13,6 +13,7 @@ mod distributed_planner;
 mod networking;
 mod observability;
 mod protobuf;
+pub use protobuf::DistributedCodec;
 #[cfg(any(feature = "integration", test))]
 pub mod test_utils;
 

src/protobuf/distributed_codec.rs

@@ -60,7 +60,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
 
         fn parse_stage_proto(
             proto: Option<StageProto>,
-            _inputs: &[Arc<dyn ExecutionPlan>],
+            inputs: &[Arc<dyn ExecutionPlan>],
         ) -> Result<Stage, DataFusionError> {
             let Some(proto) = proto else {
                 return Err(proto_error("Empty StageProto"));
@@ -70,7 +70,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
                 query_id: uuid::Uuid::from_slice(proto.query_id.as_ref())
                     .map_err(|_| proto_error("Invalid query_id in StageProto"))?,
                 num: proto.num as usize,
-                plan: None,
+                plan: inputs.first().cloned(),
                 tasks: decode_tasks(proto.tasks)?,
             })
         }
@@ -556,6 +556,15 @@ mod tests {
         }
     }
 
+    fn dummy_stage_with_plan() -> Stage {
+        Stage {
+            query_id: Default::default(),
+            num: 0,
+            plan: Some(empty_exec()),
+            tasks: vec![],
+        }
+    }
+
     fn schema_i32(name: &str) -> Arc<Schema> {
         Arc::new(Schema::new(vec![Field::new(name, DataType::Int32, false)]))
     }
@@ -581,7 +590,7 @@ mod tests {
         let mut buf = Vec::new();
         codec.try_encode(plan.clone(), &mut buf)?;
 
-        let decoded = codec.try_decode(&buf, &[empty_exec()], &ctx)?;
+        let decoded = codec.try_decode(&buf, &[], &ctx)?;
         assert_eq!(repr(&plan), repr(&decoded));
 
         Ok(())
@@ -686,6 +695,49 @@ mod tests {
         let mut buf = Vec::new();
         codec.try_encode(plan.clone(), &mut buf)?;
 
+        let decoded = codec.try_decode(&buf, &[], &ctx)?;
+        assert_eq!(repr(&plan), repr(&decoded));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_roundtrip_single_flight_with_plan() -> datafusion::common::Result<()> {
+        let codec = DistributedCodec;
+        let ctx = create_context();
+
+        let schema = schema_i32("a");
+        let part = Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 4);
+        let plan: Arc<dyn ExecutionPlan> = Arc::new(new_network_hash_shuffle_exec(
+            part,
+            schema,
+            dummy_stage_with_plan(),
+        ));
+
+        let mut buf = Vec::new();
+        codec.try_encode(plan.clone(), &mut buf)?;
+
+        let decoded = codec.try_decode(&buf, &[empty_exec()], &ctx)?;
+        assert_eq!(repr(&plan), repr(&decoded));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_roundtrip_single_flight_coalesce_with_plan() -> datafusion::common::Result<()> {
+        let codec = DistributedCodec;
+        let ctx = create_context();
+
+        let schema = schema_i32("e");
+        let plan: Arc<dyn ExecutionPlan> = Arc::new(new_network_coalesce_tasks_exec(
+            Partitioning::RoundRobinBatch(3),
+            schema,
+            dummy_stage_with_plan(),
+        ));
+
+        let mut buf = Vec::new();
+        codec.try_encode(plan.clone(), &mut buf)?;
+
         let decoded = codec.try_decode(&buf, &[empty_exec()], &ctx)?;
         assert_eq!(repr(&plan), repr(&decoded));
 

src/protobuf/mod.rs

@@ -2,7 +2,7 @@ mod distributed_codec;
 mod errors;
 mod user_codec;
 
-pub(crate) use distributed_codec::DistributedCodec;
+pub use distributed_codec::DistributedCodec;
 pub(crate) use errors::{
     datafusion_error_to_tonic_status, map_flight_to_datafusion_error,
     tonic_status_to_datafusion_error,