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#[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}