isaac_sim_arrow/
odometry.rs

1// SPDX-License-Identifier: MPL-2.0
2//! Arrow encoder and decoder for the chassis odometry channel.
3use std::sync::{Arc, OnceLock};
4
5use arrow::array::{Array, ArrayRef, Float64Array, Int64Array, StringArray, StructArray};
6use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
7use arrow::record_batch::RecordBatch;
8
9/// Borrowed view of a single chassis odometry sample, used as input to [`to_record_batch`].
10#[allow(missing_docs)]
11pub struct Odometry<'a> {
12    pub chassis_frame_id: &'a str,
13    pub odom_frame_id: &'a str,
14    pub position_x: f64,
15    pub position_y: f64,
16    pub position_z: f64,
17    pub orientation_w: f64,
18    pub orientation_x: f64,
19    pub orientation_y: f64,
20    pub orientation_z: f64,
21    pub lin_vel_x: f64,
22    pub lin_vel_y: f64,
23    pub lin_vel_z: f64,
24    pub ang_vel_x: f64,
25    pub ang_vel_y: f64,
26    pub ang_vel_z: f64,
27    pub timestamp_ns: i64,
28}
29
30/// Owned variant returned by [`from_struct_array`].
31#[derive(Debug, Clone, PartialEq)]
32#[allow(missing_docs)]
33pub struct OdometryOwned {
34    pub chassis_frame_id: String,
35    pub odom_frame_id: String,
36    pub position_x: f64,
37    pub position_y: f64,
38    pub position_z: f64,
39    pub orientation_w: f64,
40    pub orientation_x: f64,
41    pub orientation_y: f64,
42    pub orientation_z: f64,
43    pub lin_vel_x: f64,
44    pub lin_vel_y: f64,
45    pub lin_vel_z: f64,
46    pub ang_vel_x: f64,
47    pub ang_vel_y: f64,
48    pub ang_vel_z: f64,
49    pub timestamp_ns: i64,
50}
51
52/// Stable Arrow schema for an `Odometry` record batch.
53pub fn schema() -> SchemaRef {
54    static SCHEMA: OnceLock<SchemaRef> = OnceLock::new();
55    SCHEMA
56        .get_or_init(|| {
57            Arc::new(Schema::new(vec![
58                Field::new("chassis_frame_id", DataType::Utf8, false),
59                Field::new("odom_frame_id", DataType::Utf8, false),
60                Field::new("position_x", DataType::Float64, false),
61                Field::new("position_y", DataType::Float64, false),
62                Field::new("position_z", DataType::Float64, false),
63                Field::new("orientation_w", DataType::Float64, false),
64                Field::new("orientation_x", DataType::Float64, false),
65                Field::new("orientation_y", DataType::Float64, false),
66                Field::new("orientation_z", DataType::Float64, false),
67                Field::new("lin_vel_x", DataType::Float64, false),
68                Field::new("lin_vel_y", DataType::Float64, false),
69                Field::new("lin_vel_z", DataType::Float64, false),
70                Field::new("ang_vel_x", DataType::Float64, false),
71                Field::new("ang_vel_y", DataType::Float64, false),
72                Field::new("ang_vel_z", DataType::Float64, false),
73                Field::new("timestamp_ns", DataType::Int64, false),
74            ]))
75        })
76        .clone()
77}
78
79/// Encode an `Odometry` sample as a single-row `RecordBatch` matching [`schema`].
80///
81/// # Example
82///
83/// ```
84/// use isaac_sim_arrow::odometry::{Odometry, to_record_batch};
85/// let odom = Odometry {
86///     chassis_frame_id: "base_link",
87///     odom_frame_id: "odom",
88///     position_x: 1.0, position_y: 2.0, position_z: 0.0,
89///     orientation_w: 1.0, orientation_x: 0.0, orientation_y: 0.0, orientation_z: 0.0,
90///     lin_vel_x: 0.4, lin_vel_y: 0.0, lin_vel_z: 0.0,
91///     ang_vel_x: 0.0, ang_vel_y: 0.0, ang_vel_z: 0.3,
92///     timestamp_ns: 7,
93/// };
94/// let batch = to_record_batch(&odom).unwrap();
95/// assert_eq!(batch.num_rows(), 1);
96/// assert_eq!(batch.num_columns(), 16);
97/// ```
98pub fn to_record_batch(odom: &Odometry) -> Result<RecordBatch, arrow::error::ArrowError> {
99    let columns: Vec<ArrayRef> = vec![
100        Arc::new(StringArray::from(vec![odom.chassis_frame_id])),
101        Arc::new(StringArray::from(vec![odom.odom_frame_id])),
102        Arc::new(Float64Array::from_iter_values(std::iter::once(
103            odom.position_x,
104        ))),
105        Arc::new(Float64Array::from_iter_values(std::iter::once(
106            odom.position_y,
107        ))),
108        Arc::new(Float64Array::from_iter_values(std::iter::once(
109            odom.position_z,
110        ))),
111        Arc::new(Float64Array::from_iter_values(std::iter::once(
112            odom.orientation_w,
113        ))),
114        Arc::new(Float64Array::from_iter_values(std::iter::once(
115            odom.orientation_x,
116        ))),
117        Arc::new(Float64Array::from_iter_values(std::iter::once(
118            odom.orientation_y,
119        ))),
120        Arc::new(Float64Array::from_iter_values(std::iter::once(
121            odom.orientation_z,
122        ))),
123        Arc::new(Float64Array::from_iter_values(std::iter::once(
124            odom.lin_vel_x,
125        ))),
126        Arc::new(Float64Array::from_iter_values(std::iter::once(
127            odom.lin_vel_y,
128        ))),
129        Arc::new(Float64Array::from_iter_values(std::iter::once(
130            odom.lin_vel_z,
131        ))),
132        Arc::new(Float64Array::from_iter_values(std::iter::once(
133            odom.ang_vel_x,
134        ))),
135        Arc::new(Float64Array::from_iter_values(std::iter::once(
136            odom.ang_vel_y,
137        ))),
138        Arc::new(Float64Array::from_iter_values(std::iter::once(
139            odom.ang_vel_z,
140        ))),
141        Arc::new(Int64Array::from_iter_values(std::iter::once(
142            odom.timestamp_ns,
143        ))),
144    ];
145    RecordBatch::try_new(schema(), columns)
146}
147
148/// Decode the first row of a `StructArray` into a heap-owned `OdometryOwned`.
149///
150/// # Example
151///
152/// ```
153/// use arrow::array::StructArray;
154/// use isaac_sim_arrow::odometry::{Odometry, to_record_batch, from_struct_array};
155/// let odom = Odometry {
156///     chassis_frame_id: "base_link",
157///     odom_frame_id: "odom",
158///     position_x: 1.0, position_y: 2.0, position_z: 0.0,
159///     orientation_w: 1.0, orientation_x: 0.0, orientation_y: 0.0, orientation_z: 0.0,
160///     lin_vel_x: 0.4, lin_vel_y: 0.0, lin_vel_z: 0.0,
161///     ang_vel_x: 0.0, ang_vel_y: 0.0, ang_vel_z: 0.3,
162///     timestamp_ns: 7,
163/// };
164/// let batch = to_record_batch(&odom).unwrap();
165/// let array = StructArray::from(batch);
166/// let owned = from_struct_array(&array).unwrap();
167/// assert_eq!(owned.chassis_frame_id, "base_link");
168/// assert_eq!(owned.timestamp_ns, 7);
169/// ```
170pub fn from_struct_array(array: &StructArray) -> Result<OdometryOwned, arrow::error::ArrowError> {
171    if array.is_empty() {
172        return Err(arrow::error::ArrowError::InvalidArgumentError(
173            "odometry struct array is empty".into(),
174        ));
175    }
176    let str_at = |idx: usize, name: &str| -> Result<String, arrow::error::ArrowError> {
177        array
178            .column(idx)
179            .as_any()
180            .downcast_ref::<StringArray>()
181            .ok_or_else(|| {
182                arrow::error::ArrowError::SchemaError(format!("odometry '{name}' not Utf8"))
183            })
184            .map(|a| a.value(0).to_string())
185    };
186    let f64_at = |idx: usize, name: &str| -> Result<f64, arrow::error::ArrowError> {
187        array
188            .column(idx)
189            .as_any()
190            .downcast_ref::<Float64Array>()
191            .ok_or_else(|| {
192                arrow::error::ArrowError::SchemaError(format!("odometry '{name}' not Float64"))
193            })
194            .map(|a| a.value(0))
195    };
196    Ok(OdometryOwned {
197        chassis_frame_id: str_at(0, "chassis_frame_id")?,
198        odom_frame_id: str_at(1, "odom_frame_id")?,
199        position_x: f64_at(2, "position_x")?,
200        position_y: f64_at(3, "position_y")?,
201        position_z: f64_at(4, "position_z")?,
202        orientation_w: f64_at(5, "orientation_w")?,
203        orientation_x: f64_at(6, "orientation_x")?,
204        orientation_y: f64_at(7, "orientation_y")?,
205        orientation_z: f64_at(8, "orientation_z")?,
206        lin_vel_x: f64_at(9, "lin_vel_x")?,
207        lin_vel_y: f64_at(10, "lin_vel_y")?,
208        lin_vel_z: f64_at(11, "lin_vel_z")?,
209        ang_vel_x: f64_at(12, "ang_vel_x")?,
210        ang_vel_y: f64_at(13, "ang_vel_y")?,
211        ang_vel_z: f64_at(14, "ang_vel_z")?,
212        timestamp_ns: array
213            .column(15)
214            .as_any()
215            .downcast_ref::<Int64Array>()
216            .ok_or_else(|| {
217                arrow::error::ArrowError::SchemaError("odometry 'timestamp_ns' not Int64".into())
218            })?
219            .value(0),
220    })
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use arrow::array::Array;
227
228    #[test]
229    fn round_trips_through_record_batch() {
230        let odom = Odometry {
231            chassis_frame_id: "base_link",
232            odom_frame_id: "odom",
233            position_x: 1.0,
234            position_y: 2.0,
235            position_z: 0.0,
236            orientation_w: 1.0,
237            orientation_x: 0.0,
238            orientation_y: 0.0,
239            orientation_z: 0.0,
240            lin_vel_x: 0.4,
241            lin_vel_y: 0.0,
242            lin_vel_z: 0.0,
243            ang_vel_x: 0.0,
244            ang_vel_y: 0.0,
245            ang_vel_z: 0.3,
246            timestamp_ns: 7,
247        };
248        let batch = to_record_batch(&odom).expect("convert");
249        assert_eq!(batch.num_rows(), 1);
250        assert_eq!(batch.num_columns(), 16);
251
252        let chassis = batch
253            .column(0)
254            .as_any()
255            .downcast_ref::<StringArray>()
256            .expect("chassis_frame_id is Utf8");
257        assert_eq!(chassis.value(0), "base_link");
258
259        let lin_x = batch
260            .column(9)
261            .as_any()
262            .downcast_ref::<Float64Array>()
263            .expect("lin_vel_x is Float64");
264        assert!((lin_x.value(0) - 0.4).abs() < 1e-9);
265    }
266
267    #[test]
268    fn from_struct_array_round_trips() {
269        let odom = Odometry {
270            chassis_frame_id: "base_link",
271            odom_frame_id: "odom",
272            position_x: 1.0,
273            position_y: 2.0,
274            position_z: 0.0,
275            orientation_w: 1.0,
276            orientation_x: 0.0,
277            orientation_y: 0.0,
278            orientation_z: 0.0,
279            lin_vel_x: 0.4,
280            lin_vel_y: 0.0,
281            lin_vel_z: 0.0,
282            ang_vel_x: 0.0,
283            ang_vel_y: 0.0,
284            ang_vel_z: 0.3,
285            timestamp_ns: 7,
286        };
287        let batch = to_record_batch(&odom).expect("to");
288        let array = StructArray::from(batch);
289        let owned = from_struct_array(&array).expect("from");
290        assert_eq!(owned.chassis_frame_id, "base_link");
291        assert_eq!(owned.odom_frame_id, "odom");
292        assert!((owned.lin_vel_x - 0.4).abs() < 1e-9);
293        assert_eq!(owned.timestamp_ns, 7);
294    }
295}
isaac_sim_arrow/odometry.rs

isaac_sim_arrow/
odometry.rs
