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_util::{Stream, stream};
9use vortex_dtype::DType;
10use vortex_error::VortexResult;
11
12use crate::{Array, ArrayRef};
13
14mod adapter;
15mod ext;
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
33impl dyn Array + '_ {
34    /// Create an [`ArrayStream`] over the array.
35    pub fn to_array_stream(&self) -> impl ArrayStream + 'static {
36        ArrayStreamAdapter::new(self.dtype().clone(), stream::iter(self.to_array_iterator()))
37    }
38}