From b713a517f1210efcc271a8a0cccdabc751d954b0 Mon Sep 17 00:00:00 2001
From: Pranav <pranavbhole@gmail.com>
Date: Mon, 6 May 2024 02:07:29 -0700
Subject: [PATCH] Fix the bug in Immutable RTree object strategy (#16389)

* Fix the bug in Immutable Node object strategy

* Adding comments in code
---
 .../data/ImmutableRTreeObjectStrategy.java    |  6 ++--
 .../spatial/ImmutableRTreeTest.java           | 34 +++++++++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/processing/src/main/java/org/apache/druid/segment/data/ImmutableRTreeObjectStrategy.java b/processing/src/main/java/org/apache/druid/segment/data/ImmutableRTreeObjectStrategy.java
index c97b69f25c0c..8015f79202af 100644
--- a/processing/src/main/java/org/apache/druid/segment/data/ImmutableRTreeObjectStrategy.java
+++ b/processing/src/main/java/org/apache/druid/segment/data/ImmutableRTreeObjectStrategy.java
@@ -63,8 +63,10 @@ public Class<ImmutableRTree> getClazz()
   @Override
   public ImmutableRTree fromByteBuffer(ByteBuffer buffer, int numBytes)
   {
-    buffer.limit(buffer.position() + numBytes);
-    return new ImmutableRTree(buffer, bitmapFactory);
+    // always create the duplicate buffer for creating the objects as original buffer may have mutations somewhere else which can corrupt objects
+    ByteBuffer duplicateBuf = buffer.duplicate();
+    duplicateBuf.limit(duplicateBuf.position() + numBytes);
+    return new ImmutableRTree(duplicateBuf, bitmapFactory);
   }
 
   @Override
diff --git a/processing/src/test/java/org/apache/druid/collections/spatial/ImmutableRTreeTest.java b/processing/src/test/java/org/apache/druid/collections/spatial/ImmutableRTreeTest.java
index 94edd78275f9..6a1228f7a0e9 100644
--- a/processing/src/test/java/org/apache/druid/collections/spatial/ImmutableRTreeTest.java
+++ b/processing/src/test/java/org/apache/druid/collections/spatial/ImmutableRTreeTest.java
@@ -48,6 +48,7 @@
 import java.util.concurrent.TimeUnit;
 
 /**
+ *
  */
 @SuppressWarnings("CheckReturnValue")
 public class ImmutableRTreeTest
@@ -702,4 +703,37 @@ public void testPreciseRadiusBoundFilter()
     ImmutableBitmap finalSet = bf.union(points);
     Assert.assertTrue(finalSet.size() == 100);
   }
+
+  @Test
+  public void testForImmutableRTreeImmutability()
+  {
+    BitmapFactory bf = new RoaringBitmapFactory();
+    RTree tree = new RTree(2, new LinearGutmanSplitStrategy(0, 50, bf), bf);
+    float centerLat = 37.4133961f;
+    float centerLong = -122.1224665f;
+    float[][] insidePoints = SpatialUtils.generateGeoCoordinatesAroundCircle(centerLat, centerLong, 100, 100, true);
+    for (int i = 0; i < insidePoints.length; i++) {
+      tree.insert(insidePoints[i], i);
+    }
+    float[][] outsidePoints = SpatialUtils.generateGeoCoordinatesAroundCircle(centerLat, centerLong, 100, 100, false);
+    for (int i = 0; i < outsidePoints.length; i++) {
+      tree.insert(outsidePoints[i], i);
+    }
+    final ImmutableRTree searchTree = ImmutableRTree.newImmutableFromMutable(tree);
+    ByteBuffer buffer = ByteBuffer.wrap(searchTree.toBytes());
+    ImmutableRTreeObjectStrategy objectStrategy = new ImmutableRTreeObjectStrategy(bf);
+    ImmutableRTree builtTree = objectStrategy.fromByteBuffer(buffer, searchTree.toBytes().length);
+    buffer.position(10);
+    buffer.limit(100);
+
+    for (float[] insidePoint : insidePoints) {
+      Iterable<ImmutableBitmap> points = builtTree.search(new RadiusBound(
+          new float[]{insidePoint[0], insidePoint[1]},
+          100,
+          2,
+          RadiusBound.RadiusUnit.meters
+      ));
+      bf.union(points);
+    }
+  }
 }