Stop the segment preprocess from building a Schema per segment

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java

@@ -248,7 +248,7 @@ Map<String, List<Operation>> computeOperations(SegmentDirectory.Reader segmentRe
     }
 
     Map<String, List<Operation>> columnOperationsMap = new HashMap<>();
-    Set<String> existingAllColumns = segmentMetadata.getSchema().getPhysicalColumnNames();
+    Set<String> existingAllColumns = segmentMetadata.getPhysicalColumnNames();
     Set<String> existingDictColumns = _segmentDirectory.getColumnsWithIndex(StandardIndexes.dictionary());
     Set<String> existingForwardIndexColumns = _segmentDirectory.getColumnsWithIndex(StandardIndexes.forward());
     Set<String> existingInvertedIndexColumns =

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java

@@ -23,6 +23,7 @@
 import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.Consumer;
 import org.apache.commons.configuration2.PropertiesConfiguration;
 import org.apache.commons.lang3.Strings;
 import org.apache.pinot.segment.local.segment.creator.impl.SegmentColumnarIndexCreator;
@@ -44,7 +45,6 @@
 import org.apache.pinot.segment.spi.store.SegmentDirectory;
 import org.apache.pinot.segment.spi.utils.SegmentMetadataUtils;
 import org.apache.pinot.spi.data.FieldSpec;
-import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.utils.ByteArray;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -75,74 +75,60 @@ public ColumnMinMaxValueGenerator(SegmentMetadata segmentMetadata, SegmentDirect
   /// Returns the list of columns that need min/max values to be updated
   public List<String> columnMinMaxValueUpdates() {
     List<String> columns = new ArrayList<>();
-    for (String column : getColumnsToAddMinMaxValue()) {
-      if (needAddColumnMinMaxValueForColumn(column)) {
-        columns.add(column);
+    forEachSelectedColumn(columnMetadata -> {
+      if (needAddColumnMinMaxValueForColumn(columnMetadata)) {
+        columns.add(columnMetadata.getColumnName());
       }
-    }
+    });
     return columns;
   }
 
   public void addColumnMinMaxValue()
       throws Exception {
     Preconditions.checkState(_columnMinMaxValueGeneratorMode != ColumnMinMaxValueGeneratorMode.NONE);
     _segmentProperties = SegmentMetadataUtils.getPropertiesConfiguration(_segmentMetadata);
-    for (String column : getColumnsToAddMinMaxValue()) {
-      addColumnMinMaxValueForColumn(column);
-    }
+    forEachSelectedColumn(this::addColumnMinMaxValueForColumn);
     if (_minMaxValueAdded) {
       SegmentMetadataUtils.savePropertiesConfiguration(_segmentProperties, _segmentMetadata.getIndexDir());
     }
   }
 
-  private List<String> getColumnsToAddMinMaxValue() {
-    Schema schema = _segmentMetadata.getSchema();
-    List<String> columnsToAddMinMaxValue = new ArrayList<>();
+  /// Runs `action` on every column the generator mode selects.
+  ///
+  /// The selection reads the field specs off the column metadata rather than off `_segmentMetadata.getSchema()`,
+  /// which is the same data (the schema is derived from the column metadata) but costs a `Schema` per segment. This
+  /// runs on every segment load — the default mode is `ALL` — so a schema built here would be cached for the
+  /// segment's whole life, and a server holding tens of thousands of wide segments would keep one per segment.
+  private void forEachSelectedColumn(Consumer<ColumnMetadata> action) {
+    for (ColumnMetadata columnMetadata : _segmentMetadata.getAllColumnMetadata()) {
+      FieldSpec fieldSpec = columnMetadata.getFieldSpec();
+      if (!fieldSpec.isVirtualColumn() && isSelected(fieldSpec.getFieldType())) {
+        action.accept(columnMetadata);
+      }
+    }
+  }
 
-    // mode ALL - use all columns
-    // mode NON_METRIC - use all dimensions and time columns
-    // mode TIME - use only time columns
+  /// Whether the generator mode covers the given field type: `ALL` takes every column, `NON_METRIC` every column but
+  /// the metrics, `TIME` only the time columns.
+  private boolean isSelected(FieldSpec.FieldType fieldType) {
     switch (_columnMinMaxValueGeneratorMode) {
       case ALL:
-        for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) {
-          if (!fieldSpec.isVirtualColumn()) {
-            columnsToAddMinMaxValue.add(fieldSpec.getName());
-          }
-        }
-        break;
+        return true;
       case NON_METRIC:
-        for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) {
-          if (!fieldSpec.isVirtualColumn() && fieldSpec.getFieldType() != FieldSpec.FieldType.METRIC) {
-            columnsToAddMinMaxValue.add(fieldSpec.getName());
-          }
-        }
-        break;
+        return fieldType != FieldSpec.FieldType.METRIC;
       case TIME:
-        for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) {
-          if (!fieldSpec.isVirtualColumn() && (fieldSpec.getFieldType() == FieldSpec.FieldType.TIME
-              || fieldSpec.getFieldType() == FieldSpec.FieldType.DATE_TIME)) {
-            columnsToAddMinMaxValue.add(fieldSpec.getName());
-          }
-        }
-        break;
+        return fieldType == FieldSpec.FieldType.TIME || fieldType == FieldSpec.FieldType.DATE_TIME;
       default:
         throw new IllegalStateException("Unsupported generator mode: " + _columnMinMaxValueGeneratorMode);
     }
-
-    return columnsToAddMinMaxValue;
-  }
-
-  private boolean needAddColumnMinMaxValueForColumn(String columnName) {
-    return needAddColumnMinMaxValueForColumn(_segmentMetadata.getColumnMetadataFor(columnName));
   }
 
   private boolean needAddColumnMinMaxValueForColumn(ColumnMetadata columnMetadata) {
     return columnMetadata.getMinValue() == null && columnMetadata.getMaxValue() == null
         && !columnMetadata.isMinMaxValueInvalid();
   }
 
