xiangfu0 commented on code in PR #20071: URL: https://github.com/apache/pinot/pull/20071#discussion_r4412731126
########## pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/text/PinotBufferLuceneDirectoryTest.java: ########## @@ -0,0 +1,109 @@ +/** + * 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.pinot.segment.local.segment.index.readers.text; + +import java.nio.ByteOrder; +import java.util.HashMap; +import java.util.Map; +import org.apache.lucene.analysis.standard.StandardAnalyzer; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.TextField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; + + +/// Reads a Lucene index packed into a [PinotDataBuffer], as the combined text and vector index files are, through +/// [PinotBufferLuceneDirectory] in either buffer byte order. +public class PinotBufferLuceneDirectoryTest { + private static final int NUM_DOCS = 500; + private static final int PADDING = 13; + + @DataProvider + public Object[][] byteOrders() { + return new Object[][]{{ByteOrder.BIG_ENDIAN}, {ByteOrder.LITTLE_ENDIAN}}; + } + + @Test(dataProvider = "byteOrders") + public void testSearchesIndexPackedInBuffer(ByteOrder byteOrder) + throws Exception { + try (ByteBuffersDirectory source = new ByteBuffersDirectory()) { + try (IndexWriter writer = new IndexWriter(source, new IndexWriterConfig(new StandardAnalyzer()))) { + for (int i = 0; i < NUM_DOCS; i++) { + Document document = new Document(); + document.add(new TextField("text", (i % 2 == 0 ? "even" : "odd") + " doc" + i, Field.Store.NO)); + document.add(new KnnFloatVectorField("vector", new float[]{i, 1, -i}, VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(document); + } + writer.forceMerge(1); + } + + // Pack every file at an odd offset, so views do not start on an aligned boundary + Map<String, LuceneTextIndexHeader.FileInfo> fileMap = new HashMap<>(); + long size = PADDING; + for (String name : source.listAll()) { + long length = source.fileLength(name); + fileMap.put(name, new LuceneTextIndexHeader.FileInfo(name, size, length)); + size += length + PADDING; + } + try (PinotDataBuffer buffer = PinotDataBuffer.allocateDirect(size, byteOrder, "test")) { Review Comment: `allocateDirect` goes through the default `SmallWithFallback` factory, so this test only covers `PinotByteBuffer`. Production column files over 2 GB (or `pinot.offheap.prioritize.bytebuffer=false`) take the `UnsafePinotBuffer` path, where `toDirectByteBuffer` builds the buffer through `ByteBufferUtil` reflection. Could you also parameterize over `PinotDataBuffer.useFactory(new UnsafePinotBufferFactory())`? I ran it locally and it passes, but a regression test would keep it covered. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/text/PinotBufferLuceneDirectory.java: ########## @@ -109,7 +114,14 @@ public IndexInput openInput(String name, IOContext context) if (fileInfo == null) { throw new IOException("File not found: " + name); } - return new PinotBufferIndexInput(name, _indexBuffer, fileInfo.getOffset(), fileInfo.getSize()); + long size = fileInfo.getSize(); + if (size <= Integer.MAX_VALUE) { + // A direct view of the buffer: Lucene reads vectors and postings with bulk, zero-copy accessors instead of + // refilling a small heap buffer for every read + ByteBuffer view = _indexBuffer.toDirectByteBuffer(fileInfo.getOffset(), (int) size, ByteOrder.LITTLE_ENDIAN); + return new ByteBuffersIndexInput(new ByteBuffersDataInput(List.of(view)), name); + } + return new PinotBufferIndexInput(name, _indexBuffer, fileInfo.getOffset(), size); Review Comment: Optional follow-up: `ByteBuffersDataInput` also accepts multiple equal power-of-two blocks. Files over 2 GB could therefore be split into 1 GB `toDirectByteBuffer` views instead of falling back here. That would make large files zero-copy too, and let `PinotBufferIndexInput` (which allocates a `byte[]` on every refill) be removed. As written, this fallback has no test. -- 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]
