epoint_io/epoint/
read.rs

1use crate::Error::{FileNotFound, InvalidFileExtension, NoFileName};
2use crate::epoint::documents::EpointInfoDocument;
3use crate::epoint::read_impl::cast_data_frame;
4use crate::epoint::{
5    EPOINT_SEPARATOR, FILE_EXTENSION_EPOINT_FORMAT, FILE_EXTENSION_EPOINT_TAR_FORMAT,
6    FILE_NAME_ECOORD_COMPRESSED, FILE_NAME_ECOORD_UNCOMPRESSED, FILE_NAME_INFO_COMPRESSED,
7    FILE_NAME_INFO_UNCOMPRESSED, FILE_NAME_POINT_DATA_COMPRESSED,
8    FILE_NAME_POINT_DATA_UNCOMPRESSED,
9};
10use crate::error::Error;
11use ecoord::TransformTree;
12use epoint_core::PointCloud;
13use epoint_core::PointCloudInfo;
14use polars::prelude::*;
15use std::fs::File;
16use std::io::{Cursor, Read};
17use std::path::Path;
18use tar::Archive;
19
20/// `EpointReader` sets up a reader for the custom reader data structure.
21///
22#[derive(Debug, Clone)]
23pub struct EpointReader<R: Read> {
24    reader: R,
25}
26
27impl<R: Read> EpointReader<R> {
28    pub fn new(reader: R) -> Self {
29        Self { reader }
30    }
31
32    pub fn finish(self) -> Result<PointCloud, Error> {
33        let mut archive = Archive::new(self.reader);
34
35        let mut info_document: Option<EpointInfoDocument> = None;
36        let mut point_data_frame: Option<DataFrame> = None;
37        let mut transform_tree: Option<TransformTree> = None;
38
39        for file in archive.entries()? {
40            let mut f = file?;
41
42            match f.path()?.to_str().unwrap() {
43                FILE_NAME_INFO_UNCOMPRESSED => {
44                    info_document = serde_json::from_reader(f)?;
45                }
46                FILE_NAME_INFO_COMPRESSED => {
47                    let mut decompressed_buffer: Vec<u8> = Vec::new();
48                    zstd::stream::copy_decode(f, &mut decompressed_buffer)?;
49                    info_document = serde_json::from_reader(Cursor::new(decompressed_buffer))?;
50                }
51                FILE_NAME_POINT_DATA_UNCOMPRESSED => {
52                    let mut buffer: Vec<u8> = Vec::new();
53                    f.read_to_end(&mut buffer)?;
54                    let reader = Cursor::new(&buffer);
55
56                    let csv_parse_options =
57                        CsvParseOptions::default().with_separator(EPOINT_SEPARATOR);
58                    let data_frame: DataFrame = CsvReadOptions::default()
59                        .with_parse_options(csv_parse_options)
60                        .into_reader_with_file_handle(reader)
61                        .finish()?;
62                    let casted_data_frame = cast_data_frame(data_frame)?;
63
64                    point_data_frame = Some(casted_data_frame);
65                }
66                FILE_NAME_POINT_DATA_COMPRESSED => {
67                    let mut buffer: Vec<u8> = Vec::new();
68                    f.read_to_end(&mut buffer)?;
69                    let reader = Cursor::new(&buffer);
70
71                    let data_frame: DataFrame = ParquetReader::new(reader).finish()?;
72                    let casted_data_frame = cast_data_frame(data_frame)?;
73
74                    point_data_frame = Some(casted_data_frame);
75                }
76                FILE_NAME_ECOORD_UNCOMPRESSED => {
77                    transform_tree = Some(ecoord::io::EcoordReader::new(f).finish()?);
78                }
79                FILE_NAME_ECOORD_COMPRESSED => {
80                    transform_tree = Some(
81                        ecoord::io::EcoordReader::new(f)
82                            .with_compression(ecoord::io::Compression::default_zstd())
83                            .finish()?,
84                    );
85                }
86                _ => {}
87            }
88        }
89
90        let info: PointCloudInfo = info_document
91            .ok_or(FileNotFound("info".to_string()))?
92            .into();
93        let point_data_frame = point_data_frame.ok_or(FileNotFound("point_data".to_string()))?;
94        let transform_tree = transform_tree.ok_or(FileNotFound("ecoord".to_string()))?;
95
96        let point_cloud = PointCloud::from_data_frame(point_data_frame, info, transform_tree)?;
97        Ok(point_cloud)
98    }
99}
100
101impl EpointReader<File> {
102    pub fn from_path(path: impl AsRef<Path>) -> Result<Self, Error> {
103        let file_name_str = path
104            .as_ref()
105            .file_name()
106            .ok_or(NoFileName())?
107            .to_string_lossy()
108            .to_lowercase();
109        if !file_name_str.ends_with(FILE_EXTENSION_EPOINT_TAR_FORMAT)
110            && !file_name_str.ends_with(FILE_EXTENSION_EPOINT_FORMAT)
111        {
112            return Err(InvalidFileExtension(file_name_str.to_string()));
113        }
114
115        let file = std::fs::File::open(path)?;
116        Ok(Self::new(file))
117    }
118}