From 2439027ce47e8421a44cc17d2828148641bebb11 Mon Sep 17 00:00:00 2001 From: tmontaigu Date: Wed, 20 Mar 2024 00:33:38 +0100 Subject: [PATCH] Add API that allows reading many points This adds to the API of the reader some functions that allows to read many points in one call. There are few variants, depending on how the user wants to be conservative with allocations. Also ideally, there would be a `read_into(&mut self, buf: &mut [Point]) -> Result` In order to avoid more allocations (especially the potential extra bytes), but for that there needs to exist a `Point::read_into(raw: &RawPoint, dest_point: &mut Point)` or similar. Which could be added later. The main point for these methods, is that it will later allow to make use of laz-rs parrallel decompressor to greatly speed-up reading LAZ files. If needed, we could add later an iterator that returns vecs of points e.g: ``` fn chunks(points_per_chunk: u64) -> impl Iterator>> ``` --- src/compression/mod.rs | 24 ++++++++++- src/reader.rs | 42 +++++++++++++++++- tests/autzen.rs | 97 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 160 insertions(+), 3 deletions(-) diff --git a/src/compression/mod.rs b/src/compression/mod.rs index 9cecfbd0..eb4afc6e 100644 --- a/src/compression/mod.rs +++ b/src/compression/mod.rs @@ -74,9 +74,8 @@ impl<'a, R: Read + Seek + Send> PointReader for CompressedPointReader<'a, R> { .decompress_one(&mut self.decompressor_output.get_mut()); if let Err(e) = res { Some(Err(e.into())) - } else if let Err(e) = self.decompressor_output.seek(SeekFrom::Start(0)) { - Some(Err(e.into())) } else { + self.decompressor_output.set_position(0); Some(read_point_from(&mut self.decompressor_output, &self.header)) } } else { @@ -84,6 +83,27 @@ impl<'a, R: Read + Seek + Send> PointReader for CompressedPointReader<'a, R> { } } + fn read_into_vec(&mut self, points: &mut Vec, n: u64) -> Result { + let points_left = self.header().number_of_points() - self.last_point_idx; + let num_points_to_read = points_left.min(n); + + self.decompressor_output.get_mut().resize( + num_points_to_read as usize * self.header.point_format().len() as usize, + 0u8, + ); + self.decompressor + .decompress_many(self.decompressor_output.get_mut())?; + self.decompressor_output.set_position(0); + points.reserve(num_points_to_read as usize); + + for _ in 0..num_points_to_read { + let point = read_point_from(&mut self.decompressor_output, &self.header)?; + self.last_point_idx += 1; + points.push(point); + } + Ok(num_points_to_read) + } + fn seek(&mut self, position: u64) -> Result<()> { self.last_point_idx = position; self.decompressor.seek(position)?; diff --git a/src/reader.rs b/src/reader.rs index 691d1cfe..fd6b5efd 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -74,9 +74,15 @@ pub(crate) fn read_point_from( point } -/// Trait to specify behaviour a a PointReader +/// Trait to specify behaviour of a PointReader pub(crate) trait PointReader: Debug + Send { fn read_next(&mut self) -> Option>; + fn read_into_vec(&mut self, points: &mut Vec, n: u64) -> Result; + fn read_next_points(&mut self, n: u64) -> Result> { + let mut points = Vec::with_capacity(n as usize); + self.read_into_vec(&mut points, n)?; + Ok(points) + } fn seek(&mut self, position: u64) -> Result<()>; fn header(&self) -> &Header; } @@ -116,6 +122,18 @@ impl PointReader for UncompressedPointRe } } + fn read_into_vec(&mut self, points: &mut Vec, n: u64) -> Result { + let points_left = self.header().number_of_points() - self.last_point_idx; + let num_points_to_read = points_left.min(n); + points.reserve(num_points_to_read as usize); + for _ in 0..num_points_to_read { + let point = read_point_from(&mut self.source, &self.header)?; + self.last_point_idx += 1; + points.push(point); + } + Ok(num_points_to_read) + } + fn seek(&mut self, position: u64) -> Result<()> { self.last_point_idx = position; self.source.seek(SeekFrom::Start( @@ -153,6 +171,15 @@ pub trait Read { /// ``` fn read(&mut self) -> Option>; + /// Reads n points. + fn read_n(&mut self, n: u64) -> Result>; + + /// Reads n points into the vec + fn read_n_into(&mut self, n: u64, points: &mut Vec) -> Result; + + /// Reads all points left into the vec + fn read_all_points(&mut self, points: &mut Vec) -> Result; + /// Seeks to the given point number, zero-indexed. /// /// Note that seeking on compressed (LAZ) data can be expensive as the reader @@ -301,6 +328,19 @@ impl<'a> Read for Reader<'a> { self.point_reader.read_next() } + fn read_n(&mut self, n: u64) -> Result> { + self.point_reader.read_next_points(n) + } + + fn read_n_into(&mut self, n: u64, points: &mut Vec) -> Result { + self.point_reader.read_into_vec(points, n) + } + + fn read_all_points(&mut self, points: &mut Vec) -> Result { + let point_count = self.point_reader.header().number_of_points(); + self.point_reader.read_into_vec(points, point_count) + } + /// Seeks to the given point number, zero-indexed. fn seek(&mut self, position: u64) -> Result<()> { self.point_reader.seek(position) diff --git a/tests/autzen.rs b/tests/autzen.rs index 96a0ec63..a4a0c1f4 100644 --- a/tests/autzen.rs +++ b/tests/autzen.rs @@ -89,3 +89,100 @@ fn test_seek_0_works_on_las() { fn test_seek_0_works_on_laz() { test_seek_0_works_on("tests/data/autzen.laz"); } + +fn test_read_n_on(path: &str) { + use las::{Point, Read, Reader}; + + let ground_truth_points = { + let mut ground_truth_reader = Reader::from_path(path).unwrap(); + ground_truth_reader + .points() + .collect::>>() + .unwrap() + }; + + let mut reader = Reader::from_path(path).unwrap(); + let n = 7; + let mut all_points = Vec::with_capacity(reader.header().number_of_points() as usize); + loop { + let mut points = reader.read_n(n).unwrap(); + if points.len() == 0 { + break; + } + all_points.append(&mut points); + } + + assert_eq!(all_points, ground_truth_points); +} + +fn test_read_n_into_on(path: &str) { + use las::{Point, Read, Reader}; + + let ground_truth_points = { + let mut ground_truth_reader = Reader::from_path(path).unwrap(); + ground_truth_reader + .points() + .collect::>>() + .unwrap() + }; + + let mut reader = Reader::from_path(path).unwrap(); + let n = 7; + let mut all_points = Vec::with_capacity(reader.header().number_of_points() as usize); + let mut points_buffer = Vec::with_capacity(n as usize); + while reader.read_n_into(n, &mut points_buffer).unwrap() != 0 { + all_points.append(&mut points_buffer); + } + + assert_eq!(all_points, ground_truth_points); +} + +#[test] +fn test_las_read_n() { + test_read_n_on("tests/data/autzen.las"); +} + +#[cfg(feature = "laz")] +#[test] +fn test_laz_read_n() { + test_read_n_on("tests/data/autzen.laz"); +} + +#[test] +fn test_las_read_n_into() { + test_read_n_into_on("tests/data/autzen.las"); +} + +#[cfg(feature = "laz")] +#[test] +fn test_laz_read_n_into() { + test_read_n_into_on("tests/data/autzen.laz"); +} + +fn test_read_all_points_on(path: &str) { + use las::{Point, Read, Reader}; + + let ground_truth_points = { + let mut ground_truth_reader = Reader::from_path(path).unwrap(); + ground_truth_reader + .points() + .collect::>>() + .unwrap() + }; + + let mut reader = Reader::from_path(path).unwrap(); + let mut all_points = vec![]; + reader.read_all_points(&mut all_points).unwrap(); + assert_eq!(all_points, ground_truth_points); +} + +#[test] +fn test_las_read_all_points() { + test_read_all_points_on("tests/data/autzen.las"); +} + +#[cfg(feature = "laz")] +#[test] +fn test_laz_read_all_points() { + test_read_all_points_on("tests/data/autzen.laz"); +}