-  private void addColumnMinMaxValueForColumn(String columnName) {
-    ColumnMetadata columnMetadata = _segmentMetadata.getColumnMetadataFor(columnName);
+  private void addColumnMinMaxValueForColumn(ColumnMetadata columnMetadata) {
     if (!needAddColumnMinMaxValueForColumn(columnMetadata)) {
       return;
     }
@@ -155,7 +141,7 @@ private void addColumnMinMaxValueForColumn(String columnName) {
       _minMaxValueAdded = true;
     } catch (Exception e) {
       LOGGER.error("Caught exception while generating min/max value for column: {} in segment: {}, continuing without "
-          + "persisting them", columnName, _segmentMetadata.getName(), e);
+          + "persisting them", columnMetadata.getColumnName(), _segmentMetadata.getName(), e);
     }
   }
 

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

@@ -40,7 +40,10 @@
 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.local.segment.index.loader.IndexLoadingConfig;
+import org.apache.pinot.segment.local.segment.index.loader.SegmentPreProcessor;
 import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.local.segment.store.SegmentLocalFSDirectory;
 import org.apache.pinot.segment.local.segment.virtualcolumn.VirtualColumnProviderFactory;
 import org.apache.pinot.segment.spi.ColumnMetadata;
 import org.apache.pinot.segment.spi.ImmutableSegment;
@@ -51,6 +54,7 @@
 import org.apache.pinot.segment.spi.index.metadata.ColumnMetadataImpl;
 import org.apache.pinot.segment.spi.index.metadata.EmptyColumnMetadata;
 import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
+import org.apache.pinot.segment.spi.store.SegmentDirectory;
 import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
 import org.apache.pinot.spi.config.table.FieldConfig;
 import org.apache.pinot.spi.config.table.OpenStructIndexConfig;
@@ -332,6 +336,37 @@ public void testSchemaDerivedLazilyFromColumnMetadata()
     assertEquals(SegmentMetadataImpl.getNumSchemaMaterializations(), materializations + 1);
   }
 
