Preserve non-null comparison column values during segment commit

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java

@@ -1232,6 +1232,11 @@ private static Map<String, Pair<String, ValueAggregator>> fromAggregationConfig(
     return columnNameToAggregator;
   }
 
+  @Override
+  public List<String> getComparisonColumnNames() {
+    return _upsertComparisonColumns;
+  }
+
   private class IndexContainer implements Closeable {
     final FieldSpec _fieldSpec;
     final PartitionFunction _partitionFunction;

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

@@ -51,6 +51,7 @@ public class PinotSegmentRecordReader implements RecordReader {
   private Map<String, PinotSegmentColumnReader> _columnReaderMap;
   private int[] _sortedDocIds;
   private boolean _skipDefaultNullValues;
+  private List<String> _comparisonColumns;
 
   private int _nextDocId = 0;
 
@@ -154,6 +155,7 @@ private void init(IndexSegment indexSegment, boolean destroySegmentOnClose, @Nul
     _indexSegment = indexSegment;
     _destroySegmentOnClose = destroySegmentOnClose;
     _numDocs = _indexSegment.getSegmentMetadata().getTotalDocs();
+    _comparisonColumns = indexSegment.getComparisonColumnNames();
 
     if (_numDocs > 0) {
       _columnReaderMap = new HashMap<>();
@@ -221,6 +223,10 @@ public void getRecord(int docId, GenericRow buffer) {
       PinotSegmentColumnReader columnReader = entry.getValue();
       if (!columnReader.isNull(docId)) {
         buffer.putValue(column, columnReader.getValue(docId));
+      } else if (_comparisonColumns != null && _comparisonColumns.contains(column)) {
+        if (columnReader.getValue(docId) != null) {
+          buffer.putValue(column, columnReader.getValue(docId));
+        }
       } else if (!_skipDefaultNullValues) {
         buffer.putDefaultNullValue(column, columnReader.getValue(docId));
       }

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

@@ -134,4 +134,13 @@ default void release(FetchContext fetchContext) {
    * Destroys segment in memory and closes file handlers if in MMAP mode.
    */
   void destroy();
+
+  /**
+   * Returns the number of rows already indexed into the segment.
+   *
+   * @return The number of rows indexed
+   */
+  default List<String> getComparisonColumnNames() {
+    return null;
+  }
 }