556 lines
19 KiB
Rust
556 lines
19 KiB
Rust
// 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 Rust implementation `RecordBatch` reader.
|
|
//! This is handled by `ArrowArrayStreamReader`.
|
|
//!
|
|
//! ```ignore
|
|
//! # use std::fs::File;
|
|
//! # use std::sync::Arc;
|
|
//! # use arrow::error::Result;
|
|
//! # use arrow::ffi_stream::{export_reader_into_raw, ArrowArrayStreamReader, FFI_ArrowArrayStream};
|
|
//! # use arrow::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 { ArrowArrayStreamReader::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, Schema, SchemaRef, ffi::FFI_ArrowSchema};
|
|
|
|
use crate::array::Array;
|
|
use crate::array::StructArray;
|
|
use crate::ffi::from_ffi_and_data_type;
|
|
use crate::record_batch::{RecordBatch, RecordBatchReader};
|
|
|
|
type Result<T> = std::result::Result<T, ArrowError>;
|
|
|
|
const ENOMEM: i32 = 12;
|
|
const EIO: i32 = 5;
|
|
const EINVAL: i32 = 22;
|
|
const ENOSYS: i32 = 78;
|
|
|
|
/// 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)]
|
|
#[allow(non_camel_case_types)]
|
|
pub struct FFI_ArrowArrayStream {
|
|
/// C function to get schema from the stream
|
|
pub 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
|
|
pub 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
|
|
pub get_last_error: Option<unsafe extern "C" fn(arg1: *mut Self) -> *const c_char>,
|
|
/// C function to release the stream
|
|
pub release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
|
|
/// Private data used by the stream
|
|
pub 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 as *mut 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) as *mut c_void,
|
|
}
|
|
}
|
|
|
|
/// 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(),
|
|
}
|
|
}
|
|
}
|
|
|
|
struct ExportedArrayStream {
|
|
stream: *mut FFI_ArrowArrayStream,
|
|
}
|
|
|
|
impl ExportedArrayStream {
|
|
fn get_private_data(&mut self) -> &mut StreamPrivateData {
|
|
unsafe { &mut *((*self.stream).private_data as *mut 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].
|
|
#[derive(Debug)]
|
|
pub struct ArrowArrayStreamReader {
|
|
stream: FFI_ArrowArrayStream,
|
|
schema: SchemaRef,
|
|
}
|
|
|
|
/// Gets schema from a raw pointer of `FFI_ArrowArrayStream`. This is used when constructing
|
|
/// `ArrowArrayStreamReader` to cache schema.
|
|
fn get_stream_schema(stream_ptr: *mut FFI_ArrowArrayStream) -> Result<SchemaRef> {
|
|
let mut schema = FFI_ArrowSchema::empty();
|
|
|
|
let ret_code = unsafe { (*stream_ptr).get_schema.unwrap()(stream_ptr, &mut schema) };
|
|
|
|
if ret_code == 0 {
|
|
let schema = Schema::try_from(&schema)?;
|
|
Ok(Arc::new(schema))
|
|
} else {
|
|
Err(ArrowError::CDataInterface(format!(
|
|
"Cannot get schema from input stream. Error code: {ret_code:?}"
|
|
)))
|
|
}
|
|
}
|
|
|
|
impl ArrowArrayStreamReader {
|
|
/// Creates a new `ArrowArrayStreamReader` from a `FFI_ArrowArrayStream`.
|
|
/// This is used to import from the C Stream Interface.
|
|
#[allow(dead_code)]
|
|
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(&mut stream)?;
|
|
|
|
Ok(Self { stream, schema })
|
|
}
|
|
|
|
/// Creates a new `ArrowArrayStreamReader` 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) })
|
|
}
|
|
|
|
/// Get the last error from `ArrowArrayStreamReader`
|
|
fn get_stream_last_error(&mut self) -> Option<String> {
|
|
let get_last_error = self.stream.get_last_error?;
|
|
|
|
let error_str = unsafe { get_last_error(&mut self.stream) };
|
|
if error_str.is_null() {
|
|
return None;
|
|
}
|
|
|
|
let error_str = unsafe { CStr::from_ptr(error_str) };
|
|
Some(error_str.to_string_lossy().to_string())
|
|
}
|
|
}
|
|
|
|
impl Iterator for ArrowArrayStreamReader {
|
|
type Item = Result<RecordBatch>;
|
|
|
|
fn next(&mut self) -> Option<Self::Item> {
|
|
let mut array = FFI_ArrowArray::empty();
|
|
|
|
let ret_code = unsafe { self.stream.get_next.unwrap()(&mut self.stream, &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| {
|
|
RecordBatch::try_new(self.schema.clone(), StructArray::from(data).into_parts().1)
|
|
}))
|
|
} else {
|
|
let last_error = self.get_stream_last_error();
|
|
let err = ArrowError::CDataInterface(last_error.unwrap());
|
|
Some(Err(err))
|
|
}
|
|
}
|
|
}
|
|
|
|
impl RecordBatchReader for ArrowArrayStreamReader {
|
|
fn schema(&self) -> SchemaRef {
|
|
self.schema.clone()
|
|
}
|
|
}
|
|
|
|
#[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>,
|
|
) -> Box<TestRecordBatchReader> {
|
|
Box::new(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(arrays: Vec<Arc<dyn Array>>) -> Result<()> {
|
|
let metadata = HashMap::from([("foo".to_owned(), "bar".to_owned())]);
|
|
let schema = Arc::new(Schema::new_with_metadata(
|
|
vec![
|
|
Field::new("a", arrays[0].data_type().clone(), true)
|
|
.with_metadata(metadata.clone()),
|
|
Field::new("b", arrays[1].data_type().clone(), true)
|
|
.with_metadata(metadata.clone()),
|
|
Field::new("c", arrays[2].data_type().clone(), true)
|
|
.with_metadata(metadata.clone()),
|
|
],
|
|
metadata,
|
|
));
|
|
let batch = RecordBatch::try_new(schema.clone(), arrays).unwrap();
|
|
let iter = Box::new(vec![batch.clone(), batch.clone()].into_iter().map(Ok)) as _;
|
|
|
|
let reader = 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(&mut ffi_stream, &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(&mut ffi_stream, &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 record_batch = RecordBatch::try_new(
|
|
SchemaRef::from(exported_schema.clone()),
|
|
StructArray::from(array).into_parts().1,
|
|
)
|
|
.unwrap();
|
|
produced_batches.push(record_batch);
|
|
}
|
|
|
|
assert_eq!(produced_batches, vec![batch.clone(), batch]);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn _test_round_trip_import(arrays: Vec<Arc<dyn Array>>) -> Result<()> {
|
|
let metadata = HashMap::from([("foo".to_owned(), "bar".to_owned())]);
|
|
let schema = Arc::new(Schema::new_with_metadata(
|
|
vec![
|
|
Field::new("a", arrays[0].data_type().clone(), true)
|
|
.with_metadata(metadata.clone()),
|
|
Field::new("b", arrays[1].data_type().clone(), true)
|
|
.with_metadata(metadata.clone()),
|
|
Field::new("c", arrays[2].data_type().clone(), true)
|
|
.with_metadata(metadata.clone()),
|
|
],
|
|
metadata,
|
|
));
|
|
let batch = RecordBatch::try_new(schema.clone(), arrays).unwrap();
|
|
let iter = Box::new(vec![batch.clone(), batch.clone()].into_iter().map(Ok)) as _;
|
|
|
|
let reader = TestRecordBatchReader::new(schema.clone(), iter);
|
|
|
|
// Import through `FFI_ArrowArrayStream` as `ArrowArrayStreamReader`
|
|
let stream = FFI_ArrowArrayStream::new(reader);
|
|
let stream_reader = ArrowArrayStreamReader::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_export() -> Result<()> {
|
|
let array = Int32Array::from(vec![Some(2), None, Some(1), None]);
|
|
let array: Arc<dyn Array> = Arc::new(array);
|
|
|
|
_test_round_trip_export(vec![array.clone(), array.clone(), array])
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_round_trip_import() -> Result<()> {
|
|
let array = Int32Array::from(vec![Some(2), None, Some(1), None]);
|
|
let array: Arc<dyn Array> = Arc::new(array);
|
|
|
|
_test_round_trip_import(vec![array.clone(), array.clone(), array])
|
|
}
|
|
|
|
#[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("".to_string()))].into_iter());
|
|
|
|
let reader = TestRecordBatchReader::new(schema.clone(), iter);
|
|
|
|
// Import through `FFI_ArrowArrayStream` as `ArrowArrayStreamReader`
|
|
let stream = FFI_ArrowArrayStream::new(reader);
|
|
let stream_reader = ArrowArrayStreamReader::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!(produced_batches[0].is_err());
|
|
|
|
Ok(())
|
|
}
|
|
}
|