Allocate ColumnMetadataImpl index sizes lazily and skip the index_map lookup without an index dir

pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/SegmentMetadataImplTest.java

@@ -20,17 +20,23 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import java.io.File;
+import java.io.FileInputStream;
 import java.io.IOException;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import org.apache.commons.configuration2.ex.ConfigurationException;
 import org.apache.commons.io.FileUtils;
 import org.apache.pinot.segment.local.segment.creator.SegmentTestUtils;
 import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.index.converter.SegmentV1V2ToV3FormatConverter;
 import org.apache.pinot.segment.spi.ColumnMetadata;
 import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
 import org.apache.pinot.segment.spi.creator.SegmentIndexCreationDriver;
+import org.apache.pinot.segment.spi.creator.SegmentVersion;
+import org.apache.pinot.segment.spi.index.StandardIndexes;
+import org.apache.pinot.segment.spi.index.metadata.ColumnMetadataImpl;
 import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
+import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
 import org.apache.pinot.util.TestUtils;
 import org.testng.Assert;
 import org.testng.annotations.AfterMethod;
@@ -100,4 +106,39 @@ public void testToJson()
       assertEquals(jsonColumnMeta.get("hasDictionary").asBoolean(), columnMeta.hasDictionary());
     }
   }
+
+  /// Index sizes come from the local `index_map`, so a segment loaded through the stream constructor (tiered
+  /// storage, no index directory) reports none while the rest of the metadata matches the directory load.
+  @Test
+  public void testIndexSizesOnlyFromIndexDir()
+      throws Exception {
+    // The fixture builds a v1 segment; index sizes exist only in the v3 index_map.
+    new SegmentV1V2ToV3FormatConverter().convert(_segmentDirectory);
+    SegmentMetadataImpl fromDir = new SegmentMetadataImpl(_segmentDirectory);
+    assertEquals(fromDir.getVersion(), SegmentVersion.v3);
+    SegmentMetadataImpl fromStreams;
+    try (FileInputStream metadataProperties =
+        new FileInputStream(SegmentDirectoryPaths.findMetadataFile(_segmentDirectory));
+        FileInputStream creationMeta =
+            new FileInputStream(SegmentDirectoryPaths.findCreationMetaFile(_segmentDirectory))) {
+      fromStreams = new SegmentMetadataImpl(metadataProperties, creationMeta);
+    }
+
+    assertEquals(fromStreams.getColumnMetadataMap().keySet(), fromDir.getColumnMetadataMap().keySet());
+    assertEquals(fromStreams.getTotalDocs(), fromDir.getTotalDocs());
+    assertEquals(fromStreams.getSchema(), fromDir.getSchema());
+    for (Map.Entry<String, ColumnMetadata> entry : fromDir.getColumnMetadataMap().entrySet()) {
+      ColumnMetadata dirColumn = entry.getValue();
+      Assert.assertTrue(dirColumn.getNumIndexes() > 0, entry.getKey());
+      long forwardSize = dirColumn.getIndexSizeFor(StandardIndexes.forward());
+      Assert.assertTrue(forwardSize > 0, entry.getKey());
+      assertEquals(((ColumnMetadataImpl) dirColumn).getIndexSizeMap().get(StandardIndexes.forward()),
+          (Long) forwardSize, entry.getKey());
+      ColumnMetadata streamColumn = fromStreams.getColumnMetadataMap().get(entry.getKey());
+      assertEquals(streamColumn.getNumIndexes(), 0, entry.getKey());
+      assertEquals(streamColumn.getIndexSizeFor(StandardIndexes.forward()), ColumnMetadata.UNAVAILABLE,
+          entry.getKey());
+      Assert.assertTrue(((ColumnMetadataImpl) streamColumn).getIndexSizeMap().isEmpty(), entry.getKey());
+    }
+  }
 }

pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java

