vortex_array/stream/
mod.rs

1use std::pin::Pin;
2
3pub use adapter::*;
4pub use ext::*;
5use futures_util::{Stream, stream};
6pub use take_rows::*;
7use vortex_dtype::DType;
8use vortex_error::VortexResult;
9
10use crate::iter::ArrayIteratorArrayExt;
11use crate::{Array, ArrayRef};
12
13mod adapter;
14mod ext;
15mod take_rows;
16
17/// A stream of array chunks along with a DType.
18///
19/// Can be thought of as equivalent to Arrow's RecordBatchReader.
20pub trait ArrayStream: Stream<Item = VortexResult<ArrayRef>> {
21    fn dtype(&self) -> &DType;
22}
23
24/// Trait for a [`Stream`] of [`ArrayRef`]s that can be passed between threads.
25pub type SendableArrayStream = Pin<Box<dyn ArrayStream + Send>>;
26
27impl ArrayStream for SendableArrayStream {
28    fn dtype(&self) -> &DType {
29        (**self).dtype()
30    }
31}
32
33pub trait ArrayStreamArrayExt: Array {
34    /// Create an [`ArrayStream`] over the array.
35    fn to_array_stream(&self) -> impl ArrayStream + 'static {
36        ArrayStreamAdapter::new(self.dtype().clone(), stream::iter(self.to_array_iterator()))
37    }
38}
39
40impl<A: ?Sized + Array> ArrayStreamArrayExt for A {}