epoint_transform/
merge.rs

1use ecoord::ReferenceFrames;
2use epoint_core::{PointCloud, PointCloudInfo, PointDataColumnType};
3use polars::datatypes::{DataType, PlHashSet};
4
5use crate::Error::{ContainsNoPoints, DifferentPointCloudInfos};
6use polars::prelude::{IntoLazy, LazyFrame, concat};
7
8use crate::error::Error;
9
10pub fn merge(point_clouds: Vec<PointCloud>) -> Result<PointCloud, Error> {
11    if point_clouds.is_empty() {
12        return Err(ContainsNoPoints);
13    }
14    let info_set: PlHashSet<&PointCloudInfo> = point_clouds.iter().map(|x| x.info()).collect();
15    if info_set.len() > 1 {
16        return Err(DifferentPointCloudInfos);
17    }
18    let point_cloud_info = point_clouds.first().expect("must contain").info().clone();
19
20    let reference_frames: Vec<ReferenceFrames> = point_clouds
21        .iter()
22        .map(|p| p.reference_frames().clone())
23        .collect();
24    let merged_reference_frames = ecoord::merge(&reference_frames)?;
25
26    let data_frame: Vec<LazyFrame> = point_clouds
27        .iter()
28        .map(|p| {
29            let mut df = p.point_data.data_frame.clone();
30            // back casting to Utf8, as something with merging the string cache doesn't work
31            if p.point_data.contains_frame_id_column() {
32                let casted = df
33                    .column(PointDataColumnType::FrameId.as_str())
34                    .unwrap()
35                    .cast(&DataType::String)
36                    .unwrap()
37                    .take_materialized_series();
38                df.replace(PointDataColumnType::FrameId.as_str(), casted)
39                    .unwrap();
40            }
41
42            df.lazy()
43        })
44        .collect();
45    let mut merged_data_frame = concat(data_frame, Default::default())
46        .unwrap()
47        .collect()
48        .unwrap();
49
50    let frame_id_column = merged_data_frame.column(PointDataColumnType::FrameId.as_str());
51    if let Ok(frame_id_column) = frame_id_column {
52        let casted = frame_id_column
53            .to_owned()
54            .cast(&DataType::Categorical(None, Default::default()))
55            .unwrap()
56            .take_materialized_series();
57        merged_data_frame
58            .replace(PointDataColumnType::FrameId.as_str(), casted)
59            .unwrap();
60    }
61
62    let merged_point_cloud =
63        PointCloud::from_data_frame(merged_data_frame, point_cloud_info, merged_reference_frames)?;
64    Ok(merged_point_cloud)
65}