1use arrow::array::ArrayData;
41use arrow::buffer::Buffer as ArrowBuffer;
42use arrow::ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
43use arrow::ipc::{Buffer as IpcBuffer, FieldNode, MessageHeader, MetadataVersion};
44use arrow_schema::{DataType, Field, Schema};
45use eyre::{Context, bail, eyre};
46
47use super::{ARROW_BUFFER_ALIGNMENT as ALIGN, IpcPayload};
48use dora_arrow_convert::{DoraArray, internal::array_ref};
49
50const CONTINUATION_MARKER: [u8; 4] = [0xff, 0xff, 0xff, 0xff];
52const PREFIX_LEN: usize = 8;
54
55#[inline]
56fn round_up(n: usize, align: usize) -> usize {
57 debug_assert!(align.is_power_of_two());
58 (n + align - 1) & !(align - 1)
59}
60
61fn is_fast_path_type(data_type: &DataType) -> bool {
65 use DataType::*;
66 matches!(
67 data_type,
68 Null | Boolean
69 | Int8
70 | Int16
71 | Int32
72 | Int64
73 | UInt8
74 | UInt16
75 | UInt32
76 | UInt64
77 | Float16
78 | Float32
79 | Float64
80 | Timestamp(_, _)
81 | Date32
82 | Date64
83 | Time32(_)
84 | Time64(_)
85 | Duration(_)
86 | Interval(_)
87 | Decimal128(_, _)
88 | Decimal256(_, _)
89 | FixedSizeBinary(_)
90 | Binary
91 | LargeBinary
92 | Utf8
93 | LargeUtf8
94 | List(_)
95 | LargeList(_)
96 | FixedSizeList(_, _)
97 | Struct(_)
98 )
99}
100
101struct Desc {
104 offset: usize,
105 len: usize,
106 src: BufferSrc,
107}
108
109enum BufferSrc {
110 Bytes(ArrowBuffer, usize),
112 AllOnes,
114}
115
116#[derive(Default)]
117struct Layout {
118 nodes: Vec<FieldNode>,
119 ipc_buffers: Vec<IpcBuffer>,
120 descs: Vec<Desc>,
121 body_len: usize,
122}
123
124impl Layout {
125 fn push_buffer(&mut self, off: &mut usize, len: usize, src: BufferSrc) -> Option<()> {
130 let padded = len.checked_add(ALIGN - 1).map(|n| n & !(ALIGN - 1))?;
132 let next = off.checked_add(padded)?;
133 if next > super::MAX_IPC_BYTES {
134 return None;
135 }
136 self.ipc_buffers
137 .push(IpcBuffer::new(*off as i64, len as i64));
138 self.descs.push(Desc {
139 offset: *off,
140 len,
141 src,
142 });
143 *off = next;
144 Some(())
145 }
146}
147
148fn build_layout(array: &ArrayData) -> Option<Layout> {
152 let mut layout = Layout::default();
153 let mut off = 0usize;
154 build_layout_rec(array, &mut layout, &mut off)?;
155 layout.body_len = off;
156 Some(layout)
157}
158
159fn build_layout_rec(array: &ArrayData, layout: &mut Layout, off: &mut usize) -> Option<()> {
160 let data_type = array.data_type();
161 if !is_fast_path_type(data_type) || array.offset() != 0 {
162 return None;
163 }
164
165 let len = array.len();
166 let null_count = if matches!(data_type, DataType::Null) {
169 len
170 } else {
171 array.null_count()
172 };
173 layout
174 .nodes
175 .push(FieldNode::new(len as i64, null_count as i64));
176
177 if !matches!(data_type, DataType::Null) {
180 match array.nulls() {
181 Some(nulls) => {
182 if nulls.inner().offset() != 0 {
183 return None;
184 }
185 let sliced = nulls.inner().sliced();
186 let bytes = sliced.len();
187 layout.push_buffer(off, bytes, BufferSrc::Bytes(sliced, 0))?;
188 }
189 None => {
190 let bytes = len.div_ceil(8);
191 layout.push_buffer(off, bytes, BufferSrc::AllOnes)?;
192 }
193 }
194 }
195
196 let offset_width = match data_type {
202 DataType::Binary | DataType::Utf8 | DataType::List(_) => Some(4),
203 DataType::LargeBinary | DataType::LargeUtf8 | DataType::LargeList(_) => Some(8),
204 _ => None,
205 };
206 let mut list_values_len = None;
207 if let Some(width) = offset_width {
208 let (offsets, values) = match array.buffers() {
209 [offsets] => (offsets, None),
210 [offsets, values] => (offsets, Some(values)),
211 _ => return None,
212 };
213 let offsets_len = len.checked_add(1)?.checked_mul(width)?;
214 let first = read_offset(offsets, 0, width)?;
215 let last = read_offset(offsets, len, width)?;
216 if first != 0 {
219 return None;
220 }
221 layout.push_buffer(off, offsets_len, BufferSrc::Bytes(offsets.clone(), 0))?;
222 match values {
223 Some(values) => {
224 if values.len() < last {
225 return None;
226 }
227 layout.push_buffer(off, last, BufferSrc::Bytes(values.clone(), 0))?;
228 }
229 None => list_values_len = Some(last),
230 }
231 } else {
232 let needed = match data_type {
236 DataType::Boolean => Some(len.div_ceil(8)),
237 DataType::FixedSizeBinary(width) => len.checked_mul(usize::try_from(*width).ok()?),
238 _ => data_type.primitive_width().and_then(|w| len.checked_mul(w)),
239 };
240 for buffer in array.buffers() {
241 let bytes = needed.map_or(buffer.len(), |n| n.min(buffer.len()));
242 layout.push_buffer(off, bytes, BufferSrc::Bytes(buffer.clone(), 0))?;
243 }
244 }
245
246 match data_type {
248 DataType::FixedSizeList(_, value_size) => {
249 let n = len.checked_mul(*value_size as usize)?;
252 let child = array.child_data().first()?;
253 if child.len() < n {
254 return None;
255 }
256 build_layout_rec(&child.slice(0, n), layout, off)?;
257 }
258 DataType::Struct(_) => {
259 for child in array.child_data() {
265 if child.len() < len {
266 return None;
267 }
268 build_layout_rec(&child.slice(0, len), layout, off)?;
269 }
270 }
271 _ => {
272 for child in array.child_data() {
276 let values_len = list_values_len?;
277 match child.len().cmp(&values_len) {
278 std::cmp::Ordering::Less => return None,
279 std::cmp::Ordering::Equal => build_layout_rec(child, layout, off)?,
281 std::cmp::Ordering::Greater => {
282 build_layout_rec(&child.slice(0, values_len), layout, off)?
283 }
284 }
285 }
286 }
287 }
288
289 Some(())
290}
291
292fn read_offset(offsets: &ArrowBuffer, index: usize, width: usize) -> Option<usize> {
295 let start = index.checked_mul(width)?;
296 let bytes = offsets.as_slice().get(start..start.checked_add(width)?)?;
297 let value = match width {
298 4 => i64::from(i32::from_ne_bytes(bytes.try_into().ok()?)),
299 8 => i64::from_ne_bytes(bytes.try_into().ok()?),
300 _ => return None,
301 };
302 usize::try_from(value).ok()
303}
304
305struct Prepared {
307 layout: Layout,
308 schema_message: Vec<u8>,
309 record_batch_message: Vec<u8>,
310 schema_block: usize,
311 record_batch_block: usize,
312 total: usize,
313}
314
315fn ipc_write_options() -> eyre::Result<IpcWriteOptions> {
316 IpcWriteOptions::try_new(ALIGN, false, MetadataVersion::V5)
317 .map_err(|e| eyre!("failed to build Arrow IPC write options: {e}"))
318}
319
320fn build_schema_message(data_type: &DataType) -> eyre::Result<Vec<u8>> {
323 let schema = Schema::new(vec![Field::new("data", data_type.clone(), true)]);
324 let options = ipc_write_options()?;
325 let mut tracker = DictionaryTracker::new(false);
326 let encoded = IpcDataGenerator {}.schema_to_bytes_with_dictionary_tracker(
327 &schema,
328 &mut tracker,
329 &options,
330 );
331 Ok(encoded.ipc_message)
332}
333
334fn build_record_batch_message(
337 num_rows: usize,
338 nodes: &[FieldNode],
339 buffers: &[IpcBuffer],
340 body_len: usize,
341) -> Vec<u8> {
342 use flatbuffers::FlatBufferBuilder;
343
344 let mut fbb = FlatBufferBuilder::new();
345 let buffers_fb = fbb.create_vector(buffers);
346 let nodes_fb = fbb.create_vector(nodes);
347
348 let record_batch = {
349 let mut builder = arrow::ipc::RecordBatchBuilder::new(&mut fbb);
350 builder.add_length(num_rows as i64);
351 builder.add_nodes(nodes_fb);
352 builder.add_buffers(buffers_fb);
353 builder.finish()
354 };
355
356 let message = {
357 let mut builder = arrow::ipc::MessageBuilder::new(&mut fbb);
358 builder.add_version(MetadataVersion::V5);
359 builder.add_header_type(MessageHeader::RecordBatch);
360 builder.add_bodyLength(body_len as i64);
361 builder.add_header(record_batch.as_union_value());
362 builder.finish()
363 };
364
365 fbb.finish(message, None);
366 fbb.finished_data().to_vec()
367}
368
369fn prepare(array: &ArrayData) -> Option<Prepared> {
370 let layout = build_layout(array)?;
371 let schema_message = build_schema_message(array.data_type()).ok()?;
372 let record_batch_message = build_record_batch_message(
373 array.len(),
374 &layout.nodes,
375 &layout.ipc_buffers,
376 layout.body_len,
377 );
378 let schema_block = round_up(PREFIX_LEN + schema_message.len(), ALIGN);
379 let record_batch_block = round_up(PREFIX_LEN + record_batch_message.len(), ALIGN);
380 let total = schema_block + record_batch_block + layout.body_len + PREFIX_LEN;
382 if total > super::MAX_IPC_BYTES {
389 return None;
390 }
391 Some(Prepared {
392 layout,
393 schema_message,
394 record_batch_message,
395 schema_block,
396 record_batch_block,
397 total,
398 })
399}
400
401pub(crate) fn ipc_fast_path_len_data(array: &ArrayData) -> Option<usize> {
411 prepare(array).map(|p| p.total)
412}
413
414pub struct PreparedIpc(Prepared);
449
450impl PreparedIpc {
451 pub fn new(array: &DoraArray) -> Option<Self> {
454 Self::from_data(&array_ref(array).to_data())
455 }
456
457 pub(crate) fn from_data(array: &ArrayData) -> Option<Self> {
459 prepare(array).map(Self)
460 }
461
462 pub fn byte_len(&self) -> usize {
465 self.0.total
466 }
467
468 pub fn encode_into(&self, dst: &mut [u8]) -> eyre::Result<()> {
471 encode_prepared_into(&self.0, dst)
472 }
473}
474
475fn write_framed_message(dst: &mut [u8], at: usize, flatbuffer: &[u8]) -> usize {
480 let block = round_up(PREFIX_LEN + flatbuffer.len(), ALIGN);
481 let metadata_len = (block - PREFIX_LEN) as i32;
482 dst[at..at + 4].copy_from_slice(&CONTINUATION_MARKER);
483 dst[at + 4..at + 8].copy_from_slice(&metadata_len.to_le_bytes());
484 dst[at + 8..at + 8 + flatbuffer.len()].copy_from_slice(flatbuffer);
485 dst[at + 8 + flatbuffer.len()..at + block].fill(0);
487 block
488}
489
490pub(crate) fn encode_ipc_into_data(array: &ArrayData, dst: &mut [u8]) -> eyre::Result<()> {
497 let prepared =
498 prepare(array).ok_or_else(|| eyre!("array is not Arrow IPC fast-path eligible"))?;
499 encode_prepared_into(&prepared, dst)
500}
501
502fn encode_prepared_into(prepared: &Prepared, dst: &mut [u8]) -> eyre::Result<()> {
505 if dst.len() != prepared.total {
506 bail!(
507 "destination size {} does not match required IPC length {}",
508 dst.len(),
509 prepared.total
510 );
511 }
512
513 let mut at = 0;
514 at += write_framed_message(dst, at, &prepared.schema_message);
515 debug_assert_eq!(at, prepared.schema_block);
516 at += write_framed_message(dst, at, &prepared.record_batch_message);
517 debug_assert_eq!(at, prepared.schema_block + prepared.record_batch_block);
518
519 let body_start = at;
520 write_body(dst, body_start, &prepared.layout);
521 at = body_start + prepared.layout.body_len;
522
523 dst[at..at + 4].copy_from_slice(&CONTINUATION_MARKER);
525 dst[at + 4..at + 8].copy_from_slice(&0i32.to_le_bytes());
526 debug_assert_eq!(at + PREFIX_LEN, prepared.total);
527
528 Ok(())
529}
530
531fn write_body(dst: &mut [u8], body_start: usize, layout: &Layout) {
535 for desc in &layout.descs {
536 let start = body_start + desc.offset;
537 let end = start + desc.len;
538 match &desc.src {
539 BufferSrc::Bytes(buffer, src_off) => {
540 dst[start..end].copy_from_slice(&buffer.as_slice()[*src_off..*src_off + desc.len]);
541 }
542 BufferSrc::AllOnes => dst[start..end].fill(0xff),
543 }
544 let padded_end = body_start + desc.offset + round_up(desc.len, ALIGN);
545 dst[end..padded_end].fill(0);
546 }
547}
548
549pub fn encode_schema_message(array: &DoraArray) -> eyre::Result<Vec<u8>> {
557 encode_schema_message_for(array_ref(array).data_type())
558}
559
560pub(crate) fn encode_schema_message_for(data_type: &DataType) -> eyre::Result<Vec<u8>> {
562 let schema_message = build_schema_message(data_type)?;
563 let block = round_up(PREFIX_LEN + schema_message.len(), ALIGN);
564 let mut dst = vec![0u8; block];
565 write_framed_message(&mut dst, 0, &schema_message);
566 Ok(dst)
567}
568
569pub(crate) fn batch_fast_path_len_data(array: &ArrayData) -> Option<usize> {
573 let prepared = prepare(array)?;
574 Some(prepared.record_batch_block + prepared.layout.body_len + PREFIX_LEN)
575}
576
577pub(crate) fn encode_batch_into_data(array: &ArrayData, dst: &mut [u8]) -> eyre::Result<()> {
584 let prepared =
585 prepare(array).ok_or_else(|| eyre!("array is not Arrow IPC fast-path eligible"))?;
586 let expected = prepared.record_batch_block + prepared.layout.body_len + PREFIX_LEN;
587 if dst.len() != expected {
588 bail!(
589 "destination size {} does not match required batch length {expected}",
590 dst.len(),
591 );
592 }
593 let at = write_framed_message(dst, 0, &prepared.record_batch_message);
594 debug_assert_eq!(at, prepared.record_batch_block);
595 write_body(dst, at, &prepared.layout);
596 let body_end = at + prepared.layout.body_len;
599 dst[body_end..body_end + 4].copy_from_slice(&CONTINUATION_MARKER);
600 dst[body_end + 4..body_end + PREFIX_LEN].copy_from_slice(&0i32.to_le_bytes());
601 Ok(())
602}
603
604pub(crate) fn encode_ipc_to_vec_data(array: &ArrayData) -> eyre::Result<Vec<u8>> {
608 super::encode_arrow_ipc_data(array).context("Arrow IPC fallback encode")
609}
610
611pub fn ipc_fast_path_len(array: &DoraArray) -> Option<usize> {
618 ipc_fast_path_len_data(&array_ref(array).to_data())
619}
620
621pub fn encode_ipc_into(array: &DoraArray, dst: &mut [u8]) -> eyre::Result<()> {
628 encode_ipc_into_data(&array_ref(array).to_data(), dst)
629}
630
631pub fn batch_fast_path_len(array: &DoraArray) -> Option<usize> {
634 batch_fast_path_len_data(&array_ref(array).to_data())
635}
636
637pub fn encode_batch_into(array: &DoraArray, dst: &mut [u8]) -> eyre::Result<()> {
640 encode_batch_into_data(&array_ref(array).to_data(), dst)
641}
642
643pub fn encode_ipc_to_vec(array: &DoraArray) -> eyre::Result<Vec<u8>> {
647 encode_ipc_to_vec_data(&array_ref(array).to_data())
648}
649
650struct Uint8Layout {
654 schema_message: Vec<u8>,
655 record_batch_message: Vec<u8>,
656 validity_len: usize,
657 validity_padded: usize,
658 body_len: usize,
659 total: usize,
660 data_offset: usize,
661}
662
663fn uint8_layout(data_len: usize) -> eyre::Result<Uint8Layout> {
664 if data_len > super::MAX_IPC_BYTES {
669 bail!(
670 "UInt8 payload too large: {data_len} bytes (max {})",
671 super::MAX_IPC_BYTES
672 );
673 }
674 let validity_len = data_len.div_ceil(8);
675 let validity_padded = round_up(validity_len, ALIGN);
676 let body_len = validity_padded + round_up(data_len, ALIGN);
677 let nodes = [FieldNode::new(data_len as i64, 0)];
680 let buffers = [
681 IpcBuffer::new(0, validity_len as i64),
682 IpcBuffer::new(validity_padded as i64, data_len as i64),
683 ];
684 let record_batch_message = build_record_batch_message(data_len, &nodes, &buffers, body_len);
685 let schema_message = build_schema_message(&DataType::UInt8)?;
686 let schema_block = round_up(PREFIX_LEN + schema_message.len(), ALIGN);
687 let record_batch_block = round_up(PREFIX_LEN + record_batch_message.len(), ALIGN);
688 let total = schema_block + record_batch_block + body_len + PREFIX_LEN;
689 if total > super::MAX_IPC_BYTES {
696 bail!(
697 "UInt8 payload too large: {data_len} bytes encodes to a {total}-byte \
698 Arrow IPC stream (max {})",
699 super::MAX_IPC_BYTES
700 );
701 }
702 let data_offset = schema_block + record_batch_block + validity_padded;
703 Ok(Uint8Layout {
704 schema_message,
705 record_batch_message,
706 validity_len,
707 validity_padded,
708 body_len,
709 total,
710 data_offset,
711 })
712}
713
714pub fn uint8_ipc_len(data_len: usize) -> eyre::Result<usize> {
722 Ok(uint8_layout(data_len)?.total)
723}
724
725pub struct PreparedUint8Ipc {
734 layout: Uint8Layout,
735 data_len: usize,
736}
737
738impl PreparedUint8Ipc {
739 pub fn new(data_len: usize) -> eyre::Result<Self> {
742 Ok(Self {
743 layout: uint8_layout(data_len)?,
744 data_len,
745 })
746 }
747
748 pub fn byte_len(&self) -> usize {
750 self.layout.total
751 }
752
753 pub fn encode_header_into(&self, dst: &mut [u8]) -> eyre::Result<usize> {
758 encode_uint8_prepared_into(&self.layout, self.data_len, dst)
759 }
760}
761
762pub fn encode_uint8_ipc_header(dst: &mut [u8], data_len: usize) -> eyre::Result<usize> {
771 let layout = uint8_layout(data_len)?;
772 encode_uint8_prepared_into(&layout, data_len, dst)
773}
774
775fn encode_uint8_prepared_into(
779 layout: &Uint8Layout,
780 data_len: usize,
781 dst: &mut [u8],
782) -> eyre::Result<usize> {
783 if dst.len() != layout.total {
784 bail!(
785 "destination size {} does not match required UInt8 IPC length {}",
786 dst.len(),
787 layout.total
788 );
789 }
790 let mut at = 0;
791 at += write_framed_message(dst, at, &layout.schema_message);
792 at += write_framed_message(dst, at, &layout.record_batch_message);
793 let body_start = at;
794
795 dst[body_start..body_start + layout.validity_len].fill(0xff);
797 dst[body_start + layout.validity_len..body_start + layout.validity_padded].fill(0);
798
799 let data_end = layout.data_offset + data_len;
802 let body_end = body_start + layout.body_len;
803 dst[data_end..body_end].fill(0);
804
805 dst[body_end..body_end + 4].copy_from_slice(&CONTINUATION_MARKER);
807 dst[body_end + 4..body_end + 8].copy_from_slice(&0i32.to_le_bytes());
808 debug_assert_eq!(body_end + PREFIX_LEN, layout.total);
809
810 Ok(layout.data_offset)
811}
812
813pub fn schema_block_len(stream: &[u8]) -> Option<usize> {
817 if stream.len() < PREFIX_LEN || stream[0..4] != CONTINUATION_MARKER {
818 return None;
819 }
820 let metadata_len = i32::from_le_bytes(stream[4..8].try_into().ok()?);
821 let block = PREFIX_LEN.checked_add(usize::try_from(metadata_len).ok()?)?;
822 (block <= stream.len()).then_some(block)
823}
824
825pub fn schema_block_and_hash(stream: &[u8]) -> Option<(u64, &[u8])> {
832 let block = schema_block_len(stream)?;
833 let schema = stream.get(..block)?;
834 Some((dora_message::metadata::fnv1a(schema), schema))
835}
836
837pub fn batch_slice(stream: &[u8]) -> Option<&[u8]> {
843 let block = schema_block_len(stream)?;
844 (stream.len() >= block + PREFIX_LEN).then(|| &stream[block..])
847}
848
849const MAX_RETAINED_SCHEMAS: usize = 8;
854
855pub struct InputDecoder {
861 decoder: arrow::ipc::reader::StreamDecoder,
863 schema_hash: Option<u64>,
865 schemas: Vec<(u64, ArrowBuffer)>,
872}
873
874impl Default for InputDecoder {
875 fn default() -> Self {
876 Self::new()
877 }
878}
879
880impl InputDecoder {
881 pub fn new() -> Self {
885 Self {
886 decoder: arrow::ipc::reader::StreamDecoder::new(),
887 schema_hash: None,
888 schemas: Vec::new(),
889 }
890 }
891
892 pub fn reset(&mut self) {
896 self.decoder = arrow::ipc::reader::StreamDecoder::new();
897 self.schema_hash = None;
898 self.schemas.clear();
899 }
900
901 pub fn knows_schema(&self, hash: u64) -> bool {
909 self.schema_hash == Some(hash) || self.schemas.iter().any(|(h, _)| *h == hash)
910 }
911
912 pub fn set_schema(&mut self, hash: u64, schema: IpcPayload) -> eyre::Result<()> {
918 self.set_schema_raw(hash, schema.into_arrow())
919 }
920
921 pub(crate) fn set_schema_raw(&mut self, hash: u64, schema: ArrowBuffer) -> eyre::Result<()> {
923 check_ipc_size(schema.len())?;
924 if self.schema_hash == Some(hash) {
925 return Ok(());
926 }
927 self.prime(hash, schema)
928 }
929
930 fn prime(&mut self, hash: u64, schema: ArrowBuffer) -> eyre::Result<()> {
933 let mut decoder = arrow::ipc::reader::StreamDecoder::new();
937 prime_with_schema(&mut decoder, schema.clone())?;
938 self.decoder = decoder;
939 self.schema_hash = Some(hash);
940 self.schemas.retain(|(h, _)| *h != hash);
941 self.schemas.push((hash, schema));
942 if self.schemas.len() > MAX_RETAINED_SCHEMAS {
943 self.schemas.remove(0);
944 }
945 Ok(())
946 }
947
948 pub fn decode_batch(
958 &mut self,
959 buffer: IpcPayload,
960 hash: u64,
961 ) -> eyre::Result<Option<DoraArray>> {
962 Ok(self
963 .decode_batch_raw(buffer.into_arrow(), hash)?
964 .map(dora_arrow_convert::internal::from_array_data))
965 }
966
967 pub(crate) fn decode_batch_raw(
969 &mut self,
970 buffer: ArrowBuffer,
971 hash: u64,
972 ) -> eyre::Result<Option<arrow::array::ArrayData>> {
973 check_ipc_size(buffer.len())?;
974 if self.schema_hash != Some(hash) {
975 match self.schemas.iter().find(|(h, _)| *h == hash) {
976 Some((_, schema)) => {
977 let schema = schema.clone();
978 self.prime(hash, schema)?;
979 }
980 None => return Ok(None),
982 }
983 }
984 match decode_one_batch(&mut self.decoder, buffer) {
985 Ok(array) => {
986 Ok(Some(array))
996 }
997 Err(e) => {
998 self.decoder = arrow::ipc::reader::StreamDecoder::new();
1006 self.schema_hash = None;
1007 Err(e)
1008 }
1009 }
1010 }
1011}
1012
1013fn check_ipc_size(len: usize) -> eyre::Result<()> {
1018 if len > super::MAX_IPC_BYTES {
1019 bail!(
1020 "Arrow IPC payload too large: {len} bytes (max {})",
1021 super::MAX_IPC_BYTES
1022 );
1023 }
1024 Ok(())
1025}
1026
1027fn prime_with_schema(
1029 decoder: &mut arrow::ipc::reader::StreamDecoder,
1030 mut buffer: ArrowBuffer,
1031) -> eyre::Result<()> {
1032 while !buffer.is_empty() {
1033 let before = buffer.len();
1034 if decoder
1035 .decode(&mut buffer)
1036 .map_err(|e| eyre!("failed to decode IPC schema message: {e}"))?
1037 .is_some()
1038 {
1039 bail!("expected a schema message but got a record batch");
1040 }
1041 if buffer.len() == before {
1044 bail!("IPC schema decoder made no progress on a partial/corrupt message");
1045 }
1046 }
1047 Ok(())
1048}
1049
1050fn decode_one_batch(
1066 decoder: &mut arrow::ipc::reader::StreamDecoder,
1067 mut buffer: ArrowBuffer,
1068) -> eyre::Result<arrow::array::ArrayData> {
1069 while !buffer.is_empty() {
1070 let before = buffer.len();
1071 if let Some(batch) = decoder
1072 .decode(&mut buffer)
1073 .map_err(|e| eyre!("failed to decode IPC record batch: {e}"))?
1074 {
1075 if batch.num_columns() != 1 {
1076 bail!(
1077 "expected 1 column in IPC record batch, got {}",
1078 batch.num_columns()
1079 );
1080 }
1081 return Ok(batch.column(0).to_data());
1082 }
1083 if buffer.len() == before {
1086 bail!("IPC batch decoder made no progress on a partial/corrupt message");
1087 }
1088 }
1089 bail!("IPC batch message yielded no record batch")
1090}
1091
1092#[cfg(test)]
1093mod tests {
1094 use super::*;
1095 use crate::arrow_utils::decode_arrow_ipc_zero_copy_raw;
1096 use arrow::array::{
1097 Array, ArrayRef, BooleanArray, FixedSizeBinaryArray, Float32Array, Int32Array,
1098 LargeStringArray, ListArray, NullArray, StringArray, StructArray, UInt8Array, UInt64Array,
1099 };
1100 use arrow::buffer::Buffer;
1101 use arrow::ipc::reader::{StreamDecoder, StreamReader};
1102 use arrow_schema::{DataType, Field};
1103 use std::io::Cursor;
1104 use std::sync::Arc;
1105
1106 fn fast_encode(array: &ArrayData) -> Vec<u8> {
1108 let len = ipc_fast_path_len_data(array).expect("array should be fast-path eligible");
1109 let mut buf = vec![0u8; len];
1110 encode_ipc_into_data(array, &mut buf).expect("fast-path encode");
1111 buf
1112 }
1113
1114 fn read_official(bytes: &[u8]) -> ArrayData {
1117 let mut reader = StreamReader::try_new(Cursor::new(bytes), None).expect("open IPC stream");
1118 let batch = reader
1119 .next()
1120 .expect("one batch")
1121 .expect("batch decodes via official reader");
1122 assert_eq!(batch.num_columns(), 1);
1123 batch.column(0).to_data()
1124 }
1125
1126 fn aligned_buffer(bytes: &[u8]) -> (Buffer, usize, usize) {
1129 use aligned_vec::{AVec, ConstAlign};
1130 use std::ptr::NonNull;
1131 let mut aligned: AVec<u8, ConstAlign<128>> = AVec::__from_elem(128, 0, bytes.len());
1132 aligned.copy_from_slice(bytes);
1133 let base = aligned.as_ptr() as usize;
1134 let len = aligned.len();
1135 let ptr = NonNull::new(aligned.as_ptr() as *mut u8).unwrap();
1136 let buffer =
1138 unsafe { Buffer::from_custom_allocation(ptr, len, std::sync::Arc::new(aligned)) };
1139 (buffer, base, len)
1140 }
1141
1142 fn assert_fast_roundtrip(array: &ArrayData) {
1146 let encoded = fast_encode(array);
1147 let decoded = read_official(&encoded);
1148 assert_eq!(array, &decoded, "fast-path stream must decode to the input");
1149 let (buffer, _, _) = aligned_buffer(&encoded);
1151 let zc = decode_arrow_ipc_zero_copy_raw(buffer).expect("zero-copy decode");
1152 assert_eq!(array, &zc, "zero-copy decode must equal the input");
1153 }
1154
1155 fn batch_bytes(array: &ArrayData) -> Vec<u8> {
1158 let len = batch_fast_path_len_data(array).unwrap();
1159 let mut buf = vec![0u8; len];
1160 encode_batch_into_data(array, &mut buf).unwrap();
1161 buf
1162 }
1163
1164 fn batch_buf(array: &ArrayData) -> Buffer {
1165 Buffer::from_vec(batch_bytes(array))
1166 }
1167
1168 #[test]
1173 fn input_decoder_schema_then_batches() {
1174 let f32_schema =
1175 || Buffer::from_vec(encode_schema_message_for(&DataType::Float32).unwrap());
1176
1177 let mut dec = InputDecoder::new();
1178
1179 let early = Float32Array::from(vec![9.0]).into_data();
1181 assert!(
1182 dec.decode_batch_raw(batch_buf(&early), 7)
1183 .unwrap()
1184 .is_none()
1185 );
1186
1187 dec.set_schema_raw(7, f32_schema()).unwrap();
1189 for vals in [vec![1.0f32, 2.0, 3.0], vec![4.0], vec![5.0, 6.0, 7.0]] {
1190 let array = Float32Array::from(vals).into_data();
1191 assert_eq!(
1192 dec.decode_batch_raw(batch_buf(&array), 7).unwrap().unwrap(),
1193 array
1194 );
1195 }
1196
1197 let other = Float32Array::from(vec![8.0]).into_data();
1199 assert!(
1200 dec.decode_batch_raw(batch_buf(&other), 99)
1201 .unwrap()
1202 .is_none()
1203 );
1204
1205 dec.set_schema_raw(99, f32_schema()).unwrap();
1207 let after = Float32Array::from(vec![10.0, 11.0]).into_data();
1208 assert_eq!(
1209 dec.decode_batch_raw(batch_buf(&after), 99)
1210 .unwrap()
1211 .unwrap(),
1212 after
1213 );
1214 }
1215
1216 #[test]
1223 fn input_decoder_retains_multiple_schemas() {
1224 let schema_msg = |dt: &DataType| Buffer::from_vec(encode_schema_message_for(dt).unwrap());
1225
1226 let mut dec = InputDecoder::new();
1227 dec.set_schema_raw(1, schema_msg(&DataType::Float32))
1228 .unwrap();
1229 dec.set_schema_raw(2, schema_msg(&DataType::Int32)).unwrap();
1230
1231 let ints = Int32Array::from(vec![1, 2, 3]).into_data();
1233 assert_eq!(
1234 dec.decode_batch_raw(batch_buf(&ints), 2).unwrap().unwrap(),
1235 ints
1236 );
1237
1238 let floats = Float32Array::from(vec![4.0, 5.0]).into_data();
1240 assert_eq!(
1241 dec.decode_batch_raw(batch_buf(&floats), 1)
1242 .unwrap()
1243 .unwrap(),
1244 floats,
1245 "a schema installed earlier must be retained across later primes"
1246 );
1247
1248 let more_ints = Int32Array::from(vec![6]).into_data();
1250 assert_eq!(
1251 dec.decode_batch_raw(batch_buf(&more_ints), 2)
1252 .unwrap()
1253 .unwrap(),
1254 more_ints
1255 );
1256 }
1257
1258 #[test]
1263 fn input_decoder_handles_sequential_batches_arrow_59_terminal_state() {
1264 let f32_schema =
1265 || Buffer::from_vec(encode_schema_message_for(&DataType::Float32).unwrap());
1266
1267 let mut dec = InputDecoder::new();
1268 dec.set_schema_raw(7, f32_schema()).unwrap();
1270
1271 let batches = [
1275 Float32Array::from(vec![1.0, 2.0]).into_data(),
1276 Float32Array::from(vec![3.0]).into_data(),
1277 Float32Array::from(vec![4.0, 5.0, 6.0]).into_data(),
1278 Float32Array::from(vec![7.0, 8.0]).into_data(),
1279 Float32Array::from(vec![9.0, 10.0, 11.0, 12.0]).into_data(),
1280 ];
1281
1282 for (i, batch) in batches.iter().enumerate() {
1283 let result = dec.decode_batch_raw(batch_buf(batch), 7);
1284 assert!(
1285 result.is_ok(),
1286 "batch {} decode failed: {:?}",
1287 i,
1288 result.err()
1289 );
1290 let decoded = result.unwrap().unwrap();
1291 assert_eq!(
1292 &decoded, batch,
1293 "batch {} mismatch: expected {:?}, got {:?}",
1294 i, batch, decoded
1295 );
1296 }
1297
1298 assert_eq!(dec.schema_hash, Some(7));
1301 }
1302
1303 #[test]
1309 fn input_decoder_handles_empty_and_nonempty_batches() {
1310 let f32_schema =
1311 || Buffer::from_vec(encode_schema_message_for(&DataType::Float32).unwrap());
1312
1313 let mut dec = InputDecoder::new();
1314 dec.set_schema_raw(7, f32_schema()).unwrap();
1315
1316 let empty = Float32Array::from(Vec::<f32>::new()).into_data();
1317 let batches = [
1318 empty.clone(), Float32Array::from(vec![1.0, 2.0]).into_data(), empty.clone(),
1321 empty.clone(), Float32Array::from(vec![3.0]).into_data(),
1323 ];
1324
1325 for (i, batch) in batches.iter().enumerate() {
1326 let decoded = dec
1327 .decode_batch_raw(batch_buf(batch), 7)
1328 .unwrap_or_else(|e| panic!("batch {i} decode failed: {e:?}"))
1329 .unwrap_or_else(|| panic!("batch {i} was dropped"));
1330 assert_eq!(&decoded, batch, "batch {i} mismatch");
1331 assert_eq!(
1336 dec.schema_hash,
1337 Some(7),
1338 "batch {i} must not reset the decoder"
1339 );
1340 }
1341 }
1342
1343 #[test]
1346 fn input_decoder_evicts_oldest_schema_beyond_cap() {
1347 let f32_schema =
1348 || Buffer::from_vec(encode_schema_message_for(&DataType::Float32).unwrap());
1349
1350 let mut dec = InputDecoder::new();
1351 for hash in 0..=(MAX_RETAINED_SCHEMAS as u64) {
1354 dec.set_schema_raw(hash, f32_schema()).unwrap();
1355 }
1356 let array = Float32Array::from(vec![1.0]).into_data();
1357 assert!(
1358 dec.decode_batch_raw(batch_buf(&array), 0)
1359 .unwrap()
1360 .is_none(),
1361 "the oldest schema must be evicted beyond the cap"
1362 );
1363 assert_eq!(
1364 dec.decode_batch_raw(batch_buf(&array), 1).unwrap().unwrap(),
1365 array,
1366 "schemas within the cap must be retained"
1367 );
1368 }
1369
1370 #[test]
1375 fn decode_batch_resets_on_error_then_reprimes() {
1376 let mut dec = InputDecoder::new();
1377 dec.set_schema_raw(
1378 7,
1379 Buffer::from_vec(encode_schema_message_for(&DataType::Float32).unwrap()),
1380 )
1381 .unwrap();
1382
1383 let good = Float32Array::from(vec![4.0, 5.0]).into_data();
1385 assert_eq!(
1386 dec.decode_batch_raw(batch_buf(&good), 7).unwrap().unwrap(),
1387 good
1388 );
1389
1390 let dropped = Float32Array::from(vec![6.0, 7.0, 8.0, 9.0]).into_data();
1395 let mut truncated = batch_bytes(&dropped);
1396 truncated.truncate(truncated.len() / 2);
1397 assert!(
1398 dec.decode_batch_raw(Buffer::from_vec(truncated), 7)
1399 .is_err()
1400 );
1401
1402 let after = Float32Array::from(vec![10.0, 11.0]).into_data();
1405 assert_eq!(
1406 dec.decode_batch_raw(batch_buf(&after), 7).unwrap().unwrap(),
1407 after,
1408 "after a failed batch the decoder must re-prime from the retained schema"
1409 );
1410
1411 dec.reset();
1413 let final_batch = Float32Array::from(vec![12.0]).into_data();
1414 assert!(
1415 dec.decode_batch_raw(batch_buf(&final_batch), 7)
1416 .unwrap()
1417 .is_none(),
1418 "after a full reset the decoder must drop until a schema is re-installed"
1419 );
1420 }
1421
1422 #[test]
1432 fn input_decoder_dictionary_fallback_batch_sequence() {
1433 use arrow::array::DictionaryArray;
1434 use arrow::datatypes::Int32Type;
1435
1436 fn dict(words: &[&str]) -> ArrayData {
1439 let mut values: Vec<&str> = Vec::new();
1440 let mut keys: Vec<i32> = Vec::new();
1441 for w in words {
1442 let idx = values.iter().position(|v| v == w).unwrap_or_else(|| {
1443 values.push(*w);
1444 values.len() - 1
1445 });
1446 keys.push(idx as i32);
1447 }
1448 DictionaryArray::<Int32Type>::try_new(
1449 Int32Array::from(keys),
1450 Arc::new(StringArray::from(values)),
1451 )
1452 .unwrap()
1453 .into_data()
1454 }
1455
1456 let first = dict(&["a", "b", "a", "c", "b"]);
1457 assert!(
1459 ipc_fast_path_len_data(&first).is_none(),
1460 "dictionary must route to the official-writer fallback"
1461 );
1462
1463 let mut dec = InputDecoder::new();
1466 let full0 = encode_ipc_to_vec_data(&first).unwrap();
1467 let block = schema_block_len(&full0).unwrap();
1468 dec.set_schema_raw(1, Buffer::from(&full0[..block]))
1469 .unwrap();
1470
1471 let slice0 = batch_slice(&full0).expect("fallback stream is a valid IPC stream");
1476 assert_eq!(
1477 dec.decode_batch_raw(Buffer::from(slice0), 1)
1478 .unwrap()
1479 .expect("first batch decodes against the primed decoder"),
1480 first
1481 );
1482 for words in [
1483 ["x", "y", "x", "z"].as_slice(),
1484 ["b", "b"].as_slice(),
1485 ["new", "values", "entirely"].as_slice(),
1486 ] {
1487 let arr = dict(words);
1488 let full = encode_ipc_to_vec_data(&arr).unwrap();
1489 let slice = batch_slice(&full).expect("fallback stream is a valid IPC stream");
1490 let got = dec
1491 .decode_batch_raw(Buffer::from(slice), 1)
1492 .unwrap()
1493 .expect("batch must decode against the primed decoder");
1494 assert_eq!(
1495 got, arr,
1496 "replacement-dictionary batch must decode correctly"
1497 );
1498 }
1499 }
1500
1501 #[test]
1506 fn schema_primed_decoder_decodes_batch_sequence() {
1507 let schema = encode_schema_message_for(&DataType::Float32).unwrap();
1508 let mut decoder = StreamDecoder::new();
1509
1510 let mut sbuf = Buffer::from_vec(schema);
1512 while !sbuf.is_empty() {
1513 assert!(
1514 decoder.decode(&mut sbuf).unwrap().is_none(),
1515 "schema message must not yield a batch"
1516 );
1517 }
1518
1519 for vals in [vec![1.0f32, 2.0, 3.0], vec![4.0, 5.0], vec![6.0]] {
1520 let array = Float32Array::from(vals).into_data();
1521 let len = batch_fast_path_len_data(&array).unwrap();
1522 let mut buf = vec![0u8; len];
1523 encode_batch_into_data(&array, &mut buf).unwrap();
1524
1525 let mut bbuf = Buffer::from_vec(buf);
1526 let mut got = None;
1527 while !bbuf.is_empty() {
1528 if let Some(b) = decoder.decode(&mut bbuf).unwrap() {
1529 got = Some(b);
1530 break;
1531 }
1532 }
1533 assert_eq!(
1534 got.expect("batch message must decode against the primed decoder")
1535 .column(0)
1536 .to_data(),
1537 array
1538 );
1539 }
1540 }
1541
1542 #[test]
1546 fn uint8_ipc_header_constructs_in_place() {
1547 for data_len in [0usize, 1, 7, 8, 9, 1000] {
1548 let bytes: Vec<u8> = (0..data_len).map(|i| (i % 251) as u8).collect();
1549 let total = uint8_ipc_len(data_len).unwrap();
1550 let mut dst = vec![0u8; total];
1551 let offset = encode_uint8_ipc_header(&mut dst, data_len).unwrap();
1552 dst[offset..offset + data_len].copy_from_slice(&bytes);
1554
1555 let arr = arrow::array::make_array(read_official(&dst));
1557 let u8 = arr.as_any().downcast_ref::<UInt8Array>().unwrap();
1558 assert_eq!(u8.values(), bytes.as_slice(), "len {data_len}");
1559 assert_eq!(u8.null_count(), 0);
1560
1561 let array = UInt8Array::from(bytes).into_data();
1563 assert_eq!(dst, fast_encode(&array), "len {data_len}");
1564 }
1565 }
1566
1567 #[test]
1573 fn uint8_ipc_len_bounds_resulting_stream_not_input() {
1574 let max = super::super::MAX_IPC_BYTES;
1575
1576 let err = uint8_ipc_len(max).unwrap_err().to_string();
1579 assert!(
1580 err.contains("too large"),
1581 "data_len == MAX_IPC_BYTES should be rejected, got: {err}"
1582 );
1583
1584 let mut lo = 0usize;
1587 let mut hi = max;
1588 while lo < hi {
1589 let mid = lo + (hi - lo).div_ceil(2);
1590 match uint8_ipc_len(mid) {
1591 Ok(_) => lo = mid,
1592 Err(_) => hi = mid - 1,
1593 }
1594 }
1595 let largest_ok = lo;
1596 let total = uint8_ipc_len(largest_ok).unwrap();
1597 assert!(
1598 total <= max,
1599 "accepted payload of {largest_ok} bytes encodes to {total} > {max}"
1600 );
1601 assert!(uint8_ipc_len(largest_ok + 1).is_err());
1604 assert!(
1605 largest_ok < max,
1606 "boundary should sit below the input limit"
1607 );
1608 }
1609
1610 #[test]
1611 fn roundtrip_primitive_no_nulls() {
1612 let array = Float32Array::from((0..1000).map(|i| i as f32).collect::<Vec<_>>()).into_data();
1613 assert_fast_roundtrip(&array);
1614 }
1615
1616 #[test]
1617 fn roundtrip_primitive_with_nulls() {
1618 let array = UInt64Array::from(vec![Some(1), None, Some(3), None, Some(5)]).into_data();
1619 assert_fast_roundtrip(&array);
1620 }
1621
1622 #[test]
1623 fn roundtrip_empty_primitive() {
1624 let array = Int32Array::from(Vec::<i32>::new()).into_data();
1625 assert_fast_roundtrip(&array);
1626 }
1627
1628 #[test]
1634 fn uint8_header_zero_len_full_stream_roundtrip() {
1635 let len = uint8_ipc_len(0).unwrap();
1636 let mut buf = vec![0u8; len];
1637 let off = encode_uint8_ipc_header(&mut buf, 0).unwrap();
1638 assert_eq!(off, len - PREFIX_LEN);
1640 assert_eq!(read_official(&buf).len(), 0);
1641 let (buffer, _, _) = aligned_buffer(&buf);
1642 let decoded = decode_arrow_ipc_zero_copy_raw(buffer).unwrap();
1643 assert_eq!(decoded.data_type(), &DataType::UInt8);
1644 assert_eq!(decoded.len(), 0);
1645 }
1646
1647 #[test]
1653 fn schema_once_zero_len_batch_roundtrip() {
1654 let schema = || Buffer::from_vec(encode_schema_message_for(&DataType::UInt8).unwrap());
1655 let batch = |vals: &[u8]| {
1656 let a = UInt8Array::from(vals.to_vec()).into_data();
1657 let len = batch_fast_path_len_data(&a).unwrap();
1658 let mut b = vec![0u8; len];
1659 encode_batch_into_data(&a, &mut b).unwrap();
1660 Buffer::from_vec(b)
1661 };
1662
1663 let mut dec = InputDecoder::new();
1664 dec.set_schema_raw(1, schema()).unwrap();
1665
1666 let decoded = dec
1668 .decode_batch_raw(batch(&[]), 1)
1669 .unwrap()
1670 .expect("0-row batch must decode, not drop");
1671 assert_eq!(decoded.data_type(), &DataType::UInt8);
1672 assert_eq!(decoded.len(), 0);
1673
1674 for vals in [vec![1u8, 2, 3], vec![], vec![9u8], vec![]] {
1677 let d = dec.decode_batch_raw(batch(&vals), 1).unwrap().unwrap();
1678 assert_eq!(d.len(), vals.len());
1679 assert_eq!(d.data_type(), &DataType::UInt8);
1680 }
1681 }
1682
1683 #[test]
1688 fn schema_once_zero_len_via_batch_slice_roundtrip() {
1689 let total = uint8_ipc_len(0).unwrap();
1690 let mut full = vec![0u8; total];
1691 encode_uint8_ipc_header(&mut full, 0).unwrap();
1692
1693 let sblock = schema_block_len(&full).unwrap();
1694 let schema = Buffer::from(&full[..sblock]);
1695 let batch = batch_slice(&full).expect("batch slice of a valid stream");
1696
1697 let mut dec = InputDecoder::new();
1698 dec.set_schema_raw(7, schema).unwrap();
1699 let decoded = dec
1700 .decode_batch_raw(Buffer::from(batch), 7)
1701 .unwrap()
1702 .expect("0-row batch via batch_slice must decode, not drop");
1703 assert_eq!(decoded.data_type(), &DataType::UInt8);
1704 assert_eq!(decoded.len(), 0);
1705 }
1706
1707 #[test]
1713 fn schema_block_plus_batch_slice_reconstructs_full_stream() {
1714 let array = Int32Array::from(vec![1, 2, 3]).into_data();
1715 let full = fast_encode(&array);
1716 let sblock = schema_block_len(&full).unwrap();
1717 let schema = &full[..sblock];
1718 let batch = batch_slice(&full).expect("batch slice of a valid stream");
1719 let rebuilt = [schema, batch].concat();
1720 assert_eq!(
1721 rebuilt, full,
1722 "schema_block ++ batch_slice must equal the original stream"
1723 );
1724 assert_eq!(read_official(&rebuilt), array);
1725 }
1726
1727 #[test]
1728 fn roundtrip_boolean() {
1729 let array =
1730 BooleanArray::from(vec![true, false, true, true, false, false, true]).into_data();
1731 assert_fast_roundtrip(&array);
1732 }
1733
1734 #[test]
1735 fn roundtrip_utf8() {
1736 let array =
1737 StringArray::from(vec![Some("hello"), None, Some(""), Some("world!")]).into_data();
1738 assert_fast_roundtrip(&array);
1739 }
1740
1741 #[test]
1742 fn roundtrip_large_utf8_64bit_offsets() {
1743 let array = LargeStringArray::from(vec!["a", "bb", "ccc"]).into_data();
1744 assert_fast_roundtrip(&array);
1745 }
1746
1747 #[test]
1748 fn roundtrip_fixed_size_binary() {
1749 let values = vec![vec![1u8, 2, 3], vec![4, 5, 6], vec![7, 8, 9]];
1750 let array = FixedSizeBinaryArray::try_from_iter(values.into_iter())
1751 .unwrap()
1752 .into_data();
1753 assert_fast_roundtrip(&array);
1754 }
1755
1756 #[test]
1760 fn roundtrip_fixed_size_list() {
1761 use arrow::array::FixedSizeListArray;
1762 let values = Int32Array::from((0..12).collect::<Vec<_>>());
1763 let field = Arc::new(Field::new("item", DataType::Int32, true));
1764 let array = FixedSizeListArray::try_new(field, 3, Arc::new(values), None)
1765 .unwrap()
1766 .into_data();
1767 assert_fast_roundtrip(&array);
1768 }
1769
1770 #[test]
1773 fn roundtrip_fixed_size_list_with_nulls() {
1774 use arrow::array::FixedSizeListArray;
1775 use arrow::buffer::NullBuffer;
1776 let values = Int32Array::from((0..12).collect::<Vec<_>>());
1777 let field = Arc::new(Field::new("item", DataType::Int32, true));
1778 let nulls = NullBuffer::from(vec![true, false, true, true]);
1779 let array = FixedSizeListArray::try_new(field, 3, Arc::new(values), Some(nulls))
1780 .unwrap()
1781 .into_data();
1782 assert_fast_roundtrip(&array);
1783 }
1784
1785 #[test]
1787 fn roundtrip_decimal128() {
1788 use arrow::array::Decimal128Array;
1789 let array = Decimal128Array::from(vec![Some(12_345i128), None, Some(-9_876), Some(0)])
1790 .with_precision_and_scale(20, 4)
1791 .unwrap()
1792 .into_data();
1793 assert_fast_roundtrip(&array);
1794 }
1795
1796 #[test]
1798 fn roundtrip_timestamp_temporal() {
1799 use arrow::array::TimestampMicrosecondArray;
1800 let array =
1801 TimestampMicrosecondArray::from(vec![Some(1_000_000i64), None, Some(2_500_000)])
1802 .into_data();
1803 assert_fast_roundtrip(&array);
1804 }
1805
1806 #[test]
1807 fn roundtrip_struct_with_multilevel_nulls() {
1808 let array = StructArray::from(vec![
1809 (
1810 Arc::new(Field::new("a", DataType::UInt64, true)),
1811 Arc::new(UInt64Array::from(vec![Some(1), None, Some(3)])) as ArrayRef,
1812 ),
1813 (
1814 Arc::new(Field::new("b", DataType::Utf8, true)),
1815 Arc::new(StringArray::from(vec![Some("x"), Some("yy"), None])) as ArrayRef,
1816 ),
1817 ])
1818 .into_data();
1819 assert_fast_roundtrip(&array);
1820 }
1821
1822 #[test]
1827 fn roundtrip_struct_with_oversized_child() {
1828 use arrow_schema::Fields;
1829 let child = Int32Array::from(vec![10, 20, 30]).into_data(); let fields: Fields = vec![Field::new("v", DataType::Int32, false)].into();
1831 let struct_data = ArrayData::builder(DataType::Struct(fields))
1832 .len(2) .add_child_data(child)
1834 .build()
1835 .unwrap();
1836
1837 assert!(
1838 ipc_fast_path_len_data(&struct_data).is_some(),
1839 "struct with an oversized child should stay on the fast path"
1840 );
1841 let decoded = read_official(&fast_encode(&struct_data));
1842 assert_eq!(decoded.len(), 2);
1843 let arr = arrow::array::make_array(decoded);
1844 let sa = arr.as_any().downcast_ref::<StructArray>().unwrap();
1845 let col = sa.column(0).as_any().downcast_ref::<Int32Array>().unwrap();
1846 assert_eq!(
1847 col.values(),
1848 &[10, 20],
1849 "child must be truncated to the struct's len"
1850 );
1851 }
1852
1853 #[test]
1854 fn roundtrip_list_of_primitive() {
1855 let data = vec![
1856 Some(vec![Some(0), Some(1), Some(2)]),
1857 None,
1858 Some(vec![Some(3), None, Some(5)]),
1859 Some(vec![]),
1860 ];
1861 let array =
1862 ListArray::from_iter_primitive::<arrow::datatypes::Int32Type, _, _>(data).into_data();
1863 assert_fast_roundtrip(&array);
1864 }
1865
1866 #[test]
1867 fn roundtrip_nullarray_zero_and_n() {
1868 assert_fast_roundtrip(&NullArray::new(0).into_data());
1869 assert_fast_roundtrip(&NullArray::new(7).into_data());
1870 }
1871
1872 #[test]
1876 fn fast_path_decodes_zero_copy() {
1877 let array = UInt64Array::from((0..50_000u64).collect::<Vec<_>>()).into_data();
1878 let encoded = fast_encode(&array);
1879
1880 {
1882 let (mut buffer, _, _) = aligned_buffer(&encoded);
1883 let mut decoder = StreamDecoder::new().with_require_alignment(true);
1884 let mut got = None;
1885 while !buffer.is_empty() {
1886 if let Some(b) = decoder
1887 .decode(&mut buffer)
1888 .expect("aligned fast-path stream must decode without realignment")
1889 {
1890 got = Some(b);
1891 break;
1892 }
1893 }
1894 assert_eq!(got.unwrap().column(0).to_data(), array);
1895 }
1896
1897 {
1899 let (buffer, base, len) = aligned_buffer(&encoded);
1900 let decoded = decode_arrow_ipc_zero_copy_raw(buffer).unwrap();
1901 let ptr = decoded.buffers()[0].as_ptr() as usize;
1902 assert!(
1903 ptr >= base && ptr < base + len,
1904 "decoded data buffer at {ptr:#x} is outside input [{base:#x}, {:#x}) — a copy happened",
1905 base + len
1906 );
1907 }
1908 }
1909
1910 #[test]
1911 fn fast_path_len_matches_official_decode() {
1912 let array = Float32Array::from(vec![1.0, 2.0, 3.0, 4.0]).into_data();
1915 let encoded = fast_encode(&array);
1916 let mut reader = StreamReader::try_new(Cursor::new(&encoded[..]), None).unwrap();
1917 let _ = reader.next().unwrap().unwrap();
1918 assert!(reader.next().is_none(), "exactly one batch, fully consumed");
1919 }
1920
1921 #[test]
1924 fn sliced_array_routes_to_fallback_and_roundtrips() {
1925 let array = UInt64Array::from(vec![10, 20, 30, 40, 50])
1926 .into_data()
1927 .slice(2, 2); assert_eq!(array.offset(), 2);
1929 assert!(ipc_fast_path_len_data(&array).is_none());
1930
1931 let encoded = encode_ipc_to_vec_data(&array).unwrap();
1932 let decoded = read_official(&encoded);
1933 assert_eq!(array.len(), decoded.len());
1934 let dec = arrow::array::make_array(decoded);
1935 let dec = dec.as_any().downcast_ref::<UInt64Array>().unwrap();
1936 assert_eq!(dec.values(), &[30, 40]);
1937 }
1938
1939 #[test]
1944 fn sliced_offset_arrays_do_not_ship_the_parent_values() {
1945 use arrow::array::{Array, LargeBinaryArray, ListArray, StringArray};
1946 use arrow::datatypes::Float32Type;
1947
1948 let big = "x".repeat(1024);
1949 let strings = StringArray::from(vec![big.as_str(); 1000]);
1950 let head = strings.slice(0, 1).into_data();
1951 assert_eq!(head.offset(), 0, "array slice keeps offset 0");
1952 let len = ipc_fast_path_len_data(&head).expect("offsets start at 0: fast path");
1953 assert!(
1954 len < 4096,
1955 "one 1 KiB string must not encode to {len} bytes"
1956 );
1957 assert_fast_roundtrip(&head);
1958
1959 let middle = strings.slice(500, 2).into_data();
1962 assert!(ipc_fast_path_len_data(&middle).is_none());
1963 let encoded = encode_ipc_to_vec_data(&middle).unwrap();
1964 assert!(encoded.len() < 4096, "encoded {} bytes", encoded.len());
1965 assert_eq!(middle, read_official(&encoded));
1966
1967 let blobs = vec![vec![7u8; 64 * 1024]; 64];
1968 let blobs = LargeBinaryArray::from_iter_values(blobs.iter());
1969 let head = blobs.slice(0, 2).into_data();
1970 let len = ipc_fast_path_len_data(&head).expect("fast path");
1971 assert!(
1972 len < 3 * 64 * 1024,
1973 "two 64 KiB blobs encoded to {len} bytes"
1974 );
1975 assert_fast_roundtrip(&head);
1976
1977 let lists = ListArray::from_iter_primitive::<Float32Type, _, _>(
1978 (0..100).map(|i| Some((0..1000).map(move |j| Some((i * j) as f32)))),
1979 );
1980 let head = lists.slice(0, 1).into_data();
1981 let len = ipc_fast_path_len_data(&head).expect("fast path");
1982 assert!(len < 8 * 1024, "one 1000-float list encoded to {len} bytes");
1983 assert_fast_roundtrip(&head);
1984 }
1985
1986 #[test]
1988 fn view_type_routes_to_fallback() {
1989 use arrow::array::StringViewArray;
1990 let array = StringViewArray::from(vec!["a", "bb", "ccc"]).into_data();
1991 assert!(
1992 ipc_fast_path_len_data(&array).is_none(),
1993 "Utf8View is not fast-path eligible"
1994 );
1995 let encoded = encode_ipc_to_vec_data(&array).unwrap();
1996 let decoded = read_official(&encoded);
1997 assert_eq!(array, decoded);
1998 }
1999
2000 #[test]
2002 fn fuzz_roundtrip_many_shapes() {
2003 let mut state: u64 = 0x1234_5678_9abc_def0;
2005 let mut next = || {
2006 state = state
2007 .wrapping_mul(6364136223846793005)
2008 .wrapping_add(1442695040888963407);
2009 state
2010 };
2011
2012 for _ in 0..200 {
2013 let len = (next() % 64) as usize;
2014 let kind = next() % 6;
2015 let array: ArrayData = match kind {
2016 0 => UInt8Array::from(
2017 (0..len)
2018 .map(|i| {
2019 if next().is_multiple_of(4) {
2020 None
2021 } else {
2022 Some((i as u8).wrapping_add(1))
2023 }
2024 })
2025 .collect::<Vec<_>>(),
2026 )
2027 .into_data(),
2028 1 => Float32Array::from((0..len).map(|i| i as f32 * 0.5).collect::<Vec<_>>())
2029 .into_data(),
2030 2 => BooleanArray::from(
2031 (0..len)
2032 .map(|i| (i + next() as usize).is_multiple_of(2))
2033 .collect::<Vec<_>>(),
2034 )
2035 .into_data(),
2036 3 => StringArray::from(
2037 (0..len)
2038 .map(|i| {
2039 if next().is_multiple_of(5) {
2040 None
2041 } else {
2042 Some("x".repeat(i % 7))
2043 }
2044 })
2045 .collect::<Vec<_>>(),
2046 )
2047 .into_data(),
2048 4 => Int32Array::from(
2049 (0..len)
2050 .map(|i| {
2051 if next().is_multiple_of(3) {
2052 None
2053 } else {
2054 Some(i as i32 - 10)
2055 }
2056 })
2057 .collect::<Vec<_>>(),
2058 )
2059 .into_data(),
2060 _ => StructArray::from(vec![(
2061 Arc::new(Field::new("v", DataType::Int32, true)),
2062 Arc::new(Int32Array::from(
2063 (0..len).map(|i| Some(i as i32)).collect::<Vec<_>>(),
2064 )) as ArrayRef,
2065 )])
2066 .into_data(),
2067 };
2068 assert_fast_roundtrip(&array);
2069 }
2070 }
2071}