execution failing test

Author: robtandy

Cargo.lock

@@ -13,13 +13,15 @@ tokio = { version = "1.46.1", features = ["full"] }
 tonic = { version = "0.12.3", features = ["transport"] }
 tower = "0.5.2"
 http = "1.3.1"
+itertools = "0.14.0"
 futures = "0.3.31"
 url = "2.5.4"
 uuid = "1.17.0"
 delegate = "0.13.4"
 dashmap = "6.1.0"
 prost = "0.13.5"
+rand = "0.8.5"
 object_store = "0.12.3"
 
 [dev-dependencies]
-insta = { version = "1.43.1", features = ["filters"] }
+insta = { version = "1.43.1", features = ["filters"] }

Cargo.toml

@@ -0,0 +1,20 @@
+use url::Url;
+use uuid::Uuid;
+
+#[derive(Debug, Clone)]
+pub struct StageContext {
+    /// Unique identifier of the Stage.
+    pub id: Uuid,
+    /// Number of tasks involved in the query.
+    pub n_tasks: usize,
+    /// Unique identifier of the input Stage.
+    pub input_id: Uuid,
+    /// Urls from which the current stage will need to read data.
+    pub input_urls: Vec<Url>,
+}
+
+#[derive(Debug, Clone)]
+pub struct StageTaskContext {
+    /// Index of the current task in a stage
+    pub task_idx: usize,
+}

context.rs

@@ -0,0 +1,3 @@
+[toolchain]
+channel = "1.83.0"
+profile = "default"

rust-toolchain.toml

@@ -1,12 +1,14 @@
 use async_trait::async_trait;
 use datafusion::common::internal_datafusion_err;
 use datafusion::error::DataFusionError;
-use datafusion::prelude::SessionConfig;
+use datafusion::execution::TaskContext;
+use datafusion::prelude::{SessionConfig, SessionContext};
 use delegate::delegate;
 use std::sync::Arc;
 use tonic::body::BoxBody;
 use url::Url;
 
+#[derive(Clone)]
 pub struct ChannelManager(Arc<dyn ChannelResolver + Send + Sync>);
 
 impl ChannelManager {
@@ -21,29 +23,49 @@ pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService<
     tonic::transport::Error,
 >;
 
-#[derive(Clone, Debug)]
-pub struct ArrowFlightChannel {
-    pub url: Url,
-    pub channel: BoxCloneSyncChannel,
-}
-
+/// Abstracts networking details so that users can implement their own network resolution
+/// mechanism.
 #[async_trait]
 pub trait ChannelResolver {
-    async fn get_n_channels(&self, n: usize) -> Result<Vec<ArrowFlightChannel>, DataFusionError>;
-    async fn get_channel_for_url(&self, url: &Url) -> Result<ArrowFlightChannel, DataFusionError>;
+    /// Gets all available worker URLs. Used during stage assignment.
+    fn get_urls(&self) -> Result<Vec<Url>, DataFusionError>;
+    /// For a given URL, get a channel for communicating to it.
+    async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError>;
 }
 
 impl ChannelManager {
-    pub fn try_from_session(session: &SessionConfig) -> Result<Arc<Self>, DataFusionError> {
-        session
-            .get_extension::<ChannelManager>()
-            .ok_or_else(|| internal_datafusion_err!("No extension ChannelManager"))
-    }
-
     delegate! {
         to self.0 {
-            pub async fn get_n_channels(&self, n: usize) -> Result<Vec<ArrowFlightChannel>, DataFusionError>;
-            pub async fn get_channel_for_url(&self, url: &Url) -> Result<ArrowFlightChannel, DataFusionError>;
+            pub fn get_urls(&self) -> Result<Vec<Url>, DataFusionError>;
+            pub async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError>;
         }
     }
 }
+
+impl TryInto<ChannelManager> for &SessionConfig {
+    type Error = DataFusionError;
+
+    fn try_into(self) -> Result<ChannelManager, Self::Error> {
+        Ok(self
+            .get_extension::<ChannelManager>()
+            .ok_or_else(|| internal_datafusion_err!("No extension ChannelManager"))?
+            .as_ref()
+            .clone())
+    }
+}
+
+impl TryInto<ChannelManager> for &TaskContext {
+    type Error = DataFusionError;
+
+    fn try_into(self) -> Result<ChannelManager, Self::Error> {
+        self.session_config().try_into()
+    }
+}
+
+impl TryInto<ChannelManager> for &SessionContext {
+    type Error = DataFusionError;
+
+    fn try_into(self) -> Result<ChannelManager, Self::Error> {
+        self.task_ctx().as_ref().try_into()
+    }
+}

src/channel_manager.rs

@@ -40,19 +40,17 @@ impl IoErrorProto {
                 ErrorKind::WriteZero => 22,
                 ErrorKind::StorageFull => 23,
                 ErrorKind::NotSeekable => 24,
-                ErrorKind::QuotaExceeded => 25,
-                ErrorKind::FileTooLarge => 26,
-                ErrorKind::ResourceBusy => 27,
-                ErrorKind::ExecutableFileBusy => 28,
-                ErrorKind::Deadlock => 29,
-                ErrorKind::CrossesDevices => 30,
-                ErrorKind::TooManyLinks => 31,
-                ErrorKind::ArgumentListTooLong => 32,
-                ErrorKind::Interrupted => 33,
-                ErrorKind::Unsupported => 34,
-                ErrorKind::UnexpectedEof => 35,
-                ErrorKind::OutOfMemory => 36,
-                ErrorKind::Other => 37,
+                ErrorKind::FileTooLarge => 25,
+                ErrorKind::ResourceBusy => 26,
+                ErrorKind::ExecutableFileBusy => 27,
+                ErrorKind::Deadlock => 28,
+                ErrorKind::TooManyLinks => 29,
+                ErrorKind::ArgumentListTooLong => 30,
+                ErrorKind::Interrupted => 31,
+                ErrorKind::Unsupported => 32,
+                ErrorKind::UnexpectedEof => 33,
+                ErrorKind::OutOfMemory => 34,
+                ErrorKind::Other => 35,
                 _ => -1,
             },
             err: err.to_string(),
@@ -86,19 +84,17 @@ impl IoErrorProto {
             22 => ErrorKind::WriteZero,
             23 => ErrorKind::StorageFull,
             24 => ErrorKind::NotSeekable,
-            25 => ErrorKind::QuotaExceeded,
-            26 => ErrorKind::FileTooLarge,
-            27 => ErrorKind::ResourceBusy,
-            28 => ErrorKind::ExecutableFileBusy,
-            29 => ErrorKind::Deadlock,
-            30 => ErrorKind::CrossesDevices,
-            31 => ErrorKind::TooManyLinks,
-            32 => ErrorKind::ArgumentListTooLong,
-            33 => ErrorKind::Interrupted,
-            34 => ErrorKind::Unsupported,
-            35 => ErrorKind::UnexpectedEof,
-            36 => ErrorKind::OutOfMemory,
-            37 => ErrorKind::Other,
+            25 => ErrorKind::FileTooLarge,
+            26 => ErrorKind::ResourceBusy,
+            27 => ErrorKind::ExecutableFileBusy,
+            28 => ErrorKind::Deadlock,
+            29 => ErrorKind::TooManyLinks,
+            30 => ErrorKind::ArgumentListTooLong,
+            31 => ErrorKind::Interrupted,
+            32 => ErrorKind::Unsupported,
+            33 => ErrorKind::UnexpectedEof,
+            34 => ErrorKind::OutOfMemory,
+            35 => ErrorKind::Other,
             _ => ErrorKind::Other,
         };
         (