Revert "PendingSegment" (#2710)

Reverts #2699

Going to take a different approach

Author: gatesn

vortex-error/src/lib.rs

@@ -16,7 +16,7 @@ use std::error::Error;
 use std::fmt::{Debug, Display, Formatter};
 use std::num::TryFromIntError;
 use std::ops::Deref;
-use std::sync::{Arc, PoisonError};
+use std::sync::PoisonError;
 use std::{env, fmt, io};
 
 pub use ext::*;
@@ -100,11 +100,6 @@ pub enum VortexError {
     /// A wrapper for other errors, carrying additional context.
     #[error("{0}: {1}")]
     Context(ErrString, #[source] Box<VortexError>),
-    /// A wrapper for errors that need to be shared or cloneable.
-    /// This is useful when the same underlying [`VortexError`] needs to be returned to multiple
-    /// receivers.
-    #[error(transparent)]
-    Shared(Arc<VortexError>),
     /// A wrapper for errors from the Arrow library.
     #[error("{0}\nBacktrace:\n{1}")]
     ArrowError(arrow_schema::ArrowError, Backtrace),
@@ -247,20 +242,6 @@ impl Debug for VortexError {
 /// A type alias for Results that return VortexErrors as their error type.
 pub type VortexResult<T> = Result<T, VortexError>;
 
-/// A vortex result that can be shared or cloned.
-pub type SharedVortexResult<T> = Result<T, Arc<VortexError>>;
-
-impl From<&Arc<VortexError>> for VortexError {
-    fn from(e: &Arc<VortexError>) -> Self {
-        if let VortexError::Shared(e_inner) = e.as_ref() {
-            // don't re-wrap
-            VortexError::Shared(Arc::clone(e_inner))
-        } else {
-            VortexError::Shared(Arc::clone(e))
-        }
-    }
-}
-
 /// A trait for unwrapping a VortexResult.
 pub trait VortexUnwrap {
     /// The type of the value being unwrapped.

vortex-file/src/generic.rs

@@ -16,7 +16,7 @@ use vortex_error::{VortexExpect, VortexResult, vortex_err, vortex_panic};
 use vortex_io::VortexReadAt;
 use vortex_layout::instrument;
 use vortex_layout::scan::ScanDriver;
-use vortex_layout::segments::{SegmentEvent, SegmentId, SegmentReader, SegmentStream};
+use vortex_layout::segments::{AsyncSegmentReader, SegmentEvent, SegmentId, SegmentStream};
 use vortex_metrics::{Counter, VortexMetrics};
 
 use crate::footer::{Footer, Segment};
@@ -157,7 +157,7 @@ impl InflightSegments {
 impl<R: VortexReadAt> GenericScanDriver<R> {}
 
 impl<R: VortexReadAt> ScanDriver for GenericScanDriver<R> {
-    fn segment_reader(&self) -> Arc<dyn SegmentReader> {
+    fn segment_reader(&self) -> Arc<dyn AsyncSegmentReader> {
         self.segment_channel.reader()
     }
 

vortex-file/src/memory.rs

@@ -1,10 +1,11 @@
 use std::sync::Arc;
 
+use async_trait::async_trait;
 use futures::{Stream, stream};
 use vortex_buffer::ByteBuffer;
 use vortex_error::{VortexResult, vortex_err};
 use vortex_layout::scan::ScanDriver;
-use vortex_layout::segments::{PendingSegment, SegmentId, SegmentReader, SegmentStream};
+use vortex_layout::segments::{AsyncSegmentReader, SegmentId, SegmentStream};
 use vortex_metrics::VortexMetrics;
 
 use crate::segments::SegmentCache;
@@ -47,7 +48,7 @@ impl FileType for InMemoryVortexFile {
 }
 
 impl ScanDriver for InMemoryVortexFile {
-    fn segment_reader(&self) -> Arc<dyn SegmentReader> {
+    fn segment_reader(&self) -> Arc<dyn AsyncSegmentReader> {
         Arc::new(self.clone())
     }
 
@@ -56,8 +57,9 @@ impl ScanDriver for InMemoryVortexFile {
     }
 }
 
-impl SegmentReader for InMemoryVortexFile {
-    fn get(&self, id: SegmentId) -> VortexResult<Arc<dyn PendingSegment>> {
+#[async_trait]
+impl AsyncSegmentReader for InMemoryVortexFile {
+    async fn get(&self, id: SegmentId) -> VortexResult<ByteBuffer> {
         let segment: &Segment = self
             .footer
             .segment_map()
@@ -67,6 +69,6 @@ impl SegmentReader for InMemoryVortexFile {
         let start = usize::try_from(segment.offset).map_err(|_| vortex_err!("offset too large"))?;
         let end = start + segment.length as usize;
 
-        Ok(Arc::new(self.buffer.slice(start..end)))
+        Ok(self.buffer.slice(start..end))
     }
 }

vortex-file/src/segments/channel.rs

@@ -1,10 +1,12 @@
 use std::sync::Arc;
 
+use async_trait::async_trait;
 use futures::channel::mpsc;
-use futures::{FutureExt, Stream};
+use futures::{SinkExt, Stream};
 use oneshot;
+use vortex_buffer::ByteBuffer;
 use vortex_error::{VortexResult, vortex_err};
-use vortex_layout::segments::{PendingSegment, SegmentId, SegmentReader, SharedPendingSegment};
+use vortex_layout::segments::{AsyncSegmentReader, SegmentId};
 
 use crate::segments::SegmentRequest;
 
@@ -29,7 +31,7 @@ impl SegmentChannel {
     }
 
     /// Returns a reader for the segment cache.
-    pub fn reader(&self) -> Arc<dyn SegmentReader + 'static> {
+    pub fn reader(&self) -> Arc<dyn AsyncSegmentReader + 'static> {
         Arc::new(SegmentChannelReader(self.request_send.clone()))
     }
 
@@ -41,8 +43,9 @@ impl SegmentChannel {
 
 struct SegmentChannelReader(mpsc::UnboundedSender<SegmentRequest>);
 
-impl SegmentReader for SegmentChannelReader {
-    fn get(&self, id: SegmentId) -> VortexResult<Arc<dyn PendingSegment>> {
+#[async_trait]
+impl AsyncSegmentReader for SegmentChannelReader {
+    async fn get(&self, id: SegmentId) -> VortexResult<ByteBuffer> {
         // Set up a channel to send the segment back to the caller.
         let (send, recv) = oneshot::channel();
 
@@ -52,12 +55,17 @@ impl SegmentReader for SegmentChannelReader {
         // Send a request to the segment channel.
         self.0
             .clone()
-            .unbounded_send(SegmentRequest { id, callback: send })
+            .send(SegmentRequest { id, callback: send })
+            .await
             .map_err(|e| vortex_err!("Failed to request segment {} {:?}", id, e))?;
 
         // Await the callback
-        Ok(Arc::new(SharedPendingSegment::new(recv.map(|r| {
-            r.unwrap_or_else(|_recv| Err(vortex_err!("Segment request handler was dropped")))
-        }))))
+        match recv.await {
+            Ok(result) => result,
+            Err(_canceled) => {
+                // The sender was dropped before returning a result to us
+                Err(vortex_err!("Segment request handler was dropped {}", id,))
+            }
+        }
     }
 }

vortex-layout/src/data.rs

@@ -12,7 +12,7 @@ use vortex_flatbuffers::{FlatBuffer, FlatBufferRoot, WriteFlatBuffer, layout};
 use crate::LayoutId;
 use crate::context::LayoutContext;
 use crate::reader::LayoutReader;
-use crate::segments::{SegmentCollector, SegmentId, SegmentReader};
+use crate::segments::{AsyncSegmentReader, SegmentCollector, SegmentId};
 use crate::vtable::LayoutVTableRef;
 
 /// [`Layout`] is the lazy equivalent to [`vortex_array::ArrayRef`], providing a hierarchical
@@ -258,7 +258,7 @@ impl Layout {
     /// Create a reader for this layout.
     pub fn reader(
         &self,
-        segment_reader: Arc<dyn SegmentReader>,
+        segment_reader: Arc<dyn AsyncSegmentReader>,
         ctx: ArrayContext,
     ) -> VortexResult<Arc<dyn LayoutReader>> {
         self.vtable().reader(self.clone(), ctx, segment_reader)

vortex-layout/src/layouts/chunked/eval_expr.rs

@@ -80,14 +80,14 @@ mod test {
     use vortex_expr::Identity;
 
     use crate::layouts::chunked::writer::ChunkedLayoutWriter;
-    use crate::segments::SegmentReader;
+    use crate::segments::AsyncSegmentReader;
     use crate::segments::test::TestSegments;
     use crate::writer::LayoutWriterExt;
     use crate::{Layout, RowMask};
 
     #[fixture]
     /// Create a chunked layout with three chunks of primitive arrays.
-    fn chunked_layout() -> (ArrayContext, Arc<dyn SegmentReader>, Layout) {
+    fn chunked_layout() -> (ArrayContext, Arc<dyn AsyncSegmentReader>, Layout) {
         let ctx = ArrayContext::empty();
         let mut segments = TestSegments::default();
         let layout = ChunkedLayoutWriter::new(
@@ -111,7 +111,7 @@ mod test {
     fn test_chunked_evaluator(
         #[from(chunked_layout)] (ctx, segments, layout): (
             ArrayContext,
-            Arc<dyn SegmentReader>,
+            Arc<dyn AsyncSegmentReader>,
             Layout,
         ),
     ) {

vortex-layout/src/layouts/chunked/mod.rs

@@ -12,7 +12,7 @@ use vortex_error::VortexResult;
 use crate::data::Layout;
 use crate::layouts::chunked::reader::ChunkedReader;
 use crate::reader::{LayoutReader, LayoutReaderExt};
-use crate::segments::{SegmentCollector, SegmentReader};
+use crate::segments::{AsyncSegmentReader, SegmentCollector};
 use crate::vtable::LayoutVTable;
 use crate::{CHUNKED_LAYOUT_ID, LayoutId};
 
@@ -29,7 +29,7 @@ impl LayoutVTable for ChunkedLayout {
         &self,
         layout: Layout,
         ctx: ArrayContext,
-        segment_reader: Arc<dyn SegmentReader>,
+        segment_reader: Arc<dyn AsyncSegmentReader>,
     ) -> VortexResult<Arc<dyn LayoutReader>> {
         Ok(ChunkedReader::try_new(layout, ctx, segment_reader)?.into_arc())
     }

vortex-layout/src/layouts/chunked/reader.rs

@@ -7,14 +7,14 @@ use vortex_error::{VortexResult, vortex_panic};
 
 use crate::layouts::chunked::ChunkedLayout;
 use crate::reader::LayoutReader;
-use crate::segments::SegmentReader;
+use crate::segments::AsyncSegmentReader;
 use crate::{Layout, LayoutVTable};
 
 #[derive(Clone)]
 pub struct ChunkedReader {
     layout: Layout,
     ctx: ArrayContext,
-    segment_reader: Arc<dyn SegmentReader>,
+    segment_reader: Arc<dyn AsyncSegmentReader>,
 
     /// Shared lazy chunk scanners
     chunk_readers: Arc<[OnceLock<Arc<dyn LayoutReader>>]>,
@@ -26,7 +26,7 @@ impl ChunkedReader {
     pub(super) fn try_new(
         layout: Layout,
         ctx: ArrayContext,
-        segment_reader: Arc<dyn SegmentReader>,
+        segment_reader: Arc<dyn AsyncSegmentReader>,
     ) -> VortexResult<Self> {
         if layout.vtable().id() != ChunkedLayout.id() {
             vortex_panic!("Mismatched layout ID")

vortex-layout/src/layouts/flat/mod.rs

@@ -11,7 +11,7 @@ use vortex_error::VortexResult;
 
 use crate::layouts::flat::reader::FlatReader;
 use crate::reader::{LayoutReader, LayoutReaderExt};
-use crate::segments::{SegmentCollector, SegmentReader};
+use crate::segments::{AsyncSegmentReader, SegmentCollector};
 use crate::vtable::LayoutVTable;
 use crate::{FLAT_LAYOUT_ID, Layout, LayoutId};
 
@@ -27,7 +27,7 @@ impl LayoutVTable for FlatLayout {
         &self,
         layout: Layout,
         ctx: ArrayContext,
-        segment_reader: Arc<dyn SegmentReader>,
+        segment_reader: Arc<dyn AsyncSegmentReader>,
     ) -> VortexResult<Arc<dyn LayoutReader>> {
         Ok(FlatReader::try_new(layout, ctx, segment_reader)?.into_arc())
     }

vortex-layout/src/layouts/flat/reader.rs

@@ -7,13 +7,13 @@ use vortex_error::{VortexExpect, VortexResult, vortex_err, vortex_panic};
 
 use crate::layouts::flat::FlatLayout;
 use crate::reader::LayoutReader;
-use crate::segments::SegmentReader;
+use crate::segments::AsyncSegmentReader;
 use crate::{Layout, LayoutReaderExt, LayoutVTable, instrument};
 
 pub struct FlatReader {
     layout: Layout,
     ctx: ArrayContext,
-    segment_reader: Arc<dyn SegmentReader>,
+    segment_reader: Arc<dyn AsyncSegmentReader>,
     // TODO(ngates): we need to add an invalidate_row_range function to evict these from the
     //  cache.
     array: Arc<OnceCell<ArrayRef>>,
@@ -23,7 +23,7 @@ impl FlatReader {
     pub(crate) fn try_new(
         layout: Layout,
         ctx: ArrayContext,
-        segment_reader: Arc<dyn SegmentReader>,
+        segment_reader: Arc<dyn AsyncSegmentReader>,
     ) -> VortexResult<Self> {
         if layout.vtable().id() != FlatLayout.id() {
             vortex_panic!("Mismatched layout ID")
@@ -59,7 +59,7 @@ impl FlatReader {
                     );
 
                     // Fetch all the array segment.
-                    let buffer = self.segment_reader.get(segment_id)?.resolve().await?;
+                    let buffer = self.segment_reader.get(segment_id).await?;
                     let row_count = usize::try_from(self.layout().row_count())
                         .vortex_expect("FlatLayout row count does not fit within usize");
 

vortex-layout/src/layouts/stats/eval_expr.rs

@@ -82,14 +82,14 @@ mod test {
     use crate::layouts::chunked::writer::ChunkedLayoutWriter;
     use crate::layouts::flat::FlatLayout;
     use crate::layouts::stats::writer::{StatsLayoutOptions, StatsLayoutWriter};
-    use crate::segments::SegmentReader;
+    use crate::segments::AsyncSegmentReader;
     use crate::segments::test::TestSegments;
     use crate::writer::LayoutWriterExt;
     use crate::{ExprEvaluator, Layout, RowMask};
 
     #[fixture]
     /// Create a stats layout with three chunks of primitive arrays.
-    fn stats_layout() -> (ArrayContext, Arc<dyn SegmentReader>, Layout) {
+    fn stats_layout() -> (ArrayContext, Arc<dyn AsyncSegmentReader>, Layout) {
         let ctx = ArrayContext::empty();
         let mut segments = TestSegments::default();
         let layout = StatsLayoutWriter::try_new(
@@ -124,7 +124,7 @@ mod test {
     fn test_stats_evaluator(
         #[from(stats_layout)] (ctx, segments, layout): (
             ArrayContext,
-            Arc<dyn SegmentReader>,
+            Arc<dyn AsyncSegmentReader>,
             Layout,
         ),
     ) {
@@ -150,7 +150,7 @@ mod test {
     fn test_stats_pruning_mask(
         #[from(stats_layout)] (ctx, segments, layout): (
             ArrayContext,
-            Arc<dyn SegmentReader>,
+            Arc<dyn AsyncSegmentReader>,
             Layout,
         ),
     ) {

vortex-layout/src/layouts/stats/mod.rs

@@ -13,7 +13,7 @@ use vortex_error::VortexResult;
 use crate::data::Layout;
 use crate::layouts::stats::reader::StatsReader;
 use crate::reader::{LayoutReader, LayoutReaderExt};
-use crate::segments::{RequiredSegmentKind, SegmentCollector, SegmentReader};
+use crate::segments::{AsyncSegmentReader, RequiredSegmentKind, SegmentCollector};
 use crate::vtable::LayoutVTable;
 use crate::{LayoutId, STATS_LAYOUT_ID};
 
@@ -30,7 +30,7 @@ impl LayoutVTable for StatsLayout {
         &self,
         layout: Layout,
         ctx: ArrayContext,
-        segment_reader: Arc<dyn SegmentReader>,
+        segment_reader: Arc<dyn AsyncSegmentReader>,
     ) -> VortexResult<Arc<dyn LayoutReader>> {
         Ok(StatsReader::try_new(layout, ctx, segment_reader)?.into_arc())
     }

vortex-layout/src/layouts/stats/reader.rs

@@ -14,7 +14,7 @@ use vortex_mask::Mask;
 use crate::layouts::stats::StatsLayout;
 use crate::layouts::stats::stats_table::StatsTable;
 use crate::reader::LayoutReader;
-use crate::segments::SegmentReader;
+use crate::segments::AsyncSegmentReader;
 use crate::{ExprEvaluator, Layout, LayoutVTable, RowMask};
 
 type PruningCache = Arc<OnceCell<Option<Mask>>>;
@@ -23,7 +23,7 @@ type PruningCache = Arc<OnceCell<Option<Mask>>>;
 pub struct StatsReader {
     layout: Layout,
     ctx: ArrayContext,
-    segment_reader: Arc<dyn SegmentReader>,
+    segment_reader: Arc<dyn AsyncSegmentReader>,
 
     /// The number of blocks
     nblocks: usize,
@@ -43,7 +43,7 @@ impl StatsReader {
     pub(super) fn try_new(
         layout: Layout,
         ctx: ArrayContext,
-        segment_reader: Arc<dyn SegmentReader>,
+        segment_reader: Arc<dyn AsyncSegmentReader>,
     ) -> VortexResult<Self> {
         if layout.vtable().id() != StatsLayout.id() {
             vortex_panic!("Mismatched layout ID")