1use 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#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
80use libc::{EINVAL, EIO, ENOMEM, ENOSYS};
81
82#[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#[repr(C)]
97#[derive(Debug)]
98pub struct FFI_ArrowArrayStream {
99 get_schema: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowSchema) -> c_int>,
103 get_next: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowArray) -> c_int>,
105 get_last_error: Option<unsafe extern "C" fn(arg1: *mut Self) -> *const c_char>,
107 release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
109 private_data: *mut c_void,
111}
112
113unsafe impl Send for FFI_ArrowArrayStream {}
114
115unsafe 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
137unsafe 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
145unsafe 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
153unsafe extern "C" fn get_last_error(stream: *mut FFI_ArrowArrayStream) -> *const c_char {
155 let mut ffi_stream = ExportedArrayStream { stream };
156 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 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 pub unsafe fn from_raw(raw_stream: *mut FFI_ArrowArrayStream) -> Self {
203 unsafe { std::ptr::replace(raw_stream, Self::empty()) }
204 }
205
206 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 pub fn release(&self) -> Option<unsafe extern "C" fn(arg1: *mut Self)> {
219 self.release
220 }
221
222 pub fn private_data(&self) -> *mut c_void {
224 self.private_data
225 }
226
227 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 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 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#[derive(Debug)]
335pub struct ArrowArrayStreamReader {
336 stream: FFI_ArrowArrayStream,
337 schema: SchemaRef,
338}
339
340unsafe 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
366fn 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 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 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 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 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 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 let mut ffi_stream = FFI_ArrowArrayStream::new(reader);
508
509 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 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 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 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 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 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 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 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 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); assert!(STREAM_WRAPPER_RAN.load(Ordering::SeqCst));
785 }
786}