Use tokio::sync::oneshot to prevent FuturesUnordered reentrant drop crash (#9)

Author: sgrebnov

vortex-array/src/arrow/convert.rs

@@ -1797,6 +1797,8 @@ mod tests {
         use arrow_array::MapArray;
         use arrow_array::builder::MapBuilder;
         use arrow_array::builder::StringBuilder;
+        use crate::arrow::executor::ArrowArrayExecutor as _;
+        use crate::executor::VortexSessionExecute as _;
 
         // Build a MapArray: map<string, int32>
         let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
@@ -1826,12 +1828,12 @@ mod tests {
         let list_array = vortex_array.as_::<ListVTable>();
         assert_eq!(list_array.elements().len(), 3); // 3 total key-value pairs
         let struct_elements = list_array.elements().as_::<StructVTable>();
-        assert_eq!(struct_elements.nfields(), 2); // key and value fields
+        assert_eq!(struct_elements.names().len(), 2); // key and value fields
 
         // Convert back to Arrow as a MapArray
         let map_dtype = arrow_map.data_type().clone();
         let arrow_back = vortex_array
-            .execute_arrow(Some(&map_dtype), &mut crate::session::LEGACY_SESSION.create_execution_ctx())
+            .execute_arrow(Some(&map_dtype), &mut crate::LEGACY_SESSION.create_execution_ctx())
             .unwrap();
         let map_back = arrow_back
             .as_any()

vortex-array/src/expr/exprs/case_when.rs

@@ -37,7 +37,7 @@ use crate::compute::zip;
 use crate::expr::Arity;
 use crate::expr::ChildName;
 use crate::expr::ExecutionArgs;
-use crate::expr::ExecutionResult;
+
 use crate::expr::ExprId;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
@@ -168,7 +168,7 @@ impl VTable for CaseWhen {
         &self,
         options: &Self::Options,
         args: ExecutionArgs,
-    ) -> VortexResult<ExecutionResult> {
+    ) -> VortexResult<ArrayRef> {
         let row_count = args.row_count;
         let num_pairs = options.num_when_then_pairs as usize;
 
@@ -222,7 +222,7 @@ impl VTable for CaseWhen {
             result = zip(then_value.as_ref(), result.as_ref(), &mask)?;
         }
 
-        result.execute::<ExecutionResult>(args.ctx)
+        result.execute::<ArrayRef>(args.ctx)
     }
 
     fn is_null_sensitive(&self, _options: &Self::Options) -> bool {
@@ -236,7 +236,7 @@ impl VTable for CaseWhen {
 }
 
 /// Efficient implementation for binary CASE WHEN (single when/then pair)
-fn execute_binary_case_when(_has_else: bool, args: ExecutionArgs) -> VortexResult<ExecutionResult> {
+fn execute_binary_case_when(_has_else: bool, args: ExecutionArgs) -> VortexResult<ArrayRef> {
     let row_count = args.row_count;
 
     // Extract inputs based on arity: [condition, then_value] or [condition, then_value, else_value]
@@ -265,20 +265,20 @@ fn execute_binary_case_when(_has_else: bool, args: ExecutionArgs) -> VortexResul
 
     // Short-circuit: all true -> just return THEN value
     if mask.all_true() {
-        return then_value.execute::<ExecutionResult>(args.ctx);
+        return then_value.execute::<ArrayRef>(args.ctx);
     }
 
     // Short-circuit: all false -> return ELSE value or NULL
     if mask.all_false() {
         return match else_value {
-            Some(else_value) => else_value.execute::<ExecutionResult>(args.ctx),
+            Some(else_value) => else_value.execute::<ArrayRef>(args.ctx),
             None => {
                 // Create NULL constant of appropriate type
                 let then_dtype = then_value.dtype().as_nullable();
-                Ok(ExecutionResult::constant(
+                Ok(ConstantArray::new(
                     Scalar::null(then_dtype),
                     row_count,
-                ))
+                ).into_array())
             }
         };
     }
@@ -292,7 +292,7 @@ fn execute_binary_case_when(_has_else: bool, args: ExecutionArgs) -> VortexResul
     // Use zip to select: where mask is true, take then_value; else take else_value
     let result = zip(then_value.as_ref(), else_value.as_ref(), &mask)?;
 
-    result.execute::<ExecutionResult>(args.ctx)
+    result.execute::<ArrayRef>(args.ctx)
 }
 
 /// Creates an N-ary CASE WHEN expression from a flat list of children.

vortex-btrblocks/src/canonical_compressor.rs

@@ -88,7 +88,7 @@ pub trait CanonicalCompressor {
 ///     .exclude_int([IntCode::Dict])
 ///     .build();
 /// ```
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct BtrBlocksCompressor {
     /// Integer compressor with configured schemes.
     pub int_schemes: Vec<&'static dyn IntegerScheme>,

vortex-datafusion/src/convert/exprs.rs

@@ -28,7 +28,6 @@ use vortex::expr::Expression;
 use vortex::expr::Like;
 use vortex::expr::Operator;
 use vortex::expr::VTableExt;
-use vortex::expr::and;
 use vortex::expr::case_when;
 use vortex::expr::case_when_no_else;
 use vortex::expr::and_collect;

vortex-file/src/read/driver.rs

@@ -9,8 +9,12 @@ use std::task::Context;
 use std::task::Poll;
 
 use futures::Stream;
+#[cfg(all(test, not(feature = "tokio")))]
+use oneshot;
 use pin_project_lite::pin_project;
 use vortex_buffer::Alignment;
+#[cfg(all(test, feature = "tokio"))]
+use tokio::sync::oneshot;
 use vortex_error::VortexExpect;
 use vortex_io::CoalesceConfig;
 use vortex_metrics::Counter;

vortex-file/src/read/request.rs

@@ -8,6 +8,10 @@ use std::ops::Range;
 use std::sync::Arc;
 
 use vortex_array::buffer::BufferHandle;
+#[cfg(not(feature = "tokio"))]
+use oneshot;
+#[cfg(feature = "tokio")]
+use tokio::sync::oneshot;
 use vortex_buffer::Alignment;
 use vortex_error::VortexError;
 use vortex_error::VortexExpect;
@@ -108,7 +112,7 @@ impl Debug for ReadRequest {
 impl ReadRequest {
     pub(crate) fn resolve(self, result: VortexResult<BufferHandle>) {
         if let Err(e) = self.callback.send(result) {
-            tracing::debug!("ReadRequest {} dropped before resolving: {e}", self.id);
+            tracing::debug!("ReadRequest {} dropped before resolving: {e:?}", self.id);
         }
     }
 }

vortex-file/src/segments/source.rs

@@ -13,6 +13,10 @@ use futures::FutureExt;
 use futures::StreamExt;
 use futures::channel::mpsc;
 use futures::future;
+#[cfg(not(feature = "tokio"))]
+use oneshot;
+#[cfg(feature = "tokio")]
+use tokio::sync::oneshot;
 use vortex_array::buffer::BufferHandle;
 use vortex_buffer::Alignment;
 use vortex_error::VortexResult;

vortex-io/src/runtime/handle.rs

@@ -9,6 +9,10 @@ use std::task::Poll;
 use std::task::ready;
 
 use futures::FutureExt;
+#[cfg(not(feature = "tokio"))]
+use oneshot;
+#[cfg(feature = "tokio")]
+use tokio::sync::oneshot;
 use vortex_error::vortex_panic;
 
 use crate::runtime::AbortHandleRef;

vortex-io/src/runtime/single.rs

@@ -8,8 +8,13 @@ use futures::Stream;
 use futures::StreamExt;
 use futures::future::BoxFuture;
 use futures::stream::LocalBoxStream;
+#[cfg(not(feature = "tokio"))]
+use oneshot;
 use parking_lot::Mutex;
 use smol::LocalExecutor;
+// Prefer tokio::sync::oneshot when tokio feature is enabled
+#[cfg(feature = "tokio")]
+use tokio::sync::oneshot;
 use vortex_error::vortex_panic;
 
 use crate::runtime::AbortHandle;
@@ -18,7 +23,6 @@ use crate::runtime::BlockingRuntime;
 use crate::runtime::Executor;
 use crate::runtime::Handle;
 use crate::runtime::smol::SmolAbortHandle;
-
 /// A runtime that drives all work on the current thread.
 ///
 /// This is subtly different from using a current-thread runtime to drive a future since it is