Skip to content

Commit

Permalink
Align buffers from Python (FFI) (#6472)
Browse files Browse the repository at this point in the history
* Align buffers in RecordBatch.from_pyarrow_bound

* Update arrow/src/pyarrow.rs

* cargo fmt

---------

Co-authored-by: Andrew Lamb <[email protected]>
  • Loading branch information
EnricoMi and alamb authored Oct 2, 2024
1 parent 4389cf9 commit 2c524f1
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion arrow/src/pyarrow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -363,13 +363,19 @@ impl FromPyArrow for RecordBatch {

let schema_ptr = unsafe { schema_capsule.reference::<FFI_ArrowSchema>() };
let ffi_array = unsafe { FFI_ArrowArray::from_raw(array_capsule.pointer().cast()) };
let array_data = unsafe { ffi::from_ffi(ffi_array, schema_ptr) }.map_err(to_py_err)?;
let mut array_data =
unsafe { ffi::from_ffi(ffi_array, schema_ptr) }.map_err(to_py_err)?;
if !matches!(array_data.data_type(), DataType::Struct(_)) {
return Err(PyTypeError::new_err(
"Expected Struct type from __arrow_c_array.",
));
}
let options = RecordBatchOptions::default().with_row_count(Some(array_data.len()));
// Ensure data is aligned (by potentially copying the buffers).
// This is needed because some python code (for example the
// python flight client) produces unaligned buffers
// See https://github.com/apache/arrow/issues/43552 for details
array_data.align_buffers();
let array = StructArray::from(array_data);
// StructArray does not embed metadata from schema. We need to override
// the output schema with the schema from the capsule.
Expand Down

0 comments on commit 2c524f1

Please sign in to comment.