From b73bfd7aa5695fe77ffe90a5aecea4eefbfd2c51 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Sun, 19 Jul 2026 12:45:00 -0400 Subject: [PATCH 1/5] creating unit test to support switch to get_slice_memory_Size() --- arrow-flight/src/encode.rs | 116 ++++++++++++++++++++++++++++++++++++- 1 file changed, 115 insertions(+), 1 deletion(-) diff --git a/arrow-flight/src/encode.rs b/arrow-flight/src/encode.rs index a33b0e2d9a7e..5347168e8f6d 100644 --- a/arrow-flight/src/encode.rs +++ b/arrow-flight/src/encode.rs @@ -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,117 @@ 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() { + // Keep row counts fixed in the sink: 250,000 rows per column. + + const ROW_COUNT: usize = 250_000; + // int32 memory: 250_000 * 4 = 1_000_000 bytes + + // Create buffers sized for the desired row counts (memory size itself isn't important) + 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 + // Build short arrays: set each column to ROW_COUNT rows + let int_array_short = + build_array_from_buffer(int_buf.clone(), DataType::UInt32, ROW_COUNT / 2); + let long_array_short = + build_array_from_buffer(long_buf.clone(), DataType::Int64, ROW_COUNT / 2); + let float_array_short = + build_array_from_buffer(float_buf.clone(), DataType::Float32, ROW_COUNT / 2); + let int8_array_short = + build_array_from_buffer(int8_buf.clone(), DataType::Int8, ROW_COUNT / 2); + + 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"); + + // Long arrays / batch + // Build long arrays: also 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); + println!("Short split batches count: {}", split_short.len()); + for (i, b) in split_short.iter().enumerate() { + println!("Short batch {}: {} rows", i, b.num_rows()); + } + let split_long = split_batch_for_grpc_response(batch_long.clone(), 2 * 1024 * 1024); + println!("Long split batches count: {}", split_long.len()); + for (i, b) in split_long.iter().enumerate() { + println!("long batch {}: {} rows", i, b.num_rows()); + }*/ + + // Compare memory calculations for short and long batches (run after splitting) + let buf_sum_short = batch_short + .columns() + .iter() + .map(|col| col.get_buffer_memory_size()) + .sum::(); + let slice_sum_short = batch_short + .columns() + .iter() + .map(|col| col.to_data().get_slice_memory_size().unwrap()) + .sum::(); + println!("Short batch buffer-sum: {} bytes", buf_sum_short); + println!("Short batch slice-sum: {} bytes", slice_sum_short); + + let buf_sum_long = batch_long + .columns() + .iter() + .map(|col| col.get_buffer_memory_size()) + .sum::(); + let slice_sum_long = batch_long + .columns() + .iter() + .map(|col| col.to_data().get_slice_memory_size().unwrap()) + .sum::(); + println!("Long batch buffer-sum: {} bytes", buf_sum_long); + println!("Long batch slice-sum: {} bytes", slice_sum_long); + } } From 73d19611babb2a2b892a1d0802d286ccaaf66693 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Sun, 19 Jul 2026 17:35:52 -0400 Subject: [PATCH 2/5] add comments for test --- arrow-flight/src/encode.rs | 41 +++++++++++++++++--------------------- 1 file changed, 18 insertions(+), 23 deletions(-) diff --git a/arrow-flight/src/encode.rs b/arrow-flight/src/encode.rs index 5347168e8f6d..d5d5d8cf6c27 100644 --- a/arrow-flight/src/encode.rs +++ b/arrow-flight/src/encode.rs @@ -2183,12 +2183,11 @@ mod tests { #[test] fn test_encode_decode_reencode_flight_data() { - // Keep row counts fixed in the sink: 250,000 rows per column. + const ROW_COUNT: usize = 5_000_000; + // int32 memory: 5_000_000 * 4 = 20_000_000 bytes - const ROW_COUNT: usize = 250_000; - // int32 memory: 250_000 * 4 = 1_000_000 bytes - - // Create buffers sized for the desired row counts (memory size itself isn't important) + const SPLIT_SIZE: usize = 4; + // Create buffers sized for the desired row counts 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::()]); @@ -2196,15 +2195,14 @@ mod tests { let int8_buf = create_buffer_from_data(vec![0u8; ROW_COUNT * std::mem::size_of::()]); // Short arrays / batch - // Build short arrays: set each column to ROW_COUNT rows let int_array_short = - build_array_from_buffer(int_buf.clone(), DataType::UInt32, ROW_COUNT / 2); + 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 / 2); + 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 / 2); + 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 / 2); + 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), @@ -2214,8 +2212,7 @@ mod tests { ]) .expect("failed to create short record batch"); - // Long arrays / batch - // Build long arrays: also set each column to ROW_COUNT rows + // 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); @@ -2229,7 +2226,7 @@ mod tests { ]) .expect("failed to create long record batch"); - /*let split_short = split_batch_for_grpc_response(batch_short.clone(), 2 * 1024 * 1024); + let split_short = split_batch_for_grpc_response(batch_short.clone(), 2 * 1024 * 1024); println!("Short split batches count: {}", split_short.len()); for (i, b) in split_short.iter().enumerate() { println!("Short batch {}: {} rows", i, b.num_rows()); @@ -2238,33 +2235,31 @@ mod tests { println!("Long split batches count: {}", split_long.len()); for (i, b) in split_long.iter().enumerate() { println!("long batch {}: {} rows", i, b.num_rows()); - }*/ + } - // Compare memory calculations for short and long batches (run after splitting) - let buf_sum_short = batch_short + // 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 slice_sum_short = batch_short + let total_slice_memory_short = batch_short .columns() .iter() .map(|col| col.to_data().get_slice_memory_size().unwrap()) .sum::(); - println!("Short batch buffer-sum: {} bytes", buf_sum_short); - println!("Short batch slice-sum: {} bytes", slice_sum_short); - let buf_sum_long = batch_long + let total_buffer_memory_long = batch_long .columns() .iter() .map(|col| col.get_buffer_memory_size()) .sum::(); - let slice_sum_long = batch_long + let total_slice_memory_long = batch_long .columns() .iter() .map(|col| col.to_data().get_slice_memory_size().unwrap()) .sum::(); - println!("Long batch buffer-sum: {} bytes", buf_sum_long); - println!("Long batch slice-sum: {} bytes", slice_sum_long); + 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 } } From f15dfe60c48ef3bc0f880c16a2ddfa0f31f562d0 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Sun, 19 Jul 2026 17:39:39 -0400 Subject: [PATCH 3/5] replace get_Buffer_memory_size for get_slice_memory --- arrow-flight/src/encode.rs | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/arrow-flight/src/encode.rs b/arrow-flight/src/encode.rs index d5d5d8cf6c27..d45dbd8f6b76 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 = @@ -2184,10 +2184,8 @@ mod tests { #[test] fn test_encode_decode_reencode_flight_data() { const ROW_COUNT: usize = 5_000_000; - // int32 memory: 5_000_000 * 4 = 20_000_000 bytes - const SPLIT_SIZE: usize = 4; - // Create buffers sized for the desired row counts + 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::()]); @@ -2227,15 +2225,8 @@ mod tests { .expect("failed to create long record batch"); let split_short = split_batch_for_grpc_response(batch_short.clone(), 2 * 1024 * 1024); - println!("Short split batches count: {}", split_short.len()); - for (i, b) in split_short.iter().enumerate() { - println!("Short batch {}: {} rows", i, b.num_rows()); - } let split_long = split_batch_for_grpc_response(batch_long.clone(), 2 * 1024 * 1024); - println!("Long split batches count: {}", split_long.len()); - for (i, b) in split_long.iter().enumerate() { - println!("long batch {}: {} rows", i, b.num_rows()); - } + 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 From eed8d454795b969668c33739e41b365c23a1da81 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Fri, 24 Jul 2026 15:53:25 -0400 Subject: [PATCH 4/5] add more test --- arrow-flight/src/encode.rs | 87 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) diff --git a/arrow-flight/src/encode.rs b/arrow-flight/src/encode.rs index d45dbd8f6b76..deb46a9250bc 100644 --- a/arrow-flight/src/encode.rs +++ b/arrow-flight/src/encode.rs @@ -2253,4 +2253,91 @@ mod tests { 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 + } + + /// 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() — 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" + ); + } } From 8500b43f61595b69b5a95a5faa28c6feeb62e580 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Fri, 24 Jul 2026 15:58:27 -0400 Subject: [PATCH 5/5] fix --- arrow-flight/src/encode.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/arrow-flight/src/encode.rs b/arrow-flight/src/encode.rs index deb46a9250bc..4cd0b6269b6d 100644 --- a/arrow-flight/src/encode.rs +++ b/arrow-flight/src/encode.rs @@ -2264,6 +2264,7 @@ mod tests { 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() { @@ -2280,7 +2281,10 @@ mod tests { .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"); + 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))); @@ -2302,7 +2306,7 @@ mod tests { /// 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() — otherwise a 4MB slice of a 16MB batch produces as many + /// 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() {