From eb13d5bc8b3b0497ce2aca3d99e37884dc54599a Mon Sep 17 00:00:00 2001
From: Ignacio Vera <ivera@apache.org>
Date: Thu, 16 Jan 2020 16:11:46 +0100
Subject: [PATCH] LUCENE-9144: Fix error message on OneDimensionBKDWriter when
 too many points are added to the writer. (#1178)

---
 lucene/CHANGES.txt                            |   3 +
 .../org/apache/lucene/util/bkd/BKDWriter.java |   2 +-
 .../org/apache/lucene/util/bkd/TestBKD.java   | 111 ++++++++++++++++++
 3 files changed, 115 insertions(+), 1 deletion(-)

diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 97f4887f6ea..9bfa3668d29 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -132,6 +132,9 @@ Bug Fixes
 * LUCENE-9115: NRTCachingDirectory no longer caches files of unknown size.
   (Adrien Grand)
 
+* LUCENE-9144: Fix error message on OneDimensionBKDWriter when too many points are added to the writer.
+  (Ignacio Vera)
+
 Other
 ---------------------
 
diff --git a/lucene/core/src/java/org/apache/lucene/util/bkd/BKDWriter.java b/lucene/core/src/java/org/apache/lucene/util/bkd/BKDWriter.java
index 27211c1a7b5..f0183fa4ea4 100644
--- a/lucene/core/src/java/org/apache/lucene/util/bkd/BKDWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/util/bkd/BKDWriter.java
@@ -568,7 +568,7 @@ public class BKDWriter implements Closeable {
       leafCount++;
 
       if (valueCount + leafCount > totalPointCount) {
-        throw new IllegalStateException("totalPointCount=" + totalPointCount + " was passed when we were created, but we just hit " + pointCount + leafCount + " values");
+        throw new IllegalStateException("totalPointCount=" + totalPointCount + " was passed when we were created, but we just hit " + (valueCount + leafCount) + " values");
       }
 
       if (leafCount == maxPointsInLeafNode) {
diff --git a/lucene/core/src/test/org/apache/lucene/util/bkd/TestBKD.java b/lucene/core/src/test/org/apache/lucene/util/bkd/TestBKD.java
index 9c0b2b92eec..118be5acdfc 100644
--- a/lucene/core/src/test/org/apache/lucene/util/bkd/TestBKD.java
+++ b/lucene/core/src/test/org/apache/lucene/util/bkd/TestBKD.java
@@ -1457,4 +1457,115 @@ public class TestBKD extends LuceneTestCase {
       }
     });
   }
+
+  public void testTooManyPoints() throws Exception {
+    Directory dir = newDirectory();
+    final int numValues = 10;
+    final int numPointsAdded = 50; // exceeds totalPointCount
+    final int numBytesPerDim = TestUtil.nextInt(random(), 1, 4);
+    final byte[] pointValue = new byte[numBytesPerDim];
+    BKDWriter w = new BKDWriter(numValues, dir, "_temp", 1, 1, numBytesPerDim, 2,
+        BKDWriter.DEFAULT_MAX_MB_SORT_IN_HEAP, numValues);
+    for(int i=0;i<numValues;i++) {
+      random().nextBytes(pointValue);
+      w.add(pointValue, i);
+    }
+    random().nextBytes(pointValue);
+    IllegalStateException ex = expectThrows(IllegalStateException.class, () -> { w.add(pointValue, numValues);});
+    assertEquals("totalPointCount=10 was passed when we were created, but we just hit 11 values", ex.getMessage());
+    w.close();
+    dir.close();
+  }
+
+  public void testTooManyPoints1D() throws Exception {
+    Directory dir = newDirectory();
+    final int numValues = 10;
+    final int numPointsAdded = 50; // exceeds totalPointCount
+    final int numBytesPerDim = TestUtil.nextInt(random(), 1, 4);
+    final byte[][] pointValue = new byte[11][numBytesPerDim];
+    BKDWriter w = new BKDWriter(numValues + 1, dir, "_temp", 1, 1, numBytesPerDim, 2,
+        BKDWriter.DEFAULT_MAX_MB_SORT_IN_HEAP, numValues);
+    for(int i=0;i<numValues + 1;i++) {
+      random().nextBytes(pointValue[i]);
+    }
+    MutablePointValues val = new MutablePointValues() {
+      @Override
+      public void getValue(int i, BytesRef packedValue) {
+        packedValue.bytes = pointValue[i];
+        packedValue.offset = 0;
+        packedValue.length = numBytesPerDim;
+      }
+
+      @Override
+      public byte getByteAt(int i, int k) {
+        return pointValue[i][k];
+      }
+
+      @Override
+      public int getDocID(int i) {
+        return i;
+      }
+
+      @Override
+      public void swap(int i, int j) {
+        byte[] temp = pointValue[i];
+        pointValue[i] = pointValue[j];
+        pointValue[j] = temp;
+      }
+
+      @Override
+      public void intersect(IntersectVisitor visitor) throws IOException {
+        for (int i = 0; i < size(); i++) {
+          visitor.visit(i, pointValue[i]);
+        }
+      }
+
+      @Override
+      public long estimatePointCount(IntersectVisitor visitor) {
+        return 11;
+      }
+
+      @Override
+      public byte[] getMinPackedValue() {
+        return new byte[numBytesPerDim];
+      }
+
+      @Override
+      public byte[] getMaxPackedValue() {
+        return new byte[numBytesPerDim];
+      }
+
+      @Override
+      public int getNumDataDimensions() {
+        return 1;
+      }
+
+      @Override
+      public int getNumIndexDimensions() {
+        return 1;
+      }
+
+      @Override
+      public int getBytesPerDimension() {
+        return numBytesPerDim;
+      }
+
+      @Override
+      public long size() {
+        return 11;
+      }
+
+      @Override
+      public int getDocCount() {
+        return 11;
+      }
+    };
+    try (IndexOutput out = dir.createOutput("bkd", IOContext.DEFAULT)) {
+      IllegalStateException ex = expectThrows(IllegalStateException.class, () -> { w.writeField(out, "", val);});
+      assertEquals("totalPointCount=10 was passed when we were created, but we just hit 11 values", ex.getMessage());
+      w.close();
+    }
+    dir.close();
+  }
+
 }