This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11161-ace8a3e81b5e8132e3c25a97b94af3f0d751fb0c in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit df059af139f1dff142599a3b31a3b5fd0da472d5 Author: Kyle Barron <[email protected]> AuthorDate: Thu Oct 15 21:12:46 2026 +0000 feat: Introduce `ArrayReader` and `ArrayIterator` concepts (#11161) # Which issue does this PR close? - Closes https://github.com/apache/arrow-rs/issues/6586. # Rationale for this change Arrow-rs doesn't currently make it possible to describe sequences of _arrays_. It only really defines sequences of _Record batches_. This makes it hard to interface with libraries that use "chunked arrays" concepts or similar. This PR introduces code that I've used in [`pyo3-arrow`](https://crates.io/crates/pyo3-arrow) for a while. (It was necessary to support interop with a pyarrow `ChunkedArray`). # What changes are included in this PR? - Creates `ArrayReader` and `ArrayIterator` concepts: parallels to `RecordBatchReader` and `RecordBatchIterator` that yields `ArrayRef` instead of `RecordBatch`. - Defines `ArrayStreamReader` to describe an FFI stream of arrays, not of record batches > [!IMPORTANT] > Naming is hard here because the existing `ArrowArrayStreamReader` is a bit ambiguously named. It's the core struct for reading from the `ArrowArrayStream` FFI container. But it reads (into) _`RecordBatch`es_, not arrays. > > So what do we name the struct that yields arrays? One struct can't implement both `ArrayReader` and `RecordBatchReader` because each trait requires implementation of `Iterator` with a different yield type. > > I'm not sure what the best path forward is here. ## Todo: - [ ] We don't yet support _exporting_ an `ArrayReader` through FFI. I wanted to wait and get feedback on naming and general approach before implementing that. # Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? If this PR claims a performance improvement, please include evidence such as benchmark results. --> # Are there any user-facing changes? Yes, new APIs. <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. If there are any breaking changes to public APIs, please call them out. --> --------- Co-authored-by: Andrew Lamb <[email protected]> --- arrow-array/src/array_reader.rs | 163 +++++++++++++++++++ arrow-array/src/ffi_stream.rs | 226 ++++++++++++++++++++++++--- arrow-array/src/lib.rs | 3 + arrow-pyarrow-integration-testing/src/lib.rs | 10 +- arrow-pyarrow/src/lib.rs | 34 ++-- 5 files changed, 389 insertions(+), 47 deletions(-) diff --git a/arrow-array/src/array_reader.rs b/arrow-array/src/array_reader.rs new file mode 100644 index 0000000000..c423c3b957 --- /dev/null +++ b/arrow-array/src/array_reader.rs @@ -0,0 +1,163 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Traits for reading a stream of [`ArrayRef`] that share a single [`FieldRef`]. + +use arrow_schema::{ArrowError, FieldRef}; + +use crate::array::ArrayRef; + +/// Trait for types that can read `ArrayRef`'s. +/// +/// This is the array-level counterpart to [`RecordBatchReader`], for streams whose +/// elements are not record batches. +/// +/// To create from an iterator, see [ArrayIterator]. To read one from the C Stream Interface, +/// see [`ffi_stream::ArrayStreamReader`]. +/// +/// [`RecordBatchReader`]: crate::RecordBatchReader +/// [`ffi_stream::ArrayStreamReader`]: crate::ffi_stream::ArrayStreamReader +pub trait ArrayReader: Iterator<Item = Result<ArrayRef, ArrowError>> { + /// Returns the field of this `ArrayReader`. + /// + /// Implementation of this trait should guarantee that all `ArrayRef`'s returned by this + /// reader have the same field as returned from this method. + fn field(&self) -> &FieldRef; +} + +impl<R: ArrayReader + ?Sized> ArrayReader for Box<R> { + fn field(&self) -> &FieldRef { + self.as_ref().field() + } +} + +/// Generic implementation of [ArrayReader] that wraps an iterator. +/// +/// # Example +/// +/// ``` +/// # use std::sync::Arc; +/// # use arrow_array::{ArrayIterator, ArrayReader, ArrayRef, Int32Array}; +/// # use arrow_schema::{DataType, Field}; +/// # +/// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); +/// let b: ArrayRef = Arc::new(Int32Array::from(vec![3, 4])); +/// +/// let field = Arc::new(Field::new("values", DataType::Int32, false)); +/// let iter = vec![a.clone(), b.clone()].into_iter().map(Ok); +/// let mut reader = ArrayIterator::new(iter, field.clone()); +/// +/// assert_eq!(reader.field(), &field); +/// assert_eq!(&reader.next().unwrap().unwrap(), &a); +/// assert_eq!(&reader.next().unwrap().unwrap(), &b); +/// assert!(reader.next().is_none()); +/// ``` +pub struct ArrayIterator<I> +where + I: IntoIterator<Item = Result<ArrayRef, ArrowError>>, +{ + inner: I::IntoIter, + inner_field: FieldRef, +} + +impl<I> ArrayIterator<I> +where + I: IntoIterator<Item = Result<ArrayRef, ArrowError>>, +{ + /// Create a new [ArrayIterator]. + /// + /// If `iter` is an infallible iterator, use `.map(Ok)`. + pub fn new(iter: I, field: FieldRef) -> Self { + Self { + inner: iter.into_iter(), + inner_field: field, + } + } +} + +impl<I> Iterator for ArrayIterator<I> +where + I: IntoIterator<Item = Result<ArrayRef, ArrowError>>, +{ + type Item = I::Item; + + fn next(&mut self) -> Option<Self::Item> { + self.inner.next() + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.inner.size_hint() + } +} + +impl<I> ArrayReader for ArrayIterator<I> +where + I: IntoIterator<Item = Result<ArrayRef, ArrowError>>, +{ + fn field(&self) -> &FieldRef { + &self.inner_field + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_schema::{ArrowError, DataType, Field}; + + use crate::array::{ArrayRef, Int32Array}; + use crate::{ArrayIterator, ArrayReader}; + + fn test_field() -> Arc<Field> { + Arc::new(Field::new("values", DataType::Int32, true)) + } + + #[test] + fn array_iterator_reports_its_field_and_yields_its_arrays() { + let field = test_field(); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let b: ArrayRef = Arc::new(Int32Array::from(vec![3])); + + let mut reader = ArrayIterator::new( + vec![a.clone(), b.clone()].into_iter().map(Ok), + field.clone(), + ); + + assert_eq!(reader.field(), &field); + assert_eq!(&reader.next().unwrap().unwrap(), &a); + assert_eq!(&reader.next().unwrap().unwrap(), &b); + assert!(reader.next().is_none()); + } + + #[test] + fn array_iterator_delegates_size_hint() { + let arrays: Vec<Result<ArrayRef, ArrowError>> = + vec![Ok(Arc::new(Int32Array::from(vec![1])) as ArrayRef)]; + let reader = ArrayIterator::new(arrays, test_field()); + + assert_eq!(reader.size_hint(), (1, Some(1))); + } + + #[test] + fn boxed_array_reader_forwards_field() { + let field = test_field(); + let reader = ArrayIterator::new(std::iter::empty(), field.clone()); + + let boxed: Box<dyn ArrayReader + Send> = Box::new(reader); + assert_eq!(boxed.field(), &field); + } +} diff --git a/arrow-array/src/ffi_stream.rs b/arrow-array/src/ffi_stream.rs index 3e06e489ce..7c9547910e 100644 --- a/arrow-array/src/ffi_stream.rs +++ b/arrow-array/src/ffi_stream.rs @@ -21,14 +21,15 @@ //! One interface maps C ABI to native Rust types, i.e. convert c-pointers, c_char, to native rust. //! This is handled by [FFI_ArrowArrayStream]. //! -//! The second interface is used to import `FFI_ArrowArrayStream` as Rust implementation `RecordBatch` reader. -//! This is handled by `ArrowArrayStreamReader`. +//! The second interface is used to import `FFI_ArrowArrayStream` as a Rust reader. A stream +//! whose top-level type is `Struct` can be read as `RecordBatch` by [`RecordBatchStreamReader`]; +//! a stream of any type can be read as `ArrayRef` by [`ArrayStreamReader`]. //! //! ```ignore //! # use std::fs::File; //! # use std::sync::Arc; //! # use arrow::error::Result; -//! # use arrow::ffi_stream::{export_reader_into_raw, ArrowArrayStreamReader, FFI_ArrowArrayStream}; +//! # use arrow::ffi_stream::{export_reader_into_raw, RecordBatchStreamReader, FFI_ArrowArrayStream}; //! # use arrow::ipc::reader::FileReader; //! # use arrow::record_batch::RecordBatchReader; //! # fn main() -> Result<()> { @@ -43,7 +44,7 @@ //! // consumed and used by something else... //! //! // import it -//! let stream_reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() }; +//! let stream_reader = unsafe { RecordBatchStreamReader::from_raw(&mut stream).unwrap() }; //! let imported_schema = stream_reader.schema(); //! //! let mut produced_batches = vec![]; @@ -64,11 +65,13 @@ use std::{ }; use arrow_data::ffi::FFI_ArrowArray; -use arrow_schema::{ArrowError, Schema, SchemaRef, ffi::FFI_ArrowSchema}; +use arrow_schema::{ArrowError, Field, FieldRef, Schema, SchemaRef, ffi::FFI_ArrowSchema}; use crate::RecordBatchOptions; use crate::array::Array; use crate::array::StructArray; +use crate::array::{ArrayRef, make_array}; +use crate::array_reader::ArrayReader; use crate::ffi::from_ffi_and_data_type; use crate::record_batch::{RecordBatch, RecordBatchReader}; @@ -403,12 +406,18 @@ fn get_error_code(err: &ArrowError) -> i32 { /// Struct used to fetch `RecordBatch` from the C Stream Interface. /// Its main responsibility is to expose `RecordBatchReader` functionality /// that requires [FFI_ArrowArrayStream]. +/// +/// See [`ArrayStreamReader`] to read a stream whose arrays are not of type `Struct`. #[derive(Debug)] -pub struct ArrowArrayStreamReader { +pub struct RecordBatchStreamReader { stream: FFI_ArrowArrayStream, schema: SchemaRef, } +/// Renamed to [`RecordBatchStreamReader`], to distinguish it from [`ArrayStreamReader`]. +#[deprecated(since = "60.1.0", note = "Use `RecordBatchStreamReader` instead")] +pub type ArrowArrayStreamReader = RecordBatchStreamReader; + /// Returns the producer's message for the last failed call on a `FFI_ArrowArrayStream`. /// /// Returns `None` when the producer supplies no message, either because it installs no @@ -436,8 +445,22 @@ unsafe fn producer_error(stream_ptr: *mut FFI_ArrowArrayStream) -> Option<String } /// Gets schema from a raw pointer of `FFI_ArrowArrayStream`. This is used when constructing -/// `ArrowArrayStreamReader` to cache schema. +/// `RecordBatchStreamReader` to cache schema. fn get_stream_schema(stream_ptr: *mut FFI_ArrowArrayStream) -> Result<SchemaRef> { + let field = get_stream_field(stream_ptr)?; + match field.data_type() { + DataType::Struct(fields) => Ok(Arc::new( + Schema::new(fields.clone()).with_metadata(field.metadata().clone()), + )), + _ => Err(ArrowError::CDataInterface( + "Unable to interpret C data struct as a Schema".to_string(), + )), + } +} + +/// Gets the field from a raw pointer of `FFI_ArrowArrayStream`. This is used when constructing +/// `ArrayStreamReader` to cache the field. +fn get_stream_field(stream_ptr: *mut FFI_ArrowArrayStream) -> Result<FieldRef> { let mut schema = FFI_ArrowSchema::empty(); let Some(get_schema) = (unsafe { (*stream_ptr).get_schema }) else { @@ -449,8 +472,7 @@ fn get_stream_schema(stream_ptr: *mut FFI_ArrowArrayStream) -> Result<SchemaRef> let ret_code = unsafe { get_schema(stream_ptr, &raw mut schema) }; if ret_code == 0 { - let schema = Schema::try_from(&schema)?; - Ok(Arc::new(schema)) + Ok(Arc::new(Field::try_from(&schema)?)) } else { let message = format!("Cannot get schema from input stream. Error code: {ret_code}"); // SAFETY: `stream_ptr` is valid and unreleased, and the `get_schema` call above @@ -463,8 +485,8 @@ fn get_stream_schema(stream_ptr: *mut FFI_ArrowArrayStream) -> Result<SchemaRef> } } -impl ArrowArrayStreamReader { - /// Creates a new `ArrowArrayStreamReader` from a `FFI_ArrowArrayStream`. +impl RecordBatchStreamReader { + /// Creates a new `RecordBatchStreamReader` from a `FFI_ArrowArrayStream`. /// This is used to import from the C Stream Interface. pub fn try_new(mut stream: FFI_ArrowArrayStream) -> Result<Self> { if stream.release.is_none() { @@ -478,7 +500,7 @@ impl ArrowArrayStreamReader { Ok(Self { stream, schema }) } - /// Creates a new `ArrowArrayStreamReader` from a raw pointer of `FFI_ArrowArrayStream`. + /// Creates a new `RecordBatchStreamReader` from a raw pointer of `FFI_ArrowArrayStream`. /// /// Assumes that the pointer represents valid C Stream Interfaces. /// This function copies the content from the raw pointer and cleans up it to prevent @@ -493,7 +515,7 @@ impl ArrowArrayStreamReader { } } -impl Iterator for ArrowArrayStreamReader { +impl Iterator for RecordBatchStreamReader { type Item = Result<RecordBatch>; fn next(&mut self) -> Option<Self::Item> { @@ -538,12 +560,85 @@ impl Iterator for ArrowArrayStreamReader { } } -impl RecordBatchReader for ArrowArrayStreamReader { +impl RecordBatchReader for RecordBatchStreamReader { fn schema(&self) -> SchemaRef { self.schema.clone() } } +/// An [`ArrayReader`] which imports arrays from an [`FFI_ArrowArrayStream`]. +/// +/// [`RecordBatchStreamReader`] requires the stream's top-level type to be +/// [`DataType::Struct`], which it unpacks into a [`RecordBatch`]. This reader has no such +/// requirement: the stream's arrays may be of any data type. +#[derive(Debug)] +pub struct ArrayStreamReader { + stream: FFI_ArrowArrayStream, + field: FieldRef, +} + +impl ArrayStreamReader { + /// Creates a new [`ArrayStreamReader`] from an [`FFI_ArrowArrayStream`]. + /// This is used to import from the C Stream Interface. + pub fn try_new(mut stream: FFI_ArrowArrayStream) -> Result<Self> { + if stream.release.is_none() { + return Err(ArrowError::CDataInterface( + "input stream is already released".to_string(), + )); + } + + let field = get_stream_field(&raw mut stream)?; + + Ok(Self { stream, field }) + } + + /// Creates a new [`ArrayStreamReader`] from a raw pointer of [`FFI_ArrowArrayStream`]. + /// + /// # Safety + /// + /// See [`FFI_ArrowArrayStream::from_raw`] + pub unsafe fn from_raw(raw_stream: *mut FFI_ArrowArrayStream) -> Result<Self> { + Self::try_new(unsafe { FFI_ArrowArrayStream::from_raw(raw_stream) }) + } +} + +impl Iterator for ArrayStreamReader { + type Item = Result<ArrayRef>; + + fn next(&mut self) -> Option<Self::Item> { + let mut array = FFI_ArrowArray::empty(); + + let ret_code = + unsafe { self.stream.get_next.unwrap()(&raw mut self.stream, &raw mut array) }; + + if ret_code == 0 { + // The end of stream has been reached + if array.is_released() { + return None; + } + + let result = unsafe { from_ffi_and_data_type(array, self.field.data_type().clone()) }; + Some(result.map(make_array)) + } else { + let message = + format!("Cannot get next array from input stream. Error code: {ret_code}"); + // SAFETY: `self.stream` is valid and unreleased by construction, and the + // `get_next` call above returned a non-zero code. + let message = match unsafe { producer_error(&raw mut self.stream) } { + Some(producer_message) => format!("{message}. Producer error: {producer_message}"), + None => message, + }; + Some(Err(ArrowError::CDataInterface(message))) + } + } +} + +impl ArrayReader for ArrayStreamReader { + fn field(&self) -> &FieldRef { + &self.field + } +} + #[cfg(test)] mod tests { use super::*; @@ -632,9 +727,9 @@ mod tests { let reader = Box::new(TestRecordBatchReader::new(schema.clone(), iter)); - // Import through `FFI_ArrowArrayStream` as `ArrowArrayStreamReader` + // Import through `FFI_ArrowArrayStream` as `RecordBatchStreamReader` let stream = FFI_ArrowArrayStream::new(reader); - let stream_reader = ArrowArrayStreamReader::try_new(stream).unwrap(); + let stream_reader = RecordBatchStreamReader::try_new(stream).unwrap(); let imported_schema = stream_reader.schema(); assert_eq!(imported_schema, schema); @@ -695,9 +790,9 @@ mod tests { let reader = Box::new(TestRecordBatchReader::new(schema.clone(), iter)); - // Import through `FFI_ArrowArrayStream` as `ArrowArrayStreamReader` + // Import through `FFI_ArrowArrayStream` as `RecordBatchStreamReader` let stream = FFI_ArrowArrayStream::new(reader); - let stream_reader = ArrowArrayStreamReader::try_new(stream).unwrap(); + let stream_reader = RecordBatchStreamReader::try_new(stream).unwrap(); let imported_schema = stream_reader.schema(); assert_eq!(imported_schema, schema); @@ -779,7 +874,7 @@ mod tests { ) }; - let err = ArrowArrayStreamReader::try_new(stream).unwrap_err(); + let err = RecordBatchStreamReader::try_new(stream).unwrap_err(); assert!(matches!(err, ArrowError::CDataInterface(_))); assert!(err.to_string().contains("get_schema")); } @@ -797,7 +892,7 @@ mod tests { ) }; - let err = ArrowArrayStreamReader::try_new(stream) + let err = RecordBatchStreamReader::try_new(stream) .unwrap() .next() .unwrap() @@ -808,8 +903,8 @@ mod tests { #[test] fn test_import_schema_error_reports_producer_message() { - let err = - ArrowArrayStreamReader::try_new(failing_stream(Some(producer_last_error))).unwrap_err(); + let err = RecordBatchStreamReader::try_new(failing_stream(Some(producer_last_error))) + .unwrap_err(); assert_eq!( err.to_string(), format!( @@ -824,7 +919,7 @@ mod tests { // A producer need not supply a message: `get_last_error` may return NULL when no // detailed description is available. let err = - ArrowArrayStreamReader::try_new(failing_stream(Some(null_last_error))).unwrap_err(); + RecordBatchStreamReader::try_new(failing_stream(Some(null_last_error))).unwrap_err(); assert_eq!( err.to_string(), format!( @@ -835,7 +930,7 @@ mod tests { #[test] fn test_import_schema_error_without_error_callback() { - let err = ArrowArrayStreamReader::try_new(failing_stream(None)).unwrap_err(); + let err = RecordBatchStreamReader::try_new(failing_stream(None)).unwrap_err(); assert_eq!( err.to_string(), format!( @@ -851,7 +946,7 @@ mod tests { let mut stream = failing_stream(Some(null_last_error)); stream.get_schema = Some(working_get_schema); - let err = ArrowArrayStreamReader::try_new(stream) + let err = RecordBatchStreamReader::try_new(stream) .unwrap() .next() .unwrap() @@ -1027,4 +1122,85 @@ mod tests { assert!(stream.release().is_none()); assert!(stream.private_data().is_null()); } + #[test] + fn test_import_array_stream() { + let field = Field::new("a", DataType::Int32, true); + let arrays = vec![ + Int32Array::from(vec![1, 2, 3]), + Int32Array::from(vec![4, 5]), + ]; + let reader = ArrayStreamReader::try_new(array_stream(field, arrays.clone())).unwrap(); + + let produced = reader + .map(|array| Int32Array::from(array.unwrap().to_data())) + .collect::<Vec<_>>(); + + assert_eq!(produced, arrays); + } + + #[test] + fn test_import_array_stream_reports_its_field() { + let field = Field::new("a", DataType::Int32, true) + .with_metadata(HashMap::from([("key".to_string(), "value".to_string())])); + let reader = ArrayStreamReader::try_new(array_stream(field.clone(), vec![])).unwrap(); + + assert_eq!(reader.field().as_ref(), &field); + } + + #[test] + fn test_import_array_stream_rejects_released_stream() { + let err = ArrayStreamReader::try_new(FFI_ArrowArrayStream::empty()).unwrap_err(); + + assert_eq!( + err.to_string(), + "C Data interface error: input stream is already released" + ); + } + + #[test] + fn test_import_array_schema_error_reports_producer_message() { + let err = + ArrayStreamReader::try_new(failing_stream(Some(producer_last_error))).unwrap_err(); + + assert_eq!( + err.to_string(), + format!( + "C Data interface error: Cannot get schema from input stream. \ + Error code: {EIO}. Producer error: the producer failed" + ) + ); + } + + #[test] + fn test_import_array_next_error_reports_producer_message() { + let mut stream = failing_stream(Some(producer_last_error)); + stream.get_schema = Some(working_get_schema); + + let err = ArrayStreamReader::try_new(stream) + .unwrap() + .next() + .unwrap() + .unwrap_err(); + + assert_eq!( + err.to_string(), + format!( + "C Data interface error: Cannot get next array from input stream. \ + Error code: {EIO}. Producer error: the producer failed" + ) + ); + } + #[test] + fn test_import_record_batch_stream_rejects_non_struct_stream() { + // A stream of plain arrays cannot be read as record batches. Use + // `ArrayStreamReader` for that. + let field = Field::new("a", DataType::Int32, true); + + let err = RecordBatchStreamReader::try_new(array_stream(field, vec![])).unwrap_err(); + + assert_eq!( + err.to_string(), + "C Data interface error: Unable to interpret C data struct as a Schema" + ); + } } diff --git a/arrow-array/src/lib.rs b/arrow-array/src/lib.rs index 5114066c12..bee2406a55 100644 --- a/arrow-array/src/lib.rs +++ b/arrow-array/src/lib.rs @@ -242,6 +242,9 @@ pub const ARROW_VERSION: &str = env!("CARGO_PKG_VERSION"); pub mod array; pub use array::*; +mod array_reader; +pub use array_reader::{ArrayIterator, ArrayReader}; + mod record_batch; pub use record_batch::{ RecordBatch, RecordBatchIterator, RecordBatchOptions, RecordBatchReader, RecordBatchWriter, diff --git a/arrow-pyarrow-integration-testing/src/lib.rs b/arrow-pyarrow-integration-testing/src/lib.rs index a5690b3070..0e91c2d4aa 100644 --- a/arrow-pyarrow-integration-testing/src/lib.rs +++ b/arrow-pyarrow-integration-testing/src/lib.rs @@ -31,7 +31,7 @@ use arrow::array::{Array, ArrayData, ArrayRef, Int64Array, make_array}; use arrow::compute::kernels; use arrow::datatypes::{DataType, Field, Schema}; use arrow::error::ArrowError; -use arrow::ffi_stream::ArrowArrayStreamReader; +use arrow::ffi_stream::RecordBatchStreamReader; use arrow::pyarrow::{FromPyArrow, PyArrowException, PyArrowType, Table, ToPyArrow}; use arrow::record_batch::RecordBatch; @@ -135,8 +135,8 @@ fn round_trip_record_batch(obj: PyArrowType<RecordBatch>) -> PyResult<PyArrowTyp #[pyfunction] fn round_trip_record_batch_reader( - obj: PyArrowType<ArrowArrayStreamReader>, -) -> PyResult<PyArrowType<ArrowArrayStreamReader>> { + obj: PyArrowType<RecordBatchStreamReader>, +) -> PyResult<PyArrowType<RecordBatchStreamReader>> { Ok(obj) } @@ -161,7 +161,7 @@ pub fn build_table( } #[pyfunction] -fn reader_return_errors(obj: PyArrowType<ArrowArrayStreamReader>) -> PyResult<()> { +fn reader_return_errors(obj: PyArrowType<RecordBatchStreamReader>) -> PyResult<()> { // This makes sure we can correctly consume a RBR and return the error, // ensuring the error can live beyond the lifetime of the RBR. let batches = obj.0.collect::<Result<Vec<RecordBatch>, ArrowError>>(); @@ -173,7 +173,7 @@ fn reader_return_errors(obj: PyArrowType<ArrowArrayStreamReader>) -> PyResult<() #[pyfunction] fn boxed_reader_roundtrip( - obj: PyArrowType<ArrowArrayStreamReader>, + obj: PyArrowType<RecordBatchStreamReader>, ) -> PyArrowType<Box<dyn RecordBatchReader + Send>> { let schema = obj.0.schema(); let batches = obj diff --git a/arrow-pyarrow/src/lib.rs b/arrow-pyarrow/src/lib.rs index a159c9f1c4..a228ab7f3b 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/arrow-pyarrow/src/lib.rs @@ -43,11 +43,11 @@ //! | `pyarrow.Schema` | [Schema] | //! | `pyarrow.Array` | [ArrayData] | //! | `pyarrow.RecordBatch` | [RecordBatch] | -//! | `pyarrow.RecordBatchReader` | [ArrowArrayStreamReader] / `Box<dyn RecordBatchReader + Send>` (1) | +//! | `pyarrow.RecordBatchReader` | [RecordBatchStreamReader] / `Box<dyn RecordBatchReader + Send>` (1) | //! | `pyarrow.Table` | [Table] (2) | //! -//! (1) `pyarrow.RecordBatchReader` can be imported as [ArrowArrayStreamReader]. Either -//! [ArrowArrayStreamReader] or `Box<dyn RecordBatchReader + Send>` can be exported +//! (1) `pyarrow.RecordBatchReader` can be imported as [RecordBatchStreamReader]. Either +//! [RecordBatchStreamReader] or `Box<dyn RecordBatchReader + Send>` can be exported //! as `pyarrow.RecordBatchReader`. (`Box<dyn RecordBatchReader + Send>` is typically //! easier to create.) //! @@ -56,7 +56,7 @@ //! have `Vec<RecordBatch>` on the Rust side and want to export that in bulk as a `pyarrow.Table`. //! In general, it is recommended to use streaming approaches instead of dealing with data in bulk. //! For example, a `pyarrow.Table` (or any other object that implements the ArrayStream PyCapsule -//! interface) can be imported to Rust through `PyArrowType<ArrowArrayStreamReader>` instead of +//! interface) can be imported to Rust through `PyArrowType<RecordBatchStreamReader>` instead of //! forcing eager reading into `Vec<RecordBatch>`. //! //! # Type stubs @@ -83,7 +83,7 @@ use std::sync::Arc; use arrow_array::ffi; use arrow_array::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; -use arrow_array::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream}; +use arrow_array::ffi_stream::{FFI_ArrowArrayStream, RecordBatchStreamReader}; use arrow_array::{ RecordBatch, RecordBatchIterator, RecordBatchOptions, RecordBatchReader, StructArray, make_array, @@ -120,12 +120,12 @@ fn to_py_err(err: ArrowError) -> PyErr { } /// The type hint shared by every conversion that imports through the ArrowArrayStream PyCapsule -/// interface, i.e. [`ArrowArrayStreamReader`] and [`Table`]. +/// interface, i.e. [`RecordBatchStreamReader`] and [`Table`]. /// /// Both go through the same `__arrow_c_stream__` path and therefore accept exactly the same /// objects, so naming only one of the two classes would make a stub generator reject usage this /// crate's own documentation recommends — importing a `pyarrow.Table` as a -/// `PyArrowType<ArrowArrayStreamReader>`. +/// `PyArrowType<RecordBatchStreamReader>`. #[cfg(feature = "experimental-inspect")] const ARRAY_STREAM_INPUT_TYPE: PyStaticExpr = type_hint_union!( type_hint_identifier!("pyarrow", "RecordBatchReader"), @@ -513,8 +513,8 @@ impl ToPyArrow for RecordBatch { } } -/// Supports conversion from `pyarrow.RecordBatchReader` to [ArrowArrayStreamReader]. -impl FromPyArrow for ArrowArrayStreamReader { +/// Supports conversion from `pyarrow.RecordBatchReader` to [RecordBatchStreamReader]. +impl FromPyArrow for RecordBatchStreamReader { type_hint!(INPUT_TYPE = ARRAY_STREAM_INPUT_TYPE); fn from_pyarrow_bound(value: &Bound<PyAny>) -> PyResult<Self> { @@ -528,7 +528,7 @@ impl FromPyArrow for ArrowArrayStreamReader { extract_capsule(&capsule, c"arrow_array_stream", "__arrow_c_stream__")?; let stream = unsafe { FFI_ArrowArrayStream::from_raw(stream_ptr.as_ptr()) }; - let stream_reader = ArrowArrayStreamReader::try_new(stream) + let stream_reader = RecordBatchStreamReader::try_new(stream) .map_err(|err| PyValueError::new_err(err.to_string()))?; return Ok(stream_reader); @@ -547,7 +547,7 @@ impl FromPyArrow for ArrowArrayStreamReader { (&raw mut stream as Py_uintptr_t,), )?; - ArrowArrayStreamReader::try_new(stream) + RecordBatchStreamReader::try_new(stream) .map_err(|err| PyValueError::new_err(err.to_string())) } } @@ -567,8 +567,8 @@ impl IntoPyArrow for Box<dyn RecordBatchReader + Send> { } } -/// Convert a [`ArrowArrayStreamReader`] into a `pyarrow.RecordBatchReader`. -impl IntoPyArrow for ArrowArrayStreamReader { +/// Convert a [`RecordBatchStreamReader`] into a `pyarrow.RecordBatchReader`. +impl IntoPyArrow for RecordBatchStreamReader { type_hint!(OUTPUT_TYPE = type_hint_identifier!("pyarrow", "RecordBatchReader")); fn into_pyarrow(self, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> { @@ -582,7 +582,7 @@ impl IntoPyArrow for ArrowArrayStreamReader { /// /// This could be used in circumstances where you either want to consume a `pyarrow.Table` directly /// (although technically, since `pyarrow.Table` implements the ArrayStreamReader PyCapsule -/// interface, one could also consume a `PyArrowType<ArrowArrayStreamReader>` instead) or, more +/// interface, one could also consume a `PyArrowType<RecordBatchStreamReader>` instead) or, more /// importantly, where one wants to export a `pyarrow.Table` from a `Vec<RecordBatch>` from the Rust /// side. /// @@ -650,7 +650,7 @@ impl FromPyArrow for Table { fn from_pyarrow_bound(ob: &Bound<PyAny>) -> PyResult<Self> { let reader: Box<dyn RecordBatchReader> = - Box::new(ArrowArrayStreamReader::from_pyarrow_bound(ob)?); + Box::new(RecordBatchStreamReader::from_pyarrow_bound(ob)?); Self::try_from(reader).map_err(|err| PyValueError::new_err(err.to_string())) } } @@ -858,11 +858,11 @@ mod introspection_tests { #[test] fn readers_and_tables_map_to_their_pyarrow_class() { assert_eq!( - input_type::<ArrowArrayStreamReader>(), + input_type::<RecordBatchStreamReader>(), "pyarrow.RecordBatchReader | pyarrow.Table" ); assert_eq!( - output_type::<ArrowArrayStreamReader>(), + output_type::<RecordBatchStreamReader>(), "pyarrow.RecordBatchReader" ); assert_eq!(
