iprithv commented on code in PR #16214:
URL: https://github.com/apache/lucene/pull/16214#discussion_r3374974915
##########
lucene/core/src/java/org/apache/lucene/index/ReadersAndUpdates.java:
##########
@@ -460,6 +510,187 @@ public long cost() {
}
}
+ /**
+ * Writes a new generation of KNN vector files for each field that has
pending vector updates. For
+ * each such field we bump a new {@code vectorGen}, write {@code
.vec/.vemf/.vex/.vem} files at
+ * that gen suffix where the vectors are the existing field vectors overlaid
with the updates, and
+ * the HNSW graph is rebuilt eagerly (via the standard {@link
KnnVectorsWriter} flush path). The
+ * base reader (gen == -1) plus other untouched fields are left in place.
Vector analogue of
+ * {@link #handleDVUpdates}.
+ */
+ @SuppressWarnings("unchecked")
+ private synchronized void handleVectorUpdates(
+ FieldInfos infos,
+ Directory dir,
+ KnnVectorsFormat vectorsFormat,
+ final SegmentReader reader,
+ Map<Integer, Set<String>> fieldFiles,
+ long maxDelGen,
+ InfoStream infoStream,
+ boolean deferVectorGraphRebuild)
+ throws IOException {
+ for (Entry<String, List<KnnVectorFieldUpdates>> ent :
pendingVectorUpdates.entrySet()) {
+ final String field = ent.getKey();
+ final List<KnnVectorFieldUpdates> updates = ent.getValue();
+ final List<KnnVectorFieldUpdates> updatesToApply = new ArrayList<>();
+ for (KnnVectorFieldUpdates update : updates) {
+ if (update.delGen <= maxDelGen) {
+ updatesToApply.add(update);
+ }
+ }
+ if (updatesToApply.isEmpty()) {
+ continue;
+ }
+
+ // Reject quantized formats: only the unquantized Lucene99 HNSW format
is supported for
+ // in-place vector updates.
+ KnnVectorsFormat perFieldFormat = vectorsFormat;
+ if (vectorsFormat instanceof PerFieldKnnVectorsFormat perField) {
+ perFieldFormat = perField.getKnnVectorsFormatForField(field);
+ }
+ if (perFieldFormat instanceof Lucene99HnswVectorsFormat == false) {
+ throw new UnsupportedOperationException(
+ "in-place vector update is only supported for "
+ + Lucene99HnswVectorsFormat.class.getSimpleName()
+ + " but field ["
+ + field
+ + "] uses "
+ + perFieldFormat.getClass().getSimpleName());
+ }
+
+ final long nextVectorGen = info.getNextVectorGen();
+ final String segmentSuffix = Long.toString(nextVectorGen,
Character.MAX_RADIX);
+ final IOContext updatesContext = IOContext.flush(new
FlushInfo(info.info.maxDoc(), 0));
+
+ // Stamp the new vectorGen onto the (shared, cloned) FieldInfo so it is
persisted by
+ // writeFieldInfosGen and picked up at read time. A single-field
FieldInfos is handed to the
+ // vectors writer for this gen.
+ final FieldInfo fieldInfo = infos.fieldInfo(field);
+ assert fieldInfo != null && fieldInfo.hasVectorValues();
+ fieldInfo.setVectorGen(nextVectorGen);
+ final FieldInfos fieldInfos = new FieldInfos(new FieldInfo[]
{fieldInfo});
+
+ // Choose the writer format for this generation. When deferring the
graph rebuild we write
+ // only
+ // the flat vectors and skip building the HNSW graph: the gen's
".vex/.vem" carry an empty
+ // graph
+ // (vectorIndexLength == 0), so the reader falls back to an exact scan
on this segment until
+ // the
+ // next merge rebuilds the graph using the codec's normal format. The
graph build is
+ // suppressed
+ // by raising the "tiny segment" threshold above the segment size (so
shouldCreateGraph() is
+ // always false for this write). We wrap it in a
PerFieldKnnVectorsFormat so file naming and
+ // the
+ // per-field suffix attributes match what the (always per-field) reader
reconstructs.
+ final KnnVectorsFormat writeFormat;
+ if (deferVectorGraphRebuild) {
+ final KnnVectorsFormat noGraphFormat =
+ new Lucene99HnswVectorsFormat(
+ Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN,
+ Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH,
+ info.info.maxDoc() + 1);
+ writeFormat =
+ new PerFieldKnnVectorsFormat() {
+ @Override
+ public KnnVectorsFormat getKnnVectorsFormatForField(String f) {
+ return noGraphFormat;
+ }
+ };
+ } else {
+ writeFormat = vectorsFormat;
+ }
+
+ final TrackingDirectoryWrapper trackingDir = new
TrackingDirectoryWrapper(dir);
+ final SegmentWriteState state =
+ new SegmentWriteState(
+ null, trackingDir, info.info, fieldInfos, null, updatesContext,
segmentSuffix);
+
+ try (KnnVectorsWriter writer = writeFormat.fieldsWriter(state)) {
+ final KnnVectorFieldUpdates.Iterator mergedIterator =
+
KnnVectorFieldUpdates.mergedIterator(toIteratorArray(updatesToApply));
+ if (fieldInfo.getVectorEncoding() == VectorEncoding.FLOAT32) {
+ KnnFieldVectorsWriter<float[]> fieldWriter =
+ (KnnFieldVectorsWriter<float[]>) writer.addField(fieldInfo);
+ writeOverlayFloatVectors(reader, field, mergedIterator, fieldWriter);
+ } else {
+ KnnFieldVectorsWriter<byte[]> fieldWriter =
+ (KnnFieldVectorsWriter<byte[]>) writer.addField(fieldInfo);
+ writeOverlayByteVectors(reader, field, mergedIterator, fieldWriter);
+ }
+ writer.flush(info.info.maxDoc(), null);
+ writer.finish();
+ }
+
+ info.advanceVectorGen();
+ assert !fieldFiles.containsKey(fieldInfo.number);
+ fieldFiles.put(fieldInfo.number, trackingDir.getCreatedFiles());
+ }
+ }
+
+ private static KnnVectorFieldUpdates.Iterator[] toIteratorArray(
+ List<KnnVectorFieldUpdates> updatesToApply) {
+ KnnVectorFieldUpdates.Iterator[] subs =
+ new KnnVectorFieldUpdates.Iterator[updatesToApply.size()];
+ for (int i = 0; i < subs.length; i++) {
+ subs[i] = updatesToApply.get(i).iterator();
+ }
+ return subs;
+ }
+
+ /**
+ * Feeds the field writer, in docID order, every doc that currently has a
float vector for {@code
+ * field}, substituting the updated vector when the merged update iterator
has one for that doc.
+ */
+ private static void writeOverlayFloatVectors(
Review Comment:
If a user calls updateFloatVectorValue targeting a doc that doesn't actually
have a vector for that field, wouldn't the update just silently disappears? The
overlay loop only iterates docs that already have vectors, so any update
targeting a doc without one gets skipped with no error.
##########
lucene/core/src/java/org/apache/lucene/index/KnnVectorFieldUpdates.java:
##########
@@ -0,0 +1,335 @@
+/*
+ * 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.index;
+
+import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
+
+import java.util.Comparator;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.PriorityQueue;
+
+/**
+ * Holds in-place KNN vector updates of a single field, for a set of documents
within one segment.
+ * This is the vector analogue of {@link DocValuesFieldUpdates}: it
accumulates (docID -> new
+ * vector) entries during resolution, sorts them by docID at {@link #finish()}
(stable, so that the
+ * last update to a docID within a packet wins), and exposes an {@link
Iterator} that the write path
+ * overlays onto the existing vectors.
+ *
+ * @lucene.experimental
+ */
+abstract class KnnVectorFieldUpdates {
+
+ final String field;
+ final VectorEncoding encoding;
+ final long delGen;
+ final int maxDoc;
+ final int dimension;
+ protected int[] docs = new int[8];
+ protected int size;
+ private boolean finished;
+
+ protected KnnVectorFieldUpdates(
+ int maxDoc, long delGen, String field, VectorEncoding encoding, int
dimension) {
+ this.maxDoc = maxDoc;
+ this.delGen = delGen;
+ this.field = field;
+ this.encoding = encoding;
+ this.dimension = dimension;
+ }
+
+ /** Reserves a slot for the next doc and returns its ordinal. Subclasses
store the value. */
+ protected final int reserve(int doc) {
+ if (finished) {
+ throw new IllegalStateException("already finished");
+ }
+ assert doc < maxDoc;
+ if (size == docs.length) {
+ docs = ArrayUtil.grow(docs, size + 1);
+ }
+ docs[size] = doc;
+ return size++;
+ }
+
+ boolean getFinished() {
+ return finished;
+ }
+
+ boolean any() {
+ return size > 0;
+ }
+
+ /** Freezes structures and stably sorts updates by docID (last write to a
docID wins). */
+ final void finish() {
Review Comment:
This allocates a boxed Integer[] array and sorts it with a comparator, then
copies everything into new arrays. DocValuesFieldUpdates.finish() does the same
job in-place with IntroSorter and zero boxing, that's the established pattern
here. No reason to diverge from it, especially since vectors are large so the
extra copies are worse, not better.
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene94/Lucene94FieldInfosFormat.java:
##########
@@ -454,6 +458,8 @@ public void write(
output.writeVInt(fi.getVectorDimension());
output.writeByte((byte) fi.getVectorEncoding().ordinal());
output.writeByte(distFuncToOrd(fi.getVectorSimilarityFunction()));
+ // per-field KNN vector update generation (FORMAT_VECTOR_GEN+)
+ output.writeLong(fi.getVectorGen());
Review Comment:
The new vectorGen long is written unconditionally for every field, even
plain text fields that will never have vectors. That's 8 extra bytes per field
in every index. Could either only write it for fields with hasVectorValues(),
or use a flag bit to signal presence, rather than bumping the format version
and paying the cost everywhere.
##########
lucene/core/src/java/org/apache/lucene/index/BufferedUpdates.java:
##########
@@ -155,6 +159,15 @@ void addBinaryUpdate(BinaryDocValuesUpdate update, int
docIDUpto) {
numFieldUpdates.incrementAndGet();
}
+ void addVectorUpdate(KnnVectorUpdate update, int docIDUpto) {
+ FieldUpdatesBuffer buffer =
+ vectorUpdates.computeIfAbsent(
Review Comment:
The first update for a field gets added twice. computeIfAbsent creates the
buffer with the initial value already stored, then the next line
unconditionally calls buffer.addUpdate() again with the same value. Compare
with addNumericUpdate/addBinaryUpdate which have an if/else to avoid this.
Harmless due to last-write-wins but wastes memory and looks like a bug.
##########
lucene/core/src/java/org/apache/lucene/index/SegmentVectorsProducer.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.index;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.Set;
+import org.apache.lucene.codecs.KnnVectorsReader;
+import org.apache.lucene.internal.hppc.IntObjectHashMap;
+import org.apache.lucene.internal.hppc.LongArrayList;
+import org.apache.lucene.search.AcceptDocs;
+import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.store.Directory;
+
+/**
+ * Encapsulates multiple per-generation {@link KnnVectorsReader}s as one
reader when there are KNN
+ * vector updates. This is the vector analogue of {@link
SegmentDocValuesProducer}.
+ *
+ * <p>Fields with no vector updates (vectorGen == -1) are served from the
shared, immutable base
+ * reader held by {@link SegmentCoreReaders}; that base reader is never closed
by this producer (its
+ * lifecycle is owned by the core). Fields with updates are served from
per-gen readers tracked by
+ * {@link SegmentKnnVectors} via reference counting.
+ */
+class SegmentVectorsProducer extends KnnVectorsReader {
+
+ private final FieldInfos fieldInfos;
+ final IntObjectHashMap<KnnVectorsReader> readersByField = new
IntObjectHashMap<>();
+ // distinct gen readers (excluding the base core reader) for checkIntegrity
+ final Set<KnnVectorsReader> genReaders =
+ Collections.newSetFromMap(new IdentityHashMap<KnnVectorsReader,
Boolean>());
+ final LongArrayList vectorGens = new LongArrayList();
+
+ /**
+ * Creates a new producer that handles updated KNN vector fields.
+ *
+ * @param si commit point
+ * @param dir directory
+ * @param baseCoreReader the shared base (gen == -1) reader from {@link
SegmentCoreReaders}; may
+ * be {@code null} if the core has no vector fields
+ * @param allInfos all fieldinfos including updated ones
+ * @param segKnnVectors per-gen reader cache
+ */
+ SegmentVectorsProducer(
+ SegmentCommitInfo si,
+ Directory dir,
+ KnnVectorsReader baseCoreReader,
+ FieldInfos allInfos,
+ SegmentKnnVectors segKnnVectors)
+ throws IOException {
+ this.fieldInfos = allInfos;
+ try {
+ for (FieldInfo fi : allInfos) {
+ if (fi.hasVectorValues() == false) {
+ continue;
+ }
+ long vectorGen = fi.getVectorGen();
+ if (vectorGen == -1) {
+ // served from the shared base core reader; not tracked here
+ readersByField.put(fi.number, baseCoreReader);
+ } else {
+ assert !vectorGens.contains(vectorGen);
Review Comment:
vectorGens is a LongArrayList (a list, not a set). If two fields ever share
a generation in the future, you'd decRef the same gen-reader twice, which would
close it prematurely. SegmentDocValuesProducer handles dvGens with
deduplication awareness. Should at least deduplicate here or use a set to be
safe.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
