From 9402dda99bc2a87c44c202af565ede315b24a9e1 Mon Sep 17 00:00:00 2001 From: ldematte Date: Tue, 2 Sep 2025 15:28:20 +0200 Subject: [PATCH 1/8] Skip one heap copy --- .../gpu/codec/ESGpuHnswVectorsWriter.java | 88 ++++++++++++------- 1 file changed, 54 insertions(+), 34 deletions(-) diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java index 6b738dfeabce9..0020a5dafa405 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java @@ -185,22 +185,19 @@ public long ramBytesUsed() { return total; } - private static final class DatasetOrVectors { + private static final class DatasetOrVectors implements AutoCloseable { private final CuVSMatrix dataset; - private final float[][] vectors; + private final List vectors; - static DatasetOrVectors fromArray(float[][] vectors) { - return new DatasetOrVectors( - vectors.length < MIN_NUM_VECTORS_FOR_GPU_BUILD ? null : CuVSMatrix.ofArray(vectors), - vectors.length < MIN_NUM_VECTORS_FOR_GPU_BUILD ? vectors : null - ); + static DatasetOrVectors fromArray(List vectors) { + return new DatasetOrVectors(null, vectors); } static DatasetOrVectors fromDataset(CuVSMatrix dataset) { return new DatasetOrVectors(dataset, null); } - private DatasetOrVectors(CuVSMatrix dataset, float[][] vectors) { + private DatasetOrVectors(CuVSMatrix dataset, List vectors) { this.dataset = dataset; this.vectors = vectors; validateState(); @@ -213,28 +210,47 @@ private void validateState() { } int size() { - return dataset != null ? (int) dataset.size() : vectors.length; - } - - CuVSMatrix getDataset() { - return dataset; + return dataset != null ? (int) dataset.size() : vectors.size(); } - float[][] getVectors() { - return vectors; + @Override + public void close() { + if (dataset != null) { + dataset.close(); + } } } private void writeField(FieldWriter fieldWriter) throws IOException { - float[][] vectors = fieldWriter.flatFieldVectorsWriter.getVectors().toArray(float[][]::new); - writeFieldInternal(fieldWriter.fieldInfo, DatasetOrVectors.fromArray(vectors)); + var vectors = fieldWriter.flatFieldVectorsWriter.getVectors(); + final DatasetOrVectors datasetOrVectors; + if (vectors.size() < MIN_NUM_VECTORS_FOR_GPU_BUILD) { + // Use vectors/CPU + datasetOrVectors = DatasetOrVectors.fromArray(vectors); + } else { + // Avoid another heap copy (the float[][]) + + // TODO: support other data types + // TODO: another alternative is to use CuVSMatrix.deviceBuilder(), but this requires more effort + // 1. support no-copy CuVSDeviceMatrix as input in CagraIndex + // 2. ensure we are already holding a CuVSResource here + var builder = CuVSMatrix.hostBuilder(vectors.size(), vectors.getFirst().length, CuVSMatrix.DataType.FLOAT); + for (var vector : vectors) { + builder.addVector(vector); + } + datasetOrVectors = DatasetOrVectors.fromDataset(builder.build()); + } + try { + writeFieldInternal(fieldWriter.fieldInfo, datasetOrVectors); + } finally { + datasetOrVectors.close(); + } } private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { // The flatFieldVectorsWriter's flush method, called before this, has already sorted the vectors according to the sortMap. // We can now treat them as a simple, sorted list of vectors. - float[][] vectors = fieldData.flatFieldVectorsWriter.getVectors().toArray(float[][]::new); - writeFieldInternal(fieldData.fieldInfo, DatasetOrVectors.fromArray(vectors)); + writeField(fieldData); } private void writeFieldInternal(FieldInfo fieldInfo, DatasetOrVectors datasetOrVectors) throws IOException { @@ -243,11 +259,11 @@ private void writeFieldInternal(FieldInfo fieldInfo, DatasetOrVectors datasetOrV int[][] graphLevelNodeOffsets = new int[1][]; HnswGraph mockGraph; if (datasetOrVectors.vectors != null) { - float[][] vectors = datasetOrVectors.vectors; + var vectors = datasetOrVectors.vectors; if (logger.isDebugEnabled()) { logger.debug( "Skip building carga index; vectors length {} < {} (min for GPU)", - vectors.length, + vectors.size(), MIN_NUM_VECTORS_FOR_GPU_BUILD ); } @@ -341,12 +357,12 @@ private HnswGraph writeGraph(CuVSMatrix cagraGraph, int[][] levelNodeOffsets) th } // create a graph where every node is connected to every other node - private HnswGraph writeGraph(float[][] vectors, int[][] levelNodeOffsets) throws IOException { - if (vectors.length == 0) { + private HnswGraph writeGraph(List vectors, int[][] levelNodeOffsets) throws IOException { + if (vectors.isEmpty()) { return null; } - int elementCount = vectors.length; - int nodeDegree = vectors.length - 1; + int elementCount = vectors.size(); + int nodeDegree = vectors.size() - 1; levelNodeOffsets[0] = new int[elementCount]; int[] neighbors = new int[nodeDegree]; @@ -440,8 +456,10 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE } } try (IndexInput in = mergeState.segmentInfo.dir.openInput(tempRawVectorsFileName, IOContext.DEFAULT)) { - DatasetOrVectors datasetOrVectors; + var input = FilterIndexInput.unwrapOnlyTest(in); + + final DatasetOrVectors datasetOrVectors; if (input instanceof MemorySegmentAccessInput memorySegmentAccessInput && numVectors >= MIN_NUM_VECTORS_FOR_GPU_BUILD) { var ds = DatasetUtils.getInstance() .fromInput(memorySegmentAccessInput, numVectors, fieldInfo.getVectorDimension(), dataType); @@ -449,9 +467,13 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE } else { // TODO fix for byte vectors var fa = copyVectorsIntoArray(in, fieldInfo, numVectors); - datasetOrVectors = DatasetOrVectors.fromArray(fa); + datasetOrVectors = DatasetOrVectors.fromDataset(fa); + } + try { + writeFieldInternal(fieldInfo, datasetOrVectors); + } finally { + datasetOrVectors.close(); } - writeFieldInternal(fieldInfo, datasetOrVectors); } finally { org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(mergeState.segmentInfo.dir, tempRawVectorsFileName); } @@ -482,15 +504,13 @@ private static int writeByteVectorValues(IndexOutput out, ByteVectorValues vecto return numVectors; } - static float[][] copyVectorsIntoArray(IndexInput in, FieldInfo fieldInfo, int numVectors) throws IOException { + static CuVSMatrix copyVectorsIntoArray(IndexInput in, FieldInfo fieldInfo, int numVectors) throws IOException { final FloatVectorValues floatVectorValues = getFloatVectorValues(fieldInfo, in, numVectors); - float[][] vectors = new float[numVectors][fieldInfo.getVectorDimension()]; - float[] vector; + var builder = CuVSMatrix.hostBuilder(numVectors, fieldInfo.getVectorDimension(), CuVSMatrix.DataType.FLOAT); for (int i = 0; i < numVectors; i++) { - vector = floatVectorValues.vectorValue(i); - System.arraycopy(vector, 0, vectors[i], 0, vector.length); + builder.addVector(floatVectorValues.vectorValue(i)); } - return vectors; + return builder.build(); } private static int writeFloatVectorValues(FieldInfo fieldInfo, IndexOutput out, FloatVectorValues floatVectorValues) From c9fb6b0abf2597c1ae01ed9a478094e39042d4d3 Mon Sep 17 00:00:00 2001 From: ldematte Date: Wed, 3 Sep 2025 11:22:34 +0200 Subject: [PATCH 2/8] Adjust signature to new cuvs-java --- .../elasticsearch/xpack/gpu/codec/CuVSResourceManager.java | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/CuVSResourceManager.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/CuVSResourceManager.java index e7977f28c9c22..81a792195f900 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/CuVSResourceManager.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/CuVSResourceManager.java @@ -135,6 +135,11 @@ public ScopedAccess access() { return delegate.access(); } + @Override + public int deviceId() { + return delegate.deviceId(); + } + @Override public void close() { throw new UnsupportedOperationException("this resource is managed, cannot be closed by clients"); From 4f9b98ae706248b537c791ee4701df1ad081ce3a Mon Sep 17 00:00:00 2001 From: ldematte Date: Wed, 3 Sep 2025 14:53:17 +0200 Subject: [PATCH 3/8] Use memory-mapped MemorySegment from the vectorData IndexOutput/Input --- .../gpu/codec/ESGpuHnswVectorsWriter.java | 74 ++++++++++--------- .../reflect/VectorsFormatReflectionUtils.java | 58 +++++++++++++++ 2 files changed, 99 insertions(+), 33 deletions(-) create mode 100644 x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/reflect/VectorsFormatReflectionUtils.java diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java index 2f41a0236e03d..713b8aed2a8a5 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java @@ -16,6 +16,7 @@ import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter; import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; @@ -42,6 +43,7 @@ import org.elasticsearch.index.codec.vectors.ES814ScalarQuantizedVectorsFormat; import org.elasticsearch.logging.LogManager; import org.elasticsearch.logging.Logger; +import org.elasticsearch.xpack.gpu.reflect.VectorsFormatReflectionUtils; import java.io.IOException; import java.nio.ByteBuffer; @@ -73,6 +75,7 @@ final class ESGpuHnswVectorsWriter extends KnnVectorsWriter { private final CuVSResourceManager cuVSResourceManager; private final SegmentWriteState segmentWriteState; private final IndexOutput meta, vectorIndex; + private final IndexOutput vectorData; private final int M; private final int beamWidth; private final FlatVectorsWriter flatVectorWriter; @@ -94,8 +97,11 @@ final class ESGpuHnswVectorsWriter extends KnnVectorsWriter { this.beamWidth = beamWidth; this.flatVectorWriter = flatVectorWriter; if (flatVectorWriter instanceof ES814ScalarQuantizedVectorsFormat.ES814ScalarQuantizedVectorsWriter) { + vectorData = VectorsFormatReflectionUtils.getQuantizedVectorDataIndexOutput(flatVectorWriter); dataType = CuVSMatrix.DataType.BYTE; } else { + assert flatVectorWriter instanceof Lucene99FlatVectorsWriter; + vectorData = VectorsFormatReflectionUtils.getVectorDataIndexOutput(flatVectorWriter); dataType = CuVSMatrix.DataType.FLOAT; } this.segmentWriteState = state; @@ -148,11 +154,38 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException @Override public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { flatVectorWriter.flush(maxDoc, sortMap); - for (FieldWriter field : fields) { - if (sortMap == null) { - writeField(field); - } else { - writeSortingField(field, sortMap); + + try (IndexInput in = segmentWriteState.segmentInfo.dir.openInput(vectorData.getName(), IOContext.DEFAULT)) { + var input = FilterIndexInput.unwrapOnlyTest(in); + + for (FieldWriter fieldWriter : fields) { + // TODO: is this inefficient? Can we get "size" in another way? + var numVectors = fieldWriter.flatFieldVectorsWriter.getVectors().size(); + + final DatasetOrVectors datasetOrVectors; + if (input instanceof MemorySegmentAccessInput memorySegmentAccessInput && numVectors >= MIN_NUM_VECTORS_FOR_GPU_BUILD) { + // TODO: we are iterating over multiple fields, we probably need to memorySegmentAccessInput.segmentSliceOrNull()? + var ds = DatasetUtils.getInstance() + .fromInput(memorySegmentAccessInput, numVectors, fieldWriter.fieldInfo.getVectorDimension(), dataType); + datasetOrVectors = DatasetOrVectors.fromDataset(ds); + } else { + var builder = CuVSMatrix.hostBuilder(numVectors, fieldWriter.fieldInfo.getVectorDimension(), dataType); + for (var vector : fieldWriter.flatFieldVectorsWriter.getVectors()) { + builder.addVector(vector); + } + + datasetOrVectors = DatasetOrVectors.fromDataset(builder.build()); + } + + try { + if (sortMap == null) { + writeField(fieldWriter.fieldInfo, datasetOrVectors); + } else { + writeSortingField(fieldWriter.fieldInfo, datasetOrVectors, sortMap); + } + } finally { + datasetOrVectors.close(); + } } } } @@ -221,38 +254,13 @@ public void close() { } } - private void writeField(FieldWriter fieldWriter) throws IOException { - var vectors = fieldWriter.flatFieldVectorsWriter.getVectors(); - final DatasetOrVectors datasetOrVectors; - if (vectors.size() < MIN_NUM_VECTORS_FOR_GPU_BUILD) { - // Use vectors/CPU - datasetOrVectors = DatasetOrVectors.fromArray(vectors); - } else { - // Avoid another heap copy (the float[][]) - - // TODO: another alternative is to use CuVSMatrix.deviceBuilder(), but this requires more effort - // 1. support no-copy CuVSDeviceMatrix as input in CagraIndex - // 2. ensure we are already holding a CuVSResource here - var builder = CuVSMatrix.hostBuilder(vectors.size(), vectors.getFirst().length, dataType); - for (var vector : vectors) { - builder.addVector(vector); - } - datasetOrVectors = DatasetOrVectors.fromDataset(builder.build()); - } - try { - writeFieldInternal(fieldWriter.fieldInfo, datasetOrVectors); - } finally { - datasetOrVectors.close(); - } - } - - private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { + private void writeSortingField(FieldInfo fieldInfo, DatasetOrVectors datasetOrVectors, Sorter.DocMap sortMap) throws IOException { // The flatFieldVectorsWriter's flush method, called before this, has already sorted the vectors according to the sortMap. // We can now treat them as a simple, sorted list of vectors. - writeField(fieldData); + writeField(fieldInfo, datasetOrVectors); } - private void writeFieldInternal(FieldInfo fieldInfo, DatasetOrVectors datasetOrVectors) throws IOException { + private void writeField(FieldInfo fieldInfo, DatasetOrVectors datasetOrVectors) throws IOException { try { long vectorIndexOffset = vectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = new int[1][]; diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/reflect/VectorsFormatReflectionUtils.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/reflect/VectorsFormatReflectionUtils.java new file mode 100644 index 0000000000000..67943f5e7ab20 --- /dev/null +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/reflect/VectorsFormatReflectionUtils.java @@ -0,0 +1,58 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.gpu.reflect; + +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter; +import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsWriter; +import org.apache.lucene.store.IndexOutput; +import org.elasticsearch.index.codec.vectors.ES814ScalarQuantizedVectorsFormat; + +import java.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; + +public class VectorsFormatReflectionUtils { + + private static final VarHandle FLAT_VECTOR_DATA_HANDLE; + private static final VarHandle QUANTIZED_VECTOR_DATA_HANDLE; + private static final VarHandle DELEGATE_WRITER_HANDLE; + + static final Class L99_SQ_VW_CLS = Lucene99ScalarQuantizedVectorsWriter.class; + static final Class L99_F_VW_CLS = Lucene99FlatVectorsWriter.class; + static final Class ES814_SQ_VW_CLS = ES814ScalarQuantizedVectorsFormat.ES814ScalarQuantizedVectorsWriter.class; + + static { + try { + var lookup = MethodHandles.privateLookupIn(L99_F_VW_CLS, MethodHandles.lookup()); + FLAT_VECTOR_DATA_HANDLE = lookup.findVarHandle(L99_F_VW_CLS, "vectorData", IndexOutput.class); + + lookup = MethodHandles.privateLookupIn(L99_SQ_VW_CLS, MethodHandles.lookup()); + QUANTIZED_VECTOR_DATA_HANDLE = lookup.findVarHandle(L99_SQ_VW_CLS, "quantizedVectorData", IndexOutput.class); + + lookup = MethodHandles.privateLookupIn(ES814_SQ_VW_CLS, MethodHandles.lookup()); + DELEGATE_WRITER_HANDLE = lookup.findVarHandle(ES814_SQ_VW_CLS, "delegate", L99_SQ_VW_CLS); + + } catch (IllegalAccessException e) { + throw new AssertionError("should not happen, check opens", e); + } + catch (ReflectiveOperationException e) { + throw new AssertionError(e); + } + } + + public static IndexOutput getVectorDataIndexOutput(FlatVectorsWriter flatVectorWriter) { + assert flatVectorWriter instanceof ES814ScalarQuantizedVectorsFormat.ES814ScalarQuantizedVectorsWriter; + var delegate = (Lucene99ScalarQuantizedVectorsWriter)DELEGATE_WRITER_HANDLE.get(flatVectorWriter); + return (IndexOutput) QUANTIZED_VECTOR_DATA_HANDLE.get(delegate); + } + + public static IndexOutput getQuantizedVectorDataIndexOutput(FlatVectorsWriter flatVectorWriter) { + assert flatVectorWriter instanceof Lucene99FlatVectorsWriter; + return (IndexOutput) FLAT_VECTOR_DATA_HANDLE.get(flatVectorWriter); + } +} From c5ca52cfd008f5757ca93305cd4517a9052b2afd Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Wed, 3 Sep 2025 13:01:54 +0000 Subject: [PATCH 4/8] [CI] Auto commit changes from spotless --- .../xpack/gpu/reflect/VectorsFormatReflectionUtils.java | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/reflect/VectorsFormatReflectionUtils.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/reflect/VectorsFormatReflectionUtils.java index 67943f5e7ab20..7456a24b340c0 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/reflect/VectorsFormatReflectionUtils.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/reflect/VectorsFormatReflectionUtils.java @@ -39,15 +39,14 @@ public class VectorsFormatReflectionUtils { } catch (IllegalAccessException e) { throw new AssertionError("should not happen, check opens", e); - } - catch (ReflectiveOperationException e) { + } catch (ReflectiveOperationException e) { throw new AssertionError(e); } } public static IndexOutput getVectorDataIndexOutput(FlatVectorsWriter flatVectorWriter) { assert flatVectorWriter instanceof ES814ScalarQuantizedVectorsFormat.ES814ScalarQuantizedVectorsWriter; - var delegate = (Lucene99ScalarQuantizedVectorsWriter)DELEGATE_WRITER_HANDLE.get(flatVectorWriter); + var delegate = (Lucene99ScalarQuantizedVectorsWriter) DELEGATE_WRITER_HANDLE.get(flatVectorWriter); return (IndexOutput) QUANTIZED_VECTOR_DATA_HANDLE.get(delegate); } From 0e30555610b67f69ac7a4da9630d75969f146f62 Mon Sep 17 00:00:00 2001 From: ldematte Date: Thu, 4 Sep 2025 11:27:41 +0200 Subject: [PATCH 5/8] Mimic a hypothetical/missing FlatVectorsWriter#getReader(); separates vector count/mmap availability conditions --- .../gpu/codec/ESGpuHnswSQVectorsFormat.java | 3 +- .../gpu/codec/ESGpuHnswVectorsFormat.java | 5 +- .../gpu/codec/ESGpuHnswVectorsWriter.java | 152 ++++++++++++------ 3 files changed, 110 insertions(+), 50 deletions(-) diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswSQVectorsFormat.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswSQVectorsFormat.java index 400a855db6d6b..a7a3bd262623a 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswSQVectorsFormat.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswSQVectorsFormat.java @@ -61,7 +61,8 @@ public ESGpuHnswSQVectorsFormat(int maxConn, int beamWidth, Float confidenceInte @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return new ESGpuHnswVectorsWriter(cuVSResourceManager, state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state)); + return new ESGpuHnswVectorsWriter(cuVSResourceManager, state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state), + flatVectorsFormat::fieldsReader); } @Override diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsFormat.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsFormat.java index b06b452435c83..4fdd1c9ad9569 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsFormat.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsFormat.java @@ -55,7 +55,7 @@ public ESGpuHnswVectorsFormat() { public ESGpuHnswVectorsFormat(int maxConn, int beamWidth) { this(CuVSResourceManager.pooling(), maxConn, beamWidth); - }; + } public ESGpuHnswVectorsFormat(CuVSResourceManager cuVSResourceManager, int maxConn, int beamWidth) { super(NAME); @@ -66,7 +66,8 @@ public ESGpuHnswVectorsFormat(CuVSResourceManager cuVSResourceManager, int maxCo @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return new ESGpuHnswVectorsWriter(cuVSResourceManager, state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state)); + return new ESGpuHnswVectorsWriter(cuVSResourceManager, state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state), + flatVectorsFormat::fieldsReader); } @Override diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java index 4ce946550e2bc..62526ef437ea4 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java @@ -15,7 +15,9 @@ import org.apache.lucene.codecs.KnnFieldVectorsWriter; import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.lucene95.HasIndexSlice; import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter; import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.DocsWithFieldSet; @@ -24,6 +26,7 @@ import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.index.Sorter; import org.apache.lucene.index.VectorEncoding; @@ -38,12 +41,12 @@ import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; import org.apache.lucene.util.packed.DirectMonotonicWriter; import org.apache.lucene.util.quantization.ScalarQuantizer; +import org.elasticsearch.core.CheckedFunction; import org.elasticsearch.core.IOUtils; import org.elasticsearch.core.SuppressForbidden; import org.elasticsearch.index.codec.vectors.ES814ScalarQuantizedVectorsFormat; import org.elasticsearch.logging.LogManager; import org.elasticsearch.logging.Logger; -import org.elasticsearch.xpack.gpu.reflect.VectorsFormatReflectionUtils; import java.io.IOException; import java.nio.ByteBuffer; @@ -75,12 +78,12 @@ final class ESGpuHnswVectorsWriter extends KnnVectorsWriter { private final CuVSResourceManager cuVSResourceManager; private final SegmentWriteState segmentWriteState; private final IndexOutput meta, vectorIndex; - private final IndexOutput vectorData; private final int M; private final int beamWidth; private final FlatVectorsWriter flatVectorWriter; private final List fields = new ArrayList<>(); + private final CheckedFunction flatVectorsReaderProvider; private boolean finished; private final CuVSMatrix.DataType dataType; @@ -89,19 +92,19 @@ final class ESGpuHnswVectorsWriter extends KnnVectorsWriter { SegmentWriteState state, int M, int beamWidth, - FlatVectorsWriter flatVectorWriter + FlatVectorsWriter flatVectorWriter, + CheckedFunction flatVectorsReaderProvider ) throws IOException { + this.flatVectorsReaderProvider = flatVectorsReaderProvider; assert cuVSResourceManager != null : "CuVSResources must not be null"; this.cuVSResourceManager = cuVSResourceManager; this.M = M; this.beamWidth = beamWidth; this.flatVectorWriter = flatVectorWriter; if (flatVectorWriter instanceof ES814ScalarQuantizedVectorsFormat.ES814ScalarQuantizedVectorsWriter) { - vectorData = VectorsFormatReflectionUtils.getQuantizedVectorDataIndexOutput(flatVectorWriter); dataType = CuVSMatrix.DataType.BYTE; } else { assert flatVectorWriter instanceof Lucene99FlatVectorsWriter; - vectorData = VectorsFormatReflectionUtils.getVectorDataIndexOutput(flatVectorWriter); dataType = CuVSMatrix.DataType.FLOAT; } this.segmentWriteState = state; @@ -151,40 +154,81 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException return newField; } + private static MemorySegmentAccessInput getMemorySegmentAccessInputOrNull( + KnnVectorValues vectorValues) { + + if (vectorValues instanceof HasIndexSlice indexSlice) { + var input = FilterIndexInput.unwrapOnlyTest(indexSlice.getSlice()); + if (input instanceof MemorySegmentAccessInput memorySegmentAccessInput) { + return memorySegmentAccessInput; + } + } + return null; + } + @Override public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { flatVectorWriter.flush(maxDoc, sortMap); - try (IndexInput in = segmentWriteState.segmentInfo.dir.openInput(vectorData.getName(), IOContext.DEFAULT)) { - var input = FilterIndexInput.unwrapOnlyTest(in); - + // TODO: this "mimics" a hypothetical/missing FlatVectorsWriter#getReader() + try (FlatVectorsReader flatVectorsReader = flatVectorsReaderProvider.apply( + new SegmentReadState( + segmentWriteState.segmentInfo.dir, + segmentWriteState.segmentInfo, + segmentWriteState.fieldInfos, + segmentWriteState.context + ) + )) { for (FieldWriter fieldWriter : fields) { - // TODO: is this inefficient? Can we get "size" in another way? + // This might be inefficient if getVectors() materializes a List; however current implementations + // just return a reference to an inner, already allocated List, so we are fine for now. + // TODO: change when/if Lucene introduces a direct FlatFieldVectorsWriter#size() var numVectors = fieldWriter.flatFieldVectorsWriter.getVectors().size(); - final DatasetOrVectors datasetOrVectors; - if (input instanceof MemorySegmentAccessInput memorySegmentAccessInput && numVectors >= MIN_NUM_VECTORS_FOR_GPU_BUILD) { - // TODO: we are iterating over multiple fields, we probably need to memorySegmentAccessInput.segmentSliceOrNull()? - var ds = DatasetUtils.getInstance() - .fromInput(memorySegmentAccessInput, numVectors, fieldWriter.fieldInfo.getVectorDimension(), dataType); - datasetOrVectors = DatasetOrVectors.fromDataset(ds); - } else { - var builder = CuVSMatrix.hostBuilder(numVectors, fieldWriter.fieldInfo.getVectorDimension(), dataType); - for (var vector : fieldWriter.flatFieldVectorsWriter.getVectors()) { - builder.addVector(vector); + final CuVSMatrix dataset; + if (numVectors >= MIN_NUM_VECTORS_FOR_GPU_BUILD) { + if (dataType == CuVSMatrix.DataType.FLOAT) { + FloatVectorValues floatVectorValues = flatVectorsReader.getFloatVectorValues(fieldWriter.fieldInfo.name); + var memorySegmentAccessInput = getMemorySegmentAccessInputOrNull(floatVectorValues); + if (memorySegmentAccessInput != null) { + dataset = DatasetUtils.getInstance() + .fromInput(memorySegmentAccessInput, numVectors, fieldWriter.fieldInfo.getVectorDimension(), dataType); + } else { + var builder = CuVSMatrix.hostBuilder(numVectors, fieldWriter.fieldInfo.getVectorDimension(), dataType); + for (int i = 0; i < numVectors; ++i) { + builder.addVector(floatVectorValues.vectorValue(i)); + } + dataset = builder.build(); + } + } else { + assert dataType == CuVSMatrix.DataType.BYTE; + ByteVectorValues byteVectorValues = flatVectorsReader.getByteVectorValues(fieldWriter.fieldInfo.name); + var memorySegmentAccessInput = getMemorySegmentAccessInputOrNull(byteVectorValues); + if (memorySegmentAccessInput != null) { + dataset = DatasetUtils.getInstance() + .fromInput(memorySegmentAccessInput, numVectors, fieldWriter.fieldInfo.getVectorDimension(), dataType); + } else { + var builder = CuVSMatrix.hostBuilder(numVectors, fieldWriter.fieldInfo.getVectorDimension(), dataType); + for (int i = 0; i < numVectors; ++i) { + builder.addVector(byteVectorValues.vectorValue(i)); + } + dataset = builder.build(); + } } - - datasetOrVectors = DatasetOrVectors.fromDataset(builder.build()); + } else { + dataset = null; } try { if (sortMap == null) { - writeField(fieldWriter.fieldInfo, datasetOrVectors); + writeField(fieldWriter.fieldInfo, dataset, numVectors); } else { - writeSortingField(fieldWriter.fieldInfo, datasetOrVectors, sortMap); + writeSortingField(fieldWriter.fieldInfo, dataset, numVectors, sortMap); } } finally { - datasetOrVectors.close(); + if (dataset != null) { + dataset.close(); + } } } } @@ -254,37 +298,35 @@ public void close() { } } - private void writeSortingField(FieldInfo fieldInfo, DatasetOrVectors datasetOrVectors, Sorter.DocMap sortMap) throws IOException { + private void writeSortingField(FieldInfo fieldInfo, CuVSMatrix datasetOrVectors, int size, Sorter.DocMap sortMap) throws IOException { // The flatFieldVectorsWriter's flush method, called before this, has already sorted the vectors according to the sortMap. // We can now treat them as a simple, sorted list of vectors. - writeField(fieldInfo, datasetOrVectors); + writeField(fieldInfo, datasetOrVectors, size); } - private void writeField(FieldInfo fieldInfo, DatasetOrVectors datasetOrVectors) throws IOException { + private void writeField(FieldInfo fieldInfo, CuVSMatrix dataset, int size) throws IOException { try { long vectorIndexOffset = vectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = new int[1][]; - HnswGraph mockGraph; - if (datasetOrVectors.vectors != null) { - int size = datasetOrVectors.size(); + final HnswGraph graph; + if (dataset == null) { if (logger.isDebugEnabled()) { logger.debug("Skip building carga index; vectors length {} < {} (min for GPU)", size, MIN_NUM_VECTORS_FOR_GPU_BUILD); } - mockGraph = writeGraph(size, graphLevelNodeOffsets); + graph = writeMockGraph(size, graphLevelNodeOffsets); } else { - var dataset = datasetOrVectors.dataset; var cuVSResources = cuVSResourceManager.acquire((int) dataset.size(), (int) dataset.columns(), dataset.dataType()); try { try (var index = buildGPUIndex(cuVSResources, fieldInfo.getVectorSimilarityFunction(), dataset)) { assert index != null : "GPU index should be built for field: " + fieldInfo.name; - mockGraph = writeGraph(index.getGraph(), graphLevelNodeOffsets); + graph = writeGraph(index.getGraph(), graphLevelNodeOffsets); } } finally { cuVSResourceManager.release(cuVSResources); } } long vectorIndexLength = vectorIndex.getFilePointer() - vectorIndexOffset; - writeMeta(fieldInfo, vectorIndexOffset, vectorIndexLength, datasetOrVectors.size(), mockGraph, graphLevelNodeOffsets); + writeMeta(fieldInfo, vectorIndexOffset, vectorIndexLength, size, graph, graphLevelNodeOffsets); } catch (IOException e) { throw e; } catch (Throwable t) { @@ -360,7 +402,7 @@ private HnswGraph writeGraph(CuVSMatrix cagraGraph, int[][] levelNodeOffsets) th } // create a mock graph where every node is connected to every other node - private HnswGraph writeGraph(int elementCount, int[][] levelNodeOffsets) throws IOException { + private HnswGraph writeMockGraph(int elementCount, int[][] levelNodeOffsets) throws IOException { if (elementCount == 0) { return null; } @@ -458,26 +500,42 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE } } try (IndexInput in = mergeState.segmentInfo.dir.openInput(tempRawVectorsFileName, IOContext.DEFAULT)) { - var input = FilterIndexInput.unwrapOnlyTest(in); - final DatasetOrVectors datasetOrVectors; - if (input instanceof MemorySegmentAccessInput memorySegmentAccessInput && numVectors >= MIN_NUM_VECTORS_FOR_GPU_BUILD) { - var ds = DatasetUtils.getInstance() - .fromInput(memorySegmentAccessInput, numVectors, fieldInfo.getVectorDimension(), dataType); - datasetOrVectors = DatasetOrVectors.fromDataset(ds); + final CuVSMatrix dataset; + if (numVectors >= MIN_NUM_VECTORS_FOR_GPU_BUILD) { + if (input instanceof MemorySegmentAccessInput memorySegmentAccessInput) { + // Direct access to mmapped file + dataset = DatasetUtils.getInstance() + .fromInput(memorySegmentAccessInput, numVectors, fieldInfo.getVectorDimension(), dataType); + } else { + var builder = CuVSMatrix.hostBuilder(numVectors, fieldInfo.getVectorDimension(), dataType); + // Read vector-by-vector + if (dataType == CuVSMatrix.DataType.FLOAT) { + float[] vector = new float[fieldInfo.getVectorDimension()]; + for (int i = 0; i < numVectors; ++i) { + input.readFloats(vector, fieldInfo.getVectorDimension() * i, fieldInfo.getVectorDimension()); + } + } else { + assert dataType == CuVSMatrix.DataType.BYTE; + byte[] vector = new byte[fieldInfo.getVectorDimension()]; + for (int i = 0; i < numVectors; ++i) { + input.readBytes(vector, fieldInfo.getVectorDimension() * i, fieldInfo.getVectorDimension()); + } + } + dataset = builder.build(); + } } else { - // assert numVectors < MIN_NUM_VECTORS_FOR_GPU_BUILD : "numVectors: " + numVectors; // we don't really need real value for vectors here, // we just build a mock graph where every node is connected to every other node - datasetOrVectors = DatasetOrVectors.fromDataset( - CuVSMatrix.hostBuilder(numVectors, fieldInfo.getVectorDimension(), dataType).build() - ); + dataset = null; } try { - writeFieldInternal(fieldInfo, datasetOrVectors); + writeField(fieldInfo, dataset, numVectors); } finally { - datasetOrVectors.close(); + if (dataset != null) { + dataset.close(); + } } } finally { org.apache.lucene.util.IOUtils.deleteFilesIgnoringExceptions(mergeState.segmentInfo.dir, tempRawVectorsFileName); From 78d52dae3e2c607d3d2bd805fa287216ca02d209 Mon Sep 17 00:00:00 2001 From: ldematte Date: Thu, 4 Sep 2025 11:57:44 +0200 Subject: [PATCH 6/8] Mimic a hypothetical/missing FlatVectorsWriter#getReader(); separates vector count/mmap availability conditions --- .../gpu/codec/ESGpuHnswSQVectorsFormat.java | 10 ++++++++-- .../gpu/codec/ESGpuHnswVectorsFormat.java | 10 ++++++++-- .../gpu/codec/ESGpuHnswVectorsWriter.java | 19 ++++++++++--------- 3 files changed, 26 insertions(+), 13 deletions(-) diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswSQVectorsFormat.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswSQVectorsFormat.java index a7a3bd262623a..a6e20e081a1dc 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswSQVectorsFormat.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswSQVectorsFormat.java @@ -61,8 +61,14 @@ public ESGpuHnswSQVectorsFormat(int maxConn, int beamWidth, Float confidenceInte @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return new ESGpuHnswVectorsWriter(cuVSResourceManager, state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state), - flatVectorsFormat::fieldsReader); + return new ESGpuHnswVectorsWriter( + cuVSResourceManager, + state, + maxConn, + beamWidth, + flatVectorsFormat.fieldsWriter(state), + flatVectorsFormat::fieldsReader + ); } @Override diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsFormat.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsFormat.java index 4fdd1c9ad9569..c5d58c4a4f107 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsFormat.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsFormat.java @@ -66,8 +66,14 @@ public ESGpuHnswVectorsFormat(CuVSResourceManager cuVSResourceManager, int maxCo @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - return new ESGpuHnswVectorsWriter(cuVSResourceManager, state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state), - flatVectorsFormat::fieldsReader); + return new ESGpuHnswVectorsWriter( + cuVSResourceManager, + state, + maxConn, + beamWidth, + flatVectorsFormat.fieldsWriter(state), + flatVectorsFormat::fieldsReader + ); } @Override diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java index 62526ef437ea4..615e07a11414e 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java @@ -154,8 +154,7 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException return newField; } - private static MemorySegmentAccessInput getMemorySegmentAccessInputOrNull( - KnnVectorValues vectorValues) { + private static MemorySegmentAccessInput getMemorySegmentAccessInputOrNull(KnnVectorValues vectorValues) { if (vectorValues instanceof HasIndexSlice indexSlice) { var input = FilterIndexInput.unwrapOnlyTest(indexSlice.getSlice()); @@ -171,14 +170,16 @@ public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { flatVectorWriter.flush(maxDoc, sortMap); // TODO: this "mimics" a hypothetical/missing FlatVectorsWriter#getReader() - try (FlatVectorsReader flatVectorsReader = flatVectorsReaderProvider.apply( - new SegmentReadState( - segmentWriteState.segmentInfo.dir, - segmentWriteState.segmentInfo, - segmentWriteState.fieldInfos, - segmentWriteState.context + try ( + FlatVectorsReader flatVectorsReader = flatVectorsReaderProvider.apply( + new SegmentReadState( + segmentWriteState.segmentInfo.dir, + segmentWriteState.segmentInfo, + segmentWriteState.fieldInfos, + segmentWriteState.context + ) ) - )) { + ) { for (FieldWriter fieldWriter : fields) { // This might be inefficient if getVectors() materializes a List; however current implementations // just return a reference to an inner, already allocated List, so we are fine for now. From aab08c7d90b73f4c4a6ca054406a8e8bf50cf51e Mon Sep 17 00:00:00 2001 From: ldematte Date: Thu, 4 Sep 2025 12:09:26 +0200 Subject: [PATCH 7/8] Remove unused class --- .../gpu/codec/ESGpuHnswVectorsWriter.java | 36 ------------------- 1 file changed, 36 deletions(-) diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java index 615e07a11414e..104205a76acdc 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java @@ -263,42 +263,6 @@ public long ramBytesUsed() { return total; } - private static final class DatasetOrVectors implements AutoCloseable { - private final CuVSMatrix dataset; - private final List vectors; - - static DatasetOrVectors fromArray(List vectors) { - return new DatasetOrVectors(null, vectors); - } - - static DatasetOrVectors fromDataset(CuVSMatrix dataset) { - return new DatasetOrVectors(dataset, null); - } - - private DatasetOrVectors(CuVSMatrix dataset, List vectors) { - this.dataset = dataset; - this.vectors = vectors; - validateState(); - } - - private void validateState() { - if ((dataset == null && vectors == null) || (dataset != null && vectors != null)) { - throw new IllegalStateException("Exactly one of dataset or vectors must be non-null"); - } - } - - int size() { - return dataset != null ? (int) dataset.size() : vectors.size(); - } - - @Override - public void close() { - if (dataset != null) { - dataset.close(); - } - } - } - private void writeSortingField(FieldInfo fieldInfo, CuVSMatrix datasetOrVectors, int size, Sorter.DocMap sortMap) throws IOException { // The flatFieldVectorsWriter's flush method, called before this, has already sorted the vectors according to the sortMap. // We can now treat them as a simple, sorted list of vectors. From dea2da9f20be6e570845b3b643c97d045ad5043d Mon Sep 17 00:00:00 2001 From: ldematte Date: Thu, 4 Sep 2025 17:02:30 +0200 Subject: [PATCH 8/8] Fixes (but still not working -- need to flush inner writer meta to disk) --- .../xpack/gpu/codec/ESGpuHnswVectorsWriter.java | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java index 104205a76acdc..ab33739e7d1ca 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsWriter.java @@ -22,6 +22,7 @@ import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.index.KnnVectorValues; @@ -175,8 +176,9 @@ public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { new SegmentReadState( segmentWriteState.segmentInfo.dir, segmentWriteState.segmentInfo, - segmentWriteState.fieldInfos, - segmentWriteState.context + new FieldInfos(fields.stream().map(x -> x.fieldInfo).toArray(FieldInfo[]::new)), + IOContext.DEFAULT, + segmentWriteState.segmentSuffix ) ) ) { @@ -479,13 +481,13 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE if (dataType == CuVSMatrix.DataType.FLOAT) { float[] vector = new float[fieldInfo.getVectorDimension()]; for (int i = 0; i < numVectors; ++i) { - input.readFloats(vector, fieldInfo.getVectorDimension() * i, fieldInfo.getVectorDimension()); + input.readFloats(vector, 0, fieldInfo.getVectorDimension()); } } else { assert dataType == CuVSMatrix.DataType.BYTE; byte[] vector = new byte[fieldInfo.getVectorDimension()]; for (int i = 0; i < numVectors; ++i) { - input.readBytes(vector, fieldInfo.getVectorDimension() * i, fieldInfo.getVectorDimension()); + input.readBytes(vector, 0, fieldInfo.getVectorDimension()); } } dataset = builder.build();