clean up related to temp file

Author: john-wagster

server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsWriter.java

@@ -118,16 +118,8 @@ private static void printClusterQualityStatistics(int[][] clusters) {
     }
 
     @Override
-    CentroidSupplier createCentroidSupplier(
-        IndexInput centroidsInput,
-        int numParentCentroids,
-        int numCentroids,
-        int numClusters,
-        FieldInfo fieldInfo,
-        float[] globalCentroid,
-        IntIntMap clusterToCentroidMap
-    ) {
-        return new OffHeapCentroidSupplier(centroidsInput, numParentCentroids, numCentroids, numClusters, fieldInfo, clusterToCentroidMap);
+    CentroidSupplier createCentroidSupplier(IndexInput centroidsInput, int numClusters, FieldInfo fieldInfo) {
+        return new OffHeapCentroidSupplier(centroidsInput, numClusters, fieldInfo);
     }
 
     private static void writeQuantizedCentroid(
@@ -181,7 +173,7 @@ static IntIntMap writePartitionsAndCentroids(
                     centroidOutput
                 );
                 // TODO: put at the end of the parents region
-                centroidOutput.writeInt(centroidPartition.childOrdinal());
+                centroidOutput.writeInt(centroidPartition.childOffset());
                 centroidOutput.writeInt(centroidPartition.size());
             }
         }
@@ -252,7 +244,14 @@ CentroidAssignments calculateAndWriteCentroids(
         return calculateAndWriteCentroids(fieldInfo, floatVectorValues, centroidOutput, globalCentroid);
     }
 
