Add an error propagation mechanism to the Arrow Flight endpoint and the Arrow Flight reader

Author: gabotechs

src/flight_service/do_get.rs

@@ -1,4 +1,5 @@
 use crate::composed_extension_codec::ComposedPhysicalExtensionCodec;
+use crate::errors::datafusion_error_to_tonic_status;
 use crate::flight_service::service::ArrowFlightEndpoint;
 use crate::plan::ArrowFlightReadExecProtoCodec;
 use crate::stage_delegation::{ActorContext, StageContext};
@@ -140,22 +141,24 @@ impl ArrowFlightEndpoint {
         let stream_partitioner = self
             .partitioner_registry
             .get_or_create_stream_partitioner(stage_id, actor_idx, plan, partitioning)
-            .map_err(|err| {
-                Status::internal(format!("Could not create stream partitioner: {err}"))
-            })?;
+            .map_err(|err| datafusion_error_to_tonic_status(&err))?;
 
         let stream = stream_partitioner
             .execute(caller_actor_idx, state.task_ctx())
-            .map_err(|err| Status::internal(format!("Cannot get stream partition: {err}")))?;
+            .map_err(|err| datafusion_error_to_tonic_status(&err))?;
 
-        // TODO: error propagation
         let flight_data_stream = FlightDataEncoderBuilder::new()
             .with_schema(stream_partitioner.schema())
-            .build(stream.map_err(|err| FlightError::ExternalError(Box::new(err))));
-
-        Ok(Response::new(Box::pin(flight_data_stream.map_err(|err| {
-            Status::internal(format!("Error during flight stream: {err}"))
-        }))))
+            .build(stream.map_err(|err| {
+                FlightError::Tonic(Box::new(datafusion_error_to_tonic_status(&err)))
+            }));
+
+        Ok(Response::new(Box::pin(flight_data_stream.map_err(
+            |err| match err {
+                FlightError::Tonic(status) => *status,
+                _ => Status::internal(format!("Error during flight stream: {err}")),
+            },
+        ))))
     }
 }
 

src/plan/arrow_flight_read.rs

@@ -1,5 +1,6 @@
 use crate::channel_manager::{ArrowFlightChannel, ChannelManager};
 use crate::composed_extension_codec::ComposedPhysicalExtensionCodec;
+use crate::errors::tonic_status_to_datafusion_error;
 use crate::flight_service::{DoGet, DoPut};
 use crate::plan::arrow_flight_read_proto::ArrowFlightReadExecProtoCodec;
 use crate::stage_delegation::{ActorContext, StageContext, StageDelegation};
@@ -189,14 +190,18 @@ impl ExecutionPlan for ArrowFlightReadExec {
             let stream = client
                 .do_get(ticket.into_request())
                 .await
-                .map_err(|err| DataFusionError::External(Box::new(err)))?
+                .map_err(|err| tonic_status_to_datafusion_error(&err).unwrap_or_else(|| {
+                    DataFusionError::External(Box::new(err))
+                }))?
                 .into_inner()
                 .map_err(|err| FlightError::Tonic(Box::new(err)));
 
             Ok(FlightRecordBatchStream::new_from_flight_data(stream)
-                // TODO: propagate the error from the service to here, probably serializing it
-                //  somehow.
-                .map_err(|err| DataFusionError::External(Box::new(err))))
+                .map_err(|err| match err {
+                    FlightError::Tonic(status) => tonic_status_to_datafusion_error(&status)
+                        .unwrap_or_else(|| DataFusionError::External(Box::new(status))),
+                    err => DataFusionError::External(Box::new(err))
+                }))
         }.try_flatten_stream();
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(