arrow_array/

ffi_stream.rs

// 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.

//! Contains declarations to bind to the [C Stream Interface](https://arrow.apache.org/docs/format/CStreamInterface.html).
//!
//! This module has two main interfaces:
//! 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 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, RecordBatchStreamReader, FFI_ArrowArrayStream};
//! # use arrow::ipc::reader::FileReader;
//! # use arrow::record_batch::RecordBatchReader;
//! # fn main() -> Result<()> {
//! // create an record batch reader natively
//! let file = File::open("arrow_file").unwrap();
//! let reader = Box::new(FileReader::try_new(file).unwrap());
//!
//! // export it
//! let mut stream = FFI_ArrowArrayStream::empty();
//! unsafe { export_reader_into_raw(reader, &mut stream) };
//!
//! // consumed and used by something else...
//!
//! // import it
//! let stream_reader = unsafe { RecordBatchStreamReader::from_raw(&mut stream).unwrap() };
//! let imported_schema = stream_reader.schema();
//!
//! let mut produced_batches = vec![];
//! for batch in stream_reader {
//!      produced_batches.push(batch.unwrap());
//! }
//! Ok(())
//! }
//! ```

use arrow_schema::DataType;
use std::ffi::CStr;
use std::ptr::addr_of;
use std::{
    ffi::CString,
    os::raw::{c_char, c_int, c_void},
    sync::Arc,
};

use arrow_data::ffi::FFI_ArrowArray;
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};

type Result<T> = std::result::Result<T, ArrowError>;

// Errno values returned through the C stream interface, taken from libc so they match
// the platform the consumer interprets them against.
#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
use libc::{EINVAL, EIO, ENOMEM, ENOSYS};

// wasm32-unknown-unknown has no libc, and no OS to interpret the codes either — any
// non-zero value works there, so use Linux's.
#[cfg(all(target_family = "wasm", target_os = "unknown"))]
const ENOMEM: i32 = 12;
#[cfg(all(target_family = "wasm", target_os = "unknown"))]
const EIO: i32 = 5;
#[cfg(all(target_family = "wasm", target_os = "unknown"))]
const EINVAL: i32 = 22;
#[cfg(all(target_family = "wasm", target_os = "unknown"))]
const ENOSYS: i32 = 38;

/// ABI-compatible struct for `ArrayStream` from C Stream Interface
/// See <https://arrow.apache.org/docs/format/CStreamInterface.html#structure-definitions>
/// This was created by bindgen
#[repr(C)]
#[derive(Debug)]
pub struct FFI_ArrowArrayStream {
    // Fields are intentionally private so safety guarantees can be upheld via
    // explicit unsafe functions.
    /// C function to get schema from the stream
    get_schema: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowSchema) -> c_int>,
    /// C function to get next array from the stream
    get_next: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowArray) -> c_int>,
    /// C function to get the error from last operation on the stream
    get_last_error: Option<unsafe extern "C" fn(arg1: *mut Self) -> *const c_char>,
    /// C function to release the stream
    release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
    /// Private data used by the stream, owned by the release callback.
    private_data: *mut c_void,
}

unsafe impl Send for FFI_ArrowArrayStream {}

// callback used to drop [FFI_ArrowArrayStream] when it is exported.
unsafe extern "C" fn release_stream(stream: *mut FFI_ArrowArrayStream) {
    if stream.is_null() {
        return;
    }
    let stream = unsafe { &mut *stream };

    stream.get_schema = None;
    stream.get_next = None;
    stream.get_last_error = None;

    let private_data = unsafe { Box::from_raw(stream.private_data.cast::<StreamPrivateData>()) };
    drop(private_data);

    stream.release = None;
}

struct StreamPrivateData {
    batch_reader: Box<dyn RecordBatchReader + Send>,
    last_error: Option<CString>,
}

// The callback used to get array schema
unsafe extern "C" fn get_schema(
    stream: *mut FFI_ArrowArrayStream,
    schema: *mut FFI_ArrowSchema,
) -> c_int {
    ExportedArrayStream { stream }.get_schema(schema)
}

