SOLR-18335 : SPLITSHARD fails to migrate documents when using a numeric PointField as router.field

changelog/unreleased/SOLR-18335-splitshard-router-field-point-field.yml

@@ -0,0 +1,8 @@
+title: >
+  SPLITSHARD fails to migrate documents when using a numeric PointField as router.field
+type: fixed
+authors:
+  - name: Olivier Boudet
+links:
+  - name: SOLR-18335
+    url: https://issues.apache.org/jira/browse/SOLR-18335

solr/core/src/java/org/apache/solr/update/SolrIndexSplitter.java

@@ -28,14 +28,18 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.lucene.document.Document;
 import org.apache.lucene.index.CodecReader;
+import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.FilterCodecReader;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.IndexableField;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NoMergePolicy;
+import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.PostingsEnum;
 import org.apache.lucene.index.SlowCodecReaderWrapper;
 import org.apache.lucene.index.Terms;
@@ -70,6 +74,7 @@
 import org.apache.solr.handler.IndexFetcher;
 import org.apache.solr.handler.SnapShooter;
 import org.apache.solr.schema.IndexSchema;
+import org.apache.solr.schema.NumberType;
 import org.apache.solr.schema.SchemaField;
 import org.apache.solr.search.BitsFilteredPostingsEnum;
 import org.apache.solr.search.SolrIndexSearcher;
@@ -677,6 +682,20 @@ static FixedBitSet[] split(
       }
     }
 
+    if (field.getType().isPointField()) {
+      return splitPointField(
+          reader,
+          numPieces,
+          field,
+          rangesArr,
+          splitKey,
+          hashRouter,
+          delete,
+          docSets,
+          liveDocs,
+          currentPartition);
+    }
+
     Terms terms = reader.terms(field.getName());
     TermsEnum termsEnum = terms == null ? null : terms.iterator();
     if (termsEnum == null) return docSets;
