wp_connector_utils/arrow/
encode.rs

1//! Arrow IPC encoding (sink direction).
2//!
3//! - [`encode_batch_ipc_stream`] — bare Arrow IPC Stream bytes
4//! - [`encode_ipc_frame`] — wp_arrow frame `[tag_len][tag][IPC]`
5//! - [`encode_ipc_frame_multi`] — multiple batches in one frame
6
7use arrow::ipc::writer::StreamWriter;
8use arrow::record_batch::RecordBatch;
9use orion_error::conversion::SourceRawErr;
10use wp_connector_api::{SinkReason, SinkResult};
11
12// ---------------------------------------------------------------------------
13// Encode (sink direction)
14// ---------------------------------------------------------------------------
15
16pub fn encode_batch_ipc_stream(batch: &RecordBatch) -> SinkResult<Vec<u8>> {
17    let schema = batch.schema();
18    let mut buf = Vec::new();
19    {
20        let mut writer = StreamWriter::try_new(&mut buf, &schema)
21            .source_raw_err(SinkReason::Sink, "arrow create stream writer")?;
22        writer
23            .write(batch)
24            .source_raw_err(SinkReason::Sink, "arrow encode batch")?;
25        writer
26            .finish()
27            .source_raw_err(SinkReason::Sink, "arrow finish stream")?;
28    }
29    Ok(buf)
30}
31
32pub fn encode_ipc_frame(tag: &str, batch: &RecordBatch) -> SinkResult<Vec<u8>> {
33    let tag_bytes = tag.as_bytes();
34    let schema = batch.schema();
35    let mut buf = Vec::with_capacity(4 + tag_bytes.len() + 1024);
36    buf.extend_from_slice(&(tag_bytes.len() as u32).to_be_bytes());
37    buf.extend_from_slice(tag_bytes);
38    {
39        let mut writer = StreamWriter::try_new(&mut buf, &schema)
40            .source_raw_err(SinkReason::Sink, "arrow create framed writer")?;
41        writer
42            .write(batch)
43            .source_raw_err(SinkReason::Sink, "arrow encode framed")?;
44        writer
45            .finish()
46            .source_raw_err(SinkReason::Sink, "arrow finish framed")?;
47    }
48    Ok(buf)
49}
50
51pub fn encode_ipc_frame_multi(tag: &str, batches: &[RecordBatch]) -> SinkResult<Vec<u8>> {
52    if batches.is_empty() {
53        let tag_bytes = tag.as_bytes();
54        let mut buf = Vec::with_capacity(4 + tag_bytes.len());
55        buf.extend_from_slice(&(tag_bytes.len() as u32).to_be_bytes());
56        buf.extend_from_slice(tag_bytes);
57        return Ok(buf);
58    }
59    let tag_bytes = tag.as_bytes();
60    let schema = batches[0].schema();
61    let mut buf = Vec::with_capacity(4 + tag_bytes.len() + 1024);
62    buf.extend_from_slice(&(tag_bytes.len() as u32).to_be_bytes());
63    buf.extend_from_slice(tag_bytes);
64    {
65        let mut writer = StreamWriter::try_new(&mut buf, &schema)
66            .source_raw_err(SinkReason::Sink, "arrow create framed multi writer")?;
67        for batch in batches {
68            writer
69                .write(batch)
70                .source_raw_err(SinkReason::Sink, "arrow encode framed multi")?;
71        }
72        writer
73            .finish()
74            .source_raw_err(SinkReason::Sink, "arrow finish framed multi")?;
75    }
76    Ok(buf)
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82    use arrow::array::StringArray;
83    use arrow::datatypes::{DataType, Field, Schema};
84    use std::sync::Arc;
85
86    fn make_schema() -> Arc<Schema> {
87        Arc::new(Schema::new(vec![Field::new("x", DataType::Utf8, false)]))
88    }
89
90    fn make_batch(schema: &Arc<Schema>, values: Vec<&str>) -> RecordBatch {
91        RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))]).unwrap()
92    }
93
94    #[test]
95    fn ipc_roundtrip() {
96        let s = make_schema();
97        let b = make_batch(&s, vec!["hi"]);
98        let ipc = encode_batch_ipc_stream(&b).unwrap();
99        // Arrow IPC Stream starts with 0xFFFFFFFF (continuation marker)
100        assert!(ipc.len() > 8);
101        assert_eq!(&ipc[0..4], &[0xff, 0xff, 0xff, 0xff]);
102    }
103
104    #[test]
105    fn framed_roundtrip() {
106        let s = make_schema();
107        let b = make_batch(&s, vec!["hi"]);
108        let frame = encode_ipc_frame("my_tag", &b).unwrap();
109        // tag_len = 6, tag = "my_tag"
110        assert_eq!(&frame[0..4], &6u32.to_be_bytes());
111        assert_eq!(&frame[4..10], b"my_tag");
112    }
113
114    #[test]
115    fn ipc_frame_multi_roundtrip() {
116        let s = make_schema();
117        let b1 = make_batch(&s, vec!["a"]);
118        let b2 = make_batch(&s, vec!["b", "c"]);
119        let frame = encode_ipc_frame_multi("multi", &[b1, b2]).unwrap();
120        assert_eq!(&frame[0..4], &5u32.to_be_bytes());
121        assert_eq!(&frame[4..9], b"multi");
122    }
123
124    #[test]
125    fn ipc_frame_multi_empty() {
126        let frame = encode_ipc_frame_multi("tag", &[]).unwrap();
127        assert_eq!(&frame[0..4], &3u32.to_be_bytes());
128        assert_eq!(&frame[4..7], b"tag");
129    }
130}
wp_connector_utils/arrow/encode.rs

wp_connector_utils/arrow/
encode.rs