// The callback used to get next array
unsafe extern "C" fn get_next(
    stream: *mut FFI_ArrowArrayStream,
    array: *mut FFI_ArrowArray,
) -> c_int {
    ExportedArrayStream { stream }.get_next(array)
}

// The callback used to get the error from last operation on the `FFI_ArrowArrayStream`
unsafe extern "C" fn get_last_error(stream: *mut FFI_ArrowArrayStream) -> *const c_char {
    let mut ffi_stream = ExportedArrayStream { stream };
    // The consumer should not take ownership of this string, we should return
    // a const pointer to it.
    match ffi_stream.get_last_error() {
        Some(err_string) => err_string.as_ptr(),
        None => std::ptr::null(),
    }
}

impl Drop for FFI_ArrowArrayStream {
    fn drop(&mut self) {
        match self.release {
            None => (),
            Some(release) => unsafe { release(self) },
        }
    }
}

impl FFI_ArrowArrayStream {
    /// Creates a new [`FFI_ArrowArrayStream`].
    pub fn new(batch_reader: Box<dyn RecordBatchReader + Send>) -> Self {
        let private_data = Box::new(StreamPrivateData {
            batch_reader,
            last_error: None,
        });

        Self {
            get_schema: Some(get_schema),
            get_next: Some(get_next),
            get_last_error: Some(get_last_error),
            release: Some(release_stream),
            private_data: Box::into_raw(private_data).cast::<c_void>(),
        }
    }

    /// Creates a new [`FFI_ArrowArrayStream`] from its raw parts.
    ///
    /// # Safety
    ///
    /// The caller takes responsibility for the [C stream interface] contract:
    ///
    /// * Each callback that is `Some` must be sound to invoke with a pointer to
    ///   this struct, and must implement the semantics the C stream interface
    ///   specifies for it.
    /// * `private_data` must own everything the callbacks rely on, and must
    ///   remain valid until `release` is called.
    /// * `release`, if `Some`, is invoked by [`Drop`] with a pointer to this
    ///   struct. It must free `private_data` exactly once and then set the
    ///   release callback to `None` to mark the stream released.
    /// * If `release` is `None` the stream is already released, so dropping it
    ///   must not leak: nothing may be left for the callback to free.
    ///
    /// [C stream interface]: https://arrow.apache.org/docs/format/CStreamInterface.html
    pub unsafe fn new_unchecked(
        get_schema: Option<
            unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowSchema) -> c_int,
        >,
        get_next: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowArray) -> c_int>,
        get_last_error: Option<unsafe extern "C" fn(arg1: *mut Self) -> *const c_char>,
        release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
        private_data: *mut c_void,
    ) -> Self {
        Self {
            get_schema,
            get_next,
            get_last_error,
            release,
            private_data,
        }
    }

    /// Takes ownership of the pointed to [`FFI_ArrowArrayStream`]
    ///
    /// This acts to [move] the data out of `raw_stream`, setting the release callback to NULL
    ///
    /// # Safety
    ///
    /// * `raw_stream` must be [valid] for reads and writes
    /// * `raw_stream` must be properly aligned
    /// * `raw_stream` must point to a properly initialized value of [`FFI_ArrowArrayStream`]
    ///
    /// [move]: https://arrow.apache.org/docs/format/CDataInterface.html#moving-an-array
    /// [valid]: https://doc.rust-lang.org/std/ptr/index.html#safety
    pub unsafe fn from_raw(raw_stream: *mut FFI_ArrowArrayStream) -> Self {
        unsafe { std::ptr::replace(raw_stream, Self::empty()) }
    }

    /// Creates a new empty [FFI_ArrowArrayStream]. Used to import from the C Stream Interface.
    pub fn empty() -> Self {
        Self {
            get_schema: None,
            get_next: None,
            get_last_error: None,
            release: None,
            private_data: std::ptr::null_mut(),
        }
    }

    #[expect(clippy::unnecessary_safety_doc)]
    /// Returns the producer-provided callback that writes this stream's schema, if any.
    ///
    /// # Safety
    ///
    /// The callback must be invoked with a pointer to the stream it was read from.
    pub fn get_schema(
        &self,
    ) -> Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowSchema) -> c_int> {
        self.get_schema
    }

    #[expect(clippy::unnecessary_safety_doc)]
    /// Returns the producer-provided callback that yields the next array, if any.
    ///
    /// # Safety
    ///
    /// The callback must be invoked with a pointer to the stream it was read from.
    pub fn get_next(
        &self,
    ) -> Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowArray) -> c_int> {
        self.get_next
    }

    #[expect(clippy::unnecessary_safety_doc)]
    /// Returns the producer-provided callback that describes the last error, if any.
    ///
    /// # Safety
    ///
    /// The callback must be invoked with a pointer to the stream it was read from.
    /// The string it returns is owned by the producer and valid only until the next
    /// call on the stream.
    pub fn get_last_error(&self) -> Option<unsafe extern "C" fn(arg1: *mut Self) -> *const c_char> {
        self.get_last_error
    }

    /// Returns the producer-provided release callback, if any.
    pub fn release(&self) -> Option<unsafe extern "C" fn(arg1: *mut Self)> {
        self.release
    }

    /// Returns the opaque producer-provided private data pointer.
    pub fn private_data(&self) -> *mut c_void {
        self.private_data
    }

    /// Replaces the release callback, returning the previous one.
    ///
    /// Lets a consumer wrap release: save the old callback, install its own, and
    /// chain back on drop. See <https://github.com/apache/arrow-rs/issues/9771>.
    ///
    /// # Safety
    ///
    /// [`Drop`] calls this callback with a pointer to `self`. The new callback
    /// must correctly release this stream (usually by chaining to the returned
    /// one) and must match the [`FFI_ArrowArrayStream::private_data`] it reads.
    /// A wrong callback is undefined behavior on drop.
    pub unsafe fn set_release(
        &mut self,
        release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
    ) -> Option<unsafe extern "C" fn(arg1: *mut Self)> {
        std::mem::replace(&mut self.release, release)
    }

    /// Replaces the private data pointer, returning the previous one.
    ///
    /// # Safety
    ///
    /// The old pointer is returned without being freed; the caller owns it from
    /// here. The new pointer must match what the current
    /// [`FFI_ArrowArrayStream::release`] callback expects.
    pub unsafe fn set_private_data(&mut self, private_data: *mut c_void) -> *mut c_void {
        std::mem::replace(&mut self.private_data, private_data)
    }
}

