Add dense bulk fast path for multi-dimensional points via BinaryColumn

Author: iprithv

lucene/CHANGES.txt

@@ -100,6 +100,8 @@ API Changes
 
 New Features
 ---------------------
+* GITHUB#16209: Add dense bulk path for multi-dimensional point fields via BinaryColumn. (Prithvi S)
+
 * GITHUB#15505: Upgrade snowball to 2d2e312df56f2ede014a4ffb3e91e6dea43c24be. New stemmer: PolishStemmer (and
   PolishSnowballAnalyzer in the stempel package) (Justas Sakalauskas, Dawid Weiss)
 

lucene/core/src/java/org/apache/lucene/document/column/BinaryColumn.java

@@ -21,7 +21,8 @@
 import org.apache.lucene.util.BytesRef;
 
 /**
- * A {@link Column} that provides variable-size binary values via a tuple cursor. Used for {@link
+ * A {@link Column} that provides variable-size binary values via a tuple cursor, and dense values
+ * via a {@link BinaryValuesCursor}. Used for {@link
  * org.apache.lucene.index.DocValuesType#BINARY BINARY}, {@link
  * org.apache.lucene.index.DocValuesType#SORTED SORTED}, and {@link
  * org.apache.lucene.index.DocValuesType#SORTED_SET SORTED_SET} doc values, and for stored/indexed
@@ -54,4 +55,15 @@ public StoredValue.Type storedType() {
 
   /** Returns a fresh tuple cursor starting at the beginning of the batch. */
   public abstract ObjectTupleCursor<BytesRef> tuples();
+
+  /**
+   * Returns a fresh values cursor iterating dense {@link BytesRef} values for doc-ids {@code [0,
+   * numDocs)}. Must be overridden when {@link #density()} is {@link Column.Density#DENSE DENSE};
+   * the default implementation throws {@link UnsupportedOperationException} and is never called for
+   * {@link Column.Density#SPARSE SPARSE} columns.
+   */
+  public BinaryValuesCursor values() {
+    throw new UnsupportedOperationException(
+        "values() requires density() == DENSE for column \"" + name() + "\"");
+  }
 }

lucene/core/src/java/org/apache/lucene/document/column/BinaryValuesCursor.java

@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.document.column;
+
+import org.apache.lucene.util.BytesRef;
+
+/**
+ * A values cursor over a dense {@link BinaryColumn}. The cursor produces exactly {@link #size()}
+ * {@link BytesRef} values for consecutive batch-local doc-ids starting at 0, one per call to {@link
+ * #nextValue()}.
+ *
+ * @lucene.experimental
+ */
+public abstract class BinaryValuesCursor {
+
+  private final int size;
+
+  /**
+   * Creates a cursor that will produce exactly {@code size} values, one per batch-local doc-id in
+   * {@code [0, size)}. {@code size} is fixed for the cursor's lifetime and must equal the dense
+   * column's {@code numDocs}.
+   *
+   * <p>Lucene's internal indexing paths will not consume past {@code size}. Defensive throws on
+   * overrun are still encouraged to catch misuse from external callers.
+   */
+  protected BinaryValuesCursor(int size) {
+    this.size = size;
+  }
+
+  /** Total number of values this cursor will produce. */
+  public final int size() {
+    return size;
+  }
+
+  /**
+   * Returns the next {@link BytesRef} value. Must not be called more than {@link #size()} times. The
+   * returned {@link BytesRef} is only valid until the next call to {@link #nextValue()}.
+   */
+  public abstract BytesRef nextValue();
+}

lucene/core/src/java/org/apache/lucene/document/column/package-info.java

@@ -51,8 +51,9 @@
  *   <li>A tuple cursor (e.g. {@link org.apache.lucene.document.column.LongTupleCursor}, {@link
  *       org.apache.lucene.document.column.ObjectTupleCursor}) yields {@code (batchDocID, value)}
  *       pairs in non-decreasing doc-id order. Always available.
