Skip to content

Commit

Permalink
Optionally require alignment when reading IPC, respect alignment when…
Browse files Browse the repository at this point in the history
… writing (#5554)

* fix

* Remove redundant alignment check

* fix typo

* Add comment about randomized testing

* add doc comment on enforce_zero_copy

* address alamb feedback

* be explicit

* fix unit tests

* fix arrow-flight tests

* clippy

* hide api change + add require_alignment to StreamDecoder too

* Preserve docs

---------

Co-authored-by: Raphael Taylor-Davies <[email protected]>
  • Loading branch information
hzuo and tustvold authored Apr 4, 2024
1 parent 6306df0 commit eddef43
Show file tree
Hide file tree
Showing 5 changed files with 431 additions and 76 deletions.
2 changes: 1 addition & 1 deletion arrow-buffer/src/buffer/scalar.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ impl<T: ArrowNativeType> ScalarBuffer<T> {
/// This method will panic if
///
/// * `offset` or `len` would result in overflow
/// * `buffer` is not aligned to a multiple of `std::mem::size_of::<T>`
/// * `buffer` is not aligned to a multiple of `std::mem::align_of::<T>`
/// * `bytes` is not large enough for the requested slice
pub fn new(buffer: Buffer, offset: usize, len: usize) -> Self {
let size = std::mem::size_of::<T>();
Expand Down
6 changes: 5 additions & 1 deletion arrow-flight/src/encode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -627,6 +627,7 @@ mod tests {
use arrow_array::{cast::downcast_array, types::*};
use arrow_buffer::Buffer;
use arrow_cast::pretty::pretty_format_batches;
use arrow_ipc::MetadataVersion;
use arrow_schema::UnionMode;
use std::collections::HashMap;

Expand All @@ -638,7 +639,8 @@ mod tests {
/// ensure only the batch's used data (not the allocated data) is sent
/// <https://github.com/apache/arrow-rs/issues/208>
fn test_encode_flight_data() {
let options = IpcWriteOptions::default();
// use 8-byte alignment - default alignment is 64 which produces bigger ipc data
let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap();
let c1 = UInt32Array::from(vec![1, 2, 3, 4, 5, 6]);

let batch = RecordBatch::try_from_iter(vec![("a", Arc::new(c1) as ArrayRef)])
Expand Down Expand Up @@ -1343,6 +1345,8 @@ mod tests {

let mut stream = FlightDataEncoderBuilder::new()
.with_max_flight_data_size(max_flight_data_size)
// use 8-byte alignment - default alignment is 64 which produces bigger ipc data
.with_options(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap())
.build(futures::stream::iter([Ok(batch.clone())]));

let mut i = 0;
Expand Down
Loading

0 comments on commit eddef43

Please sign in to comment.