xiangfu0 commented on code in PR #19307: URL: https://github.com/apache/pinot/pull/19307#discussion_r3943092719
########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/forward/FixedByteChunkSVForwardIndexReaderV7.java: ########## @@ -0,0 +1,398 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.segment.local.segment.index.readers.forward; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import org.apache.pinot.segment.local.io.codec.CodecPipelineExecutor; +import org.apache.pinot.segment.local.io.writer.impl.FixedByteChunkForwardIndexWriterV7; +import org.apache.pinot.segment.spi.index.ForwardIndexConfig; +import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.apache.pinot.spi.data.FieldSpec.DataType; + + +/// Chunk-based single-value raw forward index reader for version-7 files written by +/// [FixedByteChunkForwardIndexWriterV7]. +/// +/// Reads the canonical codec spec from the file header, instantiates a +/// [CodecPipelineExecutor], and uses it to decode each chunk on demand. +/// +/// Supported data types: INT, LONG. +/// +/// **Threading:** the reader instance itself is immutable after construction and is safe to +/// share across threads, but each [ChunkReaderContext] is single-threaded. The +/// [ByteBuffer] returned by `getChunkBuffer` is the context's reusable scratch buffer — +/// its position/limit are mutated on every chunk transition and it must not be retained across +/// subsequent `getInt`/`getLong` calls. +public final class FixedByteChunkSVForwardIndexReaderV7 implements ForwardIndexReader<ChunkReaderContext> { + + public static final int VERSION = ForwardIndexConfig.CODEC_PIPELINE_WRITER_VERSION; + + private final PinotDataBuffer _dataBuffer; + private final DataType _storedType; + private final int _numChunks; + private final int _numDocsPerChunk; + private final int _shift; // log2(numDocsPerChunk) for fast chunk id calc + private final int _totalDocs; + private final int _dataHeaderStart; + private final int _chunkCapacityBytes; + private final CodecPipelineExecutor _executor; + private final String _canonicalSpec; + /// Composed pipeline encoded-size bound for a full chunk, computed once at construction; every + /// chunk except a partial final chunk reuses it instead of re-walking the pipeline stages. + private final int _maxFullChunkEncodedSize; + + /// Returns whether the buffer has the explicit header discriminator for the codec-pipeline V7 + /// format. Keep this predicate limited to the marker: once recognized, the constructor must see + /// and reject every corrupt structural field instead of letting the factory fall back to the + /// legacy version-7 reader. + public static boolean hasCodecPipelineHeader(PinotDataBuffer dataBuffer) { + return dataBuffer.size() >= 2L * Integer.BYTES && dataBuffer.getInt(0) == VERSION + && dataBuffer.getInt(Integer.BYTES) == FixedByteChunkForwardIndexWriterV7.FORMAT_MAGIC; + } + + public FixedByteChunkSVForwardIndexReaderV7(PinotDataBuffer dataBuffer, DataType storedType) { + this(dataBuffer, storedType, -1); + } + + /// Creates a V7 reader and, when `expectedTotalDocs` is non-negative, verifies that the index + /// belongs to segment metadata with the same document count. The two-argument constructor keeps + /// standalone fixture and StarTree helper reads available when no segment metadata is present. + public FixedByteChunkSVForwardIndexReaderV7(PinotDataBuffer dataBuffer, DataType storedType, + int expectedTotalDocs) { + _dataBuffer = dataBuffer; + _storedType = storedType; + + long bufferSize = dataBuffer.size(); + if (bufferSize < FixedByteChunkForwardIndexWriterV7.FIXED_HEADER_BYTES) { + throw new IllegalArgumentException( + "V7 forward index is truncated: " + bufferSize + " bytes; minimum header is " + + FixedByteChunkForwardIndexWriterV7.FIXED_HEADER_BYTES + " bytes"); + } + + if (storedType != DataType.INT && storedType != DataType.LONG) { + throw new IllegalArgumentException( + "FixedByteChunkSVForwardIndexReaderV7 only supports INT and LONG, got: " + storedType); + } + + int offset = 0; + int version = dataBuffer.getInt(offset); + if (version != VERSION) { + throw new IllegalArgumentException("Expected version " + VERSION + " but got " + version); + } + if (!hasCodecPipelineHeader(dataBuffer)) { + throw new IllegalArgumentException( + "Version " + VERSION + " buffer does not contain a valid codec-pipeline header discriminator"); + } + offset += Integer.BYTES; + + int formatMagic = dataBuffer.getInt(offset); + offset += Integer.BYTES; + if (formatMagic != FixedByteChunkForwardIndexWriterV7.FORMAT_MAGIC) { + throw new IllegalArgumentException("Invalid codec-pipeline format magic: " + formatMagic); + } + + _numChunks = dataBuffer.getInt(offset); + offset += Integer.BYTES; + if (_numChunks < 0) { + throw new IllegalArgumentException("Invalid numChunks in forward index header: " + _numChunks); + } + + _numDocsPerChunk = dataBuffer.getInt(offset); + offset += Integer.BYTES; + if (_numDocsPerChunk <= 0 || (_numDocsPerChunk & (_numDocsPerChunk - 1)) != 0) { + throw new IllegalArgumentException( + "Invalid numDocsPerChunk in forward index header: " + _numDocsPerChunk + + ". Expected a positive power of two."); + } + _shift = Integer.numberOfTrailingZeros(_numDocsPerChunk); + + int sizeOfEntry = dataBuffer.getInt(offset); + offset += Integer.BYTES; + if (sizeOfEntry != storedType.size()) { + throw new IllegalArgumentException( + "Header sizeOfEntry=" + sizeOfEntry + " does not match storedType=" + storedType + + " (expected " + storedType.size() + " bytes). Written for a different data type?"); + } + long chunkCapacity = (long) _numDocsPerChunk * sizeOfEntry; + if (chunkCapacity > FixedByteChunkForwardIndexWriterV7.MAX_DECODED_CHUNK_SIZE_BYTES) { + throw new IllegalArgumentException( + "Decoded chunk capacity " + chunkCapacity + " bytes exceeds V7 limit " + + FixedByteChunkForwardIndexWriterV7.MAX_DECODED_CHUNK_SIZE_BYTES + ". Segment may be corrupt."); + } + _chunkCapacityBytes = (int) chunkCapacity; + + _totalDocs = dataBuffer.getInt(offset); + offset += Integer.BYTES; + if (_totalDocs < 0) { + throw new IllegalArgumentException("Invalid totalDocs in forward index header: " + _totalDocs); + } + if (expectedTotalDocs >= 0 && _totalDocs != expectedTotalDocs) { + throw new IllegalArgumentException( + "V7 forward index totalDocs=" + _totalDocs + " does not match segment metadata totalDocs=" + + expectedTotalDocs); + } + + // Validate numChunks/totalDocs/numDocsPerChunk are mutually consistent. A corrupt header + // with mismatched values would otherwise let getChunkOffset() read past the chunk-offset table. + int expectedNumChunks = (int) (((long) _totalDocs + _numDocsPerChunk - 1) / _numDocsPerChunk); + if (_numChunks != expectedNumChunks) { + throw new IllegalArgumentException( + "Inconsistent header: numChunks=" + _numChunks + " but totalDocs=" + _totalDocs + " / numDocsPerChunk=" + + _numDocsPerChunk + " => expected " + expectedNumChunks + ". Segment may be corrupt."); + } + + int specLength = dataBuffer.getInt(offset); + offset += Integer.BYTES; + if (specLength <= 0 || specLength > FixedByteChunkForwardIndexWriterV7.MAX_CODEC_SPEC_LENGTH_BYTES) { + // V7 segments always embed a non-empty canonical codec spec; zero or negative is corruption. + throw new IllegalArgumentException( + "Invalid specLength in forward index header: " + specLength + "; expected [1, " + + FixedByteChunkForwardIndexWriterV7.MAX_CODEC_SPEC_LENGTH_BYTES + "]. Segment may be corrupt."); + } + + _dataHeaderStart = dataBuffer.getInt(offset); + offset += Integer.BYTES; + + // Validate dataHeaderStart, specLength, and the chunk-offset table bounds before using them. + long expectedSpecEnd = (long) offset + specLength; + long chunkOffsetTableEnd = _dataHeaderStart + (long) _numChunks * Long.BYTES; + if (specLength > bufferSize || _dataHeaderStart < 0 || _dataHeaderStart > bufferSize + || expectedSpecEnd != _dataHeaderStart || chunkOffsetTableEnd > bufferSize) { + throw new IllegalArgumentException( + "Forward index header is corrupt: specLength=" + specLength + ", dataHeaderStart=" + _dataHeaderStart + + ", numChunks=" + _numChunks + ", bufferSize=" + bufferSize); + } + + // Read codec spec bytes + byte[] specBytes = new byte[specLength]; + dataBuffer.copyTo(offset, specBytes, 0, specLength); + _canonicalSpec = new String(specBytes, StandardCharsets.UTF_8); + + try { + _executor = CodecPipelineExecutor.create(_canonicalSpec, storedType); Review Comment: Correction to my earlier reply on this thread, which is now stale: the plan cache it described no longer exists. The cache lived in `CodecPipelineExecutor`, which moved to the codec-API PR (#19397) during the split, and the handler-side `_configuredCodecExecutors` map was dropped by a later follow-up commit here. So this thread was resolved against a state the code no longer had — apologies for that. What is in `8ebcf1fa8c` now: `ForwardIndexHandler` memoizes the canonical form of each configured spec for the handler's lifetime, keyed by (spec, stored type) rather than by column, so reconciliation resolves each spec once instead of building and discarding an executor per column on each of the two `computeOperations` passes. Keying on the stored type matters because a spec that canonicalizes for one type can be rejected for another; `computeIfAbsent` does not cache exceptions, so invalid specs still throw identically every time. The other half you asked for — caching plans across segments by canonical spec and stored type — is deliberately not here. Its natural home is `CodecPipelineExecutor`, and this PR does not touch the codec package per your request to keep codec changes separate. The reader side is where it would pay off most (each V7 segment reader reparses the same plan from its header at load), and note the executor holds mutable per-use scratch, so it needs a plan/state split rather than a static map of executors. Happy to take that as the next codec-module PR. _🤖 Addressed by [Claude Code](https://claude.com/claude-code)_ -- 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]
