diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainer.java index d1a0b3a660f2..9350ebe35185 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainer.java @@ -18,10 +18,7 @@ */ package org.apache.pinot.segment.local.segment.index.column; -import it.unimi.dsi.fastutil.shorts.ShortArrayList; -import java.io.Closeable; import java.io.IOException; -import java.util.ArrayList; import java.util.List; import java.util.Set; import javax.annotation.Nullable; @@ -41,14 +38,33 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import static com.google.common.base.Preconditions.checkState; + +/// Index readers of one physical column of an immutable segment. +/// +/// Readers are keyed by the numeric index id assigned by [IndexService#getNumericId(IndexType)] and stored as a +/// presence bit mask plus a dense array holding only the readers that exist, ordered by id. A lookup is a shift, a +/// mask and a popcount, so it stays O(1) on the query path while the per-column footprint is exactly one array slot +/// per present reader. A server holding tens of thousands of wide segments creates one of these per (segment, column), +/// which is why the layout matters: the mask replaces a nested map object and a reader array spanning the whole +/// numeric-id range of the present readers. +/// +/// Thread safety: immutable after construction except for the multi-column text reader reference, which is set +/// once during segment load and cleared on [#close()]. public final class PhysicalColumnIndexContainer implements ColumnIndexContainer { private static final Logger LOGGER = LoggerFactory.getLogger(PhysicalColumnIndexContainer.class); private static final Set FORWARD_INDEX_ONLY_TYPES = Set.of(StandardIndexes.FORWARD_ID, StandardIndexes.DICTIONARY_ID, StandardIndexes.NULL_VALUE_VECTOR_ID); - - private final IndexTypeMap _indexTypeMap; + private static final IndexReader[] EMPTY_READERS = new IndexReader[0]; + + // Bit i is set when the index type with numeric id i has a reader in this column. Numeric ids are validated to fit + // in the mask at construction time. + private final long _presentMask; + // Readers ordered by numeric index id, holding only the present ones: the reader for id i sits at the number of + // bits set in _presentMask below bit i. + private final IndexReader[] _readers; @Nullable private final VectorIndexConfig _vectorIndexConfig; @@ -67,12 +83,18 @@ public PhysicalColumnIndexContainer(SegmentDirectory.Reader segmentReader, Colum } _vectorIndexConfig = fieldIndexConfigs.getConfig(StandardIndexes.vector()); - ArrayList indexTypes = new ArrayList<>(); - ArrayList readers = new ArrayList<>(); + IndexService indexService = IndexService.getInstance(); + List> allIndexes = indexService.getAllIndexes(); + int numIndexTypes = allIndexes.size(); + checkState(numIndexTypes <= Long.SIZE, + "Cannot track %s index types in a %s-bit presence mask, column: %s", numIndexTypes, Long.SIZE, columnName); + // Scratch array indexed by numeric id; compacted into the exactly-sized _readers below. + IndexReader[] readersById = new IndexReader[numIndexTypes]; + long presentMask = 0L; boolean forwardIndexOnly = indexLoadingConfig.isForwardIndexOnly(); try { - for (IndexType indexType : IndexService.getInstance().getAllIndexes()) { + for (IndexType indexType : allIndexes) { if (forwardIndexOnly && !FORWARD_INDEX_ONLY_TYPES.contains(indexType.getId())) { continue; } @@ -81,8 +103,9 @@ public PhysicalColumnIndexContainer(SegmentDirectory.Reader segmentReader, Colum try { IndexReader reader = readerProvider.createIndexReader(segmentReader, fieldIndexConfigs, metadata); if (reader != null) { - indexTypes.add(indexType); - readers.add(reader); + short indexId = indexService.getNumericId(indexType); + readersById[indexId] = reader; + presentMask |= 1L << indexId; } } catch (IndexReaderConstraintException ex) { LOGGER.warn("Constraint violation when indexing {} with {} index", columnName, indexType, ex); @@ -90,23 +113,42 @@ public PhysicalColumnIndexContainer(SegmentDirectory.Reader segmentReader, Colum } } } catch (Throwable t) { - for (IndexReader reader : readers) { - try { - reader.close(); - } catch (Throwable ct) { - LOGGER.warn("Can't close reader on init error, column: " + columnName + " reader: " + reader.getClass(), ct); + for (IndexReader reader : readersById) { + if (reader != null) { + try { + reader.close(); + } catch (Throwable ct) { + LOGGER.warn("Can't close reader on init error, column: " + columnName + " reader: " + reader.getClass(), + ct); + } } } throw t; } - _indexTypeMap = IndexTypeMap.get(indexTypes, readers); + _presentMask = presentMask; + int numReaders = Long.bitCount(presentMask); + if (numReaders == 0) { + _readers = EMPTY_READERS; + } else { + _readers = new IndexReader[numReaders]; + int pos = 0; + for (IndexReader reader : readersById) { + if (reader != null) { + _readers[pos++] = reader; + } + } + } } @Nullable @Override public > I getIndex(T indexType) { - return _indexTypeMap.getIndex(indexType); + short indexId = IndexService.getInstance().getNumericId(indexType); + if (((_presentMask >>> indexId) & 1L) == 0) { + return null; + } + return (I) _readers[Long.bitCount(_presentMask & ((1L << indexId) - 1))]; } @Nullable @@ -119,7 +161,9 @@ public VectorIndexConfig getVectorIndexConfig() { public void close() throws IOException { // TODO (index-spi): Verify that readers can be closed in any order - _indexTypeMap.close(); + for (IndexReader reader : _readers) { + reader.close(); + } // This reader is closed on segment destroy() _multiColTextReader = null; @@ -133,70 +177,4 @@ public void setMultiColumnTextIndex( MultiColumnLuceneTextIndexReader multiColTextReader) { _multiColTextReader = multiColTextReader; } - - static class IndexTypeMap implements Closeable { - private static final IndexReader[] EMPTY_READERS = new IndexReader[0]; - - public static final IndexTypeMap EMPTY = new IndexTypeMap((short) 0, EMPTY_READERS); - - private final short _shift; - //stores index readers ordered by index id, shifted by _shift to conserve memory - private final IndexReader[] _readers; - - private IndexTypeMap(short shift, IndexReader[] readers) { - _shift = shift; - _readers = readers; - } - - static IndexTypeMap get(List indexTypes, List readers) { - if (indexTypes.isEmpty()) { - return EMPTY; - } - - short min = Short.MAX_VALUE; - int max = -1; - - ShortArrayList indexIds = new ShortArrayList(indexTypes.size()); - IndexService indexService = IndexService.getInstance(); - - for (IndexType indexType : indexTypes) { - short indexId = indexService.getNumericId(indexType); - indexIds.add(indexId); - if (indexId < min) { - min = indexId; - } - if (indexId > max) { - max = indexId; - } - } - - short shift = min; - int size = max - min + 1; - IndexReader[] indexReaders = new IndexReader[size]; - for (int i = 0, n = indexIds.size(); i < n; i++) { - short indexId = indexIds.getShort(i); - indexReaders[indexId - shift] = readers.get(i); - } - return new IndexTypeMap(shift, indexReaders); - } - - @Nullable - public > I getIndex(T indexType) { - short indexId = IndexService.getInstance().getNumericId(indexType); - if (indexId >= _shift && indexId < _shift + _readers.length) { - return (I) _readers[indexId - _shift]; - } - return null; - } - - @Override - public void close() - throws IOException { - for (IndexReader index : _readers) { - if (index != null) { - index.close(); - } - } - } - } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainerTest.java index 6e2dc71b9a27..46081ab32f14 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainerTest.java @@ -20,18 +20,39 @@ import com.fasterxml.jackson.databind.node.ObjectNode; import java.io.File; +import java.io.IOException; import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; import org.apache.commons.io.FileUtils; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.json.JsonIndexType; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.local.segment.index.readers.text.MultiColumnLuceneTextIndexReader; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.creator.IndexCreationContext; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.index.FieldIndexConfigs; +import org.apache.pinot.segment.spi.index.IndexCreator; +import org.apache.pinot.segment.spi.index.IndexHandler; +import org.apache.pinot.segment.spi.index.IndexPlugin; +import org.apache.pinot.segment.spi.index.IndexReader; +import org.apache.pinot.segment.spi.index.IndexReaderConstraintException; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; +import org.apache.pinot.segment.spi.index.IndexService; +import org.apache.pinot.segment.spi.index.IndexType; import org.apache.pinot.segment.spi.index.StandardIndexes; -import org.apache.pinot.spi.config.table.FieldConfig.Builder; -import org.apache.pinot.spi.config.table.FieldConfig.IndexType; +import org.apache.pinot.segment.spi.index.creator.VectorIndexConfig; +import org.apache.pinot.segment.spi.store.SegmentDirectory; +import org.apache.pinot.spi.config.table.FieldConfig; +import org.apache.pinot.spi.config.table.IndexConfig; import org.apache.pinot.spi.config.table.JsonIndexConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; @@ -41,9 +62,22 @@ import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.spi.utils.ReadMode; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +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.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; public class PhysicalColumnIndexContainerTest { @@ -79,22 +113,48 @@ public class PhysicalColumnIndexContainerTest { TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE) .setTableName(RAW_TABLE_NAME) - .addFieldConfig(new Builder(STR_COL) + .addFieldConfig(new FieldConfig.Builder(STR_COL) .withIndexes(indexes) .build()) - .addFieldConfig(new Builder(INT_COL) - .withIndexTypes(List.of(IndexType.RANGE)) + .addFieldConfig(new FieldConfig.Builder(INT_COL) + .withIndexTypes(List.of(FieldConfig.IndexType.RANGE)) .build()) - .addFieldConfig(new Builder(LONG_COL) - .withIndexTypes(List.of(IndexType.RANGE)) + .addFieldConfig(new FieldConfig.Builder(LONG_COL) + .withIndexTypes(List.of(FieldConfig.IndexType.RANGE)) .build()) - .addFieldConfig(new Builder(FLOAT_COL) - .withIndexTypes(List.of(IndexType.RANGE, IndexType.SORTED)) + .addFieldConfig(new FieldConfig.Builder(FLOAT_COL) + .withIndexTypes(List.of(FieldConfig.IndexType.RANGE, FieldConfig.IndexType.SORTED)) .build()) .setRangeIndexColumns(List.of(LONG_COL, FLOAT_COL)) .build(); } + private static final String COLUMN = "column"; + + // Every standard index id. IndexService assigns numeric ids by sorted id, so the order here is the numeric-id order: + // bloom_filter = 0, dictionary = 1, forward_index = 2, ..., nullvalue_vector = 8, ..., vector_index = 12. + private static final List STANDARD_IDS = + List.of(StandardIndexes.BLOOM_FILTER_ID, StandardIndexes.DICTIONARY_ID, StandardIndexes.FORWARD_ID, + StandardIndexes.FST_ID, StandardIndexes.H3_ID, StandardIndexes.IFST_ID, StandardIndexes.INVERTED_ID, + StandardIndexes.JSON_ID, StandardIndexes.NULL_VALUE_VECTOR_ID, StandardIndexes.OPEN_STRUCT_ID, + StandardIndexes.RANGE_ID, StandardIndexes.TEXT_ID, StandardIndexes.VECTOR_ID); + + private IndexService _originalIndexService; + private final Map _stubTypes = new HashMap<>(); + private final Map _stubReaders = new HashMap<>(); + + @BeforeClass + public void saveIndexService() { + _originalIndexService = IndexService.getInstance(); + } + + @AfterMethod + public void restoreIndexService() { + IndexService.setInstance(_originalIndexService); + _stubTypes.clear(); + _stubReaders.clear(); + } + @Test public void testCreateSegmentAndCheckColumnIndexes() throws Exception { @@ -129,21 +189,355 @@ public void testCreateSegmentAndCheckColumnIndexes() assertNotNull(segment.getIndex(STR_COL, StandardIndexes.json())); assertNotNull(segment.getIndex(STR_COL, StandardIndexes.dictionary())); assertNotNull(segment.getIndex(STR_COL, StandardIndexes.forward())); + assertNull(segment.getIndex(STR_COL, StandardIndexes.range())); assertNotNull(segment.getIndex(FLOAT_COL, StandardIndexes.dictionary())); assertNotNull(segment.getIndex(FLOAT_COL, StandardIndexes.forward())); assertNotNull(segment.getIndex(FLOAT_COL, StandardIndexes.range())); + assertNull(segment.getIndex(FLOAT_COL, StandardIndexes.json())); assertNotNull(segment.getIndex(DOUBLE_COL, StandardIndexes.dictionary())); assertNotNull(segment.getIndex(DOUBLE_COL, StandardIndexes.forward())); + assertNull(segment.getIndex(DOUBLE_COL, StandardIndexes.range())); assertNotNull(segment.getIndex(LONG_COL, StandardIndexes.dictionary())); assertNotNull(segment.getIndex(LONG_COL, StandardIndexes.forward())); assertNotNull(segment.getIndex(LONG_COL, StandardIndexes.range())); + + // Index types that no column of this segment has resolve to null for every column. Every column of this + // fixture is sorted, so the inverted index resolves to the sorted reader rather than null. + for (String column : SCHEMA.getColumnNames()) { + assertNotNull(segment.getIndex(column, StandardIndexes.inverted())); + assertNull(segment.getIndex(column, StandardIndexes.bloomFilter())); + assertNull(segment.getIndex(column, StandardIndexes.nullValueVector())); + assertNull(segment.getIndex(column, StandardIndexes.text())); + assertNull(segment.getIndex(column, StandardIndexes.vector())); + } } finally { if (segment != null) { segment.destroy(); } } } + + @Test + public void testForwardOnlyColumn() + throws IOException { + installStandardStubs(); + PhysicalColumnIndexContainer container = newContainer(Set.of(StandardIndexes.FORWARD_ID), false); + assertPresentExactly(container, Set.of(StandardIndexes.FORWARD_ID)); + } + + @Test + public void testForwardAndNullValueVector() + throws IOException { + // Numeric ids 2 and 8: the span array of the old layout would have held 7 slots for these 2 readers. + installStandardStubs(); + Set present = Set.of(StandardIndexes.FORWARD_ID, StandardIndexes.NULL_VALUE_VECTOR_ID); + assertPresentExactly(newContainer(present, false), present); + } + + @Test + public void testLowestAndHighIds() + throws IOException { + // bloom_filter is numeric id 0 and vector_index is the highest standard id. + installStandardStubs(); + Set present = + Set.of(StandardIndexes.BLOOM_FILTER_ID, StandardIndexes.DICTIONARY_ID, StandardIndexes.FORWARD_ID, + StandardIndexes.JSON_ID, StandardIndexes.RANGE_ID, StandardIndexes.VECTOR_ID); + assertPresentExactly(newContainer(present, false), present); + } + + @Test + public void testSparseNonAdjacentIds() + throws IOException { + installStandardStubs(); + Set present = Set.of(StandardIndexes.BLOOM_FILTER_ID, StandardIndexes.JSON_ID, StandardIndexes.VECTOR_ID); + assertPresentExactly(newContainer(present, false), present); + } + + @Test + public void testAllIndexesPresent() + throws IOException { + installStandardStubs(); + assertPresentExactly(newContainer(new HashSet<>(STANDARD_IDS), false), new HashSet<>(STANDARD_IDS)); + } + + @Test + public void testNoIndexes() + throws IOException { + installStandardStubs(); + PhysicalColumnIndexContainer container = newContainer(Set.of(), false); + assertPresentExactly(container, Set.of()); + container.close(); + for (IndexReader reader : _stubReaders.values()) { + verify(reader, never()).close(); + } + } + + @Test + public void testForwardIndexOnlyFiltering() + throws IOException { + installStandardStubs(); + SegmentDirectory.Reader segmentReader = mockSegmentReader(new HashSet<>(STANDARD_IDS)); + PhysicalColumnIndexContainer container = newContainer(segmentReader, true); + Set expected = + Set.of(StandardIndexes.FORWARD_ID, StandardIndexes.DICTIONARY_ID, StandardIndexes.NULL_VALUE_VECTOR_ID); + assertPresentExactly(container, expected); + // Filtered types are skipped before the segment reader is consulted. + for (String id : STANDARD_IDS) { + if (!expected.contains(id)) { + verify(segmentReader, never()).hasIndexFor(COLUMN, _stubTypes.get(id)); + } + } + } + + @Test + public void testFactoryReturningNullIsAbsent() + throws IOException { + installStandardStubs(); + _stubTypes.get(StandardIndexes.DICTIONARY_ID).setReaderFactory((reader, configs, metadata) -> null); + Set present = + Set.of(StandardIndexes.DICTIONARY_ID, StandardIndexes.FORWARD_ID, StandardIndexes.RANGE_ID); + assertPresentExactly(newContainer(present, false), + Set.of(StandardIndexes.FORWARD_ID, StandardIndexes.RANGE_ID)); + } + + @Test + public void testConstraintViolationSkipsIndex() + throws IOException { + installStandardStubs(); + StubIndexType rangeType = _stubTypes.get(StandardIndexes.RANGE_ID); + rangeType.setReaderFactory((reader, configs, metadata) -> { + throw new IndexReaderConstraintException(COLUMN, rangeType, "no dictionary"); + }); + Set present = + Set.of(StandardIndexes.DICTIONARY_ID, StandardIndexes.FORWARD_ID, StandardIndexes.RANGE_ID, + StandardIndexes.TEXT_ID); + assertPresentExactly(newContainer(present, false), + Set.of(StandardIndexes.DICTIONARY_ID, StandardIndexes.FORWARD_ID, StandardIndexes.TEXT_ID)); + } + + @Test + public void testCloseClosesEveryReaderOnce() + throws IOException { + installStandardStubs(); + Set present = + Set.of(StandardIndexes.BLOOM_FILTER_ID, StandardIndexes.DICTIONARY_ID, StandardIndexes.FORWARD_ID, + StandardIndexes.NULL_VALUE_VECTOR_ID, StandardIndexes.VECTOR_ID); + PhysicalColumnIndexContainer container = newContainer(present, false); + MultiColumnLuceneTextIndexReader multiColTextReader = mock(MultiColumnLuceneTextIndexReader.class); + container.setMultiColumnTextIndex(multiColTextReader); + assertSame(container.getMultiColumnTextIndex(), multiColTextReader); + + container.close(); + + for (String id : STANDARD_IDS) { + if (present.contains(id)) { + verify(_stubReaders.get(id)).close(); + } else { + verify(_stubReaders.get(id), never()).close(); + } + } + // The multi-column text reader is shared across columns and closed by the segment, not the container. + verify(multiColTextReader, never()).close(); + assertNull(container.getMultiColumnTextIndex()); + } + + @Test + public void testInitFailureClosesCreatedReaders() + throws IOException { + installStandardStubs(); + IOException failure = new IOException("cannot open json index"); + _stubTypes.get(StandardIndexes.JSON_ID).setReaderFactory((reader, configs, metadata) -> { + throw failure; + }); + Set present = + Set.of(StandardIndexes.BLOOM_FILTER_ID, StandardIndexes.DICTIONARY_ID, StandardIndexes.FORWARD_ID, + StandardIndexes.JSON_ID, StandardIndexes.RANGE_ID); + SegmentDirectory.Reader segmentReader = mockSegmentReader(present); + + IOException thrown = expectThrows(IOException.class, () -> newContainer(segmentReader, false)); + assertSame(thrown, failure); + + // Readers created before the failure (lower numeric ids) are closed; the rest were never created. + verify(_stubReaders.get(StandardIndexes.BLOOM_FILTER_ID)).close(); + verify(_stubReaders.get(StandardIndexes.DICTIONARY_ID)).close(); + verify(_stubReaders.get(StandardIndexes.FORWARD_ID)).close(); + verify(segmentReader, never()).hasIndexFor(COLUMN, _stubTypes.get(StandardIndexes.RANGE_ID)); + verify(_stubReaders.get(StandardIndexes.RANGE_ID), never()).close(); + } + + @Test + public void testTooManyIndexTypesFailsConstruction() { + installIndexService(standardAndSyntheticTypes(Long.SIZE + 1)); + assertEquals(IndexService.getInstance().getAllIndexes().size(), Long.SIZE + 1); + + SegmentDirectory.Reader segmentReader = mock(SegmentDirectory.Reader.class); + IllegalStateException thrown = + expectThrows(IllegalStateException.class, () -> newContainer(segmentReader, false)); + assertTrue(thrown.getMessage().contains("65 index types"), thrown.getMessage()); + verify(segmentReader, never()).hasIndexFor(any(), any()); + } + + @Test + public void testMaximumIndexTypesFitInMask() + throws IOException { + List types = standardAndSyntheticTypes(Long.SIZE); + installIndexService(types); + IndexService indexService = IndexService.getInstance(); + assertEquals(indexService.getAllIndexes().size(), Long.SIZE); + + // The two extreme numeric ids (0 and 63) plus one in the middle. + Set present = new HashSet<>(); + for (int numericId : new int[]{0, 31, Long.SIZE - 1}) { + present.add(indexService.get(numericId).getId()); + } + PhysicalColumnIndexContainer container = newContainer(mockSegmentReader(present), false); + for (StubIndexType type : types) { + if (present.contains(type.getId())) { + assertSame(container.getIndex(type), _stubReaders.get(type.getId()), type.getId()); + } else { + assertNull(container.getIndex(type), type.getId()); + } + } + } + + /// Installs an IndexService whose types carry every standard index id, so numeric ids match production, backed by + /// one mock reader per type. + private void installStandardStubs() { + installIndexService(standardAndSyntheticTypes(STANDARD_IDS.size())); + } + + /// The standard ids (so [StandardIndexes] accessors resolve) padded with synthetic ids up to the given count. + private static List standardAndSyntheticTypes(int numTypes) { + List types = new ArrayList<>(); + for (String id : STANDARD_IDS) { + types.add(new StubIndexType(id)); + } + for (int i = STANDARD_IDS.size(); i < numTypes; i++) { + types.add(new StubIndexType(String.format("synthetic_index_%02d", i))); + } + return types; + } + + private void installIndexService(List types) { + Set> plugins = new HashSet<>(); + for (StubIndexType type : types) { + IndexReader reader = mock(IndexReader.class, type.getId()); + _stubReaders.put(type.getId(), reader); + _stubTypes.put(type.getId(), type); + type.setReaderFactory((segmentReader, configs, metadata) -> reader); + plugins.add((IndexPlugin) () -> type); + } + IndexService.setInstance(new IndexService(plugins)); + } + + private SegmentDirectory.Reader mockSegmentReader(Set presentIds) { + SegmentDirectory.Reader segmentReader = mock(SegmentDirectory.Reader.class); + when(segmentReader.hasIndexFor(eq(COLUMN), any())).thenAnswer( + invocation -> presentIds.contains(((IndexType) invocation.getArgument(1)).getId())); + return segmentReader; + } + + private PhysicalColumnIndexContainer newContainer(Set presentIds, boolean forwardIndexOnly) + throws IOException { + return newContainer(mockSegmentReader(presentIds), forwardIndexOnly); + } + + private static PhysicalColumnIndexContainer newContainer(SegmentDirectory.Reader segmentReader, + boolean forwardIndexOnly) + throws IOException { + ColumnMetadata metadata = mock(ColumnMetadata.class); + when(metadata.getColumnName()).thenReturn(COLUMN); + IndexLoadingConfig indexLoadingConfig = mock(IndexLoadingConfig.class); + when(indexLoadingConfig.isForwardIndexOnly()).thenReturn(forwardIndexOnly); + return new PhysicalColumnIndexContainer(segmentReader, metadata, indexLoadingConfig); + } + + /// Asserts that getIndex returns the created reader for every present standard id and null for the others. + private void assertPresentExactly(PhysicalColumnIndexContainer container, Set presentIds) { + for (String id : STANDARD_IDS) { + StubIndexType type = _stubTypes.get(id); + if (presentIds.contains(id)) { + assertSame(container.getIndex(type), _stubReaders.get(id), id); + } else { + assertNull(container.getIndex(type), id); + } + } + } + + private static final class StubIndexType implements IndexType { + private final String _id; + private IndexReaderFactory _factory; + + StubIndexType(String id) { + _id = id; + } + + void setReaderFactory(IndexReaderFactory factory) { + _factory = factory; + } + + @Override + public String getId() { + return _id; + } + + @Override + public Class getIndexConfigClass() { + return IndexConfig.class; + } + + @Override + public IndexConfig getDefaultConfig() { + // The container reads the vector config through StandardIndexes.vector(), typed as VectorIndexConfig. + return _id.equals(StandardIndexes.VECTOR_ID) ? VectorIndexConfig.DISABLED : IndexConfig.DISABLED; + } + + @Override + public Map getConfig(TableConfig tableConfig, Schema schema) { + return Map.of(); + } + + @Override + public IndexCreator createIndexCreator(IndexCreationContext context, IndexConfig indexConfig) { + throw new UnsupportedOperationException(); + } + + @Override + public IndexReaderFactory getReaderFactory() { + return _factory; + } + + @Override + public List getFileExtensions(@Nullable ColumnMetadata columnMetadata) { + return List.of(); + } + + @Override + public IndexHandler createIndexHandler(SegmentDirectory segmentDirectory, + Map configsByCol, Schema schema, TableConfig tableConfig) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean requiresDictionary(FieldSpec fieldSpec, IndexConfig indexConfig) { + return false; + } + + @Override + public boolean shouldInvalidateOnDictionaryChange(FieldSpec fieldSpec, IndexConfig indexConfig) { + return false; + } + + @Override + public void convertToNewFormat(TableConfig tableConfig, Schema schema) { + } + + @Override + public String toString() { + return _id; + } + } }