Add an efficient GeometryTypeAndDimensionsSet to improve the performance of geo statistics collection and ST_Analyze_Aggr (#285)

The idea is pretty simple: we use a 32-bit bitset to represent geometry type and dimensions. This replaces usages of HashSet<GeometryTypeAndDimensions> in `GeoStatistics` and reduces the overhead of updating geo statistics in `AnalyzeAccumulator`. We care about the performance geo statistics analyzer because it is applied to all geometries on the indexed side when running a spatial join. `AnalyzeAccumulator` and ST_Analyze_Aggr can be useful in some other places as well, so we'd like them to have minimal performance overhead.

Here is the benchmark result of st_analyze_aggr after applying this patch:

```
Gnuplot not found, using plotters backend
native-st_analyze_aggr-Array(Point)
                        time:   [4.1267 ms 4.2026 ms 4.3423 ms]
                        change: [-87.458% -87.216% -86.808%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 8 outliers among 100 measurements (8.00%)
  6 (6.00%) high mild
  2 (2.00%) high severe

native-st_analyze_aggr-Array(LineString(10))
                        time:   [5.6607 ms 5.6728 ms 5.6868 ms]
                        change: [-83.578% -83.529% -83.482%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 1 outliers among 100 measurements (1.00%)
  1 (1.00%) high severe
```

Co-authored-by: Copilot <[email protected]>

Author: Kontinuation

rust/sedona-expr/src/statistics.rs

@@ -14,11 +14,14 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
-use std::{collections::HashSet, str::FromStr};
+use std::str::FromStr;
 
 use datafusion_common::{stats::Precision, ColumnStatistics, DataFusionError, Result, ScalarValue};
 use sedona_geometry::interval::{Interval, IntervalTrait};
-use sedona_geometry::{bounding_box::BoundingBox, types::GeometryTypeAndDimensions};
+use sedona_geometry::{
+    bounding_box::BoundingBox,
+    types::{GeometryTypeAndDimensions, GeometryTypeAndDimensionsSet},
+};
 use serde::{Deserialize, Serialize};
 
 /// Statistics specific to spatial data types
@@ -33,7 +36,7 @@ use serde::{Deserialize, Serialize};
 pub struct GeoStatistics {
     // Core spatial statistics for pruning
     bbox: Option<BoundingBox>, // The overall bounding box (min/max coordinates) containing all geometries
-    geometry_types: Option<HashSet<GeometryTypeAndDimensions>>, // Set of all geometry types and dimensions present
+    geometry_types: Option<GeometryTypeAndDimensionsSet>, // Set of all geometry types and dimensions present
 
     // Extended statistics for analysis
     total_geometries: Option<i64>, // Total count of all geometries
@@ -73,16 +76,16 @@ impl GeoStatistics {
     pub fn empty() -> Self {
         Self {
             bbox: Some(BoundingBox::xy(Interval::empty(), Interval::empty())),
-            geometry_types: Some(HashSet::new()), // Empty set of geometry types
-            total_geometries: Some(0),            // Zero geometries
-            total_size_bytes: Some(0),            // Zero bytes
-            total_points: Some(0),                // Zero points
-            puntal_count: Some(0),                // Zero point geometries
-            lineal_count: Some(0),                // Zero line geometries
-            polygonal_count: Some(0),             // Zero polygon geometries
-            collection_count: Some(0),            // Zero collection geometries
-            total_envelope_width: Some(0.0),      // Zero width
-            total_envelope_height: Some(0.0),     // Zero height
+            geometry_types: Some(GeometryTypeAndDimensionsSet::new()), // Empty set of geometry types
+            total_geometries: Some(0),                                 // Zero geometries
+            total_size_bytes: Some(0),                                 // Zero bytes
+            total_points: Some(0),                                     // Zero points
+            puntal_count: Some(0),                                     // Zero point geometries
+            lineal_count: Some(0),                                     // Zero line geometries
+            polygonal_count: Some(0),                                  // Zero polygon geometries
+            collection_count: Some(0),                                 // Zero collection geometries
+            total_envelope_width: Some(0.0),                           // Zero width
+            total_envelope_height: Some(0.0),                          // Zero height
         }
     }
 
@@ -92,20 +95,10 @@ impl GeoStatistics {
     }
 
     /// Update the geometry types and return self
-    pub fn with_geometry_types(self, types: Option<&[GeometryTypeAndDimensions]>) -> Self {
-        match types {
-            Some(type_slice) => {
-                let type_set: HashSet<GeometryTypeAndDimensions> =
-                    type_slice.iter().cloned().collect();
-                Self {
-                    geometry_types: Some(type_set),
-                    ..self
-                }
-            }
-            None => Self {
-                geometry_types: None,
-                ..self
-            },
+    pub fn with_geometry_types(self, types: Option<GeometryTypeAndDimensionsSet>) -> Self {
+        Self {
+            geometry_types: types,
+            ..self
         }
     }
 
@@ -115,7 +108,7 @@ impl GeoStatistics {
     }
 
     /// Get the geometry types if available
-    pub fn geometry_types(&self) -> Option<&HashSet<GeometryTypeAndDimensions>> {
+    pub fn geometry_types(&self) -> Option<&GeometryTypeAndDimensionsSet> {
         self.geometry_types.as_ref()
     }
 
@@ -290,9 +283,7 @@ impl GeoStatistics {
         if let Some(other_types) = &other.geometry_types {
             match &mut self.geometry_types {
                 Some(types) => {
-                    let mut new_types = types.clone();
-                    new_types.extend(other_types.iter().cloned());
-                    self.geometry_types = Some(new_types);
+                    types.merge(other_types);
                 }
                 None => self.geometry_types = Some(other_types.clone()),
             }
@@ -374,13 +365,12 @@ impl GeoStatistics {
     pub fn try_with_str_geometry_types(self, geometry_types: Option<&[&str]>) -> Result<Self> {
         match geometry_types {
             Some(strings) => {
-                let new_geometry_types = strings
-                    .iter()
-                    .map(|string| {
-                        GeometryTypeAndDimensions::from_str(string)
-                            .map_err(|e| DataFusionError::External(Box::new(e)))
-                    })
-                    .collect::<Result<HashSet<GeometryTypeAndDimensions>>>()?;
+                let mut new_geometry_types = GeometryTypeAndDimensionsSet::new();
+                for string in strings {
+                    let type_and_dim = GeometryTypeAndDimensions::from_str(string)
+                        .map_err(|e| DataFusionError::External(Box::new(e)))?;
+                    new_geometry_types.insert_or_ignore(&type_and_dim);
+                }
 
                 Ok(Self {
                     geometry_types: Some(new_geometry_types),
@@ -442,7 +432,10 @@ mod test {
         // Test with_bbox
         let stats = GeoStatistics::empty().with_bbox(Some(bbox.clone()));
         assert_eq!(stats.bbox(), Some(&bbox));
-        assert_eq!(stats.geometry_types(), Some(HashSet::new()).as_ref());
+        assert_eq!(
+            stats.geometry_types(),
+            Some(&GeometryTypeAndDimensionsSet::new())
+        );
 
         let regular_stats = stats.to_column_statistics().unwrap();
         assert_eq!(
@@ -459,15 +452,17 @@ mod test {
 
     #[test]
     fn specified_geometry_types() {
-        let type_array = [GeometryTypeAndDimensions::new(
-            GeometryTypeId::Polygon,
-            Dimensions::Xy,
-        )];
+        let mut types = GeometryTypeAndDimensionsSet::new();
+        types
+            .insert(&GeometryTypeAndDimensions::new(
+                GeometryTypeId::Polygon,
+                Dimensions::Xy,
+            ))
+            .unwrap();
 
         // Test with_geometry_types
-        let stats = GeoStatistics::empty().with_geometry_types(Some(&type_array));
-        let expected_set: HashSet<GeometryTypeAndDimensions> = type_array.iter().cloned().collect();
-        assert_eq!(stats.geometry_types(), Some(&expected_set));
+        let stats = GeoStatistics::empty().with_geometry_types(Some(types.clone()));
+        assert_eq!(stats.geometry_types(), Some(&types));
         assert_eq!(
             stats.bbox(),
             Some(&BoundingBox::xy(Interval::empty(), Interval::empty()))
@@ -493,15 +488,19 @@ mod test {
             .try_with_str_geometry_types(Some(&["polygon", "point"]))
             .unwrap();
 
-        let mut expected_types = HashSet::new();
-        expected_types.insert(GeometryTypeAndDimensions::new(
-            GeometryTypeId::Polygon,
-            Dimensions::Xy,
-        ));
-        expected_types.insert(GeometryTypeAndDimensions::new(
-            GeometryTypeId::Point,
-            Dimensions::Xy,
-        ));
+        let mut expected_types = GeometryTypeAndDimensionsSet::new();
+        expected_types
+            .insert(&GeometryTypeAndDimensions::new(
+                GeometryTypeId::Polygon,
+                Dimensions::Xy,
+            ))
+            .unwrap();
+        expected_types
+            .insert(&GeometryTypeAndDimensions::new(
+                GeometryTypeId::Point,
+                Dimensions::Xy,
+            ))
+            .unwrap();
 
         assert_eq!(stats.geometry_types(), Some(&expected_types));
         assert_eq!(

rust/sedona-functions/src/st_analyze_aggr.rs

@@ -35,7 +35,7 @@ use sedona_expr::aggregate_udf::SedonaAggregateUDF;
 use sedona_expr::{aggregate_udf::SedonaAccumulator, statistics::GeoStatistics};
 use sedona_geometry::analyze::GeometryAnalysis;
 use sedona_geometry::interval::IntervalTrait;
-use sedona_geometry::types::GeometryTypeAndDimensions;
+use sedona_geometry::types::{GeometryTypeAndDimensions, GeometryTypeAndDimensionsSet};
 use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher};
 use wkb::reader::Wkb;
 
@@ -353,18 +353,15 @@ impl AnalyzeAccumulator {
         let current_types = stats.geometry_types();
         let types = if let Some(existing_types) = current_types {
             let mut new_types = existing_types.clone();
-            new_types.insert(geometry_type);
+            new_types.insert_or_ignore(&geometry_type);
             Some(new_types)
         } else {
-            Some(std::collections::HashSet::from([geometry_type]))
+            let mut new_set = GeometryTypeAndDimensionsSet::new();
+            new_set.insert_or_ignore(&geometry_type);
+            Some(new_set)
         };
 
-        if let Some(type_set) = &types {
-            let type_vec: Vec<GeometryTypeAndDimensions> = type_set.iter().cloned().collect();
-            stats.with_geometry_types(Some(&type_vec))
-        } else {
-            stats.with_geometry_types(None)
-        }
+        stats.with_geometry_types(types)
     }
 
     fn execute_update(&mut self, executor: WkbExecutor) -> Result<()> {
@@ -414,9 +411,10 @@ impl Accumulator for AnalyzeAccumulator {
         // Add approximate size for geometry types if present
         let types_size = match self.stats.geometry_types() {
             Some(types) => {
+                // GeometryTypeAndDimensionsSet is a u32 bitset
                 let elem_size = size_of::<GeometryTypeAndDimensions>();
-                let capacity = types.capacity();
-                capacity * elem_size
+                let count = types.size();
+                count * elem_size
             }
             None => 0,
         };