@@ -21,8 +21,8 @@
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.google.common.collect.Maps;
 import it.unimi.dsi.fastutil.ints.IntSet;
-import it.unimi.dsi.fastutil.longs.LongArrayList;
 import java.math.BigDecimal;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -88,12 +88,18 @@ public class ColumnMetadataImpl implements ColumnMetadata {
   @Nullable
   private final CompressionMetadata _compressionMetadata;
 
-  /// List of longs, each encodes:
+  /// Index sizes, one long per index of this column, each encoding:
   /// - 2 byte - numeric id of IndexType
   /// - 6 byte - index size
   ///
-  /// Use non-default size to save space for most columns
-  private final LongArrayList _indexTypeSizeList = new LongArrayList(2);
+  /// Null until the first [#addIndexSize(short, long)], which only [SegmentMetadataImpl] calls while loading a
+  /// v3 segment from a local index directory (the sizes come from `index_map`). Metadata loaded from streams, v1/v2
+  /// segments and the built-in virtual columns never have sizes, so they pay no allocation: a wide external-table
+  /// segment holds one [ColumnMetadataImpl] per column for its lifetime, and an eagerly allocated list plus its
+  /// backing array cost ~56 bytes per column. Populated at load time before the metadata is published, like the
+  /// rest of this class; not thread-safe.
+  @Nullable
+  private long[] _indexTypeSizes;
 
   private ColumnMetadataImpl(FieldSpec fieldSpec, int totalDocs, int cardinality, boolean hasDictionary,
       @Nullable EncodingType forwardIndexEncoding, boolean sorted, boolean nonNull, @Nullable Comparable minValue,
@@ -252,8 +258,11 @@ public List<String> getSparseKeys() {
 
   @Override
   public long getIndexSizeFor(IndexType type) {
+    if (_indexTypeSizes == null) {
+      return UNAVAILABLE;
+    }
     short indexId = IndexService.getInstance().getNumericId(type);
-    for (long typeAndSize : _indexTypeSizeList) {
+    for (long typeAndSize : _indexTypeSizes) {
       if (indexId == unpackIndexType(typeAndSize)) {
         return unpackIndexSize(typeAndSize);
       }
@@ -268,17 +277,21 @@ public void addIndexSize(short indexType, long size) {
           "Index size should be a non-negative integer value between 0 and " + SIZE_MASK);
     }
     long typeAndSize = ((long) indexType) << 48 | (size & SIZE_MASK);
-    _indexTypeSizeList.add(typeAndSize);
+    // A column has a handful of indexes, so grow by one rather than pre-size.
+    int numIndexes = _indexTypeSizes == null ? 0 : _indexTypeSizes.length;
+    long[] grown = _indexTypeSizes == null ? new long[1] : Arrays.copyOf(_indexTypeSizes, numIndexes + 1);
+    grown[numIndexes] = typeAndSize;
+    _indexTypeSizes = grown;
   }
 
   @Override
   public int getNumIndexes() {
-    return _indexTypeSizeList.size();
+    return _indexTypeSizes == null ? 0 : _indexTypeSizes.length;
   }
 
   @Override
   public short getIndexType(int position) {
-    return unpackIndexType(_indexTypeSizeList.getLong(position));
+    return unpackIndexType(_indexTypeSizes[position]);
   }
 
   private static short unpackIndexType(long typeAndSize) {
@@ -287,7 +300,7 @@ private static short unpackIndexType(long typeAndSize) {
 
   @Override
   public long getIndexSize(int position) {
-    return unpackIndexSize(_indexTypeSizeList.getLong(position));
+    return unpackIndexSize(_indexTypeSizes[position]);
   }
 
   private static long unpackIndexSize(long typeAndSize) {
@@ -341,15 +354,16 @@ public boolean equals(Object o) {
         && Objects.equals(_parentColumn, that._parentColumn)
         && Objects.equals(_sparseKeys, that._sparseKeys)
         && Objects.equals(_compressionMetadata, that._compressionMetadata)
-        && Objects.equals(_indexTypeSizeList, that._indexTypeSizeList);
+        && Arrays.equals(_indexTypeSizes, that._indexTypeSizes);
   }
 
   @Override
   public int hashCode() {
     return Objects.hash(_fieldSpec, _totalDocs, _cardinality, _hasDictionary, _forwardIndexEncoding, _sorted, _nonNull,
         _minValue, _maxValue, _minMaxValueInvalid, _lengthOfShortestElement, _lengthOfLongestElement, _isAscii,
         _totalNumberOfEntries, _maxNumberOfMultiValues, _maxRowLengthInBytes, _bitsPerElement, _partitionFunction,
-        _partitions, _autoGenerated, _parentColumn, _sparseKeys, _compressionMetadata, _indexTypeSizeList);
+        _partitions, _autoGenerated, _parentColumn, _sparseKeys, _compressionMetadata,
+        Arrays.hashCode(_indexTypeSizes));
   }
 
   @Override
@@ -377,7 +391,7 @@ public String toString() {
         + ", _parentColumn=" + _parentColumn
         + ", _sparseKeys=" + _sparseKeys
         + ", _compressionMetadata=" + _compressionMetadata
-        + ", _indexTypeSizeList=" + _indexTypeSizeList
+        + ", _indexTypeSizes=" + Arrays.toString(_indexTypeSizes)
         + '}';
   }
 
@@ -578,9 +592,13 @@ private static Comparable parseValue(DataType storedType, String column, String
   //       `/tables/{tableName}/segments/{segmentName}/metadata`
   @SuppressWarnings("unused")
   public Map<IndexType<?, ?, ?>, Long> getIndexSizeMap() {
+    if (_indexTypeSizes == null) {
+      // A mutable empty map, so the JSON payload stays `{}` and any caller that adds to it keeps working.
+      return new HashMap<>();
+    }
     IndexService service = IndexService.getInstance();
-    Map<IndexType<?, ?, ?>, Long> result = Maps.newHashMapWithExpectedSize(_indexTypeSizeList.size());
-    for (long typeAndSize : _indexTypeSizeList) {
+    Map<IndexType<?, ?, ?>, Long> result = Maps.newHashMapWithExpectedSize(_indexTypeSizes.length);
+    for (long typeAndSize : _indexTypeSizes) {
       short type = unpackIndexType(typeAndSize);
       long size = unpackIndexSize(typeAndSize);
       result.put(service.get(type), size);

pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SegmentMetadataImpl.java

@@ -243,8 +243,10 @@ private void init(PropertiesConfiguration segmentMetadata)
 
       // Load index metadata
       // Support V3 (e.g. SingleFileIndexDirectory only). Skip for empty segments — there is no payload to size up,
-      // and [EmptyColumnMetadata] does not support `addIndexSize`.
-      if (_segmentVersion == SegmentVersion.v3) {
+      // and [EmptyColumnMetadata] does not support `addIndexSize`. Index sizes come from the local index_map, so
+      // metadata loaded from streams (no index directory) has none; without the guard this probed a cwd-relative
+      // `v3/index_map` once per segment and would NPE on `_indexDir.getPath()` if such a file existed.
+      if (_segmentVersion == SegmentVersion.v3 && _indexDir != null) {
         File indexMapFile = new File(_indexDir, "v3" + File.separator + V1Constants.INDEX_MAP_FILE_NAME);
         if (indexMapFile.exists()) {
           IndexService indexService = IndexService.getInstance();

pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImplTest.java

@@ -205,6 +205,67 @@ public void compressionStatsDoNotExpandExistingColumnMetadataJson() {
     assertFalse(json.has("dictionaryUncompressedValueSizeInBytes"));
   }
 
+  // The index-size API works on numeric index ids so this module's tests need no index plugins registered.
+  private static final short FORWARD_ID = 2;
+  private static final short DICTIONARY_ID = 0;
+  private static final short JSON_ID = 5;
+
+  @Test
+  public void indexSizesAbsentByDefault()
+      throws Exception {
+    ColumnMetadataImpl metadata = ColumnMetadataImpl.fromPropertiesConfiguration(baseConfig("col"), 1, "col");
+    assertEquals(metadata.getNumIndexes(), 0);
+    assertTrue(metadata.getIndexSizeMap().isEmpty());
+    // The REST segment-metadata payload keeps its shape: an empty object, never null.
+    JsonNode indexSizeMap = JsonUtils.objectToJsonNode(metadata).get("indexSizeMap");
+    assertTrue(indexSizeMap.isObject() && indexSizeMap.isEmpty(), String.valueOf(indexSizeMap));
+  }
+
+  @Test
+  public void indexSizesRoundTripAfterAdd() {
+    ColumnMetadataImpl metadata = ColumnMetadataImpl.fromPropertiesConfiguration(baseConfig("col"), 1, "col");
+    metadata.addIndexSize(FORWARD_ID, 100);
+    metadata.addIndexSize(DICTIONARY_ID, 200);
+    metadata.addIndexSize(JSON_ID, 0);
+
+    assertEquals(metadata.getNumIndexes(), 3);
+    assertEquals(metadata.getIndexType(0), FORWARD_ID);
+    assertEquals(metadata.getIndexSize(0), 100);
+    assertEquals(metadata.getIndexType(1), DICTIONARY_ID);
+    assertEquals(metadata.getIndexSize(1), 200);
+    assertEquals(metadata.getIndexType(2), JSON_ID);
+    assertEquals(metadata.getIndexSize(2), 0);
+    // A 48-bit size survives the packing.
+    metadata.addIndexSize((short) 7, (1L << 48) - 1);
+    assertEquals(metadata.getIndexSize(3), (1L << 48) - 1);
+    assertEquals(metadata.getIndexType(3), 7);
+  }
+
+  @Test
+  public void indexSizesParticipateInValueObjectMethods() {
+    ColumnMetadataImpl first = ColumnMetadataImpl.fromPropertiesConfiguration(baseConfig("col"), 1, "col");
+    ColumnMetadataImpl second = ColumnMetadataImpl.fromPropertiesConfiguration(baseConfig("col"), 1, "col");
+    ColumnMetadataImpl third = ColumnMetadataImpl.fromPropertiesConfiguration(baseConfig("col"), 1, "col");
+    ColumnMetadataImpl noSizes = ColumnMetadataImpl.fromPropertiesConfiguration(baseConfig("col"), 1, "col");
+    first.addIndexSize(FORWARD_ID, 100);
+    second.addIndexSize(FORWARD_ID, 100);
+    third.addIndexSize(FORWARD_ID, 101);
+
+    assertEquals(first, second);
+    assertEquals(first.hashCode(), second.hashCode());
+    assertNotEquals(first, third);
+    assertNotEquals(first, noSizes);
+    assertTrue(first.toString().contains("_indexTypeSizes=["), first.toString());
+  }
+
+  @Test
+  public void rejectsInvalidIndexSize() {
+    ColumnMetadataImpl metadata = ColumnMetadataImpl.fromPropertiesConfiguration(baseConfig("col"), 1, "col");
+    expectThrows(IllegalArgumentException.class, () -> metadata.addIndexSize(FORWARD_ID, -1));
+    expectThrows(IllegalArgumentException.class, () -> metadata.addIndexSize(FORWARD_ID, 1L << 48));
+    assertEquals(metadata.getNumIndexes(), 0, "a rejected size must not be recorded");
+  }
+
   private static PropertiesConfiguration baseConfig(String column) {
     PropertiesConfiguration config = new PropertiesConfiguration();
     config.setProperty(Column.getKeyFor(column, Column.COLUMN_NAME), column);