rs_hosts2arrow_ipc_stream/
lib.rs

1use std::io;
2use std::sync::Arc;
3
4use arrow::record_batch::RecordBatch;
5
6use arrow_array::builder::{ListBuilder, StringBuilder};
7use arrow_schema::{DataType, Field, Schema};
8use hostfile::HostEntry;
9
10pub fn hosts2batch(hosts: Vec<HostEntry>) -> Result<RecordBatch, io::Error> {
11    let mut ip_builder = StringBuilder::new();
12    let mut names_builder = ListBuilder::new(StringBuilder::new());
13
14    for host in hosts {
15        ip_builder.append_value(host.ip.to_string());
16        let names = host.names.iter().map(|s| Some(s.as_str()));
17        names_builder.append_value(names);
18    }
19
20    let ip_array = ip_builder.finish();
21    let names_array = names_builder.finish();
22
23    let schema = Schema::new(vec![
24        Field::new("ip", DataType::Utf8, false),
25        Field::new(
26            "names",
27            DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
28            true,
29        ),
30    ]);
31
32    let batch = RecordBatch::try_new(
33        Arc::new(schema),
34        vec![Arc::new(ip_array), Arc::new(names_array)],
35    )
36    .map_err(|e| io::Error::other(e.to_string()))?;
37
38    Ok(batch)
39}
40
41#[cfg(test)]
42mod tests {
43    use super::*;
44    use arrow_array::cast::AsArray;
45    use std::net::{IpAddr, Ipv4Addr};
46
47    #[test]
48    fn test_hosts2batch() {
49        let host_entries = vec![
50            HostEntry {
51                ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
52                names: vec!["localhost".to_string(), "loopback".to_string()],
53            },
54            HostEntry {
55                ip: IpAddr::V4(Ipv4Addr::new(192, 168, 1, 100)),
56                names: vec!["myhost".to_string()],
57            },
58        ];
59
60        let batch = hosts2batch(host_entries).unwrap();
61
62        assert_eq!(batch.num_rows(), 2);
63        assert_eq!(batch.num_columns(), 2);
64
65        let ip_array = batch.column(0).as_string::<i32>();
66        let names_array = batch.column(1).as_list::<i32>();
67
68        assert_eq!(ip_array.value(0), "127.0.0.1");
69        assert_eq!(ip_array.value(1), "192.168.1.100");
70
71        let names_0 = names_array.value(0);
72        assert_eq!(names_0.as_string::<i32>().value(0), "localhost");
73        assert_eq!(names_0.as_string::<i32>().value(1), "loopback");
74
75        let names_1 = names_array.value(1);
76        assert_eq!(names_1.as_string::<i32>().value(0), "myhost");
77    }
78}