Add option for arrow-ipc reader to not get record batch with all buffers shared

[rluvaton]

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

    #[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<Schema> {
        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<Schema>, 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<Schema>, 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<Option<&str>> = (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::<Vec<_>>());
            let float64 = Float64Array::from((0..n).map(|_| Some(1.5)).collect::<Vec<_>>());
            let boolean = BooleanArray::from((0..n).map(|_| Some(true)).collect::<Vec<_>>());

            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::<Vec<_>>());
            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<Schema>,
        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:?}");
        }
    }

[ariel-miculas]

Issue observed in practice due to a single column sharing the entire RecordBatch instead of per-column Buffers: apache/datafusion#23604

[alamb]

Maybe you could just (deep) copy the RecordBatch? That would ensure the allocations are small

[rluvaton]

Yes, it is possible but add an unnecessary overhead when we can avoid sharing in the first place,
Think about 500MB record batch (we saw that in a real world workload)

You will be copying way too much

[alamb]

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

If we can figure out how to read the arrow buffers into their own allocations that does seem like a better way to go

However, if we have to read the data and then read it again in the IPC reader into multiple buffers, that seems like it would be slower (as it always forces a copy)

[EmilyMatt]

I believe a deep-copy of the data into the read column should be mandatory.
I see no other way of ensuring a small or even self-contained memory footprint for arrays otherwise.
This should be a declared part of the reader imo

[alamb]

I believe a deep-copy of the data into the read column should be mandatory. I see no other way of ensuring a small or even self-contained memory footprint for arrays otherwise. This should be a declared part of the reader imo

yes I agree

Maybe we can add some configuration option on the IPC reader that forces a copy to occur

Note I think @jayzhan211 basically did something like this downstream in datafusion as part of

• fix: frame IPC messages when reading spill files so batches pin only their own bytes datafusion#24592

Maybe we could port that to arrow-rs (along with a flag 🤔 )

[yigitcan-ozturk]

I’d like to work on this if the scope is still available.

Based on the discussion, I’d keep the first PR narrowly focused on an opt-in IPC reader configuration that copies uncompressed body buffers into independent allocations, preserving the existing zero-copy behavior by default. I’ll use the downstream DataFusion work in apache/datafusion#24592 as a reference, but keep the Arrow implementation at the IPC-reader layer rather than introducing DataFusion-specific behavior.

I’d add regression coverage that verifies (1) the default path still shares the backing body allocation, (2) the opt-in path produces independently owned buffers for projected/retained columns, and (3) compressed IPC behavior and round-trip correctness remain unchanged.

If that direction matches what you had in mind, I’m happy to take it through implementation and review.