16292500848

Committed 15 Jul 2025 11:52AM UTC coverage: 81.486% (-0.05%) from 81.533%

Build # 16292500848

Build Type

push

github

Committed by

web-flow

Commit Message

Arrow iterator conversions (#3875)

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Run Details

0 of 13 new or added lines in 1 file covered. (0.0%)

46270 of 56783 relevant lines covered (81.49%)

146722.21 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/vortex-array/src/arrow/iter.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::sync::Arc;

use arrow_array::cast::AsArray;
use arrow_array::{RecordBatch, RecordBatchReader, ffi_stream};
use arrow_schema::{ArrowError, DataType, SchemaRef};
use vortex_dtype::DType;
use vortex_dtype::arrow::FromArrowType;
use vortex_error::{VortexError, VortexResult};

use crate::ArrayRef;
use crate::arrow::FromArrowArray;
use crate::arrow::compute::to_arrow;
use crate::iter::ArrayIterator;

/// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayStream`.
pub struct ArrowArrayStreamAdapter {
    stream: ffi_stream::ArrowArrayStreamReader,
    dtype: DType,
}

impl ArrowArrayStreamAdapter {
    pub fn new(stream: ffi_stream::ArrowArrayStreamReader, dtype: DType) -> Self {
        Self { stream, dtype }
    }
}

impl ArrayIterator for ArrowArrayStreamAdapter {
    fn dtype(&self) -> &DType {
        &self.dtype
    }
}

impl Iterator for ArrowArrayStreamAdapter {
    type Item = VortexResult<ArrayRef>;

    fn next(&mut self) -> Option<Self::Item> {
        let batch = self.stream.next()?;

        Some(batch.map_err(VortexError::from).map(|b| {
            debug_assert_eq!(&self.dtype, &DType::from_arrow(b.schema()));
            ArrayRef::from_arrow(b, false)
        }))
    }
}

/// Adapter for converting a [`ArrayIterator`] into an Arrow [`RecordBatchReader`].
pub struct VortexRecordBatchReader<I> {
    iter: I,
    arrow_schema: SchemaRef,
    arrow_dtype: DataType,
}

impl<I: ArrayIterator> VortexRecordBatchReader<I> {
    pub fn try_new(iter: I) -> VortexResult<Self> {
        let arrow_schema = Arc::new(iter.dtype().to_arrow_schema()?);
        let arrow_dtype = DataType::Struct(arrow_schema.fields().clone());
        Ok(VortexRecordBatchReader {
            iter,
            arrow_schema,
            arrow_dtype,
        })
    }
}

impl<I: ArrayIterator> Iterator for VortexRecordBatchReader<I> {
    type Item = Result<RecordBatch, ArrowError>;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(|result| {
            result
                .and_then(|array| to_arrow(&array, &self.arrow_dtype))
                .map_err(|e| ArrowError::ExternalError(Box::new(e)))
                .map(|array| RecordBatch::from(array.as_struct()))
        })
    }
}

impl<I: ArrayIterator> RecordBatchReader for VortexRecordBatchReader<I> {
    fn schema(&self) -> SchemaRef {
        self.arrow_schema.clone()
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::sync::Arc;
5
6	use arrow_array::cast::AsArray;
7	use arrow_array::{RecordBatch, RecordBatchReader, ffi_stream};
8	use arrow_schema::{ArrowError, DataType, SchemaRef};
9	use vortex_dtype::DType;
10	use vortex_dtype::arrow::FromArrowType;
11	use vortex_error::{VortexError, VortexResult};
12
13	use crate::ArrayRef;
14	use crate::arrow::FromArrowArray;
15	use crate::arrow::compute::to_arrow;
16	use crate::iter::ArrayIterator;
17
18	/// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayStream`.
19	pub struct ArrowArrayStreamAdapter {
20	stream: ffi_stream::ArrowArrayStreamReader,
21	dtype: DType,
22	}
23
24	impl ArrowArrayStreamAdapter {
NEW 25	pub fn new(stream: ffi_stream::ArrowArrayStreamReader, dtype: DType) -> Self {	×
NEW 26	Self { stream, dtype }	×
NEW 27	}	×
28	}
29
30	impl ArrayIterator for ArrowArrayStreamAdapter {
NEW 31	fn dtype(&self) -> &DType {	×
NEW 32	&self.dtype	×
NEW 33	}	×
34	}
35
36	impl Iterator for ArrowArrayStreamAdapter {
37	type Item = VortexResult<ArrayRef>;
38
NEW 39	fn next(&mut self) -> Option<Self::Item> {	×
NEW 40	let batch = self.stream.next()?;	×
41
NEW 42	Some(batch.map_err(VortexError::from).map(\|b\| {	×
NEW 43	debug_assert_eq!(&self.dtype, &DType::from_arrow(b.schema()));	×
NEW 44	ArrayRef::from_arrow(b, false)	×
NEW 45	}))	×
NEW 46	}	×
47	}
48
49	/// Adapter for converting a [`ArrayIterator`] into an Arrow [`RecordBatchReader`].
50	pub struct VortexRecordBatchReader<I> {
51	iter: I,
52	arrow_schema: SchemaRef,
53	arrow_dtype: DataType,
54	}
55
56	impl<I: ArrayIterator> VortexRecordBatchReader<I> {
57	pub fn try_new(iter: I) -> VortexResult<Self> {	×
58	let arrow_schema = Arc::new(iter.dtype().to_arrow_schema()?);	×
59	let arrow_dtype = DataType::Struct(arrow_schema.fields().clone());	×
60	Ok(VortexRecordBatchReader {	×
61	iter,	×
62	arrow_schema,	×
63	arrow_dtype,	×
64	})	×
65	}	×
66	}
67
68	impl<I: ArrayIterator> Iterator for VortexRecordBatchReader<I> {
69	type Item = Result<RecordBatch, ArrowError>;
70
71	fn next(&mut self) -> Option<Self::Item> {	×
72	self.iter.next().map(\|result\| {	×
73	result	×
74	.and_then(\|array\| to_arrow(&array, &self.arrow_dtype))	×
75	.map_err(\|e\| ArrowError::ExternalError(Box::new(e)))	×
76	.map(\|array\| RecordBatch::from(array.as_struct()))	×
77	})	×
78	}	×
79	}
80
81	impl<I: ArrayIterator> RecordBatchReader for VortexRecordBatchReader<I> {
82	fn schema(&self) -> SchemaRef {	×
83	self.arrow_schema.clone()	×
84	}	×
85	}

vortex-data / vortex / 16292500848

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous