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

    #[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:?}");
        }
    }