Share segment-derived FieldSpec instances across segments through a weak interner

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

@@ -63,6 +63,9 @@
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertSame;
 
 
@@ -228,6 +231,60 @@ public void testOpenStructChildStringsShared()
     }
   }
 
+  /// Every segment of a table parses the same column definitions, so the FieldSpecs are interned: two loads alias one
+  /// instance per column, in the column metadata and in the segment Schema alike, while the Schema object itself stays
+  /// per segment and the specs stay equal to the schema the segment was built from.
+  @Test
+  public void testFieldSpecsSharedAcrossLoads()
+      throws Exception {
+    SegmentMetadataImpl first = new SegmentMetadataImpl(_segmentDirectory);
+    SegmentMetadataImpl second = new SegmentMetadataImpl(_segmentDirectory);
+    assertEquals(first.getSchema(), second.getSchema());
+    assertNotSame(first.getSchema(), second.getSchema());
+    for (String column : first.getColumnMetadataMap().keySet()) {
+      FieldSpec fieldSpec = first.getColumnMetadataFor(column).getFieldSpec();
+      assertSame(second.getColumnMetadataFor(column).getFieldSpec(), fieldSpec, column);
+      assertSame(first.getSchema().getFieldSpecFor(column), fieldSpec, column);
+      assertSame(second.getSchema().getFieldSpecFor(column), fieldSpec, column);
+    }
+    Schema inputSchema = SegmentTestUtils.extractSchemaFromAvroWithoutTime(_avroFile);
+    for (FieldSpec inputFieldSpec : inputSchema.getAllFieldSpecs()) {
+      assertEquals(first.getSchema().getFieldSpecFor(inputFieldSpec.getName()), inputFieldSpec);
+    }
+
+    // Only the specs are shared: removing a column from one segment's schema leaves the other segment intact.
+    String column = first.getColumnMetadataMap().firstKey();
+    first.getSchema().removeField(column);
+    assertNull(first.getSchema().getFieldSpecFor(column));
+    assertNotNull(second.getSchema().getFieldSpecFor(column));
+    assertSame(second.getSchema().getFieldSpecFor(column), second.getColumnMetadataFor(column).getFieldSpec());
+  }
+
+  /// A COMPLEX parent is not interned (ComplexFieldSpec does not override equals, so two structs with different
+  /// children would alias), but its children and the materialized child columns are.
+  @Test
+  public void testOpenStructChildSpecsSharedButParentIsNot()
+      throws Exception {
+    String parent = "metrics";
+    File segmentDir = buildOpenStructSegment(parent);
+    try {
+      SegmentMetadataImpl first = new SegmentMetadataImpl(segmentDir);
+      SegmentMetadataImpl second = new SegmentMetadataImpl(segmentDir);
+      ComplexFieldSpec firstParent = (ComplexFieldSpec) first.getColumnMetadataFor(parent).getFieldSpec();
+      ComplexFieldSpec secondParent = (ComplexFieldSpec) second.getColumnMetadataFor(parent).getFieldSpec();
+      assertNotSame(secondParent, firstParent);
+      assertEquals(secondParent.getChildFieldSpecs(), firstParent.getChildFieldSpecs());
+      for (Map.Entry<String, FieldSpec> entry : firstParent.getChildFieldSpecs().entrySet()) {
+        assertSame(secondParent.getChildFieldSpec(entry.getKey()), entry.getValue(), entry.getKey());
+      }
+      String child = OpenStructNaming.materializedColumnName(parent, "cpu");
+      assertSame(second.getColumnMetadataFor(child).getFieldSpec(), first.getColumnMetadataFor(child).getFieldSpec());
+      assertSame(second.getColumnMetadataFor("dim").getFieldSpec(), first.getColumnMetadataFor("dim").getFieldSpec());
+    } finally {
+      FileUtils.deleteQuietly(segmentDir);
+    }
+  }
+
   private static File buildOpenStructSegment(String parent)
       throws Exception {
     Map<String, FieldSpec> children = new HashMap<>();

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

@@ -25,13 +25,23 @@
 import org.apache.pinot.segment.spi.index.IndexType;
 import org.apache.pinot.spi.annotations.InterfaceAudience;
 import org.apache.pinot.spi.config.table.FieldConfig.EncodingType;
+import org.apache.pinot.spi.data.FieldSpec;
 
 
 /// The `ColumnMetadata` class holds the column level management information and data statistics.
 @InterfaceAudience.Private
 public interface ColumnMetadata extends ColumnShape {
   int UNAVAILABLE = -1;
 
+  /// Returns the [FieldSpec] of the column.
+  ///
+  /// A spec derived from segment metadata (`metadata.properties`) is shared: every loaded segment whose column parses
+  /// to an equal spec, in this table or any other, holds the same instance, so it must be treated as immutable. Never
+  /// call a setter on it; copy it (e.g. through a JSON round-trip) before mutating. Compare specs with
+  /// [FieldSpec#equals], never by identity.
+  @Override
+  FieldSpec getFieldSpec();
+
   /// Returns `true` when the column has a dictionary, `false` otherwise.
   @JsonProperty("hasDictionary")
   boolean hasDictionary();

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

@@ -64,6 +64,12 @@ public interface SegmentMetadata {
 
   SegmentVersion getVersion();
 
+  /// Returns the schema of the segment, one [org.apache.pinot.spi.data.FieldSpec] per column.
+  ///
+  /// The `Schema` object itself belongs to this segment, but the specs it holds are shared with every other loaded
+  /// segment (of this table or any other) whose column parses to an equal spec, and must be treated as immutable: never
+  /// call a setter on one; copy it (e.g. through a JSON round-trip) before mutating. Removing a column from this schema
+  /// does not affect other segments.
   Schema getSchema();
 
   int getTotalDocs();

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

@@ -20,6 +20,8 @@
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Interner;
+import com.google.common.collect.Interners;
 import com.google.common.collect.Maps;
 import it.unimi.dsi.fastutil.ints.IntSet;
 import java.math.BigDecimal;
@@ -64,12 +66,22 @@
 /// the per-column footprint small: column names, parent-column names, date-time formats/granularities and custom
 /// default-null literals are interned (they recur in every segment of a table), and a `defaultNullValue` that equals
 /// the type default is not handed to the [FieldSpec] at all, so the spec carries the shared static
-/// `FieldSpec.DEFAULT_*` constant and never retains the literal. Segment-derived [FieldSpec]s must therefore be
-/// treated as read-only: nothing may mutate their default null value in place (nothing ever did).
+/// `FieldSpec.DEFAULT_*` constant and never retains the literal. The [FieldSpec] itself is then interned through
+/// [#FIELD_SPEC_INTERNER], so every segment of a table (and every table with an identical column definition) shares
+/// one instance per distinct spec instead of retaining its own. Segment-derived [FieldSpec]s must therefore be treated
+/// as immutable: a setter call on one would bleed into every other segment and table that shares it, and would
+/// corrupt the interner's hash bucket (nothing ever mutated one; copy via a JSON round-trip before mutating).
 @SuppressWarnings({"rawtypes", "unchecked"})
 public class ColumnMetadataImpl implements ColumnMetadata {
   private static final long SIZE_MASK = 0xffffffffffffL;
 
+  /// Canonical instances of the [FieldSpec]s parsed from `metadata.properties`, keyed by [FieldSpec#equals] /
+  /// [FieldSpec#hashCode] (name, data type, single-value, default null value, max length, date-time format and
+  /// granularity, ...), so schema evolution yields a distinct canonical instance per version of a column. The specs
+  /// are held weakly: the canonical instance is exactly the one the loaded segments retain, so it lives as long as
+  /// any of them and is released once the last one is unloaded. Thread-safe.
+  private static final Interner<FieldSpec> FIELD_SPEC_INTERNER = Interners.newWeakInterner();
+
   private final FieldSpec _fieldSpec;
   private final int _totalDocs;
   private final int _cardinality;
@@ -504,6 +516,11 @@ private static ChunkCompressionType parseCompressionType(String column, @Nullabl
     }
   }
 
+  /// Parses the [FieldSpec] of the given column. DIMENSION, METRIC, TIME and DATE_TIME specs are returned from
+  /// [#FIELD_SPEC_INTERNER], so the instance is shared with every other segment whose column parses to an equal spec
+  /// and must not be mutated. A COMPLEX spec is not interned: [ComplexFieldSpec] does not override
+  /// [FieldSpec#equals], so two structs with different children would alias; its children are parsed through this
+  /// method and are interned.
   public static FieldSpec extractFieldSpec(String column, PropertiesConfiguration config) {
     // The name is retained by the FieldSpec, the segment Schema and every per-segment column map, and it recurs in
     // every segment of the table: intern it so all of them alias one JVM-wide instance. When COLUMN_NAME is absent
@@ -524,19 +541,20 @@ public static FieldSpec extractFieldSpec(String column, PropertiesConfiguration
         ? FieldSpec.MaxLengthExceedStrategy.valueOf(maxLengthExceedStrategyString) : null;
     switch (fieldType) {
       case DIMENSION:
-        return new DimensionFieldSpec(fieldName, dataType, isSingleValue, maxLength,
-            canonicalDefaultNullValue(fieldType, dataType, defaultNullValueString), maxLengthExceedStrategy);
+        return FIELD_SPEC_INTERNER.intern(new DimensionFieldSpec(fieldName, dataType, isSingleValue, maxLength,
+            canonicalDefaultNullValue(fieldType, dataType, defaultNullValueString), maxLengthExceedStrategy));
       case METRIC:
-        return new MetricFieldSpec(fieldName, dataType,
-            canonicalDefaultNullValue(fieldType, dataType, defaultNullValueString), maxLength, maxLengthExceedStrategy);
+        return FIELD_SPEC_INTERNER.intern(new MetricFieldSpec(fieldName, dataType,
+            canonicalDefaultNullValue(fieldType, dataType, defaultNullValueString), maxLength,
+            maxLengthExceedStrategy));
       case TIME:
         TimeUnit timeUnit = TimeUnit.valueOf(config.getString(Segment.TIME_UNIT, "DAYS").toUpperCase());
-        return new TimeFieldSpec(new TimeGranularitySpec(dataType, timeUnit, fieldName));
+        return FIELD_SPEC_INTERNER.intern(new TimeFieldSpec(new TimeGranularitySpec(dataType, timeUnit, fieldName)));
       case DATE_TIME:
         String format = intern(config.getString(Column.getKeyFor(column, Column.DATETIME_FORMAT)));
         String granularity = intern(config.getString(Column.getKeyFor(column, Column.DATETIME_GRANULARITY)));
-        return new DateTimeFieldSpec(fieldName, dataType, format, granularity,
-            canonicalDefaultNullValue(fieldType, dataType, defaultNullValueString), null);
+        return FIELD_SPEC_INTERNER.intern(new DateTimeFieldSpec(fieldName, dataType, format, granularity,
+            canonicalDefaultNullValue(fieldType, dataType, defaultNullValueString), null));
       case COMPLEX:
         List<String> childFieldNames =
             config.getList(String.class, Column.getKeyFor(column, Column.COMPLEX_CHILD_FIELD_NAMES));
@@ -547,6 +565,7 @@ public static FieldSpec extractFieldSpec(String column, PropertiesConfiguration
                 extractFieldSpec(ComplexFieldSpec.getFullChildName(column, childField), config));
           }
         }
+        // Deliberately not interned (see the method doc): only the children above are shared.
         return new ComplexFieldSpec(fieldName, dataType, true, childFieldSpecs);
       default:
         throw new IllegalStateException("Unsupported field type: " + fieldType);