struct ExportedArrayStream {
    stream: *mut FFI_ArrowArrayStream,
}

impl ExportedArrayStream {
    fn get_private_data(&mut self) -> &mut StreamPrivateData {
        unsafe { &mut *(*self.stream).private_data.cast::<StreamPrivateData>() }
    }

    pub fn get_schema(&mut self, out: *mut FFI_ArrowSchema) -> i32 {
        let private_data = self.get_private_data();
        let reader = &private_data.batch_reader;

        let schema = FFI_ArrowSchema::try_from(reader.schema().as_ref());

        match schema {
            Ok(schema) => {
                unsafe { std::ptr::copy(addr_of!(schema), out, 1) };
                std::mem::forget(schema);
                0
            }
            Err(ref err) => {
                private_data.last_error = Some(
                    CString::new(err.to_string()).expect("Error string has a null byte in it."),
                );
                get_error_code(err)
            }
        }
    }

    pub fn get_next(&mut self, out: *mut FFI_ArrowArray) -> i32 {
        let private_data = self.get_private_data();
        let reader = &mut private_data.batch_reader;

        match reader.next() {
            None => {
                // Marks ArrowArray released to indicate reaching the end of stream.
                unsafe { std::ptr::write(out, FFI_ArrowArray::empty()) }
                0
            }
            Some(next_batch) => {
                if let Ok(batch) = next_batch {
                    let struct_array = StructArray::from(batch);
                    let array = FFI_ArrowArray::new(&struct_array.to_data());

                    unsafe { std::ptr::write_unaligned(out, array) };
                    0
                } else {
                    let err = &next_batch.unwrap_err();
                    private_data.last_error = Some(
                        CString::new(err.to_string()).expect("Error string has a null byte in it."),
                    );
                    get_error_code(err)
                }
            }
        }
    }

