dora_core/
metadata.rs

1use dora_message::{
2    arrow_data::ArrayData,
3    arrow_schema::DataType,
4    metadata::{ArrowTypeInfo, BufferOffset},
5};
6use eyre::Context;
7
8pub trait ArrowTypeInfoExt {
9    fn empty() -> Self;
10    fn byte_array(data_len: usize) -> Self;
11    unsafe fn from_array(
12        array: &ArrayData,
13        region_start: *const u8,
14        region_len: usize,
15    ) -> eyre::Result<Self>
16    where
17        Self: Sized;
18}
19
20impl ArrowTypeInfoExt for ArrowTypeInfo {
21    fn empty() -> Self {
22        Self {
23            data_type: DataType::Null,
24            len: 0,
25            null_count: 0,
26            validity: None,
27            offset: 0,
28            buffer_offsets: Vec::new(),
29            child_data: Vec::new(),
30        }
31    }
32
33    fn byte_array(data_len: usize) -> Self {
34        Self {
35            data_type: DataType::UInt8,
36            len: data_len,
37            null_count: 0,
38            validity: None,
39            offset: 0,
40            buffer_offsets: vec![BufferOffset {
41                offset: 0,
42                len: data_len,
43            }],
44            child_data: Vec::new(),
45        }
46    }
47
48    unsafe fn from_array(
49        array: &ArrayData,
50        region_start: *const u8,
51        region_len: usize,
52    ) -> eyre::Result<Self> {
53        Ok(Self {
54            data_type: array.data_type().clone(),
55            len: array.len(),
56            null_count: array.null_count(),
57            validity: array.nulls().map(|b| b.validity().to_owned()),
58            offset: array.offset(),
59            buffer_offsets: array
60                .buffers()
61                .iter()
62                .map(|b| {
63                    let ptr = b.as_ptr();
64                    if ptr as usize <= region_start as usize {
65                        eyre::bail!("ptr {ptr:p} starts before region {region_start:p}");
66                    }
67                    if ptr as usize >= region_start as usize + region_len {
68                        eyre::bail!("ptr {ptr:p} starts after region {region_start:p}");
69                    }
70                    if ptr as usize + b.len() > region_start as usize + region_len {
71                        eyre::bail!("ptr {ptr:p} ends after region {region_start:p}");
72                    }
73                    let offset = usize::try_from(unsafe { ptr.offset_from(region_start) })
74                        .context("offset_from is negative")?;
75
76                    Result::<_, eyre::Report>::Ok(BufferOffset {
77                        offset,
78                        len: b.len(),
79                    })
80                })
81                .collect::<Result<_, _>>()?,
82            child_data: array
83                .child_data()
84                .iter()
85                .map(|c| unsafe { Self::from_array(c, region_start, region_len) })
86                .collect::<Result<_, _>>()?,
87        })
88    }
89}
dora_core/metadata.rs

dora_core/
metadata.rs
