fix(4981): incorrect error wrapping in OnceFut

datafusion/common/src/error.rs

@@ -36,6 +36,9 @@ use sqlparser::parser::ParserError;
 /// Result type for operations that could result in an [DataFusionError]
 pub type Result<T> = result::Result<T, DataFusionError>;
 
+/// Result type for operations that could result in an [DataFusionError] and needs to be shared (wrapped into `Arc`).
+pub type SharedResult<T> = result::Result<T, Arc<DataFusionError>>;
+
 /// Error type for generic operations that could result in DataFusionError::External
 pub type GenericError = Box<dyn error::Error + Send + Sync>;
 

datafusion/common/src/lib.rs

@@ -34,7 +34,7 @@ pub mod test_util;
 use arrow::compute::SortOptions;
 pub use column::Column;
 pub use dfschema::{DFField, DFSchema, DFSchemaRef, ExprSchema, ToDFSchema};
-pub use error::{field_not_found, DataFusionError, Result, SchemaError};
+pub use error::{field_not_found, DataFusionError, Result, SchemaError, SharedResult};
 pub use parsers::parse_interval;
 pub use scalar::{ScalarType, ScalarValue};
 pub use stats::{ColumnStatistics, Statistics};

datafusion/core/src/error.rs

@@ -16,4 +16,4 @@
 // under the License.
 
 //! DataFusion error types
-pub use datafusion_common::{DataFusionError, Result};
+pub use datafusion_common::{DataFusionError, Result, SharedResult};

datafusion/core/src/physical_plan/joins/utils.rs

@@ -17,7 +17,7 @@
 
 //! Join related functionality used both on logical and physical plans
 
-use crate::error::{DataFusionError, Result};
+use crate::error::{DataFusionError, Result, SharedResult};
 use crate::logical_expr::JoinType;
 use crate::physical_plan::expressions::Column;
 use crate::physical_plan::SchemaRef;
@@ -438,7 +438,7 @@ impl<T: 'static> OnceAsync<T> {
 }
 
 /// The shared future type used internally within [`OnceAsync`]
-type OnceFutPending<T> = Shared<BoxFuture<'static, Arc<Result<T>>>>;
+type OnceFutPending<T> = Shared<BoxFuture<'static, Arc<SharedResult<T>>>>;
 
 /// A [`OnceFut`] represents a shared asynchronous computation, that will be evaluated
 /// once for all [`Clone`]'s, with [`OnceFut::get`] providing a non-consuming interface
@@ -653,7 +653,7 @@ fn get_int_range(min: ScalarValue, max: ScalarValue) -> Option<usize> {
 
 enum OnceFutState<T> {
     Pending(OnceFutPending<T>),
-    Ready(Arc<Result<T>>),
+    Ready(Arc<SharedResult<T>>),
 }
 
 impl<T> Clone for OnceFutState<T> {
@@ -672,7 +672,11 @@ impl<T: 'static> OnceFut<T> {
         Fut: Future<Output = Result<T>> + Send + 'static,
     {
         Self {
-            state: OnceFutState::Pending(fut.map(Arc::new).boxed().shared()),
+            state: OnceFutState::Pending(
+                fut.map(|res| Arc::new(res.map_err(Arc::new)))
+                    .boxed()
+                    .shared(),
+            ),
         }
     }
 
@@ -692,7 +696,7 @@ impl<T: 'static> OnceFut<T> {
             OnceFutState::Ready(r) => Poll::Ready(
                 r.as_ref()
                     .as_ref()
-                    .map_err(|e| ArrowError::ExternalError(e.to_string().into())),
+                    .map_err(|e| ArrowError::ExternalError(Box::new(e.clone()))),
             ),
         }
     }
@@ -939,6 +943,7 @@ mod tests {
     use super::*;
     use arrow::datatypes::DataType;
     use datafusion_common::ScalarValue;
+    use std::pin::Pin;
 
     fn check(left: &[Column], right: &[Column], on: &[(Column, Column)]) -> Result<()> {
         let left = left
@@ -971,6 +976,41 @@ mod tests {
         assert!(check(&left, &right, on).is_err());
     }
 
+    #[tokio::test]
+    async fn check_error_nesting() {
+        let once_fut = OnceFut::<()>::new(async {
+            Err(DataFusionError::ArrowError(ArrowError::CsvError(
+                "some error".to_string(),
+            )))
+        });
+
+        struct TestFut(OnceFut<()>);
+        impl Future for TestFut {
+            type Output = ArrowResult<()>;
+
+            fn poll(
+                mut self: Pin<&mut Self>,
+                cx: &mut Context<'_>,
+            ) -> Poll<Self::Output> {
+                match ready!(self.0.get(cx)) {
+                    Ok(()) => Poll::Ready(Ok(())),
+                    Err(e) => Poll::Ready(Err(e)),
+                }
+            }
+        }
+
+        let res = TestFut(once_fut).await;
+        let arrow_err_from_fut = res.expect_err("once_fut always return error");
+
+        let wrapped_err = DataFusionError::from(arrow_err_from_fut);
+        let root_err = wrapped_err.find_root();
+
+        assert!(matches!(
+            root_err,
+            DataFusionError::ArrowError(ArrowError::CsvError(_))
+        ))
+    }
+
     #[test]
     fn check_not_in_left() {
         let left = vec![Column::new("b", 0)];