feat(rust/sedona-pointcloud): add optional round robin partitioning and parallel statistics extraction

rust/sedona-pointcloud/Cargo.toml

@@ -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/src/las/builder.rs

@@ -35,9 +35,10 @@ use geoarrow_array::{
 use geoarrow_schema::Dimension;
 use las::{Header, Point};
 
-use crate::{
-    las::{metadata::ExtraAttribute, options::LasExtraBytes, schema::try_schema_from_header},
-    options::GeometryEncoding,
+use crate::las::{
+    metadata::ExtraAttribute,
+    options::{GeometryEncoding, LasExtraBytes},
+    schema::try_schema_from_header,
 };
 
 #[derive(Debug)]
@@ -515,9 +516,9 @@ mod tests {
     use las::{point::Format, Builder, Writer};
     use object_store::{local::LocalFileSystem, path::Path, ObjectStore};
 
-    use crate::{
-        las::{options::LasExtraBytes, reader::LasFileReaderFactory},
-        options::PointcloudOptions,
+    use crate::las::{
+        options::{LasExtraBytes, LasOptions},
+        reader::LasFileReaderFactory,
     };
 
     #[tokio::test]
@@ -544,7 +545,7 @@ mod tests {
             let file_reader = LasFileReaderFactory::new(Arc::new(store), None)
                 .create_reader(
                     PartitionedFile::new(location, object.size),
-                    PointcloudOptions::default(),
+                    LasOptions::default(),
                 )
                 .unwrap();
             let metadata = file_reader.get_metadata().await.unwrap();
@@ -578,7 +579,7 @@ mod tests {
         let file_reader = LasFileReaderFactory::new(Arc::new(store), None)
             .create_reader(
                 PartitionedFile::new(location, object.size),
-                PointcloudOptions::default().with_las_extra_bytes(LasExtraBytes::Typed),
+                LasOptions::default().with_las_extra_bytes(LasExtraBytes::Typed),
             )
             .unwrap();
         let metadata = file_reader.get_metadata().await.unwrap();

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

@@ -33,9 +33,9 @@ use datafusion_physical_plan::ExecutionPlan;
 use futures::{StreamExt, TryStreamExt};
 use object_store::{ObjectMeta, ObjectStore};
 
-use crate::{
-    las::{metadata::LasMetadataReader, reader::LasFileReaderFactory, source::LasSource},
-    options::PointcloudOptions,
+use crate::las::{
+    metadata::LasMetadataReader, options::LasOptions, reader::LasFileReaderFactory,
+    source::LasSource,
 };
 
 #[derive(Debug, Clone, Copy)]
@@ -56,7 +56,7 @@ impl Extension {
 /// Factory struct used to create [LasFormat]
 pub struct LasFormatFactory {
     // inner options for LAS/LAZ
-    pub options: Option<PointcloudOptions>,
+    pub options: Option<LasOptions>,
     extension: Extension,
 }
 
@@ -70,7 +70,7 @@ impl LasFormatFactory {
     }
 
     /// Creates an instance of [LasFormatFactory] with customized default options
-    pub fn new_with(options: PointcloudOptions, extension: Extension) -> Self {
+    pub fn new_with(options: LasOptions, extension: Extension) -> Self {
         Self {
             options: Some(options),
             extension,
@@ -87,8 +87,8 @@ impl FileFormatFactory for LasFormatFactory {
         let mut options = state
             .config_options()
             .extensions
-            .get::<PointcloudOptions>()
-            .or_else(|| state.table_options().extensions.get::<PointcloudOptions>())
+            .get::<LasOptions>()
+            .or_else(|| state.table_options().extensions.get::<LasOptions>())
             .cloned()
             .or(self.options.clone())
             .unwrap_or_default();
@@ -129,7 +129,7 @@ impl fmt::Debug for LasFormatFactory {
 /// The LAS/LAZ `FileFormat` implementation
 #[derive(Debug)]
 pub struct LasFormat {
-    pub options: PointcloudOptions,
+    pub options: LasOptions,
     extension: Extension,
 }
 
@@ -141,7 +141,7 @@ impl LasFormat {
         }
     }
 
-    pub fn with_options(mut self, options: PointcloudOptions) -> Self {
+    pub fn with_options(mut self, options: LasOptions) -> Self {
         self.options = options;
         self
     }
@@ -195,7 +195,7 @@ impl FileFormat for LasFormat {
                 Ok::<_, DataFusionError>((loc_path, schema))
             })
             .boxed() // Workaround https://github.com/rust-lang/rust/issues/64552
-            // fetch schemas concurrently, if requested
+            // fetch schemas concurrently, if requested (note that this is not parallel)
             .buffered(state.config_options().execution.meta_fetch_concurrency)
             .try_collect()
             .await?;

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

@@ -36,12 +36,10 @@ use las::{
 use laz::laszip::ChunkTable;
 use object_store::{ObjectMeta, ObjectStore};
 
-use crate::{
-    las::{
-        schema::try_schema_from_header,
-        statistics::{chunk_statistics, LasStatistics},
-    },
-    options::PointcloudOptions,
+use crate::las::{
+    options::LasOptions,
+    schema::try_schema_from_header,
+    statistics::{chunk_statistics, LasStatistics},
 };
 
 /// LAS/LAZ chunk metadata
@@ -92,7 +90,7 @@ pub struct LasMetadataReader<'a> {
     store: &'a dyn ObjectStore,
     object_meta: &'a ObjectMeta,
     file_metadata_cache: Option<Arc<dyn FileMetadataCache>>,
-    options: PointcloudOptions,
+    options: LasOptions,
 }
 
 impl<'a> LasMetadataReader<'a> {
@@ -115,18 +113,11 @@ impl<'a> LasMetadataReader<'a> {
     }
 
     /// set table options
-    pub fn with_options(mut self, options: PointcloudOptions) -> Self {
+    pub fn with_options(mut self, options: LasOptions) -> Self {
         self.options = options;
         self
     }
 
-    /// Fetch header
-    pub async fn fetch_header(&self) -> Result<Header, DataFusionError> {
-        fetch_header(self.store, self.object_meta)
-            .await
-            .map_err(DataFusionError::External)
-    }
-
     /// Fetch LAS/LAZ metadata from the remote object store
     pub async fn fetch_metadata(&self) -> Result<Arc<LasMetadata>, DataFusionError> {
         let Self {
@@ -149,13 +140,9 @@ impl<'a> LasMetadataReader<'a> {
             return Ok(las_file_metadata);
         }
 
-        let header = self.fetch_header().await?;
+        let header = fetch_header(*store, object_meta).await?;
         let extra_attributes = extra_bytes_attributes(&header)?;
-        let chunk_table = if header.laz_vlr().is_ok() {
-            laz_chunk_table(*store, object_meta, &header).await?
-        } else {
-            las_chunk_table(&header).await?
-        };
+        let chunk_table = fetch_chunk_table(*store, object_meta, &header).await?;
         let statistics = if options.collect_statistics {
             Some(
                 chunk_statistics(
@@ -164,6 +151,7 @@ impl<'a> LasMetadataReader<'a> {
                     &chunk_table,
                     &header,
                     options.persist_statistics,
+                    options.parallel_statistics_extraction,
                 )
                 .await?,
             )
@@ -192,7 +180,7 @@ impl<'a> LasMetadataReader<'a> {
         let schema = try_schema_from_header(
             &metadata.header,
             self.options.geometry_encoding,
-            self.options.las.extra_bytes,
+            self.options.extra_bytes,
         )?;
 
         Ok(schema)
@@ -241,7 +229,8 @@ impl<'a> LasMetadataReader<'a> {
     }
 }
 
-async fn fetch_header(
+/// Fetch the [Header] of a LAS/LAZ file
+pub async fn fetch_header(
     store: &(impl ObjectStore + ?Sized),
     object_meta: &ObjectMeta,
 ) -> Result<Header, Box<dyn Error + Send + Sync>> {
@@ -296,6 +285,7 @@ async fn fetch_header(
     Ok(builder.into_header()?)
 }
 
+/// Extra attribute information (custom attributes in LAS/LAZ files)
 #[derive(Debug, Clone, PartialEq)]
 pub struct ExtraAttribute {
     pub data_type: DataType,
@@ -368,6 +358,19 @@ fn extra_bytes_attributes(
     Ok(attributes)
 }
 
+/// Fetch or generate chunk table metadata.
+pub async fn fetch_chunk_table(
+    store: &(impl ObjectStore + ?Sized),
+    object_meta: &ObjectMeta,
+    header: &Header,
+) -> Result<Vec<ChunkMeta>, Box<dyn Error + Send + Sync>> {
+    if header.laz_vlr().is_ok() {
+        laz_chunk_table(store, object_meta, header).await
+    } else {
+        las_chunk_table(header).await
+    }
+}
+
 async fn laz_chunk_table(
     store: &(impl ObjectStore + ?Sized),
     object_meta: &ObjectMeta,
@@ -482,7 +485,7 @@ mod tests {
     use las::{point::Format, Builder, Reader, Writer};
     use object_store::{local::LocalFileSystem, path::Path, ObjectStore};
 
-    use crate::las::metadata::LasMetadataReader;
+    use crate::las::metadata::fetch_header;
 
     #[tokio::test]
     async fn header_basic_e2e() {
@@ -503,14 +506,13 @@ mod tests {
         let store = LocalFileSystem::new();
         let location = Path::from_filesystem_path(&tmp_path).unwrap();
         let object_meta = store.head(&location).await.unwrap();
-        let metadata_reader = LasMetadataReader::new(&store, &object_meta);
 
         // read with las `Reader`
         let reader = Reader::from_path(&tmp_path).unwrap();
 
         assert_eq!(
             reader.header(),
-            &metadata_reader.fetch_header().await.unwrap()
+            &fetch_header(&store, &object_meta).await.unwrap()
         );
     }
 }