1use std::collections::HashMap;
32use std::io::{Read, Write};
33use std::sync::Arc;
34
35use arrow_array::{Array, LargeBinaryArray, RecordBatch};
36use arrow_buffer::Buffer as ArrowBuffer;
37use arrow_ipc::reader as ipc_reader;
38use arrow_ipc::writer::{write_message, DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
39use arrow_ipc::{convert as ipc_convert, root_as_message, MessageHeader};
40use arrow_schema::{DataType, Schema, SchemaRef};
41use flatbuffers::FlatBufferBuilder;
42
43use crate::errors::{Result, RpcError};
44
45pub type Metadata = HashMap<String, String>;
49
50pub(crate) const INVALID_UTF8_METADATA_KEY: &str = "\0vgi_rpc.invalid_utf8_metadata";
55
56#[inline]
58pub fn md_get<'a>(md: &'a Metadata, key: &str) -> Option<&'a str> {
59 md.get(key).map(String::as_str)
60}
61
62pub const MAX_IPC_SCHEMA_BYTES: usize = 16 * 1024 * 1024;
70
71pub const MAX_IPC_MESSAGE_BYTES: usize = u32::MAX as usize;
89
90const BODY_PREALLOC_LIMIT: usize = 8 * 1024 * 1024;
103
104const CONTINUATION_MARKER: [u8; 4] = [0xFF, 0xFF, 0xFF, 0xFF];
109
110const MAX_WRITE_CHUNK: usize = 1 << 30; const IPC_ALIGNMENT: usize = 64;
120const IPC_PADDING: [u8; IPC_ALIGNMENT] = [0; IPC_ALIGNMENT];
121
122struct ChunkedWriter<W: Write> {
150 inner: W,
151 limit: usize,
152}
153
154impl<W: Write> ChunkedWriter<W> {
155 fn new(inner: W) -> Self {
156 Self {
157 inner,
158 limit: MAX_WRITE_CHUNK,
159 }
160 }
161
162 #[cfg(test)]
165 fn with_limit(inner: W, limit: usize) -> Self {
166 Self { inner, limit }
167 }
168
169 fn get_mut(&mut self) -> &mut W {
170 &mut self.inner
171 }
172}
173
174impl<W: Write> Write for ChunkedWriter<W> {
175 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
176 let end = buf.len().min(self.limit);
177 self.inner.write(&buf[..end])
178 }
179
180 fn flush(&mut self) -> std::io::Result<()> {
181 self.inner.flush()
182 }
183}
184
185pub struct StreamWriter<W: Write> {
199 writer: ChunkedWriter<W>,
200 schema: SchemaRef,
201 opts: IpcWriteOptions,
202 data_gen: IpcDataGenerator,
203 dict_tracker: DictionaryTracker,
204 finished: bool,
205 fbb: FlatBufferBuilder<'static>,
209}
210
211impl<W: Write> StreamWriter<W> {
212 pub fn new(writer: W, schema: &Schema) -> Result<Self> {
214 let mut writer = ChunkedWriter::new(writer);
215 let opts = IpcWriteOptions::default();
216 let data_gen = IpcDataGenerator::default();
217 let mut dict_tracker = DictionaryTracker::new(false);
218 let encoded =
219 data_gen.schema_to_bytes_with_dictionary_tracker(schema, &mut dict_tracker, &opts);
220 write_message(&mut writer, encoded, &opts)?;
221 Ok(Self {
222 writer,
223 schema: Arc::new(schema.clone()),
224 opts,
225 data_gen,
226 dict_tracker,
227 finished: false,
228 fbb: FlatBufferBuilder::new(),
229 })
230 }
231
232 pub fn write(&mut self, batch: &RecordBatch, metadata: Option<&Metadata>) -> Result<()> {
236 if self.finished {
237 return Err(RpcError::new("IOError", "writer already finished"));
238 }
239 if write_single_large_binary_direct(
240 &mut self.writer,
241 &mut self.fbb,
242 batch,
243 metadata,
244 &self.opts,
245 )? {
246 return Ok(());
247 }
248 let mut ctx = Default::default();
249 let (dicts, data) = self
250 .data_gen
251 .encode(batch, &mut self.dict_tracker, &self.opts, &mut ctx)
252 .map_err(RpcError::from)?;
253 for d in dicts {
254 write_message(&mut self.writer, d, &self.opts).map_err(RpcError::from)?;
255 }
256 if let Some(md) = metadata.filter(|m| !m.is_empty()) {
257 self.fbb.reset();
258 repack_record_batch_message_with_metadata(&mut self.fbb, &data.ipc_message, md)?;
259 let encoded = arrow_ipc::writer::EncodedData {
260 ipc_message: self.fbb.finished_data().to_vec(),
261 arrow_data: data.arrow_data,
262 };
263 write_message(&mut self.writer, encoded, &self.opts).map_err(RpcError::from)?;
264 } else {
265 write_message(&mut self.writer, data, &self.opts).map_err(RpcError::from)?;
266 }
267 Ok(())
268 }
269
270 pub fn schema(&self) -> SchemaRef {
272 self.schema.clone()
273 }
274
275 pub fn finish(&mut self) -> Result<()> {
277 if self.finished {
278 return Ok(());
279 }
280 self.writer.write_all(&CONTINUATION_MARKER)?;
281 self.writer.write_all(&[0u8; 4])?;
282 self.writer.flush()?;
283 self.finished = true;
284 Ok(())
285 }
286
287 pub fn flush(&mut self) -> Result<()> {
289 self.writer.flush()?;
290 Ok(())
291 }
292
293 pub fn get_mut(&mut self) -> &mut W {
294 self.writer.get_mut()
295 }
296}
297
298fn write_single_large_binary_direct<W: Write>(
308 writer: &mut ChunkedWriter<W>,
309 fbb: &mut FlatBufferBuilder<'static>,
310 batch: &RecordBatch,
311 metadata: Option<&Metadata>,
312 opts: &IpcWriteOptions,
313) -> Result<bool> {
314 use arrow_ipc::{
315 Buffer as IpcBuffer, FieldNode, KeyValue, KeyValueArgs, MessageBuilder, MetadataVersion,
316 RecordBatchBuilder,
317 };
318
319 if batch.num_rows() != 1
320 || batch.num_columns() != 1
321 || batch.schema().field(0).data_type() != &DataType::LargeBinary
322 {
323 return Ok(false);
324 }
325 let Some(array) = batch.column(0).as_any().downcast_ref::<LargeBinaryArray>() else {
326 return Ok(false);
327 };
328 if array.is_null(0) {
329 return Ok(false);
330 }
331
332 let value = array.value(0);
333 let value_len = i64::try_from(value.len())
334 .map_err(|_| RpcError::value_error("large_binary value exceeds i64 offsets"))?;
335 let value_padding = padding_to_alignment(value.len(), IPC_ALIGNMENT);
336 let body_len = 2usize
337 .checked_mul(IPC_ALIGNMENT)
338 .and_then(|prefix| prefix.checked_add(value.len()))
339 .and_then(|length| length.checked_add(value_padding))
340 .ok_or_else(|| RpcError::value_error("large_binary IPC body length overflow"))?;
341 let body_len_i64 = i64::try_from(body_len)
342 .map_err(|_| RpcError::value_error("large_binary IPC body exceeds i64 length"))?;
343
344 fbb.reset();
345 let nodes = fbb.create_vector(&[FieldNode::new(1, 0)]);
346 let buffers = fbb.create_vector(&[
347 IpcBuffer::new(0, 1),
348 IpcBuffer::new(IPC_ALIGNMENT as i64, 16),
349 IpcBuffer::new((2 * IPC_ALIGNMENT) as i64, value_len),
350 ]);
351 let record_batch = {
352 let mut builder = RecordBatchBuilder::new(fbb);
353 builder.add_length(1);
354 builder.add_nodes(nodes);
355 builder.add_buffers(buffers);
356 builder.finish()
357 };
358
359 let custom_metadata = metadata
360 .filter(|entries| !entries.is_empty())
361 .map(|entries| {
362 let entries: Vec<_> = entries
363 .iter()
364 .map(|(key, value)| {
365 let key = fbb.create_string(key);
366 let value = fbb.create_string(value);
367 KeyValue::create(
368 fbb,
369 &KeyValueArgs {
370 key: Some(key),
371 value: Some(value),
372 },
373 )
374 })
375 .collect();
376 fbb.create_vector(&entries)
377 });
378
379 let message = {
380 let mut builder = MessageBuilder::new(fbb);
381 builder.add_version(MetadataVersion::V5);
382 builder.add_header_type(MessageHeader::RecordBatch);
383 builder.add_header(record_batch.as_union_value());
384 builder.add_bodyLength(body_len_i64);
385 if let Some(custom_metadata) = custom_metadata {
386 builder.add_custom_metadata(custom_metadata);
387 }
388 builder.finish()
389 };
390 fbb.finish(message, None);
391 let encoded = arrow_ipc::writer::EncodedData {
392 ipc_message: fbb.finished_data().to_vec(),
393 arrow_data: Vec::new(),
394 };
395 write_message(&mut *writer, encoded, opts).map_err(RpcError::from)?;
396
397 writer.write_all(&[0xff])?;
399 writer.write_all(&IPC_PADDING[..IPC_ALIGNMENT - 1])?;
400
401 writer.write_all(&0_i64.to_le_bytes())?;
403 writer.write_all(&value_len.to_le_bytes())?;
404 writer.write_all(&IPC_PADDING[..IPC_ALIGNMENT - 16])?;
405
406 writer.write_all(value)?;
408 writer.write_all(&IPC_PADDING[..value_padding])?;
409 Ok(true)
410}
411
412#[inline]
413fn padding_to_alignment(length: usize, alignment: usize) -> usize {
414 (alignment - (length % alignment)) % alignment
415}
416
417impl<W: Write> Drop for StreamWriter<W> {
418 fn drop(&mut self) {
419 let _ = self.finish();
420 }
421}
422
423fn repack_record_batch_message_with_metadata(
426 fbb: &mut FlatBufferBuilder<'static>,
427 msg_bytes: &[u8],
428 metadata: &Metadata,
429) -> Result<()> {
430 use arrow_ipc::{KeyValue, KeyValueArgs, MessageBuilder, RecordBatchBuilder};
431
432 let msg = root_as_message(msg_bytes)
433 .map_err(|e| RpcError::new("IPC", format!("parsing message: {e}")))?;
434 let version = msg.version();
435 let header_type = msg.header_type();
436 let body_length = msg.bodyLength();
437 if header_type != MessageHeader::RecordBatch {
438 return Err(RpcError::new(
439 "IPC",
440 format!("repack expected RecordBatch header, got {header_type:?}"),
441 ));
442 }
443 let rb = msg
444 .header_as_record_batch()
445 .ok_or_else(|| RpcError::new("IPC", "missing RecordBatch header"))?;
446
447 let src_nodes = rb
451 .nodes()
452 .ok_or_else(|| RpcError::new("IPC", "RecordBatch missing nodes"))?;
453 let nodes_vec = fbb.create_vector_from_iter(src_nodes.iter());
454
455 let src_buffers = rb
456 .buffers()
457 .ok_or_else(|| RpcError::new("IPC", "RecordBatch missing buffers"))?;
458 let buffers_vec = fbb.create_vector_from_iter(src_buffers.iter());
459
460 let variadic_vec = rb
461 .variadicBufferCounts()
462 .map(|v| fbb.create_vector_from_iter(v.iter()));
463
464 let new_rb = {
465 let mut b = RecordBatchBuilder::new(fbb);
466 b.add_length(rb.length());
467 b.add_nodes(nodes_vec);
468 b.add_buffers(buffers_vec);
469 if let Some(v) = variadic_vec {
470 b.add_variadicBufferCounts(v);
471 }
472 b.finish()
475 };
476
477 let kvs: Vec<_> = metadata
481 .iter()
482 .map(|(k, v)| {
483 let k_off = fbb.create_string(k);
484 let v_off = fbb.create_string(v);
485 KeyValue::create(
486 fbb,
487 &KeyValueArgs {
488 key: Some(k_off),
489 value: Some(v_off),
490 },
491 )
492 })
493 .collect();
494 let md_vec = fbb.create_vector(&kvs);
495
496 let mut mb = MessageBuilder::new(fbb);
497 mb.add_version(version);
498 mb.add_header_type(header_type);
499 mb.add_header(new_rb.as_union_value());
500 mb.add_bodyLength(body_length);
501 mb.add_custom_metadata(md_vec);
502 let m = mb.finish();
503 fbb.finish(m, None);
504 Ok(())
505}
506
507pub struct StreamReader<R: Read> {
517 reader: R,
518 schema: SchemaRef,
519 dictionaries: HashMap<i64, arrow_array::ArrayRef>,
520 finished: bool,
521 relaxed_schema: Option<SchemaRef>,
526}
527
528impl<R: Read> StreamReader<R> {
529 pub fn new(mut reader: R) -> Result<Self> {
536 let msg = read_message_bytes(&mut reader, MAX_IPC_SCHEMA_BYTES)?
537 .ok_or_else(|| RpcError::new("IPC", "empty IPC stream (no schema)"))?;
538 if msg.had_invalid_utf8 {
539 return Err(RpcError::protocol_error(
540 "Invalid UTF-8 in IPC schema metadata",
541 ));
542 }
543 let msg_fb = root_as_message(&msg.message_bytes)
544 .map_err(|e| RpcError::new("IPC", format!("parse schema message: {e}")))?;
545 if msg_fb.header_type() != MessageHeader::Schema {
546 return Err(RpcError::new(
547 "IPC",
548 format!("expected Schema, got {:?}", msg_fb.header_type()),
549 ));
550 }
551 let ipc_schema = msg_fb
552 .header_as_schema()
553 .ok_or_else(|| RpcError::new("IPC", "bad schema header"))?;
554 if ipc_schema.fields().is_none() {
561 return Err(RpcError::new("IPC", "schema message has no fields vector"));
562 }
563 let schema = decode_guard("schema message", || ipc_convert::fb_to_schema(ipc_schema))?;
568 Ok(Self {
569 reader,
570 schema: Arc::new(schema),
571 dictionaries: HashMap::new(),
572 finished: false,
573 relaxed_schema: None,
574 })
575 }
576
577 pub fn schema(&self) -> SchemaRef {
580 self.relaxed_schema
581 .clone()
582 .unwrap_or_else(|| self.schema.clone())
583 }
584
585 pub fn relax_nullability(mut self) -> Self {
591 self.relaxed_schema = Some(Arc::new(relax_schema_nullability(self.schema.as_ref())));
592 self
593 }
594
595 pub fn read_next(&mut self) -> Result<Option<(RecordBatch, Metadata)>> {
600 if self.finished {
601 return Ok(None);
602 }
603 loop {
604 let msg = match read_message_bytes(&mut self.reader, MAX_IPC_MESSAGE_BYTES)? {
605 Some(m) => m,
606 None => {
607 self.finished = true;
608 return Ok(None);
609 }
610 };
611 let msg_fb = root_as_message(&msg.message_bytes)
612 .map_err(|e| RpcError::new("IPC", format!("parse message: {e}")))?;
613 let version = msg_fb.version();
614 match msg_fb.header_type() {
615 MessageHeader::DictionaryBatch => {
616 let dict = msg_fb
617 .header_as_dictionary_batch()
618 .ok_or_else(|| RpcError::new("IPC", "bad dictionary header"))?;
619 let body_buf = ArrowBuffer::from_vec(msg.body);
620 if let Some(data) = dict.data() {
624 validate_record_batch_buffers(&data, body_buf.len())?;
625 }
626 decode_guard("dictionary batch", || {
630 ipc_reader::read_dictionary(
631 &body_buf,
632 dict,
633 self.schema.as_ref(),
634 &mut self.dictionaries,
635 &version,
636 )
637 })?
638 .map_err(RpcError::from)?;
639 }
640 MessageHeader::RecordBatch => {
641 let rb_fb = msg_fb
642 .header_as_record_batch()
643 .ok_or_else(|| RpcError::new("IPC", "bad record batch header"))?;
644 let body_buf = ArrowBuffer::from_vec(msg.body);
645 validate_record_batch_buffers(&rb_fb, body_buf.len())?;
646 let decode_schema = self
653 .relaxed_schema
654 .clone()
655 .unwrap_or_else(|| self.schema.clone());
656 let batch = decode_guard("record batch", || {
657 ipc_reader::read_record_batch(
658 &body_buf,
659 rb_fb,
660 decode_schema,
661 &self.dictionaries,
662 None,
663 &version,
664 )
665 })?
666 .map_err(RpcError::from)?;
667 let mut metadata = parse_custom_metadata(&msg_fb);
668 if msg.had_invalid_utf8 {
669 metadata.insert(INVALID_UTF8_METADATA_KEY.into(), "true".into());
670 }
671 return Ok(Some((batch, metadata)));
672 }
673 MessageHeader::Schema => {
674 return Err(RpcError::new("IPC", "unexpected schema message mid-stream"));
675 }
676 MessageHeader::NONE => continue,
677 other => {
678 return Err(RpcError::new(
679 "IPC",
680 format!("unsupported message type {other:?}"),
681 ));
682 }
683 }
684 }
685 }
686
687 pub fn drain(&mut self) -> Result<()> {
689 while self.read_next()?.is_some() {}
690 Ok(())
691 }
692
693 pub fn get_mut(&mut self) -> &mut R {
694 &mut self.reader
695 }
696}
697
698fn parse_custom_metadata(msg: &arrow_ipc::Message) -> Metadata {
699 let Some(md) = msg.custom_metadata() else {
700 return Metadata::new();
701 };
702 let mut out = Metadata::with_capacity(md.len());
704 for kv in md.iter() {
705 let k = kv.key().unwrap_or("").to_string();
706 let v = kv.value().unwrap_or("").to_string();
707 out.insert(k, v);
708 }
709 out
710}
711
712pub(crate) fn validate_record_batch_buffers(
720 rb: &arrow_ipc::RecordBatch,
721 body_len: usize,
722) -> Result<()> {
723 if let Some(buffers) = rb.buffers() {
724 for buf in buffers.iter() {
725 let offset = buf.offset();
726 let length = buf.length();
727 if offset < 0 || length < 0 {
728 return Err(RpcError::new("IPC", "negative IPC buffer descriptor"));
729 }
730 let end = (offset as u64)
731 .checked_add(length as u64)
732 .ok_or_else(|| RpcError::new("IPC", "IPC buffer descriptor overflows"))?;
733 if end > body_len as u64 {
734 return Err(RpcError::new(
735 "IPC",
736 "IPC buffer descriptor exceeds message body",
737 ));
738 }
739 }
740 }
741 Ok(())
742}
743
744pub(crate) fn decode_guard<T>(what: &str, f: impl FnOnce() -> T) -> Result<T> {
749 std::panic::catch_unwind(std::panic::AssertUnwindSafe(f))
750 .map_err(|_| RpcError::new("IPC", format!("panic decoding {what} (malformed frame)")))
751}
752
753struct RawMessage {
754 message_bytes: Vec<u8>,
755 body: Vec<u8>,
756 had_invalid_utf8: bool,
757}
758
759fn read_exact(r: &mut impl Read, buf: &mut [u8]) -> Result<bool> {
760 let mut read = 0;
761 while read < buf.len() {
762 match r.read(&mut buf[read..]) {
763 Ok(0) => {
764 if read == 0 {
765 return Ok(false);
766 }
767 return Err(RpcError::new("IOError", "unexpected EOF in IPC message"));
768 }
769 Ok(n) => read += n,
770 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
771 Err(e) => return Err(e.into()),
772 }
773 }
774 Ok(true)
775}
776
777fn read_message_bytes(r: &mut impl Read, max_bytes: usize) -> Result<Option<RawMessage>> {
788 let mut prefix = [0u8; 4];
789 if !read_exact(r, &mut prefix)? {
790 return Ok(None);
791 }
792 let size_bytes = if prefix == CONTINUATION_MARKER {
793 let mut sb = [0u8; 4];
794 if !read_exact(r, &mut sb)? {
795 return Ok(None);
796 }
797 sb
798 } else {
799 prefix
800 };
801 let size = u32::from_le_bytes(size_bytes) as usize;
802 if size == 0 {
803 return Ok(None);
805 }
806 if size > max_bytes {
807 return Err(RpcError::new(
808 "IPC",
809 format!(
810 "IPC message header length {size} bytes exceeds cap {max_bytes} — \
811 refusing to allocate before parsing"
812 ),
813 ));
814 }
815 let mut message_bytes = vec![0u8; size];
816 if !read_exact(r, &mut message_bytes)? {
817 return Err(RpcError::new("IOError", "unexpected EOF in message body"));
818 }
819 let mut had_invalid_utf8 = false;
825 let msg = loop {
826 match root_as_message(&message_bytes) {
827 Ok(msg) => break msg,
828 Err(flatbuffers::InvalidFlatbuffer::Utf8Error { range, .. })
829 if !range.is_empty() && range.end <= message_bytes.len() =>
830 {
831 message_bytes[range].fill(b'?');
832 had_invalid_utf8 = true;
833 }
834 Err(e) => {
835 return Err(RpcError::new("IPC", format!("parse message header: {e}")));
836 }
837 }
838 };
839 let body_length_signed = msg.bodyLength();
840 if body_length_signed < 0 {
841 return Err(RpcError::new(
842 "IPC",
843 format!("IPC message has negative bodyLength ({body_length_signed})"),
844 ));
845 }
846 if body_length_signed as u64 > max_bytes as u64 {
850 return Err(RpcError::new(
851 "IPC",
852 format!(
853 "IPC message bodyLength {body_length_signed} bytes exceeds cap {max_bytes} — \
854 refusing to allocate before parsing"
855 ),
856 ));
857 }
858 let body_length = body_length_signed as usize;
859 let mut body = Vec::with_capacity(body_length.min(BODY_PREALLOC_LIMIT));
863 if body_length > 0 {
864 let read = (&mut *r)
865 .take(body_length as u64)
866 .read_to_end(&mut body)
867 .map_err(RpcError::from)?;
868 if read != body_length {
869 return Err(RpcError::new("IOError", "unexpected EOF in message body"));
870 }
871 }
872 Ok(Some(RawMessage {
873 message_bytes,
874 body,
875 had_invalid_utf8,
876 }))
877}
878
879pub fn write_one_batch(batch: &RecordBatch, metadata: Option<&Metadata>) -> Result<Vec<u8>> {
886 write_one_batch_as(batch, batch.schema().as_ref(), metadata)
887}
888
889pub fn write_one_batch_as(
910 batch: &RecordBatch,
911 schema: &Schema,
912 metadata: Option<&Metadata>,
913) -> Result<Vec<u8>> {
914 let mut buf = Vec::new();
915 {
916 let mut w = StreamWriter::new(&mut buf, schema)?;
917 w.write(batch, metadata)?;
918 w.finish()?;
919 }
920 Ok(buf)
921}
922
923#[allow(dead_code)]
928pub(crate) fn bytes_to_hex(bytes: &[u8]) -> String {
929 const HEX: &[u8; 16] = b"0123456789abcdef";
930 let mut out = String::with_capacity(bytes.len() * 2);
931 for b in bytes {
932 out.push(HEX[(b >> 4) as usize] as char);
933 out.push(HEX[(b & 0x0f) as usize] as char);
934 }
935 out
936}
937
938fn relax_field_nullability(f: &arrow_schema::Field) -> arrow_schema::Field {
939 use arrow_schema::DataType;
940 let dt = match f.data_type() {
941 DataType::List(inner) => DataType::List(Arc::new(relax_field_nullability(inner))),
942 DataType::LargeList(inner) => DataType::LargeList(Arc::new(relax_field_nullability(inner))),
943 DataType::FixedSizeList(inner, n) => {
944 DataType::FixedSizeList(Arc::new(relax_field_nullability(inner)), *n)
945 }
946 DataType::Struct(fields) => DataType::Struct(
947 fields
948 .iter()
949 .map(|child| Arc::new(relax_field_nullability(child)))
950 .collect(),
951 ),
952 other => other.clone(),
956 };
957 #[allow(deprecated)]
958 let new_field = if let DataType::Dictionary(_, _) = f.data_type() {
959 arrow_schema::Field::new_dict(
960 f.name(),
961 dt,
962 true,
963 f.dict_id().unwrap_or(0),
964 f.dict_is_ordered().unwrap_or(false),
965 )
966 } else {
967 arrow_schema::Field::new(f.name(), dt, true)
968 };
969 new_field.with_metadata(f.metadata().clone())
970}
971
972fn relax_schema_nullability(s: &Schema) -> Schema {
973 let new_fields: Vec<arrow_schema::Field> = s
974 .fields()
975 .iter()
976 .map(|f| relax_field_nullability(f))
977 .collect();
978 Schema::new_with_metadata(new_fields, s.metadata().clone())
979}
980
981pub fn empty_batch(schema: &Schema) -> Result<RecordBatch> {
983 use arrow_array::array::new_empty_array;
984 use arrow_array::RecordBatchOptions;
985 let cols: Vec<arrow_array::ArrayRef> = schema
986 .fields()
987 .iter()
988 .map(|f| new_empty_array(f.data_type()))
989 .collect();
990 RecordBatch::try_new_with_options(
991 Arc::new(schema.clone()),
992 cols,
993 &RecordBatchOptions::new().with_row_count(Some(0)),
994 )
995 .map_err(RpcError::from)
996}
997
998#[cfg(test)]
999mod tests {
1000 use super::*;
1001 use arrow_array::{Int64Array, StringArray};
1002 use arrow_schema::{DataType, Field};
1003
1004 struct SabotageWriter {
1008 offered: Vec<usize>,
1009 accept: usize,
1010 sink: Vec<u8>,
1011 }
1012
1013 impl Write for SabotageWriter {
1014 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
1015 self.offered.push(buf.len());
1016 let n = buf.len().min(self.accept);
1017 self.sink.extend_from_slice(&buf[..n]);
1018 Ok(n)
1019 }
1020 fn flush(&mut self) -> std::io::Result<()> {
1021 Ok(())
1022 }
1023 }
1024
1025 #[test]
1026 fn chunked_writer_clamps_and_retries() {
1027 let mut w = ChunkedWriter::with_limit(
1032 SabotageWriter {
1033 offered: Vec::new(),
1034 accept: 3,
1035 sink: Vec::new(),
1036 },
1037 8,
1038 );
1039 let payload: Vec<u8> = (0..50u8).collect();
1040 w.write_all(&payload).unwrap();
1041 let inner = w.get_mut();
1042 assert!(
1043 inner.offered.iter().all(|n| *n <= 8),
1044 "a write was offered more than the clamp: {:?}",
1045 inner.offered
1046 );
1047 assert_eq!(inner.sink, payload, "short writes lost bytes");
1048 }
1049
1050 #[test]
1051 fn write_chunk_stays_under_int_max() {
1052 assert!(MAX_WRITE_CHUNK < i32::MAX as usize);
1055 }
1056
1057 #[test]
1058 fn oversized_body_claim_costs_nothing_to_refuse() {
1059 assert!(MAX_IPC_MESSAGE_BYTES as u64 > (1u64 << 31) + 1);
1063 assert!((0x4000000100000u64) > MAX_IPC_MESSAGE_BYTES as u64);
1064 }
1065
1066 #[test]
1067 fn roundtrip_with_metadata() {
1068 let schema = Schema::new(vec![
1069 Field::new("idx", DataType::Int64, false),
1070 Field::new("name", DataType::Utf8, false),
1071 ]);
1072 let batch = RecordBatch::try_new(
1073 Arc::new(schema.clone()),
1074 vec![
1075 Arc::new(Int64Array::from(vec![1, 2, 3])) as _,
1076 Arc::new(StringArray::from(vec!["a", "b", "c"])) as _,
1077 ],
1078 )
1079 .unwrap();
1080
1081 let mut buf: Vec<u8> = Vec::new();
1082 {
1083 let mut w = StreamWriter::new(&mut buf, &schema).unwrap();
1084 let mut md = Metadata::new();
1085 md.insert("vgi_rpc.method".into(), "echo_string".into());
1086 w.write(&batch, Some(&md)).unwrap();
1087 w.finish().unwrap();
1088 }
1089
1090 let mut r = StreamReader::new(buf.as_slice()).unwrap();
1091 let (rb, md) = r.read_next().unwrap().expect("batch");
1092 assert_eq!(rb.num_rows(), 3);
1093 assert_eq!(md_get(&md, "vgi_rpc.method"), Some("echo_string"));
1094 assert!(r.read_next().unwrap().is_none());
1095 }
1096
1097 #[test]
1098 fn single_large_binary_direct_write_roundtrips_sliced_value_and_metadata() {
1099 use arrow_array::LargeBinaryArray;
1100
1101 let schema = Schema::new(vec![Field::new("result", DataType::LargeBinary, false)]);
1102 let source = LargeBinaryArray::from_iter_values([
1103 b"discarded prefix".as_slice(),
1104 b"payload retained in the original Arrow allocation".as_slice(),
1105 ]);
1106 let sliced = source.slice(1, 1);
1107 let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(sliced)]).unwrap();
1108 let mut metadata = Metadata::new();
1109 metadata.insert("vgi_rpc.method".into(), "echo_large_binary".into());
1110
1111 let mut encoded = Vec::new();
1112 {
1113 let mut writer = StreamWriter::new(&mut encoded, &schema).unwrap();
1114 writer.write(&batch, Some(&metadata)).unwrap();
1115 writer.finish().unwrap();
1116 }
1117
1118 let mut reader = StreamReader::new(encoded.as_slice()).unwrap();
1119 let (decoded, decoded_metadata) = reader.read_next().unwrap().expect("batch");
1120 let decoded = decoded
1121 .column(0)
1122 .as_any()
1123 .downcast_ref::<LargeBinaryArray>()
1124 .unwrap();
1125 assert_eq!(
1126 decoded.value(0),
1127 b"payload retained in the original Arrow allocation"
1128 );
1129 assert_eq!(
1130 md_get(&decoded_metadata, "vgi_rpc.method"),
1131 Some("echo_large_binary")
1132 );
1133 assert!(reader.read_next().unwrap().is_none());
1134 }
1135
1136 #[test]
1137 fn zero_row_metadata_only() {
1138 let schema = Schema::empty();
1139 let batch = empty_batch(&schema).unwrap();
1140
1141 let mut buf: Vec<u8> = Vec::new();
1142 {
1143 let mut w = StreamWriter::new(&mut buf, &schema).unwrap();
1144 let mut md = Metadata::new();
1145 md.insert("vgi_rpc.log_level".into(), "INFO".into());
1146 w.write(&batch, Some(&md)).unwrap();
1147 w.finish().unwrap();
1148 }
1149 let mut r = StreamReader::new(buf.as_slice()).unwrap();
1150 let (rb, md) = r.read_next().unwrap().expect("batch");
1151 assert_eq!(rb.num_rows(), 0);
1152 assert_eq!(md_get(&md, "vgi_rpc.log_level"), Some("INFO"));
1153 }
1154
1155 #[test]
1156 fn rejects_oversize_schema_length_prefix() {
1157 let bomb: &[u8] = &[0x1A, 0x2C, 0xF5, 0x2C];
1161 let err = StreamReader::new(bomb).err().expect("must reject");
1162 assert!(
1163 err.message.contains("exceeds cap"),
1164 "unexpected error: {err:?}"
1165 );
1166 }
1167
1168 #[test]
1169 fn rejects_oversize_message_bodylength() {
1170 use arrow_ipc::{Buffer as FbBuffer, FieldNode, MessageBuilder, RecordBatchBuilder};
1174 let schema = Schema::new(vec![Field::new("v", DataType::Int64, false)]);
1176 let mut buf: Vec<u8> = Vec::new();
1177 {
1178 let w = StreamWriter::new(&mut buf, &schema).unwrap();
1179 std::mem::forget(w);
1183 }
1184 let mut fbb = FlatBufferBuilder::new();
1187 let nodes_vec = fbb.create_vector(&[FieldNode::new(0, 0)]);
1188 let buffers_vec = fbb.create_vector(&[FbBuffer::new(0, 0)]);
1189 let rb_off = {
1190 let mut b = RecordBatchBuilder::new(&mut fbb);
1191 b.add_length(0);
1192 b.add_nodes(nodes_vec);
1193 b.add_buffers(buffers_vec);
1194 b.finish()
1195 };
1196 let msg_off = {
1197 let mut mb = MessageBuilder::new(&mut fbb);
1198 mb.add_version(arrow_ipc::MetadataVersion::V5);
1199 mb.add_header_type(MessageHeader::RecordBatch);
1200 mb.add_header(rb_off.as_union_value());
1201 mb.add_bodyLength(MAX_IPC_MESSAGE_BYTES as i64 + 1);
1202 mb.finish()
1203 };
1204 fbb.finish(msg_off, None);
1205 let msg_bytes = fbb.finished_data();
1206 buf.extend_from_slice(&CONTINUATION_MARKER);
1208 buf.extend_from_slice(&(msg_bytes.len() as u32).to_le_bytes());
1209 buf.extend_from_slice(msg_bytes);
1210 let mut r = StreamReader::new(buf.as_slice()).unwrap();
1213 let err = r.read_next().expect_err("must reject");
1214 assert!(
1215 err.message.contains("bodyLength") && err.message.contains("exceeds cap"),
1216 "unexpected error: {err:?}"
1217 );
1218 }
1219
1220 #[test]
1221 fn malformed_schema_message_is_error_not_panic() {
1222 let crash: &[u8] = &[
1230 22, 0, 0, 0, 12, 0, 0, 0, 0, 0, 8, 0, 4, 0, 0, 0, 12, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 134,
1231 ];
1232 let _ = StreamReader::new(crash);
1235 }
1236
1237 #[test]
1238 fn rejects_buffer_descriptor_past_body() {
1239 use arrow_ipc::{Buffer as FbBuffer, FieldNode, MessageBuilder, RecordBatchBuilder};
1244 let schema = Schema::new(vec![Field::new("v", DataType::Int64, false)]);
1245 let mut buf: Vec<u8> = Vec::new();
1246 {
1247 let w = StreamWriter::new(&mut buf, &schema).unwrap();
1248 std::mem::forget(w);
1249 }
1250 let mut fbb = FlatBufferBuilder::new();
1251 let nodes_vec = fbb.create_vector(&[FieldNode::new(1, 0)]);
1252 let buffers_vec = fbb.create_vector(&[FbBuffer::new(0, 1000)]);
1254 let rb_off = {
1255 let mut b = RecordBatchBuilder::new(&mut fbb);
1256 b.add_length(1);
1257 b.add_nodes(nodes_vec);
1258 b.add_buffers(buffers_vec);
1259 b.finish()
1260 };
1261 let msg_off = {
1262 let mut mb = MessageBuilder::new(&mut fbb);
1263 mb.add_version(arrow_ipc::MetadataVersion::V5);
1264 mb.add_header_type(MessageHeader::RecordBatch);
1265 mb.add_header(rb_off.as_union_value());
1266 mb.add_bodyLength(8);
1267 mb.finish()
1268 };
1269 fbb.finish(msg_off, None);
1270 let msg_bytes = fbb.finished_data().to_vec();
1271 buf.extend_from_slice(&CONTINUATION_MARKER);
1272 buf.extend_from_slice(&(msg_bytes.len() as u32).to_le_bytes());
1273 buf.extend_from_slice(&msg_bytes);
1274 buf.extend_from_slice(&[0u8; 8]); let mut r = StreamReader::new(buf.as_slice()).unwrap();
1277 let err = r.read_next().expect_err("must reject");
1278 assert!(
1279 err.message.contains("buffer descriptor"),
1280 "unexpected error: {err:?}"
1281 );
1282 }
1283}