Add option for arrow-ipc reader to not get record batch with all buffers shared
- Dominant language
- Rust
- Stars
- 3.6k
- Forks
- 1.3k
- Avg merge
- 2d 18h
- Merged PRs (30d)
- 169
Description
### Is your feature request related to a problem or challenge?
Yes, having all columns in a record batch share the same underlying buffer reduce the chance of some columns to be drop
lets say we have this:
```
Project: a, b + c + d + e
Scan Arrow IPC: a, b, c, d, e
```
because the entire batch from Arrow IPC is a single underlying buffer, even though the output of project only need to have keep in memory the original column `a` the entire batch (a, b, c, d, e) columns are kept in memory
### Describe the solution you'd like
I would like to be able to have an option to disable reading Arrow IPC in flat buffer, so each column/null buffer/offsets/whatever in the batch is different allocation that can be free'd independently
### Describe alternatives you've considered
_No response_
### Additional context
Tests to show that
```rust
#[test]
fn test_write_read_two_batches_all_nullable_types_uncompressed() {
write_read_two_batches_all_nullable_types(crate::writer::IpcWriteOptions::default());
}
#[cfg(feature = "lz4")]
#[test]
fn test_write_read_two_batches_all_nullable_types_lz4() {
let options = crate::writer::IpcWriteOptions::default()
.try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
.unwrap();
write_read_two_batches_all_nullable_types(options);
}
// Schema: 6 nullable columns — string, 2 primitive, boolean, list, struct(string, primitive).
fn all_nullable_types_schema() -> Arc {
let struct_fields = Fields::from(vec![
Field::new("s", DataType::Utf8, true),
Field::new("n", DataType::Int32, true),
]);
Arc::new(Schema::new(vec![
Field::new("string", DataType::Utf8, true),
Field::new("int32", DataType::Int32, true),
Field::new("float64", DataType::Float64, true),
Field::new("boolean", DataType::Boolean, true),
Field::new(
"list",
DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
true,
),
Field::new("struct", DataType::Struct(struct_fields), true),
]))
}
fn write_read_two_batches_all_nullable_types(options: crate::writer::IpcWriteOptions) {
use arrow_array::builder::{Int32Builder, ListBuilder};
// Build one batch: 6 nullable columns with nulls sprinkled in each.
fn make_batch(schema: &Arc, base: i32) -> RecordBatch {
// string
let string = StringArray::from(vec![Some("a"), None, Some("c")]);
// 2 primitive
let int32 = Int32Array::from(vec![Some(base), None, Some(base + 2)]);
let float64 = Float64Array::from(vec![None, Some(base as f64 + 0.5), Some(1.0)]);
// boolean
let boolean = BooleanArray::from(vec![Some(true), Some(false), None]);
// list of int32
let list = {
let mut b = ListBuilder::new(Int32Builder::new());
b.values().append_value(1);
b.values().append_value(2);
b.append(true);
b.append(false); // null list
b.values().append_value(3);
b.append(true);
b.finish()
};
// struct of string + primitive
let struct_string = StringArray::from(vec![Some("x"), None, Some("z")]);
let struct_int = Int32Array::from(vec![None, Some(10), Some(20)]);
let struct_fields: Fields = match schema.field(5).data_type() {
DataType::Struct(f) => f.clone(),
_ => unreachable!(),
};
let struct_array = StructArray::new(
struct_fields,
vec![Arc::new(struct_string) as ArrayRef, Arc::new(struct_int)],
Some(NullBuffer::from(vec![true, true, false])), // null struct row
);
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(string),
Arc::new(int32),
Arc::new(float64),
Arc::new(boolean),
Arc::new(list),
Arc::new(struct_array),
],
)
.unwrap()
}
let schema = all_nullable_types_schema();
let batches = [make_batch(&schema, 0), make_batch(&schema, 100)];
// Buffers are too small to actually compress, so they stay uncompressed in the body
// and remain zero-copy slices of one shared allocation.
write_read_and_check_allocation(&schema, &batches, options, true);
}
// Same schema, but the string columns hold highly repetitive data (many identical long
// values) so LZ4 has plenty to compress — verifies compressed round-trip + buffer sharing.
#[cfg(feature = "lz4")]
#[test]
fn test_write_read_repetitive_strings_lz4() {
use arrow_array::builder::{Int32Builder, ListBuilder};
fn make_batch(schema: &Arc, base: i32) -> RecordBatch {
let n = 1024;
// Highly repetitive: same long string on most rows, occasional null.
let repeated = "the-quick-brown-fox-jumps-over-the-lazy-dog";
let strings: Vec> = (0..n)
.map(|i| if i % 8 == 0 { None } else { Some(repeated) })
.collect();
let string = StringArray::from(strings.clone());
let int32 = Int32Array::from((0..n).map(|_| Some(base)).collect::>());
let float64 = Float64Array::from((0..n).map(|_| Some(1.5)).collect::>());
let boolean = BooleanArray::from((0..n).map(|_| Some(true)).collect::>());
let list = {
let mut b = ListBuilder::new(Int32Builder::new());
for _ in 0..n {
b.values().append_value(7);
b.values().append_value(7);
b.append(true);
}
b.finish()
};
let struct_string = StringArray::from(strings);
let struct_int = Int32Array::from((0..n).map(|_| Some(base)).collect::>());
let struct_fields: Fields = match schema.field(5).data_type() {
DataType::Struct(f) => f.clone(),
_ => unreachable!(),
};
let struct_array = StructArray::new(
struct_fields,
vec![Arc::new(struct_string) as ArrayRef, Arc::new(struct_int)],
None,
);
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(string),
Arc::new(int32),
Arc::new(float64),
Arc::new(boolean),
Arc::new(list),
Arc::new(struct_array),
],
)
.unwrap()
}
let schema = all_nullable_types_schema();
let batches = [make_batch(&schema, 0), make_batch(&schema, 100)];
let options = crate::writer::IpcWriteOptions::default()
.try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
.unwrap();
// These buffers are large and highly repetitive, so LZ4 actually compresses them.
// On read each compressed buffer is decompressed into its own fresh allocation, so the
// buffers do NOT share one body allocation.
write_read_and_check_allocation(&schema, &batches, options, false);
}
fn write_read_and_check_allocation(
schema: &Arc,
batches: &[RecordBatch],
options: crate::writer::IpcWriteOptions,
expect_all_buffers_shared: bool,
) {
let mut buf = Vec::new();
{
let mut writer =
crate::writer::FileWriter::try_new_with_options(&mut buf, schema, options).unwrap();
for b in batches {
writer.write(b).unwrap();
}
writer.finish().unwrap();
}
let reader = FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
let read_batches: Vec<_> = reader.map(|b| b.unwrap()).collect();
assert_eq!(read_batches.as_slice(), batches);
// Walk the first read-back batch and print every buffer's data_ptr plus its role.
fn buffer_roles(dt: &DataType) -> Vec<&'static str> {
match dt {
DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
vec!["offsets", "values (bytes)"]
}
DataType::List(_) | DataType::LargeList(_) => vec!["offsets"],
DataType::Boolean => vec!["values (bits)"],
DataType::Struct(_) => vec![], // no data buffers, only child arrays
_ => vec!["values"], // primitives
}
}
fn print_buffers(name: &str, data: &ArrayData, indent: usize) {
let pad = " ".repeat(indent);
println!("{pad}{name}: {:?}", data.data_type());
if let Some(nulls) = data.nulls() {
let b = nulls.buffer();
println!(
"{pad} null buffer data_ptr = {:?} ptr = {:?}",
b.data_ptr(),
b.as_ptr()
);
}
let roles = buffer_roles(data.data_type());
for (i, b) in data.buffers().iter().enumerate() {
let role = roles.get(i).copied().unwrap_or("?");
println!(
"{pad} buffer[{i}] {role:<15} data_ptr = {:?} ptr = {:?}",
b.data_ptr(),
b.as_ptr()
);
}
for (i, child) in data.child_data().iter().enumerate() {
print_buffers(&format!("child[{i}]"), child, indent + 1);
}
}
// Collect every buffer's data_ptr (allocation base, offset-independent) recursively.
fn collect_data_ptrs(data: &ArrayData, out: &mut Vec<*const u8>) {
if let Some(nulls) = data.nulls() {
out.push(nulls.buffer().data_ptr().as_ptr());
}
for b in data.buffers() {
out.push(b.data_ptr().as_ptr());
}
for child in data.child_data() {
collect_data_ptrs(child, out);
}
}
let first = &read_batches[0];
println!("=== buffers of first read-back batch ===");
let mut ptrs = Vec::new();
for (idx, col) in first.columns().iter().enumerate() {
let field = first.schema().field(idx).name().clone();
print_buffers(&field, &col.to_data(), 0);
collect_data_ptrs(&col.to_data(), &mut ptrs);
}
let all_shared = ptrs.windows(2).all(|w| w[0] == w[1]);
if expect_all_buffers_shared {
// Uncompressed buffers are zero-copy slices into one shared body allocation.
assert!(all_shared, "expected all buffers to share one allocation, got {ptrs:?}");
} else {
// Compressed buffers are each decompressed into their own allocation.
assert!(!all_shared, "expected buffers in separate allocations, got {ptrs:?}");
}
}
```
Contributor guide
Research direction
Start with the Arrow IPC FileReader and the write_read_and_check_allocation test helper shown in the issue. Trace how uncompressed buffers are read and how buffer allocation sharing is currently verified. Done means an option controls independent buffer allocations without breaking the round-trip assertions, with tests covering the shared and separate allocation cases.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- data
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 58/100