diff --git a/arrow-flight/src/encode.rs b/arrow-flight/src/encode.rs index a33b0e2d9a7e..4cd0b6269b6d 100644 --- a/arrow-flight/src/encode.rs +++ b/arrow-flight/src/encode.rs @@ -675,7 +675,7 @@ fn split_batch_for_grpc_response( let size = batch .columns() .iter() - .map(|col| col.get_buffer_memory_size()) + .map(|col| col.to_data().get_slice_memory_size().unwrap()) .sum::(); let n_batches = @@ -794,8 +794,9 @@ mod tests { }; use arrow_array::*; use arrow_array::{cast::downcast_array, types::*}; - use arrow_buffer::ScalarBuffer; + use arrow_buffer::{Buffer, MutableBuffer, ScalarBuffer}; use arrow_cast::pretty::pretty_format_batches; + use arrow_data::ArrayData; use arrow_ipc::{CompressionType, MetadataVersion}; use arrow_schema::{UnionFields, UnionMode}; use builder::MapBuilder; @@ -2153,4 +2154,194 @@ mod tests { "Specified overage was too high" ); } + + // Helper function to write data into a MutableBuffer and return it as a Buffer + fn create_buffer_from_data(data: Vec) -> Buffer { + let mut mbuf = MutableBuffer::new(data.len()); + for v in data { + mbuf.push(v); + } + mbuf.into() + } + + // Helper function to build arrays from a Buffer + fn build_array_from_buffer(buffer: Buffer, data_type: DataType, len: usize) -> ArrayRef { + let array_data = ArrayData::builder(data_type.clone()) + .len(len) + .add_buffer(buffer) + .build() + .expect("failed to build array data"); + + match data_type { + DataType::UInt32 => Arc::new(UInt32Array::from(array_data)), + DataType::Int64 => Arc::new(Int64Array::from(array_data)), + DataType::Float32 => Arc::new(Float32Array::from(array_data)), + DataType::Int8 => Arc::new(Int8Array::from(array_data)), + _ => panic!("Unsupported data type"), + } + } + + #[test] + fn test_encode_decode_reencode_flight_data() { + const ROW_COUNT: usize = 5_000_000; + const SPLIT_SIZE: usize = 4; + + let int_buf = create_buffer_from_data(vec![1u8; ROW_COUNT * std::mem::size_of::()]); + let long_buf = create_buffer_from_data(vec![2u8; ROW_COUNT * std::mem::size_of::()]); + let float_buf = create_buffer_from_data(vec![3u8; ROW_COUNT * std::mem::size_of::()]); + // Int8 buffer: one byte per row + let int8_buf = create_buffer_from_data(vec![0u8; ROW_COUNT * std::mem::size_of::()]); + + // Short arrays / batch + let int_array_short = + build_array_from_buffer(int_buf.clone(), DataType::UInt32, ROW_COUNT / SPLIT_SIZE); + let long_array_short = + build_array_from_buffer(long_buf.clone(), DataType::Int64, ROW_COUNT / SPLIT_SIZE); + let float_array_short = + build_array_from_buffer(float_buf.clone(), DataType::Float32, ROW_COUNT / SPLIT_SIZE); + let int8_array_short = + build_array_from_buffer(int8_buf.clone(), DataType::Int8, ROW_COUNT / SPLIT_SIZE); + + let batch_short = RecordBatch::try_from_iter(vec![ + ("int_col", int_array_short), + ("long_col", long_array_short), + ("float_col", float_array_short), + ("int8_col", int8_array_short), + ]) + .expect("failed to create short record batch"); + + // Build long arrays: set each column to ROW_COUNT rows + let int_array_long = build_array_from_buffer(int_buf, DataType::UInt32, ROW_COUNT); + let long_array_long = build_array_from_buffer(long_buf, DataType::Int64, ROW_COUNT); + let float_array_long = build_array_from_buffer(float_buf, DataType::Float32, ROW_COUNT); + let int8_array_long = build_array_from_buffer(int8_buf, DataType::Int8, ROW_COUNT); + + let batch_long = RecordBatch::try_from_iter(vec![ + ("int_col", int_array_long), + ("long_col", long_array_long), + ("float_col", float_array_long), + ("int8_col", int8_array_long), + ]) + .expect("failed to create long record batch"); + + let split_short = split_batch_for_grpc_response(batch_short.clone(), 2 * 1024 * 1024); + let split_long = split_batch_for_grpc_response(batch_long.clone(), 2 * 1024 * 1024); + assert!(split_short.len() != split_long.len()); // more data split divided by the same max size should result in more sub-slices + + // Compare memory calculations for short and long batches + let total_buffer_memory_short = batch_short + .columns() + .iter() + .map(|col| col.get_buffer_memory_size()) + .sum::(); + let total_slice_memory_short = batch_short + .columns() + .iter() + .map(|col| col.to_data().get_slice_memory_size().unwrap()) + .sum::(); + + let total_buffer_memory_long = batch_long + .columns() + .iter() + .map(|col| col.get_buffer_memory_size()) + .sum::(); + let total_slice_memory_long = batch_long + .columns() + .iter() + .map(|col| col.to_data().get_slice_memory_size().unwrap()) + .sum::(); + assert!(total_buffer_memory_short > total_slice_memory_short); // 85 mb > 21mb + assert_eq!(total_buffer_memory_long, total_slice_memory_long); // 85 MB == 85 mb + } + + async fn collect_decoded_batches(decoder: &mut FlightDataDecoder) -> Vec { + let mut batches = vec![]; + while let Some(decoded) = decoder.next().await { + if let DecodedPayload::RecordBatch(b) = decoded.unwrap().payload { + batches.push(b); + } + } + batches + } + + /// https://github.com/apache/arrow-rs/issues/9388 + /// encode → decode → re-encode must not produce more FlightData messages than the initial encode. + #[tokio::test] + async fn test_roundtrip_encode_decode_reencode_no_extra_flight_data() { + let batch = RecordBatch::try_from_iter(vec![( + "data", + Arc::new(Int32Array::from_iter_values(0..1_000_000)) as ArrayRef, + )]) + .unwrap(); + + let max_size = 256 * 1024; // 256KB forces ~16 splits across the 4MB batch + + let encoder = FlightDataEncoderBuilder::new() + .with_max_flight_data_size(max_size) + .build(futures::stream::iter(vec![Ok(batch.clone())])); + let encoded: Vec = encoder.map(|r| r.unwrap()).collect().await; + let encoded_count = encoded.len(); + assert!( + encoded_count > 2, + "batch must split for this test to be meaningful" + ); + + let mut decoder = + FlightDataDecoder::new(futures::stream::iter(encoded.into_iter().map(Ok))); + let decoded_batches = collect_decoded_batches(&mut decoder).await; + + let reencoded_count = FlightDataEncoderBuilder::new() + .with_max_flight_data_size(max_size) + .build(futures::stream::iter(decoded_batches.into_iter().map(Ok))) + .map(|r| r.unwrap()) + .collect::>() + .await + .len(); + + assert!( + reencoded_count <= encoded_count, + "re-encode produced {reencoded_count} messages but initial encode produced {encoded_count}" + ); + } + + /// A zero-copy slice shares the original backing buffer, so get_buffer_memory_size() sees + /// the full allocation regardless of how many rows the slice covers. The encoder must use + /// get_slice_memory_size() instead, otherwise a 4MB slice of a 16MB batch produces as many + /// FlightData messages as the full 16MB batch. + #[tokio::test] + async fn test_slice_encodes_fewer_flight_data_than_full_buffer() { + const FULL_ROWS: usize = 4_000_000; // 16MB + const SLICE_ROWS: usize = FULL_ROWS / 4; // 4MB + + let full_batch = RecordBatch::try_from_iter(vec![( + "col", + Arc::new(UInt32Array::from_iter_values(0..FULL_ROWS as u32)) as ArrayRef, + )]) + .unwrap(); + let slice_batch = full_batch.slice(0, SLICE_ROWS); + + let max_size = 512 * 1024; // 512KB + + let full_count = FlightDataEncoderBuilder::new() + .with_max_flight_data_size(max_size) + .build(futures::stream::iter(vec![Ok(full_batch)])) + .map(|r| r.unwrap()) + .collect::>() + .await + .len(); + + let slice_count = FlightDataEncoderBuilder::new() + .with_max_flight_data_size(max_size) + .build(futures::stream::iter(vec![Ok(slice_batch)])) + .map(|r| r.unwrap()) + .collect::>() + .await + .len(); + + assert!( + slice_count < full_count, + "slice ({SLICE_ROWS} rows) produced {slice_count} messages, same as full batch \ + ({FULL_ROWS} rows, {full_count} messages) — encoder is accounting for the entire backing buffer" + ); + } }