1use arrow::array::ArrayData;
38use arrow::buffer::Buffer as ArrowBuffer;
39use arrow::ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
40use arrow::ipc::{Buffer as IpcBuffer, FieldNode, MessageHeader, MetadataVersion};
41use arrow_schema::{DataType, Field, Schema};
42use eyre::{Context, bail, eyre};
43
44use super::ARROW_BUFFER_ALIGNMENT as ALIGN;
45
46const CONTINUATION_MARKER: [u8; 4] = [0xff, 0xff, 0xff, 0xff];
48const PREFIX_LEN: usize = 8;
50
51#[inline]
52fn round_up(n: usize, align: usize) -> usize {
53 debug_assert!(align.is_power_of_two());
54 (n + align - 1) & !(align - 1)
55}
56
57fn is_fast_path_type(data_type: &DataType) -> bool {
61 use DataType::*;
62 matches!(
63 data_type,
64 Null | Boolean
65 | Int8
66 | Int16
67 | Int32
68 | Int64
69 | UInt8
70 | UInt16
71 | UInt32
72 | UInt64
73 | Float16
74 | Float32
75 | Float64
76 | Timestamp(_, _)
77 | Date32
78 | Date64
79 | Time32(_)
80 | Time64(_)
81 | Duration(_)
82 | Interval(_)
83 | Decimal128(_, _)
84 | Decimal256(_, _)
85 | FixedSizeBinary(_)
86 | Binary
87 | LargeBinary
88 | Utf8
89 | LargeUtf8
90 | List(_)
91 | LargeList(_)
92 | FixedSizeList(_, _)
93 | Struct(_)
94 )
95}
96
97struct Desc {
100 offset: usize,
101 len: usize,
102 src: BufferSrc,
103}
104
105enum BufferSrc {
106 Bytes(ArrowBuffer, usize),
108 AllOnes,
110}
111
112#[derive(Default)]
113struct Layout {
114 nodes: Vec<FieldNode>,
115 ipc_buffers: Vec<IpcBuffer>,
116 descs: Vec<Desc>,
117 body_len: usize,
118}
119
120impl Layout {
121 fn push_buffer(&mut self, off: &mut usize, len: usize, src: BufferSrc) -> Option<()> {
126 let padded = len.checked_add(ALIGN - 1).map(|n| n & !(ALIGN - 1))?;
128 let next = off.checked_add(padded)?;
129 if next > super::MAX_IPC_BYTES {
130 return None;
131 }
132 self.ipc_buffers
133 .push(IpcBuffer::new(*off as i64, len as i64));
134 self.descs.push(Desc {
135 offset: *off,
136 len,
137 src,
138 });
139 *off = next;
140 Some(())
141 }
142}
143
144fn build_layout(array: &ArrayData) -> Option<Layout> {
148 let mut layout = Layout::default();
149 let mut off = 0usize;
150 build_layout_rec(array, &mut layout, &mut off)?;
151 layout.body_len = off;
152 Some(layout)
153}
154
155fn build_layout_rec(array: &ArrayData, layout: &mut Layout, off: &mut usize) -> Option<()> {
156 let data_type = array.data_type();
157 if !is_fast_path_type(data_type) || array.offset() != 0 {
158 return None;
159 }
160
161 let len = array.len();
162 let null_count = if matches!(data_type, DataType::Null) {
165 len
166 } else {
167 array.null_count()
168 };
169 layout
170 .nodes
171 .push(FieldNode::new(len as i64, null_count as i64));
172
173 if !matches!(data_type, DataType::Null) {
176 match array.nulls() {
177 Some(nulls) => {
178 if nulls.inner().offset() != 0 {
179 return None;
180 }
181 let sliced = nulls.inner().sliced();
182 let bytes = sliced.len();
183 layout.push_buffer(off, bytes, BufferSrc::Bytes(sliced, 0))?;
184 }
185 None => {
186 let bytes = len.div_ceil(8);
187 layout.push_buffer(off, bytes, BufferSrc::AllOnes)?;
188 }
189 }
190 }
191
192 for buffer in array.buffers() {
194 layout.push_buffer(off, buffer.len(), BufferSrc::Bytes(buffer.clone(), 0))?;
195 }
196
197 match data_type {
199 DataType::FixedSizeList(_, value_size) => {
200 let n = len.checked_mul(*value_size as usize)?;
203 let child = array.child_data().first()?;
204 if child.len() < n {
205 return None;
206 }
207 build_layout_rec(&child.slice(0, n), layout, off)?;
208 }
209 DataType::Struct(_) => {
210 for child in array.child_data() {
216 if child.len() < len {
217 return None;
218 }
219 build_layout_rec(&child.slice(0, len), layout, off)?;
220 }
221 }
222 _ => {
223 for child in array.child_data() {
227 build_layout_rec(child, layout, off)?;
228 }
229 }
230 }
231
232 Some(())
233}
234
235struct Prepared {
237 layout: Layout,
238 schema_message: Vec<u8>,
239 record_batch_message: Vec<u8>,
240 schema_block: usize,
241 record_batch_block: usize,
242 total: usize,
243}
244
245fn ipc_write_options() -> eyre::Result<IpcWriteOptions> {
246 IpcWriteOptions::try_new(ALIGN, false, MetadataVersion::V5)
247 .map_err(|e| eyre!("failed to build Arrow IPC write options: {e}"))
248}
249
250fn build_schema_message(data_type: &DataType) -> eyre::Result<Vec<u8>> {
253 let schema = Schema::new(vec![Field::new("data", data_type.clone(), true)]);
254 let options = ipc_write_options()?;
255 let mut tracker = DictionaryTracker::new(false);
256 let encoded = IpcDataGenerator {}.schema_to_bytes_with_dictionary_tracker(
257 &schema,
258 &mut tracker,
259 &options,
260 );
261 Ok(encoded.ipc_message)
262}
263
264fn build_record_batch_message(
267 num_rows: usize,
268 nodes: &[FieldNode],
269 buffers: &[IpcBuffer],
270 body_len: usize,
271) -> Vec<u8> {
272 use flatbuffers::FlatBufferBuilder;
273
274 let mut fbb = FlatBufferBuilder::new();
275 let buffers_fb = fbb.create_vector(buffers);
276 let nodes_fb = fbb.create_vector(nodes);
277
278 let record_batch = {
279 let mut builder = arrow::ipc::RecordBatchBuilder::new(&mut fbb);
280 builder.add_length(num_rows as i64);
281 builder.add_nodes(nodes_fb);
282 builder.add_buffers(buffers_fb);
283 builder.finish()
284 };
285
286 let message = {
287 let mut builder = arrow::ipc::MessageBuilder::new(&mut fbb);
288 builder.add_version(MetadataVersion::V5);
289 builder.add_header_type(MessageHeader::RecordBatch);
290 builder.add_bodyLength(body_len as i64);
291 builder.add_header(record_batch.as_union_value());
292 builder.finish()
293 };
294
295 fbb.finish(message, None);
296 fbb.finished_data().to_vec()
297}
298
299fn prepare(array: &ArrayData) -> Option<Prepared> {
300 let layout = build_layout(array)?;
301 let schema_message = build_schema_message(array.data_type()).ok()?;
302 let record_batch_message = build_record_batch_message(
303 array.len(),
304 &layout.nodes,
305 &layout.ipc_buffers,
306 layout.body_len,
307 );
308 let schema_block = round_up(PREFIX_LEN + schema_message.len(), ALIGN);
309 let record_batch_block = round_up(PREFIX_LEN + record_batch_message.len(), ALIGN);
310 let total = schema_block + record_batch_block + layout.body_len + PREFIX_LEN;
312 Some(Prepared {
313 layout,
314 schema_message,
315 record_batch_message,
316 schema_block,
317 record_batch_block,
318 total,
319 })
320}
321
322pub fn ipc_fast_path_len(array: &ArrayData) -> Option<usize> {
327 prepare(array).map(|p| p.total)
328}
329
330fn write_framed_message(dst: &mut [u8], at: usize, flatbuffer: &[u8]) -> usize {
335 let block = round_up(PREFIX_LEN + flatbuffer.len(), ALIGN);
336 let metadata_len = (block - PREFIX_LEN) as i32;
337 dst[at..at + 4].copy_from_slice(&CONTINUATION_MARKER);
338 dst[at + 4..at + 8].copy_from_slice(&metadata_len.to_le_bytes());
339 dst[at + 8..at + 8 + flatbuffer.len()].copy_from_slice(flatbuffer);
340 dst[at + 8 + flatbuffer.len()..at + block].fill(0);
342 block
343}
344
345pub fn encode_ipc_into(array: &ArrayData, dst: &mut [u8]) -> eyre::Result<()> {
352 let prepared =
353 prepare(array).ok_or_else(|| eyre!("array is not Arrow IPC fast-path eligible"))?;
354 if dst.len() != prepared.total {
355 bail!(
356 "destination size {} does not match required IPC length {}",
357 dst.len(),
358 prepared.total
359 );
360 }
361
362 let mut at = 0;
363 at += write_framed_message(dst, at, &prepared.schema_message);
364 debug_assert_eq!(at, prepared.schema_block);
365 at += write_framed_message(dst, at, &prepared.record_batch_message);
366 debug_assert_eq!(at, prepared.schema_block + prepared.record_batch_block);
367
368 let body_start = at;
369 write_body(dst, body_start, &prepared.layout);
370 at = body_start + prepared.layout.body_len;
371
372 dst[at..at + 4].copy_from_slice(&CONTINUATION_MARKER);
374 dst[at + 4..at + 8].copy_from_slice(&0i32.to_le_bytes());
375 debug_assert_eq!(at + PREFIX_LEN, prepared.total);
376
377 Ok(())
378}
379
380fn write_body(dst: &mut [u8], body_start: usize, layout: &Layout) {
384 for desc in &layout.descs {
385 let start = body_start + desc.offset;
386 let end = start + desc.len;
387 match &desc.src {
388 BufferSrc::Bytes(buffer, src_off) => {
389 dst[start..end].copy_from_slice(&buffer.as_slice()[*src_off..*src_off + desc.len]);
390 }
391 BufferSrc::AllOnes => dst[start..end].fill(0xff),
392 }
393 let padded_end = body_start + desc.offset + round_up(desc.len, ALIGN);
394 dst[end..padded_end].fill(0);
395 }
396}
397
398pub fn encode_schema_message(data_type: &DataType) -> eyre::Result<Vec<u8>> {
406 let schema_message = build_schema_message(data_type)?;
407 let block = round_up(PREFIX_LEN + schema_message.len(), ALIGN);
408 let mut dst = vec![0u8; block];
409 write_framed_message(&mut dst, 0, &schema_message);
410 Ok(dst)
411}
412
413pub fn batch_fast_path_len(array: &ArrayData) -> Option<usize> {
417 let prepared = prepare(array)?;
418 Some(prepared.record_batch_block + prepared.layout.body_len + PREFIX_LEN)
419}
420
421pub fn encode_batch_into(array: &ArrayData, dst: &mut [u8]) -> eyre::Result<()> {
428 let prepared =
429 prepare(array).ok_or_else(|| eyre!("array is not Arrow IPC fast-path eligible"))?;
430 let expected = prepared.record_batch_block + prepared.layout.body_len + PREFIX_LEN;
431 if dst.len() != expected {
432 bail!(
433 "destination size {} does not match required batch length {expected}",
434 dst.len(),
435 );
436 }
437 let at = write_framed_message(dst, 0, &prepared.record_batch_message);
438 debug_assert_eq!(at, prepared.record_batch_block);
439 write_body(dst, at, &prepared.layout);
440 let body_end = at + prepared.layout.body_len;
443 dst[body_end..body_end + 4].copy_from_slice(&CONTINUATION_MARKER);
444 dst[body_end + 4..body_end + PREFIX_LEN].copy_from_slice(&0i32.to_le_bytes());
445 Ok(())
446}
447
448pub fn encode_ipc_to_vec(array: &ArrayData) -> eyre::Result<Vec<u8>> {
452 super::encode_arrow_ipc(array).context("Arrow IPC fallback encode")
453}
454
455struct Uint8Layout {
459 schema_message: Vec<u8>,
460 record_batch_message: Vec<u8>,
461 validity_len: usize,
462 validity_padded: usize,
463 body_len: usize,
464 total: usize,
465 data_offset: usize,
466}
467
468fn uint8_layout(data_len: usize) -> eyre::Result<Uint8Layout> {
469 if data_len > super::MAX_IPC_BYTES {
470 bail!(
471 "UInt8 payload too large: {data_len} bytes (max {})",
472 super::MAX_IPC_BYTES
473 );
474 }
475 let validity_len = data_len.div_ceil(8);
476 let validity_padded = round_up(validity_len, ALIGN);
477 let body_len = validity_padded + round_up(data_len, ALIGN);
478 let nodes = [FieldNode::new(data_len as i64, 0)];
481 let buffers = [
482 IpcBuffer::new(0, validity_len as i64),
483 IpcBuffer::new(validity_padded as i64, data_len as i64),
484 ];
485 let record_batch_message = build_record_batch_message(data_len, &nodes, &buffers, body_len);
486 let schema_message = build_schema_message(&DataType::UInt8)?;
487 let schema_block = round_up(PREFIX_LEN + schema_message.len(), ALIGN);
488 let record_batch_block = round_up(PREFIX_LEN + record_batch_message.len(), ALIGN);
489 let total = schema_block + record_batch_block + body_len + PREFIX_LEN;
490 let data_offset = schema_block + record_batch_block + validity_padded;
491 Ok(Uint8Layout {
492 schema_message,
493 record_batch_message,
494 validity_len,
495 validity_padded,
496 body_len,
497 total,
498 data_offset,
499 })
500}
501
502pub fn uint8_ipc_len(data_len: usize) -> eyre::Result<usize> {
506 Ok(uint8_layout(data_len)?.total)
507}
508
509pub fn encode_uint8_ipc_header(dst: &mut [u8], data_len: usize) -> eyre::Result<usize> {
518 let layout = uint8_layout(data_len)?;
519 if dst.len() != layout.total {
520 bail!(
521 "destination size {} does not match required UInt8 IPC length {}",
522 dst.len(),
523 layout.total
524 );
525 }
526 let mut at = 0;
527 at += write_framed_message(dst, at, &layout.schema_message);
528 at += write_framed_message(dst, at, &layout.record_batch_message);
529 let body_start = at;
530
531 dst[body_start..body_start + layout.validity_len].fill(0xff);
533 dst[body_start + layout.validity_len..body_start + layout.validity_padded].fill(0);
534
535 let data_end = layout.data_offset + data_len;
538 let body_end = body_start + layout.body_len;
539 dst[data_end..body_end].fill(0);
540
541 dst[body_end..body_end + 4].copy_from_slice(&CONTINUATION_MARKER);
543 dst[body_end + 4..body_end + 8].copy_from_slice(&0i32.to_le_bytes());
544 debug_assert_eq!(body_end + PREFIX_LEN, layout.total);
545
546 Ok(layout.data_offset)
547}
548
549pub fn schema_block_len(stream: &[u8]) -> Option<usize> {
553 if stream.len() < PREFIX_LEN || stream[0..4] != CONTINUATION_MARKER {
554 return None;
555 }
556 let metadata_len = i32::from_le_bytes(stream[4..8].try_into().ok()?);
557 let block = PREFIX_LEN.checked_add(usize::try_from(metadata_len).ok()?)?;
558 (block <= stream.len()).then_some(block)
559}
560
561pub fn schema_block_and_hash(stream: &[u8]) -> Option<(u64, &[u8])> {
568 let block = schema_block_len(stream)?;
569 let schema = stream.get(..block)?;
570 Some((dora_message::metadata::fnv1a(schema), schema))
571}
572
573pub fn batch_slice(stream: &[u8]) -> Option<&[u8]> {
579 let block = schema_block_len(stream)?;
580 (stream.len() >= block + PREFIX_LEN).then(|| &stream[block..])
583}
584
585const MAX_RETAINED_SCHEMAS: usize = 8;
590
591pub struct InputDecoder {
597 decoder: arrow::ipc::reader::StreamDecoder,
599 schema_hash: Option<u64>,
601 schemas: Vec<(u64, ArrowBuffer)>,
608}
609
610impl Default for InputDecoder {
611 fn default() -> Self {
612 Self::new()
613 }
614}
615
616impl InputDecoder {
617 pub fn new() -> Self {
621 Self {
622 decoder: arrow::ipc::reader::StreamDecoder::new(),
623 schema_hash: None,
624 schemas: Vec::new(),
625 }
626 }
627
628 pub fn reset(&mut self) {
632 self.decoder = arrow::ipc::reader::StreamDecoder::new();
633 self.schema_hash = None;
634 self.schemas.clear();
635 }
636
637 pub fn knows_schema(&self, hash: u64) -> bool {
645 self.schema_hash == Some(hash) || self.schemas.iter().any(|(h, _)| *h == hash)
646 }
647
648 pub fn set_schema(&mut self, hash: u64, schema: ArrowBuffer) -> eyre::Result<()> {
654 check_ipc_size(schema.len())?;
655 if self.schema_hash == Some(hash) {
656 return Ok(());
657 }
658 self.prime(hash, schema)
659 }
660
661 fn prime(&mut self, hash: u64, schema: ArrowBuffer) -> eyre::Result<()> {
664 let mut decoder = arrow::ipc::reader::StreamDecoder::new();
668 prime_with_schema(&mut decoder, schema.clone())?;
669 self.decoder = decoder;
670 self.schema_hash = Some(hash);
671 self.schemas.retain(|(h, _)| *h != hash);
672 self.schemas.push((hash, schema));
673 if self.schemas.len() > MAX_RETAINED_SCHEMAS {
674 self.schemas.remove(0);
675 }
676 Ok(())
677 }
678
679 pub fn decode_batch(
689 &mut self,
690 buffer: ArrowBuffer,
691 hash: u64,
692 ) -> eyre::Result<Option<arrow::array::ArrayData>> {
693 check_ipc_size(buffer.len())?;
694 if self.schema_hash != Some(hash) {
695 match self.schemas.iter().find(|(h, _)| *h == hash) {
696 Some((_, schema)) => {
697 let schema = schema.clone();
698 self.prime(hash, schema)?;
699 }
700 None => return Ok(None),
702 }
703 }
704 match decode_one_batch(&mut self.decoder, buffer) {
705 Ok(array) => {
706 self.decoder = arrow::ipc::reader::StreamDecoder::new();
712 self.schema_hash = None;
713 Ok(Some(array))
714 }
715 Err(e) => {
716 self.decoder = arrow::ipc::reader::StreamDecoder::new();
724 self.schema_hash = None;
725 Err(e)
726 }
727 }
728 }
729}
730
731fn check_ipc_size(len: usize) -> eyre::Result<()> {
736 if len > super::MAX_IPC_BYTES {
737 bail!(
738 "Arrow IPC payload too large: {len} bytes (max {})",
739 super::MAX_IPC_BYTES
740 );
741 }
742 Ok(())
743}
744
745fn prime_with_schema(
747 decoder: &mut arrow::ipc::reader::StreamDecoder,
748 mut buffer: ArrowBuffer,
749) -> eyre::Result<()> {
750 while !buffer.is_empty() {
751 let before = buffer.len();
752 if decoder
753 .decode(&mut buffer)
754 .map_err(|e| eyre!("failed to decode IPC schema message: {e}"))?
755 .is_some()
756 {
757 bail!("expected a schema message but got a record batch");
758 }
759 if buffer.len() == before {
762 bail!("IPC schema decoder made no progress on a partial/corrupt message");
763 }
764 }
765 Ok(())
766}
767
768fn decode_one_batch(
779 decoder: &mut arrow::ipc::reader::StreamDecoder,
780 mut buffer: ArrowBuffer,
781) -> eyre::Result<arrow::array::ArrayData> {
782 while !buffer.is_empty() {
783 let before = buffer.len();
784 if let Some(batch) = decoder
785 .decode(&mut buffer)
786 .map_err(|e| eyre!("failed to decode IPC record batch: {e}"))?
787 {
788 if batch.num_columns() != 1 {
789 bail!(
790 "expected 1 column in IPC record batch, got {}",
791 batch.num_columns()
792 );
793 }
794 return Ok(batch.column(0).to_data());
795 }
796 if buffer.len() == before {
799 bail!("IPC batch decoder made no progress on a partial/corrupt message");
800 }
801 }
802 bail!("IPC batch message yielded no record batch")
803}
804
805#[cfg(test)]
806mod tests {
807 use super::*;
808 use crate::arrow_utils::decode_arrow_ipc_zero_copy;
809 use arrow::array::{
810 Array, ArrayRef, BooleanArray, FixedSizeBinaryArray, Float32Array, Int32Array,
811 LargeStringArray, ListArray, NullArray, StringArray, StructArray, UInt8Array, UInt64Array,
812 };
813 use arrow::buffer::Buffer;
814 use arrow::ipc::reader::{StreamDecoder, StreamReader};
815 use arrow_schema::{DataType, Field};
816 use std::io::Cursor;
817 use std::sync::Arc;
818
819 fn fast_encode(array: &ArrayData) -> Vec<u8> {
821 let len = ipc_fast_path_len(array).expect("array should be fast-path eligible");
822 let mut buf = vec![0u8; len];
823 encode_ipc_into(array, &mut buf).expect("fast-path encode");
824 buf
825 }
826
827 fn read_official(bytes: &[u8]) -> ArrayData {
830 let mut reader = StreamReader::try_new(Cursor::new(bytes), None).expect("open IPC stream");
831 let batch = reader
832 .next()
833 .expect("one batch")
834 .expect("batch decodes via official reader");
835 assert_eq!(batch.num_columns(), 1);
836 batch.column(0).to_data()
837 }
838
839 fn aligned_buffer(bytes: &[u8]) -> (Buffer, usize, usize) {
842 use aligned_vec::{AVec, ConstAlign};
843 use std::ptr::NonNull;
844 let mut aligned: AVec<u8, ConstAlign<128>> = AVec::__from_elem(128, 0, bytes.len());
845 aligned.copy_from_slice(bytes);
846 let base = aligned.as_ptr() as usize;
847 let len = aligned.len();
848 let ptr = NonNull::new(aligned.as_ptr() as *mut u8).unwrap();
849 let buffer =
851 unsafe { Buffer::from_custom_allocation(ptr, len, std::sync::Arc::new(aligned)) };
852 (buffer, base, len)
853 }
854
855 fn assert_fast_roundtrip(array: &ArrayData) {
859 let encoded = fast_encode(array);
860 let decoded = read_official(&encoded);
861 assert_eq!(array, &decoded, "fast-path stream must decode to the input");
862 let (buffer, _, _) = aligned_buffer(&encoded);
864 let zc = decode_arrow_ipc_zero_copy(buffer).expect("zero-copy decode");
865 assert_eq!(array, &zc, "zero-copy decode must equal the input");
866 }
867
868 fn batch_bytes(array: &ArrayData) -> Vec<u8> {
871 let len = batch_fast_path_len(array).unwrap();
872 let mut buf = vec![0u8; len];
873 encode_batch_into(array, &mut buf).unwrap();
874 buf
875 }
876
877 fn batch_buf(array: &ArrayData) -> Buffer {
878 Buffer::from_vec(batch_bytes(array))
879 }
880
881 #[test]
886 fn input_decoder_schema_then_batches() {
887 let f32_schema = || Buffer::from_vec(encode_schema_message(&DataType::Float32).unwrap());
888
889 let mut dec = InputDecoder::new();
890
891 let early = Float32Array::from(vec![9.0]).into_data();
893 assert!(dec.decode_batch(batch_buf(&early), 7).unwrap().is_none());
894
895 dec.set_schema(7, f32_schema()).unwrap();
897 for vals in [vec![1.0f32, 2.0, 3.0], vec![4.0], vec![5.0, 6.0, 7.0]] {
898 let array = Float32Array::from(vals).into_data();
899 assert_eq!(
900 dec.decode_batch(batch_buf(&array), 7).unwrap().unwrap(),
901 array
902 );
903 }
904
905 let other = Float32Array::from(vec![8.0]).into_data();
907 assert!(dec.decode_batch(batch_buf(&other), 99).unwrap().is_none());
908
909 dec.set_schema(99, f32_schema()).unwrap();
911 let after = Float32Array::from(vec![10.0, 11.0]).into_data();
912 assert_eq!(
913 dec.decode_batch(batch_buf(&after), 99).unwrap().unwrap(),
914 after
915 );
916 }
917
918 #[test]
925 fn input_decoder_retains_multiple_schemas() {
926 let schema_msg = |dt: &DataType| Buffer::from_vec(encode_schema_message(dt).unwrap());
927
928 let mut dec = InputDecoder::new();
929 dec.set_schema(1, schema_msg(&DataType::Float32)).unwrap();
930 dec.set_schema(2, schema_msg(&DataType::Int32)).unwrap();
931
932 let ints = Int32Array::from(vec![1, 2, 3]).into_data();
934 assert_eq!(
935 dec.decode_batch(batch_buf(&ints), 2).unwrap().unwrap(),
936 ints
937 );
938
939 let floats = Float32Array::from(vec![4.0, 5.0]).into_data();
941 assert_eq!(
942 dec.decode_batch(batch_buf(&floats), 1).unwrap().unwrap(),
943 floats,
944 "a schema installed earlier must be retained across later primes"
945 );
946
947 let more_ints = Int32Array::from(vec![6]).into_data();
949 assert_eq!(
950 dec.decode_batch(batch_buf(&more_ints), 2).unwrap().unwrap(),
951 more_ints
952 );
953 }
954
955 #[test]
961 fn input_decoder_handles_sequential_batches_arrow_59_terminal_state() {
962 let f32_schema = || Buffer::from_vec(encode_schema_message(&DataType::Float32).unwrap());
963
964 let mut dec = InputDecoder::new();
965 dec.set_schema(7, f32_schema()).unwrap();
967
968 let batches = vec![
972 Float32Array::from(vec![1.0, 2.0]).into_data(),
973 Float32Array::from(vec![3.0]).into_data(),
974 Float32Array::from(vec![4.0, 5.0, 6.0]).into_data(),
975 Float32Array::from(vec![7.0, 8.0]).into_data(),
976 Float32Array::from(vec![9.0, 10.0, 11.0, 12.0]).into_data(),
977 ];
978
979 for (i, batch) in batches.iter().enumerate() {
980 let result = dec.decode_batch(batch_buf(batch), 7);
981 assert!(
982 result.is_ok(),
983 "batch {} decode failed: {:?}",
984 i,
985 result.err()
986 );
987 let decoded = result.unwrap().unwrap();
988 assert_eq!(
989 &decoded, batch,
990 "batch {} mismatch: expected {:?}, got {:?}",
991 i, batch, decoded
992 );
993 }
994 }
995
996 #[test]
997 fn input_decoder_evicts_oldest_schema_beyond_cap() {
998 let f32_schema = || Buffer::from_vec(encode_schema_message(&DataType::Float32).unwrap());
999
1000 let mut dec = InputDecoder::new();
1001 for hash in 0..=(MAX_RETAINED_SCHEMAS as u64) {
1004 dec.set_schema(hash, f32_schema()).unwrap();
1005 }
1006 let array = Float32Array::from(vec![1.0]).into_data();
1007 assert!(
1008 dec.decode_batch(batch_buf(&array), 0).unwrap().is_none(),
1009 "the oldest schema must be evicted beyond the cap"
1010 );
1011 assert_eq!(
1012 dec.decode_batch(batch_buf(&array), 1).unwrap().unwrap(),
1013 array,
1014 "schemas within the cap must be retained"
1015 );
1016 }
1017
1018 #[test]
1023 fn decode_batch_resets_on_error_then_reprimes() {
1024 let mut dec = InputDecoder::new();
1025 dec.set_schema(
1026 7,
1027 Buffer::from_vec(encode_schema_message(&DataType::Float32).unwrap()),
1028 )
1029 .unwrap();
1030
1031 let good = Float32Array::from(vec![4.0, 5.0]).into_data();
1033 assert_eq!(
1034 dec.decode_batch(batch_buf(&good), 7).unwrap().unwrap(),
1035 good
1036 );
1037
1038 let dropped = Float32Array::from(vec![6.0, 7.0, 8.0, 9.0]).into_data();
1043 let mut truncated = batch_bytes(&dropped);
1044 truncated.truncate(truncated.len() / 2);
1045 assert!(dec.decode_batch(Buffer::from_vec(truncated), 7).is_err());
1046
1047 let after = Float32Array::from(vec![10.0, 11.0]).into_data();
1050 assert_eq!(
1051 dec.decode_batch(batch_buf(&after), 7).unwrap().unwrap(),
1052 after,
1053 "after a failed batch the decoder must re-prime from the retained schema"
1054 );
1055
1056 dec.reset();
1058 let final_batch = Float32Array::from(vec![12.0]).into_data();
1059 assert!(
1060 dec.decode_batch(batch_buf(&final_batch), 7)
1061 .unwrap()
1062 .is_none(),
1063 "after a full reset the decoder must drop until a schema is re-installed"
1064 );
1065 }
1066
1067 #[test]
1077 fn input_decoder_dictionary_fallback_batch_sequence() {
1078 use arrow::array::DictionaryArray;
1079 use arrow::datatypes::Int32Type;
1080
1081 fn dict(words: &[&str]) -> ArrayData {
1084 let mut values: Vec<&str> = Vec::new();
1085 let mut keys: Vec<i32> = Vec::new();
1086 for w in words {
1087 let idx = values.iter().position(|v| v == w).unwrap_or_else(|| {
1088 values.push(*w);
1089 values.len() - 1
1090 });
1091 keys.push(idx as i32);
1092 }
1093 DictionaryArray::<Int32Type>::try_new(
1094 Int32Array::from(keys),
1095 Arc::new(StringArray::from(values)),
1096 )
1097 .unwrap()
1098 .into_data()
1099 }
1100
1101 let first = dict(&["a", "b", "a", "c", "b"]);
1102 assert!(
1104 ipc_fast_path_len(&first).is_none(),
1105 "dictionary must route to the official-writer fallback"
1106 );
1107
1108 let mut dec = InputDecoder::new();
1111 let full0 = encode_ipc_to_vec(&first).unwrap();
1112 let block = schema_block_len(&full0).unwrap();
1113 dec.set_schema(1, Buffer::from(&full0[..block])).unwrap();
1114
1115 let slice0 = batch_slice(&full0).expect("fallback stream is a valid IPC stream");
1120 assert_eq!(
1121 dec.decode_batch(Buffer::from(slice0), 1)
1122 .unwrap()
1123 .expect("first batch decodes against the primed decoder"),
1124 first
1125 );
1126 for words in [
1127 ["x", "y", "x", "z"].as_slice(),
1128 ["b", "b"].as_slice(),
1129 ["new", "values", "entirely"].as_slice(),
1130 ] {
1131 let arr = dict(words);
1132 let full = encode_ipc_to_vec(&arr).unwrap();
1133 let slice = batch_slice(&full).expect("fallback stream is a valid IPC stream");
1134 let got = dec
1135 .decode_batch(Buffer::from(slice), 1)
1136 .unwrap()
1137 .expect("batch must decode against the primed decoder");
1138 assert_eq!(
1139 got, arr,
1140 "replacement-dictionary batch must decode correctly"
1141 );
1142 }
1143 }
1144
1145 #[test]
1150 fn schema_primed_decoder_decodes_batch_sequence() {
1151 let schema = encode_schema_message(&DataType::Float32).unwrap();
1152 let mut decoder = StreamDecoder::new();
1153
1154 let mut sbuf = Buffer::from_vec(schema);
1156 while !sbuf.is_empty() {
1157 assert!(
1158 decoder.decode(&mut sbuf).unwrap().is_none(),
1159 "schema message must not yield a batch"
1160 );
1161 }
1162
1163 for vals in [vec![1.0f32, 2.0, 3.0], vec![4.0, 5.0], vec![6.0]] {
1164 let array = Float32Array::from(vals).into_data();
1165 let len = batch_fast_path_len(&array).unwrap();
1166 let mut buf = vec![0u8; len];
1167 encode_batch_into(&array, &mut buf).unwrap();
1168
1169 let mut bbuf = Buffer::from_vec(buf);
1170 let mut got = None;
1171 while !bbuf.is_empty() {
1172 if let Some(b) = decoder.decode(&mut bbuf).unwrap() {
1173 got = Some(b);
1174 break;
1175 }
1176 }
1177 assert_eq!(
1178 got.expect("batch message must decode against the primed decoder")
1179 .column(0)
1180 .to_data(),
1181 array
1182 );
1183 }
1184 }
1185
1186 #[test]
1190 fn uint8_ipc_header_constructs_in_place() {
1191 for data_len in [0usize, 1, 7, 8, 9, 1000] {
1192 let bytes: Vec<u8> = (0..data_len).map(|i| (i % 251) as u8).collect();
1193 let total = uint8_ipc_len(data_len).unwrap();
1194 let mut dst = vec![0u8; total];
1195 let offset = encode_uint8_ipc_header(&mut dst, data_len).unwrap();
1196 dst[offset..offset + data_len].copy_from_slice(&bytes);
1198
1199 let arr = arrow::array::make_array(read_official(&dst));
1201 let u8 = arr.as_any().downcast_ref::<UInt8Array>().unwrap();
1202 assert_eq!(u8.values(), bytes.as_slice(), "len {data_len}");
1203 assert_eq!(u8.null_count(), 0);
1204
1205 let array = UInt8Array::from(bytes).into_data();
1207 assert_eq!(dst, fast_encode(&array), "len {data_len}");
1208 }
1209 }
1210
1211 #[test]
1212 fn roundtrip_primitive_no_nulls() {
1213 let array = Float32Array::from((0..1000).map(|i| i as f32).collect::<Vec<_>>()).into_data();
1214 assert_fast_roundtrip(&array);
1215 }
1216
1217 #[test]
1218 fn roundtrip_primitive_with_nulls() {
1219 let array = UInt64Array::from(vec![Some(1), None, Some(3), None, Some(5)]).into_data();
1220 assert_fast_roundtrip(&array);
1221 }
1222
1223 #[test]
1224 fn roundtrip_empty_primitive() {
1225 let array = Int32Array::from(Vec::<i32>::new()).into_data();
1226 assert_fast_roundtrip(&array);
1227 }
1228
1229 #[test]
1235 fn uint8_header_zero_len_full_stream_roundtrip() {
1236 let len = uint8_ipc_len(0).unwrap();
1237 let mut buf = vec![0u8; len];
1238 let off = encode_uint8_ipc_header(&mut buf, 0).unwrap();
1239 assert_eq!(off, len - PREFIX_LEN);
1241 assert_eq!(read_official(&buf).len(), 0);
1242 let (buffer, _, _) = aligned_buffer(&buf);
1243 let decoded = decode_arrow_ipc_zero_copy(buffer).unwrap();
1244 assert_eq!(decoded.data_type(), &DataType::UInt8);
1245 assert_eq!(decoded.len(), 0);
1246 }
1247
1248 #[test]
1254 fn schema_once_zero_len_batch_roundtrip() {
1255 let schema = || Buffer::from_vec(encode_schema_message(&DataType::UInt8).unwrap());
1256 let batch = |vals: &[u8]| {
1257 let a = UInt8Array::from(vals.to_vec()).into_data();
1258 let len = batch_fast_path_len(&a).unwrap();
1259 let mut b = vec![0u8; len];
1260 encode_batch_into(&a, &mut b).unwrap();
1261 Buffer::from_vec(b)
1262 };
1263
1264 let mut dec = InputDecoder::new();
1265 dec.set_schema(1, schema()).unwrap();
1266
1267 let decoded = dec
1269 .decode_batch(batch(&[]), 1)
1270 .unwrap()
1271 .expect("0-row batch must decode, not drop");
1272 assert_eq!(decoded.data_type(), &DataType::UInt8);
1273 assert_eq!(decoded.len(), 0);
1274
1275 for vals in [vec![1u8, 2, 3], vec![], vec![9u8], vec![]] {
1278 let d = dec.decode_batch(batch(&vals), 1).unwrap().unwrap();
1279 assert_eq!(d.len(), vals.len());
1280 assert_eq!(d.data_type(), &DataType::UInt8);
1281 }
1282 }
1283
1284 #[test]
1289 fn schema_once_zero_len_via_batch_slice_roundtrip() {
1290 let total = uint8_ipc_len(0).unwrap();
1291 let mut full = vec![0u8; total];
1292 encode_uint8_ipc_header(&mut full, 0).unwrap();
1293
1294 let sblock = schema_block_len(&full).unwrap();
1295 let schema = Buffer::from(&full[..sblock]);
1296 let batch = batch_slice(&full).expect("batch slice of a valid stream");
1297
1298 let mut dec = InputDecoder::new();
1299 dec.set_schema(7, schema).unwrap();
1300 let decoded = dec
1301 .decode_batch(Buffer::from(batch), 7)
1302 .unwrap()
1303 .expect("0-row batch via batch_slice must decode, not drop");
1304 assert_eq!(decoded.data_type(), &DataType::UInt8);
1305 assert_eq!(decoded.len(), 0);
1306 }
1307
1308 #[test]
1314 fn schema_block_plus_batch_slice_reconstructs_full_stream() {
1315 let array = Int32Array::from(vec![1, 2, 3]).into_data();
1316 let full = fast_encode(&array);
1317 let sblock = schema_block_len(&full).unwrap();
1318 let schema = &full[..sblock];
1319 let batch = batch_slice(&full).expect("batch slice of a valid stream");
1320 let rebuilt = [schema, batch].concat();
1321 assert_eq!(
1322 rebuilt, full,
1323 "schema_block ++ batch_slice must equal the original stream"
1324 );
1325 assert_eq!(read_official(&rebuilt), array);
1326 }
1327
1328 #[test]
1329 fn roundtrip_boolean() {
1330 let array =
1331 BooleanArray::from(vec![true, false, true, true, false, false, true]).into_data();
1332 assert_fast_roundtrip(&array);
1333 }
1334
1335 #[test]
1336 fn roundtrip_utf8() {
1337 let array =
1338 StringArray::from(vec![Some("hello"), None, Some(""), Some("world!")]).into_data();
1339 assert_fast_roundtrip(&array);
1340 }
1341
1342 #[test]
1343 fn roundtrip_large_utf8_64bit_offsets() {
1344 let array = LargeStringArray::from(vec!["a", "bb", "ccc"]).into_data();
1345 assert_fast_roundtrip(&array);
1346 }
1347
1348 #[test]
1349 fn roundtrip_fixed_size_binary() {
1350 let values = vec![vec![1u8, 2, 3], vec![4, 5, 6], vec![7, 8, 9]];
1351 let array = FixedSizeBinaryArray::try_from_iter(values.into_iter())
1352 .unwrap()
1353 .into_data();
1354 assert_fast_roundtrip(&array);
1355 }
1356
1357 #[test]
1361 fn roundtrip_fixed_size_list() {
1362 use arrow::array::FixedSizeListArray;
1363 let values = Int32Array::from((0..12).collect::<Vec<_>>());
1364 let field = Arc::new(Field::new("item", DataType::Int32, true));
1365 let array = FixedSizeListArray::try_new(field, 3, Arc::new(values), None)
1366 .unwrap()
1367 .into_data();
1368 assert_fast_roundtrip(&array);
1369 }
1370
1371 #[test]
1374 fn roundtrip_fixed_size_list_with_nulls() {
1375 use arrow::array::FixedSizeListArray;
1376 use arrow::buffer::NullBuffer;
1377 let values = Int32Array::from((0..12).collect::<Vec<_>>());
1378 let field = Arc::new(Field::new("item", DataType::Int32, true));
1379 let nulls = NullBuffer::from(vec![true, false, true, true]);
1380 let array = FixedSizeListArray::try_new(field, 3, Arc::new(values), Some(nulls))
1381 .unwrap()
1382 .into_data();
1383 assert_fast_roundtrip(&array);
1384 }
1385
1386 #[test]
1388 fn roundtrip_decimal128() {
1389 use arrow::array::Decimal128Array;
1390 let array = Decimal128Array::from(vec![Some(12_345i128), None, Some(-9_876), Some(0)])
1391 .with_precision_and_scale(20, 4)
1392 .unwrap()
1393 .into_data();
1394 assert_fast_roundtrip(&array);
1395 }
1396
1397 #[test]
1399 fn roundtrip_timestamp_temporal() {
1400 use arrow::array::TimestampMicrosecondArray;
1401 let array =
1402 TimestampMicrosecondArray::from(vec![Some(1_000_000i64), None, Some(2_500_000)])
1403 .into_data();
1404 assert_fast_roundtrip(&array);
1405 }
1406
1407 #[test]
1408 fn roundtrip_struct_with_multilevel_nulls() {
1409 let array = StructArray::from(vec![
1410 (
1411 Arc::new(Field::new("a", DataType::UInt64, true)),
1412 Arc::new(UInt64Array::from(vec![Some(1), None, Some(3)])) as ArrayRef,
1413 ),
1414 (
1415 Arc::new(Field::new("b", DataType::Utf8, true)),
1416 Arc::new(StringArray::from(vec![Some("x"), Some("yy"), None])) as ArrayRef,
1417 ),
1418 ])
1419 .into_data();
1420 assert_fast_roundtrip(&array);
1421 }
1422
1423 #[test]
1428 fn roundtrip_struct_with_oversized_child() {
1429 use arrow_schema::Fields;
1430 let child = Int32Array::from(vec![10, 20, 30]).into_data(); let fields: Fields = vec![Field::new("v", DataType::Int32, false)].into();
1432 let struct_data = ArrayData::builder(DataType::Struct(fields))
1433 .len(2) .add_child_data(child)
1435 .build()
1436 .unwrap();
1437
1438 assert!(
1439 ipc_fast_path_len(&struct_data).is_some(),
1440 "struct with an oversized child should stay on the fast path"
1441 );
1442 let decoded = read_official(&fast_encode(&struct_data));
1443 assert_eq!(decoded.len(), 2);
1444 let arr = arrow::array::make_array(decoded);
1445 let sa = arr.as_any().downcast_ref::<StructArray>().unwrap();
1446 let col = sa.column(0).as_any().downcast_ref::<Int32Array>().unwrap();
1447 assert_eq!(
1448 col.values(),
1449 &[10, 20],
1450 "child must be truncated to the struct's len"
1451 );
1452 }
1453
1454 #[test]
1455 fn roundtrip_list_of_primitive() {
1456 let data = vec![
1457 Some(vec![Some(0), Some(1), Some(2)]),
1458 None,
1459 Some(vec![Some(3), None, Some(5)]),
1460 Some(vec![]),
1461 ];
1462 let array =
1463 ListArray::from_iter_primitive::<arrow::datatypes::Int32Type, _, _>(data).into_data();
1464 assert_fast_roundtrip(&array);
1465 }
1466
1467 #[test]
1468 fn roundtrip_nullarray_zero_and_n() {
1469 assert_fast_roundtrip(&NullArray::new(0).into_data());
1470 assert_fast_roundtrip(&NullArray::new(7).into_data());
1471 }
1472
1473 #[test]
1477 fn fast_path_decodes_zero_copy() {
1478 let array = UInt64Array::from((0..50_000u64).collect::<Vec<_>>()).into_data();
1479 let encoded = fast_encode(&array);
1480
1481 {
1483 let (mut buffer, _, _) = aligned_buffer(&encoded);
1484 let mut decoder = StreamDecoder::new().with_require_alignment(true);
1485 let mut got = None;
1486 while !buffer.is_empty() {
1487 if let Some(b) = decoder
1488 .decode(&mut buffer)
1489 .expect("aligned fast-path stream must decode without realignment")
1490 {
1491 got = Some(b);
1492 break;
1493 }
1494 }
1495 assert_eq!(got.unwrap().column(0).to_data(), array);
1496 }
1497
1498 {
1500 let (buffer, base, len) = aligned_buffer(&encoded);
1501 let decoded = decode_arrow_ipc_zero_copy(buffer).unwrap();
1502 let ptr = decoded.buffers()[0].as_ptr() as usize;
1503 assert!(
1504 ptr >= base && ptr < base + len,
1505 "decoded data buffer at {ptr:#x} is outside input [{base:#x}, {:#x}) — a copy happened",
1506 base + len
1507 );
1508 }
1509 }
1510
1511 #[test]
1512 fn fast_path_len_matches_official_decode() {
1513 let array = Float32Array::from(vec![1.0, 2.0, 3.0, 4.0]).into_data();
1516 let encoded = fast_encode(&array);
1517 let mut reader = StreamReader::try_new(Cursor::new(&encoded[..]), None).unwrap();
1518 let _ = reader.next().unwrap().unwrap();
1519 assert!(reader.next().is_none(), "exactly one batch, fully consumed");
1520 }
1521
1522 #[test]
1525 fn sliced_array_routes_to_fallback_and_roundtrips() {
1526 let array = UInt64Array::from(vec![10, 20, 30, 40, 50])
1527 .into_data()
1528 .slice(2, 2); assert_eq!(array.offset(), 2);
1530 assert!(ipc_fast_path_len(&array).is_none());
1531
1532 let encoded = encode_ipc_to_vec(&array).unwrap();
1533 let decoded = read_official(&encoded);
1534 assert_eq!(array.len(), decoded.len());
1535 let dec = arrow::array::make_array(decoded);
1536 let dec = dec.as_any().downcast_ref::<UInt64Array>().unwrap();
1537 assert_eq!(dec.values(), &[30, 40]);
1538 }
1539
1540 #[test]
1542 fn view_type_routes_to_fallback() {
1543 use arrow::array::StringViewArray;
1544 let array = StringViewArray::from(vec!["a", "bb", "ccc"]).into_data();
1545 assert!(
1546 ipc_fast_path_len(&array).is_none(),
1547 "Utf8View is not fast-path eligible"
1548 );
1549 let encoded = encode_ipc_to_vec(&array).unwrap();
1550 let decoded = read_official(&encoded);
1551 assert_eq!(array, decoded);
1552 }
1553
1554 #[test]
1556 fn fuzz_roundtrip_many_shapes() {
1557 let mut state: u64 = 0x1234_5678_9abc_def0;
1559 let mut next = || {
1560 state = state
1561 .wrapping_mul(6364136223846793005)
1562 .wrapping_add(1442695040888963407);
1563 state
1564 };
1565
1566 for _ in 0..200 {
1567 let len = (next() % 64) as usize;
1568 let kind = next() % 6;
1569 let array: ArrayData = match kind {
1570 0 => UInt8Array::from(
1571 (0..len)
1572 .map(|i| {
1573 if next().is_multiple_of(4) {
1574 None
1575 } else {
1576 Some((i as u8).wrapping_add(1))
1577 }
1578 })
1579 .collect::<Vec<_>>(),
1580 )
1581 .into_data(),
1582 1 => Float32Array::from((0..len).map(|i| i as f32 * 0.5).collect::<Vec<_>>())
1583 .into_data(),
1584 2 => BooleanArray::from(
1585 (0..len)
1586 .map(|i| (i + next() as usize).is_multiple_of(2))
1587 .collect::<Vec<_>>(),
1588 )
1589 .into_data(),
1590 3 => StringArray::from(
1591 (0..len)
1592 .map(|i| {
1593 if next().is_multiple_of(5) {
1594 None
1595 } else {
1596 Some("x".repeat(i % 7))
1597 }
1598 })
1599 .collect::<Vec<_>>(),
1600 )
1601 .into_data(),
1602 4 => Int32Array::from(
1603 (0..len)
1604 .map(|i| {
1605 if next().is_multiple_of(3) {
1606 None
1607 } else {
1608 Some(i as i32 - 10)
1609 }
1610 })
1611 .collect::<Vec<_>>(),
1612 )
1613 .into_data(),
1614 _ => StructArray::from(vec![(
1615 Arc::new(Field::new("v", DataType::Int32, true)),
1616 Arc::new(Int32Array::from(
1617 (0..len).map(|i| Some(i as i32)).collect::<Vec<_>>(),
1618 )) as ArrayRef,
1619 )])
1620 .into_data(),
1621 };
1622 assert_fast_roundtrip(&array);
1623 }
1624 }
1625}