-    record CentroidPartition(float[] centroid, int childOrdinal, int size, int[] assignments) {}
+    /**
+     *
+     * @param centroid the parent centroid of some set of children
+     * @param childOffset the offset of the first child within the partition defined by this parent
+     * @param size the number of children in this partition
+     * @param assignments the set of centroid ordinals (potentially duplicative) of child centroids that belong to this parent
+     */
+    record CentroidPartition(float[] centroid, int childOffset, int size, int[] assignments) {}
 
     /**
      * Calculate the centroids for the given field and write them to the given centroid output.
@@ -426,30 +425,16 @@ static void writeQuantizedValue(IndexOutput indexOutput, byte[] binaryValue, Opt
 
     static class OffHeapCentroidSupplier implements CentroidSupplier {
         private final IndexInput centroidsInput;
-        private final int numClusters;
         private final int dimension;
         private final float[] scratch;
-        private final long rawCentroidOffset;
         private int currOrd = -1;
-        private final IntIntMap clusterToCentroidMap;
-
-        OffHeapCentroidSupplier(
-            IndexInput centroidsInput,
-            int numParentCentroids,
-            int numCentroids,
-            int numClusters,
-            FieldInfo info,
-            IntIntMap clusterToCentroidMap
-        ) {
+        private final int numClusters;
+
+        OffHeapCentroidSupplier(IndexInput centroidsInput, int numClusters, FieldInfo info) {
             this.centroidsInput = centroidsInput;
-            this.numClusters = numClusters;
             this.dimension = info.getVectorDimension();
             this.scratch = new float[dimension];
-            long quantizedVectorByteSize = dimension + 3 * Float.BYTES + Short.BYTES;
-            long quantizedVectorNodeByteSize = quantizedVectorByteSize + Integer.BYTES;
-            long parentNodeByteSize = quantizedVectorByteSize + 2 * Integer.BYTES;
-            this.rawCentroidOffset = numParentCentroids * parentNodeByteSize + numCentroids * quantizedVectorNodeByteSize;
-            this.clusterToCentroidMap = clusterToCentroidMap;
+            this.numClusters = numClusters;
         }
 
         @Override
@@ -458,19 +443,13 @@ public int size() {
         }
 
         @Override
-        public float[] centroid(int clusterOrdinal) throws IOException {
-            if (clusterOrdinal == currOrd) {
+        public float[] centroid(int centroidOrdinal) throws IOException {
+            if (centroidOrdinal == currOrd) {
                 return scratch;
             }
-            int centroidOrdinal;
-            if (clusterToCentroidMap != null) {
-                centroidOrdinal = clusterToCentroidMap.get(clusterOrdinal);
-            } else {
-                centroidOrdinal = clusterOrdinal;
-            }
-            centroidsInput.seek(rawCentroidOffset + (long) centroidOrdinal * dimension * Float.BYTES);
+            centroidsInput.seek((long) centroidOrdinal * dimension * Float.BYTES);
             centroidsInput.readFloats(scratch, 0, dimension);
-            this.currOrd = clusterOrdinal;
+            this.currOrd = centroidOrdinal;
             return scratch;
         }
     }

server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsReader.java

@@ -259,7 +259,9 @@ public final void search(String field, float[] target, KnnCollector knnCollector
 
         FieldEntry entry = fields.get(fieldInfo.number);
 
+        // TODO: pass numClusters via meta rather than computing it here?
         int numCentroids = entry.childCentroidCount();
+        int numClusters = entry.childCentroidCount() / 2;
 
         CentroidWClusterOffsetQueryScorer centroidQueryScorer = getChildCentroidScorer(
             fieldInfo,
@@ -273,9 +275,9 @@ public final void search(String field, float[] target, KnnCollector knnCollector
             // scaling by the number of centroids vs. the nearest neighbors requested
             // not perfect, but a comparative heuristic.
             // we might want to utilize the total vector count as well, but this is a good start
-            nProbe = (int) Math.round(Math.log10(centroidQueryScorer.size()) * Math.sqrt(knnCollector.k()));
+            nProbe = (int) Math.round(Math.log10(numClusters) * Math.sqrt(knnCollector.k()));
             // clip to be between 1 and the number of centroids
-            nProbe = Math.max(Math.min(nProbe, centroidQueryScorer.size()), 1);
+            nProbe = Math.max(Math.min(nProbe, numClusters), 1);
         }
 
         int centroidsVisited = 0;

server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsWriter.java

@@ -142,15 +142,7 @@ abstract long[] buildAndWritePostingsLists(
         int[][] assignmentsByCluster
     ) throws IOException;
 
-    abstract CentroidSupplier createCentroidSupplier(
-        IndexInput centroidsInput,
-        int numParentCentroids,
-        int numCentroids,
-        int numClusters,
-        FieldInfo fieldInfo,
-        float[] globalCentroid,
-        IntIntMap clusterToCentroidMap
-    ) throws IOException;
+    abstract CentroidSupplier createCentroidSupplier(IndexInput centroidsInput, int numCentroids, FieldInfo fieldInfo) throws IOException;
 
     @Override
     public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
@@ -299,25 +291,35 @@ private void mergeOneFieldIVF(FieldInfo fieldInfo, MergeState mergeState) throws
             final int numClusters;
             final int[][] assignmentsByCluster;
             final float[] calculatedGlobalCentroid = new float[fieldInfo.getVectorDimension()];
-            IntIntMap clusterToCentroidMap;
             String centroidTempName = null;
             IndexOutput centroidTemp = null;
             success = false;
             try {
                 centroidTemp = mergeState.segmentInfo.dir.createTempOutput(mergeState.segmentInfo.name, "civf_", IOContext.DEFAULT);
                 centroidTempName = centroidTemp.getName();
+
+                centroidOffset = ivfCentroids.alignFilePointer(Float.BYTES);
                 CentroidAssignments centroidAssignments = calculateAndWriteCentroids(
                     fieldInfo,
                     getFloatVectorValues(fieldInfo, docs, vectors, numVectors),
-                    centroidTemp,
+                    ivfCentroids,
                     mergeState,
                     calculatedGlobalCentroid
                 );
+                centroidLength = ivfCentroids.getFilePointer() - centroidOffset;
+
+                final ByteBuffer buffer = ByteBuffer.allocate(fieldInfo.getVectorDimension() * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN);
+                for (int i = 0; i < centroidAssignments.centroids().length; i++) {
+                    float[] centroid = centroidAssignments.centroids()[i];
+                    buffer.asFloatBuffer().put(centroid);
+                    centroidTemp.writeBytes(buffer.array(), buffer.array().length);
+                }
+
                 numParentCentroids = centroidAssignments.numParentCentroids();
+                // TODO: remove this from meta when not used on the read side; we currently use it in degenerate cases to read all centroids
                 numCentroids = centroidAssignments.numCentroids();
                 numClusters = centroidAssignments.centroids().length;
                 assignmentsByCluster = centroidAssignments.assignmentsByCluster();
-                clusterToCentroidMap = centroidAssignments.clusterToCentroidMap();
                 success = true;
             } finally {
                 if (success == false && centroidTempName != null) {
@@ -327,28 +329,16 @@ private void mergeOneFieldIVF(FieldInfo fieldInfo, MergeState mergeState) throws
             }
             try {
                 if (numCentroids == 0) {
-                    centroidOffset = ivfCentroids.getFilePointer();
                     writeMeta(fieldInfo, numParentCentroids, numCentroids, centroidOffset, 0, new long[0], null);
                     CodecUtil.writeFooter(centroidTemp);
                     IOUtils.close(centroidTemp);
                     return;
                 }
                 CodecUtil.writeFooter(centroidTemp);
                 IOUtils.close(centroidTemp);
-                centroidOffset = ivfCentroids.alignFilePointer(Float.BYTES);
                 try (IndexInput centroidsInput = mergeState.segmentInfo.dir.openInput(centroidTempName, IOContext.DEFAULT)) {
-                    ivfCentroids.copyBytes(centroidsInput, centroidsInput.length() - CodecUtil.footerLength());
-                    centroidLength = ivfCentroids.getFilePointer() - centroidOffset;
 
-                    CentroidSupplier centroidSupplier = createCentroidSupplier(
-                        centroidsInput,
-                        numParentCentroids,
-                        numCentroids,
-                        numClusters,
-                        fieldInfo,
-                        calculatedGlobalCentroid,
-                        clusterToCentroidMap
-                    );
+                    CentroidSupplier centroidSupplier = createCentroidSupplier(centroidsInput, numClusters, fieldInfo);
 
                     // build a float vector values with random access
                     // build centroids