Increment nextDocId even if geo indexing fails (apache#11158)

deemoliu · Jul 25, 2023 · 7231ecf · 7231ecf
1 parent a753025
commit 7231ecf
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/...src/main/java/org/apache/pinot/segment/local/realtime/impl/geospatial/MutableH3Index.java b/...src/main/java/org/apache/pinot/segment/local/realtime/impl/geospatial/MutableH3Index.java
@@ -55,8 +55,12 @@ public MutableH3Index(H3IndexResolution resolution)
 
   @Override
   public void add(@Nonnull Object value, int dictId, int docId) {
-    Geometry geometry = GeometrySerializer.deserialize((byte[]) value);
-    add(geometry);
+    try {
+      Geometry geometry = GeometrySerializer.deserialize((byte[]) value);
+      add(geometry);
+    } finally {
+      _nextDocId++;
+    }
   }
 
   @Override
@@ -73,7 +77,7 @@ public void add(Geometry geometry) {
     Coordinate coordinate = geometry.getCoordinate();
     // TODO: support multiple resolutions
     long h3Id = H3Utils.H3_CORE.geoToH3(coordinate.y, coordinate.x, _lowestResolution);
-    _bitmaps.computeIfAbsent(h3Id, k -> new ThreadSafeMutableRoaringBitmap()).add(_nextDocId++);
+    _bitmaps.computeIfAbsent(h3Id, k -> new ThreadSafeMutableRoaringBitmap()).add(_nextDocId);
   }
 
   @Override

diff --git a/...segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/H3IndexTest.java b/...segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/H3IndexTest.java
@@ -33,6 +33,7 @@
 import org.apache.pinot.segment.local.segment.creator.impl.inv.geospatial.OnHeapH3IndexCreator;
 import org.apache.pinot.segment.local.segment.index.h3.H3IndexType;
 import org.apache.pinot.segment.local.segment.index.readers.geospatial.ImmutableH3IndexReader;
+import org.apache.pinot.segment.local.utils.GeometrySerializer;
 import org.apache.pinot.segment.local.utils.GeometryUtils;
 import org.apache.pinot.segment.local.utils.H3Utils;
 import org.apache.pinot.segment.spi.V1Constants;
@@ -89,13 +90,14 @@ public void testH3Index()
           h3IndexResolution);
           GeoSpatialIndexCreator offHeapCreator = new OffHeapH3IndexCreator(TEMP_DIR, offHeapColumnName,
               h3IndexResolution)) {
+        int docId = 0;
         while (expectedCardinalities.size() < numUniqueH3Ids) {
           double longitude = RANDOM.nextDouble() * 360 - 180;
           double latitude = RANDOM.nextDouble() * 180 - 90;
           Point point = GeometryUtils.GEOMETRY_FACTORY.createPoint(new Coordinate(longitude, latitude));
           onHeapCreator.add(point);
           offHeapCreator.add(point);
-          mutableH3Index.add(point);
+          mutableH3Index.add(GeometrySerializer.serialize(point), -1, docId++);
           long h3Id = H3Utils.H3_CORE.geoToH3(latitude, longitude, resolution);
           expectedCardinalities.merge(h3Id, 1, Integer::sum);
         }