iridis_message/
helper.rs

1//! Some helper functions for working with Arrow data types and arrays.
2
3use std::{collections::HashMap, sync::Arc};
4
5use crate::prelude::*;
6
7use arrow_array::*;
8use arrow_buffer::*;
9use arrow_data::*;
10use arrow_schema::*;
11
12/// Helper function to create a union field.
13pub fn make_union_fields(name: impl Into<String>, fields: Vec<Field>) -> Field {
14    Field::new(
15        name,
16        DataType::Union(
17            UnionFields::new(0..fields.len() as i8, fields),
18            UnionMode::Dense,
19        ),
20        false,
21    )
22}
23
24/// Helper function to unpack a union array.
25pub fn unpack_union(data: ArrayData) -> (HashMap<String, usize>, Vec<ArrayRef>) {
26    let (fields, _, _, children) = UnionArray::from(data).into_parts();
27
28    let map = fields
29        .iter()
30        .map(|(id, field)| (field.name().into(), id as usize))
31        .collect::<HashMap<String, usize>>();
32
33    (map, children)
34}
35
36/// Helper function to extract data from a union array.
37pub fn extract_union_data<T: ArrowMessage>(
38    field: &str,
39    map: &HashMap<String, usize>,
40    children: &[ArrayRef],
41) -> Result<T> {
42    T::try_from_arrow(
43        children
44            .get(
45                *map.get(field)
46                    .ok_or(eyre::eyre!("Field {} not found", field))?,
47            )
48            .ok_or(eyre::eyre!("Field {} not found", field))?
49            .into_data(),
50    )
51}
52
53/// Helper function to get the union fields of a message.
54pub fn get_union_fields<T: ArrowMessage>() -> Result<UnionFields> {
55    match T::field("").data_type() {
56        DataType::Union(fields, _) => Ok(fields.clone()),
57        _ => Err(eyre::eyre!("Expected Union data type")),
58    }
59}
60
61/// Helper function to create a union array.
62pub fn make_union_array(union_fields: UnionFields, children: Vec<ArrayRef>) -> Result<ArrayRef> {
63    UnionArray::try_new(
64        union_fields,
65        ScalarBuffer::from(vec![]),
66        Some(ScalarBuffer::from(vec![])),
67        children,
68    )
69    .map(|union| Arc::new(union) as ArrayRef)
70    .map_err(eyre::Report::msg)
71}
iridis_message/helper.rs

iridis_message/
helper.rs
