maedhroz commented on code in PR #2409: URL: https://github.com/apache/cassandra/pull/2409#discussion_r1228445739
########## src/java/org/apache/cassandra/index/sai/disk/v1/bbtree/TraversingBlockBalancedTreeReader.java: ########## @@ -0,0 +1,332 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.cassandra.index.sai.disk.v1.bbtree; + +import java.io.Closeable; + +import org.agrona.collections.IntArrayList; +import org.apache.cassandra.index.sai.disk.io.IndexInputReader; +import org.apache.cassandra.index.sai.disk.v1.SAICodecUtils; +import org.apache.cassandra.io.util.FileHandle; +import org.apache.cassandra.io.util.FileUtils; +import org.apache.cassandra.io.util.RandomAccessReader; +import org.apache.cassandra.utils.ObjectSizes; +import org.apache.cassandra.utils.Throwables; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.store.ByteArrayDataInput; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.FutureArrays; +import org.apache.lucene.util.MathUtil; + +/** + * Base reader for a block balanced tree previously written with {@link BlockBalancedTreeWriter}. + * + * Holds the index tree on heap and enables its traversal via {@link #traverse(IndexTreeTraversalCallback)}. + */ +public class TraversingBlockBalancedTreeReader implements Closeable +{ + final FileHandle indexFile; + final int bytesPerValue; + final int numLeaves; + final byte[] minPackedValue; + final byte[] maxPackedValue; + // Packed array of byte[] holding all split values in the full binary tree: + final byte[] packedIndex; + final long pointCount; + final int leafNodeOffset; + final int maxPointsInLeafNode; + final int packedBytesLength; + + TraversingBlockBalancedTreeReader(FileHandle indexFile, long root) + { + this.indexFile = indexFile; + + try (RandomAccessReader reader = indexFile.createReader(); + IndexInputReader in = IndexInputReader.create(reader)) + { + SAICodecUtils.validate(in); + in.seek(root); + + maxPointsInLeafNode = in.readVInt(); + bytesPerValue = in.readVInt(); + packedBytesLength = bytesPerValue; + + // Read index: + numLeaves = in.readVInt(); + assert numLeaves > 0; + leafNodeOffset = numLeaves; + + minPackedValue = new byte[packedBytesLength]; + maxPackedValue = new byte[packedBytesLength]; + + in.readBytes(minPackedValue, 0, packedBytesLength); + in.readBytes(maxPackedValue, 0, packedBytesLength); + + if (FutureArrays.compareUnsigned(minPackedValue, 0, bytesPerValue, maxPackedValue, 0, bytesPerValue) > 0) + { + String message = String.format("Min packed value %s is > max packed value %s.", + new BytesRef(minPackedValue), new BytesRef(maxPackedValue)); + throw new CorruptIndexException(message, in); + } + + pointCount = in.readVLong(); + + int numBytes = in.readVInt(); + packedIndex = new byte[numBytes]; + in.readBytes(packedIndex, 0, numBytes); + } + catch (Throwable t) + { + FileUtils.closeQuietly(indexFile); + throw Throwables.unchecked(t); + } + } + + public long memoryUsage() + { + return ObjectSizes.sizeOfArray(packedIndex) + + ObjectSizes.sizeOfArray(minPackedValue) + + ObjectSizes.sizeOfArray(maxPackedValue); + } + + @Override + public void close() + { + FileUtils.closeQuietly(indexFile); + } + + void traverse(IndexTreeTraversalCallback callback) + { + traverse(callback, new PackedIndexTree(), new IntArrayList()); + } + + private void traverse(IndexTreeTraversalCallback callback, PackedIndexTree index, IntArrayList pathToRoot) + { + if (index.isLeafNode()) + { + // In the unbalanced case it's possible the left most node only has one child: + if (index.nodeExists()) + { + callback.onLeaf(index.getNodeID(), index.getLeafBlockFP(), pathToRoot); + } + } + else + { + int nodeID = index.getNodeID(); + IntArrayList currentPath = new IntArrayList(); + currentPath.addAll(pathToRoot); + currentPath.add(nodeID); + + index.pushLeft(); + traverse(callback, index, currentPath); + index.pop(); + + index.pushRight(); + traverse(callback, index, currentPath); + index.pop(); + } + } + + /** + * Copy of BKDReader#getTreeDepth() Review Comment: ```suggestion * Copy of {@link BKDReader#getTreeDepth()} ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