- *   <li>A bulk values cursor (e.g. {@link org.apache.lucene.document.column.LongValuesCursor})
- *       feeds dense data directly into the underlying writer. Required when {@link
+ *   <li>A bulk values cursor (e.g. {@link org.apache.lucene.document.column.LongValuesCursor},
+ *       {@link org.apache.lucene.document.column.BinaryValuesCursor}) feeds dense data directly
+ *       into the underlying writer. Required when {@link
  *       org.apache.lucene.document.column.Column#density()} is {@link
  *       org.apache.lucene.document.column.Column.Density#DENSE DENSE} and consulted only in that
  *       case.

lucene/core/src/java/org/apache/lucene/index/IndexingChain.java

@@ -44,6 +44,7 @@
 import org.apache.lucene.document.NumericDocValuesField;
 import org.apache.lucene.document.StoredValue;
 import org.apache.lucene.document.column.BinaryColumn;
+import org.apache.lucene.document.column.BinaryValuesCursor;
 import org.apache.lucene.document.column.Column;
 import org.apache.lucene.document.column.ColumnBatch;
 import org.apache.lucene.document.column.ColumnFieldAdapter;
@@ -1083,6 +1084,20 @@ private static void processBinaryColumn(
     final DocValuesType dvType = fieldType.docValuesType();
     final boolean hasPoints = fieldType.pointDimensionCount() != 0;
     final PointValuesWriter pointWriter = hasPoints ? pf.pointValuesWriter : null;
+
+    // Dense fast path: points-only with no doc values, stored, or index options.
+    if (column.density() == Column.Density.DENSE
+        && dvType == DocValuesType.NONE
+        && hasPoints
+        && fieldType.stored() == false
+        && fieldType.indexOptions() == IndexOptions.NONE) {
+      BinaryValuesCursor cursor = column.values();
+      ColumnValidation.checkDenseCount(column, cursor.size(), numDocs);
+      int packedLength = fieldType.pointDimensionCount() * fieldType.pointNumBytes();
+      pointWriter.addDenseNDValues(baseDocID, cursor, packedLength);
+      return;
+    }
+
     final ObjectTupleCursor<BytesRef> cursor = column.tuples();
 
     if (dvType == DocValuesType.NONE) {

lucene/core/src/java/org/apache/lucene/index/PointValuesWriter.java

@@ -20,6 +20,7 @@
 import org.apache.lucene.codecs.MutablePointTree;
 import org.apache.lucene.codecs.PointsReader;
 import org.apache.lucene.codecs.PointsWriter;
+import org.apache.lucene.document.column.BinaryValuesCursor;
 import org.apache.lucene.document.column.LongValuesCursor;
 import org.apache.lucene.store.DataOutput;
 import org.apache.lucene.util.ArrayUtil;
@@ -48,6 +49,9 @@ class PointValuesWriter {
   private static final int POINTS_BUFFER_INT_VALUES = POINTS_BUFFER_BYTES / Integer.BYTES;
   private static final int POINTS_BUFFER_LONG_VALUES = POINTS_BUFFER_BYTES / Long.BYTES;
 
+  /** Minimum number of values to process per chunk in the dense N-D bulk path. */
+  private static final int MIN_VALUES_PER_CHUNK = 64;
+
   private byte[] densePointsBuffer;
 
   PointValuesWriter(Counter bytesUsed, FieldInfo fieldInfo) {
@@ -98,7 +102,7 @@ void addDense1DIntValues(int firstDocID, LongValuesCursor cursor) throws IOExcep
     if (size == 0) {
       return;
     }
-    final long ramBefore = reserveDense1D(firstDocID, size);
+    final long ramBefore = reserveDense(firstDocID, size);
     final byte[] buffer = pointsDenseBuffer();
     int remaining = size;
     while (remaining > 0) {
@@ -107,7 +111,7 @@ void addDense1DIntValues(int firstDocID, LongValuesCursor cursor) throws IOExcep
       bytesOut.writeBytes(buffer, 0, chunk * Integer.BYTES);
       remaining -= chunk;
     }
-    commitDense1D(firstDocID, size, ramBefore);
+    commitDense(firstDocID, size, ramBefore);
   }
 
   void addDense1DLongValues(int firstDocID, LongValuesCursor cursor) throws IOException {
@@ -116,7 +120,7 @@ void addDense1DLongValues(int firstDocID, LongValuesCursor cursor) throws IOExce
     if (size == 0) {
       return;
     }
-    final long ramBefore = reserveDense1D(firstDocID, size);
+    final long ramBefore = reserveDense(firstDocID, size);
     final byte[] dense = pointsDenseBuffer();
     int remaining = size;
     while (remaining > 0) {
@@ -125,7 +129,53 @@ void addDense1DLongValues(int firstDocID, LongValuesCursor cursor) throws IOExce
       bytesOut.writeBytes(dense, 0, chunk * Long.BYTES);
       remaining -= chunk;
     }
-    commitDense1D(firstDocID, size, ramBefore);
+    commitDense(firstDocID, size, ramBefore);
+  }
+
+  /**
+   * Bulk-adds dense N-dimensional packed point values from a {@link BinaryValuesCursor}. Each value
+   * is a pre-encoded packed byte array of {@code packedLength} bytes.
+   */
+  void addDenseNDValues(int firstDocID, BinaryValuesCursor cursor, int packedLength)
+      throws IOException {
+    if (fieldInfo.getPointDimensionCount() * fieldInfo.getPointNumBytes() != packedLength) {
+      throw new IllegalArgumentException(
+          "field="
+              + fieldInfo.name
+              + ": packedLength="
+              + packedLength
+              + " does not match pointDimensionCount="
+              + fieldInfo.getPointDimensionCount()
+              + " * pointNumBytes="
+              + fieldInfo.getPointNumBytes());
+    }
+    final int size = cursor.size();
+    if (size == 0) {
+      return;
+    }
+    final long ramBefore = reserveDense(firstDocID, size);
+    final byte[] buffer = pointsDenseBuffer(packedLength);
+    final int valuesPerBuffer = buffer.length / packedLength;
+    int remaining = size;
+    while (remaining > 0) {
+      int chunk = Math.min(valuesPerBuffer, remaining);
+      for (int i = 0; i < chunk; i++) {
+        BytesRef value = cursor.nextValue();
+        if (value.length != packedLength) {
+          throw new IllegalArgumentException(
+              "field="
+                  + fieldInfo.name
+                  + ": point value has length="
+                  + value.length
+                  + " but should be "
+                  + packedLength);
+        }
+        System.arraycopy(value.bytes, value.offset, buffer, i * packedLength, packedLength);
+      }
+      bytesOut.writeBytes(buffer, 0, chunk * packedLength);
+      remaining -= chunk;
+    }
+    commitDense(firstDocID, size, ramBefore);
   }
 
   private byte[] pointsDenseBuffer() {
@@ -135,6 +185,21 @@ private byte[] pointsDenseBuffer() {
     return densePointsBuffer;
   }
 
+  /**
+   * Returns a dense buffer sized to fit at least {@code POINTS_BUFFER_BYTES} worth of packed
+   * values, or a larger buffer if {@code packedLength} requires it to hold at least {@code
+   * MIN_VALUES_PER_CHUNK} values per chunk.
+   */
+  private byte[] pointsDenseBuffer(int packedLength) {
+    final int minBytes = packedLength * MIN_VALUES_PER_CHUNK;
+    if (densePointsBuffer == null) {
+      densePointsBuffer = new byte[Math.max(POINTS_BUFFER_BYTES, minBytes)];
+    } else if (densePointsBuffer.length < minBytes) {
+      densePointsBuffer = new byte[minBytes];
+    }
+    return densePointsBuffer;
+  }
+
   private void validate1DPacked(int byteWidth) {
     if (fieldInfo.getPointDimensionCount() != 1 || fieldInfo.getPointNumBytes() != byteWidth) {
       throw new IllegalArgumentException(
@@ -149,7 +214,7 @@ private void validate1DPacked(int byteWidth) {
     }
   }
 
-  private long reserveDense1D(int firstDocID, int size) {
+  private long reserveDense(int firstDocID, int size) {
     assert firstDocID > lastDocID
         : "firstDocID=" + firstDocID + " must be > lastDocID=" + lastDocID;
     final int oldLength = docIDs.length;
@@ -163,7 +228,7 @@ private long reserveDense1D(int firstDocID, int size) {
     return bytes.ramBytesUsed();
   }
 
-  private void commitDense1D(int firstDocID, int size, long ramBefore) {
+  private void commitDense(int firstDocID, int size, long ramBefore) {
     iwBytesUsed.addAndGet(bytes.ramBytesUsed() - ramBefore);
     numDocs += size;
     lastDocID = firstDocID + size - 1;

lucene/core/src/test/org/apache/lucene/document/column/ColumnBatchTestUtil.java

@@ -578,6 +578,49 @@ public int ordValue() {
     }
   }
 
+  /** Dense {@link BinaryColumn} backed by a contiguous BytesRef array. */
+  public static class ArrayDenseBinaryColumn extends BinaryColumn {
+    private final BytesRef[] values;
+
+    public ArrayDenseBinaryColumn(String name, IndexableFieldType fieldType, BytesRef[] values) {
+      super(name, fieldType, Density.DENSE);
+      this.values = values;
+    }
+
+    @Override
+    public ObjectTupleCursor<BytesRef> tuples() {
+      return new ObjectTupleCursor<>() {
+        int pos = -1;
+
+        @Override
+        public int nextDoc() {
+          pos++;
+          return pos < values.length ? pos : DocIdSetIterator.NO_MORE_DOCS;
+        }
+
+        @Override
+        public BytesRef value() {
+          return values[pos];
+        }
+      };
+    }
+
+    @Override
+    public BinaryValuesCursor values() {
+      return new BinaryValuesCursor(values.length) {
+        int pos = 0;
+
+        @Override
+        public BytesRef nextValue() {
+          if (pos >= values.length) {
+            throw new IllegalStateException("BinaryValuesCursor exhausted: size=" + values.length);
+          }
+          return values[pos++];
+        }
+      };
+    }
+  }
+
   /** Dense {@link DictionaryColumn} backed by a contiguous ordinal array. */
   public static class ArrayDenseDictionaryColumn extends DictionaryColumn {
     private final int[] ords;