This is an automated email from the ASF dual-hosted git repository. xiangfu0 pushed a commit to branch xiangfu0/codex/codec-handler-context-lifecycle in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 8910f232c69d93f782484a1b56cc0f1f12815d23 Author: Xiang Fu <[email protected]> AuthorDate: Mon Aug 17 16:07:23 2026 -0700 Fix forward index reader context lifecycle --- .../segment/index/loader/ForwardIndexHandler.java | 580 +++++++++++---------- .../loader/ForwardIndexHandlerContextTest.java | 88 ++++ 2 files changed, 382 insertions(+), 286 deletions(-) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java index b9d2e29dc94..9c53152cdee 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java @@ -656,7 +656,8 @@ public class ForwardIndexHandler extends BaseIndexHandler { private void forwardIndexRewriteHelper(String column, ColumnMetadata existingColumnMetadata, ForwardIndexReader<?> reader, ForwardIndexCreator creator, int numDocs, - @Nullable SegmentDictionaryCreator dictionaryCreator, @Nullable Dictionary dictionaryReader) { + @Nullable SegmentDictionaryCreator dictionaryCreator, @Nullable Dictionary dictionaryReader) + throws IOException { if (dictionaryReader == null && dictionaryCreator == null) { if (reader.isDictionaryEncoded()) { Preconditions.checkState(creator.isDictionaryEncoded(), "Cannot change dictionary based forward index to raw " @@ -684,294 +685,299 @@ public class ForwardIndexHandler extends BaseIndexHandler { ColumnMetadata columnMetadata, ForwardIndexReader<C> reader, ForwardIndexCreator creator, Dictionary dictionary) { DataType storedType = dictionary.getValueType().getStoredType(); - C readerContext = reader.createContext(); int numDocs = columnMetadata.getTotalDocs(); if (storedType.isFixedWidth()) { long numEntries = forwardIndexReadDictWriteDictHelper(reader, creator, numDocs); return numEntries * storedType.size(); } - long uncompressedValueSizeInBytes = 0; - if (reader.isSingleValue()) { - for (int docId = 0; docId < numDocs; docId++) { - int dictId = reader.getDictId(docId, readerContext); - creator.putDictId(dictId); - uncompressedValueSizeInBytes += dictionary.getValueSize(dictId); - } - } else { - for (int docId = 0; docId < numDocs; docId++) { - int[] dictIds = reader.getDictIdMV(docId, readerContext); - creator.putDictIdMV(dictIds); - for (int dictId : dictIds) { + try (C readerContext = reader.createContext()) { + long uncompressedValueSizeInBytes = 0; + if (reader.isSingleValue()) { + for (int docId = 0; docId < numDocs; docId++) { + int dictId = reader.getDictId(docId, readerContext); + creator.putDictId(dictId); uncompressedValueSizeInBytes += dictionary.getValueSize(dictId); } + } else { + for (int docId = 0; docId < numDocs; docId++) { + int[] dictIds = reader.getDictIdMV(docId, readerContext); + creator.putDictIdMV(dictIds); + for (int dictId : dictIds) { + uncompressedValueSizeInBytes += dictionary.getValueSize(dictId); + } + } } + return uncompressedValueSizeInBytes; } - return uncompressedValueSizeInBytes; } private static <C extends ForwardIndexReaderContext> long forwardIndexReadDictWriteDictHelper( ForwardIndexReader<C> reader, ForwardIndexCreator creator, int numDocs) { - C readerContext = reader.createContext(); - if (reader.isSingleValue()) { - for (int i = 0; i < numDocs; i++) { - creator.putDictId(reader.getDictId(i, readerContext)); - } - return numDocs; - } else { - long numEntries = 0; - for (int i = 0; i < numDocs; i++) { - int[] dictIds = reader.getDictIdMV(i, readerContext); - creator.putDictIdMV(dictIds); - numEntries += dictIds.length; + try (C readerContext = reader.createContext()) { + if (reader.isSingleValue()) { + for (int i = 0; i < numDocs; i++) { + creator.putDictId(reader.getDictId(i, readerContext)); + } + return numDocs; + } else { + long numEntries = 0; + for (int i = 0; i < numDocs; i++) { + int[] dictIds = reader.getDictIdMV(i, readerContext); + creator.putDictIdMV(dictIds); + numEntries += dictIds.length; + } + return numEntries; } - return numEntries; } } private <C extends ForwardIndexReaderContext> void forwardIndexReadRawWriteRawHelper(String column, ColumnMetadata existingColumnMetadata, ForwardIndexReader<C> reader, ForwardIndexCreator creator, int numDocs) { - C readerContext = reader.createContext(); - boolean isSVColumn = reader.isSingleValue(); - - switch (reader.getStoredType()) { - // JSON fields are either stored as string or bytes. No special handling is needed because we make this - // decision based on the storedType of the reader. - case INT: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - int val = reader.getInt(i, readerContext); - creator.putInt(val); - } else { - int[] ints = reader.getIntMV(i, readerContext); - creator.putIntMV(ints); + try (C readerContext = reader.createContext()) { + boolean isSVColumn = reader.isSingleValue(); + + switch (reader.getStoredType()) { + // JSON fields are either stored as string or bytes. No special handling is needed because we make this + // decision based on the storedType of the reader. + case INT: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + int val = reader.getInt(i, readerContext); + creator.putInt(val); + } else { + int[] ints = reader.getIntMV(i, readerContext); + creator.putIntMV(ints); + } } + break; } - break; - } - case LONG: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - long val = reader.getLong(i, readerContext); - creator.putLong(val); - } else { - long[] longs = reader.getLongMV(i, readerContext); - creator.putLongMV(longs); + case LONG: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + long val = reader.getLong(i, readerContext); + creator.putLong(val); + } else { + long[] longs = reader.getLongMV(i, readerContext); + creator.putLongMV(longs); + } } + break; } - break; - } - case FLOAT: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - float val = reader.getFloat(i, readerContext); - creator.putFloat(val); - } else { - float[] floats = reader.getFloatMV(i, readerContext); - creator.putFloatMV(floats); + case FLOAT: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + float val = reader.getFloat(i, readerContext); + creator.putFloat(val); + } else { + float[] floats = reader.getFloatMV(i, readerContext); + creator.putFloatMV(floats); + } } + break; } - break; - } - case DOUBLE: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - double val = reader.getDouble(i, readerContext); - creator.putDouble(val); - } else { - double[] doubles = reader.getDoubleMV(i, readerContext); - creator.putDoubleMV(doubles); + case DOUBLE: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + double val = reader.getDouble(i, readerContext); + creator.putDouble(val); + } else { + double[] doubles = reader.getDoubleMV(i, readerContext); + creator.putDoubleMV(doubles); + } } + break; } - break; - } - case BIG_DECIMAL: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - BigDecimal val = reader.getBigDecimal(i, readerContext); - creator.putBigDecimal(val); - } else { - BigDecimal[] bigDecimals = reader.getBigDecimalMV(i, readerContext); - creator.putBigDecimalMV(bigDecimals); + case BIG_DECIMAL: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + BigDecimal val = reader.getBigDecimal(i, readerContext); + creator.putBigDecimal(val); + } else { + BigDecimal[] bigDecimals = reader.getBigDecimalMV(i, readerContext); + creator.putBigDecimalMV(bigDecimals); + } } + break; } - break; - } - case STRING: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - String val = reader.getString(i, readerContext); - creator.putString(val); - } else { - String[] strings = reader.getStringMV(i, readerContext); - creator.putStringMV(strings); + case STRING: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + String val = reader.getString(i, readerContext); + creator.putString(val); + } else { + String[] strings = reader.getStringMV(i, readerContext); + creator.putStringMV(strings); + } } + break; } - break; - } - case BYTES: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - byte[] val = reader.getBytes(i, readerContext); - creator.putBytes(val); - } else { - byte[][] bytesArray = reader.getBytesMV(i, readerContext); - creator.putBytesMV(bytesArray); + case BYTES: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + byte[] val = reader.getBytes(i, readerContext); + creator.putBytes(val); + } else { + byte[][] bytesArray = reader.getBytesMV(i, readerContext); + creator.putBytesMV(bytesArray); + } } + break; } - break; - } - case MAP: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - byte[] val = reader.getBytes(i, readerContext); - creator.putBytes(val); - } else { - throw new IllegalStateException("Map is not supported for MV columns"); + case MAP: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + byte[] val = reader.getBytes(i, readerContext); + creator.putBytes(val); + } else { + throw new IllegalStateException("Map is not supported for MV columns"); + } } + break; } - break; + default: + throw new IllegalStateException("Unsupported storedType=" + reader.getStoredType() + " for column=" + column); } - default: - throw new IllegalStateException("Unsupported storedType=" + reader.getStoredType() + " for column=" + column); } } private <C extends ForwardIndexReaderContext> void forwardIndexReadDictWriteRawHelper(String column, ColumnMetadata existingColumnMetadata, ForwardIndexReader<C> reader, ForwardIndexCreator creator, int numDocs, Dictionary dictionaryReader) { - C readerContext = reader.createContext(); - boolean isSVColumn = reader.isSingleValue(); - DataType storedType = dictionaryReader.getValueType().getStoredType(); + try (C readerContext = reader.createContext()) { + boolean isSVColumn = reader.isSingleValue(); + DataType storedType = dictionaryReader.getValueType().getStoredType(); - switch (storedType) { - case INT: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - int dictId = reader.getDictId(i, readerContext); - int val = dictionaryReader.getIntValue(dictId); - creator.putInt(val); - } else { - int[] dictIds = reader.getDictIdMV(i, readerContext); - int[] ints = new int[dictIds.length]; - dictionaryReader.readIntValues(dictIds, dictIds.length, ints); - creator.putIntMV(ints); + switch (storedType) { + case INT: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + int dictId = reader.getDictId(i, readerContext); + int val = dictionaryReader.getIntValue(dictId); + creator.putInt(val); + } else { + int[] dictIds = reader.getDictIdMV(i, readerContext); + int[] ints = new int[dictIds.length]; + dictionaryReader.readIntValues(dictIds, dictIds.length, ints); + creator.putIntMV(ints); + } } + break; } - break; - } - case LONG: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - int dictId = reader.getDictId(i, readerContext); - long val = dictionaryReader.getLongValue(dictId); - creator.putLong(val); - } else { - int[] dictIds = reader.getDictIdMV(i, readerContext); - long[] longs = new long[dictIds.length]; - dictionaryReader.readLongValues(dictIds, dictIds.length, longs); - creator.putLongMV(longs); + case LONG: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + int dictId = reader.getDictId(i, readerContext); + long val = dictionaryReader.getLongValue(dictId); + creator.putLong(val); + } else { + int[] dictIds = reader.getDictIdMV(i, readerContext); + long[] longs = new long[dictIds.length]; + dictionaryReader.readLongValues(dictIds, dictIds.length, longs); + creator.putLongMV(longs); + } } + break; } - break; - } - case FLOAT: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - int dictId = reader.getDictId(i, readerContext); - float val = dictionaryReader.getFloatValue(dictId); - creator.putFloat(val); - } else { - int[] dictIds = reader.getDictIdMV(i, readerContext); - float[] floats = new float[dictIds.length]; - dictionaryReader.readFloatValues(dictIds, dictIds.length, floats); - creator.putFloatMV(floats); + case FLOAT: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + int dictId = reader.getDictId(i, readerContext); + float val = dictionaryReader.getFloatValue(dictId); + creator.putFloat(val); + } else { + int[] dictIds = reader.getDictIdMV(i, readerContext); + float[] floats = new float[dictIds.length]; + dictionaryReader.readFloatValues(dictIds, dictIds.length, floats); + creator.putFloatMV(floats); + } } + break; } - break; - } - case DOUBLE: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - int dictId = reader.getDictId(i, readerContext); - double val = dictionaryReader.getDoubleValue(dictId); - creator.putDouble(val); - } else { - int[] dictIds = reader.getDictIdMV(i, readerContext); - double[] doubles = new double[dictIds.length]; - dictionaryReader.readDoubleValues(dictIds, dictIds.length, doubles); - creator.putDoubleMV(doubles); + case DOUBLE: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + int dictId = reader.getDictId(i, readerContext); + double val = dictionaryReader.getDoubleValue(dictId); + creator.putDouble(val); + } else { + int[] dictIds = reader.getDictIdMV(i, readerContext); + double[] doubles = new double[dictIds.length]; + dictionaryReader.readDoubleValues(dictIds, dictIds.length, doubles); + creator.putDoubleMV(doubles); + } } + break; } - break; - } - case BIG_DECIMAL: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - int dictId = reader.getDictId(i, readerContext); - BigDecimal val = dictionaryReader.getBigDecimalValue(dictId); - creator.putBigDecimal(val); - } else { - int[] dictIds = reader.getDictIdMV(i, readerContext); - BigDecimal[] bigDecimals = new BigDecimal[dictIds.length]; - dictionaryReader.readBigDecimalValues(dictIds, dictIds.length, bigDecimals); - creator.putBigDecimalMV(bigDecimals); + case BIG_DECIMAL: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + int dictId = reader.getDictId(i, readerContext); + BigDecimal val = dictionaryReader.getBigDecimalValue(dictId); + creator.putBigDecimal(val); + } else { + int[] dictIds = reader.getDictIdMV(i, readerContext); + BigDecimal[] bigDecimals = new BigDecimal[dictIds.length]; + dictionaryReader.readBigDecimalValues(dictIds, dictIds.length, bigDecimals); + creator.putBigDecimalMV(bigDecimals); + } } + break; } - break; - } - case STRING: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - int dictId = reader.getDictId(i, readerContext); - String val = dictionaryReader.getStringValue(dictId); - creator.putString(val); - } else { - int[] dictIds = reader.getDictIdMV(i, readerContext); - String[] strings = new String[dictIds.length]; - dictionaryReader.readStringValues(dictIds, dictIds.length, strings); - creator.putStringMV(strings); + case STRING: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + int dictId = reader.getDictId(i, readerContext); + String val = dictionaryReader.getStringValue(dictId); + creator.putString(val); + } else { + int[] dictIds = reader.getDictIdMV(i, readerContext); + String[] strings = new String[dictIds.length]; + dictionaryReader.readStringValues(dictIds, dictIds.length, strings); + creator.putStringMV(strings); + } } + break; } - break; - } - case BYTES: { - for (int i = 0; i < numDocs; i++) { - if (isSVColumn) { - int dictId = reader.getDictId(i, readerContext); - byte[] val = dictionaryReader.getBytesValue(dictId); - creator.putBytes(val); - } else { - int[] dictIds = reader.getDictIdMV(i, readerContext); - byte[][] bytes = new byte[dictIds.length][]; - dictionaryReader.readBytesValues(dictIds, dictIds.length, bytes); - creator.putBytesMV(bytes); + case BYTES: { + for (int i = 0; i < numDocs; i++) { + if (isSVColumn) { + int dictId = reader.getDictId(i, readerContext); + byte[] val = dictionaryReader.getBytesValue(dictId); + creator.putBytes(val); + } else { + int[] dictIds = reader.getDictIdMV(i, readerContext); + byte[][] bytes = new byte[dictIds.length][]; + dictionaryReader.readBytesValues(dictIds, dictIds.length, bytes); + creator.putBytesMV(bytes); + } } + break; } - break; + default: + throw new IllegalStateException("Unsupported storedType=" + storedType + " for column=" + column); } - default: - throw new IllegalStateException("Unsupported storedType=" + storedType + " for column=" + column); } } private void forwardIndexReadRawWriteDictHelper(String column, ColumnMetadata existingColumnMetadata, ForwardIndexReader<?> reader, ForwardIndexCreator creator, int numDocs, - SegmentDictionaryCreator dictionaryCreator) { + SegmentDictionaryCreator dictionaryCreator) + throws IOException { boolean isSVColumn = reader.isSingleValue(); int maxNumValuesPerEntry = existingColumnMetadata.getMaxNumberOfMultiValues(); - PinotSegmentColumnReader columnReader = - new PinotSegmentColumnReader(column, reader, null, null, maxNumValuesPerEntry); - - for (int i = 0; i < numDocs; i++) { - Object obj = columnReader.getValue(i); + try (PinotSegmentColumnReader columnReader = + new PinotSegmentColumnReader(column, reader, null, null, maxNumValuesPerEntry)) { + for (int i = 0; i < numDocs; i++) { + Object obj = columnReader.getValue(i); - if (isSVColumn) { - int dictId = dictionaryCreator.indexOfSV(obj); - creator.putDictId(dictId); - } else { - int[] dictIds = dictionaryCreator.indexOfMV(obj); - creator.putDictIdMV(dictIds); + if (isSVColumn) { + int dictId = dictionaryCreator.indexOfSV(obj); + creator.putDictId(dictId); + } else { + int[] dictIds = dictionaryCreator.indexOfMV(obj); + creator.putDictIdMV(dictIds); + } } } } @@ -1363,68 +1369,69 @@ public class ForwardIndexHandler extends BaseIndexHandler { return collectShapeStatsWithCollector(column, columnMetadata, forwardIndex); } - ShapeStats stats = new ShapeStats(); - C context = forwardIndex.createContext(); - int numDocs = columnMetadata.getTotalDocs(); - boolean singleValue = columnMetadata.isSingleValue(); - int maxNumMultiValues = Math.max(columnMetadata.getMaxNumberOfMultiValues(), 1); - String[] stringBuffer = !singleValue && storedType == DataType.STRING - ? new String[maxNumMultiValues] : null; - byte[][] bytesBuffer = !singleValue && storedType == DataType.BYTES - ? new byte[maxNumMultiValues][] : null; - BigDecimal[] bigDecimalBuffer = !singleValue && storedType == DataType.BIG_DECIMAL - ? new BigDecimal[maxNumMultiValues] : null; - for (int docId = 0; docId < numDocs; docId++) { - int rowLength = 0; - switch (storedType) { - case STRING: - if (singleValue) { - String value = forwardIndex.getString(docId, context); - rowLength = Utf8Utils.encodedLengthWithReplacement(value); - stats.addElement(rowLength, isAscii(value)); - } else { - int numValues = forwardIndex.getStringMV(docId, stringBuffer, context); - for (int i = 0; i < numValues; i++) { - int valueLength = Utf8Utils.encodedLengthWithReplacement(stringBuffer[i]); - rowLength += valueLength; - stats.addElement(valueLength, isAscii(stringBuffer[i])); + try (C context = forwardIndex.createContext()) { + ShapeStats stats = new ShapeStats(); + int numDocs = columnMetadata.getTotalDocs(); + boolean singleValue = columnMetadata.isSingleValue(); + int maxNumMultiValues = Math.max(columnMetadata.getMaxNumberOfMultiValues(), 1); + String[] stringBuffer = !singleValue && storedType == DataType.STRING + ? new String[maxNumMultiValues] : null; + byte[][] bytesBuffer = !singleValue && storedType == DataType.BYTES + ? new byte[maxNumMultiValues][] : null; + BigDecimal[] bigDecimalBuffer = !singleValue && storedType == DataType.BIG_DECIMAL + ? new BigDecimal[maxNumMultiValues] : null; + for (int docId = 0; docId < numDocs; docId++) { + int rowLength = 0; + switch (storedType) { + case STRING: + if (singleValue) { + String value = forwardIndex.getString(docId, context); + rowLength = Utf8Utils.encodedLengthWithReplacement(value); + stats.addElement(rowLength, isAscii(value)); + } else { + int numValues = forwardIndex.getStringMV(docId, stringBuffer, context); + for (int i = 0; i < numValues; i++) { + int valueLength = Utf8Utils.encodedLengthWithReplacement(stringBuffer[i]); + rowLength += valueLength; + stats.addElement(valueLength, isAscii(stringBuffer[i])); + } } - } - break; - case BYTES: - if (singleValue) { - rowLength = forwardIndex.getBytes(docId, context).length; - stats.addElement(rowLength, true); - } else { - int numValues = forwardIndex.getBytesMV(docId, bytesBuffer, context); - for (int i = 0; i < numValues; i++) { - rowLength += bytesBuffer[i].length; - stats.addElement(bytesBuffer[i].length, true); + break; + case BYTES: + if (singleValue) { + rowLength = forwardIndex.getBytes(docId, context).length; + stats.addElement(rowLength, true); + } else { + int numValues = forwardIndex.getBytesMV(docId, bytesBuffer, context); + for (int i = 0; i < numValues; i++) { + rowLength += bytesBuffer[i].length; + stats.addElement(bytesBuffer[i].length, true); + } } - } - break; - case BIG_DECIMAL: - if (singleValue) { - rowLength = BigDecimalUtils.byteSize(forwardIndex.getBigDecimal(docId, context)); - stats.addElement(rowLength, true); - } else { - int numValues = forwardIndex.getBigDecimalMV(docId, bigDecimalBuffer, context); - for (int i = 0; i < numValues; i++) { - int valueLength = BigDecimalUtils.byteSize(bigDecimalBuffer[i]); - rowLength += valueLength; - stats.addElement(valueLength, true); + break; + case BIG_DECIMAL: + if (singleValue) { + rowLength = BigDecimalUtils.byteSize(forwardIndex.getBigDecimal(docId, context)); + stats.addElement(rowLength, true); + } else { + int numValues = forwardIndex.getBigDecimalMV(docId, bigDecimalBuffer, context); + for (int i = 0; i < numValues; i++) { + int valueLength = BigDecimalUtils.byteSize(bigDecimalBuffer[i]); + rowLength += valueLength; + stats.addElement(valueLength, true); + } } - } - break; - default: - throw new IllegalStateException("Unsupported variable-width type: " + storedType); - } - if (!singleValue) { - stats._maxRowLengthInBytes = Math.max(stats._maxRowLengthInBytes, rowLength); + break; + default: + throw new IllegalStateException("Unsupported variable-width type: " + storedType); + } + if (!singleValue) { + stats._maxRowLengthInBytes = Math.max(stats._maxRowLengthInBytes, rowLength); + } } + stats.finish(); + return stats; } - stats.finish(); - return stats; } private static <C extends ForwardIndexReaderContext> ShapeStats collectDictionaryShapeStats( @@ -1436,14 +1443,15 @@ public class ForwardIndexHandler extends BaseIndexHandler { } if (!columnMetadata.isSingleValue()) { int[] dictIdBuffer = new int[Math.max(columnMetadata.getMaxNumberOfMultiValues(), 1)]; - C context = forwardIndex.createContext(); - for (int docId = 0; docId < columnMetadata.getTotalDocs(); docId++) { - int numValues = forwardIndex.getDictIdMV(docId, dictIdBuffer, context); - int rowLength = 0; - for (int i = 0; i < numValues; i++) { - rowLength += dictionary.getValueSize(dictIdBuffer[i]); + try (C context = forwardIndex.createContext()) { + for (int docId = 0; docId < columnMetadata.getTotalDocs(); docId++) { + int numValues = forwardIndex.getDictIdMV(docId, dictIdBuffer, context); + int rowLength = 0; + for (int i = 0; i < numValues; i++) { + rowLength += dictionary.getValueSize(dictIdBuffer[i]); + } + stats._maxRowLengthInBytes = Math.max(stats._maxRowLengthInBytes, rowLength); } - stats._maxRowLengthInBytes = Math.max(stats._maxRowLengthInBytes, rowLength); } } stats.finish(); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerContextTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerContextTest.java new file mode 100644 index 00000000000..aad41af8707 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerContextTest.java @@ -0,0 +1,88 @@ +/** + * 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.loader; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import org.apache.pinot.segment.spi.index.creator.ForwardIndexCreator; +import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; +import org.apache.pinot.segment.spi.index.reader.ForwardIndexReaderContext; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.expectThrows; + + +/// Unit tests for the resource lifecycle of [ForwardIndexReaderContext] instances created inside +/// [ForwardIndexHandler] rewrite helpers. Verifies that a helper closes the context it creates on both the success +/// and the failure path, without closing the reader or creator it does not own. Single-threaded; no shared state. +public class ForwardIndexHandlerContextTest { + + @Test + public void testRewriteHelperClosesOwnedContextOnSuccess() throws Exception { + ForwardIndexReaderContext context = mock(ForwardIndexReaderContext.class); + ForwardIndexReader<ForwardIndexReaderContext> reader = mockReader(context); + ForwardIndexCreator creator = mock(ForwardIndexCreator.class); + when(reader.getDictId(0, context)).thenReturn(7); + + assertEquals(invokeDictRewriteHelper(reader, creator), 1L); + + verify(context).close(); + verify(reader, never()).close(); + verify(creator, never()).close(); + } + + @Test + public void testRewriteHelperClosesOwnedContextOnFailure() throws Exception { + ForwardIndexReaderContext context = mock(ForwardIndexReaderContext.class); + ForwardIndexReader<ForwardIndexReaderContext> reader = mockReader(context); + ForwardIndexCreator creator = mock(ForwardIndexCreator.class); + IllegalStateException failure = new IllegalStateException("read failed"); + when(reader.getDictId(0, context)).thenThrow(failure); + + InvocationTargetException exception = expectThrows(InvocationTargetException.class, + () -> invokeDictRewriteHelper(reader, creator)); + assertSame(exception.getCause(), failure); + + verify(context).close(); + verify(reader, never()).close(); + verify(creator, never()).close(); + } + + @SuppressWarnings("unchecked") + private static ForwardIndexReader<ForwardIndexReaderContext> mockReader(ForwardIndexReaderContext context) { + ForwardIndexReader<ForwardIndexReaderContext> reader = mock(ForwardIndexReader.class); + when(reader.createContext()).thenReturn(context); + when(reader.isSingleValue()).thenReturn(true); + return reader; + } + + private static long invokeDictRewriteHelper(ForwardIndexReader<ForwardIndexReaderContext> reader, + ForwardIndexCreator creator) throws Exception { + Method method = ForwardIndexHandler.class.getDeclaredMethod("forwardIndexReadDictWriteDictHelper", + ForwardIndexReader.class, ForwardIndexCreator.class, int.class); + method.setAccessible(true); + return (long) method.invoke(null, reader, creator, 1); + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
