Skip to main content

arrow_array/
ffi_stream.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Contains declarations to bind to the [C Stream Interface](https://arrow.apache.org/docs/format/CStreamInterface.html).
19//!
20//! This module has two main interfaces:
21//! One interface maps C ABI to native Rust types, i.e. convert c-pointers, c_char, to native rust.
22//! This is handled by [FFI_ArrowArrayStream].
23//!
24//! The second interface is used to import `FFI_ArrowArrayStream` as Rust implementation `RecordBatch` reader.
25//! This is handled by `ArrowArrayStreamReader`.
26//!
27//! ```ignore
28//! # use std::fs::File;
29//! # use std::sync::Arc;
30//! # use arrow::error::Result;
31//! # use arrow::ffi_stream::{export_reader_into_raw, ArrowArrayStreamReader, FFI_ArrowArrayStream};
32//! # use arrow::ipc::reader::FileReader;
33//! # use arrow::record_batch::RecordBatchReader;
34//! # fn main() -> Result<()> {
35//! // create an record batch reader natively
36//! let file = File::open("arrow_file").unwrap();
37//! let reader = Box::new(FileReader::try_new(file).unwrap());
38//!
39//! // export it
40//! let mut stream = FFI_ArrowArrayStream::empty();
41//! unsafe { export_reader_into_raw(reader, &mut stream) };
42//!
43//! // consumed and used by something else...
44//!
45//! // import it
46//! let stream_reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() };
47//! let imported_schema = stream_reader.schema();
48//!
49//! let mut produced_batches = vec![];
50//! for batch in stream_reader {
51//!      produced_batches.push(batch.unwrap());
52//! }
53//! Ok(())
54//! }
55//! ```
56
57use arrow_schema::DataType;
58use std::ffi::CStr;
59use std::ptr::addr_of;
60use std::{
61    ffi::CString,
62    os::raw::{c_char, c_int, c_void},
63    sync::Arc,
64};
65
66use arrow_data::ffi::FFI_ArrowArray;
67use arrow_schema::{ArrowError, Schema, SchemaRef, ffi::FFI_ArrowSchema};
68
69use crate::RecordBatchOptions;
70use crate::array::Array;
71use crate::array::StructArray;
72use crate::ffi::from_ffi_and_data_type;
73use crate::record_batch::{RecordBatch, RecordBatchReader};
74
75type Result<T> = std::result::Result<T, ArrowError>;
76
77// Errno values returned through the C stream interface, taken from libc so they match
78// the platform the consumer interprets them against.
79#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
80use libc::{EINVAL, EIO, ENOMEM, ENOSYS};
81
82// wasm32-unknown-unknown has no libc, and no OS to interpret the codes either — any
83// non-zero value works there, so use Linux's.
84#[cfg(all(target_family = "wasm", target_os = "unknown"))]
85const ENOMEM: i32 = 12;
86#[cfg(all(target_family = "wasm", target_os = "unknown"))]
87const EIO: i32 = 5;
88#[cfg(all(target_family = "wasm", target_os = "unknown"))]
89const EINVAL: i32 = 22;
90#[cfg(all(target_family = "wasm", target_os = "unknown"))]
91const ENOSYS: i32 = 38;
92
93/// ABI-compatible struct for `ArrayStream` from C Stream Interface
94/// See <https://arrow.apache.org/docs/format/CStreamInterface.html#structure-definitions>
95/// This was created by bindgen
96#[repr(C)]
97#[derive(Debug)]
98pub struct FFI_ArrowArrayStream {
99    // Fields are intentionally private so safety guarantees can be upheld via
100    // explicit unsafe functions.
101    /// C function to get schema from the stream
102    get_schema: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowSchema) -> c_int>,
103    /// C function to get next array from the stream
104    get_next: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowArray) -> c_int>,
105    /// C function to get the error from last operation on the stream
106    get_last_error: Option<unsafe extern "C" fn(arg1: *mut Self) -> *const c_char>,
107    /// C function to release the stream
108    release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
109    /// Private data used by the stream, owned by the release callback.
110    private_data: *mut c_void,
111}
112
113unsafe impl Send for FFI_ArrowArrayStream {}
114
115// callback used to drop [FFI_ArrowArrayStream] when it is exported.
116unsafe extern "C" fn release_stream(stream: *mut FFI_ArrowArrayStream) {
117    if stream.is_null() {
118        return;
119    }
120    let stream = unsafe { &mut *stream };
121
122    stream.get_schema = None;
123    stream.get_next = None;
124    stream.get_last_error = None;
125
126    let private_data = unsafe { Box::from_raw(stream.private_data.cast::<StreamPrivateData>()) };
127    drop(private_data);
128
129    stream.release = None;
130}
131
132struct StreamPrivateData {
133    batch_reader: Box<dyn RecordBatchReader + Send>,
134    last_error: Option<CString>,
135}
136
137// The callback used to get array schema
138unsafe extern "C" fn get_schema(
139    stream: *mut FFI_ArrowArrayStream,
140    schema: *mut FFI_ArrowSchema,
141) -> c_int {
142    ExportedArrayStream { stream }.get_schema(schema)
143}
144
145// The callback used to get next array
146unsafe extern "C" fn get_next(
147    stream: *mut FFI_ArrowArrayStream,
148    array: *mut FFI_ArrowArray,
149) -> c_int {
150    ExportedArrayStream { stream }.get_next(array)
151}
152
153// The callback used to get the error from last operation on the `FFI_ArrowArrayStream`
154unsafe extern "C" fn get_last_error(stream: *mut FFI_ArrowArrayStream) -> *const c_char {
155    let mut ffi_stream = ExportedArrayStream { stream };
156    // The consumer should not take ownership of this string, we should return
157    // a const pointer to it.
158    match ffi_stream.get_last_error() {
159        Some(err_string) => err_string.as_ptr(),
160        None => std::ptr::null(),
161    }
162}
163
164impl Drop for FFI_ArrowArrayStream {
165    fn drop(&mut self) {
166        match self.release {
167            None => (),
168            Some(release) => unsafe { release(self) },
169        }
170    }
171}
172
173impl FFI_ArrowArrayStream {
174    /// Creates a new [`FFI_ArrowArrayStream`].
175    pub fn new(batch_reader: Box<dyn RecordBatchReader + Send>) -> Self {
176        let private_data = Box::new(StreamPrivateData {
177            batch_reader,
178            last_error: None,
179        });
180
181        Self {
182            get_schema: Some(get_schema),
183            get_next: Some(get_next),
184            get_last_error: Some(get_last_error),
185            release: Some(release_stream),
186            private_data: Box::into_raw(private_data).cast::<c_void>(),
187        }
188    }
189
190    /// Takes ownership of the pointed to [`FFI_ArrowArrayStream`]
191    ///
192    /// This acts to [move] the data out of `raw_stream`, setting the release callback to NULL
193    ///
194    /// # Safety
195    ///
196    /// * `raw_stream` must be [valid] for reads and writes
197    /// * `raw_stream` must be properly aligned
198    /// * `raw_stream` must point to a properly initialized value of [`FFI_ArrowArrayStream`]
199    ///
200    /// [move]: https://arrow.apache.org/docs/format/CDataInterface.html#moving-an-array
201    /// [valid]: https://doc.rust-lang.org/std/ptr/index.html#safety
202    pub unsafe fn from_raw(raw_stream: *mut FFI_ArrowArrayStream) -> Self {
203        unsafe { std::ptr::replace(raw_stream, Self::empty()) }
204    }
205
206    /// Creates a new empty [FFI_ArrowArrayStream]. Used to import from the C Stream Interface.
207    pub fn empty() -> Self {
208        Self {
209            get_schema: None,
210            get_next: None,
211            get_last_error: None,
212            release: None,
213            private_data: std::ptr::null_mut(),
214        }
215    }
216
217    /// Returns the producer-provided release callback, if any.
218    pub fn release(&self) -> Option<unsafe extern "C" fn(arg1: *mut Self)> {
219        self.release
220    }
221
222    /// Returns the opaque producer-provided private data pointer.
223    pub fn private_data(&self) -> *mut c_void {
224        self.private_data
225    }
226
227    /// Replaces the release callback, returning the previous one.
228    ///
229    /// Lets a consumer wrap release: save the old callback, install its own, and
230    /// chain back on drop. See <https://github.com/apache/arrow-rs/issues/9771>.
231    ///
232    /// # Safety
233    ///
234    /// [`Drop`] calls this callback with a pointer to `self`. The new callback
235    /// must correctly release this stream (usually by chaining to the returned
236    /// one) and must match the [`FFI_ArrowArrayStream::private_data`] it reads.
237    /// A wrong callback is undefined behavior on drop.
238    pub unsafe fn set_release(
239        &mut self,
240        release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
241    ) -> Option<unsafe extern "C" fn(arg1: *mut Self)> {
242        std::mem::replace(&mut self.release, release)
243    }
244
245    /// Replaces the private data pointer, returning the previous one.
246    ///
247    /// # Safety
248    ///
249    /// The old pointer is returned without being freed; the caller owns it from
250    /// here. The new pointer must match what the current
251    /// [`FFI_ArrowArrayStream::release`] callback expects.
252    pub unsafe fn set_private_data(&mut self, private_data: *mut c_void) -> *mut c_void {
253        std::mem::replace(&mut self.private_data, private_data)
254    }
255}
256
257struct ExportedArrayStream {
258    stream: *mut FFI_ArrowArrayStream,
259}
260
261impl ExportedArrayStream {
262    fn get_private_data(&mut self) -> &mut StreamPrivateData {
263        unsafe { &mut *(*self.stream).private_data.cast::<StreamPrivateData>() }
264    }
265
266    pub fn get_schema(&mut self, out: *mut FFI_ArrowSchema) -> i32 {
267        let private_data = self.get_private_data();
268        let reader = &private_data.batch_reader;
269
270        let schema = FFI_ArrowSchema::try_from(reader.schema().as_ref());
271
272        match schema {
273            Ok(schema) => {
274                unsafe { std::ptr::copy(addr_of!(schema), out, 1) };
275                std::mem::forget(schema);
276                0
277            }
278            Err(ref err) => {
279                private_data.last_error = Some(
280                    CString::new(err.to_string()).expect("Error string has a null byte in it."),
281                );
282                get_error_code(err)
283            }
284        }
285    }
286
287    pub fn get_next(&mut self, out: *mut FFI_ArrowArray) -> i32 {
288        let private_data = self.get_private_data();
289        let reader = &mut private_data.batch_reader;
290
291        match reader.next() {
292            None => {
293                // Marks ArrowArray released to indicate reaching the end of stream.
294                unsafe { std::ptr::write(out, FFI_ArrowArray::empty()) }
295                0
296            }
297            Some(next_batch) => {
298                if let Ok(batch) = next_batch {
299                    let struct_array = StructArray::from(batch);
300                    let array = FFI_ArrowArray::new(&struct_array.to_data());
301
302                    unsafe { std::ptr::write_unaligned(out, array) };
303                    0
304                } else {
305                    let err = &next_batch.unwrap_err();
306                    private_data.last_error = Some(
307                        CString::new(err.to_string()).expect("Error string has a null byte in it."),
308                    );
309                    get_error_code(err)
310                }
311            }
312        }
313    }
314
315    pub fn get_last_error(&mut self) -> Option<&CString> {
316        self.get_private_data().last_error.as_ref()
317    }
318}
319
320fn get_error_code(err: &ArrowError) -> i32 {
321    match err {
322        ArrowError::NotYetImplemented(_) => ENOSYS,
323        ArrowError::MemoryError(_) => ENOMEM,
324        ArrowError::IoError(_, _) => EIO,
325        _ => EINVAL,
326    }
327}
328
329/// A `RecordBatchReader` which imports Arrays from `FFI_ArrowArrayStream`.
330///
331/// Struct used to fetch `RecordBatch` from the C Stream Interface.
332/// Its main responsibility is to expose `RecordBatchReader` functionality
333/// that requires [FFI_ArrowArrayStream].
334#[derive(Debug)]
335pub struct ArrowArrayStreamReader {
336    stream: FFI_ArrowArrayStream,
337    schema: SchemaRef,
338}
339
340/// Returns the producer's message for the last failed call on a `FFI_ArrowArrayStream`.
341///
342/// Returns `None` when the producer supplies no message, either because it installs no
343/// `get_last_error` callback or because that callback returns NULL: the C Stream Interface
344/// lets `get_last_error` return NULL when no detailed description is available.
345///
346/// # Safety
347///
348/// `stream_ptr` must point to a valid, not yet released [`FFI_ArrowArrayStream`], and the last
349/// operation on it must have returned an error: the C Stream Interface forbids calling
350/// `get_last_error` in any other case.
351unsafe fn producer_error(stream_ptr: *mut FFI_ArrowArrayStream) -> Option<String> {
352    let get_last_error = unsafe { (*stream_ptr).get_last_error }?;
353
354    let error_str = unsafe { get_last_error(stream_ptr) };
355    if error_str.is_null() {
356        return None;
357    }
358
359    Some(
360        unsafe { CStr::from_ptr(error_str) }
361            .to_string_lossy()
362            .into_owned(),
363    )
364}
365
366/// Gets schema from a raw pointer of `FFI_ArrowArrayStream`. This is used when constructing
367/// `ArrowArrayStreamReader` to cache schema.
368fn get_stream_schema(stream_ptr: *mut FFI_ArrowArrayStream) -> Result<SchemaRef> {
369    let mut schema = FFI_ArrowSchema::empty();
370
371    let ret_code = unsafe { (*stream_ptr).get_schema.unwrap()(stream_ptr, &mut schema) };
372
373    if ret_code == 0 {
374        let schema = Schema::try_from(&schema)?;
375        Ok(Arc::new(schema))
376    } else {
377        let message = format!("Cannot get schema from input stream. Error code: {ret_code}");
378        // SAFETY: `stream_ptr` is valid and unreleased, and the `get_schema` call above
379        // returned a non-zero code.
380        let message = match unsafe { producer_error(stream_ptr) } {
381            Some(producer_message) => format!("{message}. Producer error: {producer_message}"),
382            None => message,
383        };
384        Err(ArrowError::CDataInterface(message))
385    }
386}
387
388impl ArrowArrayStreamReader {
389    /// Creates a new `ArrowArrayStreamReader` from a `FFI_ArrowArrayStream`.
390    /// This is used to import from the C Stream Interface.
391    pub fn try_new(mut stream: FFI_ArrowArrayStream) -> Result<Self> {
392        if stream.release.is_none() {
393            return Err(ArrowError::CDataInterface(
394                "input stream is already released".to_string(),
395            ));
396        }
397
398        let schema = get_stream_schema(&mut stream)?;
399
400        Ok(Self { stream, schema })
401    }
402
403    /// Creates a new `ArrowArrayStreamReader` from a raw pointer of `FFI_ArrowArrayStream`.
404    ///
405    /// Assumes that the pointer represents valid C Stream Interfaces.
406    /// This function copies the content from the raw pointer and cleans up it to prevent
407    /// double-dropping. The caller is responsible for freeing up the memory allocated for
408    /// the pointer.
409    ///
410    /// # Safety
411    ///
412    /// See [`FFI_ArrowArrayStream::from_raw`]
413    pub unsafe fn from_raw(raw_stream: *mut FFI_ArrowArrayStream) -> Result<Self> {
414        Self::try_new(unsafe { FFI_ArrowArrayStream::from_raw(raw_stream) })
415    }
416}
417
418impl Iterator for ArrowArrayStreamReader {
419    type Item = Result<RecordBatch>;
420
421    fn next(&mut self) -> Option<Self::Item> {
422        let mut array = FFI_ArrowArray::empty();
423
424        let ret_code = unsafe { self.stream.get_next.unwrap()(&mut self.stream, &mut array) };
425
426        if ret_code == 0 {
427            // The end of stream has been reached
428            if array.is_released() {
429                return None;
430            }
431
432            let result = unsafe {
433                from_ffi_and_data_type(array, DataType::Struct(self.schema().fields().clone()))
434            };
435            Some(result.and_then(|data| {
436                let len = data.len();
437                RecordBatch::try_new_with_options(
438                    self.schema.clone(),
439                    StructArray::from(data).into_parts().1,
440                    &RecordBatchOptions::new().with_row_count(Some(len)),
441                )
442            }))
443        } else {
444            let message =
445                format!("Cannot get next batch from input stream. Error code: {ret_code}");
446            // SAFETY: `self.stream` is valid and unreleased by construction, and the
447            // `get_next` call above returned a non-zero code.
448            let message = match unsafe { producer_error(&mut self.stream) } {
449                Some(producer_message) => format!("{message}. Producer error: {producer_message}"),
450                None => message,
451            };
452            Some(Err(ArrowError::CDataInterface(message)))
453        }
454    }
455}
456
457impl RecordBatchReader for ArrowArrayStreamReader {
458    fn schema(&self) -> SchemaRef {
459        self.schema.clone()
460    }
461}
462
463#[cfg(test)]
464mod tests {
465    use super::*;
466    use std::collections::HashMap;
467
468    use arrow_schema::Field;
469
470    use crate::array::Int32Array;
471    use crate::ffi::from_ffi;
472
473    struct TestRecordBatchReader {
474        schema: SchemaRef,
475        iter: Box<dyn Iterator<Item = Result<RecordBatch>> + Send>,
476    }
477
478    impl TestRecordBatchReader {
479        pub fn new(
480            schema: SchemaRef,
481            iter: Box<dyn Iterator<Item = Result<RecordBatch>> + Send>,
482        ) -> TestRecordBatchReader {
483            TestRecordBatchReader { schema, iter }
484        }
485    }
486
487    impl Iterator for TestRecordBatchReader {
488        type Item = Result<RecordBatch>;
489
490        fn next(&mut self) -> Option<Self::Item> {
491            self.iter.next()
492        }
493    }
494
495    impl RecordBatchReader for TestRecordBatchReader {
496        fn schema(&self) -> SchemaRef {
497            self.schema.clone()
498        }
499    }
500
501    fn _test_round_trip_export(batch: RecordBatch, schema: Arc<Schema>) -> Result<()> {
502        let iter = Box::new(vec![batch.clone(), batch.clone()].into_iter().map(Ok)) as _;
503
504        let reader = Box::new(TestRecordBatchReader::new(schema.clone(), iter));
505
506        // Export a `RecordBatchReader` through `FFI_ArrowArrayStream`
507        let mut ffi_stream = FFI_ArrowArrayStream::new(reader);
508
509        // Get schema from `FFI_ArrowArrayStream`
510        let mut ffi_schema = FFI_ArrowSchema::empty();
511        let ret_code = unsafe { get_schema(&mut ffi_stream, &mut ffi_schema) };
512        assert_eq!(ret_code, 0);
513
514        let exported_schema = Schema::try_from(&ffi_schema).unwrap();
515        assert_eq!(&exported_schema, schema.as_ref());
516
517        // Get array from `FFI_ArrowArrayStream`
518        let mut produced_batches = vec![];
519        loop {
520            let mut ffi_array = FFI_ArrowArray::empty();
521            let ret_code = unsafe { get_next(&mut ffi_stream, &mut ffi_array) };
522            assert_eq!(ret_code, 0);
523
524            // The end of stream has been reached
525            if ffi_array.is_released() {
526                break;
527            }
528
529            let array = unsafe { from_ffi(ffi_array, &ffi_schema) }.unwrap();
530            let len = array.len();
531
532            let record_batch = RecordBatch::try_new_with_options(
533                SchemaRef::from(exported_schema.clone()),
534                StructArray::from(array).into_parts().1,
535                &RecordBatchOptions::new().with_row_count(Some(len)),
536            )
537            .unwrap();
538            produced_batches.push(record_batch);
539        }
540
541        assert_eq!(produced_batches, vec![batch.clone(), batch]);
542
543        Ok(())
544    }
545
546    fn _test_round_trip_import(batch: RecordBatch, schema: Arc<Schema>) -> Result<()> {
547        let iter = Box::new(vec![batch.clone(), batch.clone()].into_iter().map(Ok)) as _;
548
549        let reader = Box::new(TestRecordBatchReader::new(schema.clone(), iter));
550
551        // Import through `FFI_ArrowArrayStream` as `ArrowArrayStreamReader`
552        let stream = FFI_ArrowArrayStream::new(reader);
553        let stream_reader = ArrowArrayStreamReader::try_new(stream).unwrap();
554
555        let imported_schema = stream_reader.schema();
556        assert_eq!(imported_schema, schema);
557
558        let mut produced_batches = vec![];
559        for batch in stream_reader {
560            produced_batches.push(batch.unwrap());
561        }
562
563        assert_eq!(produced_batches, vec![batch.clone(), batch]);
564
565        Ok(())
566    }
567
568    #[test]
569    fn test_stream_round_trip() {
570        let array = Int32Array::from(vec![Some(2), None, Some(1), None]);
571        let array: Arc<dyn Array> = Arc::new(array);
572        let metadata = HashMap::from([("foo".to_owned(), "bar".to_owned())]);
573
574        let schema = Arc::new(Schema::new_with_metadata(
575            vec![
576                Field::new("a", array.data_type().clone(), true).with_metadata(metadata.clone()),
577                Field::new("b", array.data_type().clone(), true).with_metadata(metadata.clone()),
578                Field::new("c", array.data_type().clone(), true).with_metadata(metadata.clone()),
579            ],
580            metadata,
581        ));
582        let batch = RecordBatch::try_new(schema.clone(), vec![array.clone(), array.clone(), array])
583            .unwrap();
584
585        _test_round_trip_export(batch.clone(), schema.clone()).unwrap();
586        _test_round_trip_import(batch, schema).unwrap();
587    }
588
589    #[test]
590    fn test_stream_round_trip_no_columns() {
591        let metadata = HashMap::from([("foo".to_owned(), "bar".to_owned())]);
592
593        let schema = Arc::new(Schema::new_with_metadata(Vec::<Field>::new(), metadata));
594        let batch = RecordBatch::try_new_with_options(
595            schema.clone(),
596            Vec::<Arc<dyn Array>>::new(),
597            &RecordBatchOptions::new().with_row_count(Some(10)),
598        )
599        .unwrap();
600
601        _test_round_trip_export(batch.clone(), schema.clone()).unwrap();
602        _test_round_trip_import(batch, schema).unwrap();
603    }
604
605    #[test]
606    fn test_error_import() -> Result<()> {
607        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
608
609        let iter =
610            Box::new(vec![Err(ArrowError::MemoryError("out of memory".to_string()))].into_iter());
611
612        let reader = Box::new(TestRecordBatchReader::new(schema.clone(), iter));
613
614        // Import through `FFI_ArrowArrayStream` as `ArrowArrayStreamReader`
615        let stream = FFI_ArrowArrayStream::new(reader);
616        let stream_reader = ArrowArrayStreamReader::try_new(stream).unwrap();
617
618        let imported_schema = stream_reader.schema();
619        assert_eq!(imported_schema, schema);
620
621        let mut produced_batches = vec![];
622        for batch in stream_reader {
623            produced_batches.push(batch);
624        }
625
626        // The results should outlive the lifetime of the stream itself.
627        assert_eq!(produced_batches.len(), 1);
628        assert_eq!(
629            produced_batches[0].as_ref().unwrap_err().to_string(),
630            format!(
631                "C Data interface error: Cannot get next batch from input stream. \
632                 Error code: {ENOMEM}. Producer error: Memory error: out of memory"
633            )
634        );
635
636        Ok(())
637    }
638
639    unsafe extern "C" fn failing_get_schema(
640        _stream: *mut FFI_ArrowArrayStream,
641        _out: *mut FFI_ArrowSchema,
642    ) -> c_int {
643        EIO
644    }
645
646    unsafe extern "C" fn working_get_schema(
647        _stream: *mut FFI_ArrowArrayStream,
648        out: *mut FFI_ArrowSchema,
649    ) -> c_int {
650        let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
651        unsafe { std::ptr::write(out, FFI_ArrowSchema::try_from(&schema).unwrap()) };
652        0
653    }
654
655    unsafe extern "C" fn failing_get_next(
656        _stream: *mut FFI_ArrowArrayStream,
657        _out: *mut FFI_ArrowArray,
658    ) -> c_int {
659        EIO
660    }
661
662    unsafe extern "C" fn producer_last_error(_stream: *mut FFI_ArrowArrayStream) -> *const c_char {
663        c"the producer failed".as_ptr()
664    }
665
666    unsafe extern "C" fn null_last_error(_stream: *mut FFI_ArrowArrayStream) -> *const c_char {
667        std::ptr::null()
668    }
669
670    unsafe extern "C" fn mark_released(stream: *mut FFI_ArrowArrayStream) {
671        unsafe { (*stream).release = None };
672    }
673
674    fn failing_stream(
675        get_last_error: Option<unsafe extern "C" fn(*mut FFI_ArrowArrayStream) -> *const c_char>,
676    ) -> FFI_ArrowArrayStream {
677        let mut stream = FFI_ArrowArrayStream::empty();
678        stream.get_schema = Some(failing_get_schema);
679        stream.get_next = Some(failing_get_next);
680        stream.get_last_error = get_last_error;
681        stream.release = Some(mark_released);
682        stream
683    }
684
685    #[test]
686    fn test_import_schema_error_reports_producer_message() {
687        let err =
688            ArrowArrayStreamReader::try_new(failing_stream(Some(producer_last_error))).unwrap_err();
689        assert_eq!(
690            err.to_string(),
691            format!(
692                "C Data interface error: Cannot get schema from input stream. \
693                 Error code: {EIO}. Producer error: the producer failed"
694            )
695        );
696    }
697
698    #[test]
699    fn test_import_schema_error_without_producer_message() {
700        // A producer need not supply a message: `get_last_error` may return NULL when no
701        // detailed description is available.
702        let err =
703            ArrowArrayStreamReader::try_new(failing_stream(Some(null_last_error))).unwrap_err();
704        assert_eq!(
705            err.to_string(),
706            format!(
707                "C Data interface error: Cannot get schema from input stream. Error code: {EIO}"
708            )
709        );
710    }
711
712    #[test]
713    fn test_import_schema_error_without_error_callback() {
714        let err = ArrowArrayStreamReader::try_new(failing_stream(None)).unwrap_err();
715        assert_eq!(
716            err.to_string(),
717            format!(
718                "C Data interface error: Cannot get schema from input stream. Error code: {EIO}"
719            )
720        );
721    }
722
723    #[test]
724    fn test_import_next_error_without_producer_message() {
725        // Previously panicked: the message was unwrapped without checking that the producer
726        // supplied one.
727        let mut stream = failing_stream(Some(null_last_error));
728        stream.get_schema = Some(working_get_schema);
729
730        let err = ArrowArrayStreamReader::try_new(stream)
731            .unwrap()
732            .next()
733            .unwrap()
734            .unwrap_err();
735        assert_eq!(
736            err.to_string(),
737            format!(
738                "C Data interface error: Cannot get next batch from input stream. Error code: {EIO}"
739            )
740        );
741    }
742
743    // A consumer wraps the release callback with its own, then chains back to
744    // the original on drop. This is the same wrap-release pattern the
745    // release/private_data accessors exist for (#9771).
746    static STREAM_WRAPPER_RAN: std::sync::atomic::AtomicBool =
747        std::sync::atomic::AtomicBool::new(false);
748
749    struct StreamWrapperData {
750        original_release: Option<unsafe extern "C" fn(*mut FFI_ArrowArrayStream)>,
751        original_private_data: *mut c_void,
752    }
753
754    unsafe extern "C" fn wrapping_release(stream: *mut FFI_ArrowArrayStream) {
755        use std::sync::atomic::Ordering;
756        let stream = unsafe { &mut *stream };
757        let data = unsafe { Box::from_raw(stream.private_data().cast::<StreamWrapperData>()) };
758        STREAM_WRAPPER_RAN.store(true, Ordering::SeqCst);
759        unsafe { stream.set_release(data.original_release) };
760        unsafe { stream.set_private_data(data.original_private_data) };
761        if let Some(release) = stream.release() {
762            unsafe { release(stream) };
763        }
764    }
765
766    #[test]
767    fn test_wrap_release_callback() {
768        use std::sync::atomic::Ordering;
769
770        let batch_reader = Box::new(TestRecordBatchReader::new(
771            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])),
772            Box::new(std::iter::empty()),
773        ));
774        let mut stream = FFI_ArrowArrayStream::new(batch_reader);
775
776        let data = Box::new(StreamWrapperData {
777            original_release: stream.release(),
778            original_private_data: stream.private_data(),
779        });
780        unsafe { stream.set_release(Some(wrapping_release)) };
781        unsafe { stream.set_private_data(Box::into_raw(data).cast::<c_void>()) };
782
783        drop(stream); // runs wrapping_release, which chains to the original
784        assert!(STREAM_WRAPPER_RAN.load(Ordering::SeqCst));
785    }
786}