Slim ColumnMetadataImpl to 72 bytes and derive the per-segment Schema lazily

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegment.java

@@ -19,6 +19,7 @@
 package org.apache.pinot.segment.local.indexsegment.immutable;
 
 import com.google.common.base.Preconditions;
+import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -79,14 +80,16 @@ public SegmentMetadataImpl getSegmentMetadata() {
     return _segmentMetadata;
   }
 
+  // Both are views of the column metadata map, so neither builds the segment schema (see SegmentMetadataImpl)
+
   @Override
   public Set<String> getColumnNames() {
-    return _segmentMetadata.getSchema().getColumnNames();
+    return Collections.unmodifiableSet(_segmentMetadata.getColumnMetadataMap().keySet());
   }
 
   @Override
   public Set<String> getPhysicalColumnNames() {
-    return _segmentMetadata.getSchema().getPhysicalColumnNames();
+    return new PhysicalColumnNames(_segmentMetadata.getColumnMetadataMap());
   }
 
   @Override

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java

@@ -25,11 +25,13 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -111,6 +113,10 @@ public class ImmutableSegmentImpl implements ImmutableSegment {
   private final StarTreeIndexContainer _starTreeIndexContainer;
   private final TextIndexReader _multiColumnTextIndex;
   private final Map<String, DataSource> _dataSources;
+  // Views of the column metadata map's keys (all columns / the physical ones), so listing columns never builds the
+  // segment schema, which SegmentMetadataImpl derives on demand and a wide segment must not retain per column
+  private final Set<String> _columnNames;
+  private final Set<String> _physicalColumnNames;
 
   // Lazy column materialization; all null in eager mode. See the class documentation.
   @Nullable
@@ -150,14 +156,16 @@ public ImmutableSegmentImpl(
     _columnMaterializer = null;
     _openStructChildren = null;
     _materializationLock = null;
-    _dataSources =
-        new Object2ObjectOpenHashMap<>(segmentMetadata.getColumnMetadataMap().size());
+    TreeMap<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap();
+    _columnNames = Collections.unmodifiableSet(columnMetadataMap.keySet());
+    _physicalColumnNames = new PhysicalColumnNames(columnMetadataMap);
+    _dataSources = new Object2ObjectOpenHashMap<>(columnMetadataMap.size());
 
     Map<String, Map<String, DataSource>> openStructDenseChildren = new HashMap<>();
     Map<String, DataSource> openStructSparseChildren = new HashMap<>();
     Set<String> openStructParents = new HashSet<>();
 
-    for (Map.Entry<String, ColumnMetadata> entry : segmentMetadata.getColumnMetadataMap().entrySet()) {
+    for (Map.Entry<String, ColumnMetadata> entry : columnMetadataMap.entrySet()) {
       String colName = entry.getKey();
       ColumnMetadata columnMetadata = entry.getValue();
 
@@ -181,20 +189,18 @@ public ImmutableSegmentImpl(
       }
     }
 
-    if (!openStructParents.isEmpty()) {
-      Schema schema = segmentMetadata.getSchema();
-      for (String parent : openStructParents) {
-        FieldSpec fieldSpec = schema != null ? schema.getFieldSpecFor(parent) : null;
-        if (!(fieldSpec instanceof ComplexFieldSpec)) {
-          continue;
-        }
-        ColumnMetadata parentMetadata = segmentMetadata.getColumnMetadataMap().get(parent);
-        List<String> sparseKeys =
-            parentMetadata instanceof ColumnMetadataImpl impl ? impl.getSparseKeys() : null;
-        _dataSources.put(parent, new ImmutableOpenStructDataSource((ComplexFieldSpec) fieldSpec,
-            openStructDenseChildren.getOrDefault(parent, Map.of()),
-            openStructSparseChildren.get(parent), segmentMetadata.getTotalDocs(), sparseKeys));
+    for (String parent : openStructParents) {
+      // The parent's spec comes from its column metadata, not from the segment schema (see _columnNames)
+      ColumnMetadata parentMetadata = columnMetadataMap.get(parent);
+      FieldSpec fieldSpec = parentMetadata != null ? parentMetadata.getFieldSpec() : null;
+      if (!(fieldSpec instanceof ComplexFieldSpec)) {
+        continue;
       }
+      List<String> sparseKeys =
+          parentMetadata instanceof ColumnMetadataImpl impl ? impl.getSparseKeys() : null;
+      _dataSources.put(parent, new ImmutableOpenStructDataSource((ComplexFieldSpec) fieldSpec,
+          openStructDenseChildren.getOrDefault(parent, Map.of()),
+          openStructSparseChildren.get(parent), segmentMetadata.getTotalDocs(), sparseKeys));
     }
 
     _multiColumnTextIndex = multiColumnTextIndex;
@@ -225,18 +231,22 @@ public ImmutableSegmentImpl(
     _columnMaterializer = columnMaterializer;
     _openStructChildren = groupOpenStructChildren(segmentMetadata);
     _materializationLock = new ReentrantReadWriteLock();
+    TreeMap<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap();
+    _columnNames = Collections.unmodifiableSet(columnMetadataMap.keySet());
+    _physicalColumnNames = new PhysicalColumnNames(columnMetadataMap);
     _dataSources = new ConcurrentHashMap<>();
     for (String column : materializedIndexContainers.keySet()) {
       materializeDataSource(column);
     }
   }
 
-  /// Groups the materialized OPEN_STRUCT child columns under their parent, keeping only the parents the segment schema
-  /// declares as complex (the same rule the eager constructor applies).
+  /// Groups the materialized OPEN_STRUCT child columns under their parent, keeping only the parents whose column
+  /// metadata declares them complex (the same rule the eager constructor applies).
   @Nullable
   private static Map<String, List<String>> groupOpenStructChildren(SegmentMetadataImpl segmentMetadata) {
     Map<String, List<String>> children = null;
-    for (Map.Entry<String, ColumnMetadata> entry : segmentMetadata.getColumnMetadataMap().entrySet()) {
+    Map<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap();
+    for (Map.Entry<String, ColumnMetadata> entry : columnMetadataMap.entrySet()) {
       if (entry.getValue() instanceof ColumnMetadataImpl impl && impl.isMaterializedChild()) {
         if (children == null) {
           children = new HashMap<>();
@@ -247,9 +257,10 @@ private static Map<String, List<String>> groupOpenStructChildren(SegmentMetadata
     if (children == null) {
       return null;
     }
-    Schema schema = segmentMetadata.getSchema();
-    children.keySet()
-        .removeIf(parent -> !(schema != null && schema.getFieldSpecFor(parent) instanceof ComplexFieldSpec));
+    children.keySet().removeIf(parent -> {
+      ColumnMetadata parentMetadata = columnMetadataMap.get(parent);
+      return parentMetadata == null || !(parentMetadata.getFieldSpec() instanceof ComplexFieldSpec);
+    });
     return children.isEmpty() ? null : children;
   }
 
@@ -296,9 +307,9 @@ private DataSource createOpenStructDataSource(String parent) {
         denseChildren.put(OpenStructNaming.parseKey(child), childDataSource);
       }
     }
-    ComplexFieldSpec fieldSpec = (ComplexFieldSpec) _segmentMetadata.getSchema().getFieldSpecFor(parent);
-    List<String> sparseKeys =
-        columnMetadataMap.get(parent) instanceof ColumnMetadataImpl impl ? impl.getSparseKeys() : null;
+    ColumnMetadata parentMetadata = columnMetadataMap.get(parent);
+    ComplexFieldSpec fieldSpec = (ComplexFieldSpec) parentMetadata.getFieldSpec();
+    List<String> sparseKeys = parentMetadata instanceof ColumnMetadataImpl impl ? impl.getSparseKeys() : null;
     return new ImmutableOpenStructDataSource(fieldSpec, denseChildren, sparseChild, _segmentMetadata.getTotalDocs(),
         sparseKeys);
   }
@@ -470,12 +481,12 @@ public DataSource getDataSource(String column, Schema schema) {
 
   @Override
   public Set<String> getColumnNames() {
-    return _segmentMetadata.getSchema().getColumnNames();
+    return _columnNames;
   }
 
   @Override
   public Set<String> getPhysicalColumnNames() {
-    return _segmentMetadata.getSchema().getPhysicalColumnNames();
+    return _physicalColumnNames;
   }
 
   @Override

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentLoader.java

@@ -50,6 +50,7 @@
 import org.apache.pinot.segment.spi.loader.SegmentDirectoryLoaderRegistry;
 import org.apache.pinot.segment.spi.store.SegmentDirectory;
 import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
+import org.apache.pinot.spi.data.BuiltInVirtualColumnDefinitions;
 import org.apache.pinot.spi.data.FieldSpec;
 import org.apache.pinot.spi.data.OpenStructNaming;
 import org.apache.pinot.spi.data.Schema;
@@ -311,21 +312,26 @@ private static ImmutableSegmentImpl loadWithLazyColumns(SegmentDirectory segment
         starTreeIndexContainer, mcTextReader);
   }
 
-  /// Adds the built-in virtual columns to the segment schema and creates their index containers and metadata.
+  /// Creates the index containers and column metadata of the built-in virtual columns and registers them in the
+  /// segment metadata. Registering the metadata is what makes the segment schema include the virtual columns: the
+  /// schema is derived from the column metadata map on demand ([SegmentMetadataImpl#getSchema()]) and is deliberately
+  /// not built here, so a loaded segment retains no per-column schema entries until something asks for its schema.
+  /// A physical column of the same name wins, as in the schema-based registration this replaces.
   private static void instantiateVirtualColumns(SegmentMetadataImpl segmentMetadata,
       Map<String, ColumnIndexContainer> indexContainerMap) {
     Map<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap();
-    Schema segmentSchema = segmentMetadata.getSchema();
-    VirtualColumnProviderFactory.addBuiltInVirtualColumnsToSegmentSchema(segmentSchema, segmentMetadata.getName());
-    for (FieldSpec fieldSpec : segmentSchema.getAllFieldSpecs()) {
-      if (fieldSpec.isVirtualColumn()) {
-        String columnName = fieldSpec.getName();
-        VirtualColumnContext context =
-            new VirtualColumnContext(fieldSpec, segmentMetadata.getTotalDocs(), segmentMetadata);
-        VirtualColumnProvider provider = VirtualColumnProviderFactory.buildProvider(context);
-        indexContainerMap.put(columnName, provider.buildColumnIndexContainer(context));
-        columnMetadataMap.put(columnName, provider.buildMetadata(context));
+    String segmentName = segmentMetadata.getName();
+    for (BuiltInVirtualColumnDefinitions.Definition definition : BuiltInVirtualColumnDefinitions.DEFINITIONS) {
+      String columnName = definition.getName();
+      if (columnMetadataMap.containsKey(columnName)) {
+        continue;
       }
+      FieldSpec fieldSpec = VirtualColumnProviderFactory.createBuiltInFieldSpec(definition, segmentName);
+      VirtualColumnContext context =
+          new VirtualColumnContext(fieldSpec, segmentMetadata.getTotalDocs(), segmentMetadata);
+      VirtualColumnProvider provider = VirtualColumnProviderFactory.buildProvider(context);
+      indexContainerMap.put(columnName, provider.buildColumnIndexContainer(context));
+      columnMetadataMap.put(columnName, provider.buildMetadata(context));
     }
   }
 

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/PhysicalColumnNames.java

@@ -0,0 +1,103 @@
+/**
+ * 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.indexsegment.immutable;
+
+import java.util.AbstractSet;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.SortedMap;
+import org.apache.pinot.segment.spi.ColumnMetadata;
+
+
+/// Unmodifiable view of the physical columns of an immutable segment: the keys of its column metadata map whose field
+/// spec is not produced by a virtual column provider, in the map's (sorted) key order.
+///
+/// It is a view rather than a copy so a segment retains nothing per column for it: the segment schema this replaces
+/// held a `TreeMap` entry per column, and a cached `TreeSet` would hold the same. `contains` is one map lookup and
+/// iteration is a filtered pass over the map. The virtual column count is taken once at construction, which is sound
+/// because the column metadata map is fixed once the segment is loaded.
+///
+/// Thread-safe for reads, like the underlying map once loaded.
+final class PhysicalColumnNames extends AbstractSet<String> {
+  private final SortedMap<String, ColumnMetadata> _columnMetadataMap;
+  private final int _numVirtualColumns;
+
+  PhysicalColumnNames(SortedMap<String, ColumnMetadata> columnMetadataMap) {
+    _columnMetadataMap = columnMetadataMap;
+    int numVirtualColumns = 0;
+    for (ColumnMetadata columnMetadata : columnMetadataMap.values()) {
+      if (!isPhysical(columnMetadata)) {
+        numVirtualColumns++;
+      }
+    }
+    _numVirtualColumns = numVirtualColumns;
+  }
+
+  private static boolean isPhysical(ColumnMetadata columnMetadata) {
+    return !columnMetadata.getFieldSpec().isVirtualColumn();
+  }
+
+  @Override
+  public boolean contains(Object o) {
+    if (!(o instanceof String)) {
+      return false;
+    }
+    ColumnMetadata columnMetadata = _columnMetadataMap.get(o);
+    return columnMetadata != null && isPhysical(columnMetadata);
+  }
+
+  @Override
+  public int size() {
+    return _columnMetadataMap.size() - _numVirtualColumns;
+  }
+
+  @Override
+  public Iterator<String> iterator() {
+    Iterator<Map.Entry<String, ColumnMetadata>> entries = _columnMetadataMap.entrySet().iterator();
+    return new Iterator<>() {
+      private String _next = advance();
+
+      private String advance() {
+        while (entries.hasNext()) {
+          Map.Entry<String, ColumnMetadata> entry = entries.next();
+          if (isPhysical(entry.getValue())) {
+            return entry.getKey();
+          }
+        }
+        return null;
+      }
+
+      @Override
+      public boolean hasNext() {
+        return _next != null;
+      }
+
+      @Override
+      public String next() {
+        String next = _next;
+        if (next == null) {
+          throw new NoSuchElementException();
+        }
+        _next = advance();
+        return next;
+      }
+    };
+  }
+}

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImpl.java

@@ -59,10 +59,11 @@ public PinotSegmentColumnReaderImpl(IndexSegment indexSegment, String columnName
   ///                              the segment's stored value (which contains the default).
   public PinotSegmentColumnReaderImpl(IndexSegment indexSegment, String columnName,
       boolean skipDefaultNullValues) {
+    // The data source's field spec is the column's own, so this never builds the segment schema (which an immutable
+    // segment derives on demand) and also covers a mutable segment, whose metadata has no column metadata map
     this(new PinotSegmentColumnReader(indexSegment, columnName), columnName,
         indexSegment.getSegmentMetadata().getTotalDocs(),
-        indexSegment.getSegmentMetadata().getSchema().getFieldSpecFor(columnName).getDataType(),
-        skipDefaultNullValues);
+        indexSegment.getDataSource(columnName).getDataSourceMetadata().getDataType(), skipDefaultNullValues);
   }
 
   /// Constructor for subclasses that need to provide their own PinotSegmentColumnReader.

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/virtualcolumn/VirtualColumnProviderFactory.java

@@ -86,26 +86,32 @@ public static VirtualColumnProvider buildProvider(VirtualColumnContext virtualCo
   /// [BuiltInVirtualColumnDefinitions#DEFINITIONS], which the broker side uses as well, so the two can never
   /// disagree on a type.
   /// This method only layers on the provider class, and the constant value for the columns whose value is already
-  /// known here.
+  /// known here (see [#createBuiltInFieldSpec(BuiltInVirtualColumnDefinitions.Definition, String)]).
   public static void addBuiltInVirtualColumnsToSegmentSchema(Schema schema, String segmentName) {
     for (BuiltInVirtualColumnDefinitions.Definition definition : BuiltInVirtualColumnDefinitions.DEFINITIONS) {
-      String column = definition.getName();
-      if (schema.hasColumn(column)) {
-        continue;
+      if (!schema.hasColumn(definition.getName())) {
+        schema.addField(createBuiltInFieldSpec(definition, segmentName));
       }
-      DimensionFieldSpec fieldSpec = definition.createFieldSpec();
-      fieldSpec.setVirtualColumnProvider(getProviderClass(column).getName());
-      // $hostName and $segmentName are constants known at schema construction time, and are carried as the field's
-      // default null value, which DefaultNullValueVirtualColumnProvider reads back.
-      if (BuiltInVirtualColumn.HOSTNAME.equals(column)) {
-        fieldSpec.setDefaultNullValue(NetUtils.getHostnameOrAddress());
-      } else if (BuiltInVirtualColumn.SEGMENTNAME.equals(column)) {
-        fieldSpec.setDefaultNullValue(segmentName);
-      }
-      schema.addField(fieldSpec);
     }
   }
 
+  /// Creates the field spec a segment gets for one built-in virtual column: the shape from `definition`, the provider
+  /// class that produces its values, and for `$hostName` / `$segmentName` the constant value, carried as the field's
+  /// default null value, which `DefaultNullValueVirtualColumnProvider` reads back. Fresh per call: the spec is
+  /// mutable and the `$segmentName` value differs per segment, so it must not be shared across segments.
+  public static DimensionFieldSpec createBuiltInFieldSpec(BuiltInVirtualColumnDefinitions.Definition definition,
+      String segmentName) {
+    String column = definition.getName();
+    DimensionFieldSpec fieldSpec = definition.createFieldSpec();
+    fieldSpec.setVirtualColumnProvider(getProviderClass(column).getName());
+    if (BuiltInVirtualColumn.HOSTNAME.equals(column)) {
+      fieldSpec.setDefaultNullValue(NetUtils.getHostnameOrAddress());
+    } else if (BuiltInVirtualColumn.SEGMENTNAME.equals(column)) {
+      fieldSpec.setDefaultNullValue(segmentName);
+    }
+    return fieldSpec;
+  }
+
   private static Class<? extends VirtualColumnProvider> getProviderClass(String column) {
     Class<? extends VirtualColumnProvider> providerClass = PROVIDER_CLASSES.get(column);
     Preconditions.checkState(providerClass != null, "No virtual column provider registered for built-in column: %s",