You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@sedona.apache.org by "zhangfengcdt (via GitHub)" <gi...@apache.org> on 2025/09/30 15:21:14 UTC
[PR] [WIP] feat: optimize KNN join with lock-free shared geometry cache [sedona-db]
zhangfengcdt opened a new pull request, #169:
URL: https://github.com/apache/sedona-db/pull/169
Optimizes KNN join performance by implementing a lock-free shared geometry cache that eliminates redundant WKB-to-geometry conversions across multiple queries.
Changes:
- Use the new indexed distance metric and adapter pattern
- Replace per-query geometry cache with shared `Vec<OnceCell<Geometry<f64>>>`
- Use atomic operations for thread-safe lazy geometry initialization
- Simplify memory tracking using WKB buffer size as proxy
NOTE:
This is PR is pending for the refactor changes be merged:
https://github.com/wherobots/geo-index/pull/6
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: refactor KNN join with new geo-index trait and lock-free shared geometry cache [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #169:
URL: https://github.com/apache/sedona-db/pull/169#discussion_r2400020061
##########
.github/workflows/rust.yml:
##########
@@ -55,9 +55,9 @@ jobs:
runs-on: ubuntu-latest
env:
CARGO_INCREMENTAL: 0
- # Reduce debug info to save disk space
- CARGO_PROFILE_DEV_DEBUG: 1
- CARGO_PROFILE_TEST_DEBUG: 1
+ # Disable debug info completely to save disk space
+ CARGO_PROFILE_DEV_DEBUG: 0
+ CARGO_PROFILE_TEST_DEBUG: 0
Review Comment:
Can you see if bumping the cache id for the Rust cache fixes this instead of these changes?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: refactor KNN join with new geo-index trait and lock-free shared geometry cache [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #169:
URL: https://github.com/apache/sedona-db/pull/169#discussion_r2409216079
##########
.github/workflows/rust.yml:
##########
@@ -106,33 +109,50 @@ jobs:
- uses: Swatinem/rust-cache@v2
with:
# Update this key to force a new cache
- prefix-key: "rust-${{ matrix.name }}-v2"
+ prefix-key: "rust-${{ matrix.name }}-v3"
- name: Install dependencies
shell: bash
- run: sudo apt-get update && sudo apt-get install -y libgeos-dev
+ run: |
+ sudo apt-get update && sudo apt-get install -y libgeos-dev
- name: Build
if: matrix.name == 'build'
run: |
cargo build --workspace --all-targets --all-features
+ # Clean up build artifacts aggressively
+ rm -rf target/debug/deps
+ rm -rf target/debug/incremental
+ rm -rf target/debug/build
- name: Clippy
if: matrix.name == 'clippy'
run: |
cargo clippy --workspace --all-targets --all-features -- -Dwarnings
+ # Clean up clippy artifacts aggressively
+ rm -rf target/debug/deps
+ rm -rf target/debug/incremental
+ rm -rf target/debug/build
- name: Test
if: matrix.name == 'test'
run: |
cargo test --workspace --all-targets --all-features
- # Clean up intermediate build artifacts to free disk space
+ # Clean up intermediate build artifacts to free disk space aggressively
cargo clean -p sedona-s2geography
+ rm -rf target/debug/deps
+ rm -rf target/debug/incremental
+ rm -rf target/debug/build
+ # Also clean target/.rustc_info.json and other cache files
+ find target/debug -name "*.rlib" -delete
+ find target/debug -name "*.rmeta" -delete
- name: Doctests
if: matrix.name == 'test'
run: |
cargo test --workspace --doc --all-features
+ # Clean up after doctests
+ cargo clean
Review Comment:
Can we bring https://github.com/apache/sedona-db/pull/33/commits/facc7e7076d80f6b03d761d1196ec0795f97b345 back if the available disk space is tight for this workflow?
##########
rust/sedona-spatial-join/src/index.rs:
##########
@@ -937,6 +903,112 @@ async fn collect_build_partition(
/// Rough estimate for in-memory size of the rtree per rect in bytes
const RTREE_MEMORY_ESTIMATE_PER_RECT: usize = 60;
+/// Shared KNN components that can be reused across queries
+struct KnnComponents {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance,
+ /// Pre-allocated vector for geometry cache - lock-free access
+ /// Indexed by rtree data index for O(1) access
+ geometry_cache: Vec<OnceCell<Geometry<f64>>>,
+ /// Memory reservation to track geometry cache memory usage
+ _reservation: MemoryReservation,
+}
+
+impl KnnComponents {
+ fn new(
+ cache_size: usize,
+ indexed_batches: &[IndexedBatch],
+ memory_pool: Arc<dyn MemoryPool>,
+ ) -> datafusion_common::Result<Self> {
+ // Create memory consumer and reservation for geometry cache
+ let consumer = MemoryConsumer::new("SpatialJoinKnnGeometryCache");
+ let mut reservation = consumer.register(&memory_pool);
+
+ // Estimate maximum possible memory usage based on WKB sizes
+ let estimated_memory = Self::estimate_max_memory_usage(indexed_batches);
+ reservation.try_grow(estimated_memory)?;
+
+ // Pre-allocate OnceCell vector
+ let geometry_cache = (0..cache_size).map(|_| OnceCell::new()).collect();
+
+ Ok(Self {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance::default(),
+ geometry_cache,
+ _reservation: reservation,
+ })
+ }
+
+ /// Estimate the maximum memory usage for decoded geometries based on WKB sizes
+ fn estimate_max_memory_usage(indexed_batches: &[IndexedBatch]) -> usize {
+ let mut total_wkb_size = 0;
+
+ for batch in indexed_batches {
+ for wkb in batch.geom_array.wkbs().iter().flatten() {
+ total_wkb_size += wkb.buf().len();
+ }
+ }
+ total_wkb_size
+ }
+}
+
+/// Geometry accessor for SedonaDB KNN queries.
+/// This accessor provides on-demand WKB decoding and geometry caching for efficient
+/// KNN queries with support for both Euclidean and Haversine distance metrics.
+struct SedonaKnnAdapter<'a> {
+ indexed_batches: &'a [IndexedBatch],
+ data_id_to_batch_pos: &'a [(i32, i32)],
+ // Reference to KNN components for cache and memory tracking
+ knn_components: &'a KnnComponents,
+}
+
+impl<'a> SedonaKnnAdapter<'a> {
+ /// Create a new adapter
+ fn new(
+ indexed_batches: &'a [IndexedBatch],
+ data_id_to_batch_pos: &'a [(i32, i32)],
+ knn_components: &'a KnnComponents,
+ ) -> Self {
+ Self {
+ indexed_batches,
+ data_id_to_batch_pos,
+ knn_components,
+ }
+ }
+}
+
+impl<'a> GeometryAccessor for SedonaKnnAdapter<'a> {
+ /// Get geometry for the given item index with lock-free caching
+ fn get_geometry(&self, item_index: usize) -> Option<Geometry<f64>> {
Review Comment:
Can we avoid returning a clone? Given that all geo_type::Geometry values will be cached and owned by self, we can return `Option<&Geometry<f64>>` here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [WIP] feat: optimize KNN join with lock-free shared geometry cache [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #169:
URL: https://github.com/apache/sedona-db/pull/169#discussion_r2392216585
##########
rust/sedona-spatial-join/src/index.rs:
##########
@@ -557,32 +537,15 @@ impl SpatialIndex {
}
};
- // Use pre-computed cached geometries for performance
- let geometries = &self.cached_geometries;
-
- if geometries.is_empty() {
- return Ok(JoinResultMetrics {
- count: 0,
- candidate_count: 0,
- });
- }
-
- // Choose distance metric based on use_spheroid parameter
- let distance_metric: Box<dyn DistanceMetric<f32>> = if use_spheroid {
- // For spheroid (geodesic) distance, we use the Haversine formula as an approximation for now.
- // The distance metric will be used to calculate distances between geometries for ranking purposes.
- Box::new(HaversineDistance::default())
- } else {
- Box::new(EuclideanDistance)
- };
+ // Create an indexed distance metric adapter for KNN queries with cached distance metrics
+ let indexed_metric = self.create_knn_adapter(use_spheroid);
Review Comment:
The geometry cache in the KNN adapter is not reused across multiple `query_knn` calls. This renders the cache mostly useless.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: refactor KNN join with new geo-index trait and lock-free shared geometry cache [sedona-db]
Posted by "zhangfengcdt (via GitHub)" <gi...@apache.org>.
zhangfengcdt commented on code in PR #169:
URL: https://github.com/apache/sedona-db/pull/169#discussion_r2399198849
##########
rust/sedona-spatial-join/src/index.rs:
##########
@@ -937,6 +903,109 @@ async fn collect_build_partition(
/// Rough estimate for in-memory size of the rtree per rect in bytes
const RTREE_MEMORY_ESTIMATE_PER_RECT: usize = 60;
+/// Shared KNN components that can be reused across queries
+struct KnnComponents {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance,
+ /// Pre-allocated vector for geometry cache - lock-free access
+ /// Indexed by rtree data index for O(1) access
+ geometry_cache: Vec<OnceCell<Geometry<f64>>>,
+}
+
+impl KnnComponents {
+ fn new(
+ cache_size: usize,
+ indexed_batches: &[IndexedBatch],
+ memory_pool: Arc<dyn MemoryPool>,
+ ) -> datafusion_common::Result<Self> {
+ // Create memory consumer and reservation for geometry cache
+ let consumer = MemoryConsumer::new("SpatialJoinKnnGeometryCache");
+ let mut memory_reservation = consumer.register(&memory_pool);
+
+ // Estimate maximum possible memory usage based on WKB sizes
+ let estimated_memory = Self::estimate_max_memory_usage(indexed_batches);
+ memory_reservation.try_grow(estimated_memory)?;
+
+ // Pre-allocate OnceCell vector
+ let geometry_cache = (0..cache_size).map(|_| OnceCell::new()).collect();
+
+ Ok(Self {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance::default(),
+ geometry_cache,
+ })
+ }
+
+ /// Estimate the maximum memory usage for decoded geometries based on WKB sizes
+ fn estimate_max_memory_usage(indexed_batches: &[IndexedBatch]) -> usize {
+ let mut total_wkb_size = 0;
+
+ for batch in indexed_batches {
+ for wkb in batch.geom_array.wkbs().iter().flatten() {
+ total_wkb_size += wkb.buf().len();
+ }
+ }
+ total_wkb_size
Review Comment:
Using the wkb size is conservative and we are using the same approach as before.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: refactor KNN join with new geo-index trait and lock-free shared geometry cache [sedona-db]
Posted by "zhangfengcdt (via GitHub)" <gi...@apache.org>.
zhangfengcdt commented on code in PR #169:
URL: https://github.com/apache/sedona-db/pull/169#discussion_r2411439407
##########
rust/sedona-spatial-join/src/index.rs:
##########
@@ -937,6 +903,112 @@ async fn collect_build_partition(
/// Rough estimate for in-memory size of the rtree per rect in bytes
const RTREE_MEMORY_ESTIMATE_PER_RECT: usize = 60;
+/// Shared KNN components that can be reused across queries
+struct KnnComponents {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance,
+ /// Pre-allocated vector for geometry cache - lock-free access
+ /// Indexed by rtree data index for O(1) access
+ geometry_cache: Vec<OnceCell<Geometry<f64>>>,
+ /// Memory reservation to track geometry cache memory usage
+ _reservation: MemoryReservation,
+}
+
+impl KnnComponents {
+ fn new(
+ cache_size: usize,
+ indexed_batches: &[IndexedBatch],
+ memory_pool: Arc<dyn MemoryPool>,
+ ) -> datafusion_common::Result<Self> {
+ // Create memory consumer and reservation for geometry cache
+ let consumer = MemoryConsumer::new("SpatialJoinKnnGeometryCache");
+ let mut reservation = consumer.register(&memory_pool);
+
+ // Estimate maximum possible memory usage based on WKB sizes
+ let estimated_memory = Self::estimate_max_memory_usage(indexed_batches);
+ reservation.try_grow(estimated_memory)?;
+
+ // Pre-allocate OnceCell vector
+ let geometry_cache = (0..cache_size).map(|_| OnceCell::new()).collect();
+
+ Ok(Self {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance::default(),
+ geometry_cache,
+ _reservation: reservation,
+ })
+ }
+
+ /// Estimate the maximum memory usage for decoded geometries based on WKB sizes
+ fn estimate_max_memory_usage(indexed_batches: &[IndexedBatch]) -> usize {
+ let mut total_wkb_size = 0;
+
+ for batch in indexed_batches {
+ for wkb in batch.geom_array.wkbs().iter().flatten() {
+ total_wkb_size += wkb.buf().len();
+ }
+ }
+ total_wkb_size
+ }
+}
+
+/// Geometry accessor for SedonaDB KNN queries.
+/// This accessor provides on-demand WKB decoding and geometry caching for efficient
+/// KNN queries with support for both Euclidean and Haversine distance metrics.
+struct SedonaKnnAdapter<'a> {
+ indexed_batches: &'a [IndexedBatch],
+ data_id_to_batch_pos: &'a [(i32, i32)],
+ // Reference to KNN components for cache and memory tracking
+ knn_components: &'a KnnComponents,
+}
+
+impl<'a> SedonaKnnAdapter<'a> {
+ /// Create a new adapter
+ fn new(
+ indexed_batches: &'a [IndexedBatch],
+ data_id_to_batch_pos: &'a [(i32, i32)],
+ knn_components: &'a KnnComponents,
+ ) -> Self {
+ Self {
+ indexed_batches,
+ data_id_to_batch_pos,
+ knn_components,
+ }
+ }
+}
+
+impl<'a> GeometryAccessor for SedonaKnnAdapter<'a> {
+ /// Get geometry for the given item index with lock-free caching
+ fn get_geometry(&self, item_index: usize) -> Option<Geometry<f64>> {
Review Comment:
The clone should have been removed by using the new geo-index refactoring.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [WIP] feat: optimize KNN join with lock-free shared geometry cache [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #169:
URL: https://github.com/apache/sedona-db/pull/169#discussion_r2392227047
##########
rust/sedona-spatial-join/src/index.rs:
##########
@@ -937,6 +905,150 @@ async fn collect_build_partition(
/// Rough estimate for in-memory size of the rtree per rect in bytes
const RTREE_MEMORY_ESTIMATE_PER_RECT: usize = 60;
+/// Shared KNN components that can be reused across queries
+struct KnnComponents {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance,
+ /// Pre-allocated vector for geometry cache - lock-free access
+ /// Indexed by rtree data index for O(1) access
+ geometry_cache: Vec<OnceCell<Geometry<f64>>>,
+}
+
+impl KnnComponents {
+ fn new(
+ cache_size: usize,
+ indexed_batches: &[IndexedBatch],
+ memory_pool: Arc<dyn MemoryPool>,
+ ) -> datafusion_common::Result<Self> {
+ // Create memory consumer and reservation for geometry cache
+ let consumer = MemoryConsumer::new("SpatialJoinKnnGeometryCache");
+ let mut memory_reservation = consumer.register(&memory_pool);
+
+ // Estimate maximum possible memory usage based on WKB sizes
+ let estimated_memory = Self::estimate_max_memory_usage(indexed_batches);
+ memory_reservation.try_grow(estimated_memory)?;
+
+ // Pre-allocate OnceCell vector
+ let geometry_cache = (0..cache_size).map(|_| OnceCell::new()).collect();
+
+ Ok(Self {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance::default(),
+ geometry_cache,
+ })
+ }
+
+ /// Estimate the maximum memory usage for decoded geometries based on WKB sizes
+ fn estimate_max_memory_usage(indexed_batches: &[IndexedBatch]) -> usize {
+ let mut total_wkb_size = 0;
+
+ for batch in indexed_batches {
+ for wkb in batch.geom_array.wkbs().iter().flatten() {
+ total_wkb_size += wkb.buf().len();
+ }
+ }
+ total_wkb_size
+ }
+}
+
+/// Enhanced indexed distance metric adapter for SedonaDB KNN queries.
+/// This adapter provides on-demand WKB decoding and geometry caching for efficient
+/// distance calculations with support for both Euclidean and Haversine distance metrics.
+struct SedonaKnnAdapter<'a> {
+ indexed_batches: &'a [IndexedBatch],
+ data_id_to_batch_pos: &'a [(i32, i32)],
+ distance_metric: &'a dyn DistanceMetric<f32>,
+ // Reference to KNN components for cache and memory tracking
+ knn_components: &'a KnnComponents,
Review Comment:
This is a shared reference owned by the spatial index, so the geometry cache was correctly reused. My previous review comment was incorrect. This should be OK.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: refactor KNN join with new geo-index trait and lock-free shared geometry cache [sedona-db]
Posted by "zhangfengcdt (via GitHub)" <gi...@apache.org>.
zhangfengcdt commented on code in PR #169:
URL: https://github.com/apache/sedona-db/pull/169#discussion_r2411163917
##########
rust/sedona-spatial-join/src/index.rs:
##########
@@ -937,6 +903,112 @@ async fn collect_build_partition(
/// Rough estimate for in-memory size of the rtree per rect in bytes
const RTREE_MEMORY_ESTIMATE_PER_RECT: usize = 60;
+/// Shared KNN components that can be reused across queries
+struct KnnComponents {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance,
+ /// Pre-allocated vector for geometry cache - lock-free access
+ /// Indexed by rtree data index for O(1) access
+ geometry_cache: Vec<OnceCell<Geometry<f64>>>,
+ /// Memory reservation to track geometry cache memory usage
+ _reservation: MemoryReservation,
+}
+
+impl KnnComponents {
+ fn new(
+ cache_size: usize,
+ indexed_batches: &[IndexedBatch],
+ memory_pool: Arc<dyn MemoryPool>,
+ ) -> datafusion_common::Result<Self> {
+ // Create memory consumer and reservation for geometry cache
+ let consumer = MemoryConsumer::new("SpatialJoinKnnGeometryCache");
+ let mut reservation = consumer.register(&memory_pool);
+
+ // Estimate maximum possible memory usage based on WKB sizes
+ let estimated_memory = Self::estimate_max_memory_usage(indexed_batches);
+ reservation.try_grow(estimated_memory)?;
+
+ // Pre-allocate OnceCell vector
+ let geometry_cache = (0..cache_size).map(|_| OnceCell::new()).collect();
+
+ Ok(Self {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance::default(),
+ geometry_cache,
+ _reservation: reservation,
+ })
+ }
+
+ /// Estimate the maximum memory usage for decoded geometries based on WKB sizes
+ fn estimate_max_memory_usage(indexed_batches: &[IndexedBatch]) -> usize {
+ let mut total_wkb_size = 0;
+
+ for batch in indexed_batches {
+ for wkb in batch.geom_array.wkbs().iter().flatten() {
+ total_wkb_size += wkb.buf().len();
+ }
+ }
+ total_wkb_size
+ }
+}
+
+/// Geometry accessor for SedonaDB KNN queries.
+/// This accessor provides on-demand WKB decoding and geometry caching for efficient
+/// KNN queries with support for both Euclidean and Haversine distance metrics.
+struct SedonaKnnAdapter<'a> {
+ indexed_batches: &'a [IndexedBatch],
+ data_id_to_batch_pos: &'a [(i32, i32)],
+ // Reference to KNN components for cache and memory tracking
+ knn_components: &'a KnnComponents,
+}
+
+impl<'a> SedonaKnnAdapter<'a> {
+ /// Create a new adapter
+ fn new(
+ indexed_batches: &'a [IndexedBatch],
+ data_id_to_batch_pos: &'a [(i32, i32)],
+ knn_components: &'a KnnComponents,
+ ) -> Self {
+ Self {
+ indexed_batches,
+ data_id_to_batch_pos,
+ knn_components,
+ }
+ }
+}
+
+impl<'a> GeometryAccessor for SedonaKnnAdapter<'a> {
+ /// Get geometry for the given item index with lock-free caching
+ fn get_geometry(&self, item_index: usize) -> Option<Geometry<f64>> {
Review Comment:
@Kontinuation Here is the change to geo-index: https://github.com/wherobots/geo-index/pull/8
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: refactor KNN join with new geo-index trait and lock-free shared geometry cache [sedona-db]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #169:
URL: https://github.com/apache/sedona-db/pull/169#discussion_r2397360544
##########
rust/sedona-spatial-join/src/index.rs:
##########
@@ -937,6 +903,109 @@ async fn collect_build_partition(
/// Rough estimate for in-memory size of the rtree per rect in bytes
const RTREE_MEMORY_ESTIMATE_PER_RECT: usize = 60;
+/// Shared KNN components that can be reused across queries
+struct KnnComponents {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance,
+ /// Pre-allocated vector for geometry cache - lock-free access
+ /// Indexed by rtree data index for O(1) access
+ geometry_cache: Vec<OnceCell<Geometry<f64>>>,
+}
+
+impl KnnComponents {
+ fn new(
+ cache_size: usize,
+ indexed_batches: &[IndexedBatch],
+ memory_pool: Arc<dyn MemoryPool>,
+ ) -> datafusion_common::Result<Self> {
+ // Create memory consumer and reservation for geometry cache
+ let consumer = MemoryConsumer::new("SpatialJoinKnnGeometryCache");
+ let mut memory_reservation = consumer.register(&memory_pool);
+
+ // Estimate maximum possible memory usage based on WKB sizes
+ let estimated_memory = Self::estimate_max_memory_usage(indexed_batches);
+ memory_reservation.try_grow(estimated_memory)?;
Review Comment:
Memory reservation is grown but never used or tracked. The `memory_reservation` variable is created and grown but then dropped without being stored in the struct, leading to immediate memory deallocation.
##########
rust/sedona-spatial-join/src/index.rs:
##########
@@ -937,6 +903,109 @@ async fn collect_build_partition(
/// Rough estimate for in-memory size of the rtree per rect in bytes
const RTREE_MEMORY_ESTIMATE_PER_RECT: usize = 60;
+/// Shared KNN components that can be reused across queries
+struct KnnComponents {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance,
+ /// Pre-allocated vector for geometry cache - lock-free access
+ /// Indexed by rtree data index for O(1) access
+ geometry_cache: Vec<OnceCell<Geometry<f64>>>,
+}
+
+impl KnnComponents {
+ fn new(
+ cache_size: usize,
+ indexed_batches: &[IndexedBatch],
+ memory_pool: Arc<dyn MemoryPool>,
+ ) -> datafusion_common::Result<Self> {
+ // Create memory consumer and reservation for geometry cache
+ let consumer = MemoryConsumer::new("SpatialJoinKnnGeometryCache");
+ let mut memory_reservation = consumer.register(&memory_pool);
+
+ // Estimate maximum possible memory usage based on WKB sizes
+ let estimated_memory = Self::estimate_max_memory_usage(indexed_batches);
+ memory_reservation.try_grow(estimated_memory)?;
+
+ // Pre-allocate OnceCell vector
+ let geometry_cache = (0..cache_size).map(|_| OnceCell::new()).collect();
+
+ Ok(Self {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance::default(),
+ geometry_cache,
+ })
+ }
+
+ /// Estimate the maximum memory usage for decoded geometries based on WKB sizes
+ fn estimate_max_memory_usage(indexed_batches: &[IndexedBatch]) -> usize {
+ let mut total_wkb_size = 0;
+
+ for batch in indexed_batches {
+ for wkb in batch.geom_array.wkbs().iter().flatten() {
+ total_wkb_size += wkb.buf().len();
+ }
+ }
+ total_wkb_size
Review Comment:
The memory estimation only considers WKB size but doesn't account for the overhead of `OnceCell` containers and decoded `Geometry` objects, which could lead to underestimation of actual memory usage.
##########
.github/workflows/rust.yml:
##########
@@ -55,11 +55,14 @@ jobs:
runs-on: ubuntu-latest
env:
CARGO_INCREMENTAL: 0
- # Reduce debug info to save disk space
- CARGO_PROFILE_DEV_DEBUG: 1
- CARGO_PROFILE_TEST_DEBUG: 1
+ # Disable debug info completely to save disk space
+ CARGO_PROFILE_DEV_DEBUG: 0
+ CARGO_PROFILE_TEST_DEBUG: 0
# Limit parallel compilation to reduce memory pressure
CARGO_BUILD_JOBS: 2
+ # Strip symbols to reduce binary size
+ CARGO_PROFILE_DEV_STRIP: symbols
+ CARGO_PROFILE_TEST_STRIP: symbols
Review Comment:
These environment variables are not valid Cargo configuration. The correct way to strip symbols is through `Cargo.toml` profiles or using `--config` flag, not environment variables.
```suggestion
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [WIP] feat: optimize KNN join with lock-free shared geometry cache [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #169:
URL: https://github.com/apache/sedona-db/pull/169#discussion_r2392216585
##########
rust/sedona-spatial-join/src/index.rs:
##########
@@ -557,32 +537,15 @@ impl SpatialIndex {
}
};
- // Use pre-computed cached geometries for performance
- let geometries = &self.cached_geometries;
-
- if geometries.is_empty() {
- return Ok(JoinResultMetrics {
- count: 0,
- candidate_count: 0,
- });
- }
-
- // Choose distance metric based on use_spheroid parameter
- let distance_metric: Box<dyn DistanceMetric<f32>> = if use_spheroid {
- // For spheroid (geodesic) distance, we use the Haversine formula as an approximation for now.
- // The distance metric will be used to calculate distances between geometries for ranking purposes.
- Box::new(HaversineDistance::default())
- } else {
- Box::new(EuclideanDistance)
- };
+ // Create an indexed distance metric adapter for KNN queries with cached distance metrics
+ let indexed_metric = self.create_knn_adapter(use_spheroid);
Review Comment:
The geometry cache in the KNN adapter is not reused across multiple `query_knn` calls. This renders the cache mostly useless.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: refactor KNN join with new geo-index trait and lock-free shared geometry cache [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu merged PR #169:
URL: https://github.com/apache/sedona-db/pull/169
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: refactor KNN join with new geo-index trait and lock-free shared geometry cache [sedona-db]
Posted by "zhangfengcdt (via GitHub)" <gi...@apache.org>.
zhangfengcdt commented on code in PR #169:
URL: https://github.com/apache/sedona-db/pull/169#discussion_r2410706134
##########
rust/sedona-spatial-join/src/index.rs:
##########
@@ -937,6 +903,112 @@ async fn collect_build_partition(
/// Rough estimate for in-memory size of the rtree per rect in bytes
const RTREE_MEMORY_ESTIMATE_PER_RECT: usize = 60;
+/// Shared KNN components that can be reused across queries
+struct KnnComponents {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance,
+ /// Pre-allocated vector for geometry cache - lock-free access
+ /// Indexed by rtree data index for O(1) access
+ geometry_cache: Vec<OnceCell<Geometry<f64>>>,
+ /// Memory reservation to track geometry cache memory usage
+ _reservation: MemoryReservation,
+}
+
+impl KnnComponents {
+ fn new(
+ cache_size: usize,
+ indexed_batches: &[IndexedBatch],
+ memory_pool: Arc<dyn MemoryPool>,
+ ) -> datafusion_common::Result<Self> {
+ // Create memory consumer and reservation for geometry cache
+ let consumer = MemoryConsumer::new("SpatialJoinKnnGeometryCache");
+ let mut reservation = consumer.register(&memory_pool);
+
+ // Estimate maximum possible memory usage based on WKB sizes
+ let estimated_memory = Self::estimate_max_memory_usage(indexed_batches);
+ reservation.try_grow(estimated_memory)?;
+
+ // Pre-allocate OnceCell vector
+ let geometry_cache = (0..cache_size).map(|_| OnceCell::new()).collect();
+
+ Ok(Self {
+ euclidean_metric: EuclideanDistance,
+ haversine_metric: HaversineDistance::default(),
+ geometry_cache,
+ _reservation: reservation,
+ })
+ }
+
+ /// Estimate the maximum memory usage for decoded geometries based on WKB sizes
+ fn estimate_max_memory_usage(indexed_batches: &[IndexedBatch]) -> usize {
+ let mut total_wkb_size = 0;
+
+ for batch in indexed_batches {
+ for wkb in batch.geom_array.wkbs().iter().flatten() {
+ total_wkb_size += wkb.buf().len();
+ }
+ }
+ total_wkb_size
+ }
+}
+
+/// Geometry accessor for SedonaDB KNN queries.
+/// This accessor provides on-demand WKB decoding and geometry caching for efficient
+/// KNN queries with support for both Euclidean and Haversine distance metrics.
+struct SedonaKnnAdapter<'a> {
+ indexed_batches: &'a [IndexedBatch],
+ data_id_to_batch_pos: &'a [(i32, i32)],
+ // Reference to KNN components for cache and memory tracking
+ knn_components: &'a KnnComponents,
+}
+
+impl<'a> SedonaKnnAdapter<'a> {
+ /// Create a new adapter
+ fn new(
+ indexed_batches: &'a [IndexedBatch],
+ data_id_to_batch_pos: &'a [(i32, i32)],
+ knn_components: &'a KnnComponents,
+ ) -> Self {
+ Self {
+ indexed_batches,
+ data_id_to_batch_pos,
+ knn_components,
+ }
+ }
+}
+
+impl<'a> GeometryAccessor for SedonaKnnAdapter<'a> {
+ /// Get geometry for the given item index with lock-free caching
+ fn get_geometry(&self, item_index: usize) -> Option<Geometry<f64>> {
Review Comment:
Yes, I will need to change the trait definition in geo-index first before changing this to return reference.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org