+  /// The preprocess that runs on every segment load asks the forward-index handler which physical columns exist. That
+  /// question must not build the per-segment schema: doing so once per segment pins one [Schema] per loaded segment
+  /// for its whole life, which on a server holding tens of thousands of wide segments is hundreds of megabytes.
+  @Test
+  public void testPreprocessDoesNotBuildTheSegmentSchema()
+      throws Exception {
+    // The forward-index handler skips segments older than v3, so the preprocess only reaches it on a v3 segment.
+    new SegmentV1V2ToV3FormatConverter().convert(_segmentDirectory);
+
+    long materializations = SegmentMetadataImpl.getNumSchemaMaterializations();
+    SegmentMetadataImpl metadata = new SegmentMetadataImpl(_segmentDirectory);
+    Set<String> physical = metadata.getPhysicalColumnNames();
+    assertFalse(metadata.isSchemaMaterialized(), "listing physical columns must not build the segment schema");
+    assertEquals(SegmentMetadataImpl.getNumSchemaMaterializations(), materializations);
+    assertEquals(List.copyOf(physical), List.copyOf(metadata.getSchema().getPhysicalColumnNames()),
+        "the derived names must equal what the schema reports, in the same order");
+    assertFalse(physical.contains(BuiltInVirtualColumn.DOCID));
+
+    TableConfig tableConfig =
+        new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable").setTimeColumnName("daysSinceEpoch").build();
+    IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig(tableConfig, metadata.getSchema());
+    indexLoadingConfig.setReadMode(ReadMode.mmap);
+    long beforePreprocess = SegmentMetadataImpl.getNumSchemaMaterializations();
+    try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(_segmentDirectory, ReadMode.mmap);
+        SegmentPreProcessor preProcessor = new SegmentPreProcessor(segmentDirectory, indexLoadingConfig)) {
+      preProcessor.process();
+    }
+    assertEquals(SegmentMetadataImpl.getNumSchemaMaterializations(), beforePreprocess,
+        "segment preprocess must not build any segment schema");
+  }
+
   /// Loading a segment registers the built-in virtual columns in the column metadata, so the schema derived afterwards
   /// includes them exactly as the schema the loader used to build eagerly did, while neither the load nor serving the
   /// segment (column listings, data sources, the metadata JSON) builds any schema.

pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/SegmentMetadata.java

@@ -25,14 +25,17 @@
 import java.util.Map;
 import java.util.NavigableSet;
 import java.util.Set;
+import java.util.SortedSet;
 import java.util.TreeMap;
+import java.util.TreeSet;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BiConsumer;
 import javax.annotation.Nullable;
 import org.apache.pinot.segment.spi.creator.SegmentVersion;
 import org.apache.pinot.segment.spi.index.multicolumntext.MultiColumnTextMetadata;
 import org.apache.pinot.segment.spi.index.startree.StarTreeV2Metadata;
 import org.apache.pinot.spi.annotations.InterfaceAudience;
+import org.apache.pinot.spi.data.FieldSpec;
 import org.apache.pinot.spi.data.Schema;
 import org.joda.time.Duration;
 import org.joda.time.Interval;
