Optional parallel statistics extraction

Author: b4l

Cargo.lock

@@ -36,6 +36,7 @@ arrow-ipc = { workspace = true }
 arrow-schema = { workspace = true }
 async-stream = "0.3.6"
 async-trait = { workspace = true }
+byteorder = { workspace = true }
 bytes = { workspace = true }
 datafusion-catalog = { workspace = true }
 datafusion-common = { workspace = true }
@@ -51,6 +52,7 @@ las = { version = "0.9.10", features = ["laz"] }
 las-crs = { version = "1.0.0" }
 laz = "0.12.0"
 object_store = { workspace = true }
+rayon = "1.11.0"
 sedona-expr = { workspace = true }
 sedona-geometry = { workspace = true }
 

rust/sedona-pointcloud/Cargo.toml

@@ -153,6 +153,7 @@ impl<'a> LasMetadataReader<'a> {
                     &chunk_table,
                     &header,
                     options.persist_statistics,
+                    options.parallel_statistics_extraction,
                 )
                 .await?,
             )
@@ -366,9 +367,9 @@ pub async fn fetch_chunk_table(
     header: &Header,
 ) -> Result<Vec<ChunkMeta>, Box<dyn Error + Send + Sync>> {
     if header.laz_vlr().is_ok() {
-        laz_chunk_table(store, object_meta, &header).await
+        laz_chunk_table(store, object_meta, header).await
     } else {
-        las_chunk_table(&header).await
+        las_chunk_table(header).await
     }
 }
 

rust/sedona-pointcloud/src/las/metadata.rs

@@ -187,7 +187,7 @@ pub fn record_decompressor(
     Ok(decompressor)
 }
 
-pub(crate) fn read_point<R: Read>(buffer: R, header: &Header) -> Result<Point, DataFusionError> {
+fn read_point<R: Read>(buffer: R, header: &Header) -> Result<Point, DataFusionError> {
     RawPoint::read_from(buffer, header.point_format())
         .map(|raw_point| Point::new(raw_point, header.transforms()))
         .map_err(|e| DataFusionError::External(Box::new(e)))

rust/sedona-pointcloud/src/las/reader.rs

@@ -15,7 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::{collections::HashSet, io::Cursor, sync::Arc};
+use std::{
+    collections::HashSet,
+    io::{Cursor, Read},
+    sync::Arc,
+};
 
 use arrow_array::{
     builder::PrimitiveBuilder,
@@ -25,17 +29,15 @@ use arrow_array::{
 };
 use arrow_ipc::{reader::FileReader, writer::FileWriter};
 use arrow_schema::{DataType, Field, Schema};
+use byteorder::{LittleEndian, ReadBytesExt};
 use datafusion_common::{arrow::compute::concat_batches, Column, DataFusionError, ScalarValue};
 use datafusion_pruning::PruningStatistics;
-use las::{Header, Point};
+use las::Header;
 use object_store::{path::Path, ObjectMeta, ObjectStore, PutPayload};
-
+use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use sedona_geometry::bounding_box::BoundingBox;
 
-use crate::las::{
-    metadata::ChunkMeta,
-    reader::{read_point, record_decompressor},
-};
+use crate::las::{metadata::ChunkMeta, reader::record_decompressor};
 
 /// Spatial statistics (extent) of LAS/LAZ chunks for pruning.
 ///
@@ -208,6 +210,7 @@ pub async fn chunk_statistics(
     chunk_table: &[ChunkMeta],
     header: &Header,
     persist: bool,
+    parallel: bool,
 ) -> Result<LasStatistics, DataFusionError> {
     let stats_path = Path::parse(format!("{}.stats", object_meta.location.as_ref()))?;
 
@@ -234,9 +237,27 @@ pub async fn chunk_statistics(
             // extract statistics
             let mut builder = LasStatisticsBuilder::new_with_capacity(chunk_table.len());
 
-            for chunk_meta in chunk_table {
-                let stats = extract_chunk_stats(store, object_meta, chunk_meta, header).await?;
-                builder.add_values(&stats, chunk_meta.num_points);
+            if parallel {
+                let stats: Vec<[f64; 6]> = chunk_table
+                    .par_iter()
+                    .map(|chunk_meta| {
+                        futures::executor::block_on(extract_chunk_stats(
+                            store,
+                            object_meta,
+                            chunk_meta,
+                            header,
+                        ))
+                    })
+                    .collect::<Result<Vec<[f64; 6]>, DataFusionError>>()?;
+
+                for (stat, meta) in stats.iter().zip(chunk_table) {
+                    builder.add_values(stat, meta.num_points);
+                }
+            } else {
+                for chunk_meta in chunk_table {
+                    let stats = extract_chunk_stats(store, object_meta, chunk_meta, header).await?;
+                    builder.add_values(&stats, chunk_meta.num_points);
+                }
             }
 
             let stats = builder.finish();
@@ -274,14 +295,14 @@ async fn extract_chunk_stats(
         f64::NEG_INFINITY,
     ];
 
-    let extend = |stats: &mut [f64; 6], point: Point| {
+    let extend = |stats: &mut [f64; 6], point: [f64; 3]| {
         *stats = [
-            stats[0].min(point.x),
-            stats[1].max(point.x),
-            stats[2].min(point.y),
-            stats[3].max(point.y),
-            stats[4].min(point.z),
-            stats[5].max(point.z),
+            stats[0].min(point[0]),
+            stats[1].max(point[0]),
+            stats[2].min(point[1]),
+            stats[3].max(point[1]),
+            stats[4].min(point[2]),
+            stats[5].max(point[2]),
         ];
     };
 
@@ -301,21 +322,29 @@ async fn extract_chunk_stats(
         for _ in 0..chunk_meta.num_points {
             buffer.set_position(0);
             decompressor.decompress_next(buffer.get_mut())?;
-            let point = read_point(&mut buffer, header)?;
+            let point = parse_coords(&mut buffer, header)?;
             extend(&mut stats, point);
         }
     } else {
         let mut buffer = Cursor::new(bytes);
 
         for _ in 0..chunk_meta.num_points {
-            let point = read_point(&mut buffer, header)?;
+            let point = parse_coords(&mut buffer, header)?;
             extend(&mut stats, point);
         }
     }
 
     Ok(stats)
 }
 
+fn parse_coords<R: Read>(mut buffer: R, header: &Header) -> Result<[f64; 3], DataFusionError> {
+    let transforms = header.transforms();
+    let x = transforms.x.direct(buffer.read_i32::<LittleEndian>()?);
+    let y = transforms.y.direct(buffer.read_i32::<LittleEndian>()?);
+    let z = transforms.z.direct(buffer.read_i32::<LittleEndian>()?);
+    Ok([x, y, z])
+}
+
 #[cfg(test)]
 mod tests {
     use std::fs::File;

rust/sedona-pointcloud/src/las/statistics.rs

@@ -80,6 +80,7 @@ extensions_options! {
     pub struct PointcloudOptions {
         pub geometry_encoding: GeometryEncoding, default = GeometryEncoding::default()
         pub collect_statistics: bool, default = false
+        pub parallel_statistics_extraction: bool, default = false
         pub persist_statistics: bool, default = false
         pub round_robin_partitioning: bool, default = false
         pub las: LasOptions, default = LasOptions::default()