    pub fn get_last_error(&mut self) -> Option<&CString> {
        self.get_private_data().last_error.as_ref()
    }
}

fn get_error_code(err: &ArrowError) -> i32 {
    match err {
        ArrowError::NotYetImplemented(_) => ENOSYS,
        ArrowError::MemoryError(_) => ENOMEM,
        ArrowError::IoError(_, _) => EIO,
        _ => EINVAL,
    }
}

/// A `RecordBatchReader` which imports Arrays from `FFI_ArrowArrayStream`.
///
/// 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 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
/// `get_last_error` callback or because that callback returns NULL: the C Stream Interface
/// lets `get_last_error` return NULL when no detailed description is available.
///
/// # Safety
///
/// `stream_ptr` must point to a valid, not yet released [`FFI_ArrowArrayStream`], and the last
/// operation on it must have returned an error: the C Stream Interface forbids calling
/// `get_last_error` in any other case.
unsafe fn producer_error(stream_ptr: *mut FFI_ArrowArrayStream) -> Option<String> {
    let get_last_error = unsafe { (*stream_ptr).get_last_error }?;

    let error_str = unsafe { get_last_error(stream_ptr) };
    if error_str.is_null() {
        return None;
    }

    Some(
        unsafe { CStr::from_ptr(error_str) }
            .to_string_lossy()
            .into_owned(),
    )
}

