Share FieldSpecs and repeated values across segment metadata loads

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

@@ -19,16 +19,23 @@
 package org.apache.pinot.segment.local.segment.index;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 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.local.segment.readers.GenericRowRecordReader;
 import org.apache.pinot.segment.spi.ColumnMetadata;
 import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
 import org.apache.pinot.segment.spi.creator.SegmentIndexCreationDriver;
@@ -37,30 +44,48 @@
 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.spi.config.table.FieldConfig;
+import org.apache.pinot.spi.config.table.OpenStructIndexConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.ComplexFieldSpec;
+import org.apache.pinot.spi.data.DimensionFieldSpec;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.OpenStructNaming;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
 import org.apache.pinot.util.TestUtils;
 import org.testng.Assert;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 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;
 import static org.testng.Assert.assertTrue;
 
 
 public class SegmentMetadataImplTest {
   private static final String AVRO_DATA = "data/test_data-mv.avro";
   private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "SegmentMetadataImplTest");
+  private File _avroFile;
   private File _segmentDirectory;
 
   @BeforeMethod
   public void setUp()
       throws Exception {
     final String filePath =
         TestUtils.getFileFromResourceUrl(SegmentMetadataImplTest.class.getClassLoader().getResource(AVRO_DATA));
+    _avroFile = new File(filePath);
 
     // intentionally changed this to TimeUnit.Hours to make it non-default for testing
     final SegmentGeneratorConfig config = SegmentTestUtils
-        .getSegmentGenSpecWithSchemAndProjectedColumns(new File(filePath), INDEX_DIR, "daysSinceEpoch", TimeUnit.HOURS,
+        .getSegmentGenSpecWithSchemAndProjectedColumns(_avroFile, INDEX_DIR, "daysSinceEpoch", TimeUnit.HOURS,
             "testTable");
     config.setSegmentNamePostfix("1");
     config.setCustomProperties(Map.of("custom.k1", "v1", "custom.k2", "v2"));
@@ -128,6 +153,7 @@ public void testIndexSizesOnlyFromIndexDir()
     assertEquals(fromStreams.getColumnMetadataMap().keySet(), fromDir.getColumnMetadataMap().keySet());
     assertEquals(fromStreams.getTotalDocs(), fromDir.getTotalDocs());
     assertEquals(fromStreams.getSchema(), fromDir.getSchema());
+    assertSame(fromStreams.getTimeColumn(), fromDir.getTimeColumn());
     for (Map.Entry<String, ColumnMetadata> entry : fromDir.getColumnMetadataMap().entrySet()) {
       ColumnMetadata dirColumn = entry.getValue();
       assertTrue(dirColumn.getNumIndexes() > 0, entry.getKey());
@@ -142,4 +168,163 @@ public void testIndexSizesOnlyFromIndexDir()
       assertTrue(((ColumnMetadataImpl) streamColumn).getIndexSizeMap().isEmpty(), entry.getKey());
     }
   }
+
+  /// A server holds one column-metadata graph per loaded segment, so the per-column strings and default null values
+  /// are shared: two loads of the same metadata alias one interned column name (as the map key, the FieldSpec name and
+  /// the Schema entry) and hold the static FieldSpec default constant rather than a box parsed from the literal the
+  /// segment creator wrote, while the specs stay equal to the schema the segment was built from.
+  @Test
+  public void testColumnStringsAndDefaultsSharedAcrossLoads()
+      throws Exception {
+    SegmentMetadataImpl first = new SegmentMetadataImpl(_segmentDirectory);
+    SegmentMetadataImpl second = new SegmentMetadataImpl(_segmentDirectory);
+    assertEquals(first.getTimeColumn(), "daysSinceEpoch");
+    assertSame(first.getTimeColumn(), second.getTimeColumn());
+    assertSame(first.getTimeColumn(), first.getColumnMetadataMap().ceilingKey(first.getTimeColumn()));
+    assertEquals(first.getColumnMetadataMap().keySet(), second.getColumnMetadataMap().keySet());
+    assertSame(first.getColumnMetadataMap().firstKey(), second.getColumnMetadataMap().firstKey());
+    assertSame(first.getSchema().getDimensionNames().get(0), second.getSchema().getDimensionNames().get(0));
+    Iterator<String> secondKeys = second.getColumnMetadataMap().keySet().iterator();
+    for (String column : first.getColumnMetadataMap().keySet()) {
+      assertSame(column, secondKeys.next());
+      FieldSpec fieldSpec = first.getColumnMetadataFor(column).getFieldSpec();
+      assertSame(fieldSpec.getName(), column, column);
+      assertSame(fieldSpec.getName(), second.getColumnMetadataFor(column).getFieldSpec().getName(), column);
+      assertSame(first.getSchema().getFieldSpecFor(column).getName(), column, column);
+      assertSame(fieldSpec.getDefaultNullValue(),
+          FieldSpec.getDefaultNullValue(fieldSpec.getFieldType(), fieldSpec.getDataType(), null), column);
+    }
+    Schema inputSchema = SegmentTestUtils.extractSchemaFromAvroWithoutTime(_avroFile);
+    for (FieldSpec inputFieldSpec : inputSchema.getAllFieldSpecs()) {
+      assertEquals(first.getSchema().getFieldSpecFor(inputFieldSpec.getName()), inputFieldSpec);
+    }
+  }
+
+  /// OPEN_STRUCT children carry an explicit column name and a parent column in `metadata.properties`; both come back
+  /// as the interned instances, so the child's parent name is the very String that keys the parent column.
+  @Test
+  public void testOpenStructChildStringsShared()
+      throws Exception {
+    String parent = "metrics";
+    File segmentDir = buildOpenStructSegment(parent);
+    try {
+      SegmentMetadataImpl first = new SegmentMetadataImpl(segmentDir);
+      SegmentMetadataImpl second = new SegmentMetadataImpl(segmentDir);
+      assertNull(first.getTimeColumn());
+      assertNull(second.getTimeColumn());
+      String parentKey = first.getColumnMetadataMap().ceilingKey(parent);
+      assertEquals(parentKey, parent);
+      assertSame(parentKey, second.getColumnMetadataMap().ceilingKey(parent));
+      String child = OpenStructNaming.materializedColumnName(parent, "cpu");
+      ColumnMetadataImpl firstChild = (ColumnMetadataImpl) first.getColumnMetadataFor(child);
+      ColumnMetadataImpl secondChild = (ColumnMetadataImpl) second.getColumnMetadataFor(child);
+      assertEquals(firstChild.getParentColumn(), parent);
+      assertSame(firstChild.getParentColumn(), parentKey);
+      assertSame(firstChild.getParentColumn(), secondChild.getParentColumn());
+      assertSame(firstChild.getFieldSpec().getName(), secondChild.getFieldSpec().getName());
+      assertSame(firstChild.getFieldSpec().getDefaultNullValue(), FieldSpec.DEFAULT_DIMENSION_NULL_VALUE_OF_DOUBLE);
+      ComplexFieldSpec firstParent = (ComplexFieldSpec) first.getColumnMetadataFor(parent).getFieldSpec();
+      ComplexFieldSpec secondParent = (ComplexFieldSpec) second.getColumnMetadataFor(parent).getFieldSpec();
+      assertEquals(firstParent.getChildFieldSpecs().keySet(), Set.of("views", "cpu", "host"));
+      for (Map.Entry<String, FieldSpec> entry : firstParent.getChildFieldSpecs().entrySet()) {
+        FieldSpec childSpec = entry.getValue();
+        assertSame(childSpec.getName(), entry.getKey());
+        assertSame(childSpec.getName(), secondParent.getChildFieldSpec(entry.getKey()).getName());
+        assertSame(childSpec.getDefaultNullValue(),
+            FieldSpec.getDefaultNullValue(childSpec.getFieldType(), childSpec.getDataType(), null));
+      }
+    } finally {
+      FileUtils.deleteQuietly(segmentDir);
+    }
+  }
+
+  /// 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<>();
+    children.put("views", new DimensionFieldSpec("views", FieldSpec.DataType.LONG, true));
+    children.put("cpu", new DimensionFieldSpec("cpu", FieldSpec.DataType.DOUBLE, true));
+    children.put("host", new DimensionFieldSpec("host", FieldSpec.DataType.STRING, true));
+    Schema schema = new Schema.SchemaBuilder().setSchemaName("testOpenStruct")
+        .addField(new ComplexFieldSpec(parent, FieldSpec.DataType.OPEN_STRUCT, true, children))
+        .addSingleValueDimension("dim", FieldSpec.DataType.STRING)
+        .build();
+    OpenStructIndexConfig openStructConfig =
+        new OpenStructIndexConfig(false, null, 3, Set.of("views", "cpu", "host"), 0.5, List.of(), null);
+    ObjectNode indexes = JsonUtils.newObjectNode();
+    indexes.set("open_struct", JsonUtils.objectToJsonNode(openStructConfig));
+    TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("testOpenStruct")
+        .setFieldConfigList(List.of(new FieldConfig.Builder(parent).withIndexes(indexes).build())).build();
+    SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig, schema);
+    config.setOutDir(new File(INDEX_DIR, "openStruct").getAbsolutePath());
+    config.setSegmentName("openStructSegment");
+    List<GenericRow> rows = new ArrayList<>();
+    for (int i = 0; i < 10; i++) {
+      GenericRow row = new GenericRow();
+      Map<String, Object> metrics = new HashMap<>();
+      metrics.put("views", (long) i);
+      metrics.put("cpu", i * 0.5);
+      metrics.put("host", "host-" + i);
+      row.putValue(parent, metrics);
+      row.putValue("dim", "val-" + i);
+      rows.add(row);
+    }
+    SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
+    driver.init(config, new GenericRowRecordReader(rows));
+    driver.build();
+    return driver.getOutputDirectory();
+  }
 }

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`) may be shared across loaded segments and tables
+  /// whose columns parse to an equal spec, 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 may be shared with other loaded
+  /// segments whose columns parse 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();