xiangfu0 commented on code in PR #20072:
URL: https://github.com/apache/pinot/pull/20072#discussion_r4410202424
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/vector/HnswVectorIndexCreator.java:
##########
@@ -63,6 +81,12 @@ public HnswVectorIndexCreator(String column, File
segmentIndexDir, VectorIndexCo
_vectorSimilarityFunction =
VectorIndexUtils.toSimilarityFunction(vectorIndexConfig.getVectorDistanceFunction());
_storeInSegmentFile = vectorIndexConfig.isStoreInSegmentFile();
_segmentIndexDir = segmentIndexDir;
+ Map<String, String> properties = vectorIndexConfig.getProperties();
+ _buildThreads = Math.max(1,
Integer.parseInt(properties.getOrDefault("buildThreads", "1")));
Review Comment:
**MAJOR:** `VectorIndexConfigValidator` doesn't know about `buildThreads`,
`forceMerge` or `parallelBuildMaxDocs`. A table config with `buildThreads:
"four"` passes validation, then this line throws `NumberFormatException` at
segment build (commit or reload). The validator's own comment warns this can
stop ingestion.
`buildThreads` also has no upper bound, and every column build creates
`min(buildThreads, numDocs)` threads and IndexWriters. I verified
`buildThreads=2000` on 5k docs fails `seal()` with *Too many open files*.
Please add these keys to `HNSW_PROPERTIES` with parse and bounds checks, cap
the thread count (e.g. at available processors), and document them in
`VectorBackendType`.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/vector/HnswVectorIndexCreator.java:
##########
@@ -100,24 +124,94 @@ public void add(Object[] values, @Nullable int[] dictIds)
{
@Override
public void add(float[] document) {
+ if (_bufferedVectors != null) {
+ // Callers may reuse the array for the next row
+ _bufferedVectors.add(document.clone());
Review Comment:
**MINOR:** The buffer is limited by doc count, not bytes. With the default
50k docs, a 1536-dim column holds about 300 MB on heap, plus one
`maxBufferSizeMB` per part writer. Large segments pay this too, until they pass
the limit and fall back. The baseline stays bounded by `maxBufferSizeMB` (16 MB
by default). Consider a byte-based limit, or document the cost, since the PR
says large segments keep the existing memory profile.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/vector/HnswVectorIndexCreator.java:
##########
@@ -100,24 +124,94 @@ public void add(Object[] values, @Nullable int[] dictIds)
{
@Override
public void add(float[] document) {
+ if (_bufferedVectors != null) {
+ // Callers may reuse the array for the next row
+ _bufferedVectors.add(document.clone());
+ if (_bufferedVectors.size() <= _parallelBuildMaxDocs) {
+ return;
+ }
+ for (float[] buffered : _bufferedVectors) {
+ addDocument(buffered, _nextDocId++);
+ }
+ _bufferedVectors = null;
+ return;
+ }
+ addDocument(document, _nextDocId++);
+ }
+
+ private Document toDocument(float[] vector, int docId) {
Document docToIndex = new Document();
- XKnnFloatVectorField xKnnFloatVectorField =
- new XKnnFloatVectorField(_vectorColumn, document,
_vectorSimilarityFunction);
- docToIndex.add(xKnnFloatVectorField);
- docToIndex.add(new StoredField(VECTOR_INDEX_DOC_ID_COLUMN_NAME,
_nextDocId++));
+ docToIndex.add(new XKnnFloatVectorField(_vectorColumn, vector,
_vectorSimilarityFunction));
+ docToIndex.add(new StoredField(VECTOR_INDEX_DOC_ID_COLUMN_NAME, docId));
+ return docToIndex;
+ }
+
+ private void addDocument(float[] vector, int docId) {
try {
- _indexWriter.addDocument(docToIndex);
+ _indexWriter.addDocument(toDocument(vector, docId));
} catch (Exception e) {
throw new RuntimeException(
"Caught exception while adding a new document to the HNSW index for
column: " + _vectorColumn, e);
}
}
+ /// Indexes the buffered vectors from `_buildThreads` threads, each a
contiguous range of doc ids written by its own
+ /// IndexWriter, so each graph is also finished (flushed) on its own thread;
the per-thread indexes are then attached
+ /// to the main index as separate Lucene segments without rebuilding them.
+ private void addBufferedInParallel()
+ throws Exception {
+ int numDocs = _bufferedVectors.size();
+ int numThreads = Math.min(_buildThreads, Math.max(1, numDocs));
+ ExecutorService executor = Executors.newFixedThreadPool(numThreads);
+ Directory[] partDirectories = new Directory[numThreads];
+ File[] partDirs = new File[numThreads];
+ try {
+ List<Future<?>> futures = new ArrayList<>(numThreads);
+ for (int t = 0; t < numThreads; t++) {
+ int from = (int) ((long) numDocs * t / numThreads);
+ int to = (int) ((long) numDocs * (t + 1) / numThreads);
+ partDirs[t] = new File(_segmentIndexDir, _vectorColumn + ".hnsw.part"
+ t);
+ partDirectories[t] = FSDirectory.open(partDirs[t].toPath());
+ Directory partDirectory = partDirectories[t];
+ futures.add(executor.submit(() -> {
+ try (IndexWriter partWriter = new IndexWriter(partDirectory,
+ VectorIndexUtils.getIndexWriterConfig(_vectorIndexConfig))) {
+ for (int docId = from; docId < to; docId++) {
+ partWriter.addDocument(toDocument(_bufferedVectors.get(docId),
docId));
+ }
+ partWriter.commit();
+ }
+ return null;
+ }));
+ }
+ for (Future<?> future : futures) {
+ future.get();
+ }
+ _indexWriter.addIndexes(partDirectories);
+ } finally {
+ executor.shutdownNow();
+ IOUtils.close(partDirectories);
+ for (File partDir : partDirs) {
+ if (partDir != null) {
+ FileUtils.deleteQuietly(partDir);
+ }
+ }
+ }
+ _nextDocId = numDocs;
+ _bufferedVectors.clear();
+ }
+
@Override
public void seal() {
try {
- LOGGER.info("Sealing HNSW index for column: {}", _vectorColumn);
- _indexWriter.forceMerge(1);
+ LOGGER.info("Sealing HNSW index for column: {} (buildThreads: {},
forceMerge: {})", _vectorColumn,
+ _buildThreads, _forceMerge);
+ if (_bufferedVectors != null) {
+ addBufferedInParallel();
+ }
+ if (_forceMerge) {
+ _indexWriter.forceMerge(1);
Review Comment:
**MINOR:** With the default `forceMerge=true`, Lucene 9.12's merge reuses
one part's graph and re-inserts the other ~(T−1)/T vectors on a single thread.
That's why the speedup was modest in my runs and 20k×128 was slower than base
(see summary). It may be worth trying Lucene's concurrent HNSW merge
(`numMergeWorkers` + `mergeExec` on `Lucene99HnswVectorsFormat`). Please also
document the query-latency cost of `forceMerge=false`. Its segment count also
depends on `maxBufferSizeMB`: I saw 8 segments, not 4, at 20k×768.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/vector/HnswVectorIndexCreator.java:
##########
@@ -100,24 +124,94 @@ public void add(Object[] values, @Nullable int[] dictIds)
{
@Override
public void add(float[] document) {
+ if (_bufferedVectors != null) {
+ // Callers may reuse the array for the next row
+ _bufferedVectors.add(document.clone());
+ if (_bufferedVectors.size() <= _parallelBuildMaxDocs) {
+ return;
+ }
+ for (float[] buffered : _bufferedVectors) {
+ addDocument(buffered, _nextDocId++);
+ }
+ _bufferedVectors = null;
+ return;
+ }
+ addDocument(document, _nextDocId++);
+ }
+
+ private Document toDocument(float[] vector, int docId) {
Document docToIndex = new Document();
- XKnnFloatVectorField xKnnFloatVectorField =
- new XKnnFloatVectorField(_vectorColumn, document,
_vectorSimilarityFunction);
- docToIndex.add(xKnnFloatVectorField);
- docToIndex.add(new StoredField(VECTOR_INDEX_DOC_ID_COLUMN_NAME,
_nextDocId++));
+ docToIndex.add(new XKnnFloatVectorField(_vectorColumn, vector,
_vectorSimilarityFunction));
+ docToIndex.add(new StoredField(VECTOR_INDEX_DOC_ID_COLUMN_NAME, docId));
+ return docToIndex;
+ }
+
+ private void addDocument(float[] vector, int docId) {
try {
- _indexWriter.addDocument(docToIndex);
+ _indexWriter.addDocument(toDocument(vector, docId));
} catch (Exception e) {
throw new RuntimeException(
"Caught exception while adding a new document to the HNSW index for
column: " + _vectorColumn, e);
}
}
+ /// Indexes the buffered vectors from `_buildThreads` threads, each a
contiguous range of doc ids written by its own
+ /// IndexWriter, so each graph is also finished (flushed) on its own thread;
the per-thread indexes are then attached
+ /// to the main index as separate Lucene segments without rebuilding them.
+ private void addBufferedInParallel()
+ throws Exception {
+ int numDocs = _bufferedVectors.size();
+ int numThreads = Math.min(_buildThreads, Math.max(1, numDocs));
+ ExecutorService executor = Executors.newFixedThreadPool(numThreads);
+ Directory[] partDirectories = new Directory[numThreads];
+ File[] partDirs = new File[numThreads];
+ try {
+ List<Future<?>> futures = new ArrayList<>(numThreads);
+ for (int t = 0; t < numThreads; t++) {
+ int from = (int) ((long) numDocs * t / numThreads);
+ int to = (int) ((long) numDocs * (t + 1) / numThreads);
+ partDirs[t] = new File(_segmentIndexDir, _vectorColumn + ".hnsw.part"
+ t);
+ partDirectories[t] = FSDirectory.open(partDirs[t].toPath());
+ Directory partDirectory = partDirectories[t];
+ futures.add(executor.submit(() -> {
+ try (IndexWriter partWriter = new IndexWriter(partDirectory,
+ VectorIndexUtils.getIndexWriterConfig(_vectorIndexConfig))) {
+ for (int docId = from; docId < to; docId++) {
+ partWriter.addDocument(toDocument(_bufferedVectors.get(docId),
docId));
+ }
+ partWriter.commit();
+ }
+ return null;
+ }));
+ }
+ for (Future<?> future : futures) {
+ future.get();
+ }
+ _indexWriter.addIndexes(partDirectories);
+ } finally {
+ executor.shutdownNow();
Review Comment:
**MINOR:** If one part fails, `shutdownNow()` returns without waiting. The
directories are then closed and deleted while the other part writers may still
be writing, so files can be left behind. Calling `awaitTermination` before
`IOUtils.close`/`deleteQuietly` would close that race.
--
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]