/// Gets schema from a raw pointer of `FFI_ArrowArrayStream`. This is used when constructing
/// `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 {
        return Err(ArrowError::CDataInterface(
            "input stream does not provide the required get_schema callback".to_string(),
        ));
    };

    let ret_code = unsafe { get_schema(stream_ptr, &raw mut schema) };

    if ret_code == 0 {
        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
        // returned a non-zero code.
        let message = match unsafe { producer_error(stream_ptr) } {
            Some(producer_message) => format!("{message}. Producer error: {producer_message}"),
            None => message,
        };
        Err(ArrowError::CDataInterface(message))
    }
}

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() {
            return Err(ArrowError::CDataInterface(
                "input stream is already released".to_string(),
            ));
        }

        let schema = get_stream_schema(&raw mut stream)?;

        Ok(Self { stream, schema })
    }

    /// 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
    /// double-dropping. The caller is responsible for freeing up the memory allocated for
    /// the pointer.
    ///
    /// # 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 RecordBatchStreamReader {
    type Item = Result<RecordBatch>;

    fn next(&mut self) -> Option<Self::Item> {
        let mut array = FFI_ArrowArray::empty();

        let Some(get_next) = self.stream.get_next else {
            return Some(Err(ArrowError::CDataInterface(
                "input stream does not provide the required get_next callback".to_string(),
            )));
        };

        let ret_code = unsafe { get_next(&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, DataType::Struct(self.schema().fields().clone()))
            };
            Some(result.and_then(|data| {
                let len = data.len();
                RecordBatch::try_new_with_options(
                    self.schema.clone(),
                    StructArray::from(data).into_parts().1,
                    &RecordBatchOptions::new().with_row_count(Some(len)),
                )
            }))
        } else {
            let message =
                format!("Cannot get next batch 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 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::*;
    use std::collections::HashMap;

    use arrow_schema::Field;

    use crate::array::Int32Array;
    use crate::ffi::from_ffi;

    struct TestRecordBatchReader {
        schema: SchemaRef,
        iter: Box<dyn Iterator<Item = Result<RecordBatch>> + Send>,
    }

    impl TestRecordBatchReader {
        pub fn new(
            schema: SchemaRef,
            iter: Box<dyn Iterator<Item = Result<RecordBatch>> + Send>,
        ) -> TestRecordBatchReader {
            TestRecordBatchReader { schema, iter }
        }
    }

    impl Iterator for TestRecordBatchReader {
        type Item = Result<RecordBatch>;

        fn next(&mut self) -> Option<Self::Item> {
            self.iter.next()
        }
    }

    impl RecordBatchReader for TestRecordBatchReader {
        fn schema(&self) -> SchemaRef {
            self.schema.clone()
        }
    }

    fn _test_round_trip_export(batch: RecordBatch, schema: Arc<Schema>) -> Result<()> {
        let iter = Box::new(vec![batch.clone(), batch.clone()].into_iter().map(Ok)) as _;

        let reader = Box::new(TestRecordBatchReader::new(schema.clone(), iter));

        // Export a `RecordBatchReader` through `FFI_ArrowArrayStream`
        let mut ffi_stream = FFI_ArrowArrayStream::new(reader);

        // Get schema from `FFI_ArrowArrayStream`
        let mut ffi_schema = FFI_ArrowSchema::empty();
        let ret_code = unsafe { get_schema(&raw mut ffi_stream, &raw mut ffi_schema) };
        assert_eq!(ret_code, 0);

        let exported_schema = Schema::try_from(&ffi_schema).unwrap();
        assert_eq!(&exported_schema, schema.as_ref());

        // Get array from `FFI_ArrowArrayStream`
        let mut produced_batches = vec![];
        loop {
            let mut ffi_array = FFI_ArrowArray::empty();
            let ret_code = unsafe { get_next(&raw mut ffi_stream, &raw mut ffi_array) };
            assert_eq!(ret_code, 0);

            // The end of stream has been reached
            if ffi_array.is_released() {
                break;
            }

            let array = unsafe { from_ffi(ffi_array, &ffi_schema) }.unwrap();
            let len = array.len();

            let record_batch = RecordBatch::try_new_with_options(
                SchemaRef::from(exported_schema.clone()),
                StructArray::from(array).into_parts().1,
                &RecordBatchOptions::new().with_row_count(Some(len)),
            )
            .unwrap();
            produced_batches.push(record_batch);
        }

        assert_eq!(produced_batches, vec![batch.clone(), batch]);

        Ok(())
    }

    fn _test_round_trip_import(batch: RecordBatch, schema: Arc<Schema>) -> Result<()> {
        let iter = Box::new(vec![batch.clone(), batch.clone()].into_iter().map(Ok)) as _;

        let reader = Box::new(TestRecordBatchReader::new(schema.clone(), iter));

        // Import through `FFI_ArrowArrayStream` as `RecordBatchStreamReader`
        let stream = FFI_ArrowArrayStream::new(reader);
        let stream_reader = RecordBatchStreamReader::try_new(stream).unwrap();

        let imported_schema = stream_reader.schema();
        assert_eq!(imported_schema, schema);

        let mut produced_batches = vec![];
        for batch in stream_reader {
            produced_batches.push(batch.unwrap());
        }

        assert_eq!(produced_batches, vec![batch.clone(), batch]);

        Ok(())
    }

    #[test]
    fn test_stream_round_trip() {
        let array = Int32Array::from(vec![Some(2), None, Some(1), None]);
        let array: Arc<dyn Array> = Arc::new(array);
        let metadata = HashMap::from([("foo".to_owned(), "bar".to_owned())]);

        let schema = Arc::new(Schema::new_with_metadata(
            vec![
                Field::new("a", array.data_type().clone(), true).with_metadata(metadata.clone()),
                Field::new("b", array.data_type().clone(), true).with_metadata(metadata.clone()),
                Field::new("c", array.data_type().clone(), true).with_metadata(metadata.clone()),
            ],
            metadata,
        ));
        let batch = RecordBatch::try_new(schema.clone(), vec![array.clone(), array.clone(), array])
            .unwrap();

        _test_round_trip_export(batch.clone(), schema.clone()).unwrap();
        _test_round_trip_import(batch, schema).unwrap();
    }

    #[test]
    fn test_stream_round_trip_no_columns() {
        let metadata = HashMap::from([("foo".to_owned(), "bar".to_owned())]);

        let schema = Arc::new(Schema::new_with_metadata(Vec::<Field>::new(), metadata));
        let batch = RecordBatch::try_new_with_options(
            schema.clone(),
            Vec::<Arc<dyn Array>>::new(),
            &RecordBatchOptions::new().with_row_count(Some(10)),
        )
        .unwrap();

        _test_round_trip_export(batch.clone(), schema.clone()).unwrap();
        _test_round_trip_import(batch, schema).unwrap();
    }

    #[test]
    fn test_error_import() -> Result<()> {
        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));

        let iter =
            Box::new(vec![Err(ArrowError::MemoryError("out of memory".to_string()))].into_iter());

        let reader = Box::new(TestRecordBatchReader::new(schema.clone(), iter));

        // Import through `FFI_ArrowArrayStream` as `RecordBatchStreamReader`
        let stream = FFI_ArrowArrayStream::new(reader);
        let stream_reader = RecordBatchStreamReader::try_new(stream).unwrap();

        let imported_schema = stream_reader.schema();
        assert_eq!(imported_schema, schema);

        let mut produced_batches = vec![];
        for batch in stream_reader {
            produced_batches.push(batch);
        }

        // The results should outlive the lifetime of the stream itself.
        assert_eq!(produced_batches.len(), 1);
        assert_eq!(
            produced_batches[0].as_ref().unwrap_err().to_string(),
            format!(
                "C Data interface error: Cannot get next batch from input stream. \
                 Error code: {ENOMEM}. Producer error: Memory error: out of memory"
            )
        );

        Ok(())
    }

    unsafe extern "C" fn failing_get_schema(
        _stream: *mut FFI_ArrowArrayStream,
        _out: *mut FFI_ArrowSchema,
    ) -> c_int {
        EIO
    }

    unsafe extern "C" fn working_get_schema(
        _stream: *mut FFI_ArrowArrayStream,
        out: *mut FFI_ArrowSchema,
    ) -> c_int {
        let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
        unsafe { std::ptr::write(out, FFI_ArrowSchema::try_from(&schema).unwrap()) };
        0
    }

    unsafe extern "C" fn failing_get_next(
        _stream: *mut FFI_ArrowArrayStream,
        _out: *mut FFI_ArrowArray,
    ) -> c_int {
        EIO
    }

    unsafe extern "C" fn producer_last_error(_stream: *mut FFI_ArrowArrayStream) -> *const c_char {
        c"the producer failed".as_ptr()
    }

    unsafe extern "C" fn null_last_error(_stream: *mut FFI_ArrowArrayStream) -> *const c_char {
        std::ptr::null()
    }

    unsafe extern "C" fn mark_released(stream: *mut FFI_ArrowArrayStream) {
        unsafe { (*stream).release = None };
    }

    fn failing_stream(
        get_last_error: Option<unsafe extern "C" fn(*mut FFI_ArrowArrayStream) -> *const c_char>,
    ) -> FFI_ArrowArrayStream {
        let mut stream = FFI_ArrowArrayStream::empty();
        stream.get_schema = Some(failing_get_schema);
        stream.get_next = Some(failing_get_next);
        stream.get_last_error = get_last_error;
        stream.release = Some(mark_released);
        stream
    }

    #[test]
    fn test_import_stream_without_get_schema_callback() {
        // The C Stream Interface requires get_schema for an unreleased stream.
        let stream = unsafe {
            FFI_ArrowArrayStream::new_unchecked(
                None,
                Some(failing_get_next),
                None,
                Some(mark_released),
                std::ptr::null_mut(),
            )
        };

        let err = RecordBatchStreamReader::try_new(stream).unwrap_err();
        assert!(matches!(err, ArrowError::CDataInterface(_)));
        assert!(err.to_string().contains("get_schema"));
    }

    #[test]
    fn test_import_stream_without_get_next_callback() {
        // The C Stream Interface requires get_next for an unreleased stream.
        let stream = unsafe {
            FFI_ArrowArrayStream::new_unchecked(
                Some(working_get_schema),
                None,
                None,
                Some(mark_released),
                std::ptr::null_mut(),
            )
        };

        let err = RecordBatchStreamReader::try_new(stream)
            .unwrap()
            .next()
            .unwrap()
            .unwrap_err();
        assert!(matches!(err, ArrowError::CDataInterface(_)));
        assert!(err.to_string().contains("get_next"));
    }

    #[test]
    fn test_import_schema_error_reports_producer_message() {
        let err = RecordBatchStreamReader::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_schema_error_without_producer_message() {
        // A producer need not supply a message: `get_last_error` may return NULL when no
        // detailed description is available.
        let err =
            RecordBatchStreamReader::try_new(failing_stream(Some(null_last_error))).unwrap_err();
        assert_eq!(
            err.to_string(),
            format!(
                "C Data interface error: Cannot get schema from input stream. Error code: {EIO}"
            )
        );
    }

    #[test]
    fn test_import_schema_error_without_error_callback() {
        let err = RecordBatchStreamReader::try_new(failing_stream(None)).unwrap_err();
        assert_eq!(
            err.to_string(),
            format!(
                "C Data interface error: Cannot get schema from input stream. Error code: {EIO}"
            )
        );
    }

    #[test]
    fn test_import_next_error_without_producer_message() {
        // Previously panicked: the message was unwrapped without checking that the producer
        // supplied one.
        let mut stream = failing_stream(Some(null_last_error));
        stream.get_schema = Some(working_get_schema);

        let err = RecordBatchStreamReader::try_new(stream)
            .unwrap()
            .next()
            .unwrap()
            .unwrap_err();
        assert_eq!(
            err.to_string(),
            format!(
                "C Data interface error: Cannot get next batch from input stream. Error code: {EIO}"
            )
        );
    }

    // A consumer wraps the release callback with its own, then chains back to
    // the original on drop. This is the same wrap-release pattern the
    // release/private_data accessors exist for (#9771).
    static STREAM_WRAPPER_RAN: std::sync::atomic::AtomicBool =
        std::sync::atomic::AtomicBool::new(false);

    struct StreamWrapperData {
        original_release: Option<unsafe extern "C" fn(*mut FFI_ArrowArrayStream)>,
        original_private_data: *mut c_void,
    }

    unsafe extern "C" fn wrapping_release(stream: *mut FFI_ArrowArrayStream) {
        use std::sync::atomic::Ordering;
        let stream = unsafe { &mut *stream };
        let data = unsafe { Box::from_raw(stream.private_data().cast::<StreamWrapperData>()) };
        STREAM_WRAPPER_RAN.store(true, Ordering::SeqCst);
        unsafe { stream.set_release(data.original_release) };
        unsafe { stream.set_private_data(data.original_private_data) };
        if let Some(release) = stream.release() {
            unsafe { release(stream) };
        }
    }

    #[test]
    fn test_wrap_release_callback() {
        use std::sync::atomic::Ordering;

        let batch_reader = Box::new(TestRecordBatchReader::new(
            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])),
            Box::new(std::iter::empty()),
        ));
        let mut stream = FFI_ArrowArrayStream::new(batch_reader);

        let data = Box::new(StreamWrapperData {
            original_release: stream.release(),
            original_private_data: stream.private_data(),
        });
        unsafe { stream.set_release(Some(wrapping_release)) };
        unsafe { stream.set_private_data(Box::into_raw(data).cast::<c_void>()) };

        drop(stream); // runs wrapping_release, which chains to the original
        assert!(STREAM_WRAPPER_RAN.load(Ordering::SeqCst));
    }

    // A producer that exports a stream of plain arrays rather than record batches.
    // `FFI_ArrowArrayStream::new` cannot express this, since it takes a
    // `RecordBatchReader`. See <https://github.com/apache/arrow-rs/issues/6586>.
    struct ArrayStreamPrivateData {
        field: Field,
        arrays: std::vec::IntoIter<Int32Array>,
    }

    unsafe extern "C" fn array_stream_get_schema(
        stream: *mut FFI_ArrowArrayStream,
        out: *mut FFI_ArrowSchema,
    ) -> c_int {
        let private_data = unsafe { &*(*stream).private_data().cast::<ArrayStreamPrivateData>() };
        let schema = FFI_ArrowSchema::try_from(&private_data.field).unwrap();
        unsafe { std::ptr::write(out, schema) };
        0
    }

    unsafe extern "C" fn array_stream_get_next(
        stream: *mut FFI_ArrowArrayStream,
        out: *mut FFI_ArrowArray,
    ) -> c_int {
        let private_data =
            unsafe { &mut *(*stream).private_data().cast::<ArrayStreamPrivateData>() };
        match private_data.arrays.next() {
            // Marks ArrowArray released to indicate reaching the end of stream.
            None => unsafe { std::ptr::write(out, FFI_ArrowArray::empty()) },
            Some(array) => unsafe {
                std::ptr::write_unaligned(out, FFI_ArrowArray::new(&array.to_data()))
            },
        }
        0
    }

    unsafe extern "C" fn array_stream_get_last_error(
        _stream: *mut FFI_ArrowArrayStream,
    ) -> *const c_char {
        std::ptr::null()
    }

    unsafe extern "C" fn array_stream_release(stream: *mut FFI_ArrowArrayStream) {
        let private_data =
            unsafe { Box::from_raw((*stream).private_data().cast::<ArrayStreamPrivateData>()) };
        drop(private_data);
        // Clears every callback and marks the stream released, without dropping the
        // value being overwritten.
        unsafe { std::ptr::write(stream, FFI_ArrowArrayStream::empty()) };
    }

    fn array_stream(field: Field, arrays: Vec<Int32Array>) -> FFI_ArrowArrayStream {
        let private_data = Box::new(ArrayStreamPrivateData {
            field,
            arrays: arrays.into_iter(),
        });

        unsafe {
            FFI_ArrowArrayStream::new_unchecked(
                Some(array_stream_get_schema),
                Some(array_stream_get_next),
                Some(array_stream_get_last_error),
                Some(array_stream_release),
                Box::into_raw(private_data).cast::<c_void>(),
            )
        }
    }

    #[test]
    fn test_new_unchecked_exports_stream_of_arrays() {
        let field = Field::new("a", DataType::Int32, true);
        let arrays = vec![
            Int32Array::from(vec![1, 2, 3]),
            Int32Array::from(vec![4, 5]),
        ];
        let mut stream = array_stream(field.clone(), arrays.clone());

        // Drive the stream through the callbacks the constructor stored, so a
        // field the constructor put in the wrong place would surface here.
        let get_schema_fn = stream.get_schema().unwrap();
        let get_next_fn = stream.get_next().unwrap();
        let get_last_error_fn = stream.get_last_error().unwrap();

        let mut ffi_schema = FFI_ArrowSchema::empty();
        let ret_code = unsafe { get_schema_fn(&raw mut stream, &raw mut ffi_schema) };
        assert_eq!(ret_code, 0);
        assert_eq!(Field::try_from(&ffi_schema).unwrap(), field);

        let mut produced = vec![];
        loop {
            let mut ffi_array = FFI_ArrowArray::empty();
            let ret_code = unsafe { get_next_fn(&raw mut stream, &raw mut ffi_array) };
            assert_eq!(ret_code, 0);

            // The end of stream has been reached
            if ffi_array.is_released() {
                break;
            }

            let data = unsafe { from_ffi(ffi_array, &ffi_schema) }.unwrap();
            produced.push(Int32Array::from(data));
        }

        assert_eq!(produced, arrays);

        // This producer never fails, so it reports no error message.
        assert!(unsafe { get_last_error_fn(&raw mut stream) }.is_null());

        // Runs the stored release callback, freeing the private data exactly once.
        drop(stream);
    }

    #[test]
    fn test_accessors_report_an_empty_stream_as_released() {
        let stream = FFI_ArrowArrayStream::empty();
        assert!(stream.get_schema().is_none());
        assert!(stream.get_next().is_none());
        assert!(stream.get_last_error().is_none());
        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"
        );
    }
}