wp_arrow/
ipc.rs

1use arrow::array::RecordBatch;
2use arrow::ipc::reader::StreamReader;
3use arrow::ipc::writer::StreamWriter;
4
5use crate::error::WpArrowError;
6
7/// Decoded IPC data frame.
8pub struct IpcFrame {
9    pub tag: String,
10    pub batch: RecordBatch,
11}
12
13/// Encode a RecordBatch into an IPC frame: `[4B tag_len BE][tag bytes][Arrow IPC stream]`.
14pub fn encode_ipc(tag: &str, batch: &RecordBatch) -> Result<Vec<u8>, WpArrowError> {
15    let tag_bytes = tag.as_bytes();
16    let tag_len = tag_bytes.len() as u32;
17
18    let mut buf = Vec::new();
19    buf.extend_from_slice(&tag_len.to_be_bytes());
20    buf.extend_from_slice(tag_bytes);
21
22    let mut writer = StreamWriter::try_new(&mut buf, batch.schema().as_ref())
23        .map_err(|e| WpArrowError::IpcEncodeError(e.to_string()))?;
24    writer
25        .write(batch)
26        .map_err(|e| WpArrowError::IpcEncodeError(e.to_string()))?;
27    writer
28        .finish()
29        .map_err(|e| WpArrowError::IpcEncodeError(e.to_string()))?;
30
31    Ok(buf)
32}
33
34/// Decode a complete IPC frame from bytes.
35pub fn decode_ipc(data: &[u8]) -> Result<IpcFrame, WpArrowError> {
36    if data.len() < 4 {
37        return Err(WpArrowError::IpcDecodeError(format!(
38            "frame too short: {} bytes, minimum 4",
39            data.len()
40        )));
41    }
42
43    let tag_len = u32::from_be_bytes(data[0..4].try_into().unwrap()) as usize;
44    let tag_end = 4 + tag_len;
45    if data.len() < tag_end {
46        return Err(WpArrowError::IpcDecodeError(format!(
47            "frame truncated: tag_len={tag_len} but only {} bytes remain after header",
48            data.len() - 4
49        )));
50    }
51
52    let tag = String::from_utf8(data[4..tag_end].to_vec())
53        .map_err(|e| WpArrowError::IpcDecodeError(format!("invalid UTF-8 in tag: {e}")))?;
54
55    let ipc_payload = &data[tag_end..];
56    let mut reader = StreamReader::try_new(ipc_payload, None)
57        .map_err(|e| WpArrowError::IpcDecodeError(e.to_string()))?;
58
59    let batch = reader
60        .next()
61        .ok_or_else(|| WpArrowError::IpcDecodeError("no RecordBatch in IPC payload".to_string()))?
62        .map_err(|e| WpArrowError::IpcDecodeError(e.to_string()))?;
63
64    Ok(IpcFrame { tag, batch })
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use arrow::array::{Int32Array, StringArray};
71    use arrow::datatypes::{DataType, Field, Schema};
72    use std::sync::Arc;
73
74    fn make_batch(num_rows: usize) -> RecordBatch {
75        let schema = Arc::new(Schema::new(vec![
76            Field::new("id", DataType::Int32, false),
77            Field::new("name", DataType::Utf8, true),
78        ]));
79        let ids: Vec<i32> = (0..num_rows as i32).collect();
80        let names: Vec<Option<&str>> = (0..num_rows)
81            .map(|i| if i % 2 == 0 { Some("even") } else { None })
82            .collect();
83        RecordBatch::try_new(
84            schema,
85            vec![
86                Arc::new(Int32Array::from(ids)),
87                Arc::new(StringArray::from(names)),
88            ],
89        )
90        .unwrap()
91    }
92
93    #[test]
94    fn ipc_roundtrip_basic() {
95        let batch = make_batch(5);
96        let encoded = encode_ipc("test-tag", &batch).unwrap();
97        let frame = decode_ipc(&encoded).unwrap();
98        assert_eq!(frame.tag, "test-tag");
99        assert_eq!(frame.batch.num_rows(), batch.num_rows());
100        assert_eq!(frame.batch.num_columns(), batch.num_columns());
101        assert_eq!(frame.batch.schema(), batch.schema());
102        assert_eq!(frame.batch, batch);
103    }
104
105    #[test]
106    fn ipc_roundtrip_empty_batch() {
107        let batch = make_batch(0);
108        let encoded = encode_ipc("empty", &batch).unwrap();
109        let frame = decode_ipc(&encoded).unwrap();
110        assert_eq!(frame.batch.num_rows(), 0);
111        assert_eq!(frame.batch, batch);
112    }
113
114    #[test]
115    fn ipc_roundtrip_large_batch() {
116        let batch = make_batch(1000);
117        let encoded = encode_ipc("large", &batch).unwrap();
118        let frame = decode_ipc(&encoded).unwrap();
119        assert_eq!(frame.batch.num_rows(), 1000);
120        assert_eq!(frame.batch, batch);
121    }
122
123    #[test]
124    fn ipc_tag_preserved() {
125        let batch = make_batch(1);
126        let encoded = encode_ipc("my-tag", &batch).unwrap();
127        let frame = decode_ipc(&encoded).unwrap();
128        assert_eq!(frame.tag, "my-tag");
129    }
130
131    #[test]
132    fn ipc_utf8_tag() {
133        let batch = make_batch(1);
134        let tag = "数据标签-🚀";
135        let encoded = encode_ipc(tag, &batch).unwrap();
136        let frame = decode_ipc(&encoded).unwrap();
137        assert_eq!(frame.tag, tag);
138    }
139
140    #[test]
141    fn ipc_empty_tag() {
142        let batch = make_batch(1);
143        let encoded = encode_ipc("", &batch).unwrap();
144        let frame = decode_ipc(&encoded).unwrap();
145        assert_eq!(frame.tag, "");
146    }
147
148    #[test]
149    fn decode_ipc_too_short() {
150        let result = decode_ipc(&[0x00; 2]);
151        assert!(matches!(result, Err(WpArrowError::IpcDecodeError(_))));
152    }
153
154    #[test]
155    fn decode_ipc_truncated_tag() {
156        let mut data = Vec::new();
157        data.extend_from_slice(&100u32.to_be_bytes()); // tag_len = 100 but no data
158        let result = decode_ipc(&data);
159        assert!(matches!(result, Err(WpArrowError::IpcDecodeError(_))));
160    }
161}
162
163#[test]
164fn decode_roundtrip_preserves_memory() {
165    use crate::ipc::{decode_ipc, encode_ipc};
166    use arrow::array::{Array, StringArray};
167    use arrow::datatypes::{DataType, Field, Schema};
168    use std::sync::Arc;
169
170    let n = 75000;
171    let vals: Vec<Option<String>> = (0..n)
172        .map(|i| Some(format!("10.0.{}.{}", i % 4, i % 250 + 1)))
173        .collect();
174    let schema = Arc::new(Schema::new(vec![Field::new("sip", DataType::Utf8, true)]));
175    let batch = RecordBatch::try_new(schema, vec![Arc::new(StringArray::from(vals))]).unwrap();
176    eprintln!(
177        "original col bytes = {}",
178        batch.column(0).get_array_memory_size()
179    );
180
181    let enc = encode_ipc("conn_events", &batch).unwrap();
182    eprintln!("ipc payload bytes = {}", enc.len());
183    let frame = decode_ipc(&enc).unwrap();
184    eprintln!(
185        "decoded col bytes = {}",
186        frame.batch.column(0).get_array_memory_size()
187    );
188}
189
190// 12 列批次（模拟 wfgen conn 批次）往返内存
191#[test]
192fn decode_roundtrip_12col_memory() {
193    use crate::ipc::{decode_ipc, encode_ipc};
194    use arrow::array::{Array, Int64Array, StringArray};
195    use arrow::datatypes::{DataType, Field, Schema};
196    use std::sync::Arc;
197
198    let n = 75000;
199    let str_vals: Vec<Option<String>> = (0..n)
200        .map(|i| Some(format!("10.0.{}.{}", i % 4, i % 250 + 1)))
201        .collect();
202    let int_vals: Vec<i64> = (0..n).map(|i| (i % 4096) as i64).collect();
203    let schema = Arc::new(Schema::new(vec![
204        Field::new("sip", DataType::Utf8, true),
205        Field::new("dip", DataType::Utf8, true),
206        Field::new("roles_obj", DataType::Utf8, true),
207        Field::new("dport", DataType::Int64, true),
208        Field::new("bytes", DataType::Int64, true),
209        Field::new("event_time", DataType::Int64, true),
210    ]));
211    let batch = RecordBatch::try_new(
212        schema,
213        vec![
214            Arc::new(StringArray::from(str_vals.clone())) as arrow::array::ArrayRef,
215            Arc::new(StringArray::from(str_vals.clone())) as arrow::array::ArrayRef,
216            Arc::new(StringArray::from(str_vals.clone())) as arrow::array::ArrayRef,
217            Arc::new(Int64Array::from(int_vals.clone())) as arrow::array::ArrayRef,
218            Arc::new(Int64Array::from(int_vals.clone())) as arrow::array::ArrayRef,
219            Arc::new(Int64Array::from(int_vals)) as arrow::array::ArrayRef,
220        ],
221    )
222    .unwrap();
223    eprintln!("orig batch bytes = {}", batch.get_array_memory_size());
224    eprintln!("orig sip col = {}", batch.column(0).get_array_memory_size());
225    let enc = encode_ipc("conn_events", &batch).unwrap();
226    eprintln!("ipc payload = {}", enc.len());
227    let frame = decode_ipc(&enc).unwrap();
228    eprintln!(
229        "decoded batch bytes = {}",
230        frame.batch.get_array_memory_size()
231    );
232    eprintln!(
233        "decoded sip col = {}",
234        frame.batch.column(0).get_array_memory_size()
235    );
236}
wp_arrow/ipc.rs

wp_arrow/
ipc.rs
