From 5e8db2e068f2549b9619d5ac48a50c8032fc292b Mon Sep 17 00:00:00 2001 From: Mike McCandless Date: Sun, 4 Dec 2016 05:18:04 -0500 Subject: [PATCH] LUCENE-7563: use a compressed format for the in-heap BKD index --- lucene/CHANGES.txt | 4 + .../simpletext/SimpleTextBKDReader.java | 281 ++- .../simpletext/SimpleTextBKDWriter.java | 1661 +++++++++++++++++ .../simpletext/SimpleTextPointsReader.java | 5 +- .../simpletext/SimpleTextPointsWriter.java | 188 +- .../codecs/lucene60/Lucene60PointsFormat.java | 10 +- .../lucene/codecs/lucene60/package-info.java | 4 +- .../lucene/codecs/lucene62/package-info.java | 4 +- .../lucene/codecs/lucene70/package-info.java | 15 +- .../org/apache/lucene/index/CheckIndex.java | 314 ++-- .../org/apache/lucene/util/bkd/BKDReader.java | 717 ++++--- .../org/apache/lucene/util/bkd/BKDWriter.java | 293 ++- .../lucene/util/bkd/HeapPointReader.java | 7 +- .../lucene/util/bkd/HeapPointWriter.java | 22 +- .../util/bkd/MutablePointsReaderUtils.java | 21 +- .../lucene/util/bkd/OfflinePointReader.java | 8 +- .../lucene/util/bkd/OfflinePointWriter.java | 10 +- .../apache/lucene/util/bkd/PointReader.java | 14 +- .../apache/lucene/util/bkd/PointWriter.java | 6 +- .../lucene/search/TestPointQueries.java | 3 + .../lucene/util/bkd/Test2BBKDPoints.java | 11 +- .../org/apache/lucene/util/bkd/TestBKD.java | 54 + .../org/apache/lucene/util/fst/TestFSTs.java | 2 +- .../lucene/document/NearestNeighbor.java | 44 +- 24 files changed, 3030 insertions(+), 668 deletions(-) create mode 100644 lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextBKDWriter.java diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index 4afc5078fa2..79e44e112c8 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -126,6 +126,10 @@ Optimizations * LUCENE-7568: Optimize merging when index sorting is used but the index is already sorted (Jim Ferenczi via Mike McCandless) +* LUCENE-7563: The BKD in-memory index for dimensional points now uses + a compressed format, using substantially less RAM in some cases + (Adrien Grand, Mike McCandless) + Other * LUCENE-7546: Fixed references to benchmark wikipedia data and the Jenkins line-docs file diff --git a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextBKDReader.java b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextBKDReader.java index a2b784afd27..488547b4dea 100644 --- a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextBKDReader.java +++ b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextBKDReader.java @@ -16,13 +16,17 @@ */ package org.apache.lucene.codecs.simpletext; - import java.io.IOException; import java.nio.charset.StandardCharsets; +import org.apache.lucene.codecs.simpletext.SimpleTextUtil; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.index.PointValues; import org.apache.lucene.store.IndexInput; +import org.apache.lucene.util.Accountable; import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.BytesRefBuilder; +import org.apache.lucene.util.RamUsageEstimator; import org.apache.lucene.util.StringHelper; import org.apache.lucene.util.bkd.BKDReader; @@ -30,15 +34,105 @@ import static org.apache.lucene.codecs.simpletext.SimpleTextPointsWriter.BLOCK_C import static org.apache.lucene.codecs.simpletext.SimpleTextPointsWriter.BLOCK_DOC_ID; import static org.apache.lucene.codecs.simpletext.SimpleTextPointsWriter.BLOCK_VALUE; -class SimpleTextBKDReader extends BKDReader { +/** Forked from {@link BKDReader} and simplified/specialized for SimpleText's usage */ - public SimpleTextBKDReader(IndexInput datIn, int numDims, int maxPointsInLeafNode, int bytesPerDim, long[] leafBlockFPs, byte[] splitPackedValues, +final class SimpleTextBKDReader extends PointValues implements Accountable { + // Packed array of byte[] holding all split values in the full binary tree: + final private byte[] splitPackedValues; + final long[] leafBlockFPs; + final private int leafNodeOffset; + final int numDims; + final int bytesPerDim; + final int bytesPerIndexEntry; + final IndexInput in; + final int maxPointsInLeafNode; + final byte[] minPackedValue; + final byte[] maxPackedValue; + final long pointCount; + final int docCount; + final int version; + protected final int packedBytesLength; + + public SimpleTextBKDReader(IndexInput in, int numDims, int maxPointsInLeafNode, int bytesPerDim, long[] leafBlockFPs, byte[] splitPackedValues, byte[] minPackedValue, byte[] maxPackedValue, long pointCount, int docCount) throws IOException { - super(datIn, numDims, maxPointsInLeafNode, bytesPerDim, leafBlockFPs, splitPackedValues, minPackedValue, maxPackedValue, pointCount, docCount); + this.in = in; + this.numDims = numDims; + this.maxPointsInLeafNode = maxPointsInLeafNode; + this.bytesPerDim = bytesPerDim; + // no version check here because callers of this API (SimpleText) have no back compat: + bytesPerIndexEntry = numDims == 1 ? bytesPerDim : bytesPerDim + 1; + packedBytesLength = numDims * bytesPerDim; + this.leafNodeOffset = leafBlockFPs.length; + this.leafBlockFPs = leafBlockFPs; + this.splitPackedValues = splitPackedValues; + this.minPackedValue = minPackedValue; + this.maxPackedValue = maxPackedValue; + this.pointCount = pointCount; + this.docCount = docCount; + this.version = SimpleTextBKDWriter.VERSION_CURRENT; + assert minPackedValue.length == packedBytesLength; + assert maxPackedValue.length == packedBytesLength; } - @Override - protected void visitDocIDs(IndexInput in, long blockFP, IntersectVisitor visitor) throws IOException { + /** Used to track all state for a single call to {@link #intersect}. */ + public static final class IntersectState { + final IndexInput in; + final int[] scratchDocIDs; + final byte[] scratchPackedValue; + final int[] commonPrefixLengths; + + final IntersectVisitor visitor; + + public IntersectState(IndexInput in, int numDims, + int packedBytesLength, + int maxPointsInLeafNode, + IntersectVisitor visitor) { + this.in = in; + this.visitor = visitor; + this.commonPrefixLengths = new int[numDims]; + this.scratchDocIDs = new int[maxPointsInLeafNode]; + this.scratchPackedValue = new byte[packedBytesLength]; + } + } + + public void intersect(IntersectVisitor visitor) throws IOException { + intersect(getIntersectState(visitor), 1, minPackedValue, maxPackedValue); + } + + /** Fast path: this is called when the query box fully encompasses all cells under this node. */ + private void addAll(IntersectState state, int nodeID) throws IOException { + //System.out.println("R: addAll nodeID=" + nodeID); + + if (nodeID >= leafNodeOffset) { + //System.out.println("ADDALL"); + visitDocIDs(state.in, leafBlockFPs[nodeID-leafNodeOffset], state.visitor); + // TODO: we can assert that the first value here in fact matches what the index claimed? + } else { + addAll(state, 2*nodeID); + addAll(state, 2*nodeID+1); + } + } + + /** Create a new {@link IntersectState} */ + public IntersectState getIntersectState(IntersectVisitor visitor) { + return new IntersectState(in.clone(), numDims, + packedBytesLength, + maxPointsInLeafNode, + visitor); + } + + /** Visits all docIDs and packed values in a single leaf block */ + public void visitLeafBlockValues(int nodeID, IntersectState state) throws IOException { + int leafID = nodeID - leafNodeOffset; + + // Leaf node; scan and filter all points in this block: + int count = readDocIDs(state.in, leafBlockFPs[leafID], state.scratchDocIDs); + + // Again, this time reading values and checking with the visitor + visitDocValues(state.commonPrefixLengths, state.scratchPackedValue, state.in, state.scratchDocIDs, count, state.visitor); + } + + void visitDocIDs(IndexInput in, long blockFP, IntersectVisitor visitor) throws IOException { BytesRefBuilder scratch = new BytesRefBuilder(); in.seek(blockFP); readLine(in, scratch); @@ -50,8 +144,7 @@ class SimpleTextBKDReader extends BKDReader { } } - @Override - protected int readDocIDs(IndexInput in, long blockFP, int[] docIDs) throws IOException { + int readDocIDs(IndexInput in, long blockFP, int[] docIDs) throws IOException { BytesRefBuilder scratch = new BytesRefBuilder(); in.seek(blockFP); readLine(in, scratch); @@ -63,8 +156,7 @@ class SimpleTextBKDReader extends BKDReader { return count; } - @Override - protected void visitDocValues(int[] commonPrefixLengths, byte[] scratchPackedValue, IndexInput in, int[] docIDs, int count, IntersectVisitor visitor) throws IOException { + void visitDocValues(int[] commonPrefixLengths, byte[] scratchPackedValue, IndexInput in, int[] docIDs, int count, IntersectVisitor visitor) throws IOException { visitor.grow(count); // NOTE: we don't do prefix coding, so we ignore commonPrefixLengths assert scratchPackedValue.length == packedBytesLength; @@ -79,6 +171,175 @@ class SimpleTextBKDReader extends BKDReader { } } + private void visitCompressedDocValues(int[] commonPrefixLengths, byte[] scratchPackedValue, IndexInput in, int[] docIDs, int count, IntersectVisitor visitor, int compressedDim) throws IOException { + // the byte at `compressedByteOffset` is compressed using run-length compression, + // other suffix bytes are stored verbatim + final int compressedByteOffset = compressedDim * bytesPerDim + commonPrefixLengths[compressedDim]; + commonPrefixLengths[compressedDim]++; + int i; + for (i = 0; i < count; ) { + scratchPackedValue[compressedByteOffset] = in.readByte(); + final int runLen = Byte.toUnsignedInt(in.readByte()); + for (int j = 0; j < runLen; ++j) { + for(int dim=0;dim 1.1 MB with 128 points +// per leaf, and you can reduce that by putting more points per leaf +// - we could use threads while building; the higher nodes are very parallelizable + +/** Forked from {@link BKDWriter} and simplified/specialized for SimpleText's usage */ + +final class SimpleTextBKDWriter implements Closeable { + + public static final String CODEC_NAME = "BKD"; + public static final int VERSION_START = 0; + public static final int VERSION_COMPRESSED_DOC_IDS = 1; + public static final int VERSION_COMPRESSED_VALUES = 2; + public static final int VERSION_IMPLICIT_SPLIT_DIM_1D = 3; + public static final int VERSION_CURRENT = VERSION_IMPLICIT_SPLIT_DIM_1D; + + /** How many bytes each docs takes in the fixed-width offline format */ + private final int bytesPerDoc; + + /** Default maximum number of point in each leaf block */ + public static final int DEFAULT_MAX_POINTS_IN_LEAF_NODE = 1024; + + /** Default maximum heap to use, before spilling to (slower) disk */ + public static final float DEFAULT_MAX_MB_SORT_IN_HEAP = 16.0f; + + /** Maximum number of dimensions */ + public static final int MAX_DIMS = 8; + + /** How many dimensions we are indexing */ + protected final int numDims; + + /** How many bytes each value in each dimension takes. */ + protected final int bytesPerDim; + + /** numDims * bytesPerDim */ + protected final int packedBytesLength; + + final BytesRefBuilder scratch = new BytesRefBuilder(); + + final TrackingDirectoryWrapper tempDir; + final String tempFileNamePrefix; + final double maxMBSortInHeap; + + final byte[] scratchDiff; + final byte[] scratch1; + final byte[] scratch2; + final BytesRef scratchBytesRef1 = new BytesRef(); + final BytesRef scratchBytesRef2 = new BytesRef(); + final int[] commonPrefixLengths; + + protected final FixedBitSet docsSeen; + + private OfflinePointWriter offlinePointWriter; + private HeapPointWriter heapPointWriter; + + private IndexOutput tempInput; + protected final int maxPointsInLeafNode; + private final int maxPointsSortInHeap; + + /** Minimum per-dim values, packed */ + protected final byte[] minPackedValue; + + /** Maximum per-dim values, packed */ + protected final byte[] maxPackedValue; + + protected long pointCount; + + /** true if we have so many values that we must write ords using long (8 bytes) instead of int (4 bytes) */ + protected final boolean longOrds; + + /** An upper bound on how many points the caller will add (includes deletions) */ + private final long totalPointCount; + + /** True if every document has at most one value. We specialize this case by not bothering to store the ord since it's redundant with docID. */ + protected final boolean singleValuePerDoc; + + /** How much heap OfflineSorter is allowed to use */ + protected final OfflineSorter.BufferSize offlineSorterBufferMB; + + /** How much heap OfflineSorter is allowed to use */ + protected final int offlineSorterMaxTempFiles; + + private final int maxDoc; + + public SimpleTextBKDWriter(int maxDoc, Directory tempDir, String tempFileNamePrefix, int numDims, int bytesPerDim, + int maxPointsInLeafNode, double maxMBSortInHeap, long totalPointCount, boolean singleValuePerDoc) throws IOException { + this(maxDoc, tempDir, tempFileNamePrefix, numDims, bytesPerDim, maxPointsInLeafNode, maxMBSortInHeap, totalPointCount, singleValuePerDoc, + totalPointCount > Integer.MAX_VALUE, Math.max(1, (long) maxMBSortInHeap), OfflineSorter.MAX_TEMPFILES); + } + + private SimpleTextBKDWriter(int maxDoc, Directory tempDir, String tempFileNamePrefix, int numDims, int bytesPerDim, + int maxPointsInLeafNode, double maxMBSortInHeap, long totalPointCount, + boolean singleValuePerDoc, boolean longOrds, long offlineSorterBufferMB, int offlineSorterMaxTempFiles) throws IOException { + verifyParams(numDims, maxPointsInLeafNode, maxMBSortInHeap, totalPointCount); + // We use tracking dir to deal with removing files on exception, so each place that + // creates temp files doesn't need crazy try/finally/sucess logic: + this.tempDir = new TrackingDirectoryWrapper(tempDir); + this.tempFileNamePrefix = tempFileNamePrefix; + this.maxPointsInLeafNode = maxPointsInLeafNode; + this.numDims = numDims; + this.bytesPerDim = bytesPerDim; + this.totalPointCount = totalPointCount; + this.maxDoc = maxDoc; + this.offlineSorterBufferMB = OfflineSorter.BufferSize.megabytes(offlineSorterBufferMB); + this.offlineSorterMaxTempFiles = offlineSorterMaxTempFiles; + docsSeen = new FixedBitSet(maxDoc); + packedBytesLength = numDims * bytesPerDim; + + scratchDiff = new byte[bytesPerDim]; + scratch1 = new byte[packedBytesLength]; + scratch2 = new byte[packedBytesLength]; + commonPrefixLengths = new int[numDims]; + + minPackedValue = new byte[packedBytesLength]; + maxPackedValue = new byte[packedBytesLength]; + + // If we may have more than 1+Integer.MAX_VALUE values, then we must encode ords with long (8 bytes), else we can use int (4 bytes). + this.longOrds = longOrds; + + this.singleValuePerDoc = singleValuePerDoc; + + // dimensional values (numDims * bytesPerDim) + ord (int or long) + docID (int) + if (singleValuePerDoc) { + // Lucene only supports up to 2.1 docs, so we better not need longOrds in this case: + assert longOrds == false; + bytesPerDoc = packedBytesLength + Integer.BYTES; + } else if (longOrds) { + bytesPerDoc = packedBytesLength + Long.BYTES + Integer.BYTES; + } else { + bytesPerDoc = packedBytesLength + Integer.BYTES + Integer.BYTES; + } + + // As we recurse, we compute temporary partitions of the data, halving the + // number of points at each recursion. Once there are few enough points, + // we can switch to sorting in heap instead of offline (on disk). At any + // time in the recursion, we hold the number of points at that level, plus + // all recursive halves (i.e. 16 + 8 + 4 + 2) so the memory usage is 2X + // what that level would consume, so we multiply by 0.5 to convert from + // bytes to points here. Each dimension has its own sorted partition, so + // we must divide by numDims as wel. + + maxPointsSortInHeap = (int) (0.5 * (maxMBSortInHeap * 1024 * 1024) / (bytesPerDoc * numDims)); + + // Finally, we must be able to hold at least the leaf node in heap during build: + if (maxPointsSortInHeap < maxPointsInLeafNode) { + throw new IllegalArgumentException("maxMBSortInHeap=" + maxMBSortInHeap + " only allows for maxPointsSortInHeap=" + maxPointsSortInHeap + ", but this is less than maxPointsInLeafNode=" + maxPointsInLeafNode + "; either increase maxMBSortInHeap or decrease maxPointsInLeafNode"); + } + + // We write first maxPointsSortInHeap in heap, then cutover to offline for additional points: + heapPointWriter = new HeapPointWriter(16, maxPointsSortInHeap, packedBytesLength, longOrds, singleValuePerDoc); + + this.maxMBSortInHeap = maxMBSortInHeap; + } + + public static void verifyParams(int numDims, int maxPointsInLeafNode, double maxMBSortInHeap, long totalPointCount) { + // We encode dim in a single byte in the splitPackedValues, but we only expose 4 bits for it now, in case we want to use + // remaining 4 bits for another purpose later + if (numDims < 1 || numDims > MAX_DIMS) { + throw new IllegalArgumentException("numDims must be 1 .. " + MAX_DIMS + " (got: " + numDims + ")"); + } + if (maxPointsInLeafNode <= 0) { + throw new IllegalArgumentException("maxPointsInLeafNode must be > 0; got " + maxPointsInLeafNode); + } + if (maxPointsInLeafNode > ArrayUtil.MAX_ARRAY_LENGTH) { + throw new IllegalArgumentException("maxPointsInLeafNode must be <= ArrayUtil.MAX_ARRAY_LENGTH (= " + ArrayUtil.MAX_ARRAY_LENGTH + "); got " + maxPointsInLeafNode); + } + if (maxMBSortInHeap < 0.0) { + throw new IllegalArgumentException("maxMBSortInHeap must be >= 0.0 (got: " + maxMBSortInHeap + ")"); + } + if (totalPointCount < 0) { + throw new IllegalArgumentException("totalPointCount must be >=0 (got: " + totalPointCount + ")"); + } + } + + /** If the current segment has too many points then we spill over to temp files / offline sort. */ + private void spillToOffline() throws IOException { + + // For each .add we just append to this input file, then in .finish we sort this input and resursively build the tree: + offlinePointWriter = new OfflinePointWriter(tempDir, tempFileNamePrefix, packedBytesLength, longOrds, "spill", 0, singleValuePerDoc); + tempInput = offlinePointWriter.out; + PointReader reader = heapPointWriter.getReader(0, pointCount); + for(int i=0;i= maxPointsSortInHeap) { + if (offlinePointWriter == null) { + spillToOffline(); + } + offlinePointWriter.append(packedValue, pointCount, docID); + } else { + // Not too many points added yet, continue using heap: + heapPointWriter.append(packedValue, pointCount, docID); + } + + // TODO: we could specialize for the 1D case: + if (pointCount == 0) { + System.arraycopy(packedValue, 0, minPackedValue, 0, packedBytesLength); + System.arraycopy(packedValue, 0, maxPackedValue, 0, packedBytesLength); + } else { + for(int dim=0;dim 0) { + System.arraycopy(packedValue, offset, maxPackedValue, offset, bytesPerDim); + } + } + } + + pointCount++; + if (pointCount > totalPointCount) { + throw new IllegalStateException("totalPointCount=" + totalPointCount + " was passed when we were created, but we just hit " + pointCount + " values"); + } + docsSeen.set(docID); + } + + /** How many points have been added so far */ + public long getPointCount() { + return pointCount; + } + + private static class MergeReader { + final SimpleTextBKDReader bkd; + final SimpleTextBKDReader.IntersectState state; + final MergeState.DocMap docMap; + + /** Current doc ID */ + public int docID; + + /** Which doc in this block we are up to */ + private int docBlockUpto; + + /** How many docs in the current block */ + private int docsInBlock; + + /** Which leaf block we are up to */ + private int blockID; + + private final byte[] packedValues; + + public MergeReader(SimpleTextBKDReader bkd, MergeState.DocMap docMap) throws IOException { + this.bkd = bkd; + state = new SimpleTextBKDReader.IntersectState(bkd.in.clone(), + bkd.numDims, + bkd.packedBytesLength, + bkd.maxPointsInLeafNode, + null); + this.docMap = docMap; + long minFP = Long.MAX_VALUE; + //System.out.println("MR.init " + this + " bkdreader=" + bkd + " leafBlockFPs.length=" + bkd.leafBlockFPs.length); + for(long fp : bkd.leafBlockFPs) { + minFP = Math.min(minFP, fp); + //System.out.println(" leaf fp=" + fp); + } + state.in.seek(minFP); + this.packedValues = new byte[bkd.maxPointsInLeafNode * bkd.packedBytesLength]; + } + + public boolean next() throws IOException { + //System.out.println("MR.next this=" + this); + while (true) { + if (docBlockUpto == docsInBlock) { + if (blockID == bkd.leafBlockFPs.length) { + //System.out.println(" done!"); + return false; + } + //System.out.println(" new block @ fp=" + state.in.getFilePointer()); + docsInBlock = bkd.readDocIDs(state.in, state.in.getFilePointer(), state.scratchDocIDs); + assert docsInBlock > 0; + docBlockUpto = 0; + bkd.visitDocValues(state.commonPrefixLengths, state.scratchPackedValue, state.in, state.scratchDocIDs, docsInBlock, new IntersectVisitor() { + int i = 0; + + @Override + public void visit(int docID) throws IOException { + throw new UnsupportedOperationException(); + } + + @Override + public void visit(int docID, byte[] packedValue) throws IOException { + assert docID == state.scratchDocIDs[i]; + System.arraycopy(packedValue, 0, packedValues, i * bkd.packedBytesLength, bkd.packedBytesLength); + i++; + } + + @Override + public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) { + throw new UnsupportedOperationException(); + } + + }); + + blockID++; + } + + final int index = docBlockUpto++; + int oldDocID = state.scratchDocIDs[index]; + + int mappedDocID; + if (docMap == null) { + mappedDocID = oldDocID; + } else { + mappedDocID = docMap.get(oldDocID); + } + + if (mappedDocID != -1) { + // Not deleted! + docID = mappedDocID; + System.arraycopy(packedValues, index * bkd.packedBytesLength, state.scratchPackedValue, 0, bkd.packedBytesLength); + return true; + } + } + } + } + + private static class BKDMergeQueue extends PriorityQueue { + private final int bytesPerDim; + + public BKDMergeQueue(int bytesPerDim, int maxSize) { + super(maxSize); + this.bytesPerDim = bytesPerDim; + } + + @Override + public boolean lessThan(MergeReader a, MergeReader b) { + assert a != b; + + int cmp = StringHelper.compare(bytesPerDim, a.state.scratchPackedValue, 0, b.state.scratchPackedValue, 0); + if (cmp < 0) { + return true; + } else if (cmp > 0) { + return false; + } + + // Tie break by sorting smaller docIDs earlier: + return a.docID < b.docID; + } + } + + /** Write a field from a {@link MutablePointValues}. This way of writing + * points is faster than regular writes with {@link BKDWriter#add} since + * there is opportunity for reordering points before writing them to + * disk. This method does not use transient disk in order to reorder points. + */ + public long writeField(IndexOutput out, String fieldName, MutablePointValues reader) throws IOException { + if (numDims == 1) { + return writeField1Dim(out, fieldName, reader); + } else { + return writeFieldNDims(out, fieldName, reader); + } + } + + + /* In the 2+D case, we recursively pick the split dimension, compute the + * median value and partition other values around it. */ + private long writeFieldNDims(IndexOutput out, String fieldName, MutablePointValues values) throws IOException { + if (pointCount != 0) { + throw new IllegalStateException("cannot mix add and writeField"); + } + + // Catch user silliness: + if (heapPointWriter == null && tempInput == null) { + throw new IllegalStateException("already finished"); + } + + // Mark that we already finished: + heapPointWriter = null; + + long countPerLeaf = pointCount = values.size(); + long innerNodeCount = 1; + + while (countPerLeaf > maxPointsInLeafNode) { + countPerLeaf = (countPerLeaf+1)/2; + innerNodeCount *= 2; + } + + int numLeaves = Math.toIntExact(innerNodeCount); + + checkMaxLeafNodeCount(numLeaves); + + final byte[] splitPackedValues = new byte[numLeaves * (bytesPerDim + 1)]; + final long[] leafBlockFPs = new long[numLeaves]; + + // compute the min/max for this slice + Arrays.fill(minPackedValue, (byte) 0xff); + Arrays.fill(maxPackedValue, (byte) 0); + for (int i = 0; i < Math.toIntExact(pointCount); ++i) { + values.getValue(i, scratchBytesRef1); + for(int dim=0;dim 0) { + System.arraycopy(scratchBytesRef1.bytes, scratchBytesRef1.offset + offset, maxPackedValue, offset, bytesPerDim); + } + } + + docsSeen.set(values.getDocID(i)); + } + + build(1, numLeaves, values, 0, Math.toIntExact(pointCount), out, + minPackedValue, maxPackedValue, splitPackedValues, leafBlockFPs, + new int[maxPointsInLeafNode]); + + long indexFP = out.getFilePointer(); + writeIndex(out, leafBlockFPs, splitPackedValues); + return indexFP; + } + + + /* In the 1D case, we can simply sort points in ascending order and use the + * same writing logic as we use at merge time. */ + private long writeField1Dim(IndexOutput out, String fieldName, MutablePointValues reader) throws IOException { + MutablePointsReaderUtils.sort(maxDoc, packedBytesLength, reader, 0, Math.toIntExact(reader.size())); + + final OneDimensionBKDWriter oneDimWriter = new OneDimensionBKDWriter(out); + + reader.intersect(new IntersectVisitor() { + + @Override + public void visit(int docID, byte[] packedValue) throws IOException { + oneDimWriter.add(packedValue, docID); + } + + @Override + public void visit(int docID) throws IOException { + throw new IllegalStateException(); + } + + @Override + public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) { + return Relation.CELL_CROSSES_QUERY; + } + }); + + return oneDimWriter.finish(); + } + + // TODO: remove this opto: SimpleText is supposed to be simple! + + /** More efficient bulk-add for incoming {@link SimpleTextBKDReader}s. This does a merge sort of the already + * sorted values and currently only works when numDims==1. This returns -1 if all documents containing + * dimensional values were deleted. */ + public long merge(IndexOutput out, List docMaps, List readers) throws IOException { + assert docMaps == null || readers.size() == docMaps.size(); + + BKDMergeQueue queue = new BKDMergeQueue(bytesPerDim, readers.size()); + + for(int i=0;i totalPointCount) { + throw new IllegalStateException("totalPointCount=" + totalPointCount + " was passed when we were created, but we just hit " + pointCount + " values"); + } + + if (leafCount == maxPointsInLeafNode) { + // We write a block once we hit exactly the max count ... this is different from + // when we flush a new segment, where we write between max/2 and max per leaf block, + // so merged segments will behave differently from newly flushed segments: + writeLeafBlock(); + leafCount = 0; + } + + assert (lastDocID = docID) >= 0; // only assign when asserts are enabled + } + + public long finish() throws IOException { + if (leafCount > 0) { + writeLeafBlock(); + leafCount = 0; + } + + if (valueCount == 0) { + return -1; + } + + pointCount = valueCount; + + long indexFP = out.getFilePointer(); + + int numInnerNodes = leafBlockStartValues.size(); + + //System.out.println("BKDW: now rotate numInnerNodes=" + numInnerNodes + " leafBlockStarts=" + leafBlockStartValues.size()); + + byte[] index = new byte[(1+numInnerNodes) * (1+bytesPerDim)]; + rotateToTree(1, 0, numInnerNodes, index, leafBlockStartValues); + long[] arr = new long[leafBlockFPs.size()]; + for(int i=0;i 0) { + // Save the first (minimum) value in each leaf block except the first, to build the split value index in the end: + leafBlockStartValues.add(Arrays.copyOf(leafValues, packedBytesLength)); + } + leafBlockFPs.add(out.getFilePointer()); + checkMaxLeafNodeCount(leafBlockFPs.size()); + + Arrays.fill(commonPrefixLengths, bytesPerDim); + // Find per-dim common prefix: + for(int dim=0;dim packedValues = new IntFunction() { + final BytesRef scratch = new BytesRef(); + + { + scratch.length = packedBytesLength; + scratch.bytes = leafValues; + } + + @Override + public BytesRef apply(int i) { + scratch.offset = packedBytesLength * i; + return scratch; + } + }; + assert valuesInOrderAndBounds(leafCount, 0, Arrays.copyOf(leafValues, packedBytesLength), + Arrays.copyOfRange(leafValues, (leafCount - 1) * packedBytesLength, leafCount * packedBytesLength), + packedValues, leafDocs, 0); + writeLeafBlockPackedValues(out, commonPrefixLengths, leafCount, 0, packedValues); + } + + } + + // TODO: there must be a simpler way? + private void rotateToTree(int nodeID, int offset, int count, byte[] index, List leafBlockStartValues) { + //System.out.println("ROTATE: nodeID=" + nodeID + " offset=" + offset + " count=" + count + " bpd=" + bytesPerDim + " index.length=" + index.length); + if (count == 1) { + // Leaf index node + //System.out.println(" leaf index node"); + //System.out.println(" index[" + nodeID + "] = blockStartValues[" + offset + "]"); + System.arraycopy(leafBlockStartValues.get(offset), 0, index, nodeID*(1+bytesPerDim)+1, bytesPerDim); + } else if (count > 1) { + // Internal index node: binary partition of count + int countAtLevel = 1; + int totalCount = 0; + while (true) { + int countLeft = count - totalCount; + //System.out.println(" cycle countLeft=" + countLeft + " coutAtLevel=" + countAtLevel); + if (countLeft <= countAtLevel) { + // This is the last level, possibly partially filled: + int lastLeftCount = Math.min(countAtLevel/2, countLeft); + assert lastLeftCount >= 0; + int leftHalf = (totalCount-1)/2 + lastLeftCount; + + int rootOffset = offset + leftHalf; + /* + System.out.println(" last left count " + lastLeftCount); + System.out.println(" leftHalf " + leftHalf + " rightHalf=" + (count-leftHalf-1)); + System.out.println(" rootOffset=" + rootOffset); + */ + + System.arraycopy(leafBlockStartValues.get(rootOffset), 0, index, nodeID*(1+bytesPerDim)+1, bytesPerDim); + //System.out.println(" index[" + nodeID + "] = blockStartValues[" + rootOffset + "]"); + + // TODO: we could optimize/specialize, when we know it's simply fully balanced binary tree + // under here, to save this while loop on each recursion + + // Recurse left + rotateToTree(2*nodeID, offset, leftHalf, index, leafBlockStartValues); + + // Recurse right + rotateToTree(2*nodeID+1, rootOffset+1, count-leftHalf-1, index, leafBlockStartValues); + return; + } + totalCount += countAtLevel; + countAtLevel *= 2; + } + } else { + assert count == 0; + } + } + + // TODO: if we fixed each partition step to just record the file offset at the "split point", we could probably handle variable length + // encoding and not have our own ByteSequencesReader/Writer + + /** Sort the heap writer by the specified dim */ + private void sortHeapPointWriter(final HeapPointWriter writer, int dim) { + final int pointCount = Math.toIntExact(this.pointCount); + // Tie-break by docID: + + // No need to tie break on ord, for the case where the same doc has the same value in a given dimension indexed more than once: it + // can't matter at search time since we don't write ords into the index: + new MSBRadixSorter(bytesPerDim + Integer.BYTES) { + + @Override + protected int byteAt(int i, int k) { + assert k >= 0; + if (k < bytesPerDim) { + // dim bytes + int block = i / writer.valuesPerBlock; + int index = i % writer.valuesPerBlock; + return writer.blocks.get(block)[index * packedBytesLength + dim * bytesPerDim + k] & 0xff; + } else { + // doc id + int s = 3 - (k - bytesPerDim); + return (writer.docIDs[i] >>> (s * 8)) & 0xff; + } + } + + @Override + protected void swap(int i, int j) { + int docID = writer.docIDs[i]; + writer.docIDs[i] = writer.docIDs[j]; + writer.docIDs[j] = docID; + + if (singleValuePerDoc == false) { + if (longOrds) { + long ord = writer.ordsLong[i]; + writer.ordsLong[i] = writer.ordsLong[j]; + writer.ordsLong[j] = ord; + } else { + int ord = writer.ords[i]; + writer.ords[i] = writer.ords[j]; + writer.ords[j] = ord; + } + } + + byte[] blockI = writer.blocks.get(i / writer.valuesPerBlock); + int indexI = (i % writer.valuesPerBlock) * packedBytesLength; + byte[] blockJ = writer.blocks.get(j / writer.valuesPerBlock); + int indexJ = (j % writer.valuesPerBlock) * packedBytesLength; + + // scratch1 = values[i] + System.arraycopy(blockI, indexI, scratch1, 0, packedBytesLength); + // values[i] = values[j] + System.arraycopy(blockJ, indexJ, blockI, indexI, packedBytesLength); + // values[j] = scratch1 + System.arraycopy(scratch1, 0, blockJ, indexJ, packedBytesLength); + } + + }.sort(0, pointCount); + } + + private PointWriter sort(int dim) throws IOException { + assert dim >= 0 && dim < numDims; + + if (heapPointWriter != null) { + + assert tempInput == null; + + // We never spilled the incoming points to disk, so now we sort in heap: + HeapPointWriter sorted; + + if (dim == 0) { + // First dim can re-use the current heap writer + sorted = heapPointWriter; + } else { + // Subsequent dims need a private copy + sorted = new HeapPointWriter((int) pointCount, (int) pointCount, packedBytesLength, longOrds, singleValuePerDoc); + sorted.copyFrom(heapPointWriter); + } + + //long t0 = System.nanoTime(); + sortHeapPointWriter(sorted, dim); + //long t1 = System.nanoTime(); + //System.out.println("BKD: sort took " + ((t1-t0)/1000000.0) + " msec"); + + sorted.close(); + return sorted; + } else { + + // Offline sort: + assert tempInput != null; + + final int offset = bytesPerDim * dim; + + Comparator cmp; + if (dim == numDims - 1) { + // in that case the bytes for the dimension and for the doc id are contiguous, + // so we don't need a branch + cmp = new BytesRefComparator(bytesPerDim + Integer.BYTES) { + @Override + protected int byteAt(BytesRef ref, int i) { + return ref.bytes[ref.offset + offset + i] & 0xff; + } + }; + } else { + cmp = new BytesRefComparator(bytesPerDim + Integer.BYTES) { + @Override + protected int byteAt(BytesRef ref, int i) { + if (i < bytesPerDim) { + return ref.bytes[ref.offset + offset + i] & 0xff; + } else { + return ref.bytes[ref.offset + packedBytesLength + i - bytesPerDim] & 0xff; + } + } + }; + } + + OfflineSorter sorter = new OfflineSorter(tempDir, tempFileNamePrefix + "_bkd" + dim, cmp, offlineSorterBufferMB, offlineSorterMaxTempFiles, bytesPerDoc) { + + /** We write/read fixed-byte-width file that {@link OfflinePointReader} can read. */ + @Override + protected ByteSequencesWriter getWriter(IndexOutput out) { + return new ByteSequencesWriter(out) { + @Override + public void write(byte[] bytes, int off, int len) throws IOException { + assert len == bytesPerDoc: "len=" + len + " bytesPerDoc=" + bytesPerDoc; + out.writeBytes(bytes, off, len); + } + }; + } + + /** We write/read fixed-byte-width file that {@link OfflinePointReader} can read. */ + @Override + protected ByteSequencesReader getReader(ChecksumIndexInput in, String name) throws IOException { + return new ByteSequencesReader(in, name) { + final BytesRef scratch = new BytesRef(new byte[bytesPerDoc]); + @Override + public BytesRef next() throws IOException { + if (in.getFilePointer() >= end) { + return null; + } + in.readBytes(scratch.bytes, 0, bytesPerDoc); + return scratch; + } + }; + } + }; + + String name = sorter.sort(tempInput.getName()); + + return new OfflinePointWriter(tempDir, name, packedBytesLength, pointCount, longOrds, singleValuePerDoc); + } + } + + private void checkMaxLeafNodeCount(int numLeaves) { + if ((1+bytesPerDim) * (long) numLeaves > ArrayUtil.MAX_ARRAY_LENGTH) { + throw new IllegalStateException("too many nodes; increase maxPointsInLeafNode (currently " + maxPointsInLeafNode + ") and reindex"); + } + } + + /** Writes the BKD tree to the provided {@link IndexOutput} and returns the file offset where index was written. */ + public long finish(IndexOutput out) throws IOException { + // System.out.println("\nBKDTreeWriter.finish pointCount=" + pointCount + " out=" + out + " heapWriter=" + heapPointWriter); + + // TODO: specialize the 1D case? it's much faster at indexing time (no partitioning on recurse...) + + // Catch user silliness: + if (heapPointWriter == null && tempInput == null) { + throw new IllegalStateException("already finished"); + } + + if (offlinePointWriter != null) { + offlinePointWriter.close(); + } + + if (pointCount == 0) { + throw new IllegalStateException("must index at least one point"); + } + + LongBitSet ordBitSet; + if (numDims > 1) { + if (singleValuePerDoc) { + ordBitSet = new LongBitSet(maxDoc); + } else { + ordBitSet = new LongBitSet(pointCount); + } + } else { + ordBitSet = null; + } + + long countPerLeaf = pointCount; + long innerNodeCount = 1; + + while (countPerLeaf > maxPointsInLeafNode) { + countPerLeaf = (countPerLeaf+1)/2; + innerNodeCount *= 2; + } + + int numLeaves = (int) innerNodeCount; + + checkMaxLeafNodeCount(numLeaves); + + // NOTE: we could save the 1+ here, to use a bit less heap at search time, but then we'd need a somewhat costly check at each + // step of the recursion to recompute the split dim: + + // Indexed by nodeID, but first (root) nodeID is 1. We do 1+ because the lead byte at each recursion says which dim we split on. + byte[] splitPackedValues = new byte[Math.toIntExact(numLeaves*(1+bytesPerDim))]; + + // +1 because leaf count is power of 2 (e.g. 8), and innerNodeCount is power of 2 minus 1 (e.g. 7) + long[] leafBlockFPs = new long[numLeaves]; + + // Make sure the math above "worked": + assert pointCount / numLeaves <= maxPointsInLeafNode: "pointCount=" + pointCount + " numLeaves=" + numLeaves + " maxPointsInLeafNode=" + maxPointsInLeafNode; + + // Sort all docs once by each dimension: + PathSlice[] sortedPointWriters = new PathSlice[numDims]; + + // This is only used on exception; on normal code paths we close all files we opened: + List toCloseHeroically = new ArrayList<>(); + + boolean success = false; + try { + //long t0 = System.nanoTime(); + for(int dim=0;dim packedValues) throws IOException { + for (int i = 0; i < count; ++i) { + BytesRef packedValue = packedValues.apply(i); + // NOTE: we don't do prefix coding, so we ignore commonPrefixLengths + write(out, BLOCK_VALUE); + write(out, packedValue.toString()); + newline(out); + } + } + + private void writeLeafBlockPackedValuesRange(IndexOutput out, int[] commonPrefixLengths, int start, int end, IntFunction packedValues) throws IOException { + for (int i = start; i < end; ++i) { + BytesRef ref = packedValues.apply(i); + assert ref.length == packedBytesLength; + + for(int dim=0;dim packedValues, int start, int end, int byteOffset) { + BytesRef first = packedValues.apply(start); + byte b = first.bytes[first.offset + byteOffset]; + for (int i = start + 1; i < end; ++i) { + BytesRef ref = packedValues.apply(i); + byte b2 = ref.bytes[ref.offset + byteOffset]; + assert Byte.toUnsignedInt(b2) >= Byte.toUnsignedInt(b); + if (b != b2) { + return i - start; + } + } + return end - start; + } + + @Override + public void close() throws IOException { + if (tempInput != null) { + // NOTE: this should only happen on exception, e.g. caller calls close w/o calling finish: + try { + tempInput.close(); + } finally { + tempDir.deleteFile(tempInput.getName()); + tempInput = null; + } + } + } + + /** Sliced reference to points in an OfflineSorter.ByteSequencesWriter file. */ + private static final class PathSlice { + final PointWriter writer; + final long start; + final long count; + + public PathSlice(PointWriter writer, long start, long count) { + this.writer = writer; + this.start = start; + this.count = count; + } + + @Override + public String toString() { + return "PathSlice(start=" + start + " count=" + count + " writer=" + writer + ")"; + } + } + + /** Called on exception, to check whether the checksum is also corrupt in this source, and add that + * information (checksum matched or didn't) as a suppressed exception. */ + private void verifyChecksum(Throwable priorException, PointWriter writer) throws IOException { + // TODO: we could improve this, to always validate checksum as we recurse, if we shared left and + // right reader after recursing to children, and possibly within recursed children, + // since all together they make a single pass through the file. But this is a sizable re-org, + // and would mean leaving readers (IndexInputs) open for longer: + if (writer instanceof OfflinePointWriter) { + // We are reading from a temp file; go verify the checksum: + String tempFileName = ((OfflinePointWriter) writer).name; + try (ChecksumIndexInput in = tempDir.openChecksumInput(tempFileName, IOContext.READONCE)) { + CodecUtil.checkFooter(in, priorException); + } + } else { + // We are reading from heap; nothing to add: + IOUtils.reThrow(priorException); + } + } + + /** Marks bits for the ords (points) that belong in the right sub tree (those docs that have values >= the splitValue). */ + private byte[] markRightTree(long rightCount, int splitDim, PathSlice source, LongBitSet ordBitSet) throws IOException { + + // Now we mark ords that fall into the right half, so we can partition on all other dims that are not the split dim: + + // Read the split value, then mark all ords in the right tree (larger than the split value): + + // TODO: find a way to also checksum this reader? If we changed to markLeftTree, and scanned the final chunk, it could work? + try (PointReader reader = source.writer.getReader(source.start + source.count - rightCount, rightCount)) { + boolean result = reader.next(); + assert result; + System.arraycopy(reader.packedValue(), splitDim*bytesPerDim, scratch1, 0, bytesPerDim); + if (numDims > 1) { + assert ordBitSet.get(reader.ord()) == false; + ordBitSet.set(reader.ord()); + // Subtract 1 from rightCount because we already did the first value above (so we could record the split value): + reader.markOrds(rightCount-1, ordBitSet); + } + } catch (Throwable t) { + verifyChecksum(t, source.writer); + } + + return scratch1; + } + + /** Called only in assert */ + private boolean valueInBounds(BytesRef packedValue, byte[] minPackedValue, byte[] maxPackedValue) { + for(int dim=0;dim 0) { + return false; + } + } + + return true; + } + + protected int split(byte[] minPackedValue, byte[] maxPackedValue) { + // Find which dim has the largest span so we can split on it: + int splitDim = -1; + for(int dim=0;dim 0) { + System.arraycopy(scratchDiff, 0, scratch1, 0, bytesPerDim); + splitDim = dim; + } + } + + //System.out.println("SPLIT: " + splitDim); + return splitDim; + } + + /** Pull a partition back into heap once the point count is low enough while recursing. */ + private PathSlice switchToHeap(PathSlice source, List toCloseHeroically) throws IOException { + int count = Math.toIntExact(source.count); + // Not inside the try because we don't want to close it here: + PointReader reader = source.writer.getSharedReader(source.start, source.count, toCloseHeroically); + try (PointWriter writer = new HeapPointWriter(count, count, packedBytesLength, longOrds, singleValuePerDoc)) { + for(int i=0;i= leafNodeOffset) { + // leaf node + final int count = to - from; + assert count <= maxPointsInLeafNode; + + // Compute common prefixes + Arrays.fill(commonPrefixLengths, bytesPerDim); + reader.getValue(from, scratchBytesRef1); + for (int i = from + 1; i < to; ++i) { + reader.getValue(i, scratchBytesRef2); + for (int dim=0;dim packedValues = new IntFunction() { + @Override + public BytesRef apply(int i) { + reader.getValue(from + i, scratchBytesRef1); + return scratchBytesRef1; + } + }; + assert valuesInOrderAndBounds(count, sortedDim, minPackedValue, maxPackedValue, packedValues, + docIDs, 0); + writeLeafBlockPackedValues(out, commonPrefixLengths, count, sortedDim, packedValues); + + } else { + // inner node + + // compute the split dimension and partition around it + final int splitDim = split(minPackedValue, maxPackedValue); + final int mid = (from + to + 1) >>> 1; + + int commonPrefixLen = bytesPerDim; + for (int i = 0; i < bytesPerDim; ++i) { + if (minPackedValue[splitDim * bytesPerDim + i] != maxPackedValue[splitDim * bytesPerDim + i]) { + commonPrefixLen = i; + break; + } + } + MutablePointsReaderUtils.partition(maxDoc, splitDim, bytesPerDim, commonPrefixLen, + reader, from, to, mid, scratchBytesRef1, scratchBytesRef2); + + // set the split value + final int address = nodeID * (1+bytesPerDim); + splitPackedValues[address] = (byte) splitDim; + reader.getValue(mid, scratchBytesRef1); + System.arraycopy(scratchBytesRef1.bytes, scratchBytesRef1.offset + splitDim * bytesPerDim, splitPackedValues, address + 1, bytesPerDim); + + byte[] minSplitPackedValue = Arrays.copyOf(minPackedValue, packedBytesLength); + byte[] maxSplitPackedValue = Arrays.copyOf(maxPackedValue, packedBytesLength); + System.arraycopy(scratchBytesRef1.bytes, scratchBytesRef1.offset + splitDim * bytesPerDim, + minSplitPackedValue, splitDim * bytesPerDim, bytesPerDim); + System.arraycopy(scratchBytesRef1.bytes, scratchBytesRef1.offset + splitDim * bytesPerDim, + maxSplitPackedValue, splitDim * bytesPerDim, bytesPerDim); + + // recurse + build(nodeID * 2, leafNodeOffset, reader, from, mid, out, + minPackedValue, maxSplitPackedValue, splitPackedValues, leafBlockFPs, spareDocIds); + build(nodeID * 2 + 1, leafNodeOffset, reader, mid, to, out, + minSplitPackedValue, maxPackedValue, splitPackedValues, leafBlockFPs, spareDocIds); + } + } + + /** The array (sized numDims) of PathSlice describe the cell we have currently recursed to. */ + private void build(int nodeID, int leafNodeOffset, + PathSlice[] slices, + LongBitSet ordBitSet, + IndexOutput out, + byte[] minPackedValue, byte[] maxPackedValue, + byte[] splitPackedValues, + long[] leafBlockFPs, + List toCloseHeroically) throws IOException { + + for(PathSlice slice : slices) { + assert slice.count == slices[0].count; + } + + if (numDims == 1 && slices[0].writer instanceof OfflinePointWriter && slices[0].count <= maxPointsSortInHeap) { + // Special case for 1D, to cutover to heap once we recurse deeply enough: + slices[0] = switchToHeap(slices[0], toCloseHeroically); + } + + if (nodeID >= leafNodeOffset) { + + // Leaf node: write block + // We can write the block in any order so by default we write it sorted by the dimension that has the + // least number of unique bytes at commonPrefixLengths[dim], which makes compression more efficient + int sortedDim = 0; + int sortedDimCardinality = Integer.MAX_VALUE; + + for (int dim=0;dim= maxPointsInLeafNode, so we better be in heap at this point: + HeapPointWriter heapSource = (HeapPointWriter) source.writer; + + // Save the block file pointer: + leafBlockFPs[nodeID - leafNodeOffset] = out.getFilePointer(); + //System.out.println(" write leaf block @ fp=" + out.getFilePointer()); + + // Write docIDs first, as their own chunk, so that at intersect time we can add all docIDs w/o + // loading the values: + int count = Math.toIntExact(source.count); + assert count > 0: "nodeID=" + nodeID + " leafNodeOffset=" + leafNodeOffset; + writeLeafBlockDocs(out, heapSource.docIDs, Math.toIntExact(source.start), count); + + // TODO: minor opto: we don't really have to write the actual common prefixes, because BKDReader on recursing can regenerate it for us + // from the index, much like how terms dict does so from the FST: + + // Write the full values: + IntFunction packedValues = new IntFunction() { + final BytesRef scratch = new BytesRef(); + + { + scratch.length = packedBytesLength; + } + + @Override + public BytesRef apply(int i) { + heapSource.getPackedValueSlice(Math.toIntExact(source.start + i), scratch); + return scratch; + } + }; + assert valuesInOrderAndBounds(count, sortedDim, minPackedValue, maxPackedValue, packedValues, + heapSource.docIDs, Math.toIntExact(source.start)); + writeLeafBlockPackedValues(out, commonPrefixLengths, count, sortedDim, packedValues); + + } else { + // Inner node: partition/recurse + + int splitDim; + if (numDims > 1) { + splitDim = split(minPackedValue, maxPackedValue); + } else { + splitDim = 0; + } + + PathSlice source = slices[splitDim]; + + assert nodeID < splitPackedValues.length: "nodeID=" + nodeID + " splitValues.length=" + splitPackedValues.length; + + // How many points will be in the left tree: + long rightCount = source.count / 2; + long leftCount = source.count - rightCount; + + byte[] splitValue = markRightTree(rightCount, splitDim, source, ordBitSet); + int address = nodeID * (1+bytesPerDim); + splitPackedValues[address] = (byte) splitDim; + System.arraycopy(splitValue, 0, splitPackedValues, address + 1, bytesPerDim); + + // Partition all PathSlice that are not the split dim into sorted left and right sets, so we can recurse: + + PathSlice[] leftSlices = new PathSlice[numDims]; + PathSlice[] rightSlices = new PathSlice[numDims]; + + byte[] minSplitPackedValue = new byte[packedBytesLength]; + System.arraycopy(minPackedValue, 0, minSplitPackedValue, 0, packedBytesLength); + + byte[] maxSplitPackedValue = new byte[packedBytesLength]; + System.arraycopy(maxPackedValue, 0, maxSplitPackedValue, 0, packedBytesLength); + + // When we are on this dim, below, we clear the ordBitSet: + int dimToClear; + if (numDims - 1 == splitDim) { + dimToClear = numDims - 2; + } else { + dimToClear = numDims - 1; + } + + for(int dim=0;dim values, int[] docs, int docsOffset) throws IOException { + byte[] lastPackedValue = new byte[packedBytesLength]; + int lastDoc = -1; + for (int i=0;i 0) { + int cmp = StringHelper.compare(bytesPerDim, lastPackedValue, dimOffset, packedValue, packedValueOffset + dimOffset); + if (cmp > 0) { + throw new AssertionError("values out of order: last value=" + new BytesRef(lastPackedValue) + " current value=" + new BytesRef(packedValue, packedValueOffset, packedBytesLength) + " ord=" + ord); + } + if (cmp == 0 && doc < lastDoc) { + throw new AssertionError("docs out of order: last doc=" + lastDoc + " current doc=" + doc + " ord=" + ord); + } + } + System.arraycopy(packedValue, packedValueOffset, lastPackedValue, 0, packedBytesLength); + return true; + } + + PointWriter getPointWriter(long count, String desc) throws IOException { + if (count <= maxPointsSortInHeap) { + int size = Math.toIntExact(count); + return new HeapPointWriter(size, size, packedBytesLength, longOrds, singleValuePerDoc); + } else { + return new OfflinePointWriter(tempDir, tempFileNamePrefix, packedBytesLength, longOrds, desc, count, singleValuePerDoc); + } + } + + private void write(IndexOutput out, String s) throws IOException { + SimpleTextUtil.write(out, s, scratch); + } + + private void writeInt(IndexOutput out, int x) throws IOException { + SimpleTextUtil.write(out, Integer.toString(x), scratch); + } + + private void writeLong(IndexOutput out, long x) throws IOException { + SimpleTextUtil.write(out, Long.toString(x), scratch); + } + + private void write(IndexOutput out, BytesRef b) throws IOException { + SimpleTextUtil.write(out, b); + } + + private void newline(IndexOutput out) throws IOException { + SimpleTextUtil.writeNewline(out); + } +} diff --git a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextPointsReader.java b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextPointsReader.java index f7ff16ecbc2..453bd2384b2 100644 --- a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextPointsReader.java +++ b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextPointsReader.java @@ -36,7 +36,6 @@ import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.BytesRefBuilder; import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.StringHelper; -import org.apache.lucene.util.bkd.BKDReader; import static org.apache.lucene.codecs.simpletext.SimpleTextPointsWriter.BLOCK_FP; import static org.apache.lucene.codecs.simpletext.SimpleTextPointsWriter.BYTES_PER_DIM; @@ -58,7 +57,7 @@ class SimpleTextPointsReader extends PointsReader { private final IndexInput dataIn; final SegmentReadState readState; - final Map readers = new HashMap<>(); + final Map readers = new HashMap<>(); final BytesRefBuilder scratch = new BytesRefBuilder(); public SimpleTextPointsReader(SegmentReadState readState) throws IOException { @@ -98,7 +97,7 @@ class SimpleTextPointsReader extends PointsReader { this.readState = readState; } - private BKDReader initReader(long fp) throws IOException { + private SimpleTextBKDReader initReader(long fp) throws IOException { // NOTE: matches what writeIndex does in SimpleTextPointsWriter dataIn.seek(fp); readLine(dataIn); diff --git a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextPointsWriter.java b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextPointsWriter.java index c06c128d154..9d2db890fa0 100644 --- a/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextPointsWriter.java +++ b/lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextPointsWriter.java @@ -20,7 +20,6 @@ package org.apache.lucene.codecs.simpletext; import java.io.IOException; import java.util.HashMap; import java.util.Map; -import java.util.function.IntFunction; import org.apache.lucene.codecs.PointsReader; import org.apache.lucene.codecs.PointsWriter; @@ -33,29 +32,28 @@ import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.BytesRefBuilder; -import org.apache.lucene.util.bkd.BKDWriter; class SimpleTextPointsWriter extends PointsWriter { - final static BytesRef NUM_DIMS = new BytesRef("num dims "); - final static BytesRef BYTES_PER_DIM = new BytesRef("bytes per dim "); - final static BytesRef MAX_LEAF_POINTS = new BytesRef("max leaf points "); - final static BytesRef INDEX_COUNT = new BytesRef("index count "); - final static BytesRef BLOCK_COUNT = new BytesRef("block count "); - final static BytesRef BLOCK_DOC_ID = new BytesRef(" doc "); - final static BytesRef BLOCK_FP = new BytesRef(" block fp "); - final static BytesRef BLOCK_VALUE = new BytesRef(" block value "); - final static BytesRef SPLIT_COUNT = new BytesRef("split count "); - final static BytesRef SPLIT_DIM = new BytesRef(" split dim "); - final static BytesRef SPLIT_VALUE = new BytesRef(" split value "); - final static BytesRef FIELD_COUNT = new BytesRef("field count "); - final static BytesRef FIELD_FP_NAME = new BytesRef(" field fp name "); - final static BytesRef FIELD_FP = new BytesRef(" field fp "); - final static BytesRef MIN_VALUE = new BytesRef("min value "); - final static BytesRef MAX_VALUE = new BytesRef("max value "); - final static BytesRef POINT_COUNT = new BytesRef("point count "); - final static BytesRef DOC_COUNT = new BytesRef("doc count "); - final static BytesRef END = new BytesRef("END"); + public final static BytesRef NUM_DIMS = new BytesRef("num dims "); + public final static BytesRef BYTES_PER_DIM = new BytesRef("bytes per dim "); + public final static BytesRef MAX_LEAF_POINTS = new BytesRef("max leaf points "); + public final static BytesRef INDEX_COUNT = new BytesRef("index count "); + public final static BytesRef BLOCK_COUNT = new BytesRef("block count "); + public final static BytesRef BLOCK_DOC_ID = new BytesRef(" doc "); + public final static BytesRef BLOCK_FP = new BytesRef(" block fp "); + public final static BytesRef BLOCK_VALUE = new BytesRef(" block value "); + public final static BytesRef SPLIT_COUNT = new BytesRef("split count "); + public final static BytesRef SPLIT_DIM = new BytesRef(" split dim "); + public final static BytesRef SPLIT_VALUE = new BytesRef(" split value "); + public final static BytesRef FIELD_COUNT = new BytesRef("field count "); + public final static BytesRef FIELD_FP_NAME = new BytesRef(" field fp name "); + public final static BytesRef FIELD_FP = new BytesRef(" field fp "); + public final static BytesRef MIN_VALUE = new BytesRef("min value "); + public final static BytesRef MAX_VALUE = new BytesRef("max value "); + public final static BytesRef POINT_COUNT = new BytesRef("point count "); + public final static BytesRef DOC_COUNT = new BytesRef("doc count "); + public final static BytesRef END = new BytesRef("END"); private IndexOutput dataOut; final BytesRefBuilder scratch = new BytesRefBuilder(); @@ -75,105 +73,15 @@ class SimpleTextPointsWriter extends PointsWriter { boolean singleValuePerDoc = values.size() == values.getDocCount(); // We use the normal BKDWriter, but subclass to customize how it writes the index and blocks to disk: - try (BKDWriter writer = new BKDWriter(writeState.segmentInfo.maxDoc(), - writeState.directory, - writeState.segmentInfo.name, - fieldInfo.getPointDimensionCount(), - fieldInfo.getPointNumBytes(), - BKDWriter.DEFAULT_MAX_POINTS_IN_LEAF_NODE, - BKDWriter.DEFAULT_MAX_MB_SORT_IN_HEAP, - values.size(), - singleValuePerDoc) { - - @Override - protected void writeIndex(IndexOutput out, long[] leafBlockFPs, byte[] splitPackedValues) throws IOException { - write(out, NUM_DIMS); - writeInt(out, numDims); - newline(out); - - write(out, BYTES_PER_DIM); - writeInt(out, bytesPerDim); - newline(out); - - write(out, MAX_LEAF_POINTS); - writeInt(out, maxPointsInLeafNode); - newline(out); - - write(out, INDEX_COUNT); - writeInt(out, leafBlockFPs.length); - newline(out); - - write(out, MIN_VALUE); - BytesRef br = new BytesRef(minPackedValue, 0, minPackedValue.length); - write(out, br.toString()); - newline(out); - - write(out, MAX_VALUE); - br = new BytesRef(maxPackedValue, 0, maxPackedValue.length); - write(out, br.toString()); - newline(out); - - write(out, POINT_COUNT); - writeLong(out, pointCount); - newline(out); - - write(out, DOC_COUNT); - writeInt(out, docsSeen.cardinality()); - newline(out); - - for(int i=0;i packedValues) throws IOException { - for (int i = 0; i < count; ++i) { - BytesRef packedValue = packedValues.apply(i); - // NOTE: we don't do prefix coding, so we ignore commonPrefixLengths - write(out, BLOCK_VALUE); - write(out, packedValue.toString()); - newline(out); - } - } - }) { + try (SimpleTextBKDWriter writer = new SimpleTextBKDWriter(writeState.segmentInfo.maxDoc(), + writeState.directory, + writeState.segmentInfo.name, + fieldInfo.getPointDimensionCount(), + fieldInfo.getPointNumBytes(), + SimpleTextBKDWriter.DEFAULT_MAX_POINTS_IN_LEAF_NODE, + SimpleTextBKDWriter.DEFAULT_MAX_MB_SORT_IN_HEAP, + values.size(), + singleValuePerDoc)) { values.intersect(new IntersectVisitor() { @Override @@ -198,26 +106,6 @@ class SimpleTextPointsWriter extends PointsWriter { } } - private void write(IndexOutput out, String s) throws IOException { - SimpleTextUtil.write(out, s, scratch); - } - - private void writeInt(IndexOutput out, int x) throws IOException { - SimpleTextUtil.write(out, Integer.toString(x), scratch); - } - - private void writeLong(IndexOutput out, long x) throws IOException { - SimpleTextUtil.write(out, Long.toString(x), scratch); - } - - private void write(IndexOutput out, BytesRef b) throws IOException { - SimpleTextUtil.write(out, b); - } - - private void newline(IndexOutput out) throws IOException { - SimpleTextUtil.writeNewline(out); - } - @Override public void finish() throws IOException { SimpleTextUtil.write(dataOut, END); @@ -250,4 +138,24 @@ class SimpleTextPointsWriter extends PointsWriter { } } } + + private void write(IndexOutput out, String s) throws IOException { + SimpleTextUtil.write(out, s, scratch); + } + + private void writeInt(IndexOutput out, int x) throws IOException { + SimpleTextUtil.write(out, Integer.toString(x), scratch); + } + + private void writeLong(IndexOutput out, long x) throws IOException { + SimpleTextUtil.write(out, Long.toString(x), scratch); + } + + private void write(IndexOutput out, BytesRef b) throws IOException { + SimpleTextUtil.write(out, b); + } + + private void newline(IndexOutput out) throws IOException { + SimpleTextUtil.writeNewline(out); + } } diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60PointsFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60PointsFormat.java index e558d0d4fa8..1d2285c73b6 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60PointsFormat.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene60/Lucene60PointsFormat.java @@ -28,7 +28,8 @@ import org.apache.lucene.index.SegmentWriteState; /** * Lucene 6.0 point format, which encodes dimensional values in a block KD-tree structure - * for fast shape intersection filtering. See this paper for details. + * for fast 1D range and N dimesional shape intersection filtering. + * See this paper for details. * *

This data structure is written as a series of blocks on disk, with an in-memory perfectly balanced * binary tree of split values referencing those blocks at the leaves. @@ -50,10 +51,13 @@ import org.apache.lucene.index.SegmentWriteState; *

  • maxPointsInLeafNode (vInt) *
  • bytesPerDim (vInt) *
  • count (vInt) - *
  • byte[bytesPerDim]count (packed byte[] all split values) - *
  • delta-blockFP (vLong)count (delta-coded file pointers to the on-disk leaf blocks)) + *
  • packed index (byte[]) * * + *

    The packed index uses hierarchical delta and prefix coding to compactly encode the file pointer for + * all leaf blocks, once the tree is traversed, as well as the split dimension and split value for each + * inner node of the tree. + * *

    After all fields blocks + index data are written, {@link CodecUtil#writeFooter} writes the checksum. * *

    The .dii file records the file pointer in the .dim file where each field's diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene60/package-info.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene60/package-info.java index 8968a6d624c..a914001d9d2 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene60/package-info.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene60/package-info.java @@ -16,7 +16,7 @@ */ /** - * Components from the Lucene 6.0 index format. See {@link org.apache.lucene.codecs.lucene62} - * for an overview of the index format. + * Components from the Lucene 6.0 index format. See {@link org.apache.lucene.codecs.lucene70} + * for an overview of the current index format. */ package org.apache.lucene.codecs.lucene60; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene62/package-info.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene62/package-info.java index 2fe2dc74b4a..fb556732d08 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene62/package-info.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene62/package-info.java @@ -17,8 +17,8 @@ /** * Components from the Lucene 6.2 index format - * See {@link org.apache.lucene.codecs.lucene62} for an overview - * of the index format. + * See {@link org.apache.lucene.codecs.lucene70} for an overview + * of the current index format. */ package org.apache.lucene.codecs.lucene62; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene70/package-info.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene70/package-info.java index 9b432f7c4f4..cab2859766e 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene70/package-info.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene70/package-info.java @@ -185,6 +185,12 @@ * {@link org.apache.lucene.codecs.lucene50.Lucene50LiveDocsFormat Live documents}. * An optional file indicating which documents are live. *

    • + *
  • + * {@link org.apache.lucene.codecs.lucene60.Lucene60PointsFormat Point values}. + * Optional pair of files, recording dimensionally indexed fields, to enable fast + * numeric range filtering and large numeric values like BigInteger and BigDecimal (1D) + * and geographic shape intersection (2D, 3D). + *
    • * *

    Details on each of these are provided in their linked pages.

    * @@ -300,7 +306,12 @@ * * {@link org.apache.lucene.codecs.lucene50.Lucene50LiveDocsFormat Live Documents} * .liv - * Info about what files are live + * Info about what documents are live + * + * + * {@link org.apache.lucene.codecs.lucene60.Lucene60PointsFormat Point values} + * .dii, .dim + * Holds indexed points, if any * * * @@ -374,6 +385,8 @@ * that is suitable for faceting/sorting/analytics. *
  • In version 5.4, DocValues have been improved to store more information on disk: * addresses for binary fields and ord indexes for multi-valued fields. + *
  • In version 6.0, Points were added, for multi-dimensional range/distance search. + *
  • In version 6.2, new Segment info format that reads/writes the index sort, to support index sorting. *
  • In version 7.0, DocValues have been improved to better support sparse doc values * thanks to an iterator API. *
    • diff --git a/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java b/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java index 7bc08f3c4a8..fd8011d4d07 100644 --- a/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java +++ b/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java @@ -1801,161 +1801,32 @@ public final class CheckIndex implements Closeable { } for (FieldInfo fieldInfo : fieldInfos) { if (fieldInfo.getPointDimensionCount() > 0) { - FixedBitSet docsSeen = new FixedBitSet(reader.maxDoc()); - status.totalValueFields++; - int dimCount = fieldInfo.getPointDimensionCount(); - int bytesPerDim = fieldInfo.getPointNumBytes(); - int packedBytesCount = dimCount * bytesPerDim; - byte[] lastMinPackedValue = new byte[packedBytesCount]; - byte[] lastMaxPackedValue = new byte[packedBytesCount]; - BytesRef scratch = new BytesRef(); - scratch.length = bytesPerDim; - byte[] lastPackedValue = new byte[packedBytesCount]; - - long[] pointCountSeen = new long[1]; - PointValues values = pointsReader.getValues(fieldInfo.name); if (values == null) { continue; } - byte[] globalMinPackedValue = values.getMinPackedValue(); + + status.totalValueFields++; + long size = values.size(); int docCount = values.getDocCount(); - if (docCount > size) { - throw new RuntimeException("point values for field \"" + fieldInfo.name + "\" claims to have size=" + size + " points and inconsistent docCount=" + docCount); + VerifyPointsVisitor visitor = new VerifyPointsVisitor(fieldInfo.name, reader.maxDoc(), values); + values.intersect(visitor); + + if (visitor.getPointCountSeen() != size) { + throw new RuntimeException("point values for field \"" + fieldInfo.name + "\" claims to have size=" + size + " points, but in fact has " + visitor.getPointCountSeen()); } - if (docCount > reader.maxDoc()) { - throw new RuntimeException("point values for field \"" + fieldInfo.name + "\" claims to have docCount=" + docCount + " but that's greater than maxDoc=" + reader.maxDoc()); + if (visitor.getDocCountSeen() != docCount) { + throw new RuntimeException("point values for field \"" + fieldInfo.name + "\" claims to have docCount=" + docCount + " but in fact has " + visitor.getDocCountSeen()); } - if (globalMinPackedValue == null) { - if (size != 0) { - throw new RuntimeException("getMinPackedValue is null points for field \"" + fieldInfo.name + "\" yet size=" + size); - } - } else if (globalMinPackedValue.length != packedBytesCount) { - throw new RuntimeException("getMinPackedValue for field \"" + fieldInfo.name + "\" return length=" + globalMinPackedValue.length + " array, but should be " + packedBytesCount); - } - byte[] globalMaxPackedValue = values.getMaxPackedValue(); - if (globalMaxPackedValue == null) { - if (size != 0) { - throw new RuntimeException("getMaxPackedValue is null points for field \"" + fieldInfo.name + "\" yet size=" + size); - } - } else if (globalMaxPackedValue.length != packedBytesCount) { - throw new RuntimeException("getMaxPackedValue for field \"" + fieldInfo.name + "\" return length=" + globalMaxPackedValue.length + " array, but should be " + packedBytesCount); - } - - values.intersect(new PointValues.IntersectVisitor() { - - private int lastDocID = -1; - - @Override - public void visit(int docID) { - throw new RuntimeException("codec called IntersectVisitor.visit without a packed value for docID=" + docID); - } - - @Override - public void visit(int docID, byte[] packedValue) { - checkPackedValue("packed value", packedValue, docID); - pointCountSeen[0]++; - docsSeen.set(docID); - - for(int dim=0;dim 0) { - throw new RuntimeException("packed points cell minPackedValue " + Arrays.toString(minPackedValue) + - " is out-of-bounds of the cell's maxPackedValue " + Arrays.toString(maxPackedValue) + " dim=" + dim + " field=\"" + fieldInfo.name + "\""); - } - - // Make sure this cell is not outside of the global min/max: - if (StringHelper.compare(bytesPerDim, minPackedValue, offset, globalMinPackedValue, offset) < 0) { - throw new RuntimeException("packed points cell minPackedValue " + Arrays.toString(minPackedValue) + - " is out-of-bounds of the global minimum " + Arrays.toString(globalMinPackedValue) + " dim=" + dim + " field=\"" + fieldInfo.name + "\""); - } - - if (StringHelper.compare(bytesPerDim, maxPackedValue, offset, globalMinPackedValue, offset) < 0) { - throw new RuntimeException("packed points cell maxPackedValue " + Arrays.toString(maxPackedValue) + - " is out-of-bounds of the global minimum " + Arrays.toString(globalMinPackedValue) + " dim=" + dim + " field=\"" + fieldInfo.name + "\""); - } - - if (StringHelper.compare(bytesPerDim, minPackedValue, offset, globalMaxPackedValue, offset) > 0) { - throw new RuntimeException("packed points cell minPackedValue " + Arrays.toString(minPackedValue) + - " is out-of-bounds of the global maximum " + Arrays.toString(globalMaxPackedValue) + " dim=" + dim + " field=\"" + fieldInfo.name + "\""); - } - if (StringHelper.compare(bytesPerDim, maxPackedValue, offset, globalMaxPackedValue, offset) > 0) { - throw new RuntimeException("packed points cell maxPackedValue " + Arrays.toString(maxPackedValue) + - " is out-of-bounds of the global maximum " + Arrays.toString(globalMaxPackedValue) + " dim=" + dim + " field=\"" + fieldInfo.name + "\""); - } - } - - // We always pretend the query shape is so complex that it crosses every cell, so - // that packedValue is passed for every document - return PointValues.Relation.CELL_CROSSES_QUERY; - } - - private void checkPackedValue(String desc, byte[] packedValue, int docID) { - if (packedValue == null) { - throw new RuntimeException(desc + " is null for docID=" + docID + " field=\"" + fieldInfo.name + "\""); - } - - if (packedValue.length != packedBytesCount) { - throw new RuntimeException(desc + " has incorrect length=" + packedValue.length + " vs expected=" + packedBytesCount + " for docID=" + docID + " field=\"" + fieldInfo.name + "\""); - } - } - }); - - if (pointCountSeen[0] != size) { - throw new RuntimeException("point values for field \"" + fieldInfo.name + "\" claims to have size=" + size + " points, but in fact has " + pointCountSeen[0]); - } - - if (docsSeen.cardinality() != docCount) { - throw new RuntimeException("point values for field \"" + fieldInfo.name + "\" claims to have docCount=" + docCount + " but in fact has " + docsSeen.cardinality()); - } + status.totalValuePoints += visitor.getPointCountSeen(); } } } + msg(infoStream, String.format(Locale.ROOT, "OK [%d fields, %d points] [took %.3f sec]", status.totalValueFields, status.totalValuePoints, nsToSec(System.nanoTime()-startNS))); } catch (Throwable e) { @@ -1972,6 +1843,167 @@ public final class CheckIndex implements Closeable { return status; } + /** Walks the entire N-dimensional points space, verifying that all points fall within the last cell's boundaries. + * + * @lucene.internal */ + public static class VerifyPointsVisitor implements PointValues.IntersectVisitor { + private long pointCountSeen; + private int lastDocID = -1; + private final int maxDoc; + private final FixedBitSet docsSeen; + private final byte[] lastMinPackedValue; + private final byte[] lastMaxPackedValue; + private final byte[] lastPackedValue; + private final byte[] globalMinPackedValue; + private final byte[] globalMaxPackedValue; + private final int packedBytesCount; + private final int numDims; + private final int bytesPerDim; + private final String fieldName; + + /** Sole constructor */ + public VerifyPointsVisitor(String fieldName, int maxDoc, PointValues values) throws IOException { + this.maxDoc = maxDoc; + this.fieldName = fieldName; + numDims = values.getNumDimensions(); + bytesPerDim = values.getBytesPerDimension(); + packedBytesCount = numDims * bytesPerDim; + globalMinPackedValue = values.getMinPackedValue(); + globalMaxPackedValue = values.getMaxPackedValue(); + docsSeen = new FixedBitSet(maxDoc); + lastMinPackedValue = new byte[packedBytesCount]; + lastMaxPackedValue = new byte[packedBytesCount]; + lastPackedValue = new byte[packedBytesCount]; + + if (values.getDocCount() > values.size()) { + throw new RuntimeException("point values for field \"" + fieldName + "\" claims to have size=" + values.size() + " points and inconsistent docCount=" + values.getDocCount()); + } + + if (values.getDocCount() > maxDoc) { + throw new RuntimeException("point values for field \"" + fieldName + "\" claims to have docCount=" + values.getDocCount() + " but that's greater than maxDoc=" + maxDoc); + } + + if (globalMinPackedValue == null) { + if (values.size() != 0) { + throw new RuntimeException("getMinPackedValue is null points for field \"" + fieldName + "\" yet size=" + values.size()); + } + } else if (globalMinPackedValue.length != packedBytesCount) { + throw new RuntimeException("getMinPackedValue for field \"" + fieldName + "\" return length=" + globalMinPackedValue.length + " array, but should be " + packedBytesCount); + } + if (globalMaxPackedValue == null) { + if (values.size() != 0) { + throw new RuntimeException("getMaxPackedValue is null points for field \"" + fieldName + "\" yet size=" + values.size()); + } + } else if (globalMaxPackedValue.length != packedBytesCount) { + throw new RuntimeException("getMaxPackedValue for field \"" + fieldName + "\" return length=" + globalMaxPackedValue.length + " array, but should be " + packedBytesCount); + } + } + + /** Returns total number of points in this BKD tree */ + public long getPointCountSeen() { + return pointCountSeen; + } + + /** Returns total number of unique docIDs in this BKD tree */ + public long getDocCountSeen() { + return docsSeen.cardinality(); + } + + @Override + public void visit(int docID) { + throw new RuntimeException("codec called IntersectVisitor.visit without a packed value for docID=" + docID); + } + + @Override + public void visit(int docID, byte[] packedValue) { + checkPackedValue("packed value", packedValue, docID); + pointCountSeen++; + docsSeen.set(docID); + + for(int dim=0;dim 0) { + throw new RuntimeException("packed points cell minPackedValue " + Arrays.toString(minPackedValue) + + " is out-of-bounds of the cell's maxPackedValue " + Arrays.toString(maxPackedValue) + " dim=" + dim + " field=\"" + fieldName + "\""); + } + + // Make sure this cell is not outside of the global min/max: + if (StringHelper.compare(bytesPerDim, minPackedValue, offset, globalMinPackedValue, offset) < 0) { + throw new RuntimeException("packed points cell minPackedValue " + Arrays.toString(minPackedValue) + + " is out-of-bounds of the global minimum " + Arrays.toString(globalMinPackedValue) + " dim=" + dim + " field=\"" + fieldName + "\""); + } + + if (StringHelper.compare(bytesPerDim, maxPackedValue, offset, globalMinPackedValue, offset) < 0) { + throw new RuntimeException("packed points cell maxPackedValue " + Arrays.toString(maxPackedValue) + + " is out-of-bounds of the global minimum " + Arrays.toString(globalMinPackedValue) + " dim=" + dim + " field=\"" + fieldName + "\""); + } + + if (StringHelper.compare(bytesPerDim, minPackedValue, offset, globalMaxPackedValue, offset) > 0) { + throw new RuntimeException("packed points cell minPackedValue " + Arrays.toString(minPackedValue) + + " is out-of-bounds of the global maximum " + Arrays.toString(globalMaxPackedValue) + " dim=" + dim + " field=\"" + fieldName + "\""); + } + if (StringHelper.compare(bytesPerDim, maxPackedValue, offset, globalMaxPackedValue, offset) > 0) { + throw new RuntimeException("packed points cell maxPackedValue " + Arrays.toString(maxPackedValue) + + " is out-of-bounds of the global maximum " + Arrays.toString(globalMaxPackedValue) + " dim=" + dim + " field=\"" + fieldName + "\""); + } + } + + // We always pretend the query shape is so complex that it crosses every cell, so + // that packedValue is passed for every document + return PointValues.Relation.CELL_CROSSES_QUERY; + } + + private void checkPackedValue(String desc, byte[] packedValue, int docID) { + if (packedValue == null) { + throw new RuntimeException(desc + " is null for docID=" + docID + " field=\"" + fieldName + "\""); + } + + if (packedValue.length != packedBytesCount) { + throw new RuntimeException(desc + " has incorrect length=" + packedValue.length + " vs expected=" + packedBytesCount + " for docID=" + docID + " field=\"" + fieldName + "\""); + } + } + } + + /** * Test stored fields. * @lucene.experimental diff --git a/lucene/core/src/java/org/apache/lucene/util/bkd/BKDReader.java b/lucene/core/src/java/org/apache/lucene/util/bkd/BKDReader.java index 6bf7dfc1a86..6cccf4cf1d1 100644 --- a/lucene/core/src/java/org/apache/lucene/util/bkd/BKDReader.java +++ b/lucene/core/src/java/org/apache/lucene/util/bkd/BKDReader.java @@ -17,14 +17,15 @@ package org.apache.lucene.util.bkd; import java.io.IOException; -import java.util.Arrays; import org.apache.lucene.codecs.CodecUtil; import org.apache.lucene.index.CorruptIndexException; import org.apache.lucene.index.PointValues; +import org.apache.lucene.store.ByteArrayDataInput; import org.apache.lucene.store.IndexInput; import org.apache.lucene.util.Accountable; import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.MathUtil; import org.apache.lucene.util.RamUsageEstimator; import org.apache.lucene.util.StringHelper; @@ -32,14 +33,12 @@ import org.apache.lucene.util.StringHelper; * * @lucene.experimental */ -public class BKDReader extends PointValues implements Accountable { +public final class BKDReader extends PointValues implements Accountable { // Packed array of byte[] holding all split values in the full binary tree: - final private byte[] splitPackedValues; - final long[] leafBlockFPs; - final private int leafNodeOffset; + final int leafNodeOffset; final int numDims; final int bytesPerDim; - final int bytesPerIndexEntry; + final int numLeaves; final IndexInput in; final int maxPointsInLeafNode; final byte[] minPackedValue; @@ -49,6 +48,14 @@ public class BKDReader extends PointValues implements Accountable { final int version; protected final int packedBytesLength; + // Used for 6.4.0+ index format: + final byte[] packedIndex; + + // Used for Legacy (pre-6.4.0) index format, to hold a compact form of the index: + final private byte[] splitPackedValues; + final int bytesPerIndexEntry; + final long[] leafBlockFPs; + /** Caller must pre-seek the provided {@link IndexInput} to the index location that {@link BKDWriter#finish} returned */ public BKDReader(IndexInput in) throws IOException { version = CodecUtil.checkHeader(in, BKDWriter.CODEC_NAME, BKDWriter.VERSION_START, BKDWriter.VERSION_CURRENT); @@ -59,7 +66,7 @@ public class BKDReader extends PointValues implements Accountable { packedBytesLength = numDims * bytesPerDim; // Read index: - int numLeaves = in.readVInt(); + numLeaves = in.readVInt(); assert numLeaves > 0; leafNodeOffset = numLeaves; @@ -78,203 +85,378 @@ public class BKDReader extends PointValues implements Accountable { pointCount = in.readVLong(); docCount = in.readVInt(); - splitPackedValues = new byte[bytesPerIndexEntry*numLeaves]; - - // TODO: don't write split packed values[0]! - in.readBytes(splitPackedValues, 0, splitPackedValues.length); - - // Read the file pointers to the start of each leaf block: - long[] leafBlockFPs = new long[numLeaves]; - long lastFP = 0; - for(int i=0;i 1) { - //System.out.println("BKDR: numLeaves=" + numLeaves); - int levelCount = 2; - while (true) { - //System.out.println(" cycle levelCount=" + levelCount); - if (numLeaves >= levelCount && numLeaves <= 2*levelCount) { - int lastLevel = 2*(numLeaves - levelCount); - assert lastLevel >= 0; - /* - System.out.println("BKDR: lastLevel=" + lastLevel + " vs " + levelCount); - System.out.println("FPs before:"); - for(int i=0;i= maxDoc) { - throw new RuntimeException("docID=" + docID + " is out of bounds of 0.." + maxDoc); - } - for(int dim=0;dim 0) { - throw new RuntimeException("value=" + new BytesRef(packedValue, dim*bytesPerDim, bytesPerDim) + " for docID=" + docID + " dim=" + dim + " is less than this leaf block's minimum=" + new BytesRef(cellMinPacked, dim*bytesPerDim, bytesPerDim)); - } - if (StringHelper.compare(bytesPerDim, cellMaxPacked, dim*bytesPerDim, packedValue, dim*bytesPerDim) < 0) { - throw new RuntimeException("value=" + new BytesRef(packedValue, dim*bytesPerDim, bytesPerDim) + " for docID=" + docID + " dim=" + dim + " is greater than this leaf block's maximum=" + new BytesRef(cellMaxPacked, dim*bytesPerDim, bytesPerDim)); - } - } - - if (numDims == 1) { - // With only 1D, all values should always be in sorted order - if (lastPackedValue == null) { - lastPackedValue = Arrays.copyOf(packedValue, packedValue.length); - } else if (StringHelper.compare(bytesPerDim, lastPackedValue, 0, packedValue, 0) > 0) { - throw new RuntimeException("value=" + new BytesRef(packedValue) + " for docID=" + docID + " dim=0" + " sorts before last value=" + new BytesRef(lastPackedValue)); - } else { - System.arraycopy(packedValue, 0, lastPackedValue, 0, bytesPerDim); - } - } - } - - @Override - public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) { - throw new UnsupportedOperationException(); - } - } - - /** Only used for debugging, to make sure all values in each leaf block fall within the range expected by the index */ - // TODO: maybe we can get this into CheckIndex? - public void verify(int maxDoc) throws IOException { - //System.out.println("BKDR.verify this=" + this); - // Visits every doc in every leaf block and confirms that - // their values agree with the index: - byte[] rootMinPacked = new byte[packedBytesLength]; - byte[] rootMaxPacked = new byte[packedBytesLength]; - Arrays.fill(rootMaxPacked, (byte) 0xff); - verify(getIntersectState(new VerifyVisitor(numDims, bytesPerDim, maxDoc)), 1, rootMinPacked, rootMaxPacked); - } - - private void verify(IntersectState state, int nodeID, byte[] cellMinPacked, byte[] cellMaxPacked) throws IOException { - - if (nodeID >= leafNodeOffset) { - int leafID = nodeID - leafNodeOffset; - - // In the unbalanced case it's possible the left most node only has one child: - if (leafID < leafBlockFPs.length) { - //System.out.println("CHECK nodeID=" + nodeID + " leaf=" + (nodeID-leafNodeOffset) + " offset=" + leafNodeOffset + " fp=" + leafBlockFPs[leafID]); - //System.out.println("BKDR.verify leafID=" + leafID + " nodeID=" + nodeID + " fp=" + leafBlockFPs[leafID] + " min=" + new BytesRef(cellMinPacked) + " max=" + new BytesRef(cellMaxPacked)); - - // Leaf node: check that all values are in fact in bounds: - VerifyVisitor visitor = (VerifyVisitor) state.visitor; - visitor.cellMinPacked = cellMinPacked; - visitor.cellMaxPacked = cellMaxPacked; - - int count = readDocIDs(state.in, leafBlockFPs[leafID], state.scratchDocIDs); - visitDocValues(state.commonPrefixLengths, state.scratchPackedValue, state.in, state.scratchDocIDs, count, state.visitor); - } else { - //System.out.println("BKDR.verify skip leafID=" + leafID); - } + if (version >= BKDWriter.VERSION_PACKED_INDEX) { + int numBytes = in.readVInt(); + packedIndex = new byte[numBytes]; + in.readBytes(packedIndex, 0, numBytes); + leafBlockFPs = null; + splitPackedValues = null; } else { - // Non-leaf node: + // legacy un-packed index - int address = nodeID * bytesPerIndexEntry; - int splitDim; - if (numDims == 1) { - splitDim = 0; - if (version < BKDWriter.VERSION_IMPLICIT_SPLIT_DIM_1D) { - // skip over wastefully encoded 0 splitDim: - assert splitPackedValues[address] == 0; - address++; + splitPackedValues = new byte[bytesPerIndexEntry*numLeaves]; + + in.readBytes(splitPackedValues, 0, splitPackedValues.length); + + // Read the file pointers to the start of each leaf block: + long[] leafBlockFPs = new long[numLeaves]; + long lastFP = 0; + for(int i=0;i 1) { + int levelCount = 2; + while (true) { + if (numLeaves >= levelCount && numLeaves <= 2*levelCount) { + int lastLevel = 2*(numLeaves - levelCount); + assert lastLevel >= 0; + if (lastLevel != 0) { + // Last level is partially filled, so we must rotate the leaf FPs to match. We do this here, after loading + // at read-time, so that we can still delta code them on disk at write: + long[] newLeafBlockFPs = new long[numLeaves]; + System.arraycopy(leafBlockFPs, lastLevel, newLeafBlockFPs, 0, leafBlockFPs.length - lastLevel); + System.arraycopy(leafBlockFPs, 0, newLeafBlockFPs, leafBlockFPs.length - lastLevel, lastLevel); + leafBlockFPs = newLeafBlockFPs; + } + break; + } + + levelCount *= 2; } - } else { - splitDim = splitPackedValues[address++] & 0xff; } - assert splitDim < numDims; - - byte[] splitPackedValue = new byte[packedBytesLength]; - - // Recurse on left sub-tree: - System.arraycopy(cellMaxPacked, 0, splitPackedValue, 0, packedBytesLength); - System.arraycopy(splitPackedValues, address, splitPackedValue, splitDim*bytesPerDim, bytesPerDim); - verify(state, - 2*nodeID, - cellMinPacked, splitPackedValue); - - // Recurse on right sub-tree: - System.arraycopy(cellMinPacked, 0, splitPackedValue, 0, packedBytesLength); - System.arraycopy(splitPackedValues, address, splitPackedValue, splitDim*bytesPerDim, bytesPerDim); - verify(state, - 2*nodeID+1, - splitPackedValue, cellMaxPacked); + this.leafBlockFPs = leafBlockFPs; + packedIndex = null; } + + this.in = in; + } + + long getMinLeafBlockFP() { + if (packedIndex != null) { + return new ByteArrayDataInput(packedIndex).readVLong(); + } else { + long minFP = Long.MAX_VALUE; + for(long fp : leafBlockFPs) { + minFP = Math.min(minFP, fp); + } + return minFP; + } + } + + /** Used to walk the in-heap index + * + * @lucene.internal */ + public abstract class IndexTree implements Cloneable { + protected int nodeID; + // level is 1-based so that we can do level-1 w/o checking each time: + protected int level; + protected int splitDim; + protected final byte[][] splitPackedValueStack; + + protected IndexTree() { + int treeDepth = getTreeDepth(); + splitPackedValueStack = new byte[treeDepth+1][]; + nodeID = 1; + level = 1; + splitPackedValueStack[level] = new byte[packedBytesLength]; + } + + public void pushLeft() { + nodeID *= 2; + level++; + if (splitPackedValueStack[level] == null) { + splitPackedValueStack[level] = new byte[packedBytesLength]; + } + } + + /** Clone, but you are not allowed to pop up past the point where the clone happened. */ + public abstract IndexTree clone(); + + public void pushRight() { + nodeID = nodeID * 2 + 1; + level++; + if (splitPackedValueStack[level] == null) { + splitPackedValueStack[level] = new byte[packedBytesLength]; + } + } + + public void pop() { + nodeID /= 2; + level--; + splitDim = -1; + //System.out.println(" pop nodeID=" + nodeID); + } + + public boolean isLeafNode() { + return nodeID >= leafNodeOffset; + } + + public boolean nodeExists() { + return nodeID - leafNodeOffset < leafNodeOffset; + } + + public int getNodeID() { + return nodeID; + } + + public byte[] getSplitPackedValue() { + assert isLeafNode() == false; + assert splitPackedValueStack[level] != null: "level=" + level; + return splitPackedValueStack[level]; + } + + /** Only valid after pushLeft or pushRight, not pop! */ + public int getSplitDim() { + assert isLeafNode() == false; + return splitDim; + } + + /** Only valid after pushLeft or pushRight, not pop! */ + public abstract BytesRef getSplitDimValue(); + + /** Only valid after pushLeft or pushRight, not pop! */ + public abstract long getLeafBlockFP(); + } + + /** Reads the original simple yet heap-heavy index format */ + private final class LegacyIndexTree extends IndexTree { + + private long leafBlockFP; + private final byte[] splitDimValue = new byte[bytesPerDim]; + private final BytesRef scratch = new BytesRef(); + + public LegacyIndexTree() { + setNodeData(); + scratch.bytes = splitDimValue; + scratch.length = bytesPerDim; + } + + @Override + public LegacyIndexTree clone() { + LegacyIndexTree index = new LegacyIndexTree(); + index.nodeID = nodeID; + index.level = level; + index.splitDim = splitDim; + index.leafBlockFP = leafBlockFP; + index.splitPackedValueStack[index.level] = splitPackedValueStack[index.level].clone(); + + return index; + } + + @Override + public void pushLeft() { + super.pushLeft(); + setNodeData(); + } + + @Override + public void pushRight() { + super.pushRight(); + setNodeData(); + } + + private void setNodeData() { + if (isLeafNode()) { + leafBlockFP = leafBlockFPs[nodeID - leafNodeOffset]; + splitDim = -1; + } else { + leafBlockFP = -1; + int address = nodeID * bytesPerIndexEntry; + if (numDims == 1) { + splitDim = 0; + if (version < BKDWriter.VERSION_IMPLICIT_SPLIT_DIM_1D) { + // skip over wastefully encoded 0 splitDim: + assert splitPackedValues[address] == 0; + address++; + } + } else { + splitDim = splitPackedValues[address++] & 0xff; + } + System.arraycopy(splitPackedValues, address, splitDimValue, 0, bytesPerDim); + } + } + + @Override + public long getLeafBlockFP() { + assert isLeafNode(); + return leafBlockFP; + } + + @Override + public BytesRef getSplitDimValue() { + assert isLeafNode() == false; + return scratch; + } + + @Override + public void pop() { + super.pop(); + leafBlockFP = -1; + } + } + + /** Reads the new packed byte[] index format which can be up to ~63% smaller than the legacy index format on 20M NYC taxis tests. This + * format takes advantage of the limited access pattern to the BKD tree at search time, i.e. starting at the root node and recursing + * downwards one child at a time. */ + private final class PackedIndexTree extends IndexTree { + // used to read the packed byte[] + private final ByteArrayDataInput in; + // holds the minimum (left most) leaf block file pointer for each level we've recursed to: + private final long[] leafBlockFPStack; + // holds the address, in the packed byte[] index, of the left-node of each level: + private final int[] leftNodePositions; + // holds the address, in the packed byte[] index, of the right-node of each level: + private final int[] rightNodePositions; + // holds the splitDim for each level: + private final int[] splitDims; + // true if the per-dim delta we read for the node at this level is a negative offset vs. the last split on this dim; this is a packed + // 2D array, i.e. to access array[level][dim] you read from negativeDeltas[level*numDims+dim]. this will be true if the last time we + // split on this dimension, we next pushed to the left sub-tree: + private final boolean[] negativeDeltas; + // holds the packed per-level split values; the intersect method uses this to save the cell min/max as it recurses: + private final byte[][] splitValuesStack; + // scratch value to return from getPackedValue: + private final BytesRef scratch; + + public PackedIndexTree() { + int treeDepth = getTreeDepth(); + leafBlockFPStack = new long[treeDepth+1]; + leftNodePositions = new int[treeDepth+1]; + rightNodePositions = new int[treeDepth+1]; + splitValuesStack = new byte[treeDepth+1][]; + splitDims = new int[treeDepth+1]; + negativeDeltas = new boolean[numDims*(treeDepth+1)]; + + in = new ByteArrayDataInput(packedIndex); + splitValuesStack[0] = new byte[packedBytesLength]; + readNodeData(false); + scratch = new BytesRef(); + scratch.length = bytesPerDim; + } + + @Override + public PackedIndexTree clone() { + PackedIndexTree index = new PackedIndexTree(); + index.nodeID = nodeID; + index.level = level; + index.splitDim = splitDim; + System.arraycopy(negativeDeltas, level*numDims, index.negativeDeltas, level*numDims, numDims); + index.leafBlockFPStack[level] = leafBlockFPStack[level]; + index.leftNodePositions[level] = leftNodePositions[level]; + index.rightNodePositions[level] = rightNodePositions[level]; + index.splitValuesStack[index.level] = splitValuesStack[index.level].clone(); + System.arraycopy(negativeDeltas, level*numDims, index.negativeDeltas, level*numDims, numDims); + index.splitDims[level] = splitDims[level]; + return index; + } + + @Override + public void pushLeft() { + int nodePosition = leftNodePositions[level]; + super.pushLeft(); + System.arraycopy(negativeDeltas, (level-1)*numDims, negativeDeltas, level*numDims, numDims); + assert splitDim != -1; + negativeDeltas[level*numDims+splitDim] = true; + in.setPosition(nodePosition); + readNodeData(true); + } + + @Override + public void pushRight() { + int nodePosition = rightNodePositions[level]; + super.pushRight(); + System.arraycopy(negativeDeltas, (level-1)*numDims, negativeDeltas, level*numDims, numDims); + assert splitDim != -1; + negativeDeltas[level*numDims+splitDim] = false; + in.setPosition(nodePosition); + readNodeData(false); + } + + @Override + public void pop() { + super.pop(); + splitDim = splitDims[level]; + } + + @Override + public long getLeafBlockFP() { + assert isLeafNode(): "nodeID=" + nodeID + " is not a leaf"; + return leafBlockFPStack[level]; + } + + @Override + public BytesRef getSplitDimValue() { + assert isLeafNode() == false; + scratch.bytes = splitValuesStack[level]; + scratch.offset = splitDim * bytesPerDim; + return scratch; + } + + private void readNodeData(boolean isLeft) { + + leafBlockFPStack[level] = leafBlockFPStack[level-1]; + + // read leaf block FP delta + if (isLeft == false) { + leafBlockFPStack[level] += in.readVLong(); + } + + if (isLeafNode()) { + splitDim = -1; + } else { + + // read split dim, prefix, firstDiffByteDelta encoded as int: + int code = in.readVInt(); + splitDim = code % numDims; + splitDims[level] = splitDim; + code /= numDims; + int prefix = code % (1+bytesPerDim); + int suffix = bytesPerDim - prefix; + + if (splitValuesStack[level] == null) { + splitValuesStack[level] = new byte[packedBytesLength]; + } + System.arraycopy(splitValuesStack[level-1], 0, splitValuesStack[level], 0, packedBytesLength); + if (suffix > 0) { + int firstDiffByteDelta = code / (1+bytesPerDim); + if (negativeDeltas[level*numDims + splitDim]) { + firstDiffByteDelta = -firstDiffByteDelta; + } + int oldByte = splitValuesStack[level][splitDim*bytesPerDim+prefix] & 0xFF; + splitValuesStack[level][splitDim*bytesPerDim+prefix] = (byte) (oldByte + firstDiffByteDelta); + in.readBytes(splitValuesStack[level], splitDim*bytesPerDim+prefix+1, suffix-1); + } else { + // our split value is == last split value in this dim, which can happen when there are many duplicate values + } + + int leftNumBytes; + if (nodeID * 2 < leafNodeOffset) { + leftNumBytes = in.readVInt(); + } else { + leftNumBytes = 0; + } + + leftNodePositions[level] = in.getPosition(); + rightNodePositions[level] = leftNodePositions[level] + leftNumBytes; + } + } + } + + private int getTreeDepth() { + // First +1 because all the non-leave nodes makes another power + // of 2; e.g. to have a fully balanced tree with 4 leaves you + // need a depth=3 tree: + + // Second +1 because MathUtil.log computes floor of the logarithm; e.g. + // with 5 leaves you need a depth=4 tree: + return MathUtil.log(numLeaves, 2) + 2; } /** Used to track all state for a single call to {@link #intersect}. */ @@ -285,57 +467,73 @@ public class BKDReader extends PointValues implements Accountable { final int[] commonPrefixLengths; final IntersectVisitor visitor; + public final IndexTree index; public IntersectState(IndexInput in, int numDims, int packedBytesLength, int maxPointsInLeafNode, - IntersectVisitor visitor) { + IntersectVisitor visitor, + IndexTree indexVisitor) { this.in = in; this.visitor = visitor; this.commonPrefixLengths = new int[numDims]; this.scratchDocIDs = new int[maxPointsInLeafNode]; this.scratchPackedValue = new byte[packedBytesLength]; + this.index = indexVisitor; } } public void intersect(IntersectVisitor visitor) throws IOException { - intersect(getIntersectState(visitor), 1, minPackedValue, maxPackedValue); + intersect(getIntersectState(visitor), minPackedValue, maxPackedValue); } /** Fast path: this is called when the query box fully encompasses all cells under this node. */ - private void addAll(IntersectState state, int nodeID) throws IOException { + private void addAll(IntersectState state) throws IOException { //System.out.println("R: addAll nodeID=" + nodeID); - if (nodeID >= leafNodeOffset) { + if (state.index.isLeafNode()) { //System.out.println("ADDALL"); - visitDocIDs(state.in, leafBlockFPs[nodeID-leafNodeOffset], state.visitor); + if (state.index.nodeExists()) { + visitDocIDs(state.in, state.index.getLeafBlockFP(), state.visitor); + } // TODO: we can assert that the first value here in fact matches what the index claimed? } else { - addAll(state, 2*nodeID); - addAll(state, 2*nodeID+1); + state.index.pushLeft(); + addAll(state); + state.index.pop(); + + state.index.pushRight(); + addAll(state); + state.index.pop(); } } /** Create a new {@link IntersectState} */ public IntersectState getIntersectState(IntersectVisitor visitor) { + IndexTree index; + if (packedIndex != null) { + index = new PackedIndexTree(); + } else { + index = new LegacyIndexTree(); + } return new IntersectState(in.clone(), numDims, packedBytesLength, maxPointsInLeafNode, - visitor); + visitor, + index); } /** Visits all docIDs and packed values in a single leaf block */ - public void visitLeafBlockValues(int nodeID, IntersectState state) throws IOException { - int leafID = nodeID - leafNodeOffset; + public void visitLeafBlockValues(IndexTree index, IntersectState state) throws IOException { // Leaf node; scan and filter all points in this block: - int count = readDocIDs(state.in, leafBlockFPs[leafID], state.scratchDocIDs); + int count = readDocIDs(state.in, index.getLeafBlockFP(), state.scratchDocIDs); // Again, this time reading values and checking with the visitor visitDocValues(state.commonPrefixLengths, state.scratchPackedValue, state.in, state.scratchDocIDs, count, state.visitor); } - protected void visitDocIDs(IndexInput in, long blockFP, IntersectVisitor visitor) throws IOException { + private void visitDocIDs(IndexInput in, long blockFP, IntersectVisitor visitor) throws IOException { // Leaf node in.seek(blockFP); @@ -350,7 +548,7 @@ public class BKDReader extends PointValues implements Accountable { } } - protected int readDocIDs(IndexInput in, long blockFP, int[] docIDs) throws IOException { + int readDocIDs(IndexInput in, long blockFP, int[] docIDs) throws IOException { in.seek(blockFP); // How many points are stored in this leaf cell: @@ -365,7 +563,7 @@ public class BKDReader extends PointValues implements Accountable { return count; } - protected void visitDocValues(int[] commonPrefixLengths, byte[] scratchPackedValue, IndexInput in, int[] docIDs, int count, IntersectVisitor visitor) throws IOException { + void visitDocValues(int[] commonPrefixLengths, byte[] scratchPackedValue, IndexInput in, int[] docIDs, int count, IntersectVisitor visitor) throws IOException { visitor.grow(count); readCommonPrefixes(commonPrefixLengths, scratchPackedValue, in); @@ -434,13 +632,10 @@ public class BKDReader extends PointValues implements Accountable { } } - private void intersect(IntersectState state, - int nodeID, - byte[] cellMinPacked, byte[] cellMaxPacked) - throws IOException { + private void intersect(IntersectState state, byte[] cellMinPacked, byte[] cellMaxPacked) throws IOException { /* - System.out.println("\nR: intersect nodeID=" + nodeID); + System.out.println("\nR: intersect nodeID=" + state.index.getNodeID()); for(int dim=0;dim= 0 && dim < numDims; @@ -1019,46 +1034,238 @@ public class BKDWriter implements Closeable { return indexFP; } - /** Subclass can change how it writes the index. */ - protected void writeIndex(IndexOutput out, long[] leafBlockFPs, byte[] splitPackedValues) throws IOException { + /** Packs the two arrays, representing a balanced binary tree, into a compact byte[] structure. */ + private byte[] packIndex(long[] leafBlockFPs, byte[] splitPackedValues) throws IOException { + + int numLeaves = leafBlockFPs.length; + + // Possibly rotate the leaf block FPs, if the index not fully balanced binary tree (only happens + // if it was created by OneDimensionBKDWriter). In this case the leaf nodes may straddle the two bottom + // levels of the binary tree: + if (numDims == 1 && numLeaves > 1) { + int levelCount = 2; + while (true) { + if (numLeaves >= levelCount && numLeaves <= 2*levelCount) { + int lastLevel = 2*(numLeaves - levelCount); + assert lastLevel >= 0; + if (lastLevel != 0) { + // Last level is partially filled, so we must rotate the leaf FPs to match. We do this here, after loading + // at read-time, so that we can still delta code them on disk at write: + long[] newLeafBlockFPs = new long[numLeaves]; + System.arraycopy(leafBlockFPs, lastLevel, newLeafBlockFPs, 0, leafBlockFPs.length - lastLevel); + System.arraycopy(leafBlockFPs, 0, newLeafBlockFPs, leafBlockFPs.length - lastLevel, lastLevel); + leafBlockFPs = newLeafBlockFPs; + } + break; + } + + levelCount *= 2; + } + } + + /** Reused while packing the index */ + RAMOutputStream writeBuffer = new RAMOutputStream(); + + // This is the "file" we append the byte[] to: + List blocks = new ArrayList<>(); + byte[] lastSplitValues = new byte[bytesPerDim * numDims]; + //System.out.println("\npack index"); + int totalSize = recursePackIndex(writeBuffer, leafBlockFPs, splitPackedValues, 0l, blocks, 1, lastSplitValues, new boolean[numDims], false); + + // Compact the byte[] blocks into single byte index: + byte[] index = new byte[totalSize]; + int upto = 0; + for(byte[] block : blocks) { + System.arraycopy(block, 0, index, upto, block.length); + upto += block.length; + } + assert upto == totalSize; + + return index; + } + + /** Appends the current contents of writeBuffer as another block on the growing in-memory file */ + private int appendBlock(RAMOutputStream writeBuffer, List blocks) throws IOException { + int pos = Math.toIntExact(writeBuffer.getFilePointer()); + byte[] bytes = new byte[pos]; + writeBuffer.writeTo(bytes, 0); + writeBuffer.reset(); + blocks.add(bytes); + return pos; + } + + /** + * lastSplitValues is per-dimension split value previously seen; we use this to prefix-code the split byte[] on each inner node + */ + private int recursePackIndex(RAMOutputStream writeBuffer, long[] leafBlockFPs, byte[] splitPackedValues, long minBlockFP, List blocks, + int nodeID, byte[] lastSplitValues, boolean[] negativeDeltas, boolean isLeft) throws IOException { + if (nodeID >= leafBlockFPs.length) { + int leafID = nodeID - leafBlockFPs.length; + //System.out.println("recursePack leaf nodeID=" + nodeID); + + // In the unbalanced case it's possible the left most node only has one child: + if (leafID < leafBlockFPs.length) { + long delta = leafBlockFPs[leafID] - minBlockFP; + if (isLeft) { + assert delta == 0; + return 0; + } else { + assert nodeID == 1 || delta > 0: "nodeID=" + nodeID; + writeBuffer.writeVLong(delta); + return appendBlock(writeBuffer, blocks); + } + } else { + return 0; + } + } else { + long leftBlockFP; + if (isLeft == false) { + leftBlockFP = getLeftMostLeafBlockFP(leafBlockFPs, nodeID); + long delta = leftBlockFP - minBlockFP; + assert nodeID == 1 || delta > 0; + writeBuffer.writeVLong(delta); + } else { + // The left tree's left most leaf block FP is always the minimal FP: + leftBlockFP = minBlockFP; + } + + int address = nodeID * (1+bytesPerDim); + int splitDim = splitPackedValues[address++] & 0xff; + + //System.out.println("recursePack inner nodeID=" + nodeID + " splitDim=" + splitDim + " splitValue=" + new BytesRef(splitPackedValues, address, bytesPerDim)); + + // find common prefix with last split value in this dim: + int prefix = 0; + for(;prefix 0; - out.writeVInt(leafBlockFPs.length); + assert numLeaves > 0; + out.writeVInt(numLeaves); out.writeBytes(minPackedValue, 0, packedBytesLength); out.writeBytes(maxPackedValue, 0, packedBytesLength); out.writeVLong(pointCount); out.writeVInt(docsSeen.cardinality()); - - // NOTE: splitPackedValues[0] is unused, because nodeID is 1-based: - if (numDims == 1) { - // write the index, skipping the byte used to store the split dim since it is always 0 - for (int i = 1; i < splitPackedValues.length; i += 1 + bytesPerDim) { - out.writeBytes(splitPackedValues, i, bytesPerDim); - } - } else { - out.writeBytes(splitPackedValues, 0, splitPackedValues.length); - } - - long lastFP = 0; - for (int i=0;i 0: "maxPointsInLeafNode=" + maxPointsInLeafNode; out.writeVInt(count); DocIdsWriter.writeDocIds(docIDs, start, count, out); } - protected void writeLeafBlockPackedValues(IndexOutput out, int[] commonPrefixLengths, int count, int sortedDim, IntFunction packedValues) throws IOException { + private void writeLeafBlockPackedValues(IndexOutput out, int[] commonPrefixLengths, int count, int sortedDim, IntFunction packedValues) throws IOException { int prefixLenSum = Arrays.stream(commonPrefixLengths).sum(); if (prefixLenSum == packedBytesLength) { // all values in this block are equal @@ -1109,7 +1316,7 @@ public class BKDWriter implements Closeable { return end - start; } - protected void writeCommonPrefixes(IndexOutput out, int[] commonPrefixes, byte[] packedValue) throws IOException { + private void writeCommonPrefixes(IndexOutput out, int[] commonPrefixes, byte[] packedValue) throws IOException { for(int dim=0;dim(terms[idx], - outputs.newPair((long) idx, value))); + outputs.newPair((long) idx, value))); } new FSTTester<>(random(), dir, inputMode, pairs, outputs, false).doTest(true); } diff --git a/lucene/sandbox/src/java/org/apache/lucene/document/NearestNeighbor.java b/lucene/sandbox/src/java/org/apache/lucene/document/NearestNeighbor.java index 3b9f302f5eb..587c63fb7a3 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/document/NearestNeighbor.java +++ b/lucene/sandbox/src/java/org/apache/lucene/document/NearestNeighbor.java @@ -26,7 +26,10 @@ import org.apache.lucene.geo.Rectangle; import org.apache.lucene.index.PointValues.IntersectVisitor; import org.apache.lucene.index.PointValues.Relation; import org.apache.lucene.util.Bits; +import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.SloppyMath; +import org.apache.lucene.util.bkd.BKDReader.IndexTree; +import org.apache.lucene.util.bkd.BKDReader.IntersectState; import org.apache.lucene.util.bkd.BKDReader; import static org.apache.lucene.geo.GeoEncodingUtils.decodeLatitude; @@ -41,16 +44,16 @@ class NearestNeighbor { static class Cell implements Comparable { final int readerIndex; - final int nodeID; final byte[] minPacked; final byte[] maxPacked; + final IndexTree index; /** The closest possible distance of all points in this cell */ final double distanceMeters; - public Cell(int readerIndex, int nodeID, byte[] minPacked, byte[] maxPacked, double distanceMeters) { + public Cell(IndexTree index, int readerIndex, byte[] minPacked, byte[] maxPacked, double distanceMeters) { + this.index = index; this.readerIndex = readerIndex; - this.nodeID = nodeID; this.minPacked = minPacked.clone(); this.maxPacked = maxPacked.clone(); this.distanceMeters = distanceMeters; @@ -66,7 +69,7 @@ class NearestNeighbor { double minLon = decodeLongitude(minPacked, Integer.BYTES); double maxLat = decodeLatitude(maxPacked, 0); double maxLon = decodeLongitude(maxPacked, Integer.BYTES); - return "Cell(readerIndex=" + readerIndex + " lat=" + minLat + " TO " + maxLat + ", lon=" + minLon + " TO " + maxLon + "; distanceMeters=" + distanceMeters + ")"; + return "Cell(readerIndex=" + readerIndex + " nodeID=" + index.getNodeID() + " isLeaf=" + index.isLeafNode() + " lat=" + minLat + " TO " + maxLat + ", lon=" + minLon + " TO " + maxLon + "; distanceMeters=" + distanceMeters + ")"; } } @@ -219,13 +222,21 @@ class NearestNeighbor { List states = new ArrayList<>(); // Add root cell for each reader into the queue: + int bytesPerDim = -1; + for(int i=0;i