@@ -746,38 +765,145 @@ static FixedBitSet[] split(
     }
 
     if (docsMatchingRanges != null) {
-      for (int ii = 0; ii < docsMatchingRanges.length; ii++) {
-        if (0 == docsMatchingRanges[ii]) continue;
-        switch (ii) {
-          case 0:
-            // document loss
-            log.error(
-                "Splitting {}: {} documents belong to no shards and will be dropped",
-                reader,
-                docsMatchingRanges[ii]);
-            break;
-          case 1:
-            // normal case, each document moves to one of the sub-shards
-            log.info(
-                "Splitting {}: {} documents will move into a sub-shard",
-                reader,
-                docsMatchingRanges[ii]);
-            break;
-          default:
-            // document duplication
-            log.error(
-                "Splitting {}: {} documents will be moved to multiple ({}) sub-shards",
-                reader,
-                docsMatchingRanges[ii],
-                ii);
-            break;
+      logDocsMatchingRanges(reader, docsMatchingRanges);
+    }
+
+    return docSets;
+  }
+
+  private static FixedBitSet[] splitPointField(
+      LeafReader reader,
+      int numPieces,
+      SchemaField field,
+      DocRouter.Range[] rangesArr,
+      String splitKey,
+      HashBasedRouter hashRouter,
+      boolean delete,
+      FixedBitSet[] docSets,
+      Bits liveDocs,
+      AtomicInteger currentPartition)
+      throws IOException {
+    NumericDocValues numericDocValues =
+        field.hasDocValues() ? DocValues.getNumeric(reader, field.getName()) : null;
+
+    int[] docsMatchingRanges = null;
+    if (rangesArr != null) {
+      docsMatchingRanges = new int[rangesArr.length + 1];
+    }
+
+    for (int doc = 0; doc < reader.maxDoc(); doc++) {
+      if (liveDocs != null && !liveDocs.get(doc)) {
+        continue;
+      }
+
+      String routeValue = getRouteFieldValue(reader, doc, field, numericDocValues);
+      if (splitKey != null) {
+        String part1 = ((CompositeIdRouter) hashRouter).getRouteKeyNoSuffix(routeValue);
+        if (part1 == null || !splitKey.equals(part1)) {
+          continue;
         }
       }
+
+      if (rangesArr == null) {
+        if (delete) {
+          docSets[currentPartition.get()].clear(doc);
+        } else {
+          docSets[currentPartition.get()].set(doc);
+        }
+        currentPartition.set((currentPartition.get() + 1) % numPieces);
+      } else {
+        int hash = hashRouter.sliceHash(routeValue, null, null, null);
+        int matchingRangesCount = 0;
+        for (int i = 0; i < rangesArr.length; i++) {
+          if (rangesArr[i].includes(hash)) {
+            if (delete) {
+              docSets[i].clear(doc);
+            } else {
+              docSets[i].set(doc);
+            }
+            ++matchingRangesCount;
+          }
+        }
+        docsMatchingRanges[matchingRangesCount]++;
+      }
     }
 
+    if (docsMatchingRanges != null) {
+      logDocsMatchingRanges(reader, docsMatchingRanges);
+    }
     return docSets;
   }
 
+  private static String getRouteFieldValue(
+      LeafReader reader, int doc, SchemaField field, NumericDocValues numericDocValues)
+      throws IOException {
+    if (numericDocValues != null && numericDocValues.advanceExact(doc)) {
+      return numericRouteValueToString(field, numericDocValues.longValue());
+    }
+
+    if (field.stored()) {
+      Document storedDocument = reader.storedFields().document(doc);
+      IndexableField storedField = storedDocument.getField(field.getName());
+      if (storedField != null) {
+        Object routeValue = field.getType().toObject(storedField);
+        if (routeValue != null) {
+          return routeValue.toString();
+        }
+      }
+    }
+
+    throw new SolrException(
+        SolrException.ErrorCode.SERVER_ERROR,
+        "Unable to read route field '"
+            + field.getName()
+            + "' for shard splitting. Point-based route fields must expose docValues or be stored.");
+  }
+
+  private static String numericRouteValueToString(SchemaField field, long value) {
+    NumberType numberType = field.getType().getNumberType();
+    if (numberType == null) {
+      return Long.toString(value);
+    }
+
+    return switch (numberType) {
+      case INTEGER -> Integer.toString((int) value);
+      case LONG -> Long.toString(value);
+      case FLOAT -> Float.toString(Float.intBitsToFloat((int) value));
+      case DOUBLE -> Double.toString(Double.longBitsToDouble(value));
+      case DATE -> Long.toString(value);
+    };
+  }
+
+  private static void logDocsMatchingRanges(LeafReader reader, int[] docsMatchingRanges) {
+    for (int ii = 0; ii < docsMatchingRanges.length; ii++) {
+      if (0 == docsMatchingRanges[ii]) continue;
+      switch (ii) {
+        case 0:
+          // document loss
+          log.error(
+              "Splitting {}: {} documents belong to no shards and will be dropped",
+              reader,
+              docsMatchingRanges[ii]);
+          break;
+        case 1:
+          // normal case, each document moves to one of the sub-shards
+          log.info(
+              "Splitting {}: {} documents will move into a sub-shard",
+              reader,
+              docsMatchingRanges[ii]);
+          break;
+        default:
+          // document duplication
+          log.error(
+              "Splitting {}: {} documents will be moved to multiple ({}) sub-shards",
+              reader,
+              docsMatchingRanges[ii],
+              ii);
+          break;
+      }
+    }
+  }
+
   private static void checkRouterSupportsSplitKey(HashBasedRouter hashRouter, String splitKey) {
     if (splitKey != null && !(hashRouter instanceof CompositeIdRouter)) {
       throw new IllegalStateException(

solr/core/src/test-files/solr/collection1/conf/schema15.xml

@@ -35,6 +35,7 @@
   <fieldType name="float" class="${solr.tests.FloatFieldType}" docValues="${solr.tests.numeric.dv}" precisionStep="0" positionIncrementGap="0"/>
   <fieldType name="long" class="${solr.tests.LongFieldType}" docValues="${solr.tests.numeric.dv}" precisionStep="0" positionIncrementGap="0"/>
   <fieldType name="double" class="${solr.tests.DoubleFieldType}" docValues="${solr.tests.numeric.dv}" precisionStep="0" positionIncrementGap="0"/>
+  <fieldType name="plong" class="solr.LongPointField"/>
 
   <fieldType name="tint" class="${solr.tests.IntegerFieldType}" docValues="${solr.tests.numeric.dv}" precisionStep="8" positionIncrementGap="0"/>
   <fieldType name="tfloat" class="${solr.tests.FloatFieldType}" docValues="${solr.tests.numeric.dv}" precisionStep="8" positionIncrementGap="0"/>
@@ -595,6 +596,7 @@
   <dynamicField name="*_s" type="string" indexed="true" stored="true" multiValued="true"/>
   <dynamicField name="*_ss" type="string" indexed="true" stored="true" multiValued="true"/>
   <dynamicField name="*_l" type="long" indexed="true" stored="true"/>
+  <dynamicField name="*_pl" type="plong" indexed="true" stored="true"/>
   <dynamicField name="*_ll" type="long" indexed="true" stored="true" multiValued="true"/>
   <dynamicField name="*_t" type="text" indexed="true" stored="true"/>
   <dynamicField name="*_tt" type="text" indexed="true" stored="true"/>

solr/core/src/test/org/apache/solr/cloud/api/collections/ShardSplitTest.java

@@ -107,6 +107,7 @@ public void test() throws Exception {
     incompleteOrOverlappingCustomRangeTest();
     splitByUniqueKeyTest();
     splitByRouteFieldTest();
+    splitByNumericRouteFieldTest();
     splitByRouteKeyTest();
 
     // todo can't call waitForThingsToLevelOut because it looks for jettys of all shards
@@ -1009,6 +1010,75 @@ public void splitByRouteFieldTest() throws Exception {
               .query(new SolrQuery("*:*").setParam("shards", "shard1_1"))
               .getResults()
               .getNumFound());
+      assertEquals(101, collectionClient.query(new SolrQuery("*:*")).getResults().getNumFound());
+    }
+  }
+
+  public void splitByNumericRouteFieldTest() throws Exception {
+    log.info("Starting splitByNumericRouteFieldTest");
+    String collectionName = "numericRouteFieldColl";
+    int numShards = 4;
+    int replicationFactor = 2;
+
+    HashMap<String, List<Integer>> collectionInfos = new HashMap<>();
+    String shardField = "shard_pl";
+    try (CloudSolrClient client = createCloudClient(null)) {
+      Map<String, Object> props =
+          Map.of(
+              REPLICATION_FACTOR,
+              replicationFactor,
+              CollectionHandlingUtils.NUM_SLICES,
+              numShards,
+              "router.field",
+              shardField);
+
+      createCollection(collectionInfos, collectionName, props, client);
+    }
+
+    List<Integer> list = collectionInfos.get(collectionName);
+    checkForCollection(collectionName, list);
+
+    waitForRecoveriesToFinish(false);
+
+    getCommonCloudSolrClient();
+    String baseUrl = getBaseUrlFromZk(cloudClient.getClusterState(), collectionName);
+
+    try (SolrClient collectionClient = getHttpSolrClient(baseUrl, collectionName)) {
+      ClusterState clusterState = cloudClient.getClusterState();
+      final DocRouter router = clusterState.getCollection(collectionName).getRouter();
+      Slice shard1 = clusterState.getCollection(collectionName).getSlice(SHARD1);
+      DocRouter.Range shard1Range =
+          shard1.getRange() != null ? shard1.getRange() : router.fullRange();
+      final List<DocRouter.Range> ranges = router.partitionRange(2, shard1Range);
+      final int[] docCounts = new int[ranges.size()];
+
+      for (int i = 100; i <= 200; i++) {
+        collectionClient.add(getDoc(id, i, "n_ti", i, shardField, i));
+        int idx = getHashRangeIdx(router, ranges, Integer.toString(i));
+        if (idx != -1) {
+          docCounts[idx]++;
+        }
+      }
+
+      collectionClient.commit();
+
+      trySplit(collectionName, null, SHARD1, 3);
+
+      waitForRecoveriesToFinish(collectionName, false);
+
+      assertEquals(
+          docCounts[0],
+          collectionClient
+              .query(new SolrQuery("*:*").setParam("shards", "shard1_0"))
+              .getResults()
+              .getNumFound());
+      assertEquals(
+          docCounts[1],
+          collectionClient
+              .query(new SolrQuery("*:*").setParam("shards", "shard1_1"))
+              .getResults()
+              .getNumFound());
+      assertEquals(101, collectionClient.query(new SolrQuery("*:*")).getResults().getNumFound());
     }
   }