isaac_sim_arrow/lidar/
pointcloud.rs

1// SPDX-License-Identifier: MPL-2.0
2//! Arrow encoder and decoder for the 3D RTX LiDAR PointCloud channel.
3use std::sync::{Arc, OnceLock};
4
5use arrow::array::{Array, ArrayRef, Float32Array, Int32Array, ListArray, StructArray};
6use arrow::buffer::OffsetBuffer;
7use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
8use arrow::record_batch::RecordBatch;
9
10/// Borrowed view of a single 3D LiDAR PointCloud frame, used as input to [`to_record_batch`].
11#[allow(missing_docs)]
12pub struct LidarPointCloud<'a> {
13    pub points: &'a [f32],
14    pub num_points: i32,
15    pub width: i32,
16    pub height: i32,
17}
18
19/// Owned variant returned by [`from_struct_array`]. Holds heap-owned
20/// payload so a downstream node can keep the decoded value across the next event.
21#[derive(Debug, Clone, PartialEq)]
22#[allow(missing_docs)]
23pub struct LidarPointCloudOwned {
24    pub points: Vec<f32>,
25    pub num_points: i32,
26    pub width: i32,
27    pub height: i32,
28}
29
30/// Stable Arrow schema for a `LidarPointCloud` record batch.
31pub fn schema() -> SchemaRef {
32    static SCHEMA: OnceLock<SchemaRef> = OnceLock::new();
33    SCHEMA
34        .get_or_init(|| {
35            Arc::new(Schema::new(vec![
36                Field::new(
37                    "points",
38                    DataType::List(Arc::new(Field::new("item", DataType::Float32, false))),
39                    false,
40                ),
41                Field::new("num_points", DataType::Int32, false),
42                Field::new("width", DataType::Int32, false),
43                Field::new("height", DataType::Int32, false),
44            ]))
45        })
46        .clone()
47}
48
49fn list_f32(values: &[f32]) -> ListArray {
50    let inner = Float32Array::from_iter_values(values.iter().copied());
51    let offsets = OffsetBuffer::from_lengths([values.len()]);
52    ListArray::new(
53        Arc::new(Field::new("item", DataType::Float32, false)),
54        offsets,
55        Arc::new(inner),
56        None,
57    )
58}
59
60/// Encode a `LidarPointCloud` frame as a single-row `RecordBatch` matching [`schema`].
61/// The flat `[x, y, z, ...]` points buffer is stored as an Arrow `List<Float32>` column.
62///
63/// # Example
64///
65/// ```
66/// use isaac_sim_arrow::lidar::pointcloud::{LidarPointCloud, to_record_batch};
67/// let points = [1.0_f32, 0.0, 0.0, 0.0, 1.0, 0.0];
68/// let pc = LidarPointCloud { points: &points, num_points: 2, width: 2, height: 1 };
69/// let batch = to_record_batch(&pc).unwrap();
70/// assert_eq!(batch.num_rows(), 1);
71/// assert_eq!(batch.num_columns(), 4);
72/// ```
73pub fn to_record_batch(pc: &LidarPointCloud) -> Result<RecordBatch, arrow::error::ArrowError> {
74    let columns: Vec<ArrayRef> = vec![
75        Arc::new(list_f32(pc.points)),
76        Arc::new(Int32Array::from_iter_values(std::iter::once(pc.num_points))),
77        Arc::new(Int32Array::from_iter_values(std::iter::once(pc.width))),
78        Arc::new(Int32Array::from_iter_values(std::iter::once(pc.height))),
79    ];
80    RecordBatch::try_new(schema(), columns)
81}
82
83/// Zero-copy variant returned by [`from_struct_array_borrowed`]; the
84/// `points` slice borrows directly from the Arrow buffer.
85#[allow(missing_docs)]
86pub struct LidarPointCloudBorrowed<'a> {
87    pub points: &'a [f32],
88    pub num_points: i32,
89    pub width: i32,
90    pub height: i32,
91}
92
93/// Decode the first row of a `StructArray` into a borrowed view (zero-copy points slice).
94/// The returned struct borrows from `array`; use [`from_struct_array`] for an owned copy.
95pub fn from_struct_array_borrowed(
96    array: &StructArray,
97) -> Result<LidarPointCloudBorrowed<'_>, arrow::error::ArrowError> {
98    if array.is_empty() {
99        return Err(arrow::error::ArrowError::InvalidArgumentError(
100            "lidar_pointcloud struct array is empty".into(),
101        ));
102    }
103    let pts_list = array
104        .column(0)
105        .as_any()
106        .downcast_ref::<ListArray>()
107        .ok_or_else(|| {
108            arrow::error::ArrowError::SchemaError("pointcloud 'points' not ListArray".into())
109        })?;
110    let points = pts_list
111        .values()
112        .as_any()
113        .downcast_ref::<Float32Array>()
114        .ok_or_else(|| {
115            arrow::error::ArrowError::SchemaError("pointcloud 'points' inner not Float32".into())
116        })?
117        .values();
118    let i32_at = |idx: usize, name: &str| -> Result<i32, arrow::error::ArrowError> {
119        array
120            .column(idx)
121            .as_any()
122            .downcast_ref::<Int32Array>()
123            .ok_or_else(|| {
124                arrow::error::ArrowError::SchemaError(format!("pointcloud '{name}' not Int32"))
125            })
126            .map(|a| a.value(0))
127    };
128    Ok(LidarPointCloudBorrowed {
129        points,
130        num_points: i32_at(1, "num_points")?,
131        width: i32_at(2, "width")?,
132        height: i32_at(3, "height")?,
133    })
134}
135
136/// Decode the first row of a `StructArray` into a heap-owned `LidarPointCloudOwned`.
137///
138/// # Example
139///
140/// ```
141/// use arrow::array::StructArray;
142/// use isaac_sim_arrow::lidar::pointcloud::{LidarPointCloud, to_record_batch, from_struct_array};
143/// let points = [1.0_f32, 0.0, 0.0];
144/// let pc = LidarPointCloud { points: &points, num_points: 1, width: 1, height: 1 };
145/// let batch = to_record_batch(&pc).unwrap();
146/// let array = StructArray::from(batch);
147/// let owned = from_struct_array(&array).unwrap();
148/// assert_eq!(owned.num_points, 1);
149/// assert_eq!(&owned.points, &[1.0_f32, 0.0, 0.0]);
150/// ```
151pub fn from_struct_array(
152    array: &StructArray,
153) -> Result<LidarPointCloudOwned, arrow::error::ArrowError> {
154    if array.is_empty() {
155        return Err(arrow::error::ArrowError::InvalidArgumentError(
156            "lidar_pointcloud struct array is empty".into(),
157        ));
158    }
159    let pts_list = array
160        .column(0)
161        .as_any()
162        .downcast_ref::<ListArray>()
163        .ok_or_else(|| {
164            arrow::error::ArrowError::SchemaError("pointcloud 'points' not ListArray".into())
165        })?;
166    let pts_inner = pts_list
167        .values()
168        .as_any()
169        .downcast_ref::<Float32Array>()
170        .ok_or_else(|| {
171            arrow::error::ArrowError::SchemaError("pointcloud 'points' inner not Float32".into())
172        })?;
173    let i32_at = |idx: usize, name: &str| -> Result<i32, arrow::error::ArrowError> {
174        array
175            .column(idx)
176            .as_any()
177            .downcast_ref::<Int32Array>()
178            .ok_or_else(|| {
179                arrow::error::ArrowError::SchemaError(format!("pointcloud '{name}' not Int32"))
180            })
181            .map(|a| a.value(0))
182    };
183    Ok(LidarPointCloudOwned {
184        points: pts_inner.values().to_vec(),
185        num_points: i32_at(1, "num_points")?,
186        width: i32_at(2, "width")?,
187        height: i32_at(3, "height")?,
188    })
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194
195    #[test]
196    fn round_trips_through_record_batch() {
197        let points = [
198            1.0_f32, 0.0, 0.0, //
199            0.0, 1.0, 0.0, //
200            0.0, 0.0, 1.0,
201        ];
202        let pc = LidarPointCloud {
203            points: &points,
204            num_points: 3,
205            width: 3,
206            height: 1,
207        };
208        let batch = to_record_batch(&pc).expect("convert");
209        assert_eq!(batch.num_rows(), 1);
210        assert_eq!(batch.num_columns(), 4);
211
212        let pts_col = batch
213            .column(0)
214            .as_any()
215            .downcast_ref::<ListArray>()
216            .expect("points is ListArray");
217        let inner = pts_col
218            .values()
219            .as_any()
220            .downcast_ref::<Float32Array>()
221            .expect("inner is Float32Array");
222        assert_eq!(inner.len(), 9);
223        assert_eq!(inner.value(0), 1.0);
224        assert_eq!(inner.value(8), 1.0);
225    }
226
227    #[test]
228    fn from_struct_array_round_trips() {
229        let points = [1.0_f32, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
230        let pc = LidarPointCloud {
231            points: &points,
232            num_points: 3,
233            width: 3,
234            height: 1,
235        };
236        let batch = to_record_batch(&pc).expect("to");
237        let array = StructArray::from(batch);
238        let owned = from_struct_array(&array).expect("from");
239        assert_eq!(owned.points, points);
240        assert_eq!(owned.num_points, 3);
241    }
242}
isaac_sim_arrow/lidar/pointcloud.rs

isaac_sim_arrow/lidar/
pointcloud.rs