@@ -134,7 +137,8 @@ default int getNumColumns() {
   /// empty for a segment that holds none (a CONSUMING one, which answers [#getColumnMetadataFor(String)] with `null`
   /// for every column of its schema).
   default Collection<ColumnMetadata> getAllColumnMetadata() {
-    return getColumnMetadataMap().values();
+    TreeMap<String, ColumnMetadata> columnMetadataMap = getColumnMetadataMap();
+    return columnMetadataMap != null ? columnMetadataMap.values() : List.of();
   }
 
   /// Applies `action` to every (column name, column metadata) pair, in the natural column-name order of
@@ -161,6 +165,28 @@ default ColumnMetadata getColumnMetadataFor(String column) {
     return getColumnMetadataMap().get(column);
   }
 
+  /// The names of the physical (non-virtual) columns, i.e. `getSchema().getPhysicalColumnNames()` without building
+  /// the schema. Segment load runs this once per segment (the forward-index handler asks which columns exist), and on
+  /// a server holding tens of thousands of wide segments a schema built there would be cached for the segment's whole
+  /// life: one [Schema] per segment, each with a tree entry and two list slots per column. A segment that holds no
+  /// column metadata (a CONSUMING one) still answers from its schema, which it was constructed with.
+  ///
+  /// Sorted, like the [Schema#getPhysicalColumnNames()] this replaces, and the same set for both segment kinds.
+  default SortedSet<String> getPhysicalColumnNames() {
+    Collection<ColumnMetadata> columnMetadata = getAllColumnMetadata();
+    if (columnMetadata.isEmpty()) {
+      return getSchema().getPhysicalColumnNames();
+    }
+    TreeSet<String> physicalColumnNames = new TreeSet<>();
+    for (ColumnMetadata metadata : columnMetadata) {
+      FieldSpec fieldSpec = metadata.getFieldSpec();
+      if (!fieldSpec.isVirtualColumn()) {
+        physicalColumnNames.add(fieldSpec.getName());
+      }
+    }
+    return physicalColumnNames;
+  }
+
   /// Registers the metadata of a column, replacing any metadata already registered under the same name. An
   /// implementation that holds no column metadata (a CONSUMING segment) may reject this.
   default void addColumnMetadata(String column, ColumnMetadata columnMetadata) {

pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/SegmentMetadataDefaultsTest.java

@@ -0,0 +1,87 @@
+/**
+ * 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.spi;
+
+import java.util.List;
+import java.util.SortedSet;
+import org.apache.pinot.spi.data.DimensionFieldSpec;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.CALLS_REAL_METHODS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.withSettings;
+import static org.testng.Assert.assertEquals;
+
+
+/// Covers the [SegmentMetadata] default methods against an implementation that holds no column metadata, which
+/// [SegmentMetadata#getColumnMetadataMap()] is documented to answer with `null`. The only implementation in the tree
+/// overrides [SegmentMetadata#getAllColumnMetadata()], so the defaults are exercised here rather than through it.
+public class SegmentMetadataDefaultsTest {
+
+  @Test
+  public void testAllColumnMetadataIsEmptyWithoutAColumnMetadataMap() {
+    SegmentMetadata segmentMetadata = metadataHoldingNoColumns();
+    assertEquals(segmentMetadata.getAllColumnMetadata(), List.of());
+  }
+
+  @Test
+  public void testPhysicalColumnNamesFallBackToTheSchemaWithoutColumnMetadata() {
+    SegmentMetadata segmentMetadata = metadataHoldingNoColumns();
+    assertEquals(segmentMetadata.getPhysicalColumnNames(), schema().getPhysicalColumnNames());
+  }
+
+  @Test
+  public void testPhysicalColumnNamesAreSortedAndSkipVirtualColumns() {
+    ColumnMetadata zebra = columnMetadata("zebra", false);
+    ColumnMetadata apple = columnMetadata("apple", false);
+    ColumnMetadata virtual = columnMetadata("$docId", true);
+    SegmentMetadata segmentMetadata = mock(SegmentMetadata.class, CALLS_REAL_METHODS);
+    when(segmentMetadata.getAllColumnMetadata()).thenReturn(List.of(zebra, apple, virtual));
+
+    SortedSet<String> physicalColumnNames = segmentMetadata.getPhysicalColumnNames();
+    assertEquals(List.copyOf(physicalColumnNames), List.of("apple", "zebra"));
+  }
+
+  private static SegmentMetadata metadataHoldingNoColumns() {
+    SegmentMetadata segmentMetadata = mock(SegmentMetadata.class, withSettings().defaultAnswer(CALLS_REAL_METHODS));
+    when(segmentMetadata.getColumnMetadataMap()).thenReturn(null);
+    when(segmentMetadata.getSchema()).thenReturn(schema());
+    return segmentMetadata;
+  }
+
+  private static Schema schema() {
+    return new Schema.SchemaBuilder()
+        .addSingleValueDimension("zebra", FieldSpec.DataType.STRING)
+        .addSingleValueDimension("apple", FieldSpec.DataType.STRING)
+        .build();
+  }
+
+  private static ColumnMetadata columnMetadata(String column, boolean virtual) {
+    DimensionFieldSpec fieldSpec = new DimensionFieldSpec(column, FieldSpec.DataType.STRING, true);
+    if (virtual) {
+      fieldSpec.setVirtualColumnProvider("org.apache.pinot.segment.spi.virtualcolumn.DocIdVirtualColumnProvider");
+    }
+    ColumnMetadata columnMetadata = mock(ColumnMetadata.class);
+    when(columnMetadata.getFieldSpec()).thenReturn(fieldSpec);
+    return columnMetadata;
+  }
+}