Skip to main content

vortex_array/stream/
mod.rs

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