vortex_ipc/
iterator.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::io::Read;
5use std::io::Write;
6
7use bytes::Bytes;
8use bytes::BytesMut;
9use itertools::Itertools;
10use vortex_array::ArrayRef;
11use vortex_array::dtype::DType;
12use vortex_array::iter::ArrayIterator;
13use vortex_error::VortexResult;
14use vortex_error::vortex_bail;
15use vortex_error::vortex_err;
16use vortex_session::VortexSession;
17
18use crate::messages::DecoderMessage;
19use crate::messages::EncoderMessage;
20use crate::messages::MessageEncoder;
21use crate::messages::SyncMessageReader;
22
23/// An [`ArrayIterator`] for reading messages off an IPC stream.
24pub struct SyncIPCReader<R: Read> {
25    reader: SyncMessageReader<R>,
26    dtype: DType,
27    session: VortexSession,
28}
29
30impl<R: Read> SyncIPCReader<R> {
31    pub fn try_new(read: R, session: &VortexSession) -> VortexResult<Self> {
32        let mut reader = SyncMessageReader::new(read);
33        match reader.next().transpose()? {
34            Some(msg) => match msg {
35                DecoderMessage::DType(fb_dtype) => {
36                    let dtype = DType::from_flatbuffer(fb_dtype, session)?;
37                    Ok(SyncIPCReader {
38                        reader,
39                        dtype,
40                        session: session.clone(),
41                    })
42                }
43                msg => {
44                    vortex_bail!("Expected DType message, got {:?}", msg);
45                }
46            },
47            None => vortex_bail!("Expected DType message, got EOF"),
48        }
49    }
50}
51
52impl<R: Read> ArrayIterator for SyncIPCReader<R> {
53    fn dtype(&self) -> &DType {
54        &self.dtype
55    }
56}
57
58impl<R: Read> Iterator for SyncIPCReader<R> {
59    type Item = VortexResult<ArrayRef>;
60
61    fn next(&mut self) -> Option<Self::Item> {
62        match self.reader.next()? {
63            Ok(msg) => match msg {
64                DecoderMessage::Array((array_parts, ctx, row_count)) => Some(
65                    array_parts
66                        .decode(&self.dtype, row_count, &ctx, &self.session)
67                        .and_then(|array| {
68                            if array.dtype() != self.dtype() {
69                                Err(vortex_err!(
70                                    "Array data type mismatch: expected {:?}, got {:?}",
71                                    self.dtype(),
72                                    array.dtype()
73                                ))
74                            } else {
75                                Ok(array)
76                            }
77                        }),
78                ),
79                msg => Some(Err(vortex_err!("Expected Array message, got {:?}", msg))),
80            },
81            Err(e) => Some(Err(e)),
82        }
83    }
84}
85
86/// A trait for converting an [`ArrayIterator`] into an IPC stream.
87pub trait ArrayIteratorIPC {
88    fn into_ipc(self, session: &VortexSession) -> VortexResult<ArrayIteratorIPCBytes>
89    where
90        Self: Sized;
91
92    fn write_ipc<W: Write>(self, write: W, session: &VortexSession) -> VortexResult<W>
93    where
94        Self: Sized;
95}
96
97impl<I: ArrayIterator + 'static> ArrayIteratorIPC for I {
98    fn into_ipc(self, session: &VortexSession) -> VortexResult<ArrayIteratorIPCBytes>
99    where
100        Self: Sized,
101    {
102        let mut encoder = MessageEncoder::new(session.clone());
103        let buffers = encoder.encode(EncoderMessage::DType(self.dtype()))?;
104        Ok(ArrayIteratorIPCBytes {
105            inner: Box::new(self),
106            encoder,
107            buffers,
108        })
109    }
110
111    fn write_ipc<W: Write>(self, mut write: W, session: &VortexSession) -> VortexResult<W>
112    where
113        Self: Sized,
114    {
115        let mut stream = self.into_ipc(session)?;
116        for buffer in &mut stream {
117            write.write_all(buffer?.as_ref())?;
118        }
119        Ok(write)
120    }
121}
122
123pub struct ArrayIteratorIPCBytes {
124    inner: Box<dyn ArrayIterator + 'static>,
125    encoder: MessageEncoder,
126    buffers: Vec<Bytes>,
127}
128
129impl ArrayIteratorIPCBytes {
130    /// Collects the IPC bytes into a single `Bytes`.
131    pub fn collect_to_buffer(self) -> VortexResult<Bytes> {
132        let buffers: Vec<Bytes> = self.try_collect()?;
133        let mut buffer = BytesMut::with_capacity(buffers.iter().map(|b| b.len()).sum());
134        for buf in buffers {
135            buffer.extend_from_slice(buf.as_ref());
136        }
137        Ok(buffer.freeze())
138    }
139}
140
141impl Iterator for ArrayIteratorIPCBytes {
142    type Item = VortexResult<Bytes>;
143
144    fn next(&mut self) -> Option<Self::Item> {
145        // Try to flush any buffers we have
146        if !self.buffers.is_empty() {
147            return Some(Ok(self.buffers.remove(0)));
148        }
149
150        // Or else try to serialize the next array
151        match self
152            .inner
153            .next()?
154            .and_then(|chunk| self.encoder.encode(EncoderMessage::Array(&chunk)))
155        {
156            Ok(buffers) => self.buffers.extend(buffers),
157            Err(e) => return Some(Err(e)),
158        }
159
160        // Try to flush any buffers we have again
161        if !self.buffers.is_empty() {
162            return Some(Ok(self.buffers.remove(0)));
163        }
164
165        // Otherwise, we're done
166        None
167    }
168}
169
170#[cfg(test)]
171mod test {
172    use std::io::Cursor;
173
174    use vortex_array::IntoArray as _;
175    use vortex_array::VortexSessionExecute;
176    use vortex_array::assert_arrays_eq;
177    use vortex_array::iter::ArrayIterator;
178    use vortex_array::iter::ArrayIteratorExt;
179    use vortex_buffer::buffer;
180
181    use super::*;
182    use crate::test::SESSION;
183
184    #[test]
185    fn test_sync_stream() -> VortexResult<()> {
186        let array = buffer![1i32, 2, 3].into_array();
187        let ipc_buffer = array
188            .to_array_iterator()
189            .into_ipc(&SESSION)?
190            .collect_to_buffer()?;
191
192        let reader = SyncIPCReader::try_new(Cursor::new(ipc_buffer), &SESSION)?;
193
194        assert_eq!(reader.dtype(), array.dtype());
195        let result = reader.read_all()?;
196        assert_arrays_eq!(result, array, &mut SESSION.create_execution_ctx());
197
198        Ok(())
199    }
200}
vortex_ipc/iterator.rs

vortex_ipc/
iterator.rs
