fix[vortex-{scan,io,datafusion}]: propagate tracing spans (#4939)

This change correctly connects object storage spans with the expected
parent spans. Previously, object store spans were orphaned.

Signed-off-by: Alfonso Subiotto Marques <alfonso.subiotto@polarsignals.com>

Author: asubiotto

Cargo.lock

@@ -19,6 +19,7 @@ use datafusion_pruning::FilePruner;
 use futures::{FutureExt, StreamExt, TryStreamExt, stream};
 use object_store::ObjectStore;
 use object_store::path::Path;
+use tracing::Instrument;
 use vortex::dtype::FieldName;
 use vortex::error::VortexError;
 use vortex::expr::{root, select};
@@ -266,6 +267,7 @@ impl FileOpener for VortexOpener {
 
             Ok(stream)
         }
+        .in_current_span()
         .boxed())
     }
 }

vortex-datafusion/src/persistent/opener.rs

@@ -9,6 +9,7 @@ use async_compat::Compat;
 use futures::future::BoxFuture;
 use futures::stream::BoxStream;
 use futures::{FutureExt, StreamExt};
+use tracing::Instrument;
 use vortex_buffer::ByteBufferMut;
 use vortex_error::{VortexError, VortexResult};
 
@@ -145,7 +146,8 @@ impl ReadSource for ObjectStoreIoSource {
                     };
 
                     Ok(buffer.freeze())
-                };
+                }
+                .in_current_span();
 
                 async move { req.resolve(Compat::new(read).await) }
             })

vortex-io/src/file/object_store.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use futures::future::BoxFuture;
+use tracing::Instrument;
 
 use crate::runtime::{AbortHandle, AbortHandleRef, Executor, IoTask};
 
@@ -22,7 +23,7 @@ impl Executor for smol::Executor<'static> {
     }
 
     fn spawn_io(&self, task: IoTask) {
-        smol::Executor::spawn(self, task.source.drive_send(task.stream)).detach()
+        smol::Executor::spawn(self, task.source.drive_send(task.stream).in_current_span()).detach()
     }
 }
 

vortex-io/src/runtime/smol.rs

@@ -4,6 +4,7 @@
 use std::sync::{Arc, LazyLock};
 
 use futures::future::BoxFuture;
+use tracing::Instrument;
 
 use crate::runtime::{AbortHandle, AbortHandleRef, BlockingRuntime, Executor, Handle, IoTask};
 
@@ -46,7 +47,7 @@ impl Executor for tokio::runtime::Handle {
     }
 
     fn spawn_io(&self, task: IoTask) {
-        tokio::runtime::Handle::spawn(self, task.source.drive_send(task.stream));
+        tokio::runtime::Handle::spawn(self, task.source.drive_send(task.stream).in_current_span());
     }
 }
 
@@ -75,7 +76,8 @@ impl Executor for CurrentTokioRuntime {
     }
 
     fn spawn_io(&self, task: IoTask) {
-        tokio::runtime::Handle::current().spawn(task.source.drive_send(task.stream));
+        tokio::runtime::Handle::current()
+            .spawn(task.source.drive_send(task.stream).in_current_span());
     }
 }
 

vortex-io/src/runtime/tokio.rs

@@ -39,6 +39,7 @@ parking_lot = { workspace = true }
 roaring = { workspace = true, optional = true }
 sketches-ddsketch = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread"], optional = true }
+tracing = { workspace = true }
 
 [dev-dependencies]
 loom = { workspace = true }