vortex_layout/scan/
arrow.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use arrow_array::RecordBatch;
7use arrow_array::RecordBatchReader;
8use arrow_array::cast::AsArray;
9use arrow_schema::ArrowError;
10use arrow_schema::Field;
11use arrow_schema::SchemaRef;
12use futures::Stream;
13use futures::TryStreamExt;
14use vortex_array::ArrayRef;
15use vortex_array::ExecutionCtx;
16use vortex_array::VortexSessionExecute;
17use vortex_array::arrow::ArrowSessionExt;
18use vortex_error::VortexResult;
19use vortex_io::runtime::BlockingRuntime;
20
21use crate::scan::scan_builder::ScanBuilder;
22
23impl ScanBuilder<ArrayRef> {
24    /// Creates a new `RecordBatchReader` from the scan builder.
25    ///
26    /// The `schema` parameter is used to define the schema of the resulting record batches. In
27    /// general, it is not possible to exactly infer an Arrow schema from a Vortex
28    /// [`vortex_array::dtype::DType`], therefore it is required to be provided explicitly.
29    pub fn into_record_batch_reader<B: BlockingRuntime>(
30        self,
31        schema: SchemaRef,
32        runtime: &B,
33    ) -> VortexResult<impl RecordBatchReader + 'static> {
34        let struct_field = Field::new_struct("", schema.fields().clone(), false);
35        let session = self.session().clone();
36
37        let iter = self
38            .map(move |chunk| {
39                let mut ctx = session.create_execution_ctx();
40                to_record_batch(chunk, &struct_field, &mut ctx)
41            })
42            .into_iter(runtime)?
43            .map(|result| result.map_err(|e| ArrowError::ExternalError(Box::new(e))));
44
45        Ok(RecordBatchIteratorAdapter { iter, schema })
46    }
47
48    pub fn into_record_batch_stream(
49        self,
50        schema: SchemaRef,
51    ) -> VortexResult<impl Stream<Item = Result<RecordBatch, ArrowError>> + Send + 'static> {
52        let struct_field = Field::new_struct("", schema.fields().clone(), false);
53        let session = self.session().clone();
54
55        let stream = self
56            .map(move |chunk| {
57                let mut ctx = session.create_execution_ctx();
58                to_record_batch(chunk, &struct_field, &mut ctx)
59            })
60            .into_stream()?
61            .map_err(|e| ArrowError::ExternalError(Box::new(e)));
62
63        Ok(stream)
64    }
65}
66
67fn to_record_batch(
68    chunk: ArrayRef,
69    data_type: &Field,
70    ctx: &mut ExecutionCtx,
71) -> VortexResult<RecordBatch> {
72    let session = ctx.session().clone();
73    let arrow = session.arrow().execute_arrow(chunk, Some(data_type), ctx)?;
74    Ok(RecordBatch::from(arrow.as_struct().clone()))
75}
76
77/// We create an adapter for record batch iterators that supports clone.
78/// This allows us to create thread-safe [`arrow_array::RecordBatchIterator`].
79#[derive(Clone)]
80pub struct RecordBatchIteratorAdapter<I> {
81    iter: I,
82    schema: SchemaRef,
83}
84
85impl<I> RecordBatchIteratorAdapter<I> {
86    /// Creates a new `RecordBatchIteratorAdapter`.
87    pub fn new(iter: I, schema: SchemaRef) -> Self {
88        Self { iter, schema }
89    }
90}
91
92impl<I> Iterator for RecordBatchIteratorAdapter<I>
93where
94    I: Iterator<Item = Result<RecordBatch, ArrowError>>,
95{
96    type Item = Result<RecordBatch, ArrowError>;
97
98    #[inline]
99    fn next(&mut self) -> Option<Self::Item> {
100        self.iter.next()
101    }
102}
103
104impl<I> RecordBatchReader for RecordBatchIteratorAdapter<I>
105where
106    I: Iterator<Item = Result<RecordBatch, ArrowError>>,
107{
108    #[inline]
109    fn schema(&self) -> SchemaRef {
110        Arc::clone(&self.schema)
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use std::sync::Arc;
117
118    use arrow_array::Array;
119    use arrow_array::ArrayRef as ArrowArrayRef;
120    use arrow_array::Int32Array;
121    use arrow_array::RecordBatch;
122    use arrow_array::StringArray;
123    use arrow_array::StructArray;
124    use arrow_array::cast::AsArray;
125    use arrow_schema::ArrowError;
126    use arrow_schema::DataType;
127    use arrow_schema::Field;
128    use arrow_schema::Schema;
129    use vortex_array::ArrayRef;
130    use vortex_array::VortexSessionExecute;
131    use vortex_array::arrow::FromArrowArray;
132    use vortex_error::VortexResult;
133
134    use super::*;
135    use crate::scan::test::SCAN_SESSION;
136
137    fn create_test_struct_array() -> VortexResult<ArrayRef> {
138        // Create Arrow arrays
139        let id_array = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
140        let name_array = StringArray::from(vec![Some("Alice"), Some("Bob"), Some("Charlie"), None]);
141
142        // Create Arrow struct array
143        let schema = Arc::new(Schema::new(vec![
144            Field::new("id", DataType::Int32, true),
145            Field::new("name", DataType::Utf8, true),
146        ]));
147
148        let struct_array = StructArray::new(
149            schema.fields().clone(),
150            vec![
151                Arc::new(id_array) as ArrowArrayRef,
152                Arc::new(name_array) as ArrowArrayRef,
153            ],
154            None,
155        );
156
157        // Convert to Vortex
158        ArrayRef::from_arrow(&struct_array, true)
159    }
160
161    fn create_arrow_schema() -> Arc<Schema> {
162        Arc::new(Schema::new(vec![
163            Field::new("id", DataType::Int32, true),
164            Field::new("name", DataType::Utf8, true),
165        ]))
166    }
167
168    #[test]
169    fn test_record_batch_conversion() -> VortexResult<()> {
170        let vortex_array = create_test_struct_array()?;
171        let schema = create_arrow_schema();
172        let struct_field = Field::new_struct("", schema.fields().clone(), false);
173        let mut ctx = SCAN_SESSION.create_execution_ctx();
174
175        let batch = to_record_batch(vortex_array, &struct_field, &mut ctx)?;
176        assert_eq!(batch.num_columns(), 2);
177        assert_eq!(batch.num_rows(), 4);
178
179        // Check id column
180        let id_col = batch
181            .column(0)
182            .as_primitive::<arrow_array::types::Int32Type>();
183        assert_eq!(id_col.value(0), 1);
184        assert_eq!(id_col.value(1), 2);
185        assert!(id_col.is_null(2));
186        assert_eq!(id_col.value(3), 4);
187
188        // Check name column
189        let name_col = batch.column(1).as_string::<i32>();
190        assert_eq!(name_col.value(0), "Alice");
191        assert_eq!(name_col.value(1), "Bob");
192        assert_eq!(name_col.value(2), "Charlie");
193        assert!(name_col.is_null(3));
194
195        Ok(())
196    }
197
198    #[test]
199    fn test_record_batch_iterator_adapter() -> VortexResult<()> {
200        let schema = create_arrow_schema();
201        let batch1 = RecordBatch::try_new(
202            Arc::clone(&schema),
203            vec![
204                Arc::new(Int32Array::from(vec![Some(1), Some(2)])) as ArrowArrayRef,
205                Arc::new(StringArray::from(vec![Some("Alice"), Some("Bob")])) as ArrowArrayRef,
206            ],
207        )?;
208        let batch2 = RecordBatch::try_new(
209            Arc::clone(&schema),
210            vec![
211                Arc::new(Int32Array::from(vec![None, Some(4)])) as ArrowArrayRef,
212                Arc::new(StringArray::from(vec![Some("Charlie"), None])) as ArrowArrayRef,
213            ],
214        )?;
215
216        let iter = vec![Ok(batch1), Ok(batch2)].into_iter();
217        let mut adapter = RecordBatchIteratorAdapter {
218            iter,
219            schema: Arc::clone(&schema),
220        };
221
222        // Test RecordBatchReader trait
223        assert_eq!(adapter.schema(), schema);
224
225        // Test Iterator trait
226        let first = adapter.next().unwrap()?;
227        assert_eq!(first.num_rows(), 2);
228
229        let second = adapter.next().unwrap()?;
230        assert_eq!(second.num_rows(), 2);
231
232        assert!(adapter.next().is_none());
233
234        Ok(())
235    }
236
237    #[test]
238    fn test_error_in_iterator() {
239        let schema = create_arrow_schema();
240        let error = ArrowError::ComputeError("test error".to_string());
241
242        let iter = vec![Err(error)].into_iter();
243        let mut adapter = RecordBatchIteratorAdapter {
244            iter,
245            schema: Arc::clone(&schema),
246        };
247
248        // Test that error is propagated
249        assert_eq!(adapter.schema(), schema);
250        let result = adapter.next().unwrap();
251        assert!(result.is_err());
252    }
253
254    #[test]
255    fn test_mixed_success_and_error() {
256        let schema = create_arrow_schema();
257        let batch = RecordBatch::try_new(
258            Arc::clone(&schema),
259            vec![
260                Arc::new(Int32Array::from(vec![Some(1)])) as ArrowArrayRef,
261                Arc::new(StringArray::from(vec![Some("Test")])) as ArrowArrayRef,
262            ],
263        )
264        .unwrap();
265
266        let error = ArrowError::ComputeError("test error".to_string());
267
268        let iter = vec![Ok(batch.clone()), Err(error), Ok(batch)].into_iter();
269        let mut adapter = RecordBatchIteratorAdapter { iter, schema };
270
271        // First batch succeeds
272        let first = adapter.next().unwrap();
273        assert!(first.is_ok());
274
275        // Second batch errors
276        let second = adapter.next().unwrap();
277        assert!(second.is_err());
278
279        // Third batch succeeds
280        let third = adapter.next().unwrap();
281        assert!(third.is_ok());
282    }
283}
vortex_layout/scan/arrow.rs

vortex_layout/scan/
